Support spec v2 for Frozen-KV MTP; remove v1 worker (#27607)
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
co-authored by
Khoa Pham
Claude Opus 4.8
Khoa Pham
Claude Opus 4.8
parent
7f730edfdc
commit
decb88e0e3
@@ -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
|
||||
|
||||
@@ -1163,7 +1163,7 @@ class Scheduler(
|
||||
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:
|
||||
@@ -3120,9 +3120,9 @@ class Scheduler(
|
||||
)
|
||||
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)
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ from sglang.srt.utils import (
|
||||
)
|
||||
|
||||
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
|
||||
@@ -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 `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: FrozenKVMTPWorker):
|
||||
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 = {}
|
||||
|
||||
@@ -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 @@ from sglang.srt.speculative.eagle_info import (
|
||||
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)
|
||||
}
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -11,18 +11,79 @@
|
||||
# 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 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
|
||||
|
||||
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,
|
||||
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__(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
@@ -35,8 +96,659 @@ class FrozenKVMTPWorkerV2(FrozenKVMTPWorker):
|
||||
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."
|
||||
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(), (
|
||||
"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
|
||||
|
||||
@@ -146,7 +146,7 @@ class SpeculativeAlgorithm(Enum):
|
||||
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 @@ class SpeculativeAlgorithm(Enum):
|
||||
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).
|
||||
|
||||
@@ -29,7 +29,7 @@ from sglang.srt.speculative.frozen_kv_mtp_info import (
|
||||
FrozenKVMTPContext,
|
||||
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 ..attention_methods.dense_attention import (
|
||||
@@ -226,26 +226,26 @@ class _FrozenKVMTPWorkerHarness:
|
||||
)
|
||||
self.model_runner.forward = model_forward
|
||||
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(
|
||||
FrozenKVMTPWorker._frozen_kv_target_view,
|
||||
FrozenKVMTPDraftWorker._frozen_kv_target_view,
|
||||
self,
|
||||
)
|
||||
self._target_kv_pool_view = MethodType(
|
||||
FrozenKVMTPWorker._target_kv_pool_view,
|
||||
FrozenKVMTPDraftWorker._target_kv_pool_view,
|
||||
self,
|
||||
)
|
||||
self._set_positions = MethodType(FrozenKVMTPWorker._set_positions, self)
|
||||
self._set_positions = MethodType(FrozenKVMTPDraftWorker._set_positions, self)
|
||||
self._init_frozen_kv_metadata = MethodType(
|
||||
FrozenKVMTPWorker._init_frozen_kv_metadata,
|
||||
FrozenKVMTPDraftWorker._init_frozen_kv_metadata,
|
||||
self,
|
||||
)
|
||||
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._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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user