[#12359][feat] AutoDeploy: MTP performance: Integrate FI kernel for extend path

examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml

@@ -9,7 +9,7 @@ attn_backend: trtllm
 model_factory: AutoModelForCausalLM
 skip_loading_weights: false
 cuda_graph_config:
-  batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 320, 384]
+  batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128]
 kv_cache_config:
   # tunable mamba cache dtype
   # --> use float32 for accuracy and default (auto) for speed
@@ -50,10 +50,10 @@ transforms:
   fuse_mamba_a_log:
     stage: post_load_fusion
     enabled: true
-  # Triton SSM + causal conv are required for MTP as currently they are the only backends
-  # that support speculative mamba state caching.
+  # FlashInfer SSM is used for the MTP extend path; Triton causal conv is still required
+  # for speculative Mamba state caching in this configuration.
   insert_cached_ssm_attention:
-    backend: triton_ssm
+    backend: flashinfer_ssm
   insert_cached_causal_conv:
     backend: triton_causal_conv
   fuse_nvfp4_moe:

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py

@@ -16,19 +16,30 @@
 from typing import List, Optional
 
 import torch
+from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update
 from torch.fx import Node
 
 from ..._compat import KvCacheConfig
-from ..attention_interface import AttentionRegistry, BatchInfo, MHACallable, ResourceHandlerDict
+from ..attention_interface import (
+    AttentionRegistry,
+    BatchInfo,
+    MHACallable,
+    ResourceHandlerDict,
+    SpecSSMResourceHandler,
+)
 from .mamba_backend_common import (
     BaseBackendSSM,
     _flatten_ssm_inputs,
     _prepare_ssm_decode_inputs,
+    _prepare_ssm_grouped_state_update_inputs,
     _run_ssm_prefill,
 )
 
 
-@torch.library.custom_op("auto_deploy::flashinfer_cached_ssm", mutates_args=("ssm_state_cache",))
+@torch.library.custom_op(
+    "auto_deploy::flashinfer_cached_ssm",
+    mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"),
+)
 def _flashinfer_cached_ssm(
     # INPUTS (dense but may be flattened across sequences)
     hidden_states: torch.Tensor,  # [b, s, num_heads, head_dim]
@@ -50,6 +61,9 @@ def _flashinfer_cached_ssm(
     seq_idx_prefill: torch.Tensor,  # [1, num_prefill_tokens]
     # CACHES
     ssm_state_cache: torch.Tensor,  # [max_batch_size, num_heads, head_dim, ssm_state_size]
+    intermediate_ssm_state_cache: Optional[
+        torch.Tensor
+    ],  # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]
     # CONSTANTS
     time_step_limit: List[float],
     chunk_size: int,
@@ -60,9 +74,10 @@ def _flashinfer_cached_ssm(
     )
     ssm_state_size = B.shape[3]
     batch_info = BatchInfo(batch_info_host)
-    num_prefill, _, num_decode = batch_info.get_num_sequences()
-    num_prefill_tokens, _, num_decode_tokens = batch_info.get_num_tokens()
-    num_total_tokens = num_prefill_tokens + num_decode_tokens
+    num_prefill, num_extend, num_decode = batch_info.get_num_sequences()
+    num_prefill_tokens, num_extend_tokens, num_decode_tokens = batch_info.get_num_tokens()
+    num_total_tokens = num_prefill_tokens + num_extend_tokens + num_decode_tokens
+
     if out is not None:
         preallocated_ssm_out = out.view(bs, num_heads, head_dim)
     else:
@@ -72,6 +87,7 @@ def _flashinfer_cached_ssm(
             device=hidden_states.device,
         )
 
+    # PREFILL
     _run_ssm_prefill(
         hs_flat,
         B_flat,
@@ -94,6 +110,71 @@ def _flashinfer_cached_ssm(
         preallocated_ssm_out[:num_prefill_tokens].unsqueeze(0),
     )
 
+    # EXTEND: multi-token MTP verification path, writes intermediate SSM states
+    extend_inputs = _prepare_ssm_grouped_state_update_inputs(
+        hs_flat,
+        B_flat,
+        C_flat,
+        dt_flat,
+        A,
+        D,
+        dt_bias,
+        slot_idx,
+        seq_start=num_prefill,
+        token_start=num_prefill_tokens,
+        num_seq=num_extend,
+        num_tokens=num_extend_tokens,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        ssm_state_size=ssm_state_size,
+    )
+
+    if extend_inputs is not None:
+        tokens_per_extend = num_extend_tokens // num_extend
+        if intermediate_ssm_state_cache.size(1) < tokens_per_extend:
+            raise RuntimeError(
+                "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small "
+                f"for extend branch (size1={intermediate_ssm_state_cache.size(1)}, "
+                f"tokens_per_extend={tokens_per_extend})"
+            )
+        (
+            slot_idx_extend,
+            x_extend,
+            B_extend,
+            C_extend,
+            dt_extend,
+            A_full,
+            D_full,
+            dt_bias_hp,
+        ) = extend_inputs
+
+        preallocated_ssm_out_e = preallocated_ssm_out[
+            num_prefill_tokens : num_prefill_tokens + num_extend_tokens
+        ].view(num_extend, tokens_per_extend, num_heads, head_dim)
+
+        intermediate_state_indices = torch.arange(
+            num_extend, dtype=torch.int32, device=slot_idx_extend.device
+        )
+        _flashinfer_ssm_update(
+            ssm_state_cache,
+            x_extend,
+            dt_extend,
+            A_full,
+            B_extend,
+            C_extend,
+            D_full,
+            z=None,
+            dt_bias=dt_bias_hp,
+            dt_softplus=True,
+            state_batch_indices=slot_idx_extend.to(torch.int32),
+            out=preallocated_ssm_out_e,
+            disable_state_update=True,
+            intermediate_states_buffer=intermediate_ssm_state_cache,
+            cache_steps=tokens_per_extend,
+            intermediate_state_indices=intermediate_state_indices,
+        )
+
+    # DECODE: single-token autoregressive path
     decode_inputs = _prepare_ssm_decode_inputs(
         hs_flat,
         B_flat,
@@ -103,8 +184,8 @@ def _flashinfer_cached_ssm(
         D,
         dt_bias,
         slot_idx,
-        num_prefill,
-        num_prefill_tokens,
+        num_prefill + num_extend,
+        num_prefill_tokens + num_extend_tokens,
         num_decode,
         num_decode_tokens,
         num_heads,
@@ -124,28 +205,23 @@ def _flashinfer_cached_ssm(
             D_full,
         ) = decode_inputs
 
-        import flashinfer
-
-        # FlashInfer needs contiguous x/B/C with 128-byte alignment.
-        x_decode = x_decode.contiguous()
-        B_decode = B_decode.contiguous()
-        C_decode = C_decode.contiguous()
-
         slot_idx_decode_i32 = slot_idx_decode.to(torch.int32)
-        y_decode = flashinfer.mamba.selective_state_update(
+        y_decode = _flashinfer_ssm_update(
             ssm_state_cache,
             x_decode,
             dt_hp,
             A_full,
             B_decode,
             C_decode,
-            D=D_full,
+            D_full,
             z=None,
             dt_bias=dt_bias_hp,
             dt_softplus=True,
             state_batch_indices=slot_idx_decode_i32,
         )
-        preallocated_ssm_out[num_prefill_tokens:num_total_tokens].copy_(y_decode)
+        preallocated_ssm_out[num_prefill_tokens + num_extend_tokens : num_total_tokens].copy_(
+            y_decode
+        )
 
     if out is not None:
         # out is reused across CUDA graph replays with varying num_total_tokens,
@@ -179,6 +255,9 @@ def _flashinfer_cached_ssm_fake(
     seq_idx_prefill: torch.Tensor,  # [1, num_prefill_tokens]
     # CACHES
     ssm_state_cache: torch.Tensor,  # [max_batch_size, num_heads, head_dim, ssm_state_size]
+    intermediate_ssm_state_cache: Optional[
+        torch.Tensor
+    ],  # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]
     # CONSTANTS
     time_step_limit: List[float],
     chunk_size: int,
@@ -218,4 +297,7 @@ def get_cache_initializers(
                 "Consider using 'triton_ssm' backend instead."
             )
 
+        ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base(
+            ret["ssm_state_cache"]
+        )
         return ret

tests/integration/defs/accuracy/references/gsm8k.yaml

@@ -439,6 +439,10 @@ nvidia/Nemotron-Super-V3:
     accuracy: 80.85
   - spec_dec_algo: MTP
     accuracy: 92.70
+  - quant_algo: FP8
+    kv_cache_quant_algo: FP8
+    spec_dec_algo: MTP
+    accuracy: 92.61
   - quant_algo: MIXED_PRECISION
     kv_cache_quant_algo: FP8
     spec_dec_algo: MTP

tests/integration/defs/accuracy/test_llm_api_autodeploy.py

@@ -706,26 +706,62 @@ def test_functional_small(self, dtype):
     @skip_pre_hopper
     @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"])
     @pytest.mark.parametrize(
-        "world_size",
+        "model_id, world_size",
         [
             pytest.param(
+                "bf16",
                 4,
-                marks=pytest.mark.skip_less_device_memory(180000),
-                id="ws4_180gb",
+                marks=[
+                    pytest.mark.skip_less_device(4),
+                    pytest.mark.skip_less_device_memory(180000),
+                ],
+                id="bf16_ws4_180gb",
             ),
             pytest.param(
+                "fp8",
+                4,
+                marks=[
+                    pytest.mark.skip_less_device(4),
+                    pytest.mark.skip_less_device_memory(80000),
+                ],
+                id="fp8_ws4_80gb",
+            ),
+            pytest.param(
+                "nvfp4",
+                4,
+                marks=[
+                    pytest.mark.skip_less_device(4),
+                    pytest.mark.skip_less_device_memory(80000),
+                    skip_pre_blackwell,
+                ],
+                id="nvfp4_ws4_80gb",
+            ),
+            pytest.param(
+                "fp8",
+                8,
+                marks=[
+                    pytest.mark.skip_less_device(8),
+                    pytest.mark.skip_less_device_memory(80000),
+                ],
+                id="fp8_ws8_80gb",
+            ),
+            pytest.param(
+                "nvfp4",
                 8,
-                marks=pytest.mark.skip_less_device_memory(80000),
-                id="ws8_80gb",
+                marks=[
+                    pytest.mark.skip_less_device(8),
+                    pytest.mark.skip_less_device_memory(80000),
+                    skip_pre_blackwell,
+                ],
+                id="nvfp4_ws8_80gb",
             ),
         ],
     )
-    def test_mtp(self, world_size, attn_backend):
-        if get_device_count() < world_size:
-            pytest.skip(f"Not enough devices for world_size={world_size}")
+    def test_mtp(self, world_size, attn_backend, model_id):
 
-        model_path = self.MODEL_PATHS["bf16"]
+        model_path = self.MODEL_PATHS[model_id]
         kwargs = {}
+        # TODO: gate for bf16 only after replay lands
         low_memory_overrides(
             kwargs,
             max_batch_size=8,
@@ -740,8 +776,8 @@ def test_mtp(self, world_size, attn_backend):
             kwargs["compile_backend"] = "torch-simple"
 
         print(
-            f"SuperV3 MTP params: world_size={world_size}, model_path={model_path}"
-        )
+            f"SuperV3 MTP params: model_id={model_id}, world_size={world_size}, "
+            f"model_path={model_path}")
         print(f"kwargs: {kwargs}")
 
         mtp_yaml = str(
@@ -759,9 +795,17 @@ def test_mtp(self, world_size, attn_backend):
                 enable_iter_perf_stats=True,
                 **kwargs,
         ) as llm:
+            _set_quant_config(llm, model_id)
+            if model_id == "nvfp4":
+                llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
+            print_memory_usage("after engine build")
+
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)
-            self.check_acceptance_rate(llm, min_acceptance_rate=0.45)
+            # bf16 acceptance is stable; fp8/nvfp4 have higher variance due to
+            # arithmetic rounding, so use a lower threshold for quantized models.
+            min_rate = 0.50 if model_id == "bf16" else 0.40
+            self.check_acceptance_rate(llm, min_acceptance_rate=min_rate)
 
         print_memory_usage("after evaluation")
 

tests/integration/test_lists/qa/llm_function_core.txt

@@ -138,8 +138,10 @@ accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-1-
 accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm]
 accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[bf16]
 accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_functional_small[fp8]
-accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-flashinfer]
-accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-trtllm]
+accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer]
+accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm]
+accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]
+accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm]
 accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_auto_dtype[False]
 accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_auto_dtype[True]
 accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_fp8[True]

tests/integration/test_lists/test-db/l0_dgx_b200.yml

@@ -357,8 +357,9 @@ l0_dgx_b200:
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_off-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_on-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm]
-  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-flashinfer]
-  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[ws4_180gb-trtllm]
+  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer]
+  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm]
+  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-4]
   - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-NVFP4-True]
   # ------------- AutoDeploy Perf Sanity ---------------

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -374,6 +374,7 @@ l0_dgx_h100:
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_on-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_off-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_on-trtllm]
+  - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm]
   - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_attention_dp[4]
   - accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4]
   - accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16

tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py

@@ -104,6 +104,7 @@ def test_flashinfer_decode_matches_triton(mamba_env):
         None,  # seq_idx_prefill
         # CACHES
         ssm_state_cache_flashinfer,
+        None,  # intermediate_ssm_state_cache (not used in decode-only path)
         # CONSTANTS
         time_step_limit,
         chunk_size,