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:
@@ -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.
|
||||
///
|
||||
|
||||
Reference in New Issue
Block a user