[Spec] Deprecate Spec V1 (#25464)

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.",
+            "Non-overlap (synchronous) spec v2 is used for eagle/eagle3/standalone "
-            spec_v1_reason or "overlap schedule is disabled",
+            "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:

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: |

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):

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

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.
+    # Only the overlap path relays spec extras through the future_map; the
-    if batch.is_spec_v2:
+    # 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)
 
 

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
+        # Whether the V2 worker/schema is used. Independent of overlap: the
-        ret = self.enable_overlap and not self.spec_algorithm.is_none()
+        # non-overlap path also drives the V2 worker, just synchronously.
-        assert not ret or self.spec_algorithm.supports_spec_v2()
+        return self.spec_algorithm.supports_spec_v2()
-        return ret
 
     def mamba_lazy_prealloc_at_boundary(self, mamba_track_interval: int):
         """Allocate a temporary second ping-pong slot for reqs at a track boundary.

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):
+    def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool):
-        """Make SB transactional across one overlap forward.
+        """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
+        3. (overlap=True only) Pin (batch, snapshot) into batch_record_buf
-           tensors in the snapshot survive the caching allocator past the
+           for 2 iters so GPU tensors in the snapshot survive the caching
-           forward stream. Must run AFTER the sampling_info swap so the
+           allocator past the forward stream. Must run AFTER the sampling_info
-           forward-only copy gets pinned.
+           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,7 +2969,8 @@ 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
+        # 3. pin for 2-iter tensor lifetime (overlap path only)
+        if overlap:
             self.record_batch_in_overlap(batch)
 
         try:
@@ -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
+                        # Spec_v1 (NGRAM / DFLASH / FROZEN_KV_MTP, non-overlap):
-                        # req_pool_indices; relay is unused (worker rebuilds input_ids
+                        # worker shape doesn't match req_pool_indices; relay is
-                        # inside verify). Keep pre-PR behavior.
+                        # 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)
 

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

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

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 (

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,

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,

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
+    v2 draft workers (`EagleDraftWorker` and subclasses) expose the draft
-    `model_runner` (the worker itself runs the draft model);
+    model_runner as `draft_runner`; fall back to `model_runner` for workers
-    v2 (`EagleDraftWorker` and subclasses) exposes it as `draft_runner`.
+    that run the draft model directly.
     """
     return (
         worker.draft_runner if hasattr(worker, "draft_runner") else worker.model_runner

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:
+    """Frozen-KV MTP worker; same constructor shape as other TpModelWorker-based
-    :meth:`forward_batch_generation` (stubs for now).
+    spec workers. Entry: :meth:`forward_batch_generation` (stubs for now).
     """
 
     def __init__(

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -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

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,
-            )
-
-            return MultiLayerEagleWorker
-
         elif self.is_eagle():
-            if enable_overlap:
             from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
 
             return EAGLEWorkerV2
-
-            from sglang.srt.speculative.eagle_worker import EAGLEWorker
-
-            return EAGLEWorker
         elif self.is_standalone():
-            if enable_overlap:
             from sglang.srt.speculative.standalone_worker_v2 import (
                 StandaloneWorkerV2,
             )
 
             return StandaloneWorkerV2
-
-            from sglang.srt.speculative.standalone_worker import StandaloneWorker
-
-            return StandaloneWorker
         elif self.is_ngram():
             if enable_overlap:
                 raise ValueError(

python/sglang/srt/speculative/standalone_worker.py

@@ -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)

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__(

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
+    # Allowed inherited-cache offsets vs the previous turn's total. Non-overlap
-    # before max_new stops, so slot.kv_committed_len = input + output - 1.
+    # spec decode can be off by 1: the bonus token's KV is only computed by the
-    kv_inherit_offset = 0
+    # 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
+                allowed = {prev_kv_len + off for off in self.kv_inherit_offsets}
-                self.assertEqual(
+                self.assertIn(
                     cached,
-                    expected,
+                    allowed,
-                    f"Turn {turn_idx + 1}: inherited {cached} != expected {expected}",
+                    f"Turn {turn_idx + 1}: inherited {cached} not in {sorted(allowed)}",
                 )
             prev_kv_len = prompt_tokens + completion_tokens
 

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`
+    """EAGLE3 spec v2 (overlap disabled) + retract + page=256: max-pressure on
-    (spec tail + retract alloc-commit gap + page alignment)."""
+    `_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",