bench-gpt-oss: teacher-forced diagnostics + --prompt flag

Add --prompt to override the fixed prompt, and two teacher-forced
diagnostics: --forced runs prefill over prompt+oracle ids and reports
per-position top-1 agreement; --forced-decode walks the oracle trajectory
through the decode path with per-position agreement bucketed by position,
to localize long-context decode divergence from the reference.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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

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;

crates/xserv-model/src/bin/bench-gpt-oss.rs

@@ -68,7 +68,8 @@ fn main() {
     eprintln!("[rank 0] Ready.");
 
     // Prompt
-    let prompt = "What is the meaning of life?";
+    let prompt_arg = get_arg::<String>(&args, "--prompt");
+    let prompt = prompt_arg.as_deref().unwrap_or("What is the meaning of life?");
     let token_ids = tokenizer.encode(prompt);
     eprintln!("Prompt ({} tokens): {prompt}", token_ids.len());
 
@@ -77,6 +78,84 @@ fn main() {
     cache.register_sequence(slot).unwrap();
     broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Register(slot));
 
+    // Teacher-forced diagnostic: prefill (prompt + forced ids) in one shot and
+    // report, for each forced position, whether xserv's argmax == the forced
+    // (oracle) next token. Removes free-running compounding so it isolates
+    // whether per-position logits agree with the llama.cpp trajectory.
+    if let Some(forced) = get_arg::<String>(&args, "--forced") {
+        let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
+        let mut seq = token_ids.clone();
+        seq.extend_from_slice(&forced_ids);
+        // Workers must run the same prefill in lockstep (TP AllReduces match up).
+        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: seq.clone(), slot });
+        let logits = model.forward_prefill_paged(&seq, slot, &mut cache);
+        wait_workers(&worker_handles);
+        let logits_cpu = logits.to_device(Device::Cpu);
+        let vocab = logits.shape()[1];
+        let data = logits_cpu.as_slice::<half::bf16>();
+        let plen = token_ids.len();
+        let mut matches = 0usize;
+        let mut total = 0usize;
+        // position i predicts seq[i+1]; we check the forced region
+        for i in (plen - 1)..(seq.len() - 1) {
+            let row = &data[i * vocab..(i + 1) * vocab];
+            let argmax = row.iter().enumerate()
+                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+                .map(|(j, _)| j as u32).unwrap();
+            let expected = seq[i + 1];
+            let ok = argmax == expected;
+            if ok { matches += 1; }
+            total += 1;
+            eprintln!("pos {i}: xserv_argmax={argmax} oracle={expected} {}", if ok {"OK"} else {"DIFF"});
+        }
+        eprintln!("\nTeacher-forced top-1 agreement: {matches}/{total} = {:.1}%",
+            100.0 * matches as f64 / total as f64);
+        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
+        for (h, _) in worker_handles { h.join().unwrap(); }
+        return;
+    }
+
+    // Teacher-forced DECODE diagnostic: prefill the prompt, then walk the oracle
+    // trajectory through the autoregressive decode path (NOT prefill), recording
+    // per-position top-1 agreement bucketed by position. Localizes long-context
+    // decode degradation (which prefill teacher-forcing cannot see).
+    if let Some(forced) = get_arg::<String>(&args, "--forced-decode") {
+        let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
+        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: token_ids.clone(), slot });
+        let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
+        wait_workers(&worker_handles);
+        let mut pred = sample_greedy_last(&logits); // prediction for forced[0]
+        let bucket = 50usize;
+        let mut buckets: Vec<(usize, usize)> = Vec::new();
+        let (mut matches, mut total) = (0usize, 0usize);
+        for (i, &f) in forced_ids.iter().enumerate() {
+            let ok = pred == f;
+            matches += ok as usize;
+            total += 1;
+            let b = i / bucket;
+            if buckets.len() <= b { buckets.push((0, 0)); }
+            buckets[b].0 += ok as usize;
+            buckets[b].1 += 1;
+            // Teacher-force: feed the oracle token through the decode path.
+            let pos = cache.seq_len(slot);
+            broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
+                tokens: vec![f], positions: vec![pos], slots: vec![slot],
+            });
+            let logits = model.forward_decode_paged(&[f], &[pos], &[slot], &mut cache);
+            wait_workers(&worker_handles);
+            pred = sample_greedy_last(&logits);
+        }
+        eprintln!("Teacher-forced DECODE agreement: {matches}/{total} = {:.1}%",
+            100.0 * matches as f64 / total as f64);
+        for (b, (m, t)) in buckets.iter().enumerate() {
+            eprintln!("  pos[{:>4}..{:<4}]: {m:>3}/{t:<3} = {:.0}%",
+                b * bucket, b * bucket + t, 100.0 * (*m as f64) / (*t as f64));
+        }
+        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
+        for (h, _) in worker_handles { h.join().unwrap(); }
+        return;
+    }
+
     // Prefill
     let t0 = Instant::now();
     broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill {

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);

crates/xserv-model/src/lib.rs

@@ -15,7 +15,7 @@ pub use gpt_oss::GptOss;
 pub use kv_cache::GpuKVCache;
 pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
 pub use qwen3::Qwen3;
-pub use sampling::{SamplingParams, sample};
+pub use sampling::{SamplingParams, sample, sample_greedy_penalized};
 
 /// Initialize GPU kernel hooks. Called automatically by model constructors,
 /// but safe to call multiple times (idempotent via OnceLock).

crates/xserv-model/src/sampling.rs

@@ -112,10 +112,51 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
     (vocab_size - 1) as u32
 }
 
-fn argmax(data: &[f32]) -> u32 {
-    data.iter()
-        .enumerate()
-        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
-        .map(|(i, _)| i as u32)
-        .unwrap()
+/// Greedy argmax with a repetition penalty applied to `recent` token ids
+/// (HF-style: divide positive logits by `penalty`, multiply negative by it).
+/// `penalty <= 1.0` is a no-op. Mitigates greedy repetition loops on reasoning
+/// models without changing the forward pass. NaN-safe.
+pub fn sample_greedy_penalized(logits: &Tensor, recent: &[u32], penalty: f32) -> u32 {
+    assert_eq!(logits.ndim(), 2);
+    let vocab_size = logits.shape()[1];
+    let seq_len = logits.shape()[0];
+    let logits_cpu = logits.to_device(Device::Cpu);
+    let mut last_row: Vec<f32> = match logits.dtype() {
+        DType::F32 => logits_cpu.as_slice::<f32>()[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec(),
+        DType::BF16 => logits_cpu.as_slice::<bf16>()[(seq_len - 1) * vocab_size..seq_len * vocab_size]
+            .iter().map(|v| v.to_f32()).collect(),
+        _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
+    };
+    if penalty > 1.0 {
+        for &id in recent {
+            let i = id as usize;
+            if i < last_row.len() {
+                let v = last_row[i];
+                last_row[i] = if v > 0.0 { v / penalty } else { v * penalty };
+            }
+        }
+    }
+    argmax(&last_row)
+}
+
+fn argmax(data: &[f32]) -> u32 {
+    // NaN-safe: a single NaN logit must not crash the engine thread (a
+    // partial_cmp().unwrap() panics on NaN). Skip NaNs; warn once if seen.
+    let mut best_i = 0usize;
+    let mut best = f32::NEG_INFINITY;
+    let mut nan_seen = false;
+    for (i, &v) in data.iter().enumerate() {
+        if v.is_nan() {
+            nan_seen = true;
+            continue;
+        }
+        if v > best {
+            best = v;
+            best_i = i;
+        }
+    }
+    if nan_seen {
+        eprintln!("[sampling] WARNING: NaN logits encountered in argmax");
+    }
+    best_i as u32
 }

crates/xserv-server/src/api.rs

@@ -350,10 +350,23 @@ fn build_prompt(messages: &[Message], model_type: &str) -> String {
 
 fn build_prompt_gpt_oss(messages: &[Message]) -> String {
     let mut prompt = String::new();
-    // System prompt
+    // System (meta) block: channel declaration required by harmony.
     prompt.push_str("<|start|>system<|message|>");
     prompt.push_str("You are a helpful assistant.\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.");
     prompt.push_str("<|end|>");
+    // Caller-supplied system prompt(s) go in a developer instructions block
+    // (harmony puts user instructions on the `developer` role, not `system`).
+    let dev_instructions: String = messages
+        .iter()
+        .filter(|m| m.role == "system")
+        .map(|m| m.content.as_str())
+        .collect::<Vec<_>>()
+        .join("\n\n");
+    if !dev_instructions.is_empty() {
+        prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
+        prompt.push_str(&dev_instructions);
+        prompt.push_str("<|end|>");
+    }
     for msg in messages {
         match msg.role.as_str() {
             "user" => {

crates/xserv-server/src/pp_engine.rs

@@ -187,7 +187,6 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
     let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id);
     eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})");
 
-    let eos = tokenizer.eos_token_id();
     let n_workers = world - 1;
     let next_peer = 1usize;
     let broadcast = |txs: &[mpsc::Sender<PpCommand>], cmd: PpCommand| {
@@ -220,10 +219,10 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+        emit_text(&tokenizer, &req, next, &mut decode_buf);
 
         let finish = loop {
-            if eos == Some(next) {
+            if tokenizer.is_eos(next) {
                 break "stop";
             }
             if generated >= req.max_tokens {
@@ -235,7 +234,7 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
             send_hidden(&sc, &x, next_peer);
             next = token_rx.recv().expect("decode token");
             generated += 1;
-            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+            emit_text(&tokenizer, &req, next, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
@@ -253,8 +252,8 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
-    if eos == Some(token_id) {
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
+    if tokenizer.is_eos(token_id) {
         return;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);

crates/xserv-server/src/tp_engine.rs

@@ -19,7 +19,7 @@ use std::thread;
 
 use xserv_distributed::{TpContext, UniqueId};
 use xserv_model::loader;
-use xserv_model::{sample, GptOss, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_model::{sample, sample_greedy_penalized, GptOss, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;
 
@@ -149,7 +149,23 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
     let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
     eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
 
-    let eos = tokenizer.eos_token_id();
+    // Optional repetition penalty to break greedy repetition loops (reasoning
+    // models loop under pure greedy when numerics diverge from the reference).
+    // Off by default; XSERV_REP_PENALTY>1 enables it over the last
+    // XSERV_REP_WINDOW generated tokens. Applied only on the greedy path.
+    let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY").ok()
+        .and_then(|s| s.parse().ok()).unwrap_or(1.0);
+    let rep_window: usize = std::env::var("XSERV_REP_WINDOW").ok()
+        .and_then(|s| s.parse().ok()).unwrap_or(128);
+    let pick = |logits: &Tensor, sp: &xserv_model::SamplingParams, history: &[u32]| -> u32 {
+        if rep_penalty > 1.0 && sp.temperature == 0.0 {
+            let start = history.len().saturating_sub(rep_window);
+            sample_greedy_penalized(logits, &history[start..], rep_penalty)
+        } else {
+            sample(logits, sp)
+        }
+    };
+
     let n_workers = world - 1;
     let broadcast = |txs: &[mpsc::Sender<TpCommand>], cmd: TpCommand| {
         for t in txs {
@@ -172,14 +188,16 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
         broadcast(&cmd_txs, TpCommand::Prefill { tokens: req.prompt_tokens.clone(), slot });
         let logits = rc.model.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
         wait_acks(&ack_rx);
-        let mut next = sample(&logits, &req.sampling);
+        let mut gen_ids: Vec<u32> = Vec::new();
+        let mut next = pick(&logits, &req.sampling, &gen_ids);
+        gen_ids.push(next);
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+        emit_text(&tokenizer, &req, next, &mut decode_buf);
 
         let finish = loop {
-            if eos == Some(next) {
+            if tokenizer.is_eos(next) {
                 break "stop";
             }
             if generated >= req.max_tokens {
@@ -189,9 +207,10 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
             broadcast(&cmd_txs, TpCommand::Decode { tokens: vec![next], positions: vec![pos], slots: vec![slot] });
             let logits = rc.model.forward_decode_paged(&[next], &[pos], &[slot], &mut rc.cache);
             wait_acks(&ack_rx);
-            next = sample(&logits, &req.sampling);
+            next = pick(&logits, &req.sampling, &gen_ids);
+            gen_ids.push(next);
             generated += 1;
-            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
+            emit_text(&tokenizer, &req, next, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
@@ -209,8 +228,8 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
-    if eos == Some(token_id) {
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
+    if tokenizer.is_eos(token_id) {
         return;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);

crates/xserv-tokenizer/src/bpe.rs

@@ -12,6 +12,7 @@ pub struct Tokenizer {
     special_token_ids: HashMap<u32, String>,
     pre_tokenize_re: Regex,
     eos_token_id: Option<u32>,
+    eos_token_ids: Vec<u32>,
     byte_fallback: bool,
 }
 
@@ -102,11 +103,27 @@ impl Tokenizer {
             decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
             decoder[at.id as usize] = at.content.as_bytes().to_vec();
         }
-        let eos_token_id = special_tokens
-            .get("<|im_end|>")
-            .or_else(|| special_tokens.get("<|end_of_text|>"))
-            .or_else(|| special_tokens.get("<|endoftext|>"))
-            .copied();
+        // End-of-generation tokens, in priority order. Families differ:
+        // Qwen uses <|im_end|>, Llama <|end_of_text|>, GPT-2 <|endoftext|>.
+        // gpt-oss (harmony) ends the assistant turn with <|return|> and also
+        // treats <|call|> (tool call) and <|endoftext|> as terminators
+        // (see generation_config.json eos_token_id = [200002, 199999, 200012]).
+        let eos_names = [
+            "<|im_end|>",
+            "<|end_of_text|>",
+            "<|return|>",
+            "<|call|>",
+            "<|endoftext|>",
+        ];
+        let mut eos_token_ids: Vec<u32> = Vec::new();
+        for name in eos_names {
+            if let Some(&id) = special_tokens.get(name) {
+                if !eos_token_ids.contains(&id) {
+                    eos_token_ids.push(id);
+                }
+            }
+        }
+        let eos_token_id = eos_token_ids.first().copied();
 
         // Pre-tokenization regex
         let pre_tokenize_re = if byte_fallback {
@@ -125,6 +142,7 @@ impl Tokenizer {
             special_token_ids,
             pre_tokenize_re,
             eos_token_id,
+            eos_token_ids,
             byte_fallback,
         }
     }
@@ -249,6 +267,12 @@ impl Tokenizer {
         self.eos_token_id
     }
 
+    /// True if `id` is any end-of-generation token (a model may have several;
+    /// gpt-oss/harmony ends on <|return|>, <|call|>, or <|endoftext|>).
+    pub fn is_eos(&self, id: u32) -> bool {
+        self.eos_token_ids.contains(&id)
+    }
+
     pub fn vocab_size(&self) -> usize {
         self.decoder.len()
     }

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<<<grid, block, smem_bytes, (cudaStream_t)stream>>>(
+        (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,