diff --git a/crates/xserv-kernels/src/attention.rs b/crates/xserv-kernels/src/attention.rs index 4cd0e5a..3da3e03 100644 --- a/crates/xserv-kernels/src/attention.rs +++ b/crates/xserv-kernels/src/attention.rs @@ -16,6 +16,13 @@ unsafe extern "C" { 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, @@ -295,6 +302,65 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens 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, + std::ptr::null_mut(), + ); + } + + output +} + /// Paged decode attention. /// /// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU diff --git a/crates/xserv-kernels/src/lib.rs b/crates/xserv-kernels/src/lib.rs index df4933c..1b1fd80 100644 --- a/crates/xserv-kernels/src/lib.rs +++ b/crates/xserv-kernels/src/lib.rs @@ -13,7 +13,7 @@ pub mod transpose; pub use activation::{add, gelu, gpt_oss_glu, mul, scale, silu, silu_mul}; pub use argmax::{argmax_bf16_single, argmax_bf16_to_host}; pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu}; -pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention, paged_decode_attention_sinks, reshape_and_cache_bf16, reshape_and_cache_batched_bf16}; +pub use attention::{attention, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention, paged_decode_attention_sinks, reshape_and_cache_bf16, reshape_and_cache_batched_bf16}; pub use embedding::embedding; pub use gemm::{batched_matmul, matmul, GemmBackend}; pub use layernorm::layernorm; diff --git a/crates/xserv-model/src/gpt_oss.rs b/crates/xserv-model/src/gpt_oss.rs index d50540f..ce1d007 100644 --- a/crates/xserv-model/src/gpt_oss.rs +++ b/crates/xserv-model/src/gpt_oss.rs @@ -373,9 +373,8 @@ impl GptOss { paged_cache.append_tokens(slot, layer_idx, &k, &v, new_tokens, pos_offset); let (k_full, v_full) = paged_cache.gather_kv_contiguous(slot, layer_idx); - // Flash attention for prefill (sinks handled post-hoc for simplicity) - // TODO: integrate sinks into flash attention for exact match - let attn_out = flash_attention(&q, &k_full, &v_full, true); + // Flash attention with gpt-oss sinks + (per-layer) sliding window. + let attn_out = flash_attention_sinks(&q, &k_full, &v_full, &layer.sinks, layer.window_size); let attn_merged = merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim); let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt); diff --git a/csrc/attention/flash_attention.cu b/csrc/attention/flash_attention.cu index b6c3c51..a891122 100644 --- a/csrc/attention/flash_attention.cu +++ b/csrc/attention/flash_attention.cu @@ -197,6 +197,172 @@ __global__ void flash_attention_bf16_kernel( } } +// Flash Attention 2 forward with gpt-oss attention sinks + optional sliding window. +// Identical to flash_attention_bf16_kernel, plus: +// - sinks: [num_q_heads] BF16 — a per-head extra softmax logit (no value), +// folded into the denominator after the K/V tiles (exactly as the decode +// sink kernel does). +// - window_size > 0: sliding-window mask. Query at global position p attends +// to keys k with p - window_size < k <= p (matches HF gpt-oss). +__global__ void flash_attention_sinks_bf16_kernel( + const __nv_bfloat16* __restrict__ Q, + const __nv_bfloat16* __restrict__ K, + const __nv_bfloat16* __restrict__ V, + __nv_bfloat16* __restrict__ O, + const __nv_bfloat16* __restrict__ sinks, // [num_q_heads] or NULL + int num_q_heads, int num_kv_heads, + int q_len, int kv_len, int head_dim, + float scale, int causal, int window_size +) { + int q_tile_idx = blockIdx.x; + int bh = blockIdx.y; + int batch_idx = bh / num_q_heads; + int q_head = bh % num_q_heads; + + int heads_per_group = num_q_heads / num_kv_heads; + int kv_head = q_head / heads_per_group; + + int q_tile_start = q_tile_idx * BR; + if (q_tile_start >= q_len) return; + int q_tile_rows = min(BR, q_len - q_tile_start); + + const __nv_bfloat16* Q_head = Q + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim; + const __nv_bfloat16* K_head = K + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim; + const __nv_bfloat16* V_head = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim; + __nv_bfloat16* O_head = O + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim; + + int tid = threadIdx.x; + + extern __shared__ __nv_bfloat16 smem[]; + __nv_bfloat16* smem_q = smem; + __nv_bfloat16* smem_kv = smem + BR * head_dim; + + int q_elems = q_tile_rows * head_dim; + for (int i = tid; i < q_elems; i += THREADS_PER_BLOCK) { + int row = i / head_dim; + int col = i % head_dim; + smem_q[row * head_dim + col] = Q_head[(q_tile_start + row) * head_dim + col]; + } + for (int i = q_elems + tid; i < BR * head_dim; i += THREADS_PER_BLOCK) { + smem_q[i] = __float2bfloat16(0.0f); + } + __syncthreads(); + + bool owns_row = (tid < q_tile_rows); + + float O_acc[128]; + float m_val = -INFINITY; + float l_val = 0.0f; + if (owns_row) { + for (int d = 0; d < head_dim; d++) O_acc[d] = 0.0f; + } + + int kv_offset = kv_len - q_len; + int num_kv_tiles = (kv_len + BC - 1) / BC; + + for (int j = 0; j < num_kv_tiles; j++) { + int kv_tile_start = j * BC; + int kv_tile_cols = min(BC, kv_len - kv_tile_start); + + if (causal) { + int max_allowed_kv = (q_tile_start + q_tile_rows - 1) + kv_offset; + if (kv_tile_start > max_allowed_kv) continue; + } + + int kv_elems = kv_tile_cols * head_dim; + for (int i = tid; i < kv_elems; i += THREADS_PER_BLOCK) { + int row = i / head_dim; + int col = i % head_dim; + smem_kv[row * head_dim + col] = K_head[(kv_tile_start + row) * head_dim + col]; + } + for (int i = kv_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) { + smem_kv[i] = __float2bfloat16(0.0f); + } + __syncthreads(); + + float P[BC]; + + if (owns_row) { + float row_max = -INFINITY; + int q_pos = q_tile_start + tid + kv_offset; // global query position + for (int c = 0; c < kv_tile_cols; c++) { + float dot = 0.0f; + for (int d = 0; d < head_dim; d++) { + dot += __bfloat162float(smem_q[tid * head_dim + d]) + * __bfloat162float(smem_kv[c * head_dim + d]); + } + float s = dot * scale; + + int kv_pos = kv_tile_start + c; + if (causal && kv_pos > q_pos) { + s = -INFINITY; + } + // Sliding window: drop keys older than the window. + if (window_size > 0 && kv_pos <= q_pos - window_size) { + s = -INFINITY; + } + + P[c] = s; + row_max = fmaxf(row_max, s); + } + + float m_new = fmaxf(m_val, row_max); + float psum = 0.0f; + for (int c = 0; c < kv_tile_cols; c++) { + P[c] = expf(P[c] - m_new); + psum += P[c]; + } + float correction = expf(m_val - m_new); + l_val = correction * l_val + psum; + for (int d = 0; d < head_dim; d++) O_acc[d] *= correction; + m_val = m_new; + } + + __syncthreads(); + + int v_elems = kv_tile_cols * head_dim; + for (int i = tid; i < v_elems; i += THREADS_PER_BLOCK) { + int row = i / head_dim; + int col = i % head_dim; + smem_kv[row * head_dim + col] = V_head[(kv_tile_start + row) * head_dim + col]; + } + for (int i = v_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) { + smem_kv[i] = __float2bfloat16(0.0f); + } + __syncthreads(); + + if (owns_row) { + for (int c = 0; c < kv_tile_cols; c++) { + float p = P[c]; + if (p != 0.0f) { + for (int d = 0; d < head_dim; d++) { + O_acc[d] += p * __bfloat162float(smem_kv[c * head_dim + d]); + } + } + } + } + __syncthreads(); + } + + // Fold in the per-head attention sink (extra logit, no value contribution). + if (owns_row && sinks != nullptr) { + float sink_logit = __bfloat162float(sinks[q_head]); + float m_new = fmaxf(m_val, sink_logit); + float correction = expf(m_val - m_new); + l_val = correction * l_val + expf(sink_logit - m_new); + for (int d = 0; d < head_dim; d++) O_acc[d] *= correction; + m_val = m_new; + } + + if (owns_row) { + float inv_l = (l_val > 0.0f) ? (1.0f / l_val) : 0.0f; + int global_row = q_tile_start + tid; + for (int d = 0; d < head_dim; d++) { + O_head[global_row * head_dim + d] = __float2bfloat16(O_acc[d] * inv_l); + } + } +} + // ============================================================ // Decode Attention kernel: optimized for Q_len=1 (single-token decode). // Parallelizes across KV sequence dimension instead of Q rows. @@ -395,6 +561,31 @@ void launch_flash_attention_bf16( CUDA_CHECK_LAST_ERROR(); } +void launch_flash_attention_sinks_bf16( + const void* Q, const void* K, const void* V, void* O, + const void* sinks, + int batch, int num_q_heads, int num_kv_heads, + int q_len, int kv_len, int head_dim, + float scale, int causal, int window_size, void* stream +) { + int q_tiles = (q_len + BR - 1) / BR; + dim3 grid(q_tiles, batch * num_q_heads); + int block = THREADS_PER_BLOCK; + int smem_bytes = (BR + BC) * head_dim * (int)sizeof(__nv_bfloat16); + + flash_attention_sinks_bf16_kernel<<>>( + (const __nv_bfloat16*)Q, + (const __nv_bfloat16*)K, + (const __nv_bfloat16*)V, + (__nv_bfloat16*)O, + (const __nv_bfloat16*)sinks, + num_q_heads, num_kv_heads, + q_len, kv_len, head_dim, + scale, causal, window_size + ); + CUDA_CHECK_LAST_ERROR(); +} + void launch_decode_attention_bf16( const void* Q, const void* K, const void* V, void* O, int batch, int num_q_heads, int num_kv_heads,