[jhaotingc]

Overview

Reduces host overhead in the Mamba2 Triton SSD prefill path. This PR is
part of a series that brings piecewise CUDA graphs to Mamba2 prefill steps. Tested with #61293
applied, but it does not depend on it.

• Write prefill final states into ssm_state from the SSD scan when has_checkpoints is False (4 fewer launches per Mamba layer).
• Cache the autotuned config of the Triton SSD kernels, keyed on each launch's own arguments (skips Triton's per-launch key rebuild; 43 of the 84 µs saved per Mamba layer).
• Replace einops.rearrange with flatten/unflatten (same view, about 1.5 µs less host time per call).

Claims

• Up to 1.08x output throughput (OTPS/gpu) on 1x B200 with Nemotron-3.5-Lightning-NVFP4-W4A4 and
  MTP=2: 1.04-1.08x at concurrency 4-64, 1.02x at 128 and 1.01x at 1-2.
• Host time of the prefill SSD scan plus its final-state write drops from 263 us to 179 us per Mamba
  layer, about 1.9 ms less per mixed step with 23 Mamba layers. Every prefill or mixed step also
  issues 4 fewer launches per Mamba layer.

Validation

End-to-end serving throughput, before vs after, on the same node with identical flags:

• Before: main at 8c60cf0 (this PR's base, which includes [MoE] Shape-based padding for FlashInfer TRT-LLM MXFP8 experts (gated and non-gated) #61293). After: the same, plus
  this PR.
• Hardware: 1x B200, TP1.
• 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 per serve.
• Repeats: each side was served twice, alternating (after, before, after, before). Each row is the
  mean over both serves, so the concurrency 64 and 128 rows are the mean of 6 runs.

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

Concurrency	Before PR	After PR	Speedup
1	779	785	1.01x
2	1279	1296	1.01x
4	1977	2057	1.04x
8	3062	3233	1.06x
16	4609	4961	1.08x
32	7220	7643	1.06x
64	11583	12263	1.06x
128	17050	17363	1.02x

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): 75.51 ± 0.64%
• Prefix reuse (--enable-prefix-caching --mamba-cache-mode align, no ReplaySSM): 75.25 ± 0.63%

Serve and client commands

The default Triton SSD scan runs the prefills; --mamba-backend flashinfer only selects the decode
state update.

#!/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-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

The Mamba2 mixer runs eagerly in prefill and mixed prefill+decode steps, so its host-side overhead
lands directly on step latency. This matters most in two cases:

1. Short OSL: prefill and mixed steps come often, and each one interrupts the ITL of every running
   request.
2. A fast GPU with low-precision quantization (e.g. NVFP4 on B200): the prefill compute is too short
   to hide the CPU overhead.

Before this PR, a varlen SSD scan plus the final-state write spent 263 us of host time per Mamba
layer, mostly Python launch overhead, and each layer issued 4 extra launches to write final states
back. Changes, all on the default Triton SSD path:

• Write each sequence's final state into its ssm_state slot from _state_passing_fwd_kernel,
  replacing the index casts, gather and index_put. The checkpoint path and CPU path are unchanged.
• Add launch_autotuned(), which reuses the autotuner's chosen config and skips Triton's per-launch
  key rebuild. It builds its cache key from the launch's own arguments (the values of the kernel's
  autotune key arguments plus each argument's dtype, the information Triton's key holds), so callers
  keep no key list and the selected config is unchanged.
  • Same-node A/B with (two serves) vs without it (one serve): +1.0% to +2.5% OTPS/gpu at concurrency 4-64 (+2.5% at 64) and
    flat at 1-2 and 128; it accounts for 43 of the 84 us saved per layer.
• Drop 4 redundant device guards in the SSD wrappers.
• Replace einops.rearrange with flatten/unflatten.
• Cache CudaPlatformBase.is_arch_support_pdl().

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[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

[github-actions]

✅ @jhaotingc, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[jhaotingc]

/ci run

[github-actions]

❌ @jhaotingc, A reviewer with write access must run /ci run, approve the PR, or add the ready label first.

[jhaotingc]

@tomeras91 am I too slow to run the CI? TY!

[jhaotingc]

Hi @ZJY0516 do you mind taking a look at this PR? Thank you! I saw you are doing some mamba improvement.

[ZJY0516]

I don't want to maintain a triton key list manually

[jhaotingc]

Hi @ZJY0516, now there's no maintaining of a cache key. Launched_autotune derive key from launch.

I also tried VLLM_TRITON_FORCE_FIRST_CONFIG=1 to skip the rebuild but it lead to bad triton config being used.

Would you mind giving PR a look again? TY.

This is for a Nemotron Lightning ckpt.

Agent summary of the issue:

On every launch, Triton's Autotuner.run rebuilds its config cache key from scratch, checking each argument name against the key list (O(n²)) and calling str(dtype) on every tensor argument, which costs several µs of host time per launch even after the config is cached. This is tracked in triton-lang/triton#1472, which is still open; triton-lang/triton#10205 tried to fix it but was closed without merging, and Triton main still does this.

[ZJY0516]

LGTM now

[ZJY0516]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #96086 for commit b0919fffc497.

⚠️ This PR is 25 commits behind upstream main. Running CI at your own risk because --allow-stale was requested; outdated CI configuration may cause failures. Before merging, merge or rebase onto the latest main, then rerun /ci run on the latest PR commit.