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

Support speculative decoding in the trtllm_mha attention backend (#9331)

Browse Source 
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Qiaolin Yu
2025-08-21 23:53:35 -07:00
committed by GitHub
parent fedfe91c1a
commit 9ec314c6ac
3 changed files with 413 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/server_args.py
View File

@@ -500,11 +500,6 @@ class ServerArgs:
)
self.page_size = 64
if self.speculative_algorithm is not None:
raise ValueError(
"trtllm_mha backend does not support speculative decoding yet."
)
if self.attention_backend == "dual_chunk_flash_attn":
logger.warning(
"Mixed chunk, radix cache, and cuda graphs are disabled because of using dual chunk flash attention backend"
@@ -653,6 +648,16 @@ class ServerArgs:
self.speculative_num_draft_tokens,
) = auto_choose_speculative_params(self)
if (
self.attention_backend == "trtllm_mha"
or self.decode_attention_backend == "trtllm_mha"
or self.prefill_attention_backend == "trtllm_mha"
):
if self.speculative_eagle_topk > 1:
raise ValueError(
"trtllm_mha backend only supports topk = 1 for speculative decoding."
)
if (
self.speculative_eagle_topk == 1
and self.speculative_num_draft_tokens != self.speculative_num_steps + 1