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

Support spec v2 for FlashMLA speculative decoding (#24640)

Browse Source 
Co-authored-by: Jackey Hua <zhendonghua@users.noreply.github.com>
Co-authored-by: Depend <yu-depend@users.noreply.github.com>
This commit is contained in:
nagisa-kunhah
2026-05-19 15:23:17 -07:00
committed by GitHub
co-authored by Jackey Hua Depend
parent 2f70902329
commit b9d470f4a2
3 changed files with 11 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashmla_backend.py
+4 -3
View File
@@ -477,9 +477,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
forward_batch: ForwardBatch,
save_kv_cache: bool = True,
):
if (
forward_batch.forward_mode == ForwardMode.EXTEND
or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND
if forward_batch.forward_mode in (
ForwardMode.EXTEND,
ForwardMode.DRAFT_EXTEND,
ForwardMode.DRAFT_EXTEND_V2,
):
return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache)
else:
python/sglang/srt/models/deepseek_v2.py
+1 -1
View File
@@ -1558,7 +1558,7 @@ class DeepseekV2AttentionMLA(
attention_backend = get_global_server_args().decode_attention_backend
elif (
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend()
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
):
# Use the specified backend for speculative operations (both verify and draft extend)
if get_global_server_args().speculative_attention_mode == "decode":