minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Hanming Lu
2026-06-15 16:57:08 -07:00
committed by GitHub
parent cad43d3212
commit 81dcb00673
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+4 -1
View File
@@ -925,7 +925,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
page_table = self._get_layer_page_table(layer, forward_batch) 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( o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
query=q, query=q,
kv_cache=kv_cache, kv_cache=kv_cache,