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

[pst2154]

Overview

Opt-in VLLM_MAMBA_SR_SEED_PREFETCH=1 (default off): with the FlashInfer Mamba SSU backend and --enable-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.

Claims

• Removes one small RNG kernel (torch.randint, ~2.5 µs, non-PDL) per Mamba-2 layer from the critical path of every decode step, between conv1d and the SSU kernel.
• Identical seed values to stock, eagerly and over CUDA-graph replays (unit-tested).
• Measured on v0.30.0, Nemotron-H (hybrid Mamba-2 / latent-MoE, NVFP4) with MTP, TP4 on 4x B200 (out-of-tree version with identical logic): -1.4 / -2.0 / -1.9 / -1.0 % step time at BS 1 / 2 / 4 / 8 without spec decode, on top of [MoE] Overlap the runner-held gate with the routed input transform on the aux stream #59728.
• No effect without the FlashInfer SSU backend and stochastic rounding.
• On this commit, with all five related drafts enabled together (not this one alone): -4.9 to -5.6 % time per verify iteration vs stock on the model above, over 2 interleaved pairs (see Validation).

Validation

Validated on this commit: the vllm/vllm-openai:cu129-nightly-ac9126e58... image with this branch's vllm/ files overlaid read-only. The image's installed files for every touched path were first verified byte-identical to ac9126e58. Earlier validation on v0.30.0 is in pst2154#2. The port onto ac9126e58 only needed a textual conflict resolution next to the new ReplaySSM commit helper; the non-ReplaySSM SSU call path (FlashInferSSUBackend.__call__, one call per Mamba-2 layer per decode forward) is unchanged.

pytest -v tests/kernels/mamba/test_sr_seed_prefetch.py (1x B200): 4 passed. In serving, the one-time line Prefetching Mamba stochastic-rounding seeds for 40 layers per forward on a side stream. appeared.

End to end on this commit, all five related drafts (#59726, #59727, #59728, #59729, #59730) enabled together vs stock. Nemotron-H (hybrid Mamba-2 / latent-MoE, NVFP4) with MTP (7 draft tokens), TP4 on 4x B200, FP8 KV cache, MRV2, async scheduling, prefix caching, FlashInfer Mamba backend with stochastic rounding. Streaming chat at concurrency 1, 10 requests with ~10k-token prompts, temperature 1, top_p 0.95. Each arm used a fresh server, with arms interleaved on one node. Both arms set CUDA_SELECTIVE_DEVICE_CODE_RECOMPILE=FORCED, because without it the installed driver runtime-patches the image's Blackwell kernels and skews timings. Time per verify iteration = sum(request latency - TTFT) / sum(streamed chunks), one chunk per verify iteration, so it does not depend on acceptance.

arm vs stock	pair	time per verify iteration (all req / req 2-10)	ITL avg x AL	tok/s/user avg	tok/s/user p50	AL arm / stock
five PRs	1	-5.6 % / -5.6 %	-2.4 %	+1.8 %	-0.1 %	4.46 / 4.47
five PRs	2	-4.9 % / -5.3 %	+1.8 %	-0.1 %	+4.1 %	4.26 / 4.40
five PRs + drafter --speculative_config.moe_backend flashinfer_trtllm	1	-10.3 % / -10.5 %	-16.9 %	+4.2 %	-2.5 %	4.07 / 4.64
five PRs + drafter --speculative_config.moe_backend flashinfer_trtllm	2	-10.7 % / -10.6 %	-10.6 %	+4.5 %	+1.4 %	4.13 / 4.47

AL comes from the server's spec-decode metrics over the measured window. At temperature 1 it varies by several percent between runs, and tok/s/user and ITL x AL inherit that noise, so read the per-iteration time. The drafter MoE flag is not part of these PRs; it changes draft numerics, and AL was lower with it in these runs. The changes were not measured one at a time on this commit.

Server command

Stock: vllm serve <Nemotron-H NVFP4 checkpoint with MTP head, FP8 KV in its config> --tensor-parallel-size 4 --moe-backend auto --max-num-seqs 512 --max-model-len 262144 --max-num-batched-tokens 10240 --enable-prefix-caching --async-scheduling --speculative_config.method mtp --speculative_config.num_speculative_tokens 7 --speculative_config.moe_backend flashinfer_cutlass --mamba-backend flashinfer --mamba-ssm-cache-dtype float16 --enable-mamba-cache-stochastic-rounding --mamba-cache-philox-rounds 5 --mamba-cache-mode align. The PR arms add VLLM_NEMOTRON_H_MOE_ROUTER_OVERLAP=1 VLLM_MAMBA_SR_SEED_PREFETCH=1 VLLM_CUSTOM_ALL_GATHER_SMALL=1 VLLM_MAMBA_FUSED_STATE_INDEX=1 and --speculative_config.use_local_argmax_reduction true.

Earlier, on v0.30.0 (1x B200):

pytest -v tests/kernels/mamba/test_sr_seed_prefetch.py
4 passed

Eager (with and without extra calls beyond N), CUDA-graph capture + 3 replays (identical to stock, fresh values per replay), fallback outside a forward context. In serving, the one-time log line Prefetching Mamba stochastic-rounding seeds for 40 layers per forward on a side stream. appeared (40 = the model's Mamba-2 layers).

No model eval: seeds, and therefore outputs, are unchanged.

Details

• Design: StochasticRoundingSeedPrefetcher in ssu_dispatch.py. The first SSU call of a forward records _fork_event on the current stream, makes a dedicated side stream wait on it and issues the same num_layers randint calls (same shape, dtype, device, order) there, recording _events[k] after the k-th. The k-th SSU call waits on _events[k] and uses seed k. The CUDA generator's (seed, offset) advances on the host in call order, so seed k equals the stock k-th draw. Forwards are told apart by ForwardContext identity. num_layers = Mamba-2 layers in the KV cache config.
• Race-free: ordering by CUDA events only (graph edges under capture). Seed tensors are allocated on the side stream, so their blocks can only be reused by the next forward's draws, which start after a _fork_event recorded once all consumers of the previous seeds were enqueued.
• Limitations: assumes no other user of the default CUDA generator runs between the Mamba layers of a forward (true for vLLM models today). Under CUDA-graph capture, a forward that calls the SSU fewer than N times would leave the side stream unjoined, and capture fails loudly rather than silently. Disabled for ReplaySSM, Mamba-1 models, DBO and non-CUDA platforms.
• Open questions: env var vs default-on whenever stochastic rounding is enabled (values are identical); forward delimiting by ForwardContext identity vs an explicit per-forward reset hook from the model runner.
• Duplicate check: gh pr list --repo vllm-project/vllm --state open --search "stochastic rounding seed" -> [Mamba] Support quantized FlashInfer ReplaySSM state cache #52962, INT16/FP8/INT8 + SR SSM Cache Support #40012, [Spec Decode] Verify in head dtype and drop the FP32 chunk workaround #53630, [Feature][Mamba] Support batch-invariant Mamba2 prefill and recovery #54993, [Kernel] ReplaySSM: cache SSM inputs for faster Gated DeltaNet standard decode #48792, [Kernel] ReplaySSM: cache SSM inputs for faster Gated DeltaNet speculative decode #49887, [Bugfix] Bound accepted-token state lookups in GDN/KDA spec decode #50021 (SSM cache quantization / ReplaySSM / batch-invariance work; none moves the seed draw); --search "mamba stochastic rounding" -> the same set plus [Mamba] Unify FI ReplaySSM STP/MTP/Prefix Caching Lifecycle #55688, [Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 speculative decode #49847, Gdn ucache backend #50380, [Perf][MoE] Eliminate staging from NCCL symmetric reduce-scatter #49194 (unrelated to the seed draw); --search "ssu seed prefetch" -> no results. No open PR implements this change.

AI assistance: Claude Code was used to port this change, write the tests, run the validation and draft this description. The human submitter reviews every changed line before this leaves draft.

---

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.

• 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.

• 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

[pst2154]

Thanks. It doesn't change the randomness: the seeds are the same torch.randint(0, 2**32, (1,)) calls on the default CUDA generator, one per Mamba-2 layer in layer order, only in steps with decode tokens, so each layer gets exactly the value it drew inline before (unit-tested eagerly and against draws replayed from a CUDA graph). Nothing else draws from that generator in between (MRv2 sampling uses per-request seeds).

Reworked locally as you suggested: the Mamba-2 metadata builder draws them once per step into a persistent buffer (stable address for graph replay), on the current stream, with no side stream and no events. The metadata carries per-layer views that the FlashInfer, Triton and ReplaySSM paths take as rand_seed= (falling back to the inline draw). The gain is small as-is (about -0.7 % per verify iteration, close to noise), because the per-layer draw kernels remain, just moved ahead of the forward. I'm checking whether a single kernel can produce the same values in one launch; if not, I'll likely drop this one. Will update.

[pst2154]

Dropping this one: after moving the draws into the metadata build (no events), the remaining gain is within noise (about -0.7 % per verify iteration), so it isn't worth the added complexity. Thanks for the review.