[Spec] Extract stateless draft prepare helpers into eagle_worker_common (#31257)

python/sglang/srt/kv_canary/plan_input.py

@@ -118,7 +118,7 @@ def _extract_prefix_lens_and_extend_seq_lens(
         out_prefix_lens.copy_(forward_batch.seq_lens[:bs].to(torch.int64))
         out_extend_seq_lens.fill_(int(spec_info.draft_token_num))
     elif forward_mode.is_draft_extend_v2():
-        # Evidence: EagleDraftWorkerBase.prepare_for_draft_extend bumps
+        # Evidence: eagle_worker_common.prepare_for_draft_extend bumps
         # seq_lens by num_draft_tokens. FlashAttentionBackend.init_forward_metadata reads the
         # draft-extend-v2 query length from spec_info.extend_seq_lens_tensor when available.
         # CUDA-graph replay passes extend_seq_lens but omits extend_prefix_lens, so derive the

python/sglang/srt/speculative/base_spec_worker.py

@@ -1,80 +1,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
-
-import torch
-
-from sglang.srt.utils import is_cpu
-
-_is_cpu = is_cpu()
-
-if _is_cpu:
-    from sgl_kernel import assign_draft_cache_locs_contiguous_cpu
 
 if TYPE_CHECKING:
     from sglang.srt.managers.io_struct import (
         UpdateWeightFromDiskReqInput,
         UpdateWeightsFromIPCReqInput,
     )
-    from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.tp_worker import TpModelWorker
-    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
     from sglang.srt.model_executor.model_runner import ModelRunner
-    from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
-        EAGLEDraftCudaGraphRunner,
-    )
-    from sglang.srt.speculative.eagle_info import (
-        EagleDraftExtendInput,
-        EagleDraftInput,
-    )
-
-
-def duplicate_prefix_tail_to_draft_branches(
-    token_to_kv_pool,
-    rows: torch.Tensor,
-    prefix_base: torch.Tensor,
-    last_page: torch.Tensor,
-    num_new_pages: torch.Tensor,
-    topk: int,
-    page_size: int,
-) -> None:
-    """Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1).
-
-    The draft-decode expand pass reads each branch's own draft page by block id
-    (cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the
-    real prefix tail (branch 0's first page already is it). Mirrors V1 #7725.
-    """
-    if topk <= 1:
-        return
-    bs = rows.shape[0]
-    page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64)
-    branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view(
-        1, topk - 1, 1
-    )
-    # Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch.
-    src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand(
-        bs, topk - 1, page_size
-    )
-    # Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size).
-    tgt_pos = (
-        prefix_base.view(bs, 1, 1)
-        + branches * (num_new_pages.view(bs, 1, 1) * page_size)
-        + page_off.view(1, 1, page_size)
-    )
-    # Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's
-    # own draft slots and must not be overwritten.
-    vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand(
-        bs, topk - 1, page_size
-    )
-    src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape(
-        bs, topk - 1, page_size
-    )[vmask]
-    tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape(
-        bs, topk - 1, page_size
-    )[vmask]
-    if src_slots.numel() > 0:
-        token_to_kv_pool.move_kv_cache(tgt_slots, src_slots)
 
 
 class EagleDraftWorkerBase(ABC):
@@ -106,209 +41,6 @@ class EagleDraftWorkerBase(ABC):
         self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False)
         self._capture_cuda_graphs()
 
-    def prepare_for_draft_extend(
-        self,
-        draft_extend_input: EagleDraftExtendInput,
-        batch: ScheduleBatch,
-        predict: torch.Tensor,
-        num_draft_tokens: int,
-        draft_model_runner: Any,
-        cuda_graph_runner: Any,
-        *,
-        return_hidden_states_before_norm: bool,
-    ):
-        from sglang.srt.model_executor.forward_batch_info import (
-            CaptureHiddenMode,
-            ForwardBatch,
-            ForwardMode,
-        )
-        from sglang.srt.utils.async_probe import maybe_detect_oob
-        from sglang.srt.utils.common import is_npu
-
-        bs = len(batch.seq_lens)
-        extend_num_tokens = bs * num_draft_tokens
-        # When seq_lens_cpu is absent, stay on GPU-only path -- no .tolist()/.cpu().
-        gpu_only = batch.seq_lens_cpu is None
-
-        batch.spec_info = draft_extend_input
-        # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode)
-        # may run this under a plan stream; casting inside the plan stream creates a
-        # cross-stream dependency that can lead to data races and break MTP acceptance.
-        # The caller should cast to int64 before entering the plan stream context.
-        batch.input_ids = predict
-        maybe_detect_oob(
-            batch.input_ids,
-            0,
-            batch.model_config.vocab_size,
-            "v2 prepare_for_draft_extend input_ids",
-        )
-        # init_new requires both list or both Tensor;
-        # gpu_only emits device tensors to skip H2D.
-        if gpu_only:
-            batch.prefix_lens = batch.seq_lens.to(torch.int32)
-            batch.extend_lens = torch.full(
-                (bs,), num_draft_tokens, dtype=torch.int32, device=batch.seq_lens.device
-            )
-        else:
-            batch.prefix_lens = batch.seq_lens_cpu.tolist()
-            batch.extend_lens = [num_draft_tokens] * bs
-        batch.extend_num_tokens = extend_num_tokens
-        capture_mode = (
-            CaptureHiddenMode.NULL
-            if draft_model_runner.spec_algorithm.is_standalone()
-            else CaptureHiddenMode.FULL
-        )
-        batch.forward_mode = (
-            ForwardMode.IDLE
-            if batch.forward_mode.is_idle()
-            else ForwardMode.DRAFT_EXTEND_V2
-        )
-        forward_batch = ForwardBatch.init_new(
-            batch,
-            draft_model_runner,
-            capture_hidden_mode=capture_mode,
-            return_hidden_states_before_norm=return_hidden_states_before_norm,
-        )
-        # Forward sees post-write length (draft extend writes num_draft_tokens
-        # slots); mutation stays on forward_batch to preserve SB.seq_lens.
-        forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens
-        if not gpu_only:
-            forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
-            forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum())
-        else:
-            # Supply CPU mirror (extend_seq_lens are all num_draft_tokens) so
-            # backend max() reads from list without a per-iter D2H sync.
-            forward_batch.extend_seq_lens_cpu = [num_draft_tokens] * bs
-        can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
-            forward_batch
-        )
-        if not batch.forward_mode.is_idle() and not can_cuda_graph:
-            draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
-            # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer
-            # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch
-            # (the `_batch_size == batch_size` assertion, see #27091); the
-            # marked pre-pad metadata is used as-is, matching the proven
-            # skip_attn_backend_init=True behavior.
-            # On NPU with --disable-cuda-graph, block_table shape won't match
-            # after prepare_mlp_sync_batch padding; defer re-init to
-            # forward_extend (post-pad) instead.
-            if not is_npu() or can_cuda_graph:
-                forward_batch.mark_forward_metadata_ready()
-        return forward_batch
-
-    def prepare_for_draft(
-        self,
-        draft_input: EagleDraftInput,
-        req_to_token_pool: ReqToTokenPool,
-        batch: ScheduleBatch,
-        cuda_graph_runner: EAGLEDraftCudaGraphRunner,
-        draft_model_runner: ModelRunner,
-        topk: int,
-        num_steps: int,
-    ):
-        from sglang.kernels.ops.speculative.cache_locs import (
-            assign_draft_cache_locs_contiguous,
-        )
-        from sglang.srt.model_executor.forward_batch_info import (
-            CaptureHiddenMode,
-            ForwardBatch,
-        )
-
-        if not batch.forward_mode.is_idle():
-            bs = len(batch.seq_lens)
-
-            # Assign cache locations (draft-write targets).
-            page_size = batch.token_to_kv_pool_allocator.page_size
-            if page_size == 1 or topk == 1:
-                batch.out_cache_loc = torch.empty(
-                    (bs * topk * num_steps,),
-                    dtype=torch.int64,
-                    device=batch.device,
-                )
-                if _is_cpu:
-                    assign_draft_cache_locs_contiguous_cpu(
-                        batch.req_pool_indices,
-                        req_to_token_pool.req_to_token,
-                        batch.seq_lens,
-                        batch.out_cache_loc,
-                        req_to_token_pool.req_to_token.shape[1],
-                        topk,
-                        num_steps,
-                    )
-                else:
-                    # FIXME(lsyin): align with the default code path
-                    assign_draft_cache_locs_contiguous[(bs,)](
-                        batch.req_pool_indices,
-                        req_to_token_pool.req_to_token,
-                        batch.seq_lens,
-                        batch.out_cache_loc,
-                        req_to_token_pool.req_to_token.shape[1],
-                        topk,
-                        num_steps,
-                    )
-            else:
-                # page_size > 1 + topk > 1: per-branch page-aligned draft pages.
-                # Reduce out_cache_loc from the page-aligned tree region down to the
-                # dense draft slots (skip each branch's duplicated prefix-tail slots
-                # and trailing padding), matching generate_draft_decode_kv_indices'
-                # paged read formula: prefix_base + t*num_new_pages*page + last_page + s.
-                # base is batch.seq_lens (== KV-ready committed prefix at draft time;
-                # the bonus is the tree root written by verify, not part of [0:seq_lens]).
-                rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()]
-                seq_lens = batch.seq_lens.to(torch.int64)
-                last_page = seq_lens % page_size
-                prefix_base = seq_lens - last_page
-                num_new_pages = (last_page + num_steps + page_size - 1) // page_size
-                topk_ids = torch.arange(
-                    topk, device=rows.device, dtype=torch.int64
-                ).view(1, topk)
-                starts = (
-                    prefix_base.view(bs, 1)
-                    + topk_ids * (num_new_pages.view(bs, 1) * page_size)
-                    + last_page.view(bs, 1)
-                )
-                steps = torch.arange(
-                    num_steps, device=rows.device, dtype=torch.int64
-                ).view(1, 1, num_steps)
-                pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps)
-                batch.out_cache_loc = (
-                    torch.gather(rows, 1, pos).reshape(-1).contiguous()
-                )
-
-                # Each branch's page-aligned region starts with `last_page` hole slots
-                # overlapping the prefix tail page; duplicate the real prefix-tail KV
-                # into them so whole-page reads stay coherent (see helper docstring).
-                duplicate_prefix_tail_to_draft_branches(
-                    draft_model_runner.token_to_kv_pool,
-                    rows,
-                    prefix_base,
-                    last_page,
-                    num_new_pages,
-                    topk,
-                    page_size,
-                )
-
-        # Get a forward batch
-        # Actual width of the next draft-decode forward: topk tokens per req.
-        draft_input.num_tokens_per_req = topk
-        draft_input.num_tokens_for_logprob_per_req = topk
-        capture_mode = (
-            CaptureHiddenMode.NULL
-            if draft_model_runner.spec_algorithm.is_standalone()
-            else CaptureHiddenMode.LAST
-        )
-        draft_input.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
-        forward_batch = ForwardBatch.init_new(
-            batch,
-            draft_model_runner,
-            capture_hidden_mode=capture_mode,
-            return_hidden_states_before_norm=False,
-        )
-        can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
-            forward_batch
-        )
-        return forward_batch, can_cuda_graph
-
 
 class BaseSpecWorker(ABC):
     @property

python/sglang/srt/speculative/eagle_worker_common.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.kernels.ops.speculative.cache_locs import (
+    assign_draft_cache_locs_contiguous,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.utils import is_cpu
+from sglang.srt.utils.async_probe import maybe_detect_oob
+from sglang.srt.utils.common import is_npu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import assign_draft_cache_locs_contiguous_cpu
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+        EAGLEDraftCudaGraphRunner,
+    )
+    from sglang.srt.speculative.eagle_info import (
+        EagleDraftExtendInput,
+        EagleDraftInput,
+    )
+
+
+def duplicate_prefix_tail_to_draft_branches(
+    token_to_kv_pool,
+    rows: torch.Tensor,
+    prefix_base: torch.Tensor,
+    last_page: torch.Tensor,
+    num_new_pages: torch.Tensor,
+    topk: int,
+    page_size: int,
+) -> None:
+    """Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1).
+
+    The draft-decode expand pass reads each branch's own draft page by block id
+    (cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the
+    real prefix tail (branch 0's first page already is it). Mirrors V1 #7725.
+    """
+    if topk <= 1:
+        return
+    bs = rows.shape[0]
+    page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64)
+    branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view(
+        1, topk - 1, 1
+    )
+    # Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch.
+    src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand(
+        bs, topk - 1, page_size
+    )
+    # Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size).
+    tgt_pos = (
+        prefix_base.view(bs, 1, 1)
+        + branches * (num_new_pages.view(bs, 1, 1) * page_size)
+        + page_off.view(1, 1, page_size)
+    )
+    # Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's
+    # own draft slots and must not be overwritten.
+    vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand(
+        bs, topk - 1, page_size
+    )
+    src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape(
+        bs, topk - 1, page_size
+    )[vmask]
+    if src_slots.numel() > 0:
+        token_to_kv_pool.move_kv_cache(tgt_slots, src_slots)
+
+
+def prepare_for_draft_extend(
+    draft_extend_input: EagleDraftExtendInput,
+    batch: ScheduleBatch,
+    predict: torch.Tensor,
+    num_draft_tokens: int,
+    draft_model_runner: Any,
+    cuda_graph_runner: Any,
+    *,
+    return_hidden_states_before_norm: bool,
+):
+    bs = len(batch.seq_lens)
+    extend_num_tokens = bs * num_draft_tokens
+    # When seq_lens_cpu is absent, stay on GPU-only path -- no .tolist()/.cpu().
+    gpu_only = batch.seq_lens_cpu is None
+
+    batch.spec_info = draft_extend_input
+    # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode)
+    # may run this under a plan stream; casting inside the plan stream creates a
+    # cross-stream dependency that can lead to data races and break MTP acceptance.
+    # The caller should cast to int64 before entering the plan stream context.
+    batch.input_ids = predict
+    maybe_detect_oob(
+        batch.input_ids,
+        0,
+        batch.model_config.vocab_size,
+        "v2 prepare_for_draft_extend input_ids",
+    )
+    # init_new requires both list or both Tensor;
+    # gpu_only emits device tensors to skip H2D.
+    if gpu_only:
+        batch.prefix_lens = batch.seq_lens.to(torch.int32)
+        batch.extend_lens = torch.full(
+            (bs,), num_draft_tokens, dtype=torch.int32, device=batch.seq_lens.device
+        )
+    else:
+        batch.prefix_lens = batch.seq_lens_cpu.tolist()
+        batch.extend_lens = [num_draft_tokens] * bs
+    batch.extend_num_tokens = extend_num_tokens
+    capture_mode = (
+        CaptureHiddenMode.NULL
+        if draft_model_runner.spec_algorithm.is_standalone()
+        else CaptureHiddenMode.FULL
+    )
+    batch.forward_mode = (
+        ForwardMode.IDLE
+        if batch.forward_mode.is_idle()
+        else ForwardMode.DRAFT_EXTEND_V2
+    )
+    forward_batch = ForwardBatch.init_new(
+        batch,
+        draft_model_runner,
+        capture_hidden_mode=capture_mode,
+        return_hidden_states_before_norm=return_hidden_states_before_norm,
+    )
+    # Forward sees post-write length (draft extend writes num_draft_tokens
+    # slots); mutation stays on forward_batch to preserve SB.seq_lens.
+    forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens
+    if not gpu_only:
+        forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
+        forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum())
+    else:
+        # Supply CPU mirror (extend_seq_lens are all num_draft_tokens) so
+        # backend max() reads from list without a per-iter D2H sync.
+        forward_batch.extend_seq_lens_cpu = [num_draft_tokens] * bs
+    can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+        forward_batch
+    )
+    if not batch.forward_mode.is_idle() and not can_cuda_graph:
+        draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
+        # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer
+        # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch
+        # (the `_batch_size == batch_size` assertion, see #27091); the
+        # marked pre-pad metadata is used as-is, matching the proven
+        # skip_attn_backend_init=True behavior.
+        # On NPU with --disable-cuda-graph, block_table shape won't match
+        # after prepare_mlp_sync_batch padding; defer re-init to
+        # forward_extend (post-pad) instead.
+        if not is_npu() or can_cuda_graph:
+            forward_batch.mark_forward_metadata_ready()
+    return forward_batch
+
+
+def prepare_for_draft(
+    draft_input: EagleDraftInput,
+    req_to_token_pool: ReqToTokenPool,
+    batch: ScheduleBatch,
+    cuda_graph_runner: EAGLEDraftCudaGraphRunner,
+    draft_model_runner: ModelRunner,
+    topk: int,
+    num_steps: int,
+):
+
+    if not batch.forward_mode.is_idle():
+        bs = len(batch.seq_lens)
+
+        # Assign cache locations (draft-write targets).
+        page_size = batch.token_to_kv_pool_allocator.page_size
+        if page_size == 1 or topk == 1:
+            batch.out_cache_loc = torch.empty(
+                (bs * topk * num_steps,),
+                dtype=torch.int64,
+                device=batch.device,
+            )
+            if _is_cpu:
+                assign_draft_cache_locs_contiguous_cpu(
+                    batch.req_pool_indices,
+                    req_to_token_pool.req_to_token,
+                    batch.seq_lens,
+                    batch.out_cache_loc,
+                    req_to_token_pool.req_to_token.shape[1],
+                    topk,
+                    num_steps,
+                )
+            else:
+                # FIXME(lsyin): align with the default code path
+                assign_draft_cache_locs_contiguous[(bs,)](
+                    batch.req_pool_indices,
+                    req_to_token_pool.req_to_token,
+                    batch.seq_lens,
+                    batch.out_cache_loc,
+                    req_to_token_pool.req_to_token.shape[1],
+                    topk,
+                    num_steps,
+                )
+        else:
+            # page_size > 1 + topk > 1: per-branch page-aligned draft pages.
+            # Reduce out_cache_loc from the page-aligned tree region down to the
+            # dense draft slots (skip each branch's duplicated prefix-tail slots
+            # and trailing padding), matching generate_draft_decode_kv_indices'
+            # paged read formula: prefix_base + t*num_new_pages*page + last_page + s.
+            # base is batch.seq_lens (== KV-ready committed prefix at draft time;
+            # the bonus is the tree root written by verify, not part of [0:seq_lens]).
+            rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()]
+            seq_lens = batch.seq_lens.to(torch.int64)
+            last_page = seq_lens % page_size
+            prefix_base = seq_lens - last_page
+            num_new_pages = (last_page + num_steps + page_size - 1) // page_size
+            topk_ids = torch.arange(topk, device=rows.device, dtype=torch.int64).view(
+                1, topk
+            )
+            starts = (
+                prefix_base.view(bs, 1)
+                + topk_ids * (num_new_pages.view(bs, 1) * page_size)
+                + last_page.view(bs, 1)
+            )
+            steps = torch.arange(num_steps, device=rows.device, dtype=torch.int64).view(
+                1, 1, num_steps
+            )
+            pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps)
+            batch.out_cache_loc = torch.gather(rows, 1, pos).reshape(-1).contiguous()
+
+            # Each branch's page-aligned region starts with `last_page` hole slots
+            # overlapping the prefix tail page; duplicate the real prefix-tail KV
+            # into them so whole-page reads stay coherent (see helper docstring).
+            duplicate_prefix_tail_to_draft_branches(
+                draft_model_runner.token_to_kv_pool,
+                rows,
+                prefix_base,
+                last_page,
+                num_new_pages,
+                topk,
+                page_size,
+            )
+
+    # Get a forward batch
+    # Actual width of the next draft-decode forward: topk tokens per req.
+    draft_input.num_tokens_per_req = topk
+    draft_input.num_tokens_for_logprob_per_req = topk
+    capture_mode = (
+        CaptureHiddenMode.NULL
+        if draft_model_runner.spec_algorithm.is_standalone()
+        else CaptureHiddenMode.LAST
+    )
+    draft_input.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
+    forward_batch = ForwardBatch.init_new(
+        batch,
+        draft_model_runner,
+        capture_hidden_mode=capture_mode,
+        return_hidden_states_before_norm=False,
+    )
+    can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+        forward_batch
+    )
+    return forward_batch, can_cuda_graph

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -73,6 +73,10 @@ from sglang.srt.speculative.eagle_utils import (
     organize_draft_results,
     per_step_draft_out_cache_loc,
 )
+from sglang.srt.speculative.eagle_worker_common import (
+    prepare_for_draft,
+    prepare_for_draft_extend,
+)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import (
     commit_mamba_states_after_verify,
@@ -496,7 +500,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
     def draft(self, batch: ScheduleBatch):
         draft_input: EagleDraftInput = batch.spec_info
-        forward_batch, can_cuda_graph = self.prepare_for_draft(
+        forward_batch, can_cuda_graph = prepare_for_draft(
             draft_input,
             self.req_to_token_pool,
             batch,
@@ -916,7 +920,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         # Prepare for draft extend in a separate stream
         with self.plan_stream_ctx:
-            forward_batch = self.prepare_for_draft_extend(
+            forward_batch = prepare_for_draft_extend(
                 draft_extend_input,
                 batch,
                 next_token_ids,

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -55,6 +55,10 @@ from sglang.srt.speculative.eagle_utils import (
     eagle_sample,
     get_draft_recurrent_hidden_state_spec,
 )
+from sglang.srt.speculative.eagle_worker_common import (
+    prepare_for_draft,
+    prepare_for_draft_extend,
+)
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
 )
@@ -235,7 +239,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
 
     def draft(self, batch: ScheduleBatch):
         draft_input: EagleDraftInput = batch.spec_info
-        forward_batch, can_cuda_graph = self.prepare_for_draft(
+        forward_batch, can_cuda_graph = prepare_for_draft(
             draft_input,
             self.req_to_token_pool,
             batch,
@@ -526,7 +530,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         # Prepare for draft extend in a separate stream
         # Notice that here we use batch_result.next_token_ids as the input ids
         with self.plan_stream_ctx:
-            forward_batch = self.prepare_for_draft_extend(
+            forward_batch = prepare_for_draft_extend(
                 draft_extend_input,
                 batch,
                 batch_result.next_token_ids,

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -163,7 +163,6 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
                 worker.seed_dsa_topk_from_draft_extend = seed_enabled
                 worker.index_share_for_mtp_iteration = True
                 forward_batch = SimpleNamespace(forward_mode=ForwardMode.DECODE)
-                worker.prepare_for_draft = MagicMock(return_value=(forward_batch, True))
                 worker.draft_forward = MagicMock(return_value=graph_result)
                 attn_backend = SimpleNamespace(
                     get_verify_buffers_to_fill_after_draft=lambda: (None, None),
@@ -190,6 +189,9 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
                 with patch(
                     "sglang.srt.speculative.eagle_worker_v2.build_tree_kernel_efficient",
                     return_value=tree_result,
+                ), patch(
+                    "sglang.srt.speculative.eagle_worker_v2.prepare_for_draft",
+                    return_value=(forward_batch, True),
                 ):
                     worker.draft(batch)