M2.5 Phase 2: decode CUDA-graph capture/replay (single-stream) + fused-MoE residency

Wire CUDA-graph capture/replay into the fp4/CUDA pure-decode path (mirrors
vLLM's decode CUDAGraph capture: capture per batch shape, persistent inputs,
decode-only). New Qwen3_5DecodeGraph captures the forward's layer region once
per shape and replays it per token; the embedding is kept OUTSIDE the graph
(its CUDA op mallocs/frees/syncs a bounds-check flag) and writes a persistent
hidden buffer that the captured region reads. Per-step-varying inputs live in
persistent HOST vectors mutated in place (capturable on GB10 pageable memory).
A cold shape runs one eager step to pre-warm the DevicePool + residency, the
next step captures, subsequent steps replay.

Fused-MoE (MoeBlockFusedCuda): the per-layer expert pointer/scale arrays +
pair->token map are now uploaded ONCE into a resident cache instead of rebuilt
+ re-uploaded from host stack temporaries every step (those dangling host
sources are illegal inside a capture region; also pure per-step waste).

Correctness (GB10, free box): test_qwen36_paged_engine greedy gate passes
16/16 token-for-token with the graph active. All other tests pass (the sole
red, test_qwen36_weights, is pre-existing: it expects bf16 experts but the
on-disk checkpoint is NVFP4 — fails identically at baseline).

Measured decode A/B (GB10, 0% contention), graph ON vs OFF:
  num_reqs==1 (1x16x128): TPOT 65.5 vs 67.2 ms  -> ~2.6% faster (graph helps)
  num_reqs==8 (8x1024x128): TPOT 265 vs 247 ms  -> ~7% slower (graph hurts)
After Phase 1 made decode async-on-stream, host launch overhead is largely
hidden behind the GPU, so graph capture recovers little; batched decode is
GPU-bound and gains extra graph-launch overhead from the many GDN-gather nodes.
The runner therefore gates the graph to num_reqs==1 (single-stream latency win,
no batched-throughput regression). Env VLLM_CPP_CUDAGRAPH=0 disables it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UJyFKcK62CcR3imhgbiBnW

Author: mudler

include/vllm/model_executor/models/qwen3_5.h

@@ -24,6 +24,7 @@
 #pragma once
 
 #include <cstdint>
+#include <memory>
 #include <vector>
 
 #include "vllm/model_executor/models/qwen3_5_weights.h"
@@ -101,6 +102,58 @@ class Qwen3_5Model {
                                          const HfConfig& config, vt::Queue& queue);
 };
 
+// Decode-step CUDA-graph driver (M2.5 Phase 2, gate-#1 decode-launch unlock).
+// Wraps the paged forward COMPUTE body (embed -> layers -> lm_head) in a
+// capture-once / replay-per-token CUDA graph for PURE-DECODE batches, collapsing
+// the ~thousands of per-step kernel-launch + memcpy host-API calls into a single
+// cudaGraphLaunch. Mirrors vLLM's decode CUDAGraph capture: capture keyed on the
+// batch SHAPE, per-step-varying inputs threaded through PERSISTENT buffers, and
+// decode-only (prefill / mixed batches stay eager, kept off this path by the
+// runner).
+//
+// ── vt-runtime realization (deviations, so upstream ports mechanically) ──────
+//   * PERSISTENT INPUTS are the HOST step vectors (token_ids / positions / the
+//     attention+GDN metadata), held here and MUTATED IN PLACE each step. On GB10
+//     (pageable memory access) the forward's host->device input copies are
+//     capturable, so a replay re-reads each new token's inputs from the fixed
+//     host addresses — no separate device staging buffers (vLLM keeps torch
+//     tensors on-GPU; here the "buffer" is the host vector the copy reads from).
+//   * The GDN mamba-state gather offsets and the block-table column count are
+//     BAKED at capture, so the SHAPE key includes them; any change re-captures.
+//   * A cold shape runs one EAGER step first (pre-warms the DevicePool + the
+//     resident weights / fused-MoE constants) so the capture region does zero
+//     cudaMalloc; the next same-shape step captures, and subsequent ones replay.
+//   * VLLM_CPP_CUDAGRAPH=0 disables capture (always eager) for the A/B and as a
+//     safety valve. Non-CUDA devices always run eager.
+class Qwen3_5DecodeGraph {
+ public:
+  Qwen3_5DecodeGraph(const Qwen3_5MoeWeights& weights, const HfConfig& config,
+                     vt::Queue queue);
+  ~Qwen3_5DecodeGraph();
+  Qwen3_5DecodeGraph(const Qwen3_5DecodeGraph&) = delete;
+  Qwen3_5DecodeGraph& operator=(const Qwen3_5DecodeGraph&) = delete;
+
+  // One PURE-DECODE step. Returns [T, vocab] f32 logits, bit-identical to
+  // Qwen3_5Model::Forward for the same inputs/caches. attn_kv / gdn_state are the
+  // runner's persistent caches (stable addresses across steps). The caller must
+  // only route pure-decode batches here (all query_len==1, no prefill).
+  std::vector<float> Step(const std::vector<int32_t>& token_ids,
+                          const std::vector<int32_t>& positions,
+                          const v1::CommonAttentionMetadata& attn_meta,
+                          const v1::GDNAttentionMetadata& gdn_meta,
+                          const std::vector<PagedKvCache>& attn_kv,
+                          const std::vector<GdnStateCache>& gdn_state);
+
+  // Diagnostics (A/B + tests): is a graph currently captured, and how many
+  // replays have run since the last (re)capture.
+  bool captured() const;
+  int64_t replay_count() const;
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
+};
+
 // Per-layer parity replay: runs ONE decoder layer over the combined residual
 // stream. `hidden_in` is the [T*H] f32 stream INTO the layer (= residual +
 // hidden, as the pinned oracle reconstructs it); `positions` is the length-T

include/vllm/v1/worker/gpu/runner.h

@@ -54,6 +54,7 @@
 
 #include <cstdint>
 #include <map>
+#include <memory>
 #include <optional>
 #include <string>
 #include <vector>
@@ -174,6 +175,12 @@ class GPUModelRunner final : public ModelRunnerBase {
   std::vector<PagedKvCache> attn_kv_;
   std::vector<GdnStateCache> gdn_state_;
 
+  // Decode CUDA-graph driver (M2.5 Phase 2), created lazily on the first
+  // pure-decode step of an fp4/CUDA model. Captures the decode forward once per
+  // batch shape and replays it per token (mirrors vLLM's decode CUDAGraph
+  // capture). nullptr on CPU / bf16 / when disabled — those keep the eager path.
+  std::unique_ptr<Qwen3_5DecodeGraph> decode_graph_;
+
   // Stashed forward result between execute_model and sample_tokens (upstream
   // ExecuteModelState — hidden_states + input_batch handoff, here the full
   // logits + the dense-order step). num_reqs == 0 marks a 0-token flush step.