[jhaotingc]

Overview

Adds an opt-in FlashInfer fused SSD scan for Mamba2 prefills (--mamba-prefill-backend flashinfer)
and draws the ReplaySSM stochastic-rounding seed once per step instead of once per layer. This PR is
part of a series that brings piecewise CUDA graphs to Mamba2 prefill steps, and it is stacked on
#61497. Tested with #61293 applied, but it does not depend on it. Requires a FlashInfer build with
flashinfer-ai/flashinfer#6100, which is not in a release yet.

Claims

• Up to 1.11x output throughput (OTPS/gpu) over [Perf][Mamba] Reduce host overhead in the Mamba2 Triton SSD prefill path #61497 on 1x B200 with
  Nemotron-3.5-Lightning-NVFP4-W4A4 and MTP=2: 1.06-1.11x at concurrency 8-32, 1.05-1.06x at 64-128
  and 1.03x at 1-4.
• The FlashInfer scan takes 2 launches per layer instead of 5, and host time per scan drops from 174
  us to 97 us (214 us to 125 us for continued prefills with initial states).
• One seed launch per step instead of one per ReplaySSM layer (23 to 1 on this model), in both eager
  mixed steps and the decode CUDA graph.

Validation

End-to-end serving throughput, before vs after, on the same node, served back-to-back with identical
flags except the new one:

• Before: main + [MoE] Shape-based padding for FlashInfer TRT-LLM MXFP8 experts (gated and non-gated) #61293 + [Perf][Mamba] Reduce host overhead in the Mamba2 Triton SSD prefill path #61497. After: the same, plus this PR, with
  --mamba-prefill-backend flashinfer.
• Hardware: 1x B200, TP1.
• FlashInfer: upstream main at 2a2cceab0, which includes feat(cake_ssd_combined): <128-token calls, seq_idx guard, grouped epilogue loads, any chunk_size via cu_seqlens (SM100 / SM103) flashinfer-ai/flashinfer#6100
  (cu_seqlens support) and perf(cake_ssd_combined): preprocess, exact-scan and chunk-parallel prefill improvements (SM100 / SM103) flashinfer-ai/flashinfer#6206 (Cake SSD prefill performance).
• Workload: SPEED-Bench throughput_1k, low_entropy category (512 prompts),
  max_tokens 200, served through aiperf.
  • Warm-up: 256 requests at concurrency 128, discarded.
  • Then the full 512-prompt dataset at each concurrency from 128 down to 1.
  • Concurrency 64 and 128 were run 2 more times; those rows are the mean of 3 runs.

Values are OTPS/gpu (output tokens/s per GPU):

Concurrency	Before PR	After PR	Speedup
1	804	824	1.03x
2	1324	1366	1.03x
4	2033	2099	1.03x
8	3117	3305	1.06x
16	4694	5094	1.09x
32	7383	8189	1.11x
64	12110	12804	1.06x
128	17290	18176	1.05x

Mid concurrencies have the most prefill and mixed-batch interruptions, so they gain the most; the
effect is smaller at low and high concurrency.

Accuracy: GPQA-diamond through NeMo Gym on 1x GB300, TP1, MTP=2 (198 questions x 8 repeats, thinking on,
temperature 1.0, top_p 0.95), with the serve command below plus --max-model-len 262144 --reasoning-parser nemotron_v3. pass@1, mean over the 8 repeats ± standard error:

• No prefix reuse (as below): 76.45 ± 0.53%
• Prefix reuse (--enable-prefix-caching --mamba-cache-mode align, no ReplaySSM): 76.58 ± 0.93%

Serve and client commands

--mamba-prefill-backend flashinfer selects this PR's scan for prefills; --mamba-backend flashinfer only selects the decode state update. The Before run used the same command without
--mamba-prefill-backend.

#!/usr/bin/env bash
set -euo pipefail
MODEL=${MODEL:?path or HF id of the Nemotron-3.5-Lightning NVFP4 W4A4 checkpoint}
NUMA_NODE=${NUMA_NODE:?NUMA node local to the GPU}
CAPTURE_SIZES='[1,2,4,8,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256,272,288,304,320,336,352,368,384,400,416,432,448,464,480,496,512,528,544,560,576,592,608,624,640,656,672,688,704,720,736,752,768,784,800,816,832,848,864,880,896,912,928,944,960,976,992,1008,1024,1152,1280,1408,1536,1664,1792,1920,2048,2560,3072,3584,4096,5120,6144,7168,8192,10240,12288,14336,16384,18432,20480,22528,24576,26624,28672,30720,32768]'

vllm serve "$MODEL" \
  --served-model-name nemotron-3.5-lightning \
  --tensor-parallel-size 1 \
  --max-num-seqs 128 \
  --max-num-batched-tokens 32768 \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.9 \
  --trust-remote-code \
  --async-scheduling \
  --no-enable-prefix-caching \
  --speculative-config '{"method":"mtp","num_speculative_tokens":2,"moe_backend":"flashinfer_trtllm"}' \
  --mamba-cache-mode none \
  --mamba-backend flashinfer \
  --mamba-prefill-backend flashinfer \
  --mamba-ssm-cache-dtype float16 \
  --enable-mamba-cache-stochastic-rounding \
  --mamba-cache-philox-rounds 5 \
  --use-replayssm \
  --compilation-config "{\"mode\":3,\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"max_cudagraph_capture_size\":32768,\"cudagraph_capture_sizes\":${CAPTURE_SIZES}}" \
  --numa-bind --numa-bind-nodes "$NUMA_NODE" \
  --host 0.0.0.0 --port 8001

Client, run once per concurrency level:

aiperf profile \
  --model nemotron-3.5-lightning --tokenizer "$MODEL" --tokenizer-trust-remote-code \
  --endpoint-type chat --endpoint /v1/chat/completions --streaming --use-legacy-max-tokens \
  --url http://localhost:8001 \
  --input-file "$SPEED_BENCH_1K_LOW_ENTROPY_JSONL" \
  --extra-inputs max_tokens:200 --extra-inputs skip_special_tokens:false \
  --random-seed 1234 --use-server-token-count \
  --concurrency "$CONC" --request-count 512 --num-dataset-entries 512

Details

Mamba2 prefills run the chunked SSD scan as 5 Triton launches per layer (cumsum, chunk state, state
passing, bmm, chunk scan). Continued prefills add a gather and a torch.where for their initial
states. In eager prefill and mixed steps, these launches and their host overhead add directly to
step latency. FlashInfer's fused SSD (SSDCombined, Cake backend) runs the scan in 2 launches from
the packed-varlen cu_seqlens and derives its chunk metadata on the device.

Changes:

• --mamba-prefill-backend {triton,flashinfer} (MambaConfig.prefill_backend, default triton).
  Prefills that store Mamba checkpoints need every chunk state, so they keep the Triton scan.
• mamba_chunk_scan_flashinfer() (ops/ssd_flashinfer.py):
  • gathers initial states (one launch) only when some sequence has one, and otherwise starts from
    a zero state;
  • runs one SSDCombined runner per (stream, heads, groups, state dtype, initial states), so final
    states come out in the cache dtype;
  • scatters each final state into its cache slot.
• The gather and scatter kernels do not specialize indices_ptr on alignment. Its alignment changes
  with the number of decode rows, so this avoids new compiles while serving.
• The mixer keeps a bfloat16 copy of D for FlashInfer, refreshed when the parameter changes.
• The backend needs SM100/SM103, bfloat16 activations, head dim 64 and state size 128 (the shapes
  FlashInfer compiles). MambaConfig and MambaMixer2 reject other setups, and the SSD warm-up runs
  both FlashInfer variants.
• ReplaySSM seed: with --enable-mamba-cache-stochastic-rounding, each FlashInfer ReplaySSM layer
  drew its Philox seed with its own torch.randint. The layers now read their seeds from slots of
  one shared buffer, which the first layer refreshes with a single random_() per step. Every layer
  still gets an independent seed each step.

---

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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhaotingc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork