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

[Spec] Move eagle verify prepare_for_verify/sample to eagle_utils free helpers (#28117)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-12 20:37:33 -07:00
committed by GitHub
parent e02f7ca482
commit 32ef040618
7 changed files with 296 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_info.py
+15 -5
View File
@@ -8,10 +8,7 @@ from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.environ import envs
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin,
)
from sglang.srt.speculative.eagle_info_v2 import EagleDraftInputV2Mixin
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
logger = logging.getLogger(__name__)
@@ -30,7 +27,7 @@ def _draft_runner_of(worker):
@dataclass
class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
class EagleVerifyInput(SpecInput):
draft_token: torch.Tensor
custom_mask: torch.Tensor
positions: torch.Tensor
@@ -54,6 +51,19 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
if self.num_tokens_per_req < 0:
self.num_tokens_per_req = self.draft_token_num
@property
def max_tree_depth(self) -> int:
"""Longest root-to-leaf chain of the verify tree, incl. the root;
bounds the accept_index row width. EAGLE trees are depth-bounded by
the draft loop. Algorithms with other tree shapes override this."""
return self.spec_steps + 1
@property
def tree_topk(self) -> int:
"""Branching factor passed to the tree-verify kernels; -1 means an
irregular tree (no fixed per-level branching)."""
return self.topk
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
return self.draft_token_num, self.draft_token_num
python/sglang/srt/speculative/eagle_info_v2.py
-285
View File
@@ -4,14 +4,7 @@ from dataclasses import dataclass
from typing import TYPE_CHECKING
import torch
import torch.nn.functional as F
from sglang.srt.distributed import get_tp_group
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
is_dp_attention_enabled,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
@@ -19,50 +12,16 @@ from sglang.srt.mem_cache.common import (
get_alloc_reserve_per_decode,
get_last_loc,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
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
from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_LEN,
generate_simulated_accept_index,
prepare_mamba_track_for_verify,
)
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_bonus_tokens as fill_bonus_tokens,
)
from sglang.srt.utils.async_probe import (
maybe_detect_nan,
maybe_detect_oob,
sanitize_nan_logits,
)
from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_musa = is_musa()
if TYPE_CHECKING:
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.speculative.eagle_info import (
EagleDraftInput,
EagleVerifyInput,
)
if is_cuda() or is_musa():
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
@@ -164,247 +123,3 @@ class EagleDraftInputV2Mixin:
out_cache_loc,
bs,
)
@dataclass
class EagleVerifyInputV2Mixin:
@property
def max_tree_depth(self: EagleVerifyInput) -> int:
"""Longest root-to-leaf chain of the verify tree, incl. the root;
bounds the accept_index row width. EAGLE trees are depth-bounded by
the draft loop. Algorithms with other tree shapes override this."""
return self.spec_steps + 1
@property
def tree_topk(self: EagleVerifyInput) -> int:
"""Branching factor passed to the tree-verify kernels; -1 means an
irregular tree (no fixed per-level branching)."""
return self.topk
def prepare_for_verify(
self: EagleVerifyInput,
req_to_token_pool: ReqToTokenPool,
batch: ScheduleBatch,
target_worker: TpModelWorker,
):
if not batch.forward_mode.is_idle():
# Assign cache locations
bs = len(batch.req_pool_indices)
batch.input_ids = self.draft_token
maybe_detect_oob(
batch.input_ids,
0,
batch.model_config.vocab_size,
"v2 prepare_for_verify input_ids",
)
device = batch.device
batch.out_cache_loc = assign_extend_cache_locs_func(
req_pool_indices=batch.req_pool_indices,
req_to_token=req_to_token_pool.req_to_token,
start_offset=batch.seq_lens,
end_offset=batch.seq_lens + self.draft_token_num,
batch_size=bs,
draft_token_num=self.draft_token_num,
device=device,
)
prepare_mamba_track_for_verify(batch)
# TBO's split_spec_info reads these; no-verify-sync leaves both None.
self.seq_lens_cpu = batch.seq_lens_cpu
self.seq_lens_sum = (
int(batch.seq_lens_cpu.sum())
if batch.seq_lens_cpu is not None
else None
)
# Get a forward batch
batch.forward_mode = (
ForwardMode.IDLE
if batch.forward_mode.is_idle()
else ForwardMode.TARGET_VERIFY
)
capture_mode = (
CaptureHiddenMode.NULL
if target_worker.model_runner.spec_algorithm.is_standalone()
else CaptureHiddenMode.FULL
)
batch.capture_hidden_mode = capture_mode
verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
# Run attention backend plan and cuda graph preparation
can_run_cuda_graph = bool(
target_worker.model_runner.decode_cuda_graph_runner
and target_worker.model_runner.decode_cuda_graph_runner.can_run(
verify_forward_batch
)
)
if can_run_cuda_graph:
target_worker.model_runner.decode_cuda_graph_runner.replay_prepare(
verify_forward_batch
)
verify_forward_batch.mark_forward_metadata_ready()
# Non-cuda-graph: defer init to forward_extend, which runs after
# `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
# here would use pre-pad shapes and trip DSv4 indexer shape match.
return verify_forward_batch, can_run_cuda_graph
def sample(
self: EagleVerifyInput,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
vocab_mask: torch.Tensor = None,
):
"""
Verify and find accepted tokens based on logits output and batch
(which contains spec decoding information).
"""
device = batch.device
if batch.forward_mode.is_idle():
predict = torch.empty(0, dtype=torch.int32, device=device)
num_correct_drafts = torch.empty(0, dtype=torch.int32, device=device)
accept_index = torch.empty(0, dtype=torch.int32, device=device)
return predict, num_correct_drafts, accept_index
bs = len(batch.seq_lens)
sampling_info = batch.sampling_info
next_token_logits = logits_output.next_token_logits
sanitize_nan_logits(next_token_logits, "verify: target model logits")
# Apply penalty
# This is a relaxed version of penalties for speculative decoding.
if sampling_info.acc_additive_penalties is not None:
next_token_logits.add_(
torch.repeat_interleave(
sampling_info.acc_additive_penalties, self.draft_token_num, dim=0
)
)
if sampling_info.acc_scaling_penalties is not None:
apply_scaling_penalties(
next_token_logits,
torch.repeat_interleave(
sampling_info.acc_scaling_penalties, self.draft_token_num, dim=0
),
)
if sampling_info.logit_bias is not None:
next_token_logits.add_(
torch.repeat_interleave(
sampling_info.logit_bias, self.draft_token_num, dim=0
)
)
# Apply grammar mask if provided
if vocab_mask is not None:
assert self.grammar is not None
self.grammar.apply_vocab_mask(
logits=next_token_logits, vocab_mask=vocab_mask
)
candidates = self.draft_token.reshape(bs, self.draft_token_num)
predict_shape = list(next_token_logits.shape)[:-1]
predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
accept_index = torch.full(
(bs, self.max_tree_depth), -1, dtype=torch.int32, device=device
)
num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
# Sample tokens
if sampling_info.is_all_greedy or _is_npu or _is_hip:
target_predict = torch.argmax(next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
retrieve_index=self.retrieve_index,
retrieve_next_token=self.retrieve_next_token,
retrieve_next_sibling=self.retrieve_next_sibling,
target_predict=target_predict,
topk=self.tree_topk,
)
else:
# Apply temperature and get target probs
expanded_temperature = torch.repeat_interleave(
sampling_info.temperatures, self.draft_token_num, dim=0
) # (bs * num_draft_tokens, 1)
target_probs = F.softmax(
next_token_logits / expanded_temperature, dim=-1
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ks, self.draft_token_num, dim=0
),
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, self.draft_token_num, dim=0
),
)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
target_probs = target_probs.reshape(bs, self.draft_token_num, -1)
draft_probs = torch.zeros_like(target_probs)
# coins for rejection sampling
coins = torch.rand_like(candidates, dtype=torch.float32, device=device)
# coins for final sampling
coins_for_final_sampling = torch.rand(
(bs,), dtype=torch.float32, device=device
)
tree_speculative_sampling_target_only(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
retrive_index=self.retrieve_index,
retrive_next_token=self.retrieve_next_token,
retrive_next_sibling=self.retrieve_next_sibling,
uniform_samples=coins,
uniform_samples_for_final_sampling=coins_for_final_sampling,
target_probs=target_probs,
draft_probs=draft_probs,
threshold_single=get_global_server_args().speculative_accept_threshold_single,
threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
deterministic=True,
)
# Sync sampling results across TP ranks: different GPUs may
# produce slightly different target_probs due to floating-point
# non-determinism in softmax/top_k/top_p, causing different
# sampled tokens. Broadcast from rank 0 to ensure consistency.
tp_group = (
get_attention_tp_group()
if is_dp_attention_enabled()
else get_tp_group()
)
if tp_group.world_size > 1:
tp_group.broadcast(predict, src=0)
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
if SIMULATE_ACC_LEN > 0:
# Do simulation. The helper builds (and returns) a replacement
# accept_index of width spec_steps + 1, so pass max_tree_depth - 1
# to keep the simulated width identical to the real one.
accept_index = generate_simulated_accept_index(
accept_index=accept_index,
predict=predict, # mutable
num_correct_drafts=num_correct_drafts, # mutable
simulate_acc_len=SIMULATE_ACC_LEN,
bs=bs,
spec_steps=self.max_tree_depth - 1,
)
# `num_correct_drafts` stays drafts-only inside this function; the returned
# tensor includes the trailing/bonus token via out-of-place +1 so the
# name no longer flips semantics mid-function (naming doc C2).
return predict, num_correct_drafts + 1, accept_index
python/sglang/srt/speculative/eagle_utils.py
+261
View File
@@ -10,8 +10,12 @@ from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
from sglang.srt.utils.async_probe import maybe_detect_oob
if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
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_info import EagleVerifyInput
_is_cuda = is_cuda()
_is_hip = is_hip()
@@ -272,3 +276,260 @@ def get_draft_hidden_dim(model_runner: ModelRunner) -> int:
num_aux = max(len(layer_ids), 1)
return base * num_aux
return model_runner.model_config.spec_hidden_size
def eagle_prepare_for_verify(
verify_input: EagleVerifyInput,
req_to_token_pool: ReqToTokenPool,
batch: ScheduleBatch,
target_worker: TpModelWorker,
):
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify
from sglang.srt.speculative.triton_ops.cache_locs import (
assign_extend_cache_locs_func,
)
if not batch.forward_mode.is_idle():
# Assign cache locations
bs = len(batch.req_pool_indices)
batch.input_ids = verify_input.draft_token
maybe_detect_oob(
batch.input_ids,
0,
batch.model_config.vocab_size,
"v2 prepare_for_verify input_ids",
)
device = batch.device
batch.out_cache_loc = assign_extend_cache_locs_func(
req_pool_indices=batch.req_pool_indices,
req_to_token=req_to_token_pool.req_to_token,
start_offset=batch.seq_lens,
end_offset=batch.seq_lens + verify_input.draft_token_num,
batch_size=bs,
draft_token_num=verify_input.draft_token_num,
device=device,
)
prepare_mamba_track_for_verify(batch)
# TBO's split_spec_info reads these; no-verify-sync leaves both None.
verify_input.seq_lens_cpu = batch.seq_lens_cpu
verify_input.seq_lens_sum = (
int(batch.seq_lens_cpu.sum()) if batch.seq_lens_cpu is not None else None
)
# Get a forward batch
batch.forward_mode = (
ForwardMode.IDLE if batch.forward_mode.is_idle() else ForwardMode.TARGET_VERIFY
)
capture_mode = (
CaptureHiddenMode.NULL
if target_worker.model_runner.spec_algorithm.is_standalone()
else CaptureHiddenMode.FULL
)
batch.capture_hidden_mode = capture_mode
verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)
# Run attention backend plan and cuda graph preparation
can_run_cuda_graph = bool(
target_worker.model_runner.decode_cuda_graph_runner
and target_worker.model_runner.decode_cuda_graph_runner.can_run(
verify_forward_batch
)
)
if can_run_cuda_graph:
target_worker.model_runner.decode_cuda_graph_runner.replay_prepare(
verify_forward_batch
)
verify_forward_batch.mark_forward_metadata_ready()
# Non-cuda-graph: defer init to forward_extend, which runs after
# `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing
# here would use pre-pad shapes and trip DSv4 indexer shape match.
return verify_forward_batch, can_run_cuda_graph
def eagle_sample(
verify_input: EagleVerifyInput,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
vocab_mask: torch.Tensor = None,
):
"""
Verify and find accepted tokens based on logits output and batch
(which contains spec decoding information).
"""
import torch.nn.functional as F
from sglang.srt.distributed import get_tp_group
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
is_dp_attention_enabled,
)
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.spec_utils import (
SIMULATE_ACC_LEN,
generate_simulated_accept_index,
)
from sglang.srt.utils.async_probe import maybe_detect_nan, sanitize_nan_logits
device = batch.device
if batch.forward_mode.is_idle():
predict = torch.empty(0, dtype=torch.int32, device=device)
num_correct_drafts = torch.empty(0, dtype=torch.int32, device=device)
accept_index = torch.empty(0, dtype=torch.int32, device=device)
return predict, num_correct_drafts, accept_index
bs = len(batch.seq_lens)
sampling_info = batch.sampling_info
next_token_logits = logits_output.next_token_logits
sanitize_nan_logits(next_token_logits, "verify: target model logits")
# Apply penalty
# This is a relaxed version of penalties for speculative decoding.
if sampling_info.acc_additive_penalties is not None:
next_token_logits.add_(
torch.repeat_interleave(
sampling_info.acc_additive_penalties,
verify_input.draft_token_num,
dim=0,
)
)
if sampling_info.acc_scaling_penalties is not None:
apply_scaling_penalties(
next_token_logits,
torch.repeat_interleave(
sampling_info.acc_scaling_penalties, verify_input.draft_token_num, dim=0
),
)
if sampling_info.logit_bias is not None:
next_token_logits.add_(
torch.repeat_interleave(
sampling_info.logit_bias, verify_input.draft_token_num, dim=0
)
)
# Apply grammar mask if provided
if vocab_mask is not None:
assert verify_input.grammar is not None
verify_input.grammar.apply_vocab_mask(
logits=next_token_logits, vocab_mask=vocab_mask
)
candidates = verify_input.draft_token.reshape(bs, verify_input.draft_token_num)
predict_shape = list(next_token_logits.shape)[:-1]
predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()
accept_index = torch.full(
(bs, verify_input.max_tree_depth), -1, dtype=torch.int32, device=device
)
num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=device)
# Sample tokens
if sampling_info.is_all_greedy or _is_npu or _is_hip:
target_predict = torch.argmax(next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
retrieve_index=verify_input.retrieve_index,
retrieve_next_token=verify_input.retrieve_next_token,
retrieve_next_sibling=verify_input.retrieve_next_sibling,
target_predict=target_predict,
topk=verify_input.tree_topk,
)
else:
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
# Apply temperature and get target probs
expanded_temperature = torch.repeat_interleave(
sampling_info.temperatures, verify_input.draft_token_num, dim=0
) # (bs * num_draft_tokens, 1)
target_probs = F.softmax(
next_token_logits / expanded_temperature, dim=-1
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ks, verify_input.draft_token_num, dim=0
),
) # (bs * num_draft_tokens, vocab_size)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, verify_input.draft_token_num, dim=0
),
)
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1)
draft_probs = torch.zeros_like(target_probs)
# coins for rejection sampling
coins = torch.rand_like(candidates, dtype=torch.float32, device=device)
# coins for final sampling
coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device)
tree_speculative_sampling_target_only(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
retrive_index=verify_input.retrieve_index,
retrive_next_token=verify_input.retrieve_next_token,
retrive_next_sibling=verify_input.retrieve_next_sibling,
uniform_samples=coins,
uniform_samples_for_final_sampling=coins_for_final_sampling,
target_probs=target_probs,
draft_probs=draft_probs,
threshold_single=get_global_server_args().speculative_accept_threshold_single,
threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
deterministic=True,
)
# Sync sampling results across TP ranks: different GPUs may
# produce slightly different target_probs due to floating-point
# non-determinism in softmax/top_k/top_p, causing different
# sampled tokens. Broadcast from rank 0 to ensure consistency.
tp_group = (
get_attention_tp_group() if is_dp_attention_enabled() else get_tp_group()
)
if tp_group.world_size > 1:
tp_group.broadcast(predict, src=0)
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
if SIMULATE_ACC_LEN > 0:
# Do simulation. The helper builds (and returns) a replacement
# accept_index of width spec_steps + 1, so pass max_tree_depth - 1
# to keep the simulated width identical to the real one.
accept_index = generate_simulated_accept_index(
accept_index=accept_index,
predict=predict, # mutable
num_correct_drafts=num_correct_drafts, # mutable
simulate_acc_len=SIMULATE_ACC_LEN,
bs=bs,
spec_steps=verify_input.max_tree_depth - 1,
)
# `num_correct_drafts` stays drafts-only inside this function; the returned
# tensor includes the trailing/bonus token via out-of-place +1 so the
# name no longer flips semantics mid-function (naming doc C2).
return predict, num_correct_drafts + 1, accept_index
python/sglang/srt/speculative/eagle_worker_v2.py
+6 -3
View File
@@ -65,6 +65,8 @@ from sglang.srt.speculative.eagle_utils import (
TreeMaskMode,
_eagle_prefill_tail_tokens,
build_tree_kernel_efficient,
eagle_prepare_for_verify,
eagle_sample,
organize_draft_results,
per_step_draft_out_cache_loc,
)
@@ -1215,7 +1217,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
# Batch 1: Target verify
# Prepare for target verify in a separate stream
with self.plan_stream_ctx:
verify_forward_batch, can_run_cuda_graph = verify_input.prepare_for_verify(
verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify(
verify_input,
self.req_to_token_pool,
batch,
self.target_worker,
@@ -1264,7 +1267,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
# Run target verify batch in the main compute stream (GPU compute).
# Metadata init is skipped iff cuda-graph already ran replay_prepare —
# prepare_for_verify marked the batch in exactly that case; the
# eagle_prepare_for_verify marked the batch in exactly that case; the
# non-cuda-graph path stays unmarked and gets forward_extend's init
# (post-pad).
forward_batch_output = self.target_worker.forward_batch_generation(
@@ -1301,7 +1304,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
predict,
accept_lens,
accept_index,
) = verify_input.sample(batch, logits_output, vocab_mask)
) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
new_seq_lens = batch.seq_lens + accept_lens
# Update mamba state for hybrid GDN models after verification
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+10 -4
View File
@@ -49,7 +49,12 @@ from sglang.srt.speculative.eagle_info import (
EagleVerifyInput,
)
from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient
from sglang.srt.speculative.eagle_utils import (
TreeMaskMode,
build_tree_kernel_efficient,
eagle_prepare_for_verify,
eagle_sample,
)
from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
)
@@ -807,7 +812,8 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
# Batch 1: Target verify
# Prepare for target verify in a separate stream
with self.plan_stream_ctx:
verify_forward_batch, can_run_cuda_graph = verify_input.prepare_for_verify(
verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify(
verify_input,
self.req_to_token_pool,
batch,
self.target_worker,
@@ -834,7 +840,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
),
)
# NOTE: metadata init is skipped here unconditionally, although
# prepare_for_verify only plans when cuda-graph replay_prepare ran.
# eagle_prepare_for_verify only plans when cuda-graph replay_prepare ran.
# eagle_worker_v2 re-inits the non-graph path instead (post-pad); this
# worker has not adopted that fix, so preserve its behavior verbatim.
# On NPU with --disable-cuda-graph, non-graph verify needs metadata init
@@ -856,7 +862,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
predict,
accept_lens,
accept_index,
) = verify_input.sample(batch, logits_output)
) = eagle_sample(verify_input, batch, logits_output)
new_seq_lens = batch.seq_lens + accept_lens
if not batch.forward_mode.is_idle():
python/sglang/srt/speculative/ngram_info.py
+2 -5
View File
@@ -6,14 +6,11 @@ import torch
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin,
)
from sglang.srt.speculative.eagle_info_v2 import EagleDraftInputV2Mixin
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
class NgramVerifyInput(SpecInput, EagleDraftInputV2Mixin, EagleVerifyInputV2Mixin):
class NgramVerifyInput(SpecInput, EagleDraftInputV2Mixin):
def __init__(
self,
draft_token: torch.Tensor = None,
python/sglang/srt/speculative/ngram_worker.py
+2 -1
View File
@@ -14,6 +14,7 @@ from sglang.srt.observability.req_time_stats import set_time_batch
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker, EagleDraftWorkerBase
from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
from sglang.srt.speculative.eagle_utils import eagle_sample
from sglang.srt.speculative.ngram_info import NgramVerifyInput
from sglang.srt.speculative.spec_utils import (
commit_mamba_states_after_verify,
@@ -432,7 +433,7 @@ class NGRAMWorker(BaseSpecWorker):
predict,
accept_lens,
accept_index,
) = verify_input.sample(batch, logits_output, vocab_mask)
) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
new_seq_lens = batch.seq_lens + accept_lens
commit_mamba_states_after_verify(
self.target_worker,