From dc078ddd2a97c2d7aa4c7b26fe1269ed23b48e1b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 15 Jul 2026 10:54:31 -0700 Subject: [PATCH] [Spec] Extract stateless draft prepare helpers into eagle_worker_common (#31257) --- python/sglang/srt/kv_canary/plan_input.py | 2 +- .../srt/speculative/base_spec_worker.py | 270 +----------------- .../srt/speculative/eagle_worker_common.py | 267 +++++++++++++++++ .../sglang/srt/speculative/eagle_worker_v2.py | 8 +- .../multi_layer_eagle_worker_v2.py | 8 +- .../test_eagle_worker_v2_topk1_fastpath.py | 4 +- 6 files changed, 284 insertions(+), 275 deletions(-) create mode 100644 python/sglang/srt/speculative/eagle_worker_common.py diff --git a/python/sglang/srt/kv_canary/plan_input.py b/python/sglang/srt/kv_canary/plan_input.py index e8de87090..c7fdd7c21 100644 --- a/python/sglang/srt/kv_canary/plan_input.py +++ b/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 diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index fa1d9ca70..d2bddc14e 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/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 - -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 +from typing import TYPE_CHECKING, Optional 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 diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py new file mode 100644 index 000000000..d4738280e --- /dev/null +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 5934cccfd..11eb4a4a6 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 0d8762593..f0d56323c 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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, diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index a93dfb5c6..4410a8072 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/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)