gpt-oss: replay the whole batch=1 decode step as one CUDA graph

Split forward_decode_paged into host bookkeeping (decode_prepare +
ids/pos upload + advance_seq_len) and a pure-GPU decode_core. The
paged-KV and sparse-MoE designs already read every per-step variable
(block table, context lens, expert ids) from stable-address device
buffers, so decode_core captures as-is.

GptOssDecodeGraph captures lazily on the second decode step (the
first eager step warms cuBLAS) after a "retained warmup": the step
runs once with the allocator quarantine on, stocking the pool with a
dedicated block for every allocation so the capture itself never
pool-misses (a cudaMalloc while capturing is illegal — and the
capture's own quarantine is what would otherwise starve the pool).
NCCL all-reduces capture cleanly; TP=2 replays in lockstep.

Wired into tp_engine, bench-gpt-oss, and xserv-chat via
GraphedGptOssDecoder (batch>1 falls back to eager;
XSERV_DECODE_GRAPH=0 disables). Greedy tokens identical to eager.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use xserv_distributed::{TpContext, UniqueId, get_unique_id};
-use xserv_model::{loader, GptOss, ModelConfig, PagedKVCache, BLOCK_SIZE};
+use xserv_model::{loader, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, BLOCK_SIZE};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -172,6 +172,7 @@ fn main() {
     print!("{prompt}");
 
     // Decode
+    let mut decoder = GraphedGptOssDecoder::new();
     let decode_start = Instant::now();
     for _ in 1..max_tokens {
         let text = tokenizer.decode(&[next]);
@@ -183,7 +184,7 @@ fn main() {
         broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
             tokens: vec![next], positions: vec![pos], slots: vec![slot],
         });
-        let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut cache);
+        let logits = decoder.decode(&model, &[next], &[pos], &[slot], &mut cache);
         wait_workers(&worker_handles);
 
         next = sample_greedy_last(&logits);
@@ -250,6 +251,7 @@ fn worker_loop(
     eprintln!("[rank {rank}] Ready.");
     ack_tx.send(()).unwrap();
 
+    let mut decoder = GraphedGptOssDecoder::new();
     while let Ok(cmd) = rx.recv() {
         match cmd {
             WorkerCmd::Register(slot) => {
@@ -259,7 +261,7 @@ fn worker_loop(
                 let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
             }
             WorkerCmd::Decode { tokens, positions, slots } => {
-                let _ = model.forward_decode_paged(&tokens, &positions, &slots, &mut cache);
+                let _ = decoder.decode(&model, &tokens, &positions, &slots, &mut cache);
             }
             WorkerCmd::Shutdown => break,
         }
@@ -286,6 +288,11 @@ fn wait_workers(handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiv
 fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
     use half::bf16;
     assert_eq!(logits.ndim(), 2);
+    // GPU argmax fast path (4-byte D2H instead of the full logits row).
+    if logits.dtype() == xserv_tensor::DType::BF16 && logits.is_contiguous() {
+        let ids = xserv_kernels::argmax_bf16_to_host(logits);
+        return *ids.last().unwrap();
+    }
     let logits_cpu = logits.to_device(Device::Cpu);
     let vocab_size = logits.shape()[1];
     let seq_len = logits.shape()[0];

crates/xserv-model/src/bin/xserv-chat.rs

@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use std::sync::{mpsc, Arc};
 use std::thread;
 
-use xserv_model::{loader, sample, sample_greedy_penalized, GptOss, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
+use xserv_model::{GraphedGptOssDecoder, loader, sample, sample_greedy_penalized, GptOss, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -77,6 +77,7 @@ fn tp_worker_loop(
     let mut cache = PagedKVCache::new_tp(
         &config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, rank as u32,
     );
+    let mut decoder = GraphedGptOssDecoder::new();
     while let Ok(cmd) = cmd_rx.recv() {
         match cmd {
             TpCommand::Register(slot) => { let _ = cache.register_sequence(slot); }
@@ -85,7 +86,7 @@ fn tp_worker_loop(
                 let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
             }
             TpCommand::Decode { tokens, positions, slots } => {
-                let _ = model.forward_decode_paged(&tokens, &positions, &slots, &mut cache);
+                let _ = chat_decode(&model, &mut decoder, &tokens, &positions, &slots, &mut cache);
             }
         }
         let _ = ack_tx.send(());
@@ -321,6 +322,7 @@ fn main() {
     };
 
     let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
+    let mut decoder = GraphedGptOssDecoder::new();
     if let Some(h) = &tp_handle { h.send(TpCommand::Register(SLOT)); h.wait(); }
     cache.register_sequence(SLOT).expect("register chat slot");
     let use_color = opts.color && io::stdout().is_terminal();
@@ -384,7 +386,7 @@ fn main() {
             print!("assistant> ");
             io::stdout().flush().unwrap();
             let (_finish, answer) = generate_with_paged_cache(
-                &model, &mut cache, &tokenizer, &prompt_tokens, &opts.sampling,
+                &model, &mut decoder, &mut cache, &tokenizer, &prompt_tokens, &opts.sampling,
                 max_new_tokens, use_color, &tp_handle, is_moe, opts.enable_thinking,
             );
             moe_history.push((input.to_string(), answer));
@@ -421,6 +423,7 @@ fn main() {
         io::stdout().flush().unwrap();
         let (finish, _answer) = generate_with_paged_cache(
             &model,
+            &mut decoder,
             &mut cache,
             &tokenizer,
             &prompt_tokens,
@@ -679,8 +682,23 @@ fn today_ymd() -> String {
     format!("{y:04}-{m:02}-{d:02}")
 }
 
+fn chat_decode(
+    model: &ChatModel,
+    decoder: &mut GraphedGptOssDecoder,
+    tokens: &[u32],
+    positions: &[usize],
+    slots: &[usize],
+    cache: &mut PagedKVCache,
+) -> xserv_tensor::Tensor {
+    match model {
+        ChatModel::GptOss(m) => decoder.decode(m, tokens, positions, slots, cache),
+        ChatModel::Qwen3(_) => model.forward_decode_paged(tokens, positions, slots, cache),
+    }
+}
+
 fn generate_with_paged_cache(
     model: &ChatModel,
+    decoder: &mut GraphedGptOssDecoder,
     cache: &mut PagedKVCache,
     tokenizer: &Tokenizer,
     prompt_tokens: &[u32],
@@ -745,7 +763,7 @@ fn generate_with_paged_cache(
     for _ in 0..max_tokens {
         let position = cache.seq_len(SLOT);
         if let Some(h) = tp { h.send(TpCommand::Decode { tokens: vec![next], positions: vec![position], slots: vec![SLOT] }); }
-        let logits = model.forward_decode_paged(&[next], &[position], &[SLOT], cache);
+        let logits = chat_decode(model, decoder, &[next], &[position], &[SLOT], cache);
         if let Some(h) = tp { h.wait(); }
         if tokenizer.is_eos(next) {
             print_stream_text(