From bc37b752564c57e4edae1fa3bcb9b68f01444801 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Mon, 1 Sep 2025 03:27:03 -0700 Subject: [PATCH 1/3] [https://nvbugs/5483615][fix] Remove unnecessary assertion to let main model and MTP work at the same time Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/attention.py | 24 ++++++++---------------- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 8 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index abf703c11ceb..a077789dcfa5 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -338,19 +338,13 @@ def _attn_impl( output_sf: Optional[torch.Tensor] = None, attention_sinks: Optional[torch.Tensor] = None, ): - - padded_num_tokens = attn_metadata.padded_num_tokens num_tokens = attn_metadata.num_tokens - if padded_num_tokens is not None: - assert q.shape[0] == padded_num_tokens - q = q[:num_tokens, :] - if k is not None: - assert k.shape[0] == padded_num_tokens - k = k[:num_tokens, :] - if v is not None: - assert v.shape[0] == padded_num_tokens - v = v[:num_tokens, :] + q = q[:num_tokens, :] + if k is not None: + k = k[:num_tokens, :] + if v is not None: + v = v[:num_tokens, :] out_scale = None out_scale_sf = None @@ -954,12 +948,10 @@ def forward_impl(self, position_ids: Optional[torch.Tensor], num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens - padded_num_tokens = attn_metadata.padded_num_tokens - if padded_num_tokens is not None: - hidden_states = hidden_states[:num_tokens, ...] - if position_ids is not None: - position_ids = position_ids[:num_tokens, ...] + hidden_states = hidden_states[:num_tokens, ...] + if position_ids is not None: + position_ids = position_ids[:num_tokens, ...] if self.is_lite: compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index dad42b807385..1ea5c7de0050 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -336,7 +336,6 @@ test_e2e.py::test_ptp_quickstart_advanced_8gpus_chunked_prefill_sq_22k[Llama-4-M test_e2e.py::test_ptp_quickstart_advanced_8gpus_chunked_prefill_sq_22k[Llama-4-Scout-17B-16E-Instruct-FP8-llama4-models/Llama-4-Scout-17B-16E-Instruct-FP8-True] SKIP (https://nvbugs/5481094) test_e2e.py::test_ptp_quickstart_multimodal_2gpu[Phi-4-multimodal-instruct-multimodals/Phi-4-multimodal-instruct] SKIP (https://nvbugs/5480415) accuracy/test_llm_api_pytorch.py::TestLlama4MaverickInstruct::test_fp8_eagle3[tp8-torch_compile=False] SKIP (https://nvbugs/5483534) -accuracy/test_llm_api_pytorch.py::TestLlama4MaverickInstruct::test_fp8_eagle3[tp8-torch_compile=True] SKIP (https://nvbugs/5483615) accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend SKIP (https://nvbugs/5448748) accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_nixl_backend SKIP (https://nvbugs/5448748) accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-CUTLASS] SKIP (https://nvbugs/5483913) From 13c4f30a0957915b6048db7a908fc4cd1fa025a1 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Wed, 3 Sep 2025 16:32:39 +0800 Subject: [PATCH 2/3] Correct position_ids slicing Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index a077789dcfa5..a91331c4540a 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -951,7 +951,7 @@ def forward_impl(self, position_ids: Optional[torch.Tensor], hidden_states = hidden_states[:num_tokens, ...] if position_ids is not None: - position_ids = position_ids[:num_tokens, ...] + position_ids = position_ids[:, :num_tokens] if self.is_lite: compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( From 7572a21454de4ff693a27208e90c5e33e0614db8 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Thu, 4 Sep 2025 15:23:40 +0800 Subject: [PATCH 3/3] Need to handle 1D and 2D position_ids at the same time Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index a91331c4540a..4d8fc12a991f 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -951,7 +951,7 @@ def forward_impl(self, position_ids: Optional[torch.Tensor], hidden_states = hidden_states[:num_tokens, ...] if position_ids is not None: - position_ids = position_ids[:, :num_tokens] + position_ids = position_ids[..., :num_tokens] if self.is_lite: compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split(