From ed77c5b3be37d03c42b942c4fa799672322e9ce3 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 8 Jun 2026 16:36:48 -0700 Subject: [PATCH 1/6] frozen mtp spec v2; drop v1 worker Co-Authored-By: Khoa Pham Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sglang/srt/arg_groups/speculative_hook.py | 11 +- python/sglang/srt/managers/scheduler.py | 8 +- .../frozen_kv_mtp_cuda_graph_runner.py | 6 +- .../srt/speculative/frozen_kv_mtp_worker.py | 822 ------------------ .../speculative/frozen_kv_mtp_worker_v2.py | 733 +++++++++++++++- python/sglang/srt/speculative/spec_info.py | 16 +- .../runner_modes/speculative_draft_runner.py | 16 +- ...frozen_kv_mtp_all_reqs_finish_in_verify.py | 136 --- 8 files changed, 754 insertions(+), 994 deletions(-) delete mode 100644 python/sglang/srt/speculative/frozen_kv_mtp_worker.py delete mode 100644 test/registered/unit/spec/test_frozen_kv_mtp_all_reqs_finish_in_verify.py diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 055153f89bed..2ea956fa9588 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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." ) - server_args.disable_overlap_schedule = True - logger.warning( - "Overlap scheduler is disabled when using Frozen-KV MTP speculative decoding (spec v2 is not supported yet)." - ) + # SGLANG_ENABLE_SPEC_V2=False selects the non-overlap (synchronous) spec v2 + # path instead of the overlap-scheduled one; both run the V2 worker. + 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: server_args.enable_mixed_chunk = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8e32640fc6a9..e805551be08f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1163,7 +1163,7 @@ def init_overlap(self): self.device_module = torch.get_device_module(self.device) # 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 # safe decision (no accidental opt-out for unaudited shapes). if self.draft_worker is not None: @@ -3118,9 +3118,9 @@ def run_batch( ) batch.input_ids = None else: - # Spec_v1 (NGRAM / DFLASH / FROZEN_KV_MTP, non-overlap): - # worker shape doesn't match req_pool_indices; relay is - # unused (worker rebuilds input_ids inside verify). + # Spec_v1 (NGRAM / DFLASH, non-overlap): worker shape + # doesn't match req_pool_indices; relay is unused (worker + # rebuilds input_ids inside verify). batch.input_ids = batch_result.next_token_ids.to(torch.int64) self.update_cache_from_scheduler(batch, batch_result) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index ae30ece640e1..5aac2da9703d 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -34,7 +34,7 @@ ) if TYPE_CHECKING: - from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker + from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenDraftWorker @dataclass @@ -47,7 +47,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): topk_p: torch.Tensor topk_index: torch.Tensor hidden_states: torch.Tensor - # Consumed by the captured seed iter; see `FrozenKVMTPWorker.draft_forward`. + # Consumed by the captured seed iter; see `FrozenDraftWorker.draft_forward`. bonus_tokens: torch.Tensor global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] @@ -56,7 +56,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): class FrozenKVMTPCudaGraphRunner: """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: FrozenDraftWorker): self.frozen_kv_mtp_worker = frozen_kv_mtp_worker self.model_runner = model_runner = frozen_kv_mtp_worker.draft_model_runner self.graphs = {} diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py deleted file mode 100644 index adb556452a28..000000000000 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ /dev/null @@ -1,822 +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, -) -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(), - ): - 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(), - ): - 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 - 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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 958d09fb4aa7..c568468a68b4 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -11,18 +11,628 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Overlap-scheduling placeholder for frozen-KV MTP (raises until implemented).""" +"""Spec-v2 (overlap-capable) worker for Frozen-KV MTP. + +Two layers, mirroring ``eagle_worker_v2``: + +- ``FrozenDraftWorker(BaseDraftWorker, TpModelWorker)`` -- the assistant. Reads + the target KV cache read-only (it owns no KV pool of its own) and runs the + seed + recurrent draft loop. The draft-loop bodies are ported verbatim from + the former ``FrozenKVMTPWorker``. +- ``FrozenKVMTPWorkerV2(EAGLEWorkerV2)`` -- the orchestrator. The verify / + ``move_accept_tokens`` / forward skeleton is reused from EAGLE unchanged; + only the draft worker and the seed-based "draft extend" are frozen-specific. + +Unlike EAGLE, the frozen draft writes no KV, so "draft extend" is not a model +forward: it just selects the last accepted token + its target hidden state and +stashes them as the next-iter seed. The seed forward runs at the start of the +next draft (``draft_forward``'s seed iter). +""" from __future__ import annotations +import logging 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.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.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, +) +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 FrozenDraftWorker(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__( + 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(), ( + "FrozenDraftWorker should only be instantiated for " + "SpeculativeAlgorithm.FROZEN_KV_MTP, got " + f"{self.speculative_algorithm.name}." + ) + + # 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(): + # 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, + ) + + batch.maybe_evict_swa() + for req in batch.reqs: + req.decode_batch_idx += 1 + + spec_info = batch.spec_info + assert isinstance(spec_info, FrozenKVMTPDraftInput) + + # NOTE: penalty cumulation is done by the inherited + # EagleDraftInputV2Mixin.prepare_for_decode (scheduler-driven, see + # ScheduleBatch.prepare_for_decode), not here -- matching EAGLE v2. + + 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. + """ -class FrozenKVMTPWorkerV2(FrozenKVMTPWorker): def __init__( self, server_args: ServerArgs, @@ -35,8 +645,117 @@ def __init__( nccl_port: int, target_worker: TpModelWorker, ): - raise NotImplementedError( - "FrozenKVMTPWorkerV2 (overlap scheduling for Frozen-KV MTP) is " - "not yet implemented. Pass --disable-overlap-schedule to use " - "FrozenKVMTPWorker." + # 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 FrozenDraftWorker. + 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 = FrozenDraftWorker( + 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. + 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 FrozenDraftWorker'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(), + ): + 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(), + ): + 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(), + ): + self.draft_worker._draft_extend_for_decode(batch, batch_output) + + return batch_output diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 786aed78f4a4..3b12140a0221 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -146,7 +146,7 @@ def build_disagg_draft_input( return None 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( self, num_draft_tokens: int, is_draft_worker: bool @@ -177,17 +177,13 @@ def create_worker( return DFlashWorker if self.is_frozen_kv_mtp(): - if enable_overlap: - raise ValueError( - "FROZEN_KV_MTP does not support spec v2. Disable overlap " - "scheduling to use FrozenKVMTPWorker." - ) - - from sglang.srt.speculative.frozen_kv_mtp_worker import ( - FrozenKVMTPWorker, + # V2 worker drives both overlap and non-overlap (scheduler runs it + # synchronously when overlap is disabled), same as EAGLE. + from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import ( + FrozenKVMTPWorkerV2, ) - return FrozenKVMTPWorker + return FrozenKVMTPWorkerV2 # EAGLE / EAGLE3 / STANDALONE / MULTI_LAYER always use the V2 worker, # even with overlap disabled (scheduler drives it synchronously). diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index a064c8092c34..96352f8601f8 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -29,7 +29,7 @@ FrozenKVMTPContext, FrozenKVMTPDraftInput, ) -from sglang.srt.speculative.frozen_kv_mtp_worker import FrozenKVMTPWorker +from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenDraftWorker from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from ..attention_methods.dense_attention import ( @@ -226,26 +226,26 @@ def __init__( ) self.model_runner.forward = model_forward self._hidden_size = settings.hidden_size - self.draft_forward = MethodType(FrozenKVMTPWorker.draft_forward, self) + self.draft_forward = MethodType(FrozenDraftWorker.draft_forward, self) self._frozen_kv_target_view = MethodType( - FrozenKVMTPWorker._frozen_kv_target_view, + FrozenDraftWorker._frozen_kv_target_view, self, ) self._target_kv_pool_view = MethodType( - FrozenKVMTPWorker._target_kv_pool_view, + FrozenDraftWorker._target_kv_pool_view, self, ) - self._set_positions = MethodType(FrozenKVMTPWorker._set_positions, self) + self._set_positions = MethodType(FrozenDraftWorker._set_positions, self) self._init_frozen_kv_metadata = MethodType( - FrozenKVMTPWorker._init_frozen_kv_metadata, + FrozenDraftWorker._init_frozen_kv_metadata, self, ) self._init_frozen_kv_metadata_capture_cuda_graph = MethodType( - FrozenKVMTPWorker._init_frozen_kv_metadata_capture_cuda_graph, + FrozenDraftWorker._init_frozen_kv_metadata_capture_cuda_graph, self, ) self._init_frozen_kv_metadata_replay_cuda_graph = MethodType( - FrozenKVMTPWorker._init_frozen_kv_metadata_replay_cuda_graph, + FrozenDraftWorker._init_frozen_kv_metadata_replay_cuda_graph, self, ) diff --git a/test/registered/unit/spec/test_frozen_kv_mtp_all_reqs_finish_in_verify.py b/test/registered/unit/spec/test_frozen_kv_mtp_all_reqs_finish_in_verify.py deleted file mode 100644 index ce70eec57236..000000000000 --- a/test/registered/unit/spec/test_frozen_kv_mtp_all_reqs_finish_in_verify.py +++ /dev/null @@ -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) From a7266a4d580f68f4789418b75e2b37f9548ee4c2 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 8 Jun 2026 16:44:36 -0700 Subject: [PATCH 2/6] rename FrozenDraftWorker -> FrozenKVMTPDraftWorker Co-Authored-By: Khoa Pham Co-Authored-By: Claude Opus 4.8 (1M context) --- .../frozen_kv_mtp_cuda_graph_runner.py | 6 +++--- .../srt/speculative/frozen_kv_mtp_worker_v2.py | 12 ++++++------ .../runner_modes/speculative_draft_runner.py | 16 ++++++++-------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 5aac2da9703d..dffdcea45700 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -34,7 +34,7 @@ ) if TYPE_CHECKING: - from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenDraftWorker + from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenKVMTPDraftWorker @dataclass @@ -47,7 +47,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): topk_p: torch.Tensor topk_index: torch.Tensor hidden_states: torch.Tensor - # Consumed by the captured seed iter; see `FrozenDraftWorker.draft_forward`. + # Consumed by the captured seed iter; see `FrozenKVMTPDraftWorker.draft_forward`. bonus_tokens: torch.Tensor global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] @@ -56,7 +56,7 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): class FrozenKVMTPCudaGraphRunner: """CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step.""" - def __init__(self, frozen_kv_mtp_worker: FrozenDraftWorker): + def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPDraftWorker): self.frozen_kv_mtp_worker = frozen_kv_mtp_worker self.model_runner = model_runner = frozen_kv_mtp_worker.draft_model_runner self.graphs = {} diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index c568468a68b4..3b3adeea322d 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -15,7 +15,7 @@ Two layers, mirroring ``eagle_worker_v2``: -- ``FrozenDraftWorker(BaseDraftWorker, TpModelWorker)`` -- the assistant. Reads +- ``FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker)`` -- the assistant. Reads the target KV cache read-only (it owns no KV pool of its own) and runs the seed + recurrent draft loop. The draft-loop bodies are ported verbatim from the former ``FrozenKVMTPWorker``. @@ -85,7 +85,7 @@ logger = logging.getLogger(__name__) -class FrozenDraftWorker(BaseDraftWorker, TpModelWorker): +class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker): """Frozen-KV MTP draft worker. The assistant reads target KV only. It reuses EAGLE's verify input/output @@ -117,7 +117,7 @@ def __init__( server_args.speculative_algorithm ) assert self.speculative_algorithm.is_frozen_kv_mtp(), ( - "FrozenDraftWorker should only be instantiated for " + "FrozenKVMTPDraftWorker should only be instantiated for " "SpeculativeAlgorithm.FROZEN_KV_MTP, got " f"{self.speculative_algorithm.name}." ) @@ -647,7 +647,7 @@ def __init__( ): # 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 FrozenDraftWorker. + # 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 @@ -667,7 +667,7 @@ def __init__( # 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 = FrozenDraftWorker( + self._draft_worker = FrozenKVMTPDraftWorker( server_args, gpu_id, tp_rank, @@ -701,7 +701,7 @@ def spec_v2_attn_backends(self) -> tuple: 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 FrozenDraftWorker's. + # 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 diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 96352f8601f8..951b3b1412a1 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -29,7 +29,7 @@ FrozenKVMTPContext, FrozenKVMTPDraftInput, ) -from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenDraftWorker +from sglang.srt.speculative.frozen_kv_mtp_worker_v2 import FrozenKVMTPDraftWorker from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from ..attention_methods.dense_attention import ( @@ -226,26 +226,26 @@ def __init__( ) self.model_runner.forward = model_forward self._hidden_size = settings.hidden_size - self.draft_forward = MethodType(FrozenDraftWorker.draft_forward, self) + self.draft_forward = MethodType(FrozenKVMTPDraftWorker.draft_forward, self) self._frozen_kv_target_view = MethodType( - FrozenDraftWorker._frozen_kv_target_view, + FrozenKVMTPDraftWorker._frozen_kv_target_view, self, ) self._target_kv_pool_view = MethodType( - FrozenDraftWorker._target_kv_pool_view, + FrozenKVMTPDraftWorker._target_kv_pool_view, self, ) - self._set_positions = MethodType(FrozenDraftWorker._set_positions, self) + self._set_positions = MethodType(FrozenKVMTPDraftWorker._set_positions, self) self._init_frozen_kv_metadata = MethodType( - FrozenDraftWorker._init_frozen_kv_metadata, + FrozenKVMTPDraftWorker._init_frozen_kv_metadata, self, ) self._init_frozen_kv_metadata_capture_cuda_graph = MethodType( - FrozenDraftWorker._init_frozen_kv_metadata_capture_cuda_graph, + FrozenKVMTPDraftWorker._init_frozen_kv_metadata_capture_cuda_graph, self, ) self._init_frozen_kv_metadata_replay_cuda_graph = MethodType( - FrozenDraftWorker._init_frozen_kv_metadata_replay_cuda_graph, + FrozenKVMTPDraftWorker._init_frozen_kv_metadata_replay_cuda_graph, self, ) From f45b55f3ff60a7713b3606b1c524813ee2d3d1a1 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 8 Jun 2026 16:45:04 -0700 Subject: [PATCH 3/6] trim frozen_kv_mtp_worker_v2 module docstring Co-Authored-By: Khoa Pham Co-Authored-By: Claude Opus 4.8 (1M context) --- .../speculative/frozen_kv_mtp_worker_v2.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 3b3adeea322d..14611b3f6b29 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -11,22 +11,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Spec-v2 (overlap-capable) worker for Frozen-KV MTP. - -Two layers, mirroring ``eagle_worker_v2``: - -- ``FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker)`` -- the assistant. Reads - the target KV cache read-only (it owns no KV pool of its own) and runs the - seed + recurrent draft loop. The draft-loop bodies are ported verbatim from - the former ``FrozenKVMTPWorker``. -- ``FrozenKVMTPWorkerV2(EAGLEWorkerV2)`` -- the orchestrator. The verify / - ``move_accept_tokens`` / forward skeleton is reused from EAGLE unchanged; - only the draft worker and the seed-based "draft extend" are frozen-specific. - -Unlike EAGLE, the frozen draft writes no KV, so "draft extend" is not a model -forward: it just selects the last accepted token + its target hidden state and -stashes them as the next-iter seed. The seed forward runs at the start of the -next draft (``draft_forward``'s seed iter). +"""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 2e97fa72102d414c38ad05aa87ea8058b4d5459f Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 9 Jun 2026 00:38:02 -0700 Subject: [PATCH 4/6] nvtx spans on frozen v2 orchestrator --- python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 14611b3f6b29..84ecd8696f48 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -64,6 +64,7 @@ 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 ( @@ -710,6 +711,7 @@ def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None): ), 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( @@ -731,6 +733,7 @@ def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None): ), 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() @@ -745,6 +748,7 @@ def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None): ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), + spec_stage_span("draft_extend"), ): self.draft_worker._draft_extend_for_decode(batch, batch_output) From d7ae6c1caac1bde922d71b1e92f9f85cf5df8439 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 9 Jun 2026 13:16:34 -0700 Subject: [PATCH 5/6] fix double bookkeeping in frozen v2 draft --- python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 84ecd8696f48..7ab227d6202f 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -396,16 +396,15 @@ def draft(self, batch: ScheduleBatch): 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) - # NOTE: penalty cumulation is done by the inherited + # 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 From 2ac14345e259752daa7d7719853ef0d8428374f7 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 9 Jun 2026 14:03:41 -0700 Subject: [PATCH 6/6] drop v1 dead code; un-pin overlap in tests; assert no adaptive --- .../srt/speculative/frozen_kv_mtp_info.py | 26 +---------------- .../srt/speculative/frozen_kv_mtp_utils.py | 29 ++----------------- .../speculative/frozen_kv_mtp_worker_v2.py | 6 ++-- test/registered/spec/test_frozen_kv_mtp.py | 8 ----- .../spec/test_gemma4_mtp_26b_a4b_extra.py | 8 ----- .../spec/test_gemma4_mtp_31b_extra.py | 8 ----- 6 files changed, 6 insertions(+), 79 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_info.py b/python/sglang/srt/speculative/frozen_kv_mtp_info.py index 7b562b52094f..1ff0f8c13314 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_info.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_info.py @@ -13,7 +13,7 @@ # ============================================================================== from __future__ import annotations -from dataclasses import dataclass, fields +from dataclasses import dataclass from typing import Dict from sglang.srt.mem_cache.memory_pool import KVCache @@ -21,7 +21,6 @@ EagleDraftExtendInput, EagleDraftInput, EagleVerifyInput, - EagleVerifyOutput, ) from sglang.srt.speculative.spec_info import SpecInput, SpecInputType @@ -68,26 +67,3 @@ class FrozenKVMTPVerifyInput(EagleVerifyInput): def __post_init__(self): 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) - } - ) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py index dbd63c2e444c..69d121137f67 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py @@ -14,19 +14,13 @@ from __future__ import annotations from contextlib import contextmanager -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING import torch -from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.speculative.frozen_kv_mtp_info import ( - FrozenKVMTPContext, - FrozenKVMTPDraftExtendInput, - FrozenKVMTPDraftInput, -) -from sglang.srt.speculative.spec_utils import fast_topk +from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext if TYPE_CHECKING: 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) last_indices = torch.cumsum(lens, dim=0) - 1 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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 7ab227d6202f..8d9cbc19d286 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -113,9 +113,6 @@ def __init__( f"{self.speculative_algorithm.name}." ) - # 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 @@ -670,6 +667,9 @@ def __init__( ) # 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). diff --git a/test/registered/spec/test_frozen_kv_mtp.py b/test/registered/spec/test_frozen_kv_mtp.py index 9d573880a56b..c0d1705cbb85 100644 --- a/test/registered/spec/test_frozen_kv_mtp.py +++ b/test/registered/spec/test_frozen_kv_mtp.py @@ -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 @@ def _run_gsm8k_mtp(self, topk: int) -> None: "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) diff --git a/test/registered/spec/test_gemma4_mtp_26b_a4b_extra.py b/test/registered/spec/test_gemma4_mtp_26b_a4b_extra.py index 57c2634715c4..db6f4eccb7f3 100644 --- a/test/registered/spec/test_gemma4_mtp_26b_a4b_extra.py +++ b/test/registered/spec/test_gemma4_mtp_26b_a4b_extra.py @@ -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 @@ def _run_gsm8k_mtp(self, topk: int) -> None: 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) diff --git a/test/registered/spec/test_gemma4_mtp_31b_extra.py b/test/registered/spec/test_gemma4_mtp_31b_extra.py index c359a379e876..608848a562dc 100644 --- a/test/registered/spec/test_gemma4_mtp_31b_extra.py +++ b/test/registered/spec/test_gemma4_mtp_31b_extra.py @@ -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 @@ def _run_gsm8k_mtp(self, topk: int) -> None: 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)