From 81dcb0067343a801e8bacb6358f4a83dc99cec86 Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Mon, 15 Jun 2026 16:57:08 -0700 Subject: [PATCH] [Spec v2] Use decode kernel for TRT-LLM MHA draft extend (#28241) --- python/sglang/srt/layers/attention/trtllm_mha_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 62078c3bc..5792a6cc9 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -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,