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

[Spec] Naming cleanup: contiguous draft-loc kernel + accepted->accept (#27599)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-08 15:04:58 -07:00
committed by GitHub
parent c95179bc85
commit 3fe6bc390b
6 changed files with 27 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_info_v2.py
+3 -3
View File
@@ -37,7 +37,7 @@ from sglang.srt.speculative.spec_utils import (
generate_simulated_accept_index,
)
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_draft_cache_locs_page_size_1 as assign_draft_cache_locs_page_size_1,
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,
@@ -46,7 +46,7 @@ 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_accepted_out_cache_loc as fill_accepted_out_cache_loc,
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,
@@ -241,7 +241,7 @@ class EagleDraftInputV2Mixin:
device=batch.device,
)
# FIXME(lsyin): align with the default code path
assign_draft_cache_locs_page_size_1[(bs,)](
assign_draft_cache_locs_contiguous[(bs,)](
batch.req_pool_indices,
req_to_token_pool.req_to_token,
batch.seq_lens,
python/sglang/srt/speculative/eagle_worker_v2.py
+16 -20
View File
@@ -53,7 +53,7 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info_v2 import (
assign_extend_cache_locs,
fill_accepted_out_cache_loc,
fill_accept_out_cache_loc,
fill_bonus_tokens,
)
from sglang.srt.speculative.eagle_utils import (
@@ -1300,7 +1300,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
if not batch.forward_mode.is_idle() and self.topk > 1:
# topk == 1 needs nothing here: the accepted path is already the front
# chain, so the whole compaction is an identity transform.
predict = self._finalize_accepted_tree_path(
predict = self._finalize_accept_tree_path(
batch, accept_index, accept_lens, predict, logits_output, bs
)
@@ -1333,7 +1333,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
"""Update mamba state for hybrid GDN models after verification."""
# `accept_lens` already includes the bonus token (drafts + 1 per req).
if not batch.forward_mode.is_idle() and accept_index.numel() > 0:
accepted_indices_offset = torch.arange(
accept_indices_offset = torch.arange(
0,
bs * self.speculative_num_draft_tokens,
step=self.speculative_num_draft_tokens,
@@ -1345,7 +1345,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
# mamba state to commit; reduces to accept_lens - 1 for topk == 1.
last_correct_step_indices = (
accept_index[req_idx, (accept_lens - 1).to(torch.int64)]
- accepted_indices_offset
- accept_indices_offset
)
if batch.mamba_track_indices is not None:
@@ -1365,7 +1365,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
tracking_point - seq_lens_pre_verify - 1, min=0
).to(torch.int64)
candidate_track_steps = (
accept_index[req_idx, to_track_ith] - accepted_indices_offset
accept_index[req_idx, to_track_ith] - accept_indices_offset
)
mamba_steps_to_track = torch.where(
to_track_mask,
@@ -1382,7 +1382,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
model=self.target_worker.model_runner.model,
)
def _finalize_accepted_tree_path(
def _finalize_accept_tree_path(
self,
batch: ScheduleBatch,
accept_index: torch.Tensor,
@@ -1396,17 +1396,15 @@ 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_accepted_tokens_to_target_kvcache(
batch, accept_index, accept_lens - 1
)
predict = self._compact_accepted_to_front(predict, accept_index, bs)
self.move_accept_tokens_to_target_kvcache(batch, accept_index, accept_lens - 1)
predict = self._compact_accept_to_front(predict, accept_index, bs)
if logits_output.hidden_states is not None:
logits_output.hidden_states = self._compact_accepted_to_front(
logits_output.hidden_states = self._compact_accept_to_front(
logits_output.hidden_states, accept_index, bs
)
return predict
def move_accepted_tokens_to_target_kvcache(
def move_accept_tokens_to_target_kvcache(
self,
batch: ScheduleBatch,
accept_index: torch.Tensor,
@@ -1426,7 +1424,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
# tree exceeds the accepted chain, over-reading accept_index (illegal memory).
size = bs * accept_index.shape[1]
# fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
# fill_accept_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok.
maybe_detect_oob(
accept_index,
-1,
@@ -1439,9 +1437,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
dtype=torch.int64,
device=self.device,
)
accepted_out_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,
@@ -1451,17 +1447,17 @@ class EAGLEWorkerV2(BaseSpecWorker):
self.req_to_token_pool.req_to_token.shape[1],
next_power_of_2(bs),
)
fill_accepted_out_cache_loc[(size,)](
fill_accept_out_cache_loc[(size,)](
accept_index,
batch.out_cache_loc,
accepted_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, accepted_out_cache_loc
tgt_cache_loc, accept_out_cache_loc
)
def _compact_accepted_to_front(
def _compact_accept_to_front(
self, x: torch.Tensor, accept_index: torch.Tensor, bs: int
) -> torch.Tensor:
"""Gather the accepted tree path to the front of each per-req block.
python/sglang/srt/speculative/spec_utils.py
+3 -3
View File
@@ -335,13 +335,13 @@ def traverse_tree(
is_accepted = True
else:
parent_bitmask = allocate_token_bitmask[parent_pos]
curr_token_id = draft_tokens[curr]
if vocab_size and curr_token_id >= vocab_size:
current_token = draft_tokens[curr]
if vocab_size and current_token >= vocab_size:
is_accepted = False
else:
# 32 boolean bitmask values are packed into 32-bit integers
is_accepted = (
parent_bitmask[curr_token_id // 32] & (1 << (curr_token_id % 32))
parent_bitmask[current_token // 32] & (1 << (current_token % 32))
) != 0
if is_accepted:
python/sglang/srt/speculative/triton_ops/cache_locs.py
+1 -1
View File
@@ -94,7 +94,7 @@ def assign_req_to_token_pool_func(
@triton.jit
def assign_draft_cache_locs_page_size_1(
def assign_draft_cache_locs_contiguous(
req_pool_indices,
req_to_token,
seq_lens,
python/sglang/srt/speculative/triton_ops/eagle.py
+3 -3
View File
@@ -22,10 +22,10 @@ def fill_bonus_tokens(
@triton.jit
def fill_accepted_out_cache_loc(
def fill_accept_out_cache_loc(
accept_index,
out_cache_loc,
accepted_out_cache_loc,
accept_out_cache_loc,
size_upper: tl.constexpr,
):
pid = tl.program_id(axis=0)
@@ -36,4 +36,4 @@ def fill_accepted_out_cache_loc(
src = tl.load(accept_index + pid)
if src > -1:
value = tl.load(out_cache_loc + src)
tl.store(accepted_out_cache_loc + dst, value)
tl.store(accept_out_cache_loc + dst, value)