[Spec] Move prepare_for_draft to EagleDraftWorkerBase (#28105)
This commit is contained in:
Liangsheng Yin
@@ -8,7 +8,62 @@ import torch
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
||||
from sglang.srt.managers.tp_worker import TpModelWorker
|
||||
from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
|
||||
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):
|
||||
@@ -111,6 +166,101 @@ class EagleDraftWorkerBase(ABC):
|
||||
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.srt.model_executor.forward_batch_info import (
|
||||
CaptureHiddenMode,
|
||||
ForwardBatch,
|
||||
)
|
||||
from sglang.srt.speculative.triton_ops.cache_locs import (
|
||||
assign_draft_cache_locs_contiguous,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
# 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
|
||||
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)
|
||||
batch.capture_hidden_mode = capture_mode
|
||||
forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
|
||||
can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
|
||||
return forward_batch, can_cuda_graph
|
||||
|
||||
|
||||
class BaseSpecWorker(ABC):
|
||||
@property
|
||||
|
||||
@@ -25,7 +25,6 @@ from sglang.srt.model_executor.forward_batch_info import (
|
||||
ForwardBatch,
|
||||
ForwardMode,
|
||||
)
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
from sglang.srt.sampling.penaltylib.repetition_penalty import apply_scaling_penalties
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.speculative.eagle_utils import verify_tree_greedy_func
|
||||
@@ -34,9 +33,6 @@ from sglang.srt.speculative.spec_utils import (
|
||||
generate_simulated_accept_index,
|
||||
prepare_mamba_track_for_verify,
|
||||
)
|
||||
from sglang.srt.speculative.triton_ops.cache_locs import (
|
||||
assign_draft_cache_locs_contiguous as assign_draft_cache_locs_contiguous,
|
||||
)
|
||||
from sglang.srt.speculative.triton_ops.cache_locs import (
|
||||
assign_extend_cache_locs_func as assign_extend_cache_locs_func,
|
||||
)
|
||||
@@ -57,9 +53,6 @@ _is_musa = is_musa()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.tp_worker import TpModelWorker
|
||||
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
|
||||
EAGLEDraftCudaGraphRunner,
|
||||
)
|
||||
from sglang.srt.speculative.eagle_info import (
|
||||
EagleDraftInput,
|
||||
EagleVerifyInput,
|
||||
@@ -73,53 +66,6 @@ if is_cuda() or is_musa():
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EagleDraftInputV2Mixin:
|
||||
def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
|
||||
@@ -219,92 +165,6 @@ class EagleDraftInputV2Mixin:
|
||||
bs,
|
||||
)
|
||||
|
||||
def prepare_for_draft(
|
||||
self: 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,
|
||||
)
|
||||
# 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
|
||||
self.num_tokens_per_req = topk
|
||||
self.num_tokens_for_logprob_per_req = topk
|
||||
capture_mode = (
|
||||
CaptureHiddenMode.NULL
|
||||
if draft_model_runner.spec_algorithm.is_standalone()
|
||||
else CaptureHiddenMode.LAST
|
||||
)
|
||||
self.positions = batch.seq_lens.repeat_interleave(topk, dim=0)
|
||||
batch.capture_hidden_mode = capture_mode
|
||||
forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
|
||||
can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
|
||||
return forward_batch, can_cuda_graph
|
||||
|
||||
|
||||
@dataclass
|
||||
class EagleVerifyInputV2Mixin:
|
||||
|
||||
@@ -408,7 +408,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
|
||||
|
||||
def draft(self, batch: ScheduleBatch):
|
||||
draft_input: EagleDraftInput = batch.spec_info
|
||||
forward_batch, can_cuda_graph = draft_input.prepare_for_draft(
|
||||
forward_batch, can_cuda_graph = self.prepare_for_draft(
|
||||
draft_input,
|
||||
self.req_to_token_pool,
|
||||
batch,
|
||||
self.cuda_graph_runner,
|
||||
|
||||
@@ -253,7 +253,8 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
|
||||
|
||||
def draft(self, batch: ScheduleBatch):
|
||||
draft_input: EagleDraftInput = batch.spec_info
|
||||
forward_batch, can_cuda_graph = draft_input.prepare_for_draft(
|
||||
forward_batch, can_cuda_graph = self.prepare_for_draft(
|
||||
draft_input,
|
||||
self.req_to_token_pool,
|
||||
batch,
|
||||
self.cuda_graph_runner,
|
||||
|
||||
Reference in New Issue
Block a user