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

Support spec v2 for Frozen-KV MTP; remove v1 worker (#27607)

Browse Source 
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Liangsheng Yin
2026-06-09 15:30:20 -07:00
committed by GitHub
co-authored by Khoa Pham Claude Opus 4.8
parent 7f730edfdc
commit decb88e0e3
13 changed files with 750 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/spec/test_frozen_kv_mtp.py
-8
View File
@@ -1,4 +1,3 @@
import os
import unittest
from types import SimpleNamespace
from typing import Optional
@@ -43,12 +42,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestFrozenKVMTP(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST
@classmethod
def _server_env(cls) -> dict[str, str]:
env = dict(os.environ)
env["SGLANG_ENABLE_SPEC_V2"] = "0"
return env
@classmethod
def _common_server_args(cls) -> list[str]:
args = [
@@ -110,7 +103,6 @@ class TestFrozenKVMTP(CustomTestCase):
"google/gemma-4-E4B-it",
self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
env=self._server_env(),
other_args=self._server_args(topk),
)
requests.get(self.base_url + "/flush_cache", timeout=30)
test/registered/spec/test_gemma4_mtp_26b_a4b_extra.py
-8
View File
@@ -1,4 +1,3 @@
import os
import unittest
from types import SimpleNamespace
from typing import Optional
@@ -68,12 +67,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestGemma4MTP26BA4B(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST
@classmethod
def _server_env(cls) -> dict[str, str]:
env = dict(os.environ)
env["SGLANG_ENABLE_SPEC_V2"] = "0"
return env
@classmethod
def _common_server_args(cls) -> list[str]:
args = [
@@ -140,7 +133,6 @@ class TestGemma4MTP26BA4B(CustomTestCase):
TARGET_PATH,
self.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
env=self._server_env(),
other_args=self._server_args(topk),
)
requests.get(self.base_url + "/flush_cache", timeout=30)
test/registered/spec/test_gemma4_mtp_31b_extra.py
-8
View File
@@ -1,4 +1,3 @@
import os
import unittest
from types import SimpleNamespace
from typing import Optional
@@ -58,12 +57,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestGemma4MTP31B(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST
@classmethod
def _server_env(cls) -> dict[str, str]:
env = dict(os.environ)
env["SGLANG_ENABLE_SPEC_V2"] = "0"
return env
@classmethod
def _common_server_args(cls) -> list[str]:
args = [
@@ -127,7 +120,6 @@ class TestGemma4MTP31B(CustomTestCase):
TARGET_PATH,
self.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
env=self._server_env(),
other_args=self._server_args(topk),
)
requests.get(self.base_url + "/flush_cache", timeout=30)
test/registered/unit/spec/test_frozen_kv_mtp_all_reqs_finish_in_verify.py
-136
View File
@@ -1,136 +0,0 @@
import unittest
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock, patch
import torch
from sglang.srt.speculative.frozen_kv_mtp_info import (
FrozenKVMTPDraftInput,
FrozenKVMTPVerifyInput,
)
from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
HIDDEN_SIZE = 8
TOPK = 1
def _stale_verify_input() -> FrozenKVMTPVerifyInput:
"""Placeholder for the verify input installed before target verification."""
return FrozenKVMTPVerifyInput.__new__(FrozenKVMTPVerifyInput)
def _make_prefill_draft_input() -> FrozenKVMTPDraftInput:
"""A non-idle draft input shaped like a single-req prefill arrival."""
return FrozenKVMTPDraftInput(
topk_p=torch.ones(1, TOPK, dtype=torch.float32),
topk_index=torch.zeros(1, TOPK, dtype=torch.int64),
hidden_states=torch.zeros(1, HIDDEN_SIZE, dtype=torch.float32),
bonus_tokens=torch.zeros(1, dtype=torch.int32),
)
class _FakeVerifyOutput(SimpleNamespace):
"""Fake verify result for worker versions that return either an object or a tuple."""
def __iter__(self):
yield self.logits_output
yield self
yield self.can_run_cuda_graph
class TestFrozenKVMTPWorker(CustomTestCase):
def _make_worker(self):
worker = FrozenKVMTPWorker.__new__(FrozenKVMTPWorker)
worker.device = torch.device("cpu")
worker.topk = TOPK
worker.model_config = SimpleNamespace(dtype=torch.float32)
worker.server_args = SimpleNamespace(enable_dp_attention=False)
worker._model_runner = SimpleNamespace(
tp_group=None, model=SimpleNamespace(backbone_hidden_size=HIDDEN_SIZE)
)
worker.draft_tp_context = lambda _: nullcontext()
stale_verify = _stale_verify_input()
worker.draft = Mock(return_value=stale_verify)
worker.verify = Mock(
return_value=_FakeVerifyOutput(
# Empty input_ids is the verify postcondition for:
# has_finished=True and no unfinished requests remain.
draft_extend_input=SimpleNamespace(
input_ids=torch.empty((0,), dtype=torch.int64)
),
logits_output=SimpleNamespace(),
accept_tokens=torch.empty((0,), dtype=torch.int64),
num_correct_drafts_per_req_cpu=[0, 0],
can_run_cuda_graph=False,
)
)
worker.forward_draft_extend_after_decode = Mock()
return worker, stale_verify
def _make_decode_batch(self):
return SimpleNamespace(
forward_mode=SimpleNamespace(
is_extend=lambda: False,
is_idle=lambda: False,
),
is_extend_in_batch=False,
reqs=[SimpleNamespace(), SimpleNamespace()],
spec_info=None,
)
def _forward_generation(self, worker, batch):
with (
patch(
"sglang.srt.speculative.frozen_kv_mtp_worker."
"speculative_moe_backend_context",
lambda: nullcontext(),
),
patch(
"sglang.srt.speculative.frozen_kv_mtp_worker."
"speculative_moe_a2a_backend_context",
lambda: nullcontext(),
),
):
return worker.forward_batch_generation(batch)
def test_forward_generation_installs_idle_draft_when_verify_finishes_all_reqs(
self,
):
worker, stale_verify = self._make_worker()
batch = self._make_decode_batch()
result = self._forward_generation(worker, batch)
worker.forward_draft_extend_after_decode.assert_not_called()
self.assertIsNot(batch.spec_info, stale_verify)
self.assertIsInstance(batch.spec_info, FrozenKVMTPDraftInput)
self.assertEqual(batch.spec_info.topk_index.shape, (0, TOPK))
self.assertEqual(batch.spec_info.hidden_states.shape, (0, HIDDEN_SIZE))
self.assertEqual(result.num_correct_drafts, 0)
def test_idle_draft_input_accepts_next_iter_prefill_merge(self):
worker, _ = self._make_worker()
batch = self._make_decode_batch()
self._forward_generation(worker, batch)
# This mirrors the scheduler's next-iter failure mode:
# running_batch.spec_info.merge_batch(other.spec_info). Without the
# all-reqs-finished else branch, batch.spec_info is still the stale
# FrozenKVMTPVerifyInput and this raises AttributeError.
example_prefill_draft_input = _make_prefill_draft_input()
batch.spec_info.merge_batch(example_prefill_draft_input)
self.assertIsInstance(batch.spec_info, FrozenKVMTPDraftInput)
self.assertEqual(batch.spec_info.topk_index.shape, (1, TOPK))
self.assertEqual(batch.spec_info.hidden_states.shape, (1, HIDDEN_SIZE))
if __name__ == "__main__":
unittest.main(verbosity=3)