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-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, sample_greedy_penalized, GptOss, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_model::{sample, sample_greedy_penalized, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;
 
@@ -58,6 +58,16 @@ impl TpModel {
 struct RankCtx {
     model: TpModel,
     cache: PagedKVCache,
+    decoder: GraphedGptOssDecoder,
+}
+
+/// Decode one step: gpt-oss batch=1 goes through the CUDA-graph decoder
+/// (lazy capture, replay thereafter); everything else runs eager.
+fn rank_decode(rc: &mut RankCtx, tokens: &[u32], positions: &[usize], slots: &[usize]) -> Tensor {
+    match &rc.model {
+        TpModel::GptOss(m) => rc.decoder.decode(m, tokens, positions, slots, &mut rc.cache),
+        TpModel::Qwen3(_) => rc.model.forward_decode_paged(tokens, positions, slots, &mut rc.cache),
+    }
 }
 
 fn build_rank(
@@ -81,7 +91,7 @@ fn build_rank(
     let cache = PagedKVCache::new_tp(
         config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
     );
-    RankCtx { model, cache }
+    RankCtx { model, cache, decoder: GraphedGptOssDecoder::new() }
 }
 
 fn worker_loop(
@@ -106,7 +116,7 @@ fn worker_loop(
                 let _ = rc.model.forward_prefill_paged(&tokens, slot, &mut rc.cache);
             }
             TpCommand::Decode { tokens, positions, slots } => {
-                let _ = rc.model.forward_decode_paged(&tokens, &positions, &slots, &mut rc.cache);
+                let _ = rank_decode(&mut rc, &tokens, &positions, &slots);
             }
             TpCommand::Shutdown => {
                 let _ = ack_tx.send(());
@@ -207,7 +217,7 @@ pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
             }
             let pos = rc.cache.seq_len(slot);
             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);
+            let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]);
             wait_acks(&ack_rx);
             next = pick(&logits, &req.sampling, &gen_ids);
             gen_ids.push(next);