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/gpt_oss.rs

@@ -373,24 +373,62 @@ impl GptOss {
         assert_eq!(seq_slots.len(), batch);
         assert!(batch > 0);
 
-        let num_heads = self.local_num_heads;
-        let num_kv_heads = self.local_num_kv_heads;
-        let head_dim = self.config.head_dim();
-        let eps = self.norm_eps();
+        self.decode_prepare(positions, seq_slots, paged_cache);
 
+        // Upload token ids + positions, then run the pure-GPU core.
+        let ids_gpu = upload_u32(tokens);
+        let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
+        let pos_gpu = upload_u32(&positions_u32);
+        let logits = self.decode_core(
+            ids_gpu.as_ptr() as *const c_void,
+            pos_gpu.as_ptr() as *const c_void,
+            batch,
+            paged_cache,
+        );
+
+        for &slot in seq_slots {
+            paged_cache.advance_seq_len(slot, 1);
+        }
+        logits
+    }
+
+    /// Host-side per-step cache bookkeeping: block allocation + uploading
+    /// block tables / context lens to their (stable-address) GPU buffers.
+    /// Runs OUTSIDE the CUDA-graph captured region.
+    pub fn decode_prepare(
+        &self,
+        positions: &[usize],
+        seq_slots: &[usize],
+        paged_cache: &mut PagedKVCache,
+    ) {
         let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
         for (b, &slot) in seq_slots.iter().enumerate() {
             paged_cache.ensure_capacity(slot, positions[b] + 1);
         }
         paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens);
+    }
+
+    /// The pure-GPU decode step: embedding → 24 layers → final norm → logits.
+    /// Token ids and positions are read from device buffers; every other input
+    /// (weights, KV pools, block table, context lens) has a stable address —
+    /// which is exactly what makes this region CUDA-graph capturable.
+    pub fn decode_core(
+        &self,
+        ids_gpu: *const c_void,
+        pos_gpu: *const c_void,
+        batch: usize,
+        paged_cache: &mut PagedKVCache,
+    ) -> Tensor {
+        let num_heads = self.local_num_heads;
+        let num_kv_heads = self.local_num_kv_heads;
+        let head_dim = self.config.head_dim();
+        let eps = self.norm_eps();
 
         let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
         let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
         let max_blocks = paged_cache.max_blocks_per_seq();
 
-        let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
-
-        let mut x = embedding(&self.embed_tokens, tokens);
+        let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);
 
         for (layer_idx, layer) in self.layers.iter().enumerate() {
             let residual = x.clone();
@@ -407,8 +445,8 @@ impl GptOss {
             let k_3d = k_all.reshape(&[batch, num_kv_heads, head_dim]);
 
             // RoPE (no QK-norm for gpt-oss)
-            rope_inplace(&q_3d, &self.rope_cache, &positions_u32);
-            rope_inplace(&k_3d, &self.rope_cache, &positions_u32);
+            rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
+            rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
 
             let v_3d = v_all.reshape(&[batch, num_kv_heads, head_dim]);
 
@@ -445,11 +483,6 @@ impl GptOss {
             x = xserv_kernels::add(&residual, &moe_out);
         }
 
-        // Advance KV cache
-        for &slot in seq_slots {
-            paged_cache.advance_seq_len(slot, 1);
-        }
-
         let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
         matmul_2d(&x, &self.lm_head_t)
     }
@@ -673,6 +706,16 @@ impl GptOss {
 
 // --- Helpers ---
 
+/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
+fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
+    let bytes = unsafe {
+        std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4)
+    };
+    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
+    buf.copy_from_host(bytes).unwrap();
+    buf
+}
+
 /// XSERV_DENSE_MOE=1 forces the dense all-expert path (A/B benchmarking).
 fn dense_moe_forced() -> bool {
     static FORCED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();