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