{{ message }}
Repository navigation
[Mamba] Add FlashInfer ReplaySSM support for MTP - #52928
Merged
benchislett merged 30 commits intoSep 30, 2026
Merged
Conversation
3 of 6 tasks
askliar
marked this pull request as ready for review
August 19, 2026 15:13
askliar
requested review from
AndreasKaratzas,
LucasWilkinson,
MatthewBonanni,
ProExpertProg,
WoosukKwon,
heheda12345,
houseroad,
ivanium,
mgoin,
njhill,
robertgshaw2-redhat,
tdoublep,
tlrmchlsmth,
tomeras91,
yewentao256,
youkaichao and
zyongye
as code owners
August 19, 2026 15:13
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Member
|
/ci run |
|
✅ @askliar, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #86434 for commit |
Contributor
|
Hi @askliar, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
Member
|
For new FlashInfer test files, please adda global pytest skipif statement on the top of the files if platform is not cuda and if flashinfer is not found. Otherwise, some CI resources are going to be spent evaluating each function separately before they are skipped. It also makes the test files' purpose clear "Not supported since this is a FlashInfer related test file" or something for msg. |
|
✅ Triggered Buildkite CI #91968 for commit |
Member
|
/ci run |
benchislett
enabled auto-merge (squash)
September 30, 2026 15:03
|
❌ This PR is 40 commits behind upstream |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #92186 for commit |
Collaborator
CI selector (shadow): 62 test steps (109 jobs) instead of 88 (132 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (62)
Would skip (today's rules run them) (45)
Would add (today's rules do not run them) (19)
AMD mirrors: would skip (52)
AMD mirrors: would add (9)
CI results (2026-09-30 16:56 UTC)217 passed, 1 failed, 19 pending. No misses: the selector would have run every failed job.
18 changed files · base |
Contributor
Author
|
/ci retry |
askliar
pushed a commit
to askliar/vllm
that referenced
this pull request
Oct 1, 2026
Port the unified STP/MTP and align-prefix lifecycle onto current main after vllm-project#52928. Preserve upstream replay boundaries, runner interfaces, and removal of all mode. Keep accepted-state ownership and materialization in the model lifecycle. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar <askliar@nvidia.com>
njhill
added a commit
that referenced
this pull request
Oct 3, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 13, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 13, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 13, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 13, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 17, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>
njhill
added a commit
that referenced
this pull request
Oct 23, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 23, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 26, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 27, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 28, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 28, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 28, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 28, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
added a commit
that referenced
this pull request
Oct 28, 2026
…y held weights - Upstream #53934 lets Elastic EP run on Model Runner V2, so the extensible KV cache would now be enabled with it by default. Scale-up engines skip warmup and size the KV cache from the profiled estimate the running engines hand over, which under the extensible cache is the full reserved capacity; nothing would measure it before the final commit. Treat elastic EP as unsupported, as it effectively was when it forced the V1 runner. - Upstream #57298 subtracts weights held by an external loader process from the memory budget. Recompute the budget with them when the driver probe disables the extensible cache. - The ReplaySSM warmup test stub provides `max_num_batched_tokens`, which upstream #52928 now reads for the autotune prefill. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This stacked follow-up extends the native FlashInfer ReplaySSM integration from single-token decode in #52506 to multi-token prediction (MTP) verification for Mamba2 models.
With MTP, the model evaluates one target token plus several draft tokens before sampling determines how many were accepted. ReplaySSM avoids materializing a separate recurrent state for every speculative token: each layer keeps a checkpoint and a ring of inputs, the accepted length advances a shared cursor before the next forward, and the checkpoint is updated only when the next verification window would overflow the ring.
Dependencies
Important
This PR is stacked on #52506 and must not merge independently. Until #52506 lands, GitHub shows the cumulative base and MTP changes in this diff.
The FlashInfer path also depends on the compatible ReplaySSM API and native autotuning work in flashinfer-ai/flashinfer#4545.
Design
(batch, tokens, ...)inputs so FlashInfer can select its native MTP tactics. Variable-length batches use the packedcu_seqlensinterface.checkpointing_ssuAPI and native autotuning. PDL is forwarded explicitly but defaults off because it was not a consistent end-to-end win.This PR does not modify FlashInfer or add prefix-cache alignment support. ReplaySSM continues to use
mamba_cache_mode=none.Correctness
Matched GSM8K evaluation on GB200 compared FlashInfer ReplaySSM with the exact #49847 Triton implementation:
The comparison used FP32 SSM state with stochastic rounding disabled. A strict H100 B16/AL1 CUDA-graph control also produced identical complete output hashes across the two backends.
Actual-model throughput
Workload:
NVIDIA/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4H100 80GB
GB200
The allocation identified the device as
NVIDIA GB200.PDL was harmful at GB200 B16 and effectively neutral at larger batches. The production path therefore keeps it disabled while preserving the explicit argument for future FlashInfer tuning and debugging.
Production-profile findings
Nsight Systems profiles used 64 steady CUDA-graph decode steps and summed ReplaySSM work over all 23 Mamba layers:
The eager microbenchmark showed a much larger FlashInfer advantage because the Triton path issued several kernels from Python. CUDA-graph replay removes most of that dispatch and inter-launch gap. In production graphs, B16 favors the Triton kernels while B128 is effectively tied; the shared FlashInfer tracker is consistently cheaper.
On GB200 at B256, ReplaySSM accounted for approximately 16% of execute-step GPU time. Its 1.56× in-model improvement predicts a maximum whole-step gain of roughly 6.1%, close to the measured 5.65%. This explains why a much larger isolated-kernel ratio becomes a mid-single-digit end-to-end improvement.
Cache capacity
After decoupling verification length from physical state-page allocation, both backends admit the full requested concurrency:
This removes the previous four-state-pages-per-request over-allocation and brings FlashInfer capacity within 2% of #49847.
Validation