From 45e5191b1464e5116ab411954aa377d2d528a9cf Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Mon, 26 Jan 2026 06:19:58 +0000 Subject: [PATCH 1/2] fix attention_need_spec_dec_mode flag Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 8 -------- tensorrt_llm/_torch/pyexecutor/model_engine.py | 8 ++------ tensorrt_llm/_torch/speculative/interface.py | 14 ++++++-------- 3 files changed, 8 insertions(+), 22 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 32a9332a0191..845ad8084c40 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -1271,14 +1271,6 @@ int AttentionOp::mlaGeneration( mXqaDispatcher->run(xqaParams, kv_cache_buffer, kv_scale_cache_buffer); return 0; } - else if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - TLLM_CHECK_WITH_INFO(false, "No available XQA kernels are found for speculative decoding mode."); - } - else if (mFuseFp4Quant) - { - TLLM_CHECK_WITH_INFO(false, "No available kernels are found for FP4 output."); - } } // Use FMHA otherwise. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6604dfb260f5..8663ffb53b5c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3246,13 +3246,9 @@ def forward(self, no_cache=kv_cache_manager is None) # attn_metadata now depends on spec_metadata since it determines the shape/content of spec_dec parameter Tensors - enable_mla = is_mla(self.model.model_config.pretrained_config) is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( - spec_resource_manager, - self.is_draft_model, - self.attn_backend, - self.model_is_wrapped, - is_mla=enable_mla) + spec_resource_manager, self.is_draft_model, self.attn_backend, + self.model_is_wrapped) attn_metadata.update_spec_dec_param( batch_size=scheduled_requests.batch_size, is_spec_decoding_enabled=is_spec_dec_mode, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 5037057699ed..fe59ca865f72 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -146,12 +146,11 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]): TrtllmAttention) or not xqa_supported def attention_need_spec_dec_mode( - self, - spec_resource_manager: Optional[BaseResourceManager], - is_draft_model: bool, - attention_backend: Type[AttentionBackend], - use_chain_drafter: bool, # CDL - is_mla: bool, + self, + spec_resource_manager: Optional[BaseResourceManager], + is_draft_model: bool, + attention_backend: Type[AttentionBackend], + use_chain_drafter: bool, # CDL ): """ If true, the attention backend kernel needs to run in spec-dec mode (multi-token query mode). @@ -164,8 +163,7 @@ def attention_need_spec_dec_mode( is_trtllm_attention = issubclass(attention_backend, TrtllmAttention) # Always use the multi-token query mode for 1-model if the kernels are available. - xqa_supported = not is_mla or get_sm_version() < 120 - use_case_1 = self.use_one_engine() and xqa_supported + use_case_1 = self.use_one_engine() # For 2-model, we need to enable it when we process multiple tokens at once. This occurs with # the target model (verification) or on the first draft for CDL based speculation. use_case_2 = not self.use_one_engine() and ( From 23334bfc4a3eaa346dbeda62a21954d07541b68a Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Mon, 26 Jan 2026 07:48:45 +0000 Subject: [PATCH 2/2] unwaive related test Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a6c1bb652256..aa3cb1e0fb4f 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -340,7 +340,6 @@ accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo-cudagraph:without_padding-pp2tp1cp2] SKIP (https://nvbugs/5787836) accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo-cudagraph:with_padding-pp2tp1cp2] SKIP (https://nvbugs/5787836) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_ngram SKIP (https://nvbugs/5769815) -accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[eagle3_one_model=True-torch_compile=False] SKIP (https://nvbugs/5787892) accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[eagle3_one_model=False-torch_compile=False] SKIP (https://nvbugs/5787892) accuracy/test_llm_api_pytorch.py::TestLlama4MaverickInstruct::test_fp8_chunked_prefill[tp8ep8-cuda_graph=True] SKIP (https://nvbugs/5791839) accuracy/test_llm_api_pytorch.py::TestLlama4MaverickInstruct::test_fp8_chunked_prefill[tp8ep8-cuda_graph=False] SKIP (https://nvbugs/5795918)