[Spec] Deprecate Spec V1 (#25464)
This commit is contained in:
Liangsheng Yin
@@ -285,22 +285,22 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
|
||||
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
|
||||
)
|
||||
|
||||
spec_v1_reason = None
|
||||
# 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
|
||||
spec_v1_reason = "SGLANG_ENABLE_SPEC_V2=False"
|
||||
|
||||
if server_args.disable_overlap_schedule:
|
||||
logger.warning(
|
||||
"Spec v1 is used for eagle/eagle3/standalone speculative decoding because %s.",
|
||||
spec_v1_reason or "overlap schedule is disabled",
|
||||
"Non-overlap (synchronous) spec v2 is used for eagle/eagle3/standalone "
|
||||
"speculative decoding."
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Spec v2 is enabled by default for eagle/eagle3/standalone speculative decoding."
|
||||
"Overlap spec v2 is enabled by default for eagle/eagle3/standalone speculative decoding."
|
||||
)
|
||||
|
||||
if server_args.enable_mixed_chunk:
|
||||
|
||||
@@ -9,25 +9,6 @@ from sglang.srt.environ import envs
|
||||
|
||||
_PR_REVERT_YAML_25015 = """
|
||||
patches:
|
||||
- target: sglang.srt.speculative.eagle_worker.EAGLEWorker.draft_forward
|
||||
edits:
|
||||
- match: |
|
||||
forward_batch.out_cache_loc = out_cache_loc[i]
|
||||
spec_info.hidden_states = hidden_states
|
||||
replacement: |
|
||||
forward_batch.out_cache_loc = out_cache_loc[i]
|
||||
forward_batch.positions.add_(1)
|
||||
spec_info.hidden_states = hidden_states
|
||||
- match: |
|
||||
hidden_states = logits_output.hidden_states
|
||||
maybe_detect_nan(hidden_states, f"draft_forward step {i}: hidden_states")
|
||||
maybe_detect_inf(hidden_states, f"draft_forward step {i}: hidden_states")
|
||||
forward_batch.positions.add_(1)
|
||||
replacement: |
|
||||
hidden_states = logits_output.hidden_states
|
||||
maybe_detect_nan(hidden_states, f"draft_forward step {i}: hidden_states")
|
||||
maybe_detect_inf(hidden_states, f"draft_forward step {i}: hidden_states")
|
||||
|
||||
- target: sglang.srt.speculative.eagle_worker_v2.EagleDraftWorker.draft_forward
|
||||
edits:
|
||||
- match: |
|
||||
|
||||
+2
-2
@@ -27,11 +27,11 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
|
||||
|
||||
|
||||
class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
|
||||
def __init__(self, eagle_worker: EAGLEWorker):
|
||||
def __init__(self, eagle_worker: EagleDraftWorker):
|
||||
super().__init__(eagle_worker)
|
||||
|
||||
def _create_graph(self):
|
||||
|
||||
@@ -29,7 +29,7 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
|
||||
|
||||
from sglang.srt.utils import is_npu
|
||||
|
||||
@@ -45,7 +45,7 @@ if is_npu():
|
||||
|
||||
|
||||
class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
|
||||
def __init__(self, eagle_worker: EAGLEWorker):
|
||||
def __init__(self, eagle_worker: EagleDraftWorker):
|
||||
super().__init__(eagle_worker)
|
||||
self.update_attr_name = None
|
||||
self.update_attr_type = None
|
||||
|
||||
@@ -106,8 +106,9 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None:
|
||||
batch.input_ids, future_map.output_tokens_buf, batch.req_pool_indices
|
||||
)
|
||||
|
||||
# spec_v1 (non-overlap spec) doesn't relay extras; only spec_v2 does.
|
||||
if batch.is_spec_v2:
|
||||
# Only the overlap path relays spec extras through the future_map; the
|
||||
# synchronous (non-overlap) V2 path installs next_draft_input directly.
|
||||
if batch.enable_overlap and batch.is_spec_v2:
|
||||
future_map._resolve_spec_extras(batch)
|
||||
|
||||
|
||||
|
||||
@@ -2457,10 +2457,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
|
||||
@property
|
||||
def is_spec_v2(self):
|
||||
# FIXME: finally deprecate is_spec_v2
|
||||
ret = self.enable_overlap and not self.spec_algorithm.is_none()
|
||||
assert not ret or self.spec_algorithm.supports_spec_v2()
|
||||
return ret
|
||||
# Whether the V2 worker/schema is used. Independent of overlap: the
|
||||
# non-overlap path also drives the V2 worker, just synchronously.
|
||||
return self.spec_algorithm.supports_spec_v2()
|
||||
|
||||
def mamba_lazy_prealloc_at_boundary(self, mamba_track_interval: int):
|
||||
"""Allocate a temporary second ping-pong slot for reqs at a track boundary.
|
||||
|
||||
@@ -493,7 +493,6 @@ class Scheduler(
|
||||
spec_algorithm=self.spec_algorithm,
|
||||
server_args=self.server_args,
|
||||
enable_hierarchical_cache=self.enable_hierarchical_cache,
|
||||
enable_overlap=self.enable_overlap,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
|
||||
@@ -1041,7 +1040,6 @@ class Scheduler(
|
||||
draft_worker=self.draft_worker,
|
||||
spec_algorithm=self.spec_algorithm,
|
||||
server_args=self.server_args,
|
||||
enable_overlap=self.enable_overlap,
|
||||
)
|
||||
# Default to the target model_config so the MetadataBuffers branches
|
||||
# below can always access it; overridden by the draft model_config
|
||||
@@ -2941,8 +2939,8 @@ class Scheduler(
|
||||
self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot]
|
||||
|
||||
@contextmanager
|
||||
def _overlap_forward_isolation(self, batch: ScheduleBatch):
|
||||
"""Make SB transactional across one overlap forward.
|
||||
def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool):
|
||||
"""Make SB transactional across one forward (overlap and non-overlap).
|
||||
|
||||
1. Snapshot SB fields so V2's mid-forward mutations (forward_mode /
|
||||
input_ids / seq_lens / spec_info / ...) can be undone. V1 / non-spec
|
||||
@@ -2951,10 +2949,12 @@ class Scheduler(
|
||||
2. Substitute sampling_info with a forward-only copy (orchestrator=None,
|
||||
shares the pre-accumulated penalty buffer) so V2's multiple init_new
|
||||
calls don't double-accumulate penalties.
|
||||
3. Pin (batch, snapshot) into batch_record_buf for 2 iters so GPU
|
||||
tensors in the snapshot survive the caching allocator past the
|
||||
forward stream. Must run AFTER the sampling_info swap so the
|
||||
forward-only copy gets pinned.
|
||||
3. (overlap=True only) Pin (batch, snapshot) into batch_record_buf
|
||||
for 2 iters so GPU tensors in the snapshot survive the caching
|
||||
allocator past the forward stream. Must run AFTER the sampling_info
|
||||
swap so the forward-only copy gets pinned. The non-overlap (sync) path
|
||||
runs on a single stream and doesn't allocate batch_record_buf, so it
|
||||
passes overlap=False.
|
||||
"""
|
||||
# 1. snapshot
|
||||
snapshot_v2_full = batch.is_spec_v2
|
||||
@@ -2969,8 +2969,9 @@ class Scheduler(
|
||||
if sched_sampling_info is not None:
|
||||
batch.sampling_info = sched_sampling_info.copy_for_forward()
|
||||
|
||||
# 3. pin for 2-iter tensor lifetime
|
||||
self.record_batch_in_overlap(batch)
|
||||
# 3. pin for 2-iter tensor lifetime (overlap path only)
|
||||
if overlap:
|
||||
self.record_batch_in_overlap(batch)
|
||||
|
||||
try:
|
||||
yield
|
||||
@@ -3018,7 +3019,7 @@ class Scheduler(
|
||||
# post-forward must not un-consume staging.
|
||||
resolve_forward_inputs(batch, self.future_map)
|
||||
|
||||
with self._overlap_forward_isolation(batch):
|
||||
with self._forward_isolation(batch, overlap=True):
|
||||
future_indices = batch.req_pool_indices
|
||||
|
||||
# Spec_v2 fires on_publish mid-worker (between verify and
|
||||
@@ -3077,6 +3078,28 @@ class Scheduler(
|
||||
batch.req_pool_indices, batch_result.next_token_ids
|
||||
)
|
||||
batch.input_ids = None
|
||||
elif batch.is_spec_v2:
|
||||
# Non-overlap V2: drive the V2 worker synchronously (no
|
||||
# future_map relay / on_publish).
|
||||
resolve_forward_inputs(batch, self.future_map)
|
||||
with self._forward_isolation(batch, overlap=False):
|
||||
batch_result = self.model_worker.forward_batch_generation(batch)
|
||||
# The isolation restore reverted the worker's in-forward SB edits;
|
||||
# re-apply what must carry to the next iter.
|
||||
batch.spec_info = batch_result.next_draft_input
|
||||
if batch_result.new_seq_lens is not None:
|
||||
batch.seq_lens = batch_result.new_seq_lens
|
||||
if batch.seq_lens_cpu is not None:
|
||||
batch.seq_lens_cpu = batch_result.new_seq_lens.to("cpu")
|
||||
batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
|
||||
batch.input_ids = None # rebuilt next iter from draft_token
|
||||
self.update_cache_from_scheduler(batch, batch_result)
|
||||
# Sync D2H so the result processor can read CPU tensors.
|
||||
batch_result.copy_done = self.device_module.Event()
|
||||
batch_result.copy_to_cpu(
|
||||
return_logprob=batch.return_logprob,
|
||||
return_hidden_states=batch.return_hidden_states,
|
||||
)
|
||||
else:
|
||||
kwargs = (
|
||||
{"pp_proxy_tensors": pp_proxy_tensors}
|
||||
@@ -3095,9 +3118,9 @@ class Scheduler(
|
||||
)
|
||||
batch.input_ids = None
|
||||
else:
|
||||
# Spec_v1 (non-overlap spec): worker shape doesn't match
|
||||
# req_pool_indices; relay is unused (worker rebuilds input_ids
|
||||
# inside verify). Keep pre-PR behavior.
|
||||
# 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).
|
||||
batch.input_ids = batch_result.next_token_ids.to(torch.int64)
|
||||
self.update_cache_from_scheduler(batch, batch_result)
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_draft_model_runner(draft_worker):
|
||||
# EAGLEWorker (v1): draft_model_runner property -> self.model_runner
|
||||
# DFlashWorker: exposes draft_model_runner directly
|
||||
runner = getattr(draft_worker, "draft_model_runner", None)
|
||||
if runner is not None:
|
||||
return runner
|
||||
|
||||
@@ -48,14 +48,15 @@ def get_draft_kv_pool(
|
||||
draft_worker: "BaseTpWorker",
|
||||
spec_algorithm: SpeculativeAlgorithm,
|
||||
server_args: ServerArgs,
|
||||
enable_overlap: bool,
|
||||
):
|
||||
"""Return (draft_token_to_kv_pool, draft_model_config) for the current
|
||||
draft worker, or (None, None) when no draft KV pool is available."""
|
||||
if draft_worker is None or spec_algorithm.is_ngram():
|
||||
return None, None
|
||||
|
||||
if spec_algorithm.supports_spec_v2() and enable_overlap:
|
||||
# V2 (EAGLE family) nests the runner under `.draft_worker`; DFLASH /
|
||||
# FROZEN_KV_MTP expose `.model_runner` directly.
|
||||
if spec_algorithm.supports_spec_v2():
|
||||
if server_args.enable_multi_layer_eagle:
|
||||
draft_runner = draft_worker.draft_worker.draft_runner_list[0]
|
||||
else:
|
||||
@@ -75,7 +76,6 @@ def maybe_register_hicache_draft(
|
||||
spec_algorithm: SpeculativeAlgorithm,
|
||||
server_args: ServerArgs,
|
||||
enable_hierarchical_cache: bool,
|
||||
enable_overlap: bool,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
"""Register draft KV pool with HiCacheController for piggyback L2/L3 ops."""
|
||||
@@ -86,7 +86,6 @@ def maybe_register_hicache_draft(
|
||||
draft_worker=draft_worker,
|
||||
spec_algorithm=spec_algorithm,
|
||||
server_args=server_args,
|
||||
enable_overlap=enable_overlap,
|
||||
)
|
||||
if draft_kv_pool is None:
|
||||
return
|
||||
|
||||
@@ -62,7 +62,7 @@ def adaptive_unsupported_reason(server_args: ServerArgs) -> str | None:
|
||||
if server_args.enable_multi_layer_eagle:
|
||||
return (
|
||||
"enable_multi_layer_eagle=True is not supported "
|
||||
"(MultiLayerEagleWorker does not implement adaptive)"
|
||||
"(MultiLayerEagleWorkerV2 does not implement adaptive)"
|
||||
)
|
||||
if server_args.enable_two_batch_overlap:
|
||||
return (
|
||||
|
||||
@@ -37,7 +37,7 @@ from sglang.srt.utils import (
|
||||
from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -62,7 +62,7 @@ class EagleDraftInputBuffers(ForwardInputBuffers):
|
||||
class EAGLEDraftCudaGraphRunner:
|
||||
def __init__(
|
||||
self,
|
||||
eagle_worker: EAGLEWorker,
|
||||
eagle_worker: EagleDraftWorker,
|
||||
*,
|
||||
draft_attn_backend=None,
|
||||
speculative_num_steps: Optional[int] = None,
|
||||
|
||||
@@ -40,7 +40,7 @@ from sglang.srt.utils import (
|
||||
_is_hip = is_hip()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -64,7 +64,7 @@ class EagleDraftExtendInputBuffers(ForwardInputBuffers):
|
||||
class EAGLEDraftExtendCudaGraphRunner:
|
||||
def __init__(
|
||||
self,
|
||||
eagle_worker: EAGLEWorker,
|
||||
eagle_worker: EagleDraftWorker,
|
||||
*,
|
||||
draft_extend_attn_backend=None,
|
||||
speculative_num_steps: Optional[int] = None,
|
||||
|
||||
@@ -57,11 +57,11 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _draft_runner_of(worker):
|
||||
"""Draft model_runner accessor that handles v1 / v2 worker naming.
|
||||
"""Draft model_runner accessor across worker shapes.
|
||||
|
||||
v1 (`EAGLEWorker` and subclasses) exposes the draft model_runner as
|
||||
`model_runner` (the worker itself runs the draft model);
|
||||
v2 (`EagleDraftWorker` and subclasses) exposes it as `draft_runner`.
|
||||
v2 draft workers (`EagleDraftWorker` and subclasses) expose the draft
|
||||
model_runner as `draft_runner`; fall back to `model_runner` for workers
|
||||
that run the draft model directly.
|
||||
"""
|
||||
return (
|
||||
worker.draft_runner if hasattr(worker, "draft_runner") else worker.model_runner
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -83,8 +83,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FrozenKVMTPWorker(TpModelWorker):
|
||||
"""Frozen-KV MTP worker; same constructor shape as EAGLEWorker. Entry:
|
||||
:meth:`forward_batch_generation` (stubs for now).
|
||||
"""Frozen-KV MTP worker; same constructor shape as other TpModelWorker-based
|
||||
spec workers. Entry: :meth:`forward_batch_generation` (stubs for now).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -1,821 +0,0 @@
|
||||
# Copyright 2023-2024 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.
|
||||
# ==============================================================================
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_group
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.moe.utils import 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.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.draft_utils import DraftBackendFactory
|
||||
from sglang.srt.speculative.eagle_info import (
|
||||
EagleDraftExtendInput,
|
||||
EagleDraftInput,
|
||||
EagleVerifyInput,
|
||||
EagleVerifyOutput,
|
||||
)
|
||||
from sglang.srt.speculative.eagle_utils import (
|
||||
apply_eagle_prefill_input_rotation,
|
||||
build_tree_kernel_efficient,
|
||||
organize_draft_results,
|
||||
)
|
||||
from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
|
||||
MultiLayerEagleDraftExtendCudaGraphRunner,
|
||||
)
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.speculative.spec_utils import (
|
||||
draft_tp_context,
|
||||
fast_topk,
|
||||
generate_token_bitmask,
|
||||
load_token_map,
|
||||
select_top_k_tokens,
|
||||
)
|
||||
from sglang.srt.utils import empty_context, get_available_gpu_memory, is_cuda, is_npu
|
||||
from sglang.srt.utils.async_probe import maybe_detect_nan
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
|
||||
_is_npu = is_npu()
|
||||
|
||||
if is_cuda():
|
||||
from sgl_kernel import segment_packbits # noqa: F401
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MultiLayerEagleWorker(TpModelWorker):
|
||||
|
||||
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,
|
||||
):
|
||||
# Parse arguments
|
||||
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
|
||||
assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, (
|
||||
"multi-layer EAGLE requires speculative_num_draft_tokens == "
|
||||
"speculative_num_steps + 1, "
|
||||
f"got {self.speculative_num_draft_tokens} and {self.speculative_num_steps}"
|
||||
)
|
||||
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.draft_extend_attn_backend_list = []
|
||||
|
||||
# Override the context length of the draft model to be the same as the target model.
|
||||
server_args.context_length = target_worker.model_runner.model_config.context_len
|
||||
|
||||
# Do not capture cuda graph in `super().__init__()`
|
||||
# It will be captured later.
|
||||
backup_disable_cuda_graph = server_args.disable_cuda_graph
|
||||
server_args.disable_cuda_graph = True
|
||||
# Share the allocator with a target worker.
|
||||
# Draft and target worker own their own KV cache pools.
|
||||
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
|
||||
target_worker.get_memory_pool()
|
||||
)
|
||||
|
||||
# Load hot token ids
|
||||
if self.speculative_algorithm.is_eagle3():
|
||||
if server_args.speculative_token_map is not None:
|
||||
logger.warning(
|
||||
"Speculative token map specified, but EAGLE3 models already have this. Ignoring the specified token map."
|
||||
)
|
||||
self.hot_token_id = None
|
||||
elif server_args.speculative_token_map is not None:
|
||||
self.hot_token_id = load_token_map(server_args.speculative_token_map)
|
||||
server_args.json_model_override_args = (
|
||||
f'{{"hot_vocab_size": {len(self.hot_token_id)}}}'
|
||||
)
|
||||
else:
|
||||
self.hot_token_id = None
|
||||
|
||||
# Init draft worker
|
||||
if server_args.enable_dp_attention and self.speculative_algorithm.is_eagle3():
|
||||
ctx = draft_tp_context(get_attention_tp_group())
|
||||
else:
|
||||
ctx = empty_context()
|
||||
with ctx, speculative_moe_backend_context():
|
||||
super().__init__(
|
||||
server_args=server_args,
|
||||
gpu_id=gpu_id,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=0, # spec workers don't support pipeline parallelism
|
||||
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=target_worker.model_runner.memory_pool_config,
|
||||
is_multi_layer_eagle=True,
|
||||
)
|
||||
|
||||
self.eagle_use_aux_hidden_state = False
|
||||
if self.speculative_algorithm.is_eagle3():
|
||||
eagle_config = getattr(
|
||||
self.model_runner.model_config.hf_config, "eagle_config", {}
|
||||
)
|
||||
self.eagle_use_aux_hidden_state = eagle_config.get(
|
||||
"use_aux_hidden_state", True
|
||||
)
|
||||
|
||||
embed, head = self.target_worker.model_runner.model.get_embed_and_head()
|
||||
|
||||
if self.speculative_algorithm.is_eagle3():
|
||||
# most cases EAGLE3 models don't share lm_head
|
||||
# but some models (e.g. nvidia/gpt-oss-120b-Eagle3) shares
|
||||
if (
|
||||
hasattr(self.draft_model_runner.model, "load_lm_head_from_target")
|
||||
and self.draft_model_runner.model.load_lm_head_from_target
|
||||
):
|
||||
self.draft_model_runner.model.set_embed_and_head(embed, head)
|
||||
else:
|
||||
self.draft_model_runner.model.set_embed(embed)
|
||||
|
||||
# grab hot token ids
|
||||
if self.draft_model_runner.model.hot_token_id is not None:
|
||||
self.hot_token_id = self.draft_model_runner.model.hot_token_id.to(
|
||||
embed.device
|
||||
)
|
||||
|
||||
else:
|
||||
if self.hot_token_id is not None:
|
||||
head = head.clone()
|
||||
self.hot_token_id = self.hot_token_id.to(head.device)
|
||||
head.data = head.data[self.hot_token_id]
|
||||
|
||||
# Share the embedding and lm_head
|
||||
for i in range(self.speculative_num_steps):
|
||||
self.mtp_model_runner(i).model.set_embed_and_head(embed, head)
|
||||
|
||||
# Init attention backend and cuda graphs
|
||||
for i in range(self.speculative_num_steps):
|
||||
self.mtp_model_runner(i).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
|
||||
)
|
||||
with (
|
||||
self.draft_tp_context(self.mtp_model_runner(0).tp_group),
|
||||
speculative_moe_backend_context(),
|
||||
):
|
||||
self.init_attention_backend()
|
||||
self.init_cuda_graphs()
|
||||
|
||||
# Some dummy tensors
|
||||
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)
|
||||
|
||||
def init_attention_backend(self):
|
||||
# Create multi-step attn backends and cuda graph runners
|
||||
for step in range(self.speculative_num_steps):
|
||||
draft_backend_factory = DraftBackendFactory(
|
||||
self.server_args,
|
||||
self.mtp_model_runner(step),
|
||||
self.topk,
|
||||
self.speculative_num_steps,
|
||||
)
|
||||
|
||||
# Initialize draft extend attention backend (respects speculative_attention_mode setting)
|
||||
self.draft_extend_attn_backend_list.append(
|
||||
draft_backend_factory.create_draft_extend_backend()
|
||||
)
|
||||
|
||||
def init_cuda_graphs(self):
|
||||
"""Capture cuda graphs."""
|
||||
self.cuda_graph_runner_for_draft_extend_list = []
|
||||
|
||||
if self.server_args.disable_cuda_graph:
|
||||
return
|
||||
|
||||
# Capture extend
|
||||
for step in range(self.speculative_num_steps):
|
||||
if self.draft_extend_attn_backend_list[step] and not _is_npu:
|
||||
tic = time.perf_counter()
|
||||
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
|
||||
logger.info(
|
||||
f"Capture draft extend cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
|
||||
)
|
||||
self.cuda_graph_runner_for_draft_extend_list.append(
|
||||
MultiLayerEagleDraftExtendCudaGraphRunner(self, step)
|
||||
)
|
||||
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
|
||||
logger.info(
|
||||
f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB."
|
||||
)
|
||||
|
||||
def mtp_model_runner(self, layer_id: int) -> ModelRunner:
|
||||
return self.model_runner_list[layer_id]
|
||||
|
||||
def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
|
||||
"""Run speculative decoding forward.
|
||||
|
||||
NOTE: Many states of batch is modified as you go through. It is not guaranteed that
|
||||
the final output batch have the same state as the input.
|
||||
|
||||
Args:
|
||||
batch: The batch to run forward. The state of the batch is modified as it runs.
|
||||
Returns:
|
||||
A tuple of the final logit output of the target model, next tokens accepted,
|
||||
the batch id (used for overlap schedule), and number of accepted tokens.
|
||||
"""
|
||||
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.mtp_model_runner(0).tp_group),
|
||||
speculative_moe_backend_context(),
|
||||
):
|
||||
self.forward_draft_extend(
|
||||
batch, logits_output.hidden_states, next_token_ids, seq_lens_cpu
|
||||
)
|
||||
return GenerationBatchResult(
|
||||
logits_output=logits_output,
|
||||
next_token_ids=next_token_ids,
|
||||
num_correct_drafts=0,
|
||||
can_run_cuda_graph=can_run_cuda_graph,
|
||||
)
|
||||
else:
|
||||
set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
|
||||
|
||||
with (
|
||||
self.draft_tp_context(self.mtp_model_runner(0).tp_group),
|
||||
speculative_moe_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.mtp_model_runner(0).tp_group),
|
||||
speculative_moe_backend_context(),
|
||||
):
|
||||
# NOTE: We should use `check_forward_draft_extend_after_decode`
|
||||
# when DP attention is enabled, but it is slow. Skip it for now.
|
||||
draft_extend_input = verify_output.draft_extend_input
|
||||
if (
|
||||
self.server_args.enable_dp_attention
|
||||
or draft_extend_input.input_ids.shape[0] > 0
|
||||
):
|
||||
# decode is not finished; install draft_extend_input for
|
||||
# the extend forward, then install the next-iter
|
||||
# EagleDraftInput it returns.
|
||||
batch.spec_info = draft_extend_input
|
||||
next_draft_input = self.forward_draft_extend_after_decode(batch)
|
||||
batch.spec_info = next_draft_input
|
||||
else:
|
||||
# All reqs finished and dp_attention isn't forcing extend.
|
||||
# Install an idle EagleDraftInput 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]:
|
||||
"""Run the target extend.
|
||||
|
||||
Args:
|
||||
batch: The batch to run. States could be modified.
|
||||
|
||||
Returns:
|
||||
logits_output: The output of logits. It will contain the full hidden states.
|
||||
next_token_ids: Next token ids generated.
|
||||
seq_lens_cpu: CPU copy of sequence lengths for the draft prefill path.
|
||||
can_run_cuda_graph: Whether the target prefill ran with cuda graph.
|
||||
"""
|
||||
# Forward with the target model and get hidden states.
|
||||
# We need the full hidden states to prefill the KV cache of the draft model.
|
||||
capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.FULL
|
||||
)
|
||||
batch.capture_hidden_mode = capture_mode
|
||||
batch.return_hidden_states_before_norm = True
|
||||
batch_result = self.target_worker.forward_batch_generation(batch)
|
||||
logits_output, next_token_ids = (
|
||||
batch_result.logits_output,
|
||||
batch_result.next_token_ids,
|
||||
)
|
||||
return (
|
||||
logits_output,
|
||||
next_token_ids,
|
||||
batch.seq_lens_cpu,
|
||||
batch_result.can_run_cuda_graph,
|
||||
)
|
||||
|
||||
def _draft_preprocess_decode(self, batch: ScheduleBatch):
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
|
||||
# FIXME: migrate multi-layer eagle worker to eagle worker
|
||||
return EAGLEWorker._draft_preprocess_decode(self, batch)
|
||||
|
||||
def _draft_preprocess_idle(self, batch: ScheduleBatch):
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
|
||||
# FIXME: migrate multi-layer eagle worker to eagle worker
|
||||
return EAGLEWorker._draft_preprocess_idle(self, batch)
|
||||
|
||||
def draft(self, batch: ScheduleBatch):
|
||||
# Parse args
|
||||
if batch.forward_mode.is_idle():
|
||||
self._draft_preprocess_idle(batch)
|
||||
else:
|
||||
self._draft_preprocess_decode(batch)
|
||||
|
||||
spec_info = batch.spec_info
|
||||
assert isinstance(spec_info, EagleDraftInput)
|
||||
|
||||
draft_capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.LAST
|
||||
)
|
||||
spec_info.capture_hidden_mode = draft_capture_mode
|
||||
spec_info.num_tokens_per_req = self.topk
|
||||
spec_info.num_tokens_for_logprob_per_req = self.topk
|
||||
batch.return_hidden_states = False
|
||||
|
||||
# Get forward batch
|
||||
forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
|
||||
assert forward_batch.capture_hidden_mode == draft_capture_mode
|
||||
forward_batch.can_run_dp_cuda_graph = False
|
||||
forward_batch.return_hidden_states_before_norm = True
|
||||
|
||||
# Parse args
|
||||
assert isinstance(spec_info, EagleDraftInput)
|
||||
topk_p, topk_index, hidden_states = (
|
||||
spec_info.topk_p,
|
||||
spec_info.topk_index,
|
||||
spec_info.hidden_states,
|
||||
)
|
||||
|
||||
maybe_detect_nan(topk_p, "draft: NaN in initial topk_p from spec_info")
|
||||
|
||||
# Return values
|
||||
score_list: List[torch.Tensor] = []
|
||||
token_list: List[torch.Tensor] = []
|
||||
parents_list: List[torch.Tensor] = []
|
||||
|
||||
# Forward multiple steps
|
||||
scores = None
|
||||
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
|
||||
0, topk_p, topk_index, hidden_states, scores, self.topk
|
||||
)
|
||||
if self.speculative_num_steps == 1:
|
||||
score_list.append(tree_info[0])
|
||||
token_list.append(tree_info[1])
|
||||
parents_list.append(tree_info[2])
|
||||
else:
|
||||
for i in range(self.speculative_num_steps):
|
||||
score_list.append(tree_info[0][:, :, i].unsqueeze(-1))
|
||||
token_index = tree_info[1][:, i].unsqueeze(-1)
|
||||
token_list.append(token_index)
|
||||
if i == 0:
|
||||
parents_list.append(tree_info[2])
|
||||
else:
|
||||
parents_list.append(
|
||||
torch.full(
|
||||
(tree_info[2].size(0), 1),
|
||||
i,
|
||||
dtype=torch.long,
|
||||
device=self.device,
|
||||
)
|
||||
)
|
||||
|
||||
parent_list, top_scores_index, draft_tokens = organize_draft_results(
|
||||
score_list, token_list, parents_list, self.speculative_num_draft_tokens
|
||||
)
|
||||
|
||||
if batch.forward_mode.is_idle():
|
||||
return EagleVerifyInput.create_idle_input(
|
||||
self.topk,
|
||||
self.speculative_num_steps,
|
||||
self.speculative_num_draft_tokens,
|
||||
)
|
||||
|
||||
(
|
||||
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,
|
||||
)
|
||||
|
||||
target_capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.FULL
|
||||
)
|
||||
return EagleVerifyInput(
|
||||
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.server_args.speculative_num_draft_tokens,
|
||||
capture_hidden_mode=target_capture_mode,
|
||||
seq_lens_sum=forward_batch.seq_lens_sum,
|
||||
seq_lens_cpu=forward_batch.seq_lens_cpu,
|
||||
)
|
||||
|
||||
def clear_cache_pool(self):
|
||||
# allocator and kv cache pool are shared with target worker
|
||||
pass
|
||||
|
||||
def verify(self, batch: ScheduleBatch):
|
||||
spec_info: EagleVerifyInput = batch.spec_info
|
||||
spec_info.prepare_for_verify(batch, self.page_size)
|
||||
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()
|
||||
|
||||
# Forward
|
||||
batch.seq_lens_cpu_cache = spec_info.seq_lens_cpu
|
||||
batch.return_hidden_states_before_norm = True
|
||||
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:
|
||||
# Generate the logit mask for structured output.
|
||||
# Overlap the CPU operations for bitmask generation with the forward pass.
|
||||
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)
|
||||
# NOTE (sk): otherwise, this vocab mask will be the one from the previous extend stage
|
||||
# and will be applied to produce wrong results
|
||||
batch.sampling_info.vocab_mask = None
|
||||
|
||||
maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits")
|
||||
|
||||
spec_info.hidden_states = logits_output.hidden_states
|
||||
res: EagleVerifyOutput = spec_info.verify(
|
||||
batch,
|
||||
logits_output,
|
||||
self.token_to_kv_pool_allocator,
|
||||
self.page_size,
|
||||
vocab_mask,
|
||||
)
|
||||
|
||||
# Post process based on verified outputs.
|
||||
# Pick indices that we care (accepted)
|
||||
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:
|
||||
num_correct_drafts = torch.tensor(
|
||||
res.num_correct_drafts_per_req_cpu,
|
||||
device=logits_output.hidden_states.device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
|
||||
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
|
||||
# res.accept_indices.shape[0] > 0 skips DP attn idle batch
|
||||
if spec_info.topk > 1 and res.accept_indices.shape[0] > 0:
|
||||
# accept_indices=[0,2,3,4,5,7,9,10,11], num_accept_tokens=[4, 3, 2], cumulative_num_accept_tokens=[4, 7, 9]
|
||||
# first_token_indices_per_req=prepend(0, accept_indices[cumulative_num_accept_tokens[:-1]]) = [0, 5, 10]
|
||||
# last_token_indices_per_req=accept_indices[cumulative_num_accept_tokens - 1] = [4, 9, 11] (last token ID of each req)
|
||||
# last_correct_step_indices = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches
|
||||
# equivalent: last_correct_step_indices = last_token_indices_per_req - first_token_indices_per_req;
|
||||
# `accepted_indices_offset` equals `first_token_indices_per_req` because the first accepted slot of each req is its "current token" at logical position i * draft_token_num.
|
||||
cumulative_num_accept_tokens = torch.cumsum(
|
||||
num_correct_drafts + 1, dim=0
|
||||
)
|
||||
accepted_indices_offset = torch.arange(
|
||||
0,
|
||||
len(batch.seq_lens) * self.speculative_num_draft_tokens,
|
||||
step=self.speculative_num_draft_tokens,
|
||||
dtype=num_correct_drafts.dtype,
|
||||
device=num_correct_drafts.device,
|
||||
)
|
||||
last_correct_step_indices = (
|
||||
res.accept_indices[cumulative_num_accept_tokens - 1]
|
||||
- accepted_indices_offset
|
||||
)
|
||||
else:
|
||||
last_correct_step_indices = num_correct_drafts
|
||||
self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify(
|
||||
last_correct_step_indices=last_correct_step_indices,
|
||||
mamba_track_indices=None,
|
||||
mamba_steps_to_track=None,
|
||||
model=self.target_worker.model_runner.model,
|
||||
)
|
||||
|
||||
if batch.return_logprob:
|
||||
add_output_logprobs_for_spec_v1(batch, res, logits_output)
|
||||
|
||||
# Prepare the batch for the next draft forwards.
|
||||
batch.forward_mode = (
|
||||
ForwardMode.DECODE if not batch.forward_mode.is_idle() else ForwardMode.IDLE
|
||||
)
|
||||
|
||||
res.can_run_cuda_graph = can_run_cuda_graph
|
||||
return res
|
||||
|
||||
def forward_draft_extend(
|
||||
self,
|
||||
batch: ScheduleBatch,
|
||||
hidden_states: torch.Tensor,
|
||||
next_token_ids: torch.Tensor,
|
||||
seq_lens_cpu: Optional[torch.Tensor],
|
||||
):
|
||||
"""Run draft model extend. This API modifies the states of the batch.
|
||||
|
||||
Args:
|
||||
batch: The batch to run.
|
||||
hidden_states: Hidden states from the target model forward
|
||||
next_token_ids: Next token ids generated from the target forward.
|
||||
"""
|
||||
batch.spec_info = EagleDraftInput(
|
||||
hidden_states=hidden_states,
|
||||
bonus_tokens=next_token_ids,
|
||||
num_tokens_per_req=1,
|
||||
num_tokens_for_logprob_per_req=1,
|
||||
)
|
||||
batch.return_hidden_states = False
|
||||
apply_eagle_prefill_input_rotation(batch, next_token_ids)
|
||||
capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.LAST
|
||||
)
|
||||
batch.spec_info.capture_hidden_mode = capture_mode
|
||||
batch.seq_lens_cpu_cache = seq_lens_cpu
|
||||
forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
|
||||
forward_batch.return_logprob = False
|
||||
forward_batch.return_hidden_states_before_norm = True
|
||||
topk_p_list = []
|
||||
topk_index_list = []
|
||||
for step in range(self.speculative_num_steps):
|
||||
logits_output = (
|
||||
self.mtp_model_runner(step).forward(forward_batch).logits_output
|
||||
)
|
||||
maybe_detect_nan(
|
||||
logits_output.next_token_logits,
|
||||
f"draft_extend_for_prefill step {step}",
|
||||
)
|
||||
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
|
||||
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
|
||||
topk_p_list.append(topk_p)
|
||||
topk_index_list.append(topk_index)
|
||||
pt = 0
|
||||
if forward_batch.extend_seq_lens is not None:
|
||||
for i, extend_len in enumerate(forward_batch.extend_seq_lens):
|
||||
input_ids = forward_batch.input_ids[pt : pt + extend_len]
|
||||
forward_batch.input_ids[pt : pt + extend_len] = torch.cat(
|
||||
(input_ids[1:], topk_index[i].reshape(1))
|
||||
)
|
||||
pt += extend_len
|
||||
|
||||
assert isinstance(forward_batch.spec_info, EagleDraftInput)
|
||||
assert forward_batch.spec_info is batch.spec_info
|
||||
forward_batch.spec_info.topk_p = torch.cat(topk_p_list, dim=1)
|
||||
forward_batch.spec_info.topk_index = torch.cat(topk_index_list, dim=1)
|
||||
|
||||
def forward_draft_extend_after_decode(
|
||||
self, batch: ScheduleBatch
|
||||
) -> EagleDraftInput:
|
||||
draft_extend_input: EagleDraftExtendInput = batch.spec_info
|
||||
|
||||
# Backup fields that will be modified in-place
|
||||
seq_lens_backup = batch.seq_lens.clone()
|
||||
seq_lens_cpu_backup = batch.seq_lens_cpu.clone()
|
||||
req_pool_indices_backup = batch.req_pool_indices
|
||||
return_logprob_backup = batch.return_logprob
|
||||
|
||||
input_is_idle = batch.forward_mode.is_idle()
|
||||
|
||||
draft_extend_capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.LAST
|
||||
)
|
||||
if draft_extend_input.input_ids.shape[0] == 0:
|
||||
# Single source for hidden_size via hidden_size_for(self) (incl.
|
||||
# EAGLE-3 aux widening). Two stub origins from verify(): fully-idle
|
||||
# batch and active batch with all reqs finished.
|
||||
batch = batch.copy()
|
||||
batch.prepare_for_idle()
|
||||
draft_extend_input = EagleDraftExtendInput.create_idle_input(
|
||||
device=self.device,
|
||||
hidden_size=EagleDraftExtendInput.hidden_size_for(self),
|
||||
dtype=EagleDraftExtendInput.dtype_for(self),
|
||||
capture_hidden_mode=draft_extend_capture_mode,
|
||||
)
|
||||
batch.spec_info = draft_extend_input
|
||||
|
||||
# Phase 1: prepare extend (kernel writes draft_extend_input.{positions, bonus_tokens})
|
||||
draft_extend_input.num_tokens_per_req = self.speculative_num_steps + 1
|
||||
draft_extend_input.num_tokens_for_logprob_per_req = 1
|
||||
draft_extend_input.prepare_extend_after_decode(
|
||||
batch,
|
||||
speculative_num_steps=self.speculative_num_steps,
|
||||
)
|
||||
batch.forward_mode = (
|
||||
ForwardMode.DRAFT_EXTEND
|
||||
if not batch.forward_mode.is_idle()
|
||||
else ForwardMode.IDLE
|
||||
)
|
||||
|
||||
batch.return_hidden_states = False
|
||||
draft_extend_input.capture_hidden_mode = draft_extend_capture_mode
|
||||
forward_batch = ForwardBatch.init_new(batch, self.mtp_model_runner(0))
|
||||
assert forward_batch.capture_hidden_mode == draft_extend_capture_mode
|
||||
forward_batch.return_hidden_states_before_norm = True
|
||||
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 = batch.seq_lens.sum().item()
|
||||
topk_p_list = []
|
||||
topk_index_list = []
|
||||
# Run
|
||||
for step in range(self.speculative_num_steps):
|
||||
can_cuda_graph = len(
|
||||
self.cuda_graph_runner_for_draft_extend_list
|
||||
) and self.cuda_graph_runner_for_draft_extend_list[step].can_run(
|
||||
forward_batch
|
||||
)
|
||||
if can_cuda_graph:
|
||||
logits_output = self.cuda_graph_runner_for_draft_extend_list[
|
||||
step
|
||||
].replay(forward_batch)
|
||||
else:
|
||||
forward_batch.can_run_dp_cuda_graph = False
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
self.mtp_model_runner(step).attn_backend.init_forward_metadata(
|
||||
forward_batch
|
||||
)
|
||||
# Planned pre-pad; do NOT opt into post-pad re-plan — a
|
||||
# DP-padded re-plan breaks DSA's indexer schedule_meta
|
||||
# (see #27091). Use the marked pre-pad metadata as-is.
|
||||
forward_batch.mark_forward_metadata_ready()
|
||||
logits_output = (
|
||||
self.mtp_model_runner(step).forward(forward_batch).logits_output
|
||||
)
|
||||
|
||||
maybe_detect_nan(
|
||||
logits_output.next_token_logits,
|
||||
f"draft_extend_after_decode step {step} (cuda_graph={can_cuda_graph})",
|
||||
)
|
||||
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
|
||||
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
|
||||
topk_p_list.append(topk_p)
|
||||
topk_index_list.append(topk_index)
|
||||
pt = 0
|
||||
if forward_batch.extend_seq_lens is not None:
|
||||
for i, extend_len in enumerate(forward_batch.extend_seq_lens):
|
||||
input_ids = forward_batch.input_ids[pt : pt + extend_len]
|
||||
forward_batch.input_ids[pt : pt + extend_len] = torch.cat(
|
||||
(input_ids[1:], topk_index[i].reshape(1))
|
||||
)
|
||||
pt += extend_len
|
||||
|
||||
# Phase 3: assemble next-iter EagleDraftInput from extend output
|
||||
next_decode_capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if self.speculative_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.LAST
|
||||
)
|
||||
next_draft_input = EagleDraftInput(
|
||||
bonus_tokens=draft_extend_input.bonus_tokens,
|
||||
hidden_states=logits_output.hidden_states,
|
||||
topk_p=torch.cat(topk_p_list, dim=1),
|
||||
topk_index=torch.cat(topk_index_list, dim=1),
|
||||
capture_hidden_mode=next_decode_capture_mode,
|
||||
)
|
||||
|
||||
# Restore batch fields. `seq_lens` etc. were modified by
|
||||
# `prepare_extend_after_decode`. Caller installs `next_draft_input` as
|
||||
# `batch.spec_info`.
|
||||
batch.forward_mode = (
|
||||
ForwardMode.DECODE if not input_is_idle else ForwardMode.IDLE
|
||||
)
|
||||
batch.seq_lens = seq_lens_backup
|
||||
batch.seq_lens_cpu = seq_lens_cpu_backup
|
||||
batch.req_pool_indices = req_pool_indices_backup
|
||||
batch.return_logprob = return_logprob_backup
|
||||
return next_draft_input
|
||||
@@ -189,41 +189,25 @@ class SpeculativeAlgorithm(Enum):
|
||||
|
||||
return FrozenKVMTPWorker
|
||||
|
||||
# EAGLE / EAGLE3 / STANDALONE / MULTI_LAYER always use the V2 worker,
|
||||
# even with overlap disabled (scheduler drives it synchronously).
|
||||
if self.is_eagle() and server_args.enable_multi_layer_eagle:
|
||||
# FIXME: migrate to EagleWorker
|
||||
if enable_overlap:
|
||||
from sglang.srt.speculative.multi_layer_eagle_worker_v2 import (
|
||||
MultiLayerEagleWorkerV2,
|
||||
)
|
||||
|
||||
return MultiLayerEagleWorkerV2
|
||||
|
||||
from sglang.srt.speculative.multi_layer_eagle_worker import (
|
||||
MultiLayerEagleWorker,
|
||||
from sglang.srt.speculative.multi_layer_eagle_worker_v2 import (
|
||||
MultiLayerEagleWorkerV2,
|
||||
)
|
||||
|
||||
return MultiLayerEagleWorker
|
||||
return MultiLayerEagleWorkerV2
|
||||
|
||||
elif self.is_eagle():
|
||||
if enable_overlap:
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
|
||||
|
||||
return EAGLEWorkerV2
|
||||
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
|
||||
return EAGLEWorker
|
||||
return EAGLEWorkerV2
|
||||
elif self.is_standalone():
|
||||
if enable_overlap:
|
||||
from sglang.srt.speculative.standalone_worker_v2 import (
|
||||
StandaloneWorkerV2,
|
||||
)
|
||||
from sglang.srt.speculative.standalone_worker_v2 import (
|
||||
StandaloneWorkerV2,
|
||||
)
|
||||
|
||||
return StandaloneWorkerV2
|
||||
|
||||
from sglang.srt.speculative.standalone_worker import StandaloneWorker
|
||||
|
||||
return StandaloneWorker
|
||||
return StandaloneWorkerV2
|
||||
elif self.is_ngram():
|
||||
if enable_overlap:
|
||||
raise ValueError(
|
||||
|
||||
@@ -1,121 +0,0 @@
|
||||
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.tp_worker import TpModelWorker
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.speculative.adaptive_runtime_state import (
|
||||
AdaptiveController,
|
||||
)
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.speculative.spec_utils import draft_tp_context, load_token_map
|
||||
from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
|
||||
|
||||
if is_cuda():
|
||||
from sgl_kernel import segment_packbits # noqa: F401
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
|
||||
|
||||
|
||||
class StandaloneWorker(EAGLEWorker):
|
||||
|
||||
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,
|
||||
):
|
||||
# Parse arguments
|
||||
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
|
||||
)
|
||||
|
||||
# TODO: Adaptive speculative
|
||||
self.adaptive_controller: Optional[AdaptiveController] = None
|
||||
|
||||
# Override the context length of the draft model to be the same as the target model.
|
||||
server_args.context_length = target_worker.model_runner.model_config.context_len
|
||||
|
||||
# Do not capture cuda graph in `super().__init__()`
|
||||
# It will be captured later.
|
||||
backup_disable_cuda_graph = server_args.disable_cuda_graph
|
||||
server_args.disable_cuda_graph = True
|
||||
# Share the allocator with a target worker.
|
||||
# Draft and target worker own their own KV cache pools.
|
||||
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
|
||||
target_worker.get_memory_pool()
|
||||
)
|
||||
|
||||
# Load hot token ids
|
||||
if server_args.speculative_token_map is not None:
|
||||
self.hot_token_id = load_token_map(server_args.speculative_token_map)
|
||||
server_args.json_model_override_args = (
|
||||
f'{{"hot_vocab_size": {len(self.hot_token_id)}}}'
|
||||
)
|
||||
else:
|
||||
self.hot_token_id = None
|
||||
|
||||
# Init draft worker
|
||||
with (
|
||||
empty_context(),
|
||||
speculative_moe_backend_context(),
|
||||
speculative_moe_a2a_backend_context(),
|
||||
):
|
||||
TpModelWorker.__init__(
|
||||
self,
|
||||
server_args=server_args,
|
||||
gpu_id=gpu_id,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=0, # spec workers don't support pipeline parallelism
|
||||
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=target_worker.model_runner.memory_pool_config,
|
||||
)
|
||||
|
||||
# Init attention backend and cuda graphs
|
||||
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
|
||||
)
|
||||
with (
|
||||
self.draft_tp_context(self.draft_model_runner.tp_group),
|
||||
speculative_moe_backend_context(),
|
||||
speculative_moe_a2a_backend_context(),
|
||||
):
|
||||
self.init_attention_backend()
|
||||
self.init_cuda_graphs()
|
||||
|
||||
# Some dummy tensors
|
||||
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)
|
||||
+18
-2
@@ -21,7 +21,7 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
|
||||
EAGLEDraftCudaGraphRunner,
|
||||
)
|
||||
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
||||
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
||||
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
|
||||
from sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner import (
|
||||
FrozenKVMTPCudaGraphRunner,
|
||||
)
|
||||
@@ -175,12 +175,28 @@ class _EagleDraftWorkerHarness:
|
||||
self.speculative_algorithm = SpeculativeAlgorithm.EAGLE
|
||||
self.hot_token_id = None
|
||||
self.model_runner.forward = model_forward
|
||||
self.draft_forward = MethodType(EAGLEWorker.draft_forward, self)
|
||||
self.draft_forward = MethodType(EagleDraftWorker.draft_forward, self)
|
||||
# draft_forward's topk=1 fast path reads these prealloc buffers (built
|
||||
# in EagleDraftWorker.__init__, which the harness skips), so build them
|
||||
# here. _rebuild_topk1_chain_buffers asserts num_draft_tokens ==
|
||||
# num_steps + 1; the fast path never reads num_draft_tokens, so pin it.
|
||||
self.device = self.model_runner.device
|
||||
if self.topk == 1:
|
||||
self.speculative_num_draft_tokens = self.speculative_num_steps + 1
|
||||
self._topk1_parents_prealloc = None
|
||||
self._topk1_score_indices_prealloc = None
|
||||
EagleDraftWorker._rebuild_topk1_chain_buffers(self)
|
||||
|
||||
@property
|
||||
def draft_model_runner(self):
|
||||
return self.model_runner
|
||||
|
||||
@property
|
||||
def draft_runner(self):
|
||||
# V2 draft_forward reads self.draft_runner (forward / model_config /
|
||||
# canary_manager); for the harness that's the fixture runner.
|
||||
return self.model_runner
|
||||
|
||||
|
||||
class _FrozenKVMTPWorkerHarness:
|
||||
def __init__(
|
||||
|
||||
@@ -24,13 +24,14 @@ from sglang.test.server_fixtures.streaming_session_fixture import (
|
||||
class StreamingSessionKitMixin:
|
||||
"""Streaming-session KV-inheritance + retract/abort-recovery suite."""
|
||||
|
||||
# -1 for non-overlap subclasses: the last sampled token isn't committed
|
||||
# before max_new stops, so slot.kv_committed_len = input + output - 1.
|
||||
kv_inherit_offset = 0
|
||||
# Allowed inherited-cache offsets vs the previous turn's total. Non-overlap
|
||||
# spec decode can be off by 1: the bonus token's KV is only computed by the
|
||||
# next forward, which sync skips at finish (overlap drains it, so it's 0).
|
||||
kv_inherit_offsets = (0,)
|
||||
|
||||
def test_kv_cache_inheritance(self, gen_len=12):
|
||||
"""Each turn's cached_tokens must equal previous turn's prompt+completion
|
||||
(modulo kv_inherit_offset)."""
|
||||
(modulo kv_inherit_offsets)."""
|
||||
chunks = [
|
||||
"Let me tell you something about France.",
|
||||
"The capital of France is",
|
||||
@@ -75,11 +76,11 @@ class StreamingSessionKitMixin:
|
||||
else:
|
||||
# Turns 2+: cached_tokens reflects KV inherited from previous turn
|
||||
# (via inherit_kv_states, not radix tree matching).
|
||||
expected = prev_kv_len + self.kv_inherit_offset
|
||||
self.assertEqual(
|
||||
allowed = {prev_kv_len + off for off in self.kv_inherit_offsets}
|
||||
self.assertIn(
|
||||
cached,
|
||||
expected,
|
||||
f"Turn {turn_idx + 1}: inherited {cached} != expected {expected}",
|
||||
allowed,
|
||||
f"Turn {turn_idx + 1}: inherited {cached} not in {sorted(allowed)}",
|
||||
)
|
||||
prev_kv_len = prompt_tokens + completion_tokens
|
||||
|
||||
|
||||
Reference in New Issue
Block a user