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:
2026-07-01 23:09:35 +08:00
parent 40d8a29e33
commit 6da0972740
4 changed files with 143 additions and 3 deletions

View File

@@ -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];