[Spec v2] Use decode kernel for TRT-LLM MHA draft extend (#28241)

2026-06-15 16:57:08 -07:00
parent cad43d3212
commit 81dcb00673
1 changed files with 4 additions and 1 deletions
@@ -925,7 +925,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):

        page_table = self._get_layer_page_table(layer, forward_batch)

-        if forward_batch.forward_mode.is_target_verify():
+        if (
+            forward_batch.forward_mode.is_target_verify()
+            or forward_batch.forward_mode.is_draft_extend_v2()
+        ):
            o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
                query=q,
                kv_cache=kv_cache,