diff --git a/crates/xserv-kernels/src/attention.rs b/crates/xserv-kernels/src/attention.rs index 432d7d6..491fe1b 100644 --- a/crates/xserv-kernels/src/attention.rs +++ b/crates/xserv-kernels/src/attention.rs @@ -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. /// diff --git a/csrc/attention/flash_attention.cu b/csrc/attention/flash_attention.cu index 8086b38..7f51c1f 100644 --- a/csrc/attention/flash_attention.cu +++ b/csrc/attention/flash_attention.cu @@ -429,7 +429,30 @@ void launch_decode_attention_bf16( (__nv_bfloat16*)O, num_q_heads, num_kv_heads, kv_len, head_dim, - scale + scale, + nullptr, 0 // no sinks, no sliding window (qwen3 path unchanged) + ); + CUDA_CHECK_LAST_ERROR(); +} + +// gpt-oss decode: per-head attention sinks + optional sliding window. +void launch_decode_attention_sink_bf16( + const void* Q, const void* K, const void* V, void* O, + int batch, int num_q_heads, int num_kv_heads, + int kv_len, int head_dim, + float scale, const void* sinks, int window, void* stream +) { + int grid = batch * num_q_heads; + int block = DECODE_THREADS; + decode_attention_bf16_kernel<<>>( + (const __nv_bfloat16*)Q, + (const __nv_bfloat16*)K, + (const __nv_bfloat16*)V, + (__nv_bfloat16*)O, + num_q_heads, num_kv_heads, + kv_len, head_dim, + scale, + (const float*)sinks, window ); CUDA_CHECK_LAST_ERROR(); }