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, forward_batch: ForwardBatch,
save_kv_cache: bool = True, save_kv_cache: bool = True,
): ):
if ( if forward_batch.forward_mode in (
forward_batch.forward_mode == ForwardMode.EXTEND ForwardMode.EXTEND,
or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND ForwardMode.DRAFT_EXTEND,
ForwardMode.DRAFT_EXTEND_V2,
): ):
return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache) return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache)
else: 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 attention_backend = get_global_server_args().decode_attention_backend
elif ( elif (
forward_batch.forward_mode.is_target_verify() 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) # Use the specified backend for speculative operations (both verify and draft extend)
if get_global_server_args().speculative_attention_mode == "decode": if get_global_server_args().speculative_attention_mode == "decode":
test/registered/mla/test_flashmla.py
-2
View File
@@ -9,7 +9,6 @@ from types import SimpleNamespace
import requests import requests
import torch import torch
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval from sglang.test.run_eval import run_eval
@@ -54,7 +53,6 @@ class TestFlashMLAMTP(CustomTestCase):
] ]
) )
# Use longer timeout for DeepGEMM JIT compilation which can take 10-20 minutes # Use longer timeout for DeepGEMM JIT compilation which can take 10-20 minutes
with envs.SGLANG_ENABLE_SPEC_V2.override(False):
cls.process = popen_launch_server( cls.process = popen_launch_server(
cls.model, cls.model,
cls.base_url, cls.base_url,