kernels: fix sink-decode build (b4db953 didn't compile)

b4db953 committed before a green build: the existing
launch_decode_attention_bf16 call wasn't updated for the new kernel
signature (too-few-args at flash_attention.cu:433) and the Rust FFI
binding + decode_attention_sink wrapper never landed (failed string
matches). This adds the missing pieces; xserv now builds green on dash5
(release, 17s). qwen3 decode path unchanged (passes nullptr/0).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 21:31:02 +08:00
parent 1e8091a111
commit 403879959a
2 changed files with 62 additions and 1 deletions

View File

@@ -22,6 +22,12 @@ unsafe extern "C" {
kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
);
fn launch_decode_attention_sink_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, sinks: *const c_void, window: i32, stream: *mut c_void,
);
fn launch_paged_decode_attention_bf16(
q: *const c_void,
k_cache: *const c_void,
@@ -142,6 +148,38 @@ pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
output
}
/// Decode attention with gpt-oss per-head attention sinks + optional sliding
/// window. `q`:[batch,num_q_heads,1,head_dim], `k`/`v`:[batch,num_kv_heads,kv_len,
/// head_dim], all BF16 contiguous on GPU. `sinks`:[num_q_heads] f32 on GPU.
/// `window`=0 means full attention; otherwise attend only the last `window` keys.
/// `scale` is explicit (gpt-oss: head_dim**-0.5 with head_dim 64).
pub fn decode_attention_sink(
q: &Tensor, k: &Tensor, v: &Tensor, sinks: &Tensor, scale: f32, window: usize,
) -> Tensor {
assert_eq!(q.ndim(), 4);
assert_eq!(q.shape()[2], 1, "decode_attention_sink 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 output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
unsafe {
launch_decode_attention_sink_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, sinks.data_ptr() as *const c_void, window as i32,
std::ptr::null_mut(),
);
}
output
}
/// Flash Attention 2 — O(1) extra memory, supports GQA natively.
/// Auto-dispatches to decode_attention when q_len == 1.
///