[#12359][feat] AutoDeploy: Support SSM replay kernel for MTP with FlashInfer

[galagam]

Summary by CodeRabbit

Release Notes

• New Features

  • Added MTP speculative decoding configurations for SuperV3 model (single-GPU and multi-GPU setups).
  • Introduced SSM replay state-update capability for improved speculative decoding efficiency.

• Improvements

  • Updated SuperV3 + MTP configuration with optimized throughput-oriented limits.
  • Enhanced Mamba SSM kernel compatibility and cache handling.

• Tests

  • Added smoke test for SuperV3 MTP with SSM replay functionality.

Description

Integrate replay_selective_state_update (tl.dot fast-forward) into the
AD flashinfer_ssm backend as an optional path controlled by:

  insert_cached_ssm_attention:
    backend: flashinfer_ssm
    ssm_replay: true

The replay kernel replaces the FlashInfer horizontal SSM scan in the MTP extend path with the same two-kernel algorithm the pytorch backend uses:

• _replay_precompute_kernel: computes CB_scaled/decay_vec from cached B,dt,A
• _replay_state_update_kernel: fast-forwards SSM state via tl.dot on cached values, then computes output using new x/C and precomputed tensors

Test Coverage

Unit smoke (tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py):

• test_super_mtp_smoke: existing triton_ssm path (unchanged, all pass)
• test_super_mtp_ssm_replay_smoke: flashinfer_ssm + ssm_replay=True

Accuracy (B200, ws4_180gb, trtllm backend):
nvfp4 + fp8: PASSED (GSM8K ~92%, acceptance_rate ~45%)

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.

[hnover-nv]

was there anything in particular you wanted me to review? Or helpful to look at nsys traces if we aren't seeing much perf improvement?

[galagam]

was there anything in particular you wanted me to review? Or helpful to look at nsys traces if we aren't seeing much perf improvement?

@hnover-nv mostly for general awareness and to prevent dup. work. I'll ping you again when I move it out of draft. Still working on it.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces SSM replay caching support for speculative MTP decoding. It defines replay buffer resource handlers, extends both FlashInfer and Triton backends with replay-aware execution paths, wires replay resources through the cache manager via configuration-driven transforms, and adds comprehensive test coverage with updated configurations.

Changes

SSM Replay Infrastructure and Implementation

Layer / File(s)	Summary
Replay Buffer Resource Handler Contracts tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py	Six new StateResourceHandler subclasses define shape contracts for per-layer replay state (old_x, old_B, old_dt, old_dA_cumsum) and shared replay metadata (cache_buf_idx, prev_num_accepted_tokens), each enforcing backing by MambaHybridCacheManager.
FlashInfer Backend Replay Path Implementation tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py	Extended _flashinfer_cached_ssm custom op accepts optional replay buffers; hoists slot_idx conversion to int32 once; branches extend logic to invoke replay_selective_state_update when replay buffers present, otherwise uses existing FlashInfer path. Added ssm_replay config flag to FlashinferBackendSSM and conditional replay resource handler initialization for SM ≥ 80.
Triton Causal Conv Backend Optimization tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py	Hoists slot_idx to int32 once; reorders execution to run decode before extend; slices pre-converted int32 index instead of converting per-slice; refactors extend scatter-back to use .copy_() with .permute() instead of transpose+view+assign.
Replay State Resource Binding and Cache Manager tensorrt_llm/_torch/auto_deploy/shim/interface.py, tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py	Extended CachedSequenceInterface discovers replay handler lists, conditionally binds per-layer and global replay buffers to cache manager with use_replay_state_update flag. Added six public accessor methods to MambaHybridCacheManager facade and Python implementation to expose replay tensors when enabled.
SSM Cache Transform Configuration tensorrt_llm/_torch/auto_deploy/transform/library/ssm_cache.py	Introduced SSMCacheTransformConfig with ssm_replay boolean field; expanded SSMCacheTransform to config-driven transform with get_config_class() override and conditional descriptor subclassing to set ssm_replay=True.
Supporting Transform Updates tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py, tensorrt_llm/_torch/auto_deploy/transform/library/fuse_causal_conv.py	KV cache transform builds cached-op call arguments via args tuple with optional positional mask appending; handles resource_handler=None as sentinel. Fuse-causal-conv extends pattern matching to support both CUDA and Triton backend wrappers with matched wrapper selection in replacement.
Router Predicate Fix tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py	Updated DefaultADPRouter and KVCacheAwareADPRouter to treat attention_dp_rank=None as "relaxed" for sorting, in addition to missing py_scheduling_params.
FlashInfer SSD Kernel Constraint tensorrt_llm/_torch/modules/mamba/ssd_combined.py	Added dstate == chunk_size requirement to _flashinfer_ssd_supported() predicate to avoid FlashInfer upstream tiling issues.
Configuration and Test Coverage examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml, examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml, examples/configs/curated/super_v3_mtp_ws1.yaml, examples/configs/curated/super_v3_mtp_ws4.yaml, tests/integration/defs/accuracy/test_llm_api_autodeploy.py, tests/integration/test_lists/test-db/l0_b200.yml, tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py, tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py	Updated existing config with replay enablement and lower throughput limits; added three new curated/registry configs for 1-GPU and 4-GPU setups with MTP + replay. Updated test helper to use max_batch_size-keyed CUDA graph config. Added comprehensive smoke test for MTP with flashinfer_ssm and ssm_replay=True; updated existing custom-op test signature. Registered new smoke test in B200 pre-merge list.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#13711: Modifies the same _flashinfer_cached_ssm and FlashinferBackendSSM backend to implement alternative MTP extend path using FlashInfer kernel instead of replay, affecting overlapping code regions.

Suggested reviewers

• QiJune
• Shixiaowei02
• StanleySun639
• govind-ramnarayan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 32.56% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and concisely describes the main change: adding SSM replay kernel support for MTP speculative decoding with FlashInfer backend in AutoDeploy.
Description check	✅ Passed	The PR description adequately explains what the change does (integrating replay_selective_state_update), why it matters (optional path for MTP extend), includes test coverage details, and has a completed checklist.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[QiJune]

LGTM

[galagam]

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1,DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

[tensorrt-cicd]

PR_Github #51501 [ run ] triggered by Bot. Commit: 97788c2 Link to invocation

[tensorrt-cicd]

PR_Github #51501 [ run ] completed with state SUCCESS. Commit: 97788c2
/LLM/main/L0_MergeRequest_PR pipeline #40906 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[galagam]

/bot run

[tensorrt-cicd]

PR_Github #51572 [ run ] triggered by Bot. Commit: ae7c388 Link to invocation

[tensorrt-cicd]

PR_Github #51572 [ run ] completed with state SUCCESS. Commit: ae7c388
/LLM/main/L0_MergeRequest_PR pipeline #40965 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[galagam]

/bot run

[galagam]

/bot run

[tensorrt-cicd]

PR_Github #51615 [ run ] triggered by Bot. Commit: ae7c388 Link to invocation

[tensorrt-cicd]

PR_Github #51615 [ run ] completed with state SUCCESS. Commit: ae7c388
/LLM/main/L0_MergeRequest_PR pipeline #41001 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[galagam]

/bot run

[tensorrt-cicd]

PR_Github #51760 [ run ] triggered by Bot. Commit: ae7c388 Link to invocation

[tensorrt-cicd]

PR_Github #51760 [ run ] completed with state SUCCESS. Commit: ae7c388
/LLM/main/L0_MergeRequest_PR pipeline #41132 completed with status: 'SUCCESS'

CI Report

Link to invocation