Fix TRTLLM MHA routing for draft extend (#24856)

2026-05-12 15:48:32 -07:00
parent d5f3254ed1
commit a4109e87ac
1 changed files with 1 additions and 4 deletions
@@ -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,