[Mamba] Prefetch stochastic-rounding seeds on a side stream

[pst2154]

Overview

Opt-in (VLLM_MAMBA_SR_SEED_PREFETCH=1): with the FlashInfer Mamba SSU backend and Mamba cache stochastic rounding, draw a forward pass's per-layer stochastic-rounding seeds up front on a side CUDA stream; each SSU call waits on an event for its seed. Seed values are identical to stock.

One of five independent opt-in Nemotron-H / MTP decode changes (#1–#5 in this fork).

Claims

• Removes one small RNG kernel per Mamba-2 layer from the critical path of every decode step (~40 per step on the measured model).
• Identical seed values to stock, eagerly and over CUDA-graph replays (unit-tested).
• Measured (plugin with identical logic, 4x B200 TP4): -1.4 / -2.0 / -1.9 / -1.0 % step time at BS 1 / 2 / 4 / 8 on top of [Model] NemotronH: overlap MoE router gate with fc1_latent on a side stream #1 without spec decode; with [Model] NemotronH: overlap MoE router gate with fc1_latent on a side stream #1, -6.7 / -7.0 / -3.6 / -2.8 % per MTP verify iteration at temperature 1. The source port engages in serving and the five-PR combination is faster than stock end to end (below).
• Default off; no effect without the FlashInfer SSU backend and --enable-mamba-cache-stochastic-rounding.

Validation

GPU validation of this source port (4x B200, vLLM v0.30.0 image with the branch files overlaid):

• python -m pytest tests/kernels/mamba/test_sr_seed_prefetch.py -v, 1x B200: 4 passed — eager (with and without extra calls beyond N), CUDA-graph capture + 3 replays (identical to stock, fresh per replay), fallback outside a forward context.
• Engagement in serving: Prefetching Mamba stochastic-rounding seeds for 40 layers per forward on a side stream. (40 = the model's Mamba-2 layers).

End-to-end, #1–#5 combined (all five env vars, plus the drafter flag --speculative_config.moe_backend flashinfer_trtllm, a CLI flag that is not part of these PRs) vs. stock, same node and image, fresh server per arm, 2 back-to-back pairs. 4x B200, TP4, Nemotron-H NVFP4 with MTP (7 draft tokens), FP8 KV cache; aiperf BS1 streaming chat, 10 requests with ~10k-token prompts, temperature 1, top_p 0.95:

metric (combined vs stock)	pair 1	pair 2
Output tok/s/user, p50	+8.6 %	+11.8 %
Output tok/s/user, avg	+8.4 %	+7.9 %
Inter-token latency, p50	-7.6 %	-11.1 %
Inter-token latency, avg	-6.8 %	-5.6 %
Time to first token, avg	-1.1 %	-0.6 %
Acceptance length, combined / stock	4.38 / 4.48	4.36 / 4.48

Absolute numbers omitted. Acceptance length at temperature 1 varies by a few percent between runs, and tokens/s/user scales with it. Model eval was not run for these fork drafts.

Server command (both arms)

vllm serve <Nemotron-H NVFP4 checkpoint with MTP head> --tensor-parallel-size 4 --language-model-only --speculative_config.method mtp --speculative_config.num_speculative_tokens 7 --kv-cache-dtype fp8 --mamba-backend flashinfer --mamba-ssm-cache-dtype float16 --enable-mamba-cache-stochastic-rounding --mamba-cache-philox-rounds 5 --mamba-cache-mode align --enable-prefix-caching --async-scheduling --max-model-len 262144 --max-num-batched-tokens 10240

Details

Stock FlashInferSSUBackend.__call__ launches torch.randint(0, 2**32, (1,)) before every SSU kernel. That is a non-PDL kernel of about 2.5 us between conv1d and the SSU kernel, on the critical path of every Mamba layer. With about 40 Mamba-2 layers this adds up to roughly 100 us per decode step.

Limitations:

• Assumes no other user of the CUDA default generator runs between the Mamba layers of a forward (true for vLLM models today); otherwise seed values would shift.
• Under CUDA-graph capture, a forward that calls the SSU fewer than N times would leave the side stream unjoined and capture fails loudly (never silently); every Nemotron-H decode forward calls it once per Mamba-2 layer.
• Disabled for ReplaySSM, Mamba-1 models, DBO and non-CUDA platforms.

Open questions for review:

• Gating: env var vs. default-on whenever stochastic rounding is enabled (values are identical).
• Forward delimiting: forwards are told apart by ForwardContext identity; an alternative is an explicit per-forward reset hook from the model runner.

Duplicate check: searched open upstream PRs (gh pr list --repo vllm-project/vllm --state open --search ...) for "stochastic rounding seed" (SSM-cache / ReplaySSM work (vllm-project#52962, vllm-project#40012, vllm-project#54993, ...), none moves the seed draw); "mamba stochastic rounding" (same set, none related); "ssu seed prefetch" (no results). No open PR implements this change.

AI assistance: Claude Code was used to port the plugin, write the tests, run the GPU validation and draft this description. A human must review every changed line before anything goes upstream.

Design, race-freedom argument, gating, plugin-level evidence (from the original description)

Design

• StochasticRoundingSeedPrefetcher (in ssu_dispatch.py), N = number of Mamba-2 layers in the KV cache config.
  • On the first SSU call of a forward pass, it records _fork_event on the current stream and makes a dedicated side stream wait on it. It then issues the same N torch.randint(0, 2**32, (1,), device=...) calls (same shape, dtype, device and order) on the side stream, recording _events[k] after the k-th.
  • The k-th SSU call of that forward makes the current stream wait on _events[k] and passes seed k to the unchanged FlashInfer SSU kernel.
• Forward passes are told apart by their ForwardContext (held through a weakref). This is robust where a plain call counter could drift. The prefetcher returns None, and the caller draws inline exactly as stock, for:
  • calls beyond N within one forward;
  • calls outside a forward context;
  • DBO micro-batches.
• initialize_mamba_ssu_backend attaches the prefetcher to the FlashInfer backend only when all of these hold:
  • the env var is set;
  • stochastic rounding is on;
  • ReplaySSM is off (it has its own SSU path, not touched here);
  • the platform is CUDA;
  • the model has no Mamba-1 layers.
• The kernel call itself is unchanged: FlashInferSSUBackend.__call__ only changes where rand_seed comes from.

Why it is race-free

• Ordering uses CUDA events only. One fork event per forward and one event per seed; they become graph edges under CUDA graph capture. There is no polling and no timing assumption.
• Producer to consumer. _events[k] is recorded on the side stream after seed k is written. The current stream waits on it before enqueuing the SSU kernel that reads seed k.
• Buffer lifetimes across streams (the event ordering gives what record_stream would):
  • The seeds are allocated on a side stream used only by this prefetcher, so only a later side-stream allocation can reuse their blocks. That means the next forward's draws, which start after that forward's _fork_event. The current stream records that event after every consumer of the previous seeds was enqueued.
  • The seed list is replaced only after the new seeds are allocated, so the old blocks cannot be handed back within the same draw.
• Values. The CUDA generator's (seed, offset) is advanced on the host at each randint call, in call order. No other generator user runs between the Mamba layers of a forward. So the k-th pre-drawn seed equals the stock k-th draw, eagerly and under graph replay. Extra calls beyond N draw inline after the N pre-drawn seeds, which is again the stock order.
• Graph capture. All N events are waited on by the N SSU calls of the forward, so the side stream is joined before capture ends. A forward that called the SSU fewer than N times would leave it unjoined, and CUDA graph capture would fail loudly rather than silently. In Nemotron-H every decode forward calls the SSU once per Mamba-2 layer, and the MTP drafter has no Mamba layers.

Gating

• VLLM_MAMBA_SR_SEED_PREFETCH (bool, default 0) enables the prefetch. It is a normal vllm/envs.py entry and part of the compile cache key.
• It has no effect without the FlashInfer SSU backend and --enable-mamba-cache-stochastic-rounding.
• With the flag off, the only change in the hot path is one None check.

Plugin-level correctness evidence

Identity test on 1x B200, 40 layers, other work interleaved on the main stream between uses. The pre-drawn seeds equal stock's per-layer draws eagerly and over 3 CUDA graph replays, and the seeds differ across replays as in stock. All checks passed. This PR ports that test: tests/kernels/mamba/test_sr_seed_prefetch.py (eager, eager with extra calls beyond N, CUDA graph replay, fallback outside a forward context).

Plugin-level performance

Model: Nemotron-H (hybrid Mamba-2 / latent-MoE, NVFP4) with MTP, TP4 on 4x B200, vLLM v0.30.0, FlashInfer Mamba backend, fp16 SSM cache with stochastic rounding, prefix caching, async scheduling. Metric: median ITL, i.e. time per decode step, or per verify iteration with MTP. This metric does not depend on acceptance.

setting	BS1	BS2	BS4	BS8
No spec decode, step time, incremental on top of the MoE router-gate overlap (1 rep vs 2 reps)	-1.4 %	-2.0 %	-1.9 %	-1.0 %
MTP (5 spec tokens), temperature 1, per verify iteration, combined with the router-gate overlap vs stock (fresh servers, 2 reps each)	-6.7 %	-7.0 %	-3.6 %	-2.8 %

For comparison, the router-gate overlap alone measured -4.5 / -4.7 / -3.5 / -1.4 % per MTP verify iteration in a separate run at temperature 0. A BS1 streaming chat benchmark with both items (1 rep, temperature 1) gave +4.9 % and +8.7 % TPS/user p50 on two prompt sets. No isolated measurement of this item with MTP was taken.

Port notes

Deviations from the plugin:

• N is counted from the Mamba-2 layers in the KV cache config instead of the MambaMixer2 modules of the target model; both are 40 for this model.
• Forwards are delimited by ForwardContext identity instead of a modular call counter, with a stock fallback for extra calls, calls outside a context and DBO.
• Disabled for ReplaySSM and Mamba-1 models.

Seed values and stream/event structure are unchanged in the validated configuration.

🤖 Generated with Claude Code

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity. (See open questions.)

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples. (Pending human review; full pre-commit / mypy not run.)

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code