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

[ROCm] Fix EAGLE spec-decode verify silently sampling greedy on HIP (#37134)

Browse Source 
Co-authored-by: JohnQinAMD <yanyuan.qin@amd.com>
Co-authored-by: Jiejing Zhang <jiejing.zhang@amd.com>
This commit is contained in:
xiaobochen-amd
2026-09-15 18:49:49 -07:00
committed by GitHub
co-authored by JohnQinAMD Jiejing Zhang
parent 2f7d04da5d
commit 7eedd57ab0
7 changed files with 298 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/speculative/reject_sampling.py
+15 -5
View File
@@ -68,7 +68,15 @@ def speculative_sampling_classic_kernel(
coin = tl.load(uni_ptr_base + (step - 1) * stride_uni_s)
if coin * q < p:
# X was sampled from q, so q(X) has to be a positive probability.
# Anything else means this row is not the distribution X came from, and
# `coin * q < p` would then accept unconditionally -- -inf < p for an
# -inf q, 0 < p for a zero one, and the range guard the residual passes
# use lets zero through. Reject instead: the residual path resamples
# from the target, which is the safe direction to fail in.
q_is_prob = (q > 0.0) & (q <= 1.0)
if q_is_prob & (coin * q < p):
num_accept += 1
cur_prob_row = step
tl.store(Predicts + last_accepted_global_idx, draft_token)
@@ -111,8 +119,10 @@ def speculative_sampling_classic_kernel(
else:
q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v
q_val = tl.load(q_ptr, mask=mask, other=0.0)
# Treat NaN q (degenerate draft rows) as 0: residual falls back to p.
q_val = tl.where(q_val == q_val, q_val, 0.0)
# Treat any non-probability q (NaN, +-inf, negative) as 0: the
# residual falls back to p. A comparison against NaN is false, so
# the range test rejects it along with the infinities.
q_val = tl.where((q_val >= 0.0) & (q_val <= 1.0), q_val, 0.0)
diff = p_val - q_val
val = tl.where(diff > 0.0, diff, 0.0)
@@ -139,8 +149,8 @@ def speculative_sampling_classic_kernel(
else:
q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v
q_val = tl.load(q_ptr, mask=mask, other=0.0)
# Same NaN-q guard as pass 1.
q_val = tl.where(q_val == q_val, q_val, 0.0)
# Same guard as pass 1.
q_val = tl.where((q_val >= 0.0) & (q_val <= 1.0), q_val, 0.0)
diff = p_val - q_val
val = tl.where(diff > 0.0, diff, 0.0)
python/sglang/srt/arg_groups/speculative_hook.py
+59 -2
View File
@@ -24,6 +24,37 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _should_auto_enable_hip_rejection_sampling(
*,
is_hip: bool,
use_rejection_sampling: bool,
algorithm: Optional[str],
token_map: Optional[str],
eagle_topk: int,
accept_threshold_single: float,
accept_threshold_acc: float,
enable_deterministic_inference: bool,
) -> bool:
"""Whether HIP may default ``speculative_use_rejection_sampling`` on.
Rejection sampling still cannot consume a reduced / hot draft vocab
(``eagle_worker_v2`` FIXME: scatter via the d2t map). Auto-enabling there
would crash configs that previously ran greedy on HIP, including EAGLE3
stage-a ``test_basic_sanity_eagle3`` (draft 32000 vs target 128256). Skip
EAGLE3 and any EAGLE run that already has a token map.
"""
return (
is_hip
and not use_rejection_sampling
and algorithm == "EAGLE"
and token_map is None
and eagle_topk == 1
and accept_threshold_single == 1.0
and accept_threshold_acc == 1.0
and not enable_deterministic_inference
)
def _disable_overlap_schedule_for_cpu(server_args: ServerArgs) -> None:
cfg = resolving_view(server_args)
if cfg.device != "cpu" or cfg.disable_overlap_schedule:
@@ -813,7 +844,6 @@ def _handle_frozen_kv_mtp(server_args: ServerArgs) -> None:
def _handle_eagle_family(server_args: ServerArgs) -> None:
cfg = resolving_view(server_args)
if (
@@ -920,7 +950,34 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
"trtllm_mha backend only supports topk = 1 for speculative decoding."
)
if cfg.speculative_use_rejection_sampling:
# ROCm/HIP has no CUDA/MUSA sampling-verify kernels, so EAGLE verify would
# otherwise fall back to greedy (argmax) and silently ignore temperature and
# top_p. Default rejection sampling on -- it routes verify through the Triton
# chain sampler -- for configs that support it. See
# _should_auto_enable_hip_rejection_sampling for the cases we must not flip.
if _should_auto_enable_hip_rejection_sampling(
is_hip=get_platform().is_hip,
use_rejection_sampling=cfg.speculative_use_rejection_sampling,
algorithm=cfg.speculative_algorithm,
token_map=cfg.speculative_token_map,
eagle_topk=cfg.speculative_eagle_topk,
accept_threshold_single=cfg.speculative_accept_threshold_single,
accept_threshold_acc=cfg.speculative_accept_threshold_acc,
enable_deterministic_inference=cfg.enable_deterministic_inference,
):
declare_resolution(
server_args,
"_handle_eagle_family",
speculative_use_rejection_sampling=True,
)
logger.info(
"ROCm needs rejection sampling for EAGLE spec-decode to sample at all; "
"enabling speculative_use_rejection_sampling by default."
)
# resolved_view, not cfg: the block above may have just decided this field,
# and declare_resolution writes to the stash rather than the dataclass.
if resolved_view(server_args).speculative_use_rejection_sampling:
# Resolved alias by now: NEXTN -> EAGLE, Gemma4 draft -> FROZEN_KV_MTP.
# Only the EAGLE/EAGLE3 draft workers emit a target-vocab proposal that
# the rejection-sampling kernel consumes; everything else (STANDALONE,
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+11 -1
View File
@@ -42,6 +42,7 @@ from sglang.srt.runtime_context import (
get_spec,
)
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_params import TOP_K_ALL
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req
@@ -213,6 +214,14 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
)
self.temperatures = torch.ones((self.max_bs, 1), dtype=torch.float)
# Real per-request top_k, for the same reason temperatures are
# carried: the draft proposal cannot tell a greedy request from a
# T=1 one by temperature alone, because SamplingParams rewrites
# temperature 0 to temperature=1.0 with top_k=1.
# TOP_K_ALL, not -1: -1 is not a top_k this pipeline ever carries
# (SamplingParams rewrites it), and it would read as top_k <= 1, i.e.
# greedy, for the padded rows and for a run that never copies in.
self.top_ks = torch.full((self.max_bs,), TOP_K_ALL, dtype=torch.int32)
if self.require_gathered_buffer:
if self.require_mlp_tp_gather:
@@ -417,7 +426,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
sampling_info = SamplingBatchInfo(
temperatures=self.temperatures[:num_seqs],
top_ps=torch.ones((num_seqs,), dtype=torch.float),
top_ks=torch.full((num_seqs,), -1, dtype=torch.int32),
top_ks=self.top_ks[:num_seqs],
min_ps=torch.zeros((num_seqs,), dtype=torch.float),
is_all_greedy=False,
is_any_greedy=False,
@@ -624,6 +633,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
self.temperatures[:raw_bs].copy_(
forward_batch.sampling_info.temperatures[:raw_bs]
)
self.top_ks[:raw_bs].copy_(forward_batch.sampling_info.top_ks[:raw_bs])
# TODO(ch-wan): support num_token_non_padded
if self.require_gathered_buffer:
python/sglang/srt/speculative/eagle_utils.py
+45 -13
View File
@@ -685,6 +685,24 @@ def _verify_coins(
return coins, coins_for_final_sampling
def _verify_uses_greedy(
*,
is_all_greedy: bool,
is_cpu: bool,
is_hip: bool,
is_xpu: bool,
use_rejection_sampling: bool,
) -> bool:
"""Whether EAGLE verify must commit argmax instead of taking the sampling path.
HIP has no CUDA/MUSA sampling-verify kernels, so it used to be listed here
unconditionally. Rejection sampling routes it through the pure-Triton chain
sampler instead, so only a HIP run without that still has to go greedy. Every
other platform reduces to the original predicate.
"""
return is_all_greedy or is_cpu or is_xpu or (is_hip and not use_rejection_sampling)
def _can_use_sparse_uno_tree_target_sampling(
max_top_k: Optional[int],
sampling_info: SamplingBatchInfo,
@@ -781,7 +799,14 @@ def eagle_sample(
# Sample tokens
target_predict = None
if sampling_info.is_all_greedy or _is_cpu or _is_hip or _is_xpu:
use_rejection_sampling = get_spec().speculative_use_rejection_sampling
if _verify_uses_greedy(
is_all_greedy=sampling_info.is_all_greedy,
is_cpu=_is_cpu,
is_hip=_is_hip,
is_xpu=_is_xpu,
use_rejection_sampling=use_rejection_sampling,
):
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(
@@ -855,23 +880,30 @@ def eagle_sample(
tree_speculative_sampling_target_only,
)
else:
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
from sglang.kernels.ops.speculative.reject_sampling import (
chain_speculative_sampling_triton,
)
use_rejection_sampling = get_spec().speculative_use_rejection_sampling
# if/else, not a ternary: the CUDA-only name still has to resolve in the
# branch not taken, and HIP only reaches here with rejection sampling on.
if use_rejection_sampling:
sampling_fn = chain_speculative_sampling_triton
else:
if not _is_npu:
from sgl_kernel import tree_speculative_sampling_target_only
sampling_fn = (
chain_speculative_sampling_triton
if use_rejection_sampling
else tree_speculative_sampling_target_only
)
sampling_fn = tree_speculative_sampling_target_only
if _is_hip:
# Same names, same contract: dflash_utils.py aliases these too.
from sglang.kernels.ops.sampling.renorm_triton import (
top_k_renorm_probs_triton as top_k_renorm_prob,
)
from sglang.kernels.ops.sampling.renorm_triton import (
top_p_renorm_probs_triton as top_p_renorm_prob,
)
elif not _is_npu:
from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob
expanded_temperature = torch.repeat_interleave(
sampling_info.temperatures, verify_input.draft_token_num, dim=0
python/sglang/srt/speculative/eagle_worker_v2.py
+9
View File
@@ -628,6 +628,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
parent_list, top_scores_index, draft_tokens, draft_probs = (
self.cuda_graph_runner.execute(forward_batch)
)
if draft_probs is not None:
# draft_probs is the one graph output read after the target
# forward rather than by it, and it points into the graph's
# private memory pool. The pool recycles that block in the
# meantime -- in practice the DSA top-k mask lands there and
# eagle_sample sees -inf. Copy out at the boundary.
draft_probs = draft_probs.clone()
else:
if (
not forward_batch.forward_mode.is_idle()
@@ -767,6 +774,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
probs, topk_p, topk_index = sample_draft_proposal(
logits_output.next_token_logits,
forward_batch.sampling_info.temperatures,
forward_batch.sampling_info.top_ks,
)
draft_probs_list.append(probs)
forward_batch.positions.add_(1)
@@ -1117,6 +1125,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
ret_draft_probs, ret_topk_p, ret_topk_index = sample_draft_proposal(
draft_logits_output.next_token_logits,
batch.sampling_info.temperatures,
batch.sampling_info.top_ks,
)
elif self.topk == 1 and not _is_hip:
# Gated to CUDA: see #26358 — ROCm's argmax tie-break corrupts
python/sglang/srt/speculative/spec_utils.py
+32 -1
View File
@@ -165,15 +165,46 @@ def renorm_draft_probs(
return torch.softmax(next_token_logits / sampling_info.temperatures, dim=-1)
def sample_draft_proposal(next_token_logits: torch.Tensor, temperatures: torch.Tensor):
def sample_draft_proposal(
next_token_logits: torch.Tensor,
temperatures: torch.Tensor,
top_ks: Optional[torch.Tensor] = None,
):
"""Leviathan draft proposal: q = softmax(logits / T), X ~ q.
Returns (q, q(X), X). The verify's accept test coin*q(X) < p(X) is unbiased
only if q is exactly the distribution X was drawn from, so callers must hand
the returned q (not a recomputed one) to the verify.
A greedy row (``top_k == 1``) proposes its argmax instead. SamplingParams
rewrites temperature 0 to ``temperature=1.0, top_k=1``, so T alone cannot
tell a greedy request from a T=1 one, and sampling a sharp-but-not-
degenerate distribution proposes a non-argmax token often enough to cost
real accept length.
That row's X is then not drawn from the q returned beside it, which the
unbiasedness argument above otherwise rests on. It stays correct because
eagle_sample renormalises the target by the same per-row ``top_ks`` before
the accept test, so a greedy row's p is one-hot: X equal to the target
argmax accepts (p(X) = 1), any other X rejects (p(X) = 0) and the residual
(p - q)+ it resamples from is p itself. Both arms commit the target argmax,
which is what greedy means. Drop that renorm and this stops holding.
"""
probs = torch.softmax(next_token_logits / temperatures, dim=-1)
topk_p, topk_index = fast_sample(probs, num_samples=1)
if top_ks is not None:
# Assert rather than skip on a device mismatch: a host-side top_ks would
# make this correction silently vanish, and the symptom -- draft accept
# length quietly dropping about 20% -- reads as a model problem, not a
# plumbing one.
assert top_ks.device == probs.device, (
f"top_ks must be on {probs.device} to reach the draft proposal, "
f"got {top_ks.device}; the caller has to carry the real per-request "
"top_k, not a host placeholder"
)
greedy = (top_ks <= 1).view(-1, 1)
topk_index = torch.where(greedy, probs.argmax(dim=-1, keepdim=True), topk_index)
topk_p = probs.gather(1, topk_index)
return probs, topk_p, topk_index