[#12359][feat] AutoDeploy: MTP performance: Integrate FI kernel for extend path - #13711
Conversation
📝 WalkthroughWalkthroughThis PR extends the Flashinfer SSM backend for speculative decoding with multi-token (MTP) support by introducing intermediate state caching during the extend phase. The config shifts from Triton to Flashinfer SSM, the custom op gains an "EXTEND" execution path with intermediate state buffering, cache initialization is expanded, and tests are parametrized across model variants (bf16, FP8, NVFP4). ChangesFlashinfer SSM with Intermediate State Caching
Sequence DiagramsequenceDiagram
participant Test as Test / AutoDeploy
participant SSMOp as flashinfer_cached_ssm<br>Custom Op
participant FI as Flashinfer Library
participant Cache as intermediate_ssm_state_cache
Test->>SSMOp: Call with extend_inputs present
activate SSMOp
SSMOp->>SSMOp: Compute num_prefill, num_extend, num_decode
SSMOp->>SSMOp: Prefill path (standard)
Note over SSMOp: Process prefill tokens
alt EXTEND Phase (new)
SSMOp->>SSMOp: Group extend inputs by sequence
SSMOp->>FI: _flashinfer_ssm_update<br>(disable_state_update=True)
FI->>FI: Compute SSM states<br>for extend tokens
FI->>Cache: Write intermediate states
Cache-->>SSMOp: Confirm buffered
Note over SSMOp: Store intermediate states for<br>next decode round
end
SSMOp->>SSMOp: Decode path<br>(offset by prefill + extend)
SSMOp->>FI: _flashinfer_ssm_update<br>(standard decode)
FI->>Cache: Read intermediate states if needed
Cache-->>FI: Retrieve buffered states
FI->>SSMOp: Return decode outputs
SSMOp->>Test: Emit full (prefill+extend+decode)<br>output
deactivate SSMOp
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 👉 Get your free trial and get 200 agent minutes per Slack user (a $50 value). Review rate limit: 9/10 reviews remaining, refill in 6 minutes. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/accuracy/test_llm_api_autodeploy.py (1)
687-709:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSkip the new
nvfp4parameter on pre-Blackwell GPUs.
@skip_pre_hopperis not enough for the addedmodel_id="nvfp4"case. Hopper will still run this test, but NVFP4 is Blackwell-only elsewhere in this file, so this parameterization can now fail on supported non-Blackwell runners.Suggested guard
def test_mtp(self, world_size, attn_backend, model_id): if get_device_count() < world_size: pytest.skip(f"Not enough devices for world_size={world_size}") + if model_id == "nvfp4" and get_sm_version() < 100: + pytest.skip("NVFP4 requires Blackwell or later") # bf16 120B model requires at least 4 GPUs if model_id == "bf16" and world_size < 4: pytest.skip("bf16 Super V3 requires at least 4 GPUs")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py` around lines 687 - 709, The new parameter model_id="nvfp4" needs an explicit guard to skip on pre-Blackwell hardware; update the test_mtp function to detect when model_id == "nvfp4" and call pytest.skip (or use the existing skip helper) when Blackwell GPUs are not available (similar to the existing get_device_count() / bf16 guard logic); locate the test_mtp function and add the nvfp4 check so the nvfp4 case is only run on Blackwell-capable runners.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml`:
- Around line 55-56: Update the nearby comment above the
insert_cached_ssm_attention block to reflect that flashinfer_ssm is now a
supported speculative Mamba-cache backend (in addition to Triton) instead of
stating Triton is the only backend; locate the insert_cached_ssm_attention entry
and change the comment text to reference flashinfer_ssm (and any other supported
backends) so the comment matches the config selection.
In `@tests/integration/test_lists/test-db/l0_dgx_b200.yml`:
- Around line 339-342: The pre-merge autodeploy block in
tests/integration/test_lists/test-db/l0_dgx_b200.yml is missing the MTP
perf-sanity entries; add the same MTP test entries that appear in the post-merge
block into the authoritative pre-merge autodeploy list so the rewritten MTP
extend-path kernel is perf-tested before merge (add lines for
accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp with the
ws4_180gb-flashinfer-bf16, ws4_180gb-trtllm-bf16, ws4_180gb-trtllm-fp8 and
ws4_180gb-trtllm-nvfp4 variants into the pre-merge/autodeploy section).
---
Outside diff comments:
In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py`:
- Around line 687-709: The new parameter model_id="nvfp4" needs an explicit
guard to skip on pre-Blackwell hardware; update the test_mtp function to detect
when model_id == "nvfp4" and call pytest.skip (or use the existing skip helper)
when Blackwell GPUs are not available (similar to the existing
get_device_count() / bf16 guard logic); locate the test_mtp function and add the
nvfp4 check so the nvfp4 case is only run on Blackwell-capable runners.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f04452ad-f943-4201-8271-31cccea9b573
📒 Files selected for processing (7)
examples/auto_deploy/model_registry/configs/super_v3_mtp.yamltensorrt_llm/_torch/auto_deploy/custom_ops/mamba/__init__.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.yml
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #46935 [ run ] triggered by Bot. Commit: |
|
PR_Github #46935 [ run ] completed with state
|
govind-ramnarayan
left a comment
govind-ramnarayan
left a comment
There was a problem hiding this comment.
Looks good, thanks!
d6f93d3 to
f92214c
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
1 similar comment
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #47141 [ run ] triggered by Bot. Commit: |
|
PR_Github #47178 [ run ] triggered by Bot. Commit: |
|
PR_Github #47141 [ run ] completed with state |
|
PR_Github #47178 [ run ] completed with state
|
f92214c to
01c6ea3
Compare
|
/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #47604 [ run ] triggered by Bot. Commit: |
|
PR_Github #47604 [ run ] completed with state
|
01c6ea3 to
a894172
Compare
|
PR_Github #49426 [ run ] triggered by Bot. Commit: |
|
PR_Github #49426 [ run ] completed with state
|
|
/bot run |
|
PR_Github #49473 [ run ] triggered by Bot. Commit: |
|
PR_Github #49473 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49583 [ run ] triggered by Bot. Commit: |
|
/bot kill |
526ff7f to
b7b53db
Compare
|
/bot run --stage-list "DGX_H100-4_GPUs-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #49629 [ kill ] triggered by Bot. Commit: |
|
PR_Github #49634 [ run ] triggered by Bot. Commit: |
|
PR_Github #49629 [ kill ] completed with state |
|
PR_Github #49583 [ run ] completed with state |
|
/bot run --stage-list "DGX_H100-4_GPUs-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" |
|
PR_Github #49651 [ run ] triggered by Bot. Commit: |
|
PR_Github #49634 [ run ] completed with state |
|
/bot run |
|
PR_Github #49661 [ run ] triggered by Bot. Commit: |
|
PR_Github #49651 [ run ] completed with state |
…rtllm_ssm backend) Replace the Triton _selective_scan_update_kernel with FlashInfer's selective_state_update_kernel_simple_mtp for the MTP (speculative decoding) extend path in AutoDeploy. Enhance accuracy test with fp8 and nvfp4 variants. Acceptance rate threshold for bf16 tightened: 45% -> 50% (achieves ~53%) Acceptance rate for fp8/nvfp4 set to 40% due to numerical issues (achieves ~44%) The acceptance rate and GSM8K accuracy scores are consistent between Triton and FlashInfer. Triton SSM: 0.329 ms/call × 40 layers = 13.2 ms/iter (decode) FlashInfer: 0.149 ms/call × 40 layers = 5.9 ms/iter (decode) Iter speedup: 30.0 ms → 23.5 ms (1.27x) Throughput: +20% across conc=1–128 Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
b7b53db to
07a6bb9
Compare
|
/bot run |
|
PR_Github #49673 [ run ] triggered by Bot. Commit: |
|
PR_Github #49661 [ run ] completed with state |
|
PR_Github #49673 [ run ] completed with state |
… for extend path (NVIDIA#13711) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
… for extend path (NVIDIA#13711) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
4 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
Replace the Triton _selective_scan_update_kernel with FlashInfer's selective_state_update_kernel_simple_mtp for the MTP (speculative decoding) extend path in AutoDeploy.
Extend the existing flashinfer_backend_mamba.py custom op code to use the FI kernel for the extend path as well.
Prefill is still using triton.
Triton SSM: 0.329 ms/call × 40 layers = 13.2 ms/iter (decode)
FlashInfer: 0.149 ms/call × 40 layers = 5.9 ms/iter (decode)
Iter speedup: 30.0 ms → 23.5 ms (1.27x)
Throughput: +20% across conc=1–128
Test Coverage
Extend accuracy test with fp8 and nvfp4 variants.
Acceptance rate threshold for bf16 tightened: 45% -> 50% (achieves ~53%)
Acceptance rate for fp8/nvfp4 set to 40% due to numerical issues (achieves ~44%)
GSM8K score for all variants bf16/fp8/nvfp4 are similar.
The acceptance rate and GSM8K accuracy scores are consistent between Triton and FlashInfer.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.