speculative: copy_kv_position primitive for tree drafting KV remap
SGLang-style "write-all, copy-move on acceptance" approach: after tree verification, physically copy an accepted sibling's K/V from its physical cache slot to the canonical sequential position. New CUDA kernel: copy_kv_position_kernel in reshape_and_cache.cu. For one token (src_pos → dst_pos), copies head_dim × num_kv_heads BF16 elements in both K and V pools. Grid = num_kv_heads, block = head_dim. Cost for one token across 36 layers: ~5.3 MB D2D copy @ 900 GB/s = <6μs. Rust FFI: copy_kv_position(k_pool, v_pool, block_ids, src_pos, dst_pos, num_kv_heads, head_dim, block_size, stream). PagedKVCache method: copy_kv_position(slot, src_pos, dst_pos) — uploads block_ids for the sequence, calls the kernel per layer. This is the primitive needed by tree drafting: when a non-primary sibling at cache position P+2 is accepted as the "true" token for target position P+1, call copy_kv_position(slot, P+2, P+1) then truncate to P+2. Next: wire into bench-eagle3 tree drafting loop with top-2 siblings.
This commit is contained in:
@@ -145,6 +145,17 @@ unsafe extern "C" {
|
||||
max_blocks_per_seq: i32,
|
||||
stream: *mut c_void,
|
||||
);
|
||||
fn launch_copy_kv_position(
|
||||
k_pool: *mut c_void,
|
||||
v_pool: *mut c_void,
|
||||
block_ids: *const i32,
|
||||
src_pos: i32,
|
||||
dst_pos: i32,
|
||||
num_kv_heads: i32,
|
||||
head_dim: i32,
|
||||
block_size: i32,
|
||||
stream: *mut c_void,
|
||||
);
|
||||
}
|
||||
|
||||
/// Scatter `[num_kv_heads, num_tokens, head_dim]` BF16 K/V into a paged
|
||||
@@ -231,6 +242,36 @@ pub unsafe fn reshape_and_cache_batched_bf16(
|
||||
}
|
||||
}
|
||||
|
||||
/// Copy one token's K/V from `src_pos` to `dst_pos` within the same sequence's
|
||||
/// paged cache (one layer). Used by tree speculative decoding to remap
|
||||
/// accepted sibling K/V to canonical sequential positions after acceptance.
|
||||
///
|
||||
/// # Safety
|
||||
/// Pool and block_ids pointers must be valid GPU pointers for the given layer.
|
||||
pub unsafe fn copy_kv_position(
|
||||
k_pool_ptr: *mut c_void,
|
||||
v_pool_ptr: *mut c_void,
|
||||
block_ids_gpu: *const i32,
|
||||
src_pos: usize,
|
||||
dst_pos: usize,
|
||||
num_kv_heads: usize,
|
||||
head_dim: usize,
|
||||
block_size: usize,
|
||||
stream: *mut c_void,
|
||||
) {
|
||||
launch_copy_kv_position(
|
||||
k_pool_ptr,
|
||||
v_pool_ptr,
|
||||
block_ids_gpu,
|
||||
src_pos as i32,
|
||||
dst_pos as i32,
|
||||
num_kv_heads as i32,
|
||||
head_dim as i32,
|
||||
block_size as i32,
|
||||
stream,
|
||||
);
|
||||
}
|
||||
|
||||
fn apply_causal_mask(scores: &Tensor, offset: usize) {
|
||||
let ndim = scores.ndim();
|
||||
let rows = scores.shape()[ndim - 2];
|
||||
|
||||
Reference in New Issue
Block a user