From 84e6e8e18ad1486cc476e13f2a9ca66fed5f06d4 Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 9 Feb 2026 06:24:23 +0000 Subject: [PATCH 1/8] update mtp selective_state_update to flashinfer Signed-off-by: qgai --- requirements.txt | 2 +- .../_torch/modules/mamba/mamba2_mixer.py | 22 +++++++------------ 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/requirements.txt b/requirements.txt index 8c44a008a214..ceb164a723b8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -54,7 +54,7 @@ ordered-set peft patchelf einops -flashinfer-python~=0.6.2 +flashinfer-python~=0.6.3 opencv-python-headless xgrammar==0.1.25 llguidance==0.7.29 diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 192e304419d1..a0130a812be7 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -139,19 +139,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( @@ -368,7 +362,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, @@ -402,7 +396,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, From fae2f45173af1301db393d51a7eaf381c2a10e7c Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 26 Feb 2026 13:47:29 +0000 Subject: [PATCH 2/8] add ar test Signed-off-by: qgai --- .../_torch/cute_dsl_kernels/argmax.py | 15 +++-- .../defs/accuracy/test_llm_api_pytorch.py | 63 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py b/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py index 6c3a635e5a23..d5380e5ca071 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py +++ b/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): diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 84c21313bae3..449f1683a756 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5802,6 +5802,69 @@ 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 = [ + "Below is a list of sentences. I'd like you to translate a few of them into Spanish please:\n\n1. The water cycle, an essential process for life on Earth, involves the continuous movement of water through evaporation, condensation, precipitation, and runoff.\n\n2. The phenomenon of bioluminescence, where living organisms produce light, creates magical scenes in the depths of the ocean and in the night landscapes.\n\n3. Black holes, the mysterious cosmic phenomena where gravity is so strong that not even light can escape, serve as gateways to understanding the limits of our physical laws.\n\n4. The exploration of quantum superposition has led to the conceptualization of quantum bits or qubits, which, unlike their classical counterparts, can represent a 0, a 1, or any quantum superposition of these states, a property that allows quantum computers to perform complex calculations at speeds unattainable by classical computers, offering new horizons in drug discovery, materials science, and cryptography, where they could solve problems considered intractable for traditional computing systems.\n\n5. The mystery of the Tunguska event, a massive explosion in Siberia in 1908, thought to be caused by a comet or asteroid, remains one of the 20th century's great enigmas.\n\nPlease translate the following numbered sentences into Spanish: [1, 2, 3, 4, 5]. Don't repeat the sentences in English. Only translate those 5 sentences. Number them in your response. Thank you!", + "Below is a list of sentences. I'd like you to translate a few of them into French please:\n\n1. The intricate dance of planets around stars in distant solar systems, known as exoplanets, expands our horizons in the search for extraterrestrial life.\n\n2. The development of smart cities, using technology to improve the efficiency of services and meet residents' needs, represents a new frontier in urban planning.\n\n3. The melting of polar ice caps, accelerated by global warming, contributes to rising sea levels and the loss of habitat for species like the polar bear.\n\n4. The mysteries of dark matter and dark energy, making up most of the universe's mass and energy, challenge our understanding of the cosmos.\n\n5. The Great Barrier Reef, the world's largest coral reef system, is home to a vast array of marine species and is visible from space.\n\nPlease translate the following numbered sentences into French: [1, 2, 3, 4, 5]. Don't repeat the sentences in English. Only translate those 5 sentences. Number them in your response. Thank you!", + ] + 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.35, \ + f"Acceptance rate too low for prompt {i}: {accept_rate:.2f}" + @skip_pre_hopper class TestMiniMaxM2(LlmapiAccuracyTestHarness): From 0daab8451d1ec0d3d3180c4c74b6af5e0ba52d2a Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 26 Feb 2026 14:35:38 +0000 Subject: [PATCH 3/8] add test list Signed-off-by: qgai --- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 1bb60496a325..24c941bcf46b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -81,6 +81,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[attention_dp_on-trtllm] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attention_dp_on-cutlass] TIMEOUT (60) - condition: From 7599578f29eaa9b5797d51bdeea6096d579ca10b Mon Sep 17 00:00:00 2001 From: qgai Date: Thu, 26 Feb 2026 14:41:28 +0000 Subject: [PATCH 4/8] change requirements.txt Signed-off-by: qgai --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index ceb164a723b8..77df6e3a72ab 100644 --- a/requirements.txt +++ b/requirements.txt @@ -54,7 +54,7 @@ ordered-set peft patchelf einops -flashinfer-python~=0.6.3 +flashinfer-python==0.6.4 opencv-python-headless xgrammar==0.1.25 llguidance==0.7.29 From 87b44abd9872741b16d5d642968962d183116346 Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 27 Feb 2026 02:59:16 +0000 Subject: [PATCH 5/8] update unitest Signed-off-by: qgai --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 449f1683a756..36a0695bbe39 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5833,8 +5833,9 @@ def test_nvfp4_4gpu_mtp_ar(self): llm_spec = LLM(**llm_common_config, speculative_config=mtp_config) raw_prompts = [ - "Below is a list of sentences. I'd like you to translate a few of them into Spanish please:\n\n1. The water cycle, an essential process for life on Earth, involves the continuous movement of water through evaporation, condensation, precipitation, and runoff.\n\n2. The phenomenon of bioluminescence, where living organisms produce light, creates magical scenes in the depths of the ocean and in the night landscapes.\n\n3. Black holes, the mysterious cosmic phenomena where gravity is so strong that not even light can escape, serve as gateways to understanding the limits of our physical laws.\n\n4. The exploration of quantum superposition has led to the conceptualization of quantum bits or qubits, which, unlike their classical counterparts, can represent a 0, a 1, or any quantum superposition of these states, a property that allows quantum computers to perform complex calculations at speeds unattainable by classical computers, offering new horizons in drug discovery, materials science, and cryptography, where they could solve problems considered intractable for traditional computing systems.\n\n5. The mystery of the Tunguska event, a massive explosion in Siberia in 1908, thought to be caused by a comet or asteroid, remains one of the 20th century's great enigmas.\n\nPlease translate the following numbered sentences into Spanish: [1, 2, 3, 4, 5]. Don't repeat the sentences in English. Only translate those 5 sentences. Number them in your response. Thank you!", - "Below is a list of sentences. I'd like you to translate a few of them into French please:\n\n1. The intricate dance of planets around stars in distant solar systems, known as exoplanets, expands our horizons in the search for extraterrestrial life.\n\n2. The development of smart cities, using technology to improve the efficiency of services and meet residents' needs, represents a new frontier in urban planning.\n\n3. The melting of polar ice caps, accelerated by global warming, contributes to rising sea levels and the loss of habitat for species like the polar bear.\n\n4. The mysteries of dark matter and dark energy, making up most of the universe's mass and energy, challenge our understanding of the cosmos.\n\n5. The Great Barrier Reef, the world's largest coral reef system, is home to a vast array of marine species and is visible from space.\n\nPlease translate the following numbered sentences into French: [1, 2, 3, 4, 5]. Don't repeat the sentences in English. Only translate those 5 sentences. Number them in your response. Thank you!", + "The capital of France is", + "The president of the United States is", + "The future of AI is", ] prompts = [ llm_spec.tokenizer.apply_chat_template([{ @@ -5862,7 +5863,7 @@ def test_nvfp4_4gpu_mtp_ar(self): num_tokens = len(new_tokens) accept_rate = num_accepted / num_drafted - assert accept_rate > 0.35, \ + assert accept_rate > 0.2, \ f"Acceptance rate too low for prompt {i}: {accept_rate:.2f}" From cbca008edee49af4a749d76ad1af3d390fc07bea Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 27 Feb 2026 03:43:42 +0000 Subject: [PATCH 6/8] fix nemotron model mtp when enabling enable_attention_dp Signed-off-by: qgai --- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index d1a4974e744d..e6960339e177 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -250,7 +250,8 @@ def forward( assert hidden_states.shape[-1] == self.hidden_dim orig_shape = hidden_states.shape hidden_states = hidden_states.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: @@ -633,6 +634,7 @@ def forward( hidden_states: torch.Tensor, residual: torch.Tensor | None = None, attn_metadata: Optional[AttentionMetadata] = None, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: if self.has_start_projections: @@ -662,6 +664,7 @@ def forward( hidden_states = self.mixer( hidden_states=hidden_states, attn_metadata=attn_metadata, + **kwargs, ) if self.has_end_norm: @@ -768,6 +771,7 @@ def forward( hidden_states=hidden_states, residual=residual, attn_metadata=attn_metadata, + all_rank_num_tokens=all_rank_num_tokens, ) return hidden_states From 1120510bdff9b87485831859022407a6a0a8335a Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 27 Feb 2026 06:09:12 +0000 Subject: [PATCH 7/8] add test in tests/integration/test_lists/qa/llm_function_core.txt Signed-off-by: qgai --- tests/integration/test_lists/qa/llm_function_core.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 088ba536c77a..1652d3966bb7 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -267,6 +267,8 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_auto_dtype_4gpus[4-4 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off] 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 # multimodal accuracy tests accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_VL_7B::test_auto_dtype From 0c8af7e04fd408ed4ce4ae816d00636e9d4a485d Mon Sep 17 00:00:00 2001 From: qgai Date: Fri, 27 Feb 2026 07:07:09 +0000 Subject: [PATCH 8/8] small fix Signed-off-by: qgai --- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 3e9f36a52cee..dd61d81d7832 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -269,7 +269,7 @@ def forward( orig_shape = hidden_states_hp.shape hidden_states_hp_2d = hidden_states_hp.view(-1, self.hidden_dim) all_rank_num_tokens = kwargs.get('all_rank_num_tokens', - attn_metadata.all_rank_num_tokens) + attn_metadata.all_rank_num_tokens) def _compute_shared_output(): if self.shared_experts is not None: