diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index b73f3d7d0..055153f89 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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: diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 31c59c6f0..3829a84ee 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -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: | diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py index 5415f4160..22ef7e237 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py @@ -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): diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 16bc1bf37..044fc538e 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py @@ -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 diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 23b920a77..bd4c82e71 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c51a27f2f..db7ad15ff 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ea14ee492..8e32640fc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index c546a5650..17ab61793 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -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 diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 6cdcde100..09531eb96 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -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 diff --git a/python/sglang/srt/speculative/adaptive_spec_params.py b/python/sglang/srt/speculative/adaptive_spec_params.py index e714fd6bd..e701357f6 100644 --- a/python/sglang/srt/speculative/adaptive_spec_params.py +++ b/python/sglang/srt/speculative/adaptive_spec_params.py @@ -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 ( diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 95f3a5e5c..96c7286af 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index dc5f2505e..8143e60f0 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index b59e7cdc6..ef338410c 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py deleted file mode 100644 index 856c2839f..000000000 --- a/python/sglang/srt/speculative/eagle_worker.py +++ /dev/null @@ -1,1356 +0,0 @@ -import contextlib -import logging -import time -from contextlib import contextmanager -from typing import List, Optional, Tuple - -import torch - -from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import ( - EAGLEDraftNpuGraphRunner, -) -from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner -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_a2a_backend_context, - speculative_moe_backend_context, -) -from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1 -from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput -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.mem_cache.common import ( - alloc_paged_token_slots_extend, - alloc_token_slots, - get_last_loc, -) -from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner -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.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.adaptive_runtime_state import ( - AdaptiveController, - SpecRuntimeState, -) -from sglang.srt.speculative.draft_utils import DraftBackendFactory -from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( - EAGLEDraftCudaGraphRunner, -) -from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( - EAGLEDraftExtendCudaGraphRunner, -) -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.spec_info import SpeculativeAlgorithm -from sglang.srt.speculative.spec_utils import ( - assign_draft_cache_locs, - draft_tp_context, - fast_topk, - generate_token_bitmask, - get_last_loc_large_page_size_large_top_k, - load_token_map, - select_top_k_tokens, -) -from sglang.srt.utils import ( - MultiprocessingSerializer, - empty_context, - get_available_gpu_memory, - is_cuda, - is_musa, - is_npu, - log_info_on_rank0, - next_power_of_2, -) -from sglang.srt.utils.async_probe import ( - maybe_detect_inf, - maybe_detect_nan, - maybe_detect_oob, -) -from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions - -_is_npu = is_npu() -_is_musa = is_musa() - -if is_cuda(): - from sgl_kernel import segment_packbits # noqa: F401 - -logger = logging.getLogger(__name__) - - -class EAGLEWorker(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 - 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 - ) - - # Adaptive speculative - self.adaptive_controller: Optional[AdaptiveController] = None - if server_args.speculative_adaptive: - self.adaptive_controller = AdaptiveController( - self, - config_path=server_args.speculative_adaptive_config, - ) - - # 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(), speculative_moe_a2a_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, - ) - - 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 - self.draft_model_runner.model.set_embed_and_head(embed, head) - - # 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 - ) - self.eagle_use_aux_hidden_state = False - if self.speculative_algorithm.is_eagle3(): - self.eagle_use_aux_hidden_state = True - eagle_config = getattr( - self.draft_model_runner.model_config.hf_config, "eagle_config", {} - ) - self.eagle_use_aux_hidden_state = eagle_config.get( - "use_aux_hidden_state", True - ) - 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() - if self.adaptive_controller is not None: - self.adaptive_controller.register( - SpecRuntimeState( - speculative_num_steps=self.speculative_num_steps, - speculative_num_draft_tokens=self.speculative_num_draft_tokens, - draft_attn_backend=self.draft_attn_backend, - cuda_graph_runner=self.cuda_graph_runner, - target_attn_backend=self.target_worker.model_runner.attn_backend, - target_graph_runner=self.target_worker.model_runner.graph_runner, - draft_extend_attn_backend=self.draft_extend_attn_backend, - cuda_graph_runner_for_draft_extend=self.cuda_graph_runner_for_draft_extend, - ) - ) - self.adaptive_controller.init_states( - cuda_graph_bs=( - None - if self.server_args.disable_cuda_graph - else self.server_args.cuda_graph_bs - ), - ) - - # 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 - draft_backend_factory = DraftBackendFactory( - self.server_args, - self.draft_model_runner, - self.topk, - self.speculative_num_steps, - ) - - # Initialize decode attention backend - self.draft_attn_backend = draft_backend_factory.create_decode_backend() - - # Initialize draft extend attention backend (respects speculative_attention_mode setting) - self.draft_extend_attn_backend = ( - draft_backend_factory.create_draft_extend_backend() - ) - - self.draft_model_runner.draft_attn_backend = self.draft_attn_backend - - def init_cuda_graphs(self): - """Capture cuda graphs.""" - self.cuda_graph_runner = None - self.cuda_graph_runner_for_draft_extend = None - - if self.server_args.disable_cuda_graph: - return - - Device2DraftCudaGraphRunner = { - "npu": EAGLEDraftNpuGraphRunner, - "cuda": EAGLEDraftCudaGraphRunner, - "musa": EAGLEDraftCudaGraphRunner, - } - # Capture draft - if self.speculative_num_steps > 1: - tic = time.perf_counter() - before_mem = get_available_gpu_memory(self.device, self.gpu_id) - log_info_on_rank0( - logger, - f"Capture draft cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB", - ) - self.cuda_graph_runner = Device2DraftCudaGraphRunner[ - self.target_worker.device - ](self) - after_mem = get_available_gpu_memory(self.device, self.gpu_id) - log_info_on_rank0( - logger, - f"Capture draft 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.", - ) - - # Capture extend - if self.draft_extend_attn_backend and not _is_npu: - tic = time.perf_counter() - before_mem = get_available_gpu_memory(self.device, self.gpu_id) - log_info_on_rank0( - logger, - 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 = EAGLEDraftExtendCudaGraphRunner( - self - ) - after_mem = get_available_gpu_memory(self.device, self.gpu_id) - log_info_on_rank0( - logger, - 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 apply_runtime_state(self, state: SpecRuntimeState): - """Apply a pre-built runtime state to this worker.""" - if self.speculative_num_steps == state.speculative_num_steps: - return - - log_info_on_rank0( - logger, - "Switch adaptive runtime state: " - f"steps {self.speculative_num_steps} -> {state.speculative_num_steps}, " - f"draft_tokens {self.speculative_num_draft_tokens} -> " - f"{state.speculative_num_draft_tokens}", - ) - - self.speculative_num_steps = state.speculative_num_steps - self.speculative_num_draft_tokens = state.speculative_num_draft_tokens - # Draft stage - self.draft_attn_backend = state.draft_attn_backend - self.draft_model_runner.draft_attn_backend = state.draft_attn_backend - self.cuda_graph_runner = state.cuda_graph_runner - # Verify stage - self.target_worker.model_runner.attn_backend = state.target_attn_backend - self.target_worker.model_runner.graph_runner = state.target_graph_runner - # Extend stage - self.draft_extend_attn_backend = state.draft_extend_attn_backend - self.cuda_graph_runner_for_draft_extend = ( - state.cuda_graph_runner_for_draft_extend - ) - # Sync server_args - self.server_args.speculative_num_steps = state.speculative_num_steps - self.server_args.speculative_num_draft_tokens = ( - state.speculative_num_draft_tokens - ) - - def build_adaptive_runtime_state( - self, - speculative_num_steps: int, - speculative_num_draft_tokens: int, - cuda_graph_bs: list[int] | None = None, - ) -> SpecRuntimeState: - """Build a SpecRuntimeState for the given step configuration.""" - tic = time.perf_counter() - before_mem = get_available_gpu_memory(self.device, self.gpu_id) - - with self._override_worker_state( - speculative_num_steps, - speculative_num_draft_tokens, - cuda_graph_bs=cuda_graph_bs, - ): - # Reuse existing init methods for draft attention backend and cuda graphs - self.init_attention_backend() - self.init_cuda_graphs() - - # Capture target attention backend and CUDA graph - target_model_runner = self.target_worker.model_runner - backup_init = target_model_runner.init_new_workspace - try: - target_attn_backend = target_model_runner._get_attention_backend( - init_new_workspace=True - ) - finally: - target_model_runner.init_new_workspace = backup_init - - target_graph_runner = None - if not self.server_args.disable_cuda_graph: - TargetGraphRunnerCls = NPUGraphRunner if _is_npu else CudaGraphRunner - target_graph_runner = TargetGraphRunnerCls( - target_model_runner, - attn_backend=target_attn_backend, - speculative_num_steps=speculative_num_steps, - speculative_num_draft_tokens=speculative_num_draft_tokens, - ) - - state = SpecRuntimeState( - speculative_num_steps=speculative_num_steps, - speculative_num_draft_tokens=speculative_num_draft_tokens, - # Draft stage - draft_attn_backend=self.draft_attn_backend, - cuda_graph_runner=self.cuda_graph_runner, - # Verify stage - target_attn_backend=target_attn_backend, - target_graph_runner=target_graph_runner, - # Extend stage - draft_extend_attn_backend=self.draft_extend_attn_backend, - cuda_graph_runner_for_draft_extend=self.cuda_graph_runner_for_draft_extend, - ) - - after_mem = get_available_gpu_memory(self.device, self.gpu_id) - log_info_on_rank0( - logger, - f"Built adaptive runtime state steps={speculative_num_steps}: " - f"elapsed={time.perf_counter() - tic:.2f}s, " - f"mem={(before_mem - after_mem):.2f}GB", - ) - - return state - - @contextmanager - def _override_worker_state( - self, - speculative_num_steps: int, - speculative_num_draft_tokens: int, - cuda_graph_bs: list[int] | None = None, - ): - """Temporarily override server_args and worker attributes for graph capture.""" - sa = self.server_args - backup = ( - self.speculative_num_steps, - self.speculative_num_draft_tokens, - self.draft_attn_backend, - self.draft_extend_attn_backend, - self.draft_model_runner.draft_attn_backend, - self.cuda_graph_runner, - self.cuda_graph_runner_for_draft_extend, - sa.speculative_num_steps, - sa.speculative_num_draft_tokens, - sa.cuda_graph_bs, - sa.disable_cuda_graph, - ) - self.speculative_num_steps = speculative_num_steps - self.speculative_num_draft_tokens = speculative_num_draft_tokens - sa.speculative_num_steps = speculative_num_steps - sa.speculative_num_draft_tokens = speculative_num_draft_tokens - if cuda_graph_bs is not None: - sa.cuda_graph_bs = cuda_graph_bs - # BS-aware adaptive spec may prune cuda_graph_bs to an empty list - # for steps that no BS range uses (e.g. step=1). Disable graph - # capture for those steps; restore in finally so subsequent steps - # are not affected. - if not cuda_graph_bs: - sa.disable_cuda_graph = True - try: - yield - finally: - ( - self.speculative_num_steps, - self.speculative_num_draft_tokens, - self.draft_attn_backend, - self.draft_extend_attn_backend, - self.draft_model_runner.draft_attn_backend, - self.cuda_graph_runner, - self.cuda_graph_runner_for_draft_extend, - sa.speculative_num_steps, - sa.speculative_num_draft_tokens, - sa.cuda_graph_bs, - sa.disable_cuda_graph, - ) = backup - - @property - def draft_model_runner(self): - return self.model_runner - - 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.draft_model_runner.tp_group), - speculative_moe_backend_context(), - speculative_moe_a2a_backend_context(), - ): - self.forward_draft_extend( - batch, - logits_output.hidden_states, - next_token_ids, - seq_lens_cpu, - logits_output.mm_input_embeds, - ) - return GenerationBatchResult( - logits_output=logits_output, - next_token_ids=next_token_ids, - num_correct_drafts=0, - can_run_cuda_graph=can_run_cuda_graph, - ) - else: - if self.adaptive_controller is not None: - self.adaptive_controller.activate_step_by_batch(batch.batch_size()) - - set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True) - - with ( - self.draft_tp_context(self.draft_model_runner.tp_group), - speculative_moe_backend_context(), - speculative_moe_a2a_backend_context(), - ): - verify_input = self.draft(batch) - - set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) - set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) - - # Install verify_input as `batch.spec_info` for the verify forward. - batch.spec_info = verify_input - verify_output = self.verify(batch) - - if get_global_tracing_enabled(): - for idx, req in enumerate(batch.reqs): - num_correct_drafts = verify_output.num_correct_drafts_per_req_cpu[ - idx - ] - req.time_stats.set_spec_verify_end_time( - num_correct_drafts=num_correct_drafts - ) - - set_time_batch( - batch.reqs, "set_spec_draft_extend_start_time", trace_only=True - ) - - with ( - self.draft_tp_context(self.draft_model_runner.tp_group), - speculative_moe_backend_context(), - speculative_moe_a2a_backend_context(), - ): - # 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 - ) - - if self.adaptive_controller is not None: - self.adaptive_controller.on_verify_complete( - verify_output.num_correct_drafts_per_req_cpu, - batch_size=batch.batch_size(), - ) - - 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_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): - batch.maybe_evict_swa() - for req in batch.reqs: - req.decode_batch_idx += 1 - - # Parse args - num_seqs = batch.batch_size() - spec_info = batch.spec_info - - # Accumulate penalty - if batch.sampling_info.penalizer_orchestrator.is_required: - # This is a relaxed version of penalties for speculative decoding. - batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - spec_info.bonus_tokens.to(torch.int64) - ) - - # Allocate cache locations - # Layout of the out_cache_loc - # [ topk 0 ] [ topk 1 ] - # [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2] - if self.page_size == 1: - alloc_len_per_decode = self.speculative_num_steps * self.topk - # TODO: We only need self.speculative_num_steps - 1 * topk cache loc - out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots( - batch.tree_cache, - num_seqs * alloc_len_per_decode, - backup_state=True, - ) - else: - if self.topk == 1: - prefix_lens, seq_lens, last_loc = get_last_loc_large_page_size_top_k_1( - batch.req_to_token_pool.req_to_token, - batch.req_pool_indices, - batch.seq_lens, - self.speculative_num_steps, - ) - prefix_lens_cpu = batch.seq_lens_cpu - seq_lens_cpu = batch.seq_lens_cpu + self.speculative_num_steps - extend_num_tokens = num_seqs * self.speculative_num_steps - else: - # In this case, the last partial page needs to be duplicated. - # KV cache layout in batch.req_to_token_pool.req_to_token: - # - # | -------- | -- xxxx .. | -- xxxx .. | -- xxxx .. | - # prefix top-k = 0 tok-k = 1 top-k = 2 - # - # "-" means prefix tokens - # "x" means speculative draft tokens - # "." means padded tokens - - ( - prefix_lens, - seq_lens, - last_loc, - self.num_new_pages_per_topk, - self.extend_lens, - last_page_lens, - ) = get_last_loc_large_page_size_large_top_k( - batch.req_to_token_pool.req_to_token, - batch.req_pool_indices, - batch.seq_lens, - self.speculative_num_steps, - self.topk, - self.page_size, - ) - prefix_lens_cpu = batch.seq_lens_cpu - last_page_lens_cpu = prefix_lens_cpu % self.page_size - num_new_pages_per_topk = ( - last_page_lens_cpu + self.speculative_num_steps + self.page_size - 1 - ) // self.page_size - seq_lens_cpu = ( - prefix_lens_cpu // self.page_size * self.page_size - + num_new_pages_per_topk * (self.page_size * self.topk) - ) - extend_num_tokens = torch.sum((seq_lens_cpu - prefix_lens_cpu)).item() - - out_cache_loc, token_to_kv_pool_state_backup = ( - alloc_paged_token_slots_extend( - batch.tree_cache, - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - backup_state=True, - ) - ) - - if self.page_size > 1 and self.topk > 1: - last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0) - duplicate_cache_len = torch.sum(last_page_lens_cpu).item() * (self.topk - 1) - target_cache_loc = torch.zeros( - duplicate_cache_len, dtype=torch.int32, device=self.device - ) - source_cache_loc = torch.zeros( - duplicate_cache_len, dtype=torch.int32, device=self.device - ) - else: - # When source_cache_loc is not needed, simply skip - duplicate_cache_len = 0 - source_cache_loc, target_cache_loc, last_page_lens_cumsum = None, None, None - - assign_draft_cache_locs[(num_seqs,)]( - batch.req_pool_indices, - batch.req_to_token_pool.req_to_token, - batch.seq_lens, - self.extend_lens, - self.num_new_pages_per_topk, - out_cache_loc, - source_cache_loc, - target_cache_loc, - last_page_lens_cumsum, - duplicate_cache_len, - batch.req_to_token_pool.req_to_token.shape[1], - self.topk, - self.speculative_num_steps, - self.page_size, - next_power_of_2(num_seqs), - next_power_of_2(self.speculative_num_steps + self.page_size), - ) - - if self.page_size > 1 and self.topk > 1: - if duplicate_cache_len > 0: - self.draft_model_runner.token_to_kv_pool.move_kv_cache( - target_cache_loc, source_cache_loc - ) - # Remove padded slots - # TODO: We only need self.speculative_num_steps - 1 cache loc - out_cache_loc = out_cache_loc[ - : num_seqs * self.topk * self.speculative_num_steps - ] - - batch.out_cache_loc = out_cache_loc - batch.seq_lens_sum = torch.sum(batch.seq_lens).item() - batch.return_hidden_states = False - spec_info.positions = batch.seq_lens.repeat_interleave(self.topk, dim=0) - self.token_to_kv_pool_allocator.restore_state(token_to_kv_pool_state_backup) - - def _draft_preprocess_idle(self, batch: ScheduleBatch): - capture_mode = ( - CaptureHiddenMode.NULL - if self.speculative_algorithm.is_standalone() - else CaptureHiddenMode.LAST - ) - batch.spec_info = EagleDraftInput.create_idle_input( - device=self.device, - hidden_size=EagleDraftInput.hidden_size_for(self), - dtype=EagleDraftInput.dtype_for(self), - topk=self.topk, - capture_hidden_mode=capture_mode, - ) - - 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.draft_model_runner) - assert forward_batch.capture_hidden_mode == draft_capture_mode - can_cuda_graph = self.cuda_graph_runner and self.cuda_graph_runner.can_run( - forward_batch - ) - if can_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 - if ( - not forward_batch.forward_mode.is_idle() - and self.speculative_num_steps > 1 - ): - # Skip attention backend init for idle mode or 1-step draft - self.draft_attn_backend.init_forward_metadata(forward_batch) - forward_batch.mark_forward_metadata_ready() - # Run forward steps - parent_list, top_scores_index, draft_tokens = self.draft_forward( - forward_batch - ) - - 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.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 draft_forward(self, forward_batch: ForwardBatch): - # Parse args - spec_info = forward_batch.spec_info - assert isinstance(spec_info, EagleDraftInput) - out_cache_loc = forward_batch.out_cache_loc - topk_p, topk_index, hidden_states = ( - spec_info.topk_p, - spec_info.topk_index, - spec_info.hidden_states, - ) - - maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") - - if self.hot_token_id is not None: - topk_index = self.hot_token_id[topk_index] - # TODO: We only need self.speculative_num_steps - 1 cache loc - out_cache_loc = out_cache_loc.reshape( - forward_batch.batch_size, self.topk, self.speculative_num_steps - ) - out_cache_loc = out_cache_loc.permute((2, 0, 1)).reshape( - self.speculative_num_steps, -1 - ) - - # Return values - score_list: List[torch.Tensor] = [] - token_list: List[torch.Tensor] = [] - parents_list: List[torch.Tensor] = [] - - # Forward multiple steps - scores = None - # Reuse NSA/DSA topk_indices from the first draft forward step for - # subsequent steps, analogous to skip_topk in deepseek_v2.py layers. - # Only safe with topk == 1: select_top_k_tokens reorders candidate rows - # each step, which would desync the cached indices from their rows. - index_share_for_mtp_iteration = ( - getattr(self.model_config.hf_config, "index_share_for_mtp_iteration", False) - and self.topk == 1 - ) - if index_share_for_mtp_iteration: - forward_batch.reuse_mtp_topk_indices = True - forward_batch.topk_indices = 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]) - - # We don't need to run the last forward. we get 1 token from draft prefill and (#spec steps - 1) tokens here - if i == self.speculative_num_steps - 1: - break - - # Set inputs - forward_batch.input_ids = input_ids - # Some draft model RoPE kernels need cache_loc to be contiguous. - if ( - self.server_args.speculative_algorithm == "STANDALONE" - and self.model_config.hf_config.architectures[0] == "GptOssForCausalLM" - ) or self.model_config.hf_config.architectures[0] == ( - "Qwen3MoeForCausalLMMTP" - ): - out_cache_loc = out_cache_loc.contiguous() - forward_batch.out_cache_loc = out_cache_loc[i] - spec_info.hidden_states = hidden_states - - # Run forward under a per-step ForwardContext so the model layer - # reads attn_backends[i] for the i-th draft step. ``_forward_raw`` - # is no-op for the attn_backend half when a context is already - # active, so this outer wrap is what reaches RadixAttention. - with forward_context( - ForwardContext(attn_backend=self.draft_attn_backend.attn_backends[i]) - ): - logits_output = self.draft_model_runner.forward( - forward_batch - ).logits_output - maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") - maybe_detect_inf(logits_output.next_token_logits, f"draft_forward 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], - f"draft_forward step {i}: topk_index OOB vs vocab_size={logits_output.next_token_logits.shape[-1]}", - ) - if self.hot_token_id is not None: - topk_index = self.hot_token_id[topk_index] - 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) - - if index_share_for_mtp_iteration: - forward_batch.topk_indices = None - forward_batch.reuse_mtp_topk_indices = False - parent_list, top_scores_index, draft_tokens = organize_draft_results( - score_list, token_list, parents_list, self.speculative_num_draft_tokens - ) - - return parent_list, top_scores_index, draft_tokens - - 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 - 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() - - # Forward - 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: - # 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") - maybe_detect_inf(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 - ] - if logits_output.hidden_states is not None: - 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 - ): - self._mamba_verify_update( - batch, res, logits_output, spec_info, seq_lens_pre_verify - ) - - 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 _mamba_verify_update( - self, - batch: ScheduleBatch, - res: EagleVerifyOutput, - logits_output: LogitsProcessorOutput, - spec_info: EagleVerifyInput, - seq_lens_pre_verify: torch.Tensor, - ): - # Under DP attention, some ranks can be IDLE during target verify and never - # initialize mamba forward metadata for this step. - if batch.forward_mode.is_idle(): - return - - num_correct_drafts = torch.tensor( - res.num_correct_drafts_per_req_cpu, - device=logits_output.next_token_logits.device, - dtype=torch.int64, - ) - cumulative_num_accept_tokens = torch.cumsum(num_correct_drafts + 1, dim=0) - # prepend 0 to the cumulative_num_accept_tokens - accepted_indices_start = torch.cat( - [ - torch.zeros( - 1, - dtype=cumulative_num_accept_tokens.dtype, - device=cumulative_num_accept_tokens.device, - ), - cumulative_num_accept_tokens[:-1], - ] - ) - accepted_indices_offset = torch.arange( - 0, - len(batch.seq_lens) * batch.spec_info.draft_token_num, - step=batch.spec_info.draft_token_num, - dtype=accepted_indices_start.dtype, - device=accepted_indices_start.device, - ) - - # 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. - last_correct_step_indices = ( - res.accept_indices[cumulative_num_accept_tokens - 1] - - accepted_indices_offset - ) - else: - last_correct_step_indices = num_correct_drafts - - if batch.mamba_track_indices is not None: - # If after verify, the request's seq_lens has crossed a mamba track interval, - # we need to update the mamba state for the request at the crossing point. - mamba_track_interval = self.server_args.mamba_track_interval - to_track_mask = ( - seq_lens_pre_verify // mamba_track_interval - != batch.seq_lens // mamba_track_interval - ) - tracking_point = ( - batch.seq_lens // mamba_track_interval * mamba_track_interval - ) - to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0) - mamba_steps_to_track = torch.where( - to_track_mask, - res.accept_indices[to_track_ith + accepted_indices_start] - - accepted_indices_offset, - -1, - ) - else: - mamba_steps_to_track = None - - self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify( - last_correct_step_indices=last_correct_step_indices, - mamba_track_indices=batch.mamba_track_indices, - mamba_steps_to_track=mamba_steps_to_track, - model=self.target_worker.model_runner.model, - ) - - 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, - ): - """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.draft_model_runner) - forward_batch.return_logprob = False - if mm_input_embeds is not None: - forward_batch.mm_input_embeds = mm_input_embeds - logits_output = self.draft_model_runner.forward(forward_batch).logits_output - maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill") - assert isinstance(forward_batch.spec_info, EagleDraftInput) - assert forward_batch.spec_info is batch.spec_info - self.capture_for_decode(logits_output, forward_batch.spec_info) - - 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 (DP attn rank w/o reqs) and active batch with all reqs - # finished. prepare_for_idle() is idempotent on already-idle. - 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 - # Verify-time construction of EagleDraftExtendInput uses the dataclass - # default (LAST); override here so ForwardBatch.init_new picks up the - # correct mode (NULL for STANDALONE). - draft_extend_input.capture_hidden_mode = draft_extend_capture_mode - forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner) - assert forward_batch.capture_hidden_mode == draft_extend_capture_mode - 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() - - # Phase 2: run draft-extend forward - can_cuda_graph = ( - self.cuda_graph_runner_for_draft_extend - and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) - ) - if can_cuda_graph: - logits_output = self.cuda_graph_runner_for_draft_extend.replay( - forward_batch - ) - # cuda-graph replay populates logits_output.{topk_p, topk_index, hidden_states}. - topk_p = logits_output.topk_p - topk_index = logits_output.topk_index - hidden_states = logits_output.hidden_states - else: - forward_batch.can_run_dp_cuda_graph = False - attn_backend = None - if not forward_batch.forward_mode.is_idle(): - attn_backend = ( - self.draft_extend_attn_backend - or self.draft_model_runner.attn_backend - ) - attn_backend.init_forward_metadata(forward_batch) - forward_batch.mark_forward_metadata_ready() - # Publish the chosen backend via ForwardContext so model code - # picks it up for this forward (no runner-attr mutation). - if attn_backend is not None: - ctx_mgr = forward_context(ForwardContext(attn_backend=attn_backend)) - else: - ctx_mgr = contextlib.nullcontext() - with ctx_mgr: - logits_output = self.draft_model_runner.forward( - forward_batch - ).logits_output - # Non-cuda-graph path: compute topk_p / topk_index inline. - probs = torch.softmax(logits_output.next_token_logits, dim=-1) - topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) - hidden_states = logits_output.hidden_states - - maybe_detect_nan( - logits_output.next_token_logits, - f"draft_extend_after_decode (cuda_graph={can_cuda_graph})", - ) - - # 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=hidden_states, - topk_p=topk_p, - topk_index=topk_index, - 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 - - def capture_for_decode( - self, logits_output: LogitsProcessorOutput, draft_input: EagleDraftInput - ): - probs = torch.softmax(logits_output.next_token_logits, dim=-1) - draft_input.topk_p, draft_input.topk_index = fast_topk(probs, self.topk, dim=-1) - draft_input.hidden_states = logits_output.hidden_states - - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): - monkey_patch_torch_reductions() - named_tensors = MultiprocessingSerializer.deserialize( - recv_req.serialized_named_tensors[self.tp_rank] - ) - success, message = self.model_runner.update_weights_from_tensor( - named_tensors=named_tensors, - load_format=recv_req.load_format, - ) - if not success: - return success, message - - success, message = self.target_worker.model_runner.update_weights_from_tensor( - named_tensors=named_tensors, - load_format=recv_req.load_format, - ) - return success, message - - -@torch.compile(dynamic=True, disable=(_is_npu or _is_musa)) -def get_last_loc_large_page_size_top_k_1( - req_to_token: torch.Tensor, - req_pool_indices: torch.Tensor, - seq_lens, - speculative_num_steps: int, -): - prefix_lens = seq_lens - seq_lens = prefix_lens + speculative_num_steps - last_loc = get_last_loc( - req_to_token, - req_pool_indices, - prefix_lens, - ) - return prefix_lens, seq_lens, last_loc diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index 5fae03b8e..adb556452 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -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__( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py deleted file mode 100644 index 3712c8460..000000000 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ /dev/null @@ -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 diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 02cb0c3fd..786aed78f 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -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( diff --git a/python/sglang/srt/speculative/standalone_worker.py b/python/sglang/srt/speculative/standalone_worker.py deleted file mode 100644 index ea216b4ae..000000000 --- a/python/sglang/srt/speculative/standalone_worker.py +++ /dev/null @@ -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) diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 336b6419e..a064c8092 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -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__( diff --git a/python/sglang/test/kits/streaming_session_kit.py b/python/sglang/test/kits/streaming_session_kit.py index 0d4e7d6fa..4d02263e6 100644 --- a/python/sglang/test/kits/streaming_session_kit.py +++ b/python/sglang/test/kits/streaming_session_kit.py @@ -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 diff --git a/test/registered/sessions/test_streaming_session_extra.py b/test/registered/sessions/test_streaming_session_extra.py index 582bc4740..de296001c 100644 --- a/test/registered/sessions/test_streaming_session_extra.py +++ b/test/registered/sessions/test_streaming_session_extra.py @@ -51,9 +51,9 @@ _EAGLE3_SPEC_ARGS = [ class TestStreamingSessionEagle(StreamingSessionServerBase, StreamingSessionKitMixin): - """EAGLE3 spec v1 (overlap disabled); offset=-1 — see kit's note.""" + """EAGLE3 spec v2, overlap disabled; inherited count jitters {0, -1} — see kit's note.""" - kv_inherit_offset = -1 + kv_inherit_offsets = (0, -1) model = DEFAULT_TARGET_MODEL_EAGLE3 extra_args = [ "--disable-overlap-schedule", @@ -82,10 +82,11 @@ class TestStreamingSessionEagleV2(StreamingSessionServerBase, StreamingSessionKi class TestStreamingSessionEagleRetractLargePage( StreamingSessionServerBase, StreamingSessionKitMixin ): - """EAGLE3 spec v1 + retract + page=256: max-pressure on `_free_tail` - (spec tail + retract alloc-commit gap + page alignment).""" + """EAGLE3 spec v2 (overlap disabled) + retract + page=256: max-pressure on + `_free_tail` (spec tail + retract alloc-commit gap + page alignment). + Inherited count jitters {0, -1} — see kit's note.""" - kv_inherit_offset = -1 + kv_inherit_offsets = (0, -1) model = DEFAULT_TARGET_MODEL_EAGLE3 extra_args = [ "--disable-overlap-schedule",