Gahow Wang
6da0972740
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.
683 lines
20 KiB
Rust
683 lines
20 KiB
Rust
use std::ffi::c_void;
|
||
use xserv_tensor::{DType, Tensor};
|
||
|
||
use crate::activation::scale;
|
||
use crate::gemm::batched_matmul;
|
||
use crate::softmax::softmax;
|
||
|
||
unsafe extern "C" {
|
||
fn launch_causal_mask_f32(
|
||
scores: *mut c_void,
|
||
batch: i32,
|
||
rows: i32,
|
||
cols: i32,
|
||
offset: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_causal_mask_bf16(
|
||
scores: *mut c_void,
|
||
batch: i32,
|
||
rows: i32,
|
||
cols: i32,
|
||
offset: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_flash_attention_bf16(
|
||
q: *const c_void,
|
||
k: *const c_void,
|
||
v: *const c_void,
|
||
o: *mut c_void,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
q_len: i32,
|
||
kv_len: i32,
|
||
head_dim: i32,
|
||
scale: f32,
|
||
causal: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_flash_attention_sinks_bf16(
|
||
q: *const c_void,
|
||
k: *const c_void,
|
||
v: *const c_void,
|
||
o: *mut c_void,
|
||
sinks: *const c_void,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
q_len: i32,
|
||
kv_len: i32,
|
||
head_dim: i32,
|
||
scale: f32,
|
||
causal: i32,
|
||
window_size: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_decode_attention_bf16(
|
||
q: *const c_void,
|
||
k: *const c_void,
|
||
v: *const c_void,
|
||
o: *mut c_void,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
kv_len: i32,
|
||
head_dim: i32,
|
||
scale: f32,
|
||
causal: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_paged_decode_attention_bf16(
|
||
q: *const c_void,
|
||
k_cache: *const c_void,
|
||
v_cache: *const c_void,
|
||
o: *mut c_void,
|
||
block_tables: *const i32,
|
||
context_lens: *const i32,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
head_dim: i32,
|
||
max_blocks_per_seq: i32,
|
||
scale: f32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_paged_decode_attention_tree_bf16(
|
||
q: *const c_void,
|
||
k_cache: *const c_void,
|
||
v_cache: *const c_void,
|
||
o: *mut c_void,
|
||
block_tables: *const i32,
|
||
context_lens: *const i32,
|
||
tree_mask: *const i32,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
head_dim: i32,
|
||
max_blocks_per_seq: i32,
|
||
tree_start: i32,
|
||
tree_len: i32,
|
||
scale: f32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_paged_decode_attention_sinks_bf16(
|
||
q: *const c_void,
|
||
k_cache: *const c_void,
|
||
v_cache: *const c_void,
|
||
o: *mut c_void,
|
||
block_tables: *const i32,
|
||
context_lens: *const i32,
|
||
sinks: *const c_void,
|
||
batch: i32,
|
||
num_q_heads: i32,
|
||
num_kv_heads: i32,
|
||
head_dim: i32,
|
||
max_blocks_per_seq: i32,
|
||
scale: f32,
|
||
window_size: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_reshape_and_cache_bf16(
|
||
k_src: *const c_void,
|
||
v_src: *const c_void,
|
||
k_pool: *mut c_void,
|
||
v_pool: *mut c_void,
|
||
block_ids: *const c_void,
|
||
num_tokens: i32,
|
||
num_heads: i32,
|
||
head_dim: i32,
|
||
start_pos: i32,
|
||
block_size: i32,
|
||
stream: *mut c_void,
|
||
);
|
||
fn launch_reshape_and_cache_batched_bf16(
|
||
k_src: *const c_void,
|
||
v_src: *const c_void,
|
||
k_pool: *mut c_void,
|
||
v_pool: *mut c_void,
|
||
block_tables: *const c_void,
|
||
kv_lens: *const c_void,
|
||
batch: i32,
|
||
num_heads: i32,
|
||
head_dim: i32,
|
||
block_size: i32,
|
||
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
|
||
/// pool for a single sequence whose block table lives at `block_ids_gpu`
|
||
/// (int32, on device).
|
||
///
|
||
/// `k_pool_ptr`/`v_pool_ptr` point to one layer's pool, of logical shape
|
||
/// `[num_blocks_total, num_kv_heads, block_size, head_dim]`.
|
||
///
|
||
/// All pointers must be on the same GPU as the launching context.
|
||
///
|
||
/// # Safety
|
||
/// Pointers must be valid GPU pointers with the documented layouts.
|
||
/// `block_ids_gpu` must contain at least `(start_pos + num_tokens + block_size - 1) / block_size`
|
||
/// valid physical block ids.
|
||
pub unsafe fn reshape_and_cache_bf16(
|
||
k_src: *const c_void,
|
||
v_src: *const c_void,
|
||
k_pool_ptr: *mut c_void,
|
||
v_pool_ptr: *mut c_void,
|
||
block_ids_gpu: *const i32,
|
||
num_tokens: usize,
|
||
num_heads: usize,
|
||
head_dim: usize,
|
||
start_pos: usize,
|
||
block_size: usize,
|
||
stream: *mut c_void,
|
||
) {
|
||
unsafe {
|
||
launch_reshape_and_cache_bf16(
|
||
k_src,
|
||
v_src,
|
||
k_pool_ptr,
|
||
v_pool_ptr,
|
||
block_ids_gpu as *const c_void,
|
||
num_tokens as i32,
|
||
num_heads as i32,
|
||
head_dim as i32,
|
||
start_pos as i32,
|
||
block_size as i32,
|
||
stream,
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Batched scatter for the multi-sequence decode step. Reads
|
||
/// `block_tables` (`[batch, max_blocks_per_seq]` int32 — same buffer the
|
||
/// paged-attention kernel reads) and `kv_lens` (`[batch]` int32, current
|
||
/// seq_len + 1 — i.e., the index of the just-written token + 1) so the
|
||
/// caller doesn't need a separate per-step upload of block ids.
|
||
///
|
||
/// # Safety
|
||
/// All pointers must be on the same GPU. `block_tables` and `kv_lens` must
|
||
/// already be synced to the device for the active batch.
|
||
pub unsafe fn reshape_and_cache_batched_bf16(
|
||
k_src: *const c_void,
|
||
v_src: *const c_void,
|
||
k_pool_ptr: *mut c_void,
|
||
v_pool_ptr: *mut c_void,
|
||
block_tables_gpu: *const i32,
|
||
kv_lens_gpu: *const i32,
|
||
batch: usize,
|
||
num_heads: usize,
|
||
head_dim: usize,
|
||
block_size: usize,
|
||
max_blocks_per_seq: usize,
|
||
stream: *mut c_void,
|
||
) {
|
||
unsafe {
|
||
launch_reshape_and_cache_batched_bf16(
|
||
k_src,
|
||
v_src,
|
||
k_pool_ptr,
|
||
v_pool_ptr,
|
||
block_tables_gpu as *const c_void,
|
||
kv_lens_gpu as *const c_void,
|
||
batch as i32,
|
||
num_heads as i32,
|
||
head_dim as i32,
|
||
block_size as i32,
|
||
max_blocks_per_seq as i32,
|
||
stream,
|
||
);
|
||
}
|
||
}
|
||
|
||
/// 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];
|
||
let cols = scores.shape()[ndim - 1];
|
||
let batch: usize = scores.shape()[..ndim - 2].iter().product();
|
||
|
||
unsafe {
|
||
match scores.dtype() {
|
||
DType::F32 => launch_causal_mask_f32(
|
||
scores.data_ptr() as *mut c_void,
|
||
batch as i32,
|
||
rows as i32,
|
||
cols as i32,
|
||
offset as i32,
|
||
xserv_cuda::current_stream_raw(),
|
||
),
|
||
DType::BF16 => launch_causal_mask_bf16(
|
||
scores.data_ptr() as *mut c_void,
|
||
batch as i32,
|
||
rows as i32,
|
||
cols as i32,
|
||
offset as i32,
|
||
xserv_cuda::current_stream_raw(),
|
||
),
|
||
_ => panic!("unsupported dtype for causal mask"),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Multi-head attention (naive, materializes S×S score matrix).
|
||
///
|
||
/// q, k, v: [batch, num_heads, seq_len, head_dim] — contiguous, on GPU
|
||
/// Returns: [batch, num_heads, seq_len, head_dim]
|
||
pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(k.ndim(), 4);
|
||
assert_eq!(v.ndim(), 4);
|
||
assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
|
||
|
||
let batch = q.shape()[0];
|
||
let num_heads = q.shape()[1];
|
||
let q_len = q.shape()[2];
|
||
let head_dim = q.shape()[3];
|
||
let kv_len = k.shape()[2];
|
||
|
||
assert_eq!(k.shape(), &[batch, num_heads, kv_len, head_dim]);
|
||
assert_eq!(v.shape(), &[batch, num_heads, kv_len, head_dim]);
|
||
|
||
// scores = Q @ K^T → [B, H, q_len, kv_len]
|
||
let k_t = k.transpose(2, 3).contiguous();
|
||
let scores = batched_matmul(q, &k_t);
|
||
|
||
// Scale by 1/sqrt(head_dim)
|
||
let scale_factor = 1.0 / (head_dim as f32).sqrt();
|
||
let scaled_scores = scale(&scores, scale_factor);
|
||
|
||
// Causal mask
|
||
if causal {
|
||
let offset = kv_len - q_len;
|
||
apply_causal_mask(&scaled_scores, offset);
|
||
}
|
||
|
||
// Softmax
|
||
let weights = softmax(&scaled_scores);
|
||
|
||
// output = weights @ V → [B, H, q_len, head_dim]
|
||
batched_matmul(&weights, v)
|
||
}
|
||
|
||
/// Decode Attention — optimized for single-token decode (q_len=1).
|
||
///
|
||
/// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU
|
||
/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
|
||
/// v: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
|
||
///
|
||
/// Returns: [batch, num_q_heads, 1, head_dim] BF16
|
||
pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(q.shape()[2], 1, "decode_attention requires q_len == 1");
|
||
|
||
let batch = q.shape()[0];
|
||
let num_q_heads = q.shape()[1];
|
||
let head_dim = q.shape()[3];
|
||
let num_kv_heads = k.shape()[1];
|
||
let kv_len = k.shape()[2];
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
|
||
|
||
unsafe {
|
||
launch_decode_attention_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k.data_ptr() as *const c_void,
|
||
v.data_ptr() as *const c_void,
|
||
output.data_ptr() as *mut c_void,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
kv_len as i32,
|
||
head_dim as i32,
|
||
scale,
|
||
1, // causal (always 1 for decode)
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Flash Attention 2 — O(1) extra memory, supports GQA natively.
|
||
/// Auto-dispatches to decode_attention when q_len == 1.
|
||
///
|
||
/// q: [batch, num_q_heads, q_len, head_dim] BF16, contiguous, GPU
|
||
/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
|
||
/// v: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
|
||
///
|
||
/// Returns: [batch, num_q_heads, q_len, head_dim] BF16
|
||
pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(k.ndim(), 4);
|
||
assert_eq!(v.ndim(), 4);
|
||
assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
|
||
assert_eq!(q.dtype(), DType::BF16, "flash_attention requires BF16");
|
||
assert_eq!(k.dtype(), DType::BF16);
|
||
assert_eq!(v.dtype(), DType::BF16);
|
||
|
||
let batch = q.shape()[0];
|
||
let num_q_heads = q.shape()[1];
|
||
let q_len = q.shape()[2];
|
||
let head_dim = q.shape()[3];
|
||
let num_kv_heads = k.shape()[1];
|
||
let kv_len = k.shape()[2];
|
||
|
||
assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
|
||
assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
|
||
assert!(
|
||
num_q_heads % num_kv_heads == 0,
|
||
"num_q_heads must be divisible by num_kv_heads"
|
||
);
|
||
assert!(
|
||
head_dim <= 128,
|
||
"flash_attention supports head_dim up to 128"
|
||
);
|
||
|
||
// Dispatch to specialized decode kernel for single-token generation
|
||
if q_len == 1 {
|
||
return decode_attention(q, k, v);
|
||
}
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(
|
||
&[batch, num_q_heads, q_len, head_dim],
|
||
DType::BF16,
|
||
q.device(),
|
||
);
|
||
|
||
unsafe {
|
||
launch_flash_attention_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k.data_ptr() as *const c_void,
|
||
v.data_ptr() as *const c_void,
|
||
output.data_ptr() as *mut c_void,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
q_len as i32,
|
||
kv_len as i32,
|
||
head_dim as i32,
|
||
scale,
|
||
if causal { 1 } else { 0 },
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Flash attention for prefill with gpt-oss attention sinks + optional sliding window.
|
||
///
|
||
/// Same layout/contract as `flash_attention`, plus a per-head `sinks` tensor
|
||
/// ([num_q_heads] BF16, GPU) folded into the softmax denominator, and a
|
||
/// `window_size` (0 = full causal, >0 = sliding window). Always causal.
|
||
pub fn flash_attention_sinks(
|
||
q: &Tensor,
|
||
k: &Tensor,
|
||
v: &Tensor,
|
||
sinks: &Tensor,
|
||
window_size: usize,
|
||
) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(k.ndim(), 4);
|
||
assert_eq!(v.ndim(), 4);
|
||
assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
|
||
assert_eq!(q.dtype(), DType::BF16);
|
||
assert_eq!(k.dtype(), DType::BF16);
|
||
assert_eq!(v.dtype(), DType::BF16);
|
||
|
||
let batch = q.shape()[0];
|
||
let num_q_heads = q.shape()[1];
|
||
let q_len = q.shape()[2];
|
||
let head_dim = q.shape()[3];
|
||
let num_kv_heads = k.shape()[1];
|
||
let kv_len = k.shape()[2];
|
||
|
||
assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
|
||
assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
|
||
assert!(num_q_heads % num_kv_heads == 0);
|
||
assert!(head_dim <= 128);
|
||
assert_eq!(
|
||
sinks.shape()[0],
|
||
num_q_heads,
|
||
"sinks must have num_q_heads entries"
|
||
);
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(
|
||
&[batch, num_q_heads, q_len, head_dim],
|
||
DType::BF16,
|
||
q.device(),
|
||
);
|
||
|
||
unsafe {
|
||
launch_flash_attention_sinks_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k.data_ptr() as *const c_void,
|
||
v.data_ptr() as *const c_void,
|
||
output.data_ptr() as *mut c_void,
|
||
sinks.data_ptr() as *const c_void,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
q_len as i32,
|
||
kv_len as i32,
|
||
head_dim as i32,
|
||
scale,
|
||
1, // always causal
|
||
window_size as i32,
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Paged decode attention.
|
||
///
|
||
/// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU
|
||
/// k_cache_ptr / v_cache_ptr: pointers to [num_blocks, num_kv_heads, BLOCK_SIZE, head_dim] BF16 pools
|
||
/// block_tables_ptr: i32 [batch, max_blocks_per_seq] (rows already arranged for this batch)
|
||
/// context_lens_ptr: i32 [batch]
|
||
///
|
||
/// Returns: [batch, num_q_heads, 1, head_dim] BF16
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn paged_decode_attention(
|
||
q: &Tensor,
|
||
k_cache_ptr: *const c_void,
|
||
v_cache_ptr: *const c_void,
|
||
block_tables_ptr: *const i32,
|
||
context_lens_ptr: *const i32,
|
||
batch: usize,
|
||
num_q_heads: usize,
|
||
num_kv_heads: usize,
|
||
head_dim: usize,
|
||
max_blocks_per_seq: usize,
|
||
) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(
|
||
q.shape()[2],
|
||
1,
|
||
"paged_decode_attention requires q_len == 1"
|
||
);
|
||
assert_eq!(q.dtype(), DType::BF16);
|
||
assert!(
|
||
num_q_heads % num_kv_heads == 0,
|
||
"GQA: num_q_heads must be divisible by num_kv_heads"
|
||
);
|
||
assert!(head_dim <= 128);
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
|
||
|
||
unsafe {
|
||
launch_paged_decode_attention_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k_cache_ptr,
|
||
v_cache_ptr,
|
||
output.data_ptr() as *mut c_void,
|
||
block_tables_ptr,
|
||
context_lens_ptr,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
head_dim as i32,
|
||
max_blocks_per_seq as i32,
|
||
scale,
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Tree-aware paged decode attention. Adds a per-query attention mask over
|
||
/// the newly-written K/V region `[tree_start, tree_start+tree_len)`. Query i
|
||
/// attends to position tree_start+j iff tree_mask[i, j] != 0. Positions <
|
||
/// tree_start are always attended.
|
||
///
|
||
/// Used by speculative decoding with tree drafting to let sibling candidates
|
||
/// share position slots without seeing each other's K/V.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn paged_decode_attention_tree(
|
||
q: &Tensor,
|
||
k_cache_ptr: *const c_void,
|
||
v_cache_ptr: *const c_void,
|
||
block_tables_ptr: *const i32,
|
||
context_lens_ptr: *const i32,
|
||
tree_mask_ptr: *const i32,
|
||
batch: usize,
|
||
num_q_heads: usize,
|
||
num_kv_heads: usize,
|
||
head_dim: usize,
|
||
max_blocks_per_seq: usize,
|
||
tree_start: usize,
|
||
tree_len: usize,
|
||
) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(q.shape()[2], 1);
|
||
assert_eq!(q.dtype(), DType::BF16);
|
||
assert!(num_q_heads % num_kv_heads == 0);
|
||
assert!(head_dim <= 128);
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
|
||
|
||
unsafe {
|
||
launch_paged_decode_attention_tree_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k_cache_ptr,
|
||
v_cache_ptr,
|
||
output.data_ptr() as *mut c_void,
|
||
block_tables_ptr,
|
||
context_lens_ptr,
|
||
tree_mask_ptr,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
head_dim as i32,
|
||
max_blocks_per_seq as i32,
|
||
tree_start as i32,
|
||
tree_len as i32,
|
||
scale,
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Paged decode attention with attention sinks and optional sliding window.
|
||
///
|
||
/// sinks_ptr: pointer to [num_q_heads] BF16 on GPU (or null for no sinks)
|
||
/// window_size: 0 = full attention, >0 = sliding window
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn paged_decode_attention_sinks(
|
||
q: &Tensor,
|
||
k_cache_ptr: *const c_void,
|
||
v_cache_ptr: *const c_void,
|
||
block_tables_ptr: *const i32,
|
||
context_lens_ptr: *const i32,
|
||
sinks_ptr: *const c_void,
|
||
batch: usize,
|
||
num_q_heads: usize,
|
||
num_kv_heads: usize,
|
||
head_dim: usize,
|
||
max_blocks_per_seq: usize,
|
||
window_size: usize,
|
||
) -> Tensor {
|
||
assert_eq!(q.ndim(), 4);
|
||
assert_eq!(q.shape()[2], 1);
|
||
assert_eq!(q.dtype(), DType::BF16);
|
||
assert!(num_q_heads % num_kv_heads == 0);
|
||
assert!(head_dim <= 128);
|
||
|
||
let scale = 1.0 / (head_dim as f32).sqrt();
|
||
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
|
||
|
||
unsafe {
|
||
launch_paged_decode_attention_sinks_bf16(
|
||
q.data_ptr() as *const c_void,
|
||
k_cache_ptr,
|
||
v_cache_ptr,
|
||
output.data_ptr() as *mut c_void,
|
||
block_tables_ptr,
|
||
context_lens_ptr,
|
||
sinks_ptr,
|
||
batch as i32,
|
||
num_q_heads as i32,
|
||
num_kv_heads as i32,
|
||
head_dim as i32,
|
||
max_blocks_per_seq as i32,
|
||
scale,
|
||
window_size as i32,
|
||
xserv_cuda::current_stream_raw(),
|
||
);
|
||
}
|
||
|
||
output
|
||
}
|