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
python/sglang/srt/arg_groups/speculative_hook.py
+7 -4
View File
@@ -256,10 +256,13 @@ def _handle_frozen_kv_mtp(server_args: "ServerArgs") -> None:
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests." "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
) )
server_args.disable_overlap_schedule = True # SGLANG_ENABLE_SPEC_V2=False selects the non-overlap (synchronous) spec v2
logger.warning( # path instead of the overlap-scheduled one; both run the V2 worker.
"Overlap scheduler is disabled when using Frozen-KV MTP speculative decoding (spec v2 is not supported yet)." if (
) not envs.SGLANG_ENABLE_SPEC_V2.get()
and not server_args.disable_overlap_schedule
):
server_args.disable_overlap_schedule = True
if server_args.enable_mixed_chunk: if server_args.enable_mixed_chunk:
server_args.enable_mixed_chunk = False server_args.enable_mixed_chunk = False
python/sglang/srt/managers/scheduler.py
+4 -4
View File
@@ -1163,7 +1163,7 @@ class Scheduler(
self.device_module = torch.get_device_module(self.device) self.device_module = torch.get_device_module(self.device)
# FutureMap is always-on: input_ids relay used in both modes. # FutureMap is always-on: input_ids relay used in both modes.
# Workers not on BaseSpecWorker (e.g. FrozenKVMTPWorker) lack the # Workers not on BaseSpecWorker (e.g. NGRAM / DFLASH) lack the
# override; fall back to target-only so the helper still produces a # override; fall back to target-only so the helper still produces a
# safe decision (no accidental opt-out for unaudited shapes). # safe decision (no accidental opt-out for unaudited shapes).
if self.draft_worker is not None: if self.draft_worker is not None:
@@ -3120,9 +3120,9 @@ class Scheduler(
) )
batch.input_ids = None batch.input_ids = None
else: else:
# Spec_v1 (NGRAM / DFLASH / FROZEN_KV_MTP, non-overlap): # Spec_v1 (NGRAM / DFLASH, non-overlap): worker shape
# worker shape doesn't match req_pool_indices; relay is # doesn't match req_pool_indices; relay is unused (worker
# unused (worker rebuilds input_ids inside verify). # rebuilds input_ids inside verify).
batch.input_ids = batch_result.next_token_ids.to(torch.int64) batch.input_ids = batch_result.next_token_ids.to(torch.int64)
self.update_cache_from_scheduler(batch, batch_result) self.update_cache_from_scheduler(batch, batch_result)
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+3 -3
View File
@@ -34,7 +34,7 @@ from sglang.srt.utils import (
) )
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenKVMTPDraftWorker
@dataclass @dataclass
@@ -47,7 +47,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers):
topk_p: torch.Tensor topk_p: torch.Tensor
topk_index: torch.Tensor topk_index: torch.Tensor
hidden_states: torch.Tensor hidden_states: torch.Tensor
# Consumed by the captured seed iter; see `FrozenKVMTPWorker.draft_forward`. # Consumed by the captured seed iter; see `FrozenKVMTPDraftWorker.draft_forward`.
bonus_tokens: torch.Tensor bonus_tokens: torch.Tensor
global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_gpu: Optional[torch.Tensor]
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
@@ -56,7 +56,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers):
class FrozenKVMTPCudaGraphRunner: class FrozenKVMTPCudaGraphRunner:
"""CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step.""" """CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step."""
def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPWorker): def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPDraftWorker):
self.frozen_kv_mtp_worker = frozen_kv_mtp_worker self.frozen_kv_mtp_worker = frozen_kv_mtp_worker
self.model_runner = model_runner = frozen_kv_mtp_worker.draft_model_runner self.model_runner = model_runner = frozen_kv_mtp_worker.draft_model_runner
self.graphs = {} self.graphs = {}
python/sglang/srt/speculative/frozen_kv_mtp_info.py
+1 -25
View File
@@ -13,7 +13,7 @@
# ============================================================================== # ==============================================================================
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass, fields from dataclasses import dataclass
from typing import Dict from typing import Dict
from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.mem_cache.memory_pool import KVCache
@@ -21,7 +21,6 @@ from sglang.srt.speculative.eagle_info import (
EagleDraftExtendInput, EagleDraftExtendInput,
EagleDraftInput, EagleDraftInput,
EagleVerifyInput, EagleVerifyInput,
EagleVerifyOutput,
) )
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
@@ -68,26 +67,3 @@ class FrozenKVMTPVerifyInput(EagleVerifyInput):
def __post_init__(self): def __post_init__(self):
SpecInput.__init__(self, SpecInputType.FROZEN_KV_MTP_VERIFY) SpecInput.__init__(self, SpecInputType.FROZEN_KV_MTP_VERIFY)
def verify(self, *args, **kwargs) -> EagleVerifyOutput:
output = super().verify(*args, **kwargs)
output.draft_extend_input = _to_frozen_kv_mtp_draft_extend_input(
output.draft_extend_input
)
return output
FrozenKVMTPVerifyOutput = EagleVerifyOutput
def _to_frozen_kv_mtp_draft_extend_input(
draft_extend_input: EagleDraftExtendInput,
) -> FrozenKVMTPDraftExtendInput:
if isinstance(draft_extend_input, FrozenKVMTPDraftExtendInput):
return draft_extend_input
return FrozenKVMTPDraftExtendInput(
**{
field.name: getattr(draft_extend_input, field.name)
for field in fields(EagleDraftExtendInput)
}
)
python/sglang/srt/speculative/frozen_kv_mtp_utils.py
+2 -27
View File
@@ -14,19 +14,13 @@
from __future__ import annotations from __future__ import annotations
from contextlib import contextmanager from contextlib import contextmanager
from typing import TYPE_CHECKING, Tuple from typing import TYPE_CHECKING
import torch import torch
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.frozen_kv_mtp_info import ( from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext
FrozenKVMTPContext,
FrozenKVMTPDraftExtendInput,
FrozenKVMTPDraftInput,
)
from sglang.srt.speculative.spec_utils import fast_topk
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -159,22 +153,3 @@ def select_last_extend_hidden(
lens = torch.tensor(batch.extend_lens, device=hidden_states.device) lens = torch.tensor(batch.extend_lens, device=hidden_states.device)
last_indices = torch.cumsum(lens, dim=0) - 1 last_indices = torch.cumsum(lens, dim=0) - 1
return hidden_states[last_indices.to(torch.long)] return hidden_states[last_indices.to(torch.long)]
def select_last_verified_seed(
draft_input: FrozenKVMTPDraftExtendInput,
) -> Tuple[torch.Tensor, torch.Tensor]:
counts = draft_input.num_accept_tokens.to(torch.long)
last_indices = torch.cumsum(counts, dim=0) - 1
return (
draft_input.input_ids[last_indices],
draft_input.hidden_states[last_indices],
)
def capture_for_decode(
logits_output: LogitsProcessorOutput, draft_input: FrozenKVMTPDraftInput, topk: int
) -> None:
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
draft_input.topk_p, draft_input.topk_index = fast_topk(probs, topk, dim=-1)
draft_input.hidden_states = logits_output.hidden_states
python/sglang/srt/speculative/frozen_kv_mtp_worker.py
-826
View File
@@ -1,826 +0,0 @@
# Copyright 2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Frozen-KV MTP draft worker.
The assistant reads target KV only. It reuses EAGLE's verify input/output
contract, but owns the seed and recurrent draft loop because there is no
assistant-side KV extension.
"""
from __future__ import annotations
import logging
from typing import List, Optional, Tuple
import torch
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
from sglang.srt.observability.req_time_stats import set_time_batch
from sglang.srt.observability.trace import get_global_tracing_enabled
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.eagle_utils import (
build_tree_kernel_efficient,
organize_draft_results,
)
from sglang.srt.speculative.frozen_kv_mtp_info import (
FrozenKVMTPContext,
FrozenKVMTPDraftExtendInput,
FrozenKVMTPDraftInput,
FrozenKVMTPVerifyInput,
FrozenKVMTPVerifyOutput,
)
from sglang.srt.speculative.frozen_kv_mtp_utils import (
capture_for_decode,
expand_for_topk_draft,
frozen_kv_target_view,
position_for_batch,
select_last_extend_hidden,
select_last_verified_seed,
set_frozen_kv_positions,
target_kv_pool_view,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import (
draft_tp_context,
fast_topk,
generate_token_bitmask,
select_top_k_tokens,
spec_stage_span,
)
from sglang.srt.utils import empty_context
from sglang.srt.utils.async_probe import (
maybe_detect_inf,
maybe_detect_nan,
maybe_detect_oob,
)
logger = logging.getLogger(__name__)
class FrozenKVMTPWorker(TpModelWorker):
"""Frozen-KV MTP worker; same constructor shape as other TpModelWorker-based
spec workers. Entry: :meth:`forward_batch_generation` (stubs for now).
"""
def __init__(
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
self.server_args = server_args
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.gpu_id = gpu_id
self.device = server_args.device
self.target_worker = target_worker
self.page_size = server_args.page_size
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
)
assert self.speculative_algorithm.is_frozen_kv_mtp(), (
"FrozenKVMTPWorker should only be instantiated for "
"SpeculativeAlgorithm.FROZEN_KV_MTP, got "
f"{self.speculative_algorithm.name}. The dispatch happens in "
"arg_groups.speculative_hook.handle_speculative_decoding -> "
"_resolve_speculative_algorithm_alias."
)
# Assistant reads target KV directly, so its context length must match the target.
server_args.context_length = target_worker.model_runner.model_config.context_len
# Defer cuda graph capture; we do it ourselves below.
backup_disable_cuda_graph = server_args.disable_cuda_graph
server_args.disable_cuda_graph = True
# Draft attention uses target req_to_token + KV allocator (read-only).
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
target_cfg = target_worker.model_runner.memory_pool_config
draft_pool_config = MemoryPoolConfig(
max_total_num_tokens=64, # Dummy value
max_running_requests=target_cfg.max_running_requests,
)
self.hot_token_id = None
with (
empty_context()
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
super().__init__(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0,
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=draft_pool_config,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
if hasattr(self.draft_model_runner.model, "set_embed_and_head"):
self.draft_model_runner.model.set_embed_and_head(embed, head)
else:
logger.debug(
"Draft model %s does not implement set_embed_and_head; "
"skipping target-embedding bind in Frozen-KV MTP skeleton.",
type(self.draft_model_runner.model).__name__,
)
self.kv_context: Optional["FrozenKVMTPContext"] = None
if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
self._bind_kv_context()
self.draft_model_runner.server_args.disable_cuda_graph = (
backup_disable_cuda_graph
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
self.draft_attn_backend = self._init_draft_attn_backend()
self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
self.cuda_graph_runner = None
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
):
self.init_cuda_graphs()
@property
def draft_model_runner(self):
return self.model_runner
def get_attn_backend(self): # pragma: no cover - exposed for adaptive
return self.draft_attn_backend
def clear_cache_pool(self):
pass
def _resolve_draft_backend_type(self) -> str:
return (
self.server_args.speculative_draft_attention_backend
or self.server_args.decode_attention_backend
or self.server_args.attention_backend
)
def _init_draft_attn_backend(self):
if self.topk == 1:
return self.draft_model_runner.attn_backend
backend_type = self._resolve_draft_backend_type()
if backend_type != "triton":
raise ValueError(
"Frozen-KV MTP topk > 1 currently supports only the triton "
f"attention backend, got {backend_type}."
)
return self._init_triton_draft_attn_backend()
def _init_triton_draft_attn_backend(self):
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
max_bs = self.req_to_token_pool.size * self.topk
kv_indptr_buf = torch.zeros(
(max_bs + 1,), dtype=torch.int32, device=self.draft_model_runner.device
)
return TritonAttnBackend(
self.draft_model_runner,
skip_prefill=True,
kv_indptr_buf=kv_indptr_buf,
)
def _bind_kv_context(self) -> None:
draft_model = self.draft_model_runner.model
if not hasattr(draft_model, "build_frozen_kv_mtp_context") or not hasattr(
draft_model, "bind_frozen_kv_context"
):
logger.debug(
"Draft model %s does not implement Frozen-KV MTP context hooks; "
"skipping frozen-kv bind.",
type(draft_model).__name__,
)
return
ctx = draft_model.build_frozen_kv_mtp_context(
target_model=self.target_worker.model_runner.model,
target_token_to_kv_pool=self.target_worker.model_runner.token_to_kv_pool,
)
draft_model.bind_frozen_kv_context(ctx)
self.kv_context = ctx
def _frozen_kv_target_view(self, forward_batch: ForwardBatch):
return frozen_kv_target_view(
forward_batch, self.kv_context, self.draft_attn_backend
)
def _target_kv_pool_view(self, forward_batch: ForwardBatch):
return target_kv_pool_view(
forward_batch, self.kv_context, self.draft_attn_backend
)
def _set_positions(self, forward_batch: ForwardBatch) -> None:
set_frozen_kv_positions(forward_batch, self.topk)
def _expand_for_topk_draft(self, forward_batch: ForwardBatch) -> None:
expand_for_topk_draft(forward_batch, self.topk)
def _position_for_batch(self, batch: ScheduleBatch) -> torch.Tensor:
return position_for_batch(batch)
@property
def _recurrent_hidden_size(self) -> int:
return int(self.draft_model_runner.model.backbone_hidden_size)
def _init_frozen_kv_metadata(self, forward_batch: ForwardBatch) -> None:
if forward_batch.forward_mode.is_idle():
return
if forward_batch.seq_lens_cpu is not None:
forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()
else:
forward_batch.seq_lens_sum = torch.sum(forward_batch.seq_lens).item()
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata(forward_batch)
forward_batch.mark_forward_metadata_ready()
def _init_frozen_kv_metadata_capture_cuda_graph(
self, forward_batch: ForwardBatch
) -> None:
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata_out_graph(
forward_batch, in_capture=True
)
forward_batch.mark_forward_metadata_ready()
def _init_frozen_kv_metadata_replay_cuda_graph(
self, forward_batch: ForwardBatch, bs: int, seq_lens_sum: int
) -> None:
from types import SimpleNamespace
fb_view = SimpleNamespace(
batch_size=bs,
forward_mode=ForwardMode.DECODE,
input_ids=getattr(forward_batch, "input_ids", None),
req_pool_indices=forward_batch.req_pool_indices[:bs],
seq_lens=forward_batch.seq_lens[:bs],
seq_lens_sum=seq_lens_sum,
seq_lens_cpu=(
forward_batch.seq_lens_cpu[:bs]
if forward_batch.seq_lens_cpu is not None
else None
),
encoder_lens=None,
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
spec_info=None,
)
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata_out_graph(fb_view)
def init_cuda_graphs(self) -> None:
if self.server_args.disable_cuda_graph or self.speculative_num_steps <= 1:
return
if self.target_worker.device != "cuda":
logger.info(
"Frozen-KV MTP draft CUDA graph is only supported on CUDA; "
"running the draft loop eagerly on %s.",
self.target_worker.device,
)
return
from sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner import (
FrozenKVMTPCudaGraphRunner,
)
logger.info("Capture Frozen-KV MTP draft cuda graph begin.")
self.cuda_graph_runner = FrozenKVMTPCudaGraphRunner(self)
logger.info("Capture Frozen-KV MTP draft cuda graph end.")
def _select_last_extend_hidden(
self, batch: ScheduleBatch, hidden_states: torch.Tensor
) -> torch.Tensor:
return select_last_extend_hidden(batch, hidden_states)
def _select_last_verified_seed(
self, draft_input: FrozenKVMTPDraftExtendInput
) -> Tuple[torch.Tensor, torch.Tensor]:
return select_last_verified_seed(draft_input)
def _capture_for_decode(
self, logits_output: LogitsProcessorOutput, draft_input: FrozenKVMTPDraftInput
) -> None:
capture_for_decode(logits_output, draft_input, self.topk)
def _draft_preprocess_idle(self, batch: ScheduleBatch) -> None:
batch.spec_info = FrozenKVMTPDraftInput.create_idle_input(
device=self.device,
hidden_size=self._recurrent_hidden_size,
dtype=self.model_config.dtype,
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
def _run_assistant_seed_step(
self,
batch: ScheduleBatch,
last_token_ids: torch.Tensor,
last_hidden_states: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor] = None,
mm_input_embeds: Optional[torch.Tensor] = None,
draft_input: Optional[FrozenKVMTPDraftInput] = None,
) -> None:
"""Stash seed inputs on ``batch.spec_info``; the forward runs inside
the captured draft graph (see ``draft_forward``'s seed iter)."""
del seq_lens_cpu, mm_input_embeds, draft_input
if batch.forward_mode.is_idle() or last_token_ids.numel() == 0:
batch.spec_info = FrozenKVMTPDraftInput.create_idle_input(
device=batch.device,
hidden_size=self._recurrent_hidden_size,
dtype=self.model_config.dtype,
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
return
stashed = FrozenKVMTPDraftInput()
stashed.bonus_tokens = last_token_ids.to(torch.int64)
stashed.hidden_states = last_hidden_states
# Real-shaped zeros so inherited `filter_batch`/`merge_batch` can slice
# them between iters; overwritten by the captured seed iter.
bs = last_token_ids.shape[0]
device = last_token_ids.device
stashed.topk_p = torch.zeros(
(bs, self.topk), device=device, dtype=torch.float32
)
stashed.topk_index = torch.zeros(
(bs, self.topk), device=device, dtype=torch.int64
)
stashed.capture_hidden_mode = CaptureHiddenMode.LAST
stashed.num_tokens_per_req = 1
stashed.num_tokens_for_logprob_per_req = 1
batch.spec_info = stashed
def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
(
logits_output,
next_token_ids,
seq_lens_cpu,
can_run_cuda_graph,
) = self.forward_target_extend(batch)
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft_extend"),
):
self.forward_draft_extend(
batch,
logits_output.hidden_states,
next_token_ids,
seq_lens_cpu,
logits_output.mm_input_embeds,
)
return GenerationBatchResult(
logits_output=logits_output,
next_token_ids=next_token_ids,
num_correct_drafts=0,
can_run_cuda_graph=can_run_cuda_graph,
)
set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft"),
):
verify_input = self.draft(batch)
set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
# Install verify_input as `batch.spec_info` for the verify forward.
batch.spec_info = verify_input
verify_output = self.verify(batch)
if get_global_tracing_enabled():
for idx, req in enumerate(batch.reqs):
num_correct_drafts = verify_output.num_correct_drafts_per_req_cpu[idx]
req.time_stats.set_spec_verify_end_time(
num_correct_drafts=num_correct_drafts
)
set_time_batch(batch.reqs, "set_spec_draft_extend_start_time", trace_only=True)
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
):
draft_extend_input = verify_output.draft_extend_input
if (
self.server_args.enable_dp_attention
or draft_extend_input.input_ids.shape[0] > 0
):
# Install draft_extend_input as `batch.spec_info` for the seed
# step; `_run_assistant_seed_step` replaces it with a fresh
# `FrozenKVMTPDraftInput` for next iter.
batch.spec_info = draft_extend_input
with spec_stage_span("draft_extend"):
self.forward_draft_extend_after_decode(batch)
else:
# All reqs finished and dp_attention isn't forcing extend.
# Install an idle FrozenKVMTPDraftInput so next iter's scheduler
# ops (merge_batch / filter_batch) see well-typed empty
# tensors instead of None.
self._draft_preprocess_idle(batch)
set_time_batch(batch.reqs, "set_spec_draft_extend_end_time", trace_only=True)
return GenerationBatchResult(
logits_output=verify_output.logits_output,
next_token_ids=verify_output.accept_tokens,
num_correct_drafts=sum(verify_output.num_correct_drafts_per_req_cpu),
num_correct_drafts_per_req_cpu=verify_output.num_correct_drafts_per_req_cpu,
can_run_cuda_graph=verify_output.can_run_cuda_graph,
)
def forward_target_extend(
self, batch: ScheduleBatch
) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
batch.capture_hidden_mode = CaptureHiddenMode.FULL
batch_result = self.target_worker.forward_batch_generation(batch)
return (
batch_result.logits_output,
batch_result.next_token_ids,
batch.seq_lens_cpu,
batch_result.can_run_cuda_graph,
)
def forward_draft_extend(
self,
batch: ScheduleBatch,
hidden_states: torch.Tensor,
next_token_ids: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
mm_input_embeds: Optional[torch.Tensor] = None,
) -> None:
last_hidden = self._select_last_extend_hidden(batch, hidden_states)
self._run_assistant_seed_step(
batch,
next_token_ids,
last_hidden,
seq_lens_cpu=seq_lens_cpu,
mm_input_embeds=mm_input_embeds,
)
def forward_draft_extend_after_decode(self, batch: ScheduleBatch) -> None:
draft_extend_input: FrozenKVMTPDraftExtendInput = batch.spec_info
input_is_idle = batch.forward_mode.is_idle()
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
# All reqs finished. Install an idle FrozenKVMTPDraftInput so the
# next-iter draft sees a valid spec_info.
batch = batch.copy()
batch.prepare_for_idle()
batch.spec_info = FrozenKVMTPDraftInput.create_idle_input(
device=self.device,
hidden_size=self._recurrent_hidden_size,
dtype=self.model_config.dtype,
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
return
if batch.forward_mode.is_idle():
return
seq_lens_backup = batch.seq_lens.clone()
seq_lens_cpu_backup = batch.seq_lens_cpu.clone()
req_pool_indices_backup = batch.req_pool_indices
try:
# Verify may leave finished requests in ScheduleBatch; seed only
# the unfinished reqs carried by `draft_extend_input`.
batch.seq_lens = draft_extend_input.seq_lens
batch.seq_lens_cpu = draft_extend_input.seq_lens_cpu
batch.req_pool_indices = draft_extend_input.req_pool_indices
last_token_ids, last_hidden = self._select_last_verified_seed(
draft_extend_input
)
# `_run_assistant_seed_step` constructs a fresh `FrozenKVMTPDraftInput`
# and installs it on `batch.spec_info` for next iter.
self._run_assistant_seed_step(
batch,
last_token_ids,
last_hidden,
seq_lens_cpu=draft_extend_input.seq_lens_cpu,
)
finally:
batch.seq_lens = seq_lens_backup
batch.seq_lens_cpu = seq_lens_cpu_backup
batch.req_pool_indices = req_pool_indices_backup
def draft(self, batch: ScheduleBatch):
if batch.forward_mode.is_idle():
return FrozenKVMTPVerifyInput.create_idle_input(
self.topk,
self.speculative_num_steps,
self.speculative_num_draft_tokens,
)
batch.maybe_evict_swa()
for req in batch.reqs:
req.decode_batch_idx += 1
spec_info = batch.spec_info
assert isinstance(spec_info, FrozenKVMTPDraftInput)
if batch.sampling_info.penalizer_orchestrator.is_required:
batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
spec_info.bonus_tokens.to(torch.int64)
)
spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
spec_info.num_tokens_per_req = self.topk
spec_info.num_tokens_for_logprob_per_req = self.topk
spec_info.positions = self._position_for_batch(batch)
batch.seq_lens_sum = torch.sum(batch.seq_lens).item()
batch.return_hidden_states = False
forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
assert forward_batch.capture_hidden_mode == CaptureHiddenMode.LAST
self._set_positions(forward_batch)
self._expand_for_topk_draft(forward_batch)
can_run_cuda_graph = self.cuda_graph_runner and self.cuda_graph_runner.can_run(
forward_batch
)
if can_run_cuda_graph:
parent_list, top_scores_index, draft_tokens = self.cuda_graph_runner.replay(
forward_batch
)
else:
forward_batch.can_run_dp_cuda_graph = False
parent_list, top_scores_index, draft_tokens = self.draft_forward(
forward_batch
)
(
tree_mask,
position,
retrieve_index,
retrieve_next_token,
retrieve_next_sibling,
draft_tokens,
) = build_tree_kernel_efficient(
spec_info.bonus_tokens,
parent_list,
top_scores_index,
draft_tokens,
batch.seq_lens,
batch.seq_lens_sum,
self.topk,
self.speculative_num_steps,
self.speculative_num_draft_tokens,
)
return FrozenKVMTPVerifyInput(
draft_token=draft_tokens,
custom_mask=tree_mask,
positions=position,
retrieve_index=retrieve_index,
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_cum_len=None,
spec_steps=self.speculative_num_steps,
topk=self.topk,
draft_token_num=self.speculative_num_draft_tokens,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=batch.seq_lens_sum,
seq_lens_cpu=batch.seq_lens_cpu,
)
def draft_forward(self, forward_batch: ForwardBatch):
spec_info = forward_batch.spec_info
assert isinstance(spec_info, FrozenKVMTPDraftInput)
score_list: List[torch.Tensor] = []
token_list: List[torch.Tensor] = []
parents_list: List[torch.Tensor] = []
# Seed + recurrent iters share the same `seq_lens - 1` rope position,
# so one init covers the loop. Must run even at num_steps == 1.
if forward_batch.needs_forward_metadata_init():
self._init_frozen_kv_metadata(forward_batch)
# Seed iter: assistant forward on (bonus_token, target_h) to produce
# iter-0 `(topk_p, topk_index, hidden_states)`. For topk>1, replicate
# to `bs*topk` to match kernel shapes, then slice back per-req.
bonus_tokens = spec_info.bonus_tokens
target_hidden = spec_info.hidden_states
if self.topk > 1:
seed_input_ids = bonus_tokens.repeat_interleave(self.topk, dim=0)
seed_prev_hidden = target_hidden.repeat_interleave(self.topk, dim=0)
else:
seed_input_ids = bonus_tokens
seed_prev_hidden = target_hidden
forward_batch.input_ids = seed_input_ids
forward_batch.spec_info.hidden_states = seed_prev_hidden
self._set_positions(forward_batch)
with (
self._target_kv_pool_view(forward_batch),
forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
):
seed_output = self.draft_model_runner.forward(forward_batch).logits_output
maybe_detect_nan(
seed_output.next_token_logits, "frozen_kv_mtp_draft: seed iter"
)
if self.topk > 1:
seed_next_logits = seed_output.next_token_logits[:: self.topk]
seed_hidden_per_req = seed_output.hidden_states[:: self.topk]
else:
seed_next_logits = seed_output.next_token_logits
seed_hidden_per_req = seed_output.hidden_states
probs = torch.softmax(seed_next_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
maybe_detect_oob(
topk_index,
0,
seed_next_logits.shape[-1],
"frozen_kv_mtp_draft: seed topk_index OOB",
)
hidden_states = seed_hidden_per_req
scores = None
for i in range(self.speculative_num_steps):
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
i, topk_p, topk_index, hidden_states, scores, self.topk
)
score_list.append(tree_info[0])
token_list.append(tree_info[1])
parents_list.append(tree_info[2])
if i == self.speculative_num_steps - 1:
break
forward_batch.input_ids = input_ids
forward_batch.spec_info.hidden_states = hidden_states
self._set_positions(forward_batch)
with (
self._target_kv_pool_view(forward_batch),
forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
):
logits_output = self.draft_model_runner.forward(
forward_batch
).logits_output
maybe_detect_nan(
logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
)
maybe_detect_inf(
logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
)
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
maybe_detect_oob(
topk_index,
0,
logits_output.next_token_logits.shape[-1],
"frozen_kv_mtp_draft: topk_index OOB",
)
hidden_states = logits_output.hidden_states
return organize_draft_results(
score_list, token_list, parents_list, self.speculative_num_draft_tokens
)
def verify(self, batch: ScheduleBatch):
spec_info: FrozenKVMTPVerifyInput = batch.spec_info
seq_lens_pre_verify = batch.seq_lens.clone()
spec_info.prepare_for_verify(batch, self.page_size)
spec_info.num_tokens_per_req = self.speculative_num_steps + 1
batch.return_hidden_states = False
batch.forward_mode = (
ForwardMode.TARGET_VERIFY
if not batch.forward_mode.is_idle()
else ForwardMode.IDLE
)
if batch.has_grammar:
retrieve_next_token_cpu = spec_info.retrieve_next_token.cpu()
retrieve_next_sibling_cpu = spec_info.retrieve_next_sibling.cpu()
draft_tokens_cpu = spec_info.draft_token.view(
spec_info.retrieve_next_token.shape
).cpu()
batch.seq_lens_cpu_cache = spec_info.seq_lens_cpu
batch_result = self.target_worker.forward_batch_generation(
batch, is_verify=True
)
logits_output, can_run_cuda_graph = (
batch_result.logits_output,
batch_result.can_run_cuda_graph,
)
vocab_mask = None
if batch.has_grammar:
vocab_mask = generate_token_bitmask(
batch.reqs,
spec_info,
retrieve_next_token_cpu,
retrieve_next_sibling_cpu,
draft_tokens_cpu,
batch.sampling_info.vocab_size,
)
if vocab_mask is not None:
assert spec_info.grammar is not None
vocab_mask = vocab_mask.to(spec_info.retrieve_next_token.device)
batch.sampling_info.vocab_mask = None
maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_verify")
maybe_detect_inf(logits_output.next_token_logits, "frozen_kv_mtp_verify")
spec_info.hidden_states = logits_output.hidden_states
res: FrozenKVMTPVerifyOutput = spec_info.verify(
batch,
logits_output,
self.token_to_kv_pool_allocator,
self.page_size,
vocab_mask,
)
logits_output.next_token_logits = logits_output.next_token_logits[
res.accept_indices
]
logits_output.hidden_states = logits_output.hidden_states[res.accept_indices]
if (
self.target_worker.model_runner.hybrid_gdn_config is not None
or self.target_worker.model_runner.mamba2_config is not None
or self.target_worker.model_runner.hybrid_lightning_config is not None
):
logger.warning(
"Frozen-KV MTP does not implement mamba state updates; "
"targets with recurrent state should not use this path."
)
if batch.return_logprob:
add_output_logprobs_for_spec_v1(batch, res, logits_output)
batch.forward_mode = (
ForwardMode.DECODE if not batch.forward_mode.is_idle() else ForwardMode.IDLE
)
del seq_lens_pre_verify
res.can_run_cuda_graph = can_run_cuda_graph
return res
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+719 -7
View File
@@ -11,18 +11,79 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
"""Overlap-scheduling placeholder for frozen-KV MTP (raises until implemented).""" """Spec-v2 worker for Frozen-KV MTP (two layers, like ``eagle_worker_v2``).
The frozen draft reads the target KV cache read-only and owns no KV pool, so
its "draft extend" is not a model forward: it selects the last accepted token +
target hidden state as the next-iter seed, and the seed forward runs at the
start of the next draft.
"""
from __future__ import annotations from __future__ import annotations
import logging
from typing import Optional from typing import Optional
import torch
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker from sglang.srt.speculative.base_spec_worker import BaseDraftWorker
from sglang.srt.speculative.eagle_utils import (
build_tree_kernel_efficient,
organize_draft_results,
)
from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2, _get_plan_stream
from sglang.srt.speculative.frozen_kv_mtp_info import (
FrozenKVMTPContext,
FrozenKVMTPDraftInput,
FrozenKVMTPVerifyInput,
)
from sglang.srt.speculative.frozen_kv_mtp_utils import (
expand_for_topk_draft,
frozen_kv_target_view,
position_for_batch,
select_last_extend_hidden,
set_frozen_kv_positions,
target_kv_pool_view,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import (
draft_tp_context,
fast_topk,
select_top_k_tokens,
spec_stage_span,
)
from sglang.srt.utils import empty_context
from sglang.srt.utils.async_probe import (
maybe_detect_inf,
maybe_detect_nan,
maybe_detect_oob,
)
logger = logging.getLogger(__name__)
class FrozenKVMTPWorkerV2(FrozenKVMTPWorker): class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker):
"""Frozen-KV MTP draft worker.
The assistant reads target KV only. It reuses EAGLE's verify input/output
contract, but owns the seed and recurrent draft loop because there is no
assistant-side KV extension.
"""
def __init__( def __init__(
self, self,
server_args: ServerArgs, server_args: ServerArgs,
@@ -35,8 +96,659 @@ class FrozenKVMTPWorkerV2(FrozenKVMTPWorker):
nccl_port: int, nccl_port: int,
target_worker: TpModelWorker, target_worker: TpModelWorker,
): ):
raise NotImplementedError( self.server_args = server_args
"FrozenKVMTPWorkerV2 (overlap scheduling for Frozen-KV MTP) is " self.topk = server_args.speculative_eagle_topk
"not yet implemented. Pass --disable-overlap-schedule to use " self.speculative_num_steps = server_args.speculative_num_steps
"FrozenKVMTPWorker." self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.gpu_id = gpu_id
self.device = server_args.device
self.target_worker = target_worker
self.page_size = server_args.page_size
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
) )
assert self.speculative_algorithm.is_frozen_kv_mtp(), (
"FrozenKVMTPDraftWorker should only be instantiated for "
"SpeculativeAlgorithm.FROZEN_KV_MTP, got "
f"{self.speculative_algorithm.name}."
)
# Defer cuda graph capture; we do it ourselves below.
backup_disable_cuda_graph = server_args.disable_cuda_graph
server_args.disable_cuda_graph = True
# Draft attention uses target req_to_token + KV allocator (read-only).
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
target_cfg = target_worker.model_runner.memory_pool_config
draft_pool_config = MemoryPoolConfig(
max_total_num_tokens=64, # Dummy value
max_running_requests=target_cfg.max_running_requests,
)
self.hot_token_id = None
with (
empty_context()
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
# NOTE: call TpModelWorker.__init__ explicitly -- BaseDraftWorker is
# an ABC with no __init__, so cooperative super() would be ambiguous.
TpModelWorker.__init__(
self,
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0,
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
memory_pool_config=draft_pool_config,
)
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
if hasattr(self.draft_model_runner.model, "set_embed_and_head"):
self.draft_model_runner.model.set_embed_and_head(embed, head)
else:
logger.debug(
"Draft model %s does not implement set_embed_and_head; "
"skipping target-embedding bind in Frozen-KV MTP skeleton.",
type(self.draft_model_runner.model).__name__,
)
self.kv_context: Optional["FrozenKVMTPContext"] = None
if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"):
self._bind_kv_context()
self.draft_model_runner.server_args.disable_cuda_graph = (
backup_disable_cuda_graph
)
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
self.draft_attn_backend = self._init_draft_attn_backend()
self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
self.cuda_graph_runner = None
# Frozen draft has no draft-extend forward (seed-select only); keep these
# None so inherited probes (spec_v2_attn_backends, adaptive) stay typed.
self.draft_extend_attn_backend = None
self.cuda_graph_runner_for_draft_extend = None
with (
self.draft_tp_context(self.draft_model_runner.tp_group),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
):
self.init_cuda_graphs()
@property
def draft_model_runner(self):
return self.model_runner
@property
def draft_runner(self):
# Alias for the inherited EAGLEWorkerV2 forward/verify skeleton, which
# reads `draft_worker.draft_runner`.
return self.model_runner
def get_attn_backend(self): # pragma: no cover - exposed for adaptive
return self.draft_attn_backend
def clear_cache_pool(self):
pass
def _resolve_draft_backend_type(self) -> str:
return (
self.server_args.speculative_draft_attention_backend
or self.server_args.decode_attention_backend
or self.server_args.attention_backend
)
def _init_draft_attn_backend(self):
if self.topk == 1:
return self.draft_model_runner.attn_backend
backend_type = self._resolve_draft_backend_type()
if backend_type != "triton":
raise ValueError(
"Frozen-KV MTP topk > 1 currently supports only the triton "
f"attention backend, got {backend_type}."
)
return self._init_triton_draft_attn_backend()
def _init_triton_draft_attn_backend(self):
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
max_bs = self.req_to_token_pool.size * self.topk
kv_indptr_buf = torch.zeros(
(max_bs + 1,), dtype=torch.int32, device=self.draft_model_runner.device
)
return TritonAttnBackend(
self.draft_model_runner,
skip_prefill=True,
kv_indptr_buf=kv_indptr_buf,
)
def _bind_kv_context(self) -> None:
draft_model = self.draft_model_runner.model
if not hasattr(draft_model, "build_frozen_kv_mtp_context") or not hasattr(
draft_model, "bind_frozen_kv_context"
):
logger.debug(
"Draft model %s does not implement Frozen-KV MTP context hooks; "
"skipping frozen-kv bind.",
type(draft_model).__name__,
)
return
ctx = draft_model.build_frozen_kv_mtp_context(
target_model=self.target_worker.model_runner.model,
target_token_to_kv_pool=self.target_worker.model_runner.token_to_kv_pool,
)
draft_model.bind_frozen_kv_context(ctx)
self.kv_context = ctx
def _frozen_kv_target_view(self, forward_batch: ForwardBatch):
return frozen_kv_target_view(
forward_batch, self.kv_context, self.draft_attn_backend
)
def _target_kv_pool_view(self, forward_batch: ForwardBatch):
return target_kv_pool_view(
forward_batch, self.kv_context, self.draft_attn_backend
)
def _set_positions(self, forward_batch: ForwardBatch) -> None:
set_frozen_kv_positions(forward_batch, self.topk)
def _expand_for_topk_draft(self, forward_batch: ForwardBatch) -> None:
expand_for_topk_draft(forward_batch, self.topk)
def _position_for_batch(self, batch: ScheduleBatch) -> torch.Tensor:
return position_for_batch(batch)
@property
def _recurrent_hidden_size(self) -> int:
return int(self.draft_model_runner.model.backbone_hidden_size)
def _init_frozen_kv_metadata(self, forward_batch: ForwardBatch) -> None:
if forward_batch.forward_mode.is_idle():
return
if forward_batch.seq_lens_cpu is not None:
forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item()
else:
forward_batch.seq_lens_sum = torch.sum(forward_batch.seq_lens).item()
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata(forward_batch)
forward_batch.mark_forward_metadata_ready()
def _init_frozen_kv_metadata_capture_cuda_graph(
self, forward_batch: ForwardBatch
) -> None:
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata_out_graph(
forward_batch, in_capture=True
)
forward_batch.mark_forward_metadata_ready()
def _init_frozen_kv_metadata_replay_cuda_graph(
self, forward_batch: ForwardBatch, bs: int, seq_lens_sum: int
) -> None:
from types import SimpleNamespace
fb_view = SimpleNamespace(
batch_size=bs,
forward_mode=ForwardMode.DECODE,
input_ids=getattr(forward_batch, "input_ids", None),
req_pool_indices=forward_batch.req_pool_indices[:bs],
seq_lens=forward_batch.seq_lens[:bs],
seq_lens_sum=seq_lens_sum,
seq_lens_cpu=(
forward_batch.seq_lens_cpu[:bs]
if forward_batch.seq_lens_cpu is not None
else None
),
encoder_lens=None,
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
spec_info=None,
)
with self._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata_out_graph(fb_view)
def init_cuda_graphs(self) -> None:
if self.server_args.disable_cuda_graph or self.speculative_num_steps <= 1:
return
if self.target_worker.device != "cuda":
logger.info(
"Frozen-KV MTP draft CUDA graph is only supported on CUDA; "
"running the draft loop eagerly on %s.",
self.target_worker.device,
)
return
from sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner import (
FrozenKVMTPCudaGraphRunner,
)
logger.info("Capture Frozen-KV MTP draft cuda graph begin.")
self.cuda_graph_runner = FrozenKVMTPCudaGraphRunner(self)
logger.info("Capture Frozen-KV MTP draft cuda graph end.")
def _select_last_extend_hidden(
self, batch: ScheduleBatch, hidden_states: torch.Tensor
) -> torch.Tensor:
return select_last_extend_hidden(batch, hidden_states)
def _idle_seed(self) -> FrozenKVMTPDraftInput:
return FrozenKVMTPDraftInput.create_idle_input(
device=self.device,
hidden_size=self._recurrent_hidden_size,
dtype=self.model_config.dtype,
topk=self.topk,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
def _build_seed_draft_input(
self,
last_token_ids: torch.Tensor,
last_hidden_states: torch.Tensor,
) -> FrozenKVMTPDraftInput:
"""Build the next-iter seed ``FrozenKVMTPDraftInput`` from (bonus token,
target hidden). No forward here -- the seed forward runs inside the
captured draft graph (see ``draft_forward``'s seed iter)."""
if last_token_ids.numel() == 0:
return self._idle_seed()
stashed = FrozenKVMTPDraftInput()
stashed.bonus_tokens = last_token_ids.to(torch.int64)
stashed.hidden_states = last_hidden_states
# Real-shaped zeros so inherited `filter_batch`/`merge_batch` can slice
# them between iters; overwritten by the captured seed iter.
bs = last_token_ids.shape[0]
device = last_token_ids.device
stashed.topk_p = torch.zeros(
(bs, self.topk), device=device, dtype=torch.float32
)
stashed.topk_index = torch.zeros(
(bs, self.topk), device=device, dtype=torch.int64
)
stashed.capture_hidden_mode = CaptureHiddenMode.LAST
stashed.num_tokens_per_req = 1
stashed.num_tokens_for_logprob_per_req = 1
return stashed
def draft(self, batch: ScheduleBatch):
if batch.forward_mode.is_idle():
return FrozenKVMTPVerifyInput.create_idle_input(
self.topk,
self.speculative_num_steps,
self.speculative_num_draft_tokens,
)
spec_info = batch.spec_info
assert isinstance(spec_info, FrozenKVMTPDraftInput)
# NOTE: per-iter bookkeeping (penalty cumulation, maybe_evict_swa,
# decode_batch_idx tick) is done by the inherited
# EagleDraftInputV2Mixin.prepare_for_decode (scheduler-driven, see
# ScheduleBatch.prepare_for_decode), not here -- matching EAGLE v2.
# Repeating evict/tick here would double-run them: the idx clock
# gates SWA eviction timing and the SWA prefix-lock release.
spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
spec_info.num_tokens_per_req = self.topk
spec_info.num_tokens_for_logprob_per_req = self.topk
spec_info.positions = self._position_for_batch(batch)
batch.seq_lens_sum = torch.sum(batch.seq_lens).item()
batch.return_hidden_states = False
forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner)
assert forward_batch.capture_hidden_mode == CaptureHiddenMode.LAST
self._set_positions(forward_batch)
self._expand_for_topk_draft(forward_batch)
can_run_cuda_graph = self.cuda_graph_runner and self.cuda_graph_runner.can_run(
forward_batch
)
if can_run_cuda_graph:
parent_list, top_scores_index, draft_tokens = self.cuda_graph_runner.replay(
forward_batch
)
else:
forward_batch.can_run_dp_cuda_graph = False
parent_list, top_scores_index, draft_tokens = self.draft_forward(
forward_batch
)
(
tree_mask,
position,
retrieve_index,
retrieve_next_token,
retrieve_next_sibling,
draft_tokens,
) = build_tree_kernel_efficient(
spec_info.bonus_tokens,
parent_list,
top_scores_index,
draft_tokens,
batch.seq_lens,
batch.seq_lens_sum,
self.topk,
self.speculative_num_steps,
self.speculative_num_draft_tokens,
)
return FrozenKVMTPVerifyInput(
draft_token=draft_tokens,
custom_mask=tree_mask,
positions=position,
retrieve_index=retrieve_index,
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_cum_len=None,
spec_steps=self.speculative_num_steps,
topk=self.topk,
draft_token_num=self.speculative_num_draft_tokens,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=batch.seq_lens_sum,
seq_lens_cpu=batch.seq_lens_cpu,
)
def draft_forward(self, forward_batch: ForwardBatch):
spec_info = forward_batch.spec_info
assert isinstance(spec_info, FrozenKVMTPDraftInput)
score_list: list[torch.Tensor] = []
token_list: list[torch.Tensor] = []
parents_list: list[torch.Tensor] = []
# Seed + recurrent iters share the same `seq_lens - 1` rope position,
# so one init covers the loop. Must run even at num_steps == 1.
if forward_batch.needs_forward_metadata_init():
self._init_frozen_kv_metadata(forward_batch)
# Seed iter: assistant forward on (bonus_token, target_h) to produce
# iter-0 `(topk_p, topk_index, hidden_states)`. For topk>1, replicate
# to `bs*topk` to match kernel shapes, then slice back per-req.
bonus_tokens = spec_info.bonus_tokens
target_hidden = spec_info.hidden_states
if self.topk > 1:
seed_input_ids = bonus_tokens.repeat_interleave(self.topk, dim=0)
seed_prev_hidden = target_hidden.repeat_interleave(self.topk, dim=0)
else:
seed_input_ids = bonus_tokens
seed_prev_hidden = target_hidden
forward_batch.input_ids = seed_input_ids
forward_batch.spec_info.hidden_states = seed_prev_hidden
self._set_positions(forward_batch)
with (
self._target_kv_pool_view(forward_batch),
forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
):
seed_output = self.draft_model_runner.forward(forward_batch).logits_output
maybe_detect_nan(
seed_output.next_token_logits, "frozen_kv_mtp_draft: seed iter"
)
if self.topk > 1:
seed_next_logits = seed_output.next_token_logits[:: self.topk]
seed_hidden_per_req = seed_output.hidden_states[:: self.topk]
else:
seed_next_logits = seed_output.next_token_logits
seed_hidden_per_req = seed_output.hidden_states
probs = torch.softmax(seed_next_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
maybe_detect_oob(
topk_index,
0,
seed_next_logits.shape[-1],
"frozen_kv_mtp_draft: seed topk_index OOB",
)
hidden_states = seed_hidden_per_req
scores = None
for i in range(self.speculative_num_steps):
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
i, topk_p, topk_index, hidden_states, scores, self.topk
)
score_list.append(tree_info[0])
token_list.append(tree_info[1])
parents_list.append(tree_info[2])
if i == self.speculative_num_steps - 1:
break
forward_batch.input_ids = input_ids
forward_batch.spec_info.hidden_states = hidden_states
self._set_positions(forward_batch)
with (
self._target_kv_pool_view(forward_batch),
forward_context(ForwardContext(attn_backend=self.draft_attn_backend)),
):
logits_output = self.draft_model_runner.forward(
forward_batch
).logits_output
maybe_detect_nan(
logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
)
maybe_detect_inf(
logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}"
)
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
maybe_detect_oob(
topk_index,
0,
logits_output.next_token_logits.shape[-1],
"frozen_kv_mtp_draft: topk_index OOB",
)
hidden_states = logits_output.hidden_states
return organize_draft_results(
score_list, token_list, parents_list, self.speculative_num_draft_tokens
)
def draft_extend(self):
# BaseDraftWorker contract. Frozen has no draft-KV extend forward; the
# orchestrator calls `_draft_extend_for_{prefill,decode}` directly.
pass
def _draft_extend_for_prefill(
self,
batch: ScheduleBatch,
target_hidden_states: torch.Tensor,
next_token_ids: torch.Tensor,
mm_input_embeds: Optional[torch.Tensor] = None,
) -> FrozenKVMTPDraftInput:
"""Seed for the first decode iter after prefill. Frozen draft writes no
KV (reads target KV), so unlike EAGLE there is no draft-extend forward:
just select the last prompt hidden + bonus token and stash the seed."""
del mm_input_embeds # frozen seed needs no input embeds
if batch.forward_mode.is_idle():
return self._idle_seed()
last_hidden = self._select_last_extend_hidden(batch, target_hidden_states)
return self._build_seed_draft_input(next_token_ids, last_hidden)
def _draft_extend_for_decode(self, batch: ScheduleBatch, batch_result) -> None:
"""Frozen 'draft extend': no forward. Pull the last accepted token's
target hidden from the verify output and stash it as the next-iter seed.
Replaces verify's `EagleDraftInput` with a `FrozenKVMTPDraftInput` so the
next draft passes the FROZEN_KV_MTP attn-backend assertions.
"""
if batch.forward_mode.is_idle():
batch_result.next_draft_input = self._idle_seed()
return
bs = len(batch.seq_lens)
# Same per-req select_index EAGLE uses on its draft-extend output: the
# last accepted node (accept_lens - 1) in each per-req block of width
# num_draft_tokens. Verify already compacted the accepted path to the
# front (topk > 1) / it is the front chain (topk == 1).
select_index = (
torch.arange(bs, device=self.device) * self.speculative_num_draft_tokens
+ batch_result.accept_lens
- 1
)
last_hidden = batch_result.logits_output.hidden_states[select_index]
bonus_tokens = batch_result.next_draft_input.bonus_tokens
batch_result.next_draft_input = self._build_seed_draft_input(
bonus_tokens, last_hidden
)
class FrozenKVMTPWorkerV2(EAGLEWorkerV2):
"""Spec-v2 (overlap) orchestrator for Frozen-KV MTP.
Reuses ``EAGLEWorkerV2``'s verify / ``move_accept_tokens`` / forward
skeleton verbatim; only the draft worker and the seed-based draft-extend
are frozen-specific.
"""
def __init__(
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
# NOTE: intentionally does NOT call EAGLEWorkerV2.__init__ -- that builds
# an EagleDraftWorker (with its own draft KV pool). The frozen draft owns
# no KV, so we mirror the relevant setup and build a FrozenKVMTPDraftWorker.
self.server_args = server_args
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.tp_rank = tp_rank
self.gpu_id = gpu_id
self.device = server_args.device
self._target_worker = target_worker
self.page_size = server_args.page_size
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Match the draft context length to the target (assistant reads target KV).
server_args.context_length = target_worker.model_runner.model_config.context_len
self._draft_worker = FrozenKVMTPDraftWorker(
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Frozen MTP does not wire the adaptive controller yet.
assert (
not server_args.speculative_adaptive
), "Frozen-KV MTP does not support adaptive speculative decoding yet."
self.adaptive_controller = None
# Some dummy tensors (parity with EAGLEWorkerV2 init).
self.num_new_pages_per_topk = torch.empty(
(), dtype=torch.int64, device=self.device
)
self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
@property
def spec_v2_attn_backends(self) -> tuple:
# Frozen draft touches no draft-extend backend; only target + draft.
return (
self._target_worker.model_runner.attn_backend,
self._draft_worker.draft_attn_backend,
)
def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None):
# Mirrors EAGLEWorkerV2.forward_batch_generation; the only frozen-specific
# change is the idle draft-input (FrozenKVMTPDraftInput + recurrent hidden
# size). The draft / seed-based draft-extend hooks are FrozenKVMTPDraftWorker's.
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
# Target prefill (frozen is never standalone -> capture FULL hidden).
batch.capture_hidden_mode = CaptureHiddenMode.FULL
batch_output = self.target_worker.forward_batch_generation(batch)
# Spec_v2 convention: batch.seq_lens = length BEFORE this iter's tokens.
batch_output.new_seq_lens = batch.seq_lens
# Publish before draft-extend so the fence is at target-end.
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
# Draft prefill seed (no forward).
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft_extend"),
):
batch_output.next_draft_input = (
self.draft_worker._draft_extend_for_prefill(
batch,
batch_output.logits_output.hidden_states,
batch_output.next_token_ids,
batch_output.logits_output.mm_input_embeds,
)
)
return batch_output
else:
self.activate_step_by_batch(batch.seq_lens.shape[0])
if batch.spec_info is None:
batch.spec_info = self.draft_worker._idle_seed()
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft"),
):
verify_input = self.draft_worker.draft(batch)
assert verify_input.is_verify_input()
batch.spec_info = verify_input
batch_output = self.verify(batch)
# Publish before draft-extend so the fence is at verify-end.
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft_extend"),
):
self.draft_worker._draft_extend_for_decode(batch, batch_output)
return batch_output
python/sglang/srt/speculative/spec_info.py
+6 -10
View File
@@ -146,7 +146,7 @@ class SpeculativeAlgorithm(Enum):
return None return None
def supports_spec_v2(self) -> bool: def supports_spec_v2(self) -> bool:
return (self.is_eagle() and not self.is_frozen_kv_mtp()) or self.is_standalone() return self.is_eagle() or self.is_standalone()
def get_num_tokens_per_bs_for_target_verify( def get_num_tokens_per_bs_for_target_verify(
self, num_draft_tokens: int, is_draft_worker: bool self, num_draft_tokens: int, is_draft_worker: bool
@@ -177,17 +177,13 @@ class SpeculativeAlgorithm(Enum):
return DFlashWorker return DFlashWorker
if self.is_frozen_kv_mtp(): if self.is_frozen_kv_mtp():
if enable_overlap: # V2 worker drives both overlap and non-overlap (scheduler runs it
raise ValueError( # synchronously when overlap is disabled), same as EAGLE.
"FROZEN_KV_MTP does not support spec v2. Disable overlap " from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import (
"scheduling to use FrozenKVMTPWorker." FrozenKVMTPWorkerV2,
)
from sglang.srt.speculative.frozen_kv_mtp_worker import (
FrozenKVMTPWorker,
) )
return FrozenKVMTPWorker return FrozenKVMTPWorkerV2
# EAGLE / EAGLE3 / STANDALONE / MULTI_LAYER always use the V2 worker, # EAGLE / EAGLE3 / STANDALONE / MULTI_LAYER always use the V2 worker,
# even with overlap disabled (scheduler drives it synchronously). # even with overlap disabled (scheduler drives it synchronously).
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py
+8 -8
View File
@@ -29,7 +29,7 @@ from sglang.srt.speculative.frozen_kv_mtp_info import (
FrozenKVMTPContext, FrozenKVMTPContext,
FrozenKVMTPDraftInput, FrozenKVMTPDraftInput,
) )
from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenKVMTPDraftWorker
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from ..attention_methods.dense_attention import ( from ..attention_methods.dense_attention import (
@@ -226,26 +226,26 @@ class _FrozenKVMTPWorkerHarness:
) )
self.model_runner.forward = model_forward self.model_runner.forward = model_forward
self._hidden_size = settings.hidden_size self._hidden_size = settings.hidden_size
self.draft_forward = MethodType(FrozenKVMTPWorker.draft_forward, self) self.draft_forward = MethodType(FrozenKVMTPDraftWorker.draft_forward, self)
self._frozen_kv_target_view = MethodType( self._frozen_kv_target_view = MethodType(
FrozenKVMTPWorker._frozen_kv_target_view, FrozenKVMTPDraftWorker._frozen_kv_target_view,
self, self,
) )
self._target_kv_pool_view = MethodType( self._target_kv_pool_view = MethodType(
FrozenKVMTPWorker._target_kv_pool_view, FrozenKVMTPDraftWorker._target_kv_pool_view,
self, self,
) )
self._set_positions = MethodType(FrozenKVMTPWorker._set_positions, self) self._set_positions = MethodType(FrozenKVMTPDraftWorker._set_positions, self)
self._init_frozen_kv_metadata = MethodType( self._init_frozen_kv_metadata = MethodType(
FrozenKVMTPWorker._init_frozen_kv_metadata, FrozenKVMTPDraftWorker._init_frozen_kv_metadata,
self, self,
) )
self._init_frozen_kv_metadata_capture_cuda_graph = MethodType( self._init_frozen_kv_metadata_capture_cuda_graph = MethodType(
FrozenKVMTPWorker._init_frozen_kv_metadata_capture_cuda_graph, FrozenKVMTPDraftWorker._init_frozen_kv_metadata_capture_cuda_graph,
self, self,
) )
self._init_frozen_kv_metadata_replay_cuda_graph = MethodType( self._init_frozen_kv_metadata_replay_cuda_graph = MethodType(
FrozenKVMTPWorker._init_frozen_kv_metadata_replay_cuda_graph, FrozenKVMTPDraftWorker._init_frozen_kv_metadata_replay_cuda_graph,
self, self,
) )
test/registered/spec/test_frozen_kv_mtp.py
-8
View File
@@ -1,4 +1,3 @@
import os
import unittest import unittest
from types import SimpleNamespace from types import SimpleNamespace
from typing import Optional from typing import Optional
@@ -43,12 +42,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestFrozenKVMTP(CustomTestCase): class TestFrozenKVMTP(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST 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 @classmethod
def _common_server_args(cls) -> list[str]: def _common_server_args(cls) -> list[str]:
args = [ args = [
@@ -110,7 +103,6 @@ class TestFrozenKVMTP(CustomTestCase):
"google/gemma-4-E4B-it", "google/gemma-4-E4B-it",
self.base_url, self.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
env=self._server_env(),
other_args=self._server_args(topk), other_args=self._server_args(topk),
) )
requests.get(self.base_url + "/flush_cache", timeout=30) 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 import unittest
from types import SimpleNamespace from types import SimpleNamespace
from typing import Optional from typing import Optional
@@ -68,12 +67,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestGemma4MTP26BA4B(CustomTestCase): class TestGemma4MTP26BA4B(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST 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 @classmethod
def _common_server_args(cls) -> list[str]: def _common_server_args(cls) -> list[str]:
args = [ args = [
@@ -140,7 +133,6 @@ class TestGemma4MTP26BA4B(CustomTestCase):
TARGET_PATH, TARGET_PATH,
self.base_url, self.base_url,
timeout=SERVER_LAUNCH_TIMEOUT, timeout=SERVER_LAUNCH_TIMEOUT,
env=self._server_env(),
other_args=self._server_args(topk), other_args=self._server_args(topk),
) )
requests.get(self.base_url + "/flush_cache", timeout=30) 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 import unittest
from types import SimpleNamespace from types import SimpleNamespace
from typing import Optional from typing import Optional
@@ -58,12 +57,6 @@ def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
class TestGemma4MTP31B(CustomTestCase): class TestGemma4MTP31B(CustomTestCase):
base_url = DEFAULT_URL_FOR_TEST 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 @classmethod
def _common_server_args(cls) -> list[str]: def _common_server_args(cls) -> list[str]:
args = [ args = [
@@ -127,7 +120,6 @@ class TestGemma4MTP31B(CustomTestCase):
TARGET_PATH, TARGET_PATH,
self.base_url, self.base_url,
timeout=SERVER_LAUNCH_TIMEOUT, timeout=SERVER_LAUNCH_TIMEOUT,
env=self._server_env(),
other_args=self._server_args(topk), other_args=self._server_args(topk),
) )
requests.get(self.base_url + "/flush_cache", timeout=30) 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)