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/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(