From a4109e87ac0e14cd62cdf25a0da1c4001ee2d400 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Tue, 12 May 2026 15:48:32 -0700 Subject: [PATCH] Fix TRTLLM MHA routing for draft extend (#24856) --- python/sglang/srt/layers/attention/trtllm_mha_backend.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 253bb92d9..1ce665583 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -848,10 +848,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): page_table = self._get_layer_page_table(layer, forward_batch) - if ( - forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend_v2() - ): + if forward_batch.forward_mode.is_target_verify(): o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q, kv_cache=kv_cache,