[TRTLLM-10062][feat] Update mtp selective_state_update to flashinfer

tensorrt_llm/_torch/cute_dsl_kernels/argmax.py

@@ -597,9 +597,10 @@ def argmax(x: torch.Tensor) -> torch.Tensor:
             x: Input tensor of shape (M, N)
 
         Returns:
-            Output tensor of shape (M, 2) where:
-            - Column 0: Maximum value in each row
-            - Column 1: Index of maximum value in each row (argmax)
+            Output tensor of shape (M, 2) in float32 dtype where:
+            - Column 0: Maximum value in each row (converted to float32)
+            - Column 1: Index of maximum value in each row (argmax, stored as float32)
+
         """
         assert x.dim() == 2, "Input must be 2D"
         assert x.is_cuda, "Tensor must be on CUDA device"
@@ -609,9 +610,13 @@ def argmax(x: torch.Tensor) -> torch.Tensor:
 
         if _should_use_torch_fallback(N, x.dtype):
             max_vals, max_indices = torch.max(x, dim=-1, keepdim=True)
-            return torch.cat([max_vals, max_indices.to(x.dtype)], dim=-1)
+            # Use float32 for indices to avoid precision loss with large vocab sizes
+            return torch.cat([max_vals.to(torch.float32), max_indices.to(torch.float32)], dim=-1)
 
-        out = torch.empty((M, 2), dtype=x.dtype, device=x.device)
+        # Use float32 for output to preserve argmax index precision
+        # Float32 can exactly represent all integers up to 16M (vocab size 131072 is safe)
+        out = torch.empty((M, 2), dtype=torch.float32, device=x.device)
+        # Input dtype for the kernel (input logits)
         dtype = torch2cute_dtype_map[x.dtype]
 
         def convert_from_dlpack(tensor):

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -268,7 +268,8 @@ def forward(
         assert hidden_states_hp.shape[-1] == self.hidden_dim
         orig_shape = hidden_states_hp.shape
         hidden_states_hp_2d = hidden_states_hp.view(-1, self.hidden_dim)
-        all_rank_num_tokens = attn_metadata.all_rank_num_tokens
+        all_rank_num_tokens = kwargs.get('all_rank_num_tokens',
+                                         attn_metadata.all_rank_num_tokens)
 
         def _compute_shared_output():
             if self.shared_experts is not None:
@@ -725,6 +726,7 @@ def forward(
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None = None,
         attn_metadata: AttentionMetadata | None = None,
+        **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         if self.has_start_projections:
             assert inputs_embeds is not None
@@ -753,6 +755,7 @@ def forward(
         hidden_states = self.mixer(
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
+            **kwargs,
         )
 
         if self.has_end_norm:
@@ -859,6 +862,7 @@ def forward(
                 hidden_states=hidden_states,
                 residual=residual,
                 attn_metadata=attn_metadata,
+                all_rank_num_tokens=all_rank_num_tokens,
             )
         return hidden_states
 

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -143,19 +143,13 @@ def __init__(
         self._mamba_ssm_cache_dtype = config.quant_config.mamba_ssm_cache_dtype
         supported_head_dim_in_flashinfer = [64, 128]
         if head_dim in supported_head_dim_in_flashinfer:
-            logger.info_once(
-                "Using flashinfer for selective state update for no MTP",
-                key="selective_state_update_no_mtp")
-            self.selective_state_update_func_no_mtp = selective_state_update_fi
+            logger.info_once("Using flashinfer for selective state update",
+                             key="selective_state_update")
+            self.selective_state_update_func = selective_state_update_fi
         else:
-            logger.info_once(
-                "Using native for selective state update for no MTP",
-                key="selective_state_update_no_mtp")
-            self.selective_state_update_func_no_mtp = selective_state_update_native
-        # TODO: support MTP selective state update in flashinfer.
-        logger.info_once("Using native for selective state update for MTP",
-                         key="selective_state_update_mtp")
-        self.selective_state_update_func_mtp = selective_state_update_native
+            logger.info_once("Using native for selective state update",
+                             key="selective_state_update")
+            self.selective_state_update_func = selective_state_update_native
 
         # D
         self.D = nn.Parameter(
@@ -388,7 +382,7 @@ def forward(
             D = repeat(self.D, "h -> h p", p=self.head_dim)
             if is_target_verify:
                 intermediate_ssm_states = layer_cache.intermediate_ssm
-                self.selective_state_update_func_mtp(
+                self.selective_state_update_func(
                     ssm_states,
                     x_d.view(
                         num_decodes,
@@ -422,7 +416,7 @@ def forward(
                     intermediate_state_indices=self.intermediate_state_indices,
                 )
             else:
-                self.selective_state_update_func_no_mtp(
+                self.selective_state_update_func(
                     ssm_states,
                     x_d,
                     dt_d,

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -5905,6 +5905,70 @@ def test_nvfp4_8gpus_mtp(self):
             task.evaluate(llm,
                           extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS)
 
+    @skip_pre_blackwell
+    @pytest.mark.skip_less_device(4)
+    @pytest.mark.skip_less_device_memory(80000)
+    def test_nvfp4_4gpu_mtp_ar(self):
+        max_draft_len = 7
+        mtp_config = MTPDecodingConfig(
+            num_nextn_predict_layers=max_draft_len,
+            mtp_eagle_one_model=True,
+        )
+        model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-NVFP4-FP8KV-011526"
+
+        llm_common_config = dict(
+            model=model_path,
+            tensor_parallel_size=4,
+            moe_expert_parallel_size=4,
+            kv_cache_config=KvCacheConfig(
+                enable_block_reuse=False,
+                mamba_ssm_cache_dtype="float16",
+                free_gpu_memory_fraction=0.5,
+            ),
+            max_batch_size=4,
+            enable_attention_dp=False,
+            cuda_graph_config=CudaGraphConfig(max_batch_size=32,
+                                              enable_padding=True),
+            disable_overlap_scheduler=False,
+            moe_config=MoeConfig(backend="CUTLASS"),
+        )
+
+        llm_spec = LLM(**llm_common_config, speculative_config=mtp_config)
+
+        raw_prompts = [
+            "The capital of France is",
+            "The president of the United States is",
+            "The future of AI is",
+        ]
+        prompts = [
+            llm_spec.tokenizer.apply_chat_template([{
+                "role": "user",
+                "content": p
+            }],
+                                                   tokenize=False,
+                                                   add_generation_prompt=True)
+            for p in raw_prompts
+        ]
+        tok_ids = [llm_spec.tokenizer.encode(p) for p in prompts]
+
+        sampling_params = SamplingParams(max_tokens=128, temperature=0)
+
+        for i in range(len(tok_ids)):
+            num_tokens = 0
+            num_drafted = 0
+            num_accepted = 0
+            for output in llm_spec.generate_async(tok_ids[i],
+                                                  sampling_params,
+                                                  streaming=True):
+                new_tokens = output.outputs[0].token_ids
+                num_drafted += max_draft_len
+                num_accepted += len(new_tokens) - num_tokens - 1
+                num_tokens = len(new_tokens)
+
+            accept_rate = num_accepted / num_drafted
+            assert accept_rate > 0.2, \
+                f"Acceptance rate too low for prompt {i}: {accept_rate:.2f}"
+
 
 @skip_pre_hopper
 class TestMiniMaxM2(LlmapiAccuracyTestHarness):

tests/integration/test_lists/qa/llm_function_core.txt

@@ -285,6 +285,8 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-cpp_mamba_cache]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attention_dp_on-trtllm]
+accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpu_mtp_ar
+accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TEP4_PP2]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP8_PP1]

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

@@ -114,6 +114,7 @@ l0_dgx_b200:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] TIMEOUT (60)
   - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60)
+  - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpu_mtp_ar TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attention_dp_on-trtllm] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attention_dp_on-cutlass] TIMEOUT (60)