minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Spec] Extract move_accept_tokens_to_target_kvcache into spec_utils (#27764)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-09 21:55:26 -07:00
committed by GitHub
parent 2495c02c2c
commit d1895cb60d
3 changed files with 70 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_info_v2.py
-6
View File
@@ -39,15 +39,9 @@ from sglang.srt.speculative.spec_utils import (
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 as assign_extend_cache_locs,
)
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_extend_cache_locs_func as assign_extend_cache_locs_func,
)
from sglang.srt.speculative.triton_ops.eagle import (
fill_accept_out_cache_loc as fill_accept_out_cache_loc,
)
from sglang.srt.speculative.triton_ops.eagle import (
fill_bonus_tokens as fill_bonus_tokens,
)
python/sglang/srt/speculative/eagle_worker_v2.py
+5 -60
View File
@@ -56,11 +56,7 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info_v2 import (
assign_extend_cache_locs,
fill_accept_out_cache_loc,
fill_bonus_tokens,
)
from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
from sglang.srt.speculative.eagle_utils import (
TreeMaskMode,
_eagle_prefill_tail_tokens,
@@ -73,6 +69,7 @@ from sglang.srt.speculative.spec_utils import (
draft_tp_context,
generate_token_bitmask,
load_token_map,
move_accept_tokens_to_target_kvcache,
record_stream_each,
record_stream_for_v2_verify,
select_top_k_tokens,
@@ -93,7 +90,6 @@ from sglang.srt.utils.common import (
is_musa,
is_npu,
log_info_on_rank0,
next_power_of_2,
)
from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
@@ -1414,7 +1410,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
downstream chain-layout code (draft-extend select_index, committed-KV reads)
assumes. Returns compacted predict; mutates logits_output.hidden_states
(moved only when present)."""
self.move_accept_tokens_to_target_kvcache(batch, accept_index, accept_lens - 1)
move_accept_tokens_to_target_kvcache(
batch, accept_index, accept_lens - 1, self.token_to_kv_pool_allocator
)
predict = self._compact_accept_to_front(predict, accept_index, bs)
if logits_output.hidden_states is not None:
logits_output.hidden_states = self._compact_accept_to_front(
@@ -1422,59 +1420,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
)
return predict
def move_accept_tokens_to_target_kvcache(
self,
batch: ScheduleBatch,
accept_index: torch.Tensor,
num_correct_drafts: torch.Tensor,
):
"""
Move accepted tokens (drafts + bonus) to the target KV cache.
Args:
batch: The batch to run.
accept_index: The index of the accepted tokens (incl. bonus).
num_correct_drafts: Per-req count of correct drafts (excludes bonus);
seq_lens is advanced by num_correct_drafts + 1 to cover the bonus slot.
"""
bs = len(batch.seq_lens)
# accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the
# tree exceeds the accepted chain, over-reading accept_index (illegal memory).
size = bs * accept_index.shape[1]
# fill_accept_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
maybe_detect_oob(
accept_index,
-1,
batch.out_cache_loc.size(0),
"eagle v2 move_accepted_tokens accept_index",
)
tgt_cache_loc = torch.zeros(
size,
dtype=torch.int64,
device=self.device,
)
accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=self.device)
assign_extend_cache_locs[(bs,)](
batch.req_pool_indices,
self.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + num_correct_drafts + 1,
tgt_cache_loc,
self.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
)
fill_accept_out_cache_loc[(size,)](
accept_index,
batch.out_cache_loc,
accept_out_cache_loc,
next_power_of_2(size),
)
self.token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
tgt_cache_loc, accept_out_cache_loc
)
def _compact_accept_to_front(
self, x: torch.Tensor, accept_index: torch.Tensor, bs: int
) -> torch.Tensor:
python/sglang/srt/speculative/spec_utils.py
+65 -2
View File
@@ -18,6 +18,9 @@ from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.triton_ops.cache_locs import (
align_evict_mask_to_page_size as align_evict_mask_to_page_size,
)
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_extend_cache_locs as assign_extend_cache_locs,
)
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_req_to_token_pool as assign_req_to_token_pool,
)
@@ -39,7 +42,11 @@ from sglang.srt.speculative.triton_ops.cache_locs import (
from sglang.srt.speculative.triton_ops.cache_locs import (
get_target_cache_loc as get_target_cache_loc,
)
from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
from sglang.srt.speculative.triton_ops.eagle import (
fill_accept_out_cache_loc as fill_accept_out_cache_loc,
)
from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, next_power_of_2
from sglang.srt.utils.async_probe import maybe_detect_oob
_is_cuda = is_cuda()
_is_hip = is_hip()
@@ -48,7 +55,8 @@ _is_musa = is_musa()
if TYPE_CHECKING:
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.eagle_info import EagleVerifyInput
@@ -473,3 +481,58 @@ def spec_stage_span(name: str):
if torch.autograd._profiler_enabled():
return torch.profiler.record_function(name)
return nullcontext()
def move_accept_tokens_to_target_kvcache(
batch: ScheduleBatch,
accept_index: torch.Tensor,
num_correct_drafts: torch.Tensor,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
):
"""
Move accepted tokens (drafts + bonus) to the target KV cache.
Args:
batch: The batch to run.
accept_index: The index of the accepted tokens (incl. bonus).
num_correct_drafts: Per-req count of correct drafts (excludes bonus);
seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot.
"""
bs = len(batch.seq_lens)
device = batch.seq_lens.device
# accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the
# tree exceeds the accepted chain, over-reading accept_index (illegal memory).
size = bs * accept_index.shape[1]
# fill_accept_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
maybe_detect_oob(
accept_index,
-1,
batch.out_cache_loc.size(0),
"spec v2 move_accept_tokens accept_index",
)
tgt_cache_loc = torch.zeros(
size,
dtype=torch.int64,
device=device,
)
accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device)
assign_extend_cache_locs[(bs,)](
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + num_correct_drafts + 1,
tgt_cache_loc,
batch.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
)
fill_accept_out_cache_loc[(size,)](
accept_index,
batch.out_cache_loc,
accept_out_cache_loc,
next_power_of_2(size),
)
token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
tgt_cache_loc, accept_out_cache_loc
)