[Spec] Extract stateless draft prepare helpers into eagle_worker_common (#31257)
This commit is contained in:
Liangsheng Yin
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user