From ea9690440af972b96e7f1580f065728bd070948a Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Tue, 14 Jul 2026 12:03:19 -0500 Subject: [PATCH 1/2] [ROCm] Fix EAGLE spec-decode verify silently sampling greedy on HIP The sgl_kernel sampling-verify ops (tree_speculative_sampling_target_only, top_k/top_p_renorm_prob) are CUDA/MUSA-only, so HIP fell through eagle_sample's gate to argmax(greedy), silently ignoring temperature/top_p and degenerating into repetition loops at temp>0. Route HIP through the pure-Triton chain sampler (reject_sampling.py) when rejection sampling is active and the batch isn't all-greedy, with a Triton pivot-search top-p renorm (_renorm_top_k_top_p_hip) replacing the CUDA-only renorm ops. The gate is factored into a pure helper (_verify_uses_greedy) that reduces to the original predicate byte-for-byte on every non-HIP platform; CUDA-only imports are guarded so HIP never ImportErrors. top_p>=1.0 is an exact no-op to stay bit-exact with CUDA's top_p_renorm_prob(p, 1.0)==p. Tests: gate-routing truth table (CPU) and renorm-vs-nucleus numerical equivalence (GPU). Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/srt/speculative/eagle_utils.py | 185 +++++++++++++++--- .../spec/eagle/test_rocm_eagle_renorm.py | 129 ++++++++++++ .../unit/spec/test_eagle_gate_routing.py | 91 +++++++++ 3 files changed, 378 insertions(+), 27 deletions(-) create mode 100644 test/registered/spec/eagle/test_rocm_eagle_renorm.py create mode 100644 test/registered/unit/spec/test_eagle_gate_routing.py diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 28cf4610c791..1b34612a89c8 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, List, Optional, Tuple import torch +import triton +import triton.language as tl from sglang.kernels.ops.speculative.spec_tree import ( sgl_build_tree_kernel_efficient_triton, @@ -650,6 +652,105 @@ def _verify_coins( return coins, coins_for_final_sampling +@triton.jit +def _top_p_renorm_kernel( + probs_ptr, + out_ptr, + top_p_ptr, + vocab: tl.constexpr, + BLOCK: tl.constexpr, + N_ITER: tl.constexpr, +): + # Per-row top-p renorm by pivot binary-search (no sort): keep the largest tau with + # sum(p >= tau) >= top_p, then renormalize p >= tau. N_ITER=30 pins tau to + # max_prob*2^-30, far below realistic prob gaps, so it matches a sort to <1e-6. + row = tl.program_id(0) + base = row * vocab + tp = tl.load(top_p_ptr + row) + offs = tl.arange(0, BLOCK) + hi = 0.0 + for start in range(0, vocab, BLOCK): + idx = start + offs + m = idx < vocab + p = tl.load(probs_ptr + base + idx, mask=m, other=0.0) + hi = tl.maximum(hi, tl.max(p, axis=0)) + lo = 0.0 + for _ in range(N_ITER): + mid = (lo + hi) * 0.5 + s = 0.0 + for start in range(0, vocab, BLOCK): + idx = start + offs + m = idx < vocab + p = tl.load(probs_ptr + base + idx, mask=m, other=0.0) + s += tl.sum(tl.where(p >= mid, p, 0.0), axis=0) + take = s >= tp + lo = tl.where(take, mid, lo) + hi = tl.where(take, hi, mid) + Z = 0.0 + for start in range(0, vocab, BLOCK): + idx = start + offs + m = idx < vocab + p = tl.load(probs_ptr + base + idx, mask=m, other=0.0) + Z += tl.sum(tl.where(p >= lo, p, 0.0), axis=0) + Z = tl.maximum(Z, 1e-12) + for start in range(0, vocab, BLOCK): + idx = start + offs + m = idx < vocab + p = tl.load(probs_ptr + base + idx, mask=m, other=0.0) + tl.store(out_ptr + base + idx, tl.where(p >= lo, p / Z, 0.0), mask=m) + + +def _renorm_top_k_top_p_hip( + probs: torch.Tensor, top_ks: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """ROCm replacement for sgl_kernel top_k/top_p_renorm_prob (CUDA/MUSA-only). + + top-p runs the Triton pivot kernel; top-k uses a sort only for rows that actually + restrict it (a no-op when top_k >= vocab). Rows with top_p >= 1.0 pass through + unchanged to stay bit-exact with CUDA's top_p_renorm_prob(p, 1.0) == p. + """ + vocab = probs.shape[-1] + probs = probs.contiguous() + if bool((top_ks < vocab).any()): + sorted_probs, sorted_indices = probs.sort(dim=-1, descending=True) + order = torch.arange(vocab, device=probs.device).view(1, -1) + sorted_probs = sorted_probs.masked_fill(order >= top_ks.reshape(-1, 1), 0.0) + sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp_min( + 1e-12 + ) + probs = torch.zeros_like(sorted_probs).scatter_( + -1, sorted_indices, sorted_probs + ) + top_ps = top_ps.reshape(-1).to(torch.float32).contiguous() + out = torch.empty_like(probs) + _top_p_renorm_kernel[(probs.shape[0],)]( + probs, out, top_ps, vocab, BLOCK=4096, N_ITER=30 + ) + return torch.where((top_ps >= 1.0).view(-1, 1), probs, out) + + +def _verify_uses_greedy( + *, + is_all_greedy: bool, + is_cpu: bool, + is_npu: bool, + is_hip: bool, + is_xpu: bool, + use_rejection_sampling: bool, +) -> bool: + """Whether EAGLE verify must commit argmax(greedy) instead of the sampling path. + + HIP has no CUDA/MUSA sampling-verify kernels, so it historically forced greedy here; + route it through the pure-Triton chain sampler when rejection sampling supplies the + draft proposal and the batch isn't all-greedy. Reduces to the original predicate on + non-HIP platforms. + """ + hip_can_sample = is_hip and use_rejection_sampling and not is_all_greedy + return ( + is_all_greedy or is_cpu or is_npu or is_xpu or (is_hip and not hip_can_sample) + ) + + def eagle_sample( verify_input: EagleVerifyInput, batch: ScheduleBatch, @@ -727,7 +828,15 @@ def eagle_sample( # Sample tokens target_predict = None - if sampling_info.is_all_greedy or _is_cpu or _is_npu or _is_hip or _is_xpu: + _use_rej = get_spec().speculative_use_rejection_sampling + if _verify_uses_greedy( + is_all_greedy=sampling_info.is_all_greedy, + is_cpu=_is_cpu, + is_npu=_is_npu, + is_hip=_is_hip, + is_xpu=_is_xpu, + use_rejection_sampling=_use_rej, + ): 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( @@ -757,17 +866,21 @@ def eagle_sample( tp_group.broadcast(accept_index, src=0) tp_group.broadcast(num_correct_drafts, src=0) else: - from sgl_kernel import ( - top_k_renorm_prob, - top_p_renorm_prob, - tree_speculative_sampling_target_only, - ) + # sgl_kernel renorm/tree-verify ops are CUDA/MUSA-only; import lazily so the HIP + # branch never ImportErrors. The gate forces rejection sampling on HIP, so the + # CUDA-only tree_speculative_sampling_target_only is never referenced there. + if not _is_hip: + 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 + use_rejection_sampling = _use_rej sampling_fn = ( chain_speculative_sampling_triton @@ -787,26 +900,44 @@ def eagle_sample( next_token_logits / expanded_temperature, dim=-1 ) # (bs * num_draft_tokens, vocab_size) maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") - if sampling_info.need_top_k_sampling: - 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" - ) - if sampling_info.need_top_p_sampling: - 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" - ) + if _is_hip: + if ( + sampling_info.need_top_k_sampling + or sampling_info.need_top_p_sampling + ): + target_probs = _renorm_top_k_top_p_hip( + target_probs, + torch.repeat_interleave( + sampling_info.top_ks, verify_input.draft_token_num, dim=0 + ), + 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_k/top_p renorm" + ) + else: + if sampling_info.need_top_k_sampling: + 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" + ) + if sampling_info.need_top_p_sampling: + 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 = ( verify_input.draft_probs diff --git a/test/registered/spec/eagle/test_rocm_eagle_renorm.py b/test/registered/spec/eagle/test_rocm_eagle_renorm.py new file mode 100644 index 000000000000..2b3e2297905b --- /dev/null +++ b/test/registered/spec/eagle/test_rocm_eagle_renorm.py @@ -0,0 +1,129 @@ +"""Numerical-equivalence tests for the ROCm EAGLE top-k/top-p renorm. + +Guards ``eagle_utils._renorm_top_k_top_p_hip`` (the Triton replacement for the +CUDA/MUSA-only ``sgl_kernel`` ``top_k_renorm_prob`` + ``top_p_renorm_prob``) +against an independent torch nucleus reference. A regression in the pivot search +(wrong iteration count, wrong threshold, broken top-k, or a ``top_p == 1.0`` that +stops being an exact no-op) turns a case red. The kernel is device-agnostic, so +this runs on any triton-capable GPU (CUDA in CI); skipped without one. +""" + +import unittest + +import torch + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +# Runs on any triton GPU; AMD registration exercises the actual ROCm target path. +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd") + +try: + import triton # noqa: F401 + + _HAS_TRITON = True +except Exception: + _HAS_TRITON = False + +_RUNNABLE = torch.cuda.is_available() and _HAS_TRITON + + +def _torch_nucleus_reference(probs, top_ks, top_ps): + """top-k renorm then nucleus (top-p) renorm, mirroring sgl_kernel semantics. + + Always keeps the argmax token (like the pivot kernel and sgl_kernel), so it + agrees on the top_p -> 0 edge instead of zeroing the row. + """ + vocab = probs.shape[-1] + order = torch.arange(vocab, device=probs.device).view(1, -1) + # top-k + sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True) + sorted_probs = torch.where( + order < top_ks.view(-1, 1), sorted_probs, torch.zeros_like(sorted_probs) + ) + sorted_probs = sorted_probs / sorted_probs.sum(-1, keepdim=True).clamp_min(1e-12) + p = torch.zeros_like(probs).scatter_(-1, sorted_idx, sorted_probs) + # top-p (nucleus) + sorted_probs, sorted_idx = p.sort(dim=-1, descending=True) + csum = sorted_probs.cumsum(dim=-1) + keep = (csum - sorted_probs) < top_ps.view(-1, 1) + keep[..., 0] = True + sorted_probs = torch.where(keep, sorted_probs, torch.zeros_like(sorted_probs)) + sorted_probs = sorted_probs / sorted_probs.sum(-1, keepdim=True).clamp_min(1e-12) + return torch.zeros_like(p).scatter_(-1, sorted_idx, sorted_probs) + + +@unittest.skipUnless(_RUNNABLE, "requires a triton-capable GPU") +class TestRocmEagleRenorm(CustomTestCase): + VOCAB = 4096 + ROWS = 16 + + def setUp(self): + torch.manual_seed(0) + + def _rand_probs(self, rows=None): + rows = rows or self.ROWS + return torch.softmax(torch.randn(rows, self.VOCAB, device="cuda"), dim=-1) + + def _unrestricted_top_ks(self, rows=None): + rows = rows or self.ROWS + return torch.full((rows,), self.VOCAB, dtype=torch.int32, device="cuda") + + def test_top_p_matches_nucleus(self): + from sglang.srt.speculative.eagle_utils import _renorm_top_k_top_p_hip + + probs = self._rand_probs() + top_ks = self._unrestricted_top_ks() + for top_p in (0.5, 0.9, 0.95): + top_ps = torch.full((self.ROWS,), top_p, device="cuda") + got = _renorm_top_k_top_p_hip(probs.clone(), top_ks, top_ps) + ref = _torch_nucleus_reference(probs.clone(), top_ks, top_ps) + self.assertLess((got - ref).abs().max().item(), 1e-6, f"top_p={top_p}") + + def test_top_p_one_is_exact_noop(self): + from sglang.srt.speculative.eagle_utils import _renorm_top_k_top_p_hip + + probs = self._rand_probs() + top_ks = self._unrestricted_top_ks() + top_ps = torch.ones(self.ROWS, device="cuda") + got = _renorm_top_k_top_p_hip(probs.clone(), top_ks, top_ps) + self.assertEqual((got - probs).abs().max().item(), 0.0) + + def test_top_k_restriction(self): + from sglang.srt.speculative.eagle_utils import _renorm_top_k_top_p_hip + + probs = self._rand_probs() + top_ks = torch.full((self.ROWS,), 8, dtype=torch.int32, device="cuda") + top_ps = torch.ones(self.ROWS, device="cuda") # top_p no-op -> isolate top-k + got = _renorm_top_k_top_p_hip(probs.clone(), top_ks, top_ps) + self.assertTrue(((got > 0).sum(-1) == 8).all().item()) + ref = _torch_nucleus_reference(probs.clone(), top_ks, top_ps) + self.assertLess((got - ref).abs().max().item(), 1e-6) + + def test_edge_top_p_to_zero_keeps_top1(self): + from sglang.srt.speculative.eagle_utils import _renorm_top_k_top_p_hip + + probs = self._rand_probs(rows=1) + top_ks = self._unrestricted_top_ks(rows=1) + top_ps = torch.full((1,), 1e-4, device="cuda") + got = _renorm_top_k_top_p_hip(probs.clone(), top_ks, top_ps) + self.assertEqual((got > 0).sum().item(), 1) + self.assertEqual(got.argmax().item(), probs.argmax().item()) + + def test_edge_single_nonmasked_token(self): + from sglang.srt.speculative.eagle_utils import _renorm_top_k_top_p_hip + + one_hot = torch.zeros(1, self.VOCAB, device="cuda") + one_hot[0, 123] = 1.0 + got = _renorm_top_k_top_p_hip( + one_hot.clone(), + self._unrestricted_top_ks(rows=1), + torch.full((1,), 0.9, device="cuda"), + ) + self.assertEqual((got > 0).sum().item(), 1) + self.assertAlmostEqual(got[0, 123].item(), 1.0, places=5) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_eagle_gate_routing.py b/test/registered/unit/spec/test_eagle_gate_routing.py new file mode 100644 index 000000000000..6a0565ca5713 --- /dev/null +++ b/test/registered/unit/spec/test_eagle_gate_routing.py @@ -0,0 +1,91 @@ +"""Gate-routing tests for EAGLE verify: sampling vs. greedy (argmax). + +Guards the correctness contract of the ROCm fix in +``eagle_utils._verify_uses_greedy``: on every non-HIP platform the gate must +reduce byte-for-byte to the pre-patch predicate +(``is_all_greedy or is_cpu or is_npu or is_hip or is_xpu``), and HIP may take the +sampling path only when rejection sampling is on and the batch isn't all-greedy. +A regression that forced greedy on CUDA, or that let HIP sample without rejection +sampling, would turn a case here red. Pure-boolean logic, so it runs on CPU CI. +""" + +import itertools +import unittest + +from sglang.srt.speculative.eagle_utils import _verify_uses_greedy +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _pre_patch_gate(is_all_greedy, is_cpu, is_npu, is_hip, is_xpu): + # eagle_utils.py gate before this PR. Non-HIP behavior must match this exactly. + return is_all_greedy or is_cpu or is_npu or is_hip or is_xpu + + +# (name, is_cpu, is_npu, is_hip, is_xpu); CUDA == no platform flag set. +_PLATFORMS = { + "cuda": (False, False, False, False), + "hip": (False, False, True, False), + "cpu": (True, False, False, False), + "npu": (False, True, False, False), + "xpu": (False, False, False, True), +} + + +class TestEagleGateRouting(CustomTestCase): + def _gate(self, is_all_greedy, platform, use_rej): + is_cpu, is_npu, is_hip, is_xpu = _PLATFORMS[platform] + return _verify_uses_greedy( + is_all_greedy=is_all_greedy, + is_cpu=is_cpu, + is_npu=is_npu, + is_hip=is_hip, + is_xpu=is_xpu, + use_rejection_sampling=use_rej, + ) + + def test_non_hip_is_byte_identical_to_pre_patch(self): + for platform, is_all_greedy, use_rej in itertools.product( + _PLATFORMS, (False, True), (False, True) + ): + is_cpu, is_npu, is_hip, is_xpu = _PLATFORMS[platform] + if is_hip: + continue + got = self._gate(is_all_greedy, platform, use_rej) + expected = _pre_patch_gate(is_all_greedy, is_cpu, is_npu, is_hip, is_xpu) + self.assertEqual( + got, + expected, + f"{platform} greedy={is_all_greedy} rej={use_rej}: gate diverged " + f"from pre-patch ({got} != {expected})", + ) + + def test_hip_samples_only_with_rejection_and_non_greedy(self): + # The single new sampling entry: HIP + rejection sampling + not all-greedy. + self.assertFalse(self._gate(False, "hip", True)) + # Every other HIP combination still commits greedy (argmax). + self.assertTrue(self._gate(True, "hip", True)) + self.assertTrue(self._gate(True, "hip", False)) + self.assertTrue(self._gate(False, "hip", False)) + + def test_cuda_keyed_on_all_greedy_only(self): + # CUDA samples whenever the batch isn't all-greedy, regardless of the flag. + self.assertFalse(self._gate(False, "cuda", False)) + self.assertFalse(self._gate(False, "cuda", True)) + self.assertTrue(self._gate(True, "cuda", False)) + self.assertTrue(self._gate(True, "cuda", True)) + + def test_cpu_npu_xpu_always_greedy(self): + for platform in ("cpu", "npu", "xpu"): + for is_all_greedy in (False, True): + for use_rej in (False, True): + self.assertTrue( + self._gate(is_all_greedy, platform, use_rej), + f"{platform} must force greedy", + ) + + +if __name__ == "__main__": + unittest.main() From 17184fd60d9c10356ac16f88114e895035a53cbe Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Tue, 14 Jul 2026 12:03:19 -0500 Subject: [PATCH 2/2] [ROCm] Default rejection sampling on for HIP EAGLE verify Without this, the ROCm sampling fix only helps HIP users who pass --speculative-use-rejection-sampling; everyone else still silently gets argmax verify. Default the flag on in _handle_eagle_family when running on HIP with an EAGLE/EAGLE3 topk=1 config whose options are compatible with rejection sampling (default accept thresholds, non-deterministic), unless the user opted in explicitly. The conditions mirror the validation below so the flip never turns a previously-working config into a raise, and CUDA behavior is untouched. Setting the flag makes the draft worker emit the target-vocab proposal (draft_probs) that the Triton chain sampler consumes. Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/srt/arg_groups/speculative_hook.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 05cb2cbe517a..5146ab37e5be 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -764,6 +764,30 @@ def _handle_eagle_family(server_args: ServerArgs) -> None: "trtllm_mha backend only supports topk = 1 for speculative decoding." ) + # ROCm/HIP lacks the CUDA/MUSA sampling-verify kernels, so EAGLE verify would + # otherwise silently fall back to greedy (argmax), ignoring temperature/top_p. + # Default rejection sampling on (routes verify through the Triton chain sampler) + # for configs that support it, unless the user opted in or picked an incompatible + # option. Conditions mirror the validation below so the flip never triggers a raise. + # Read through cfg so the algorithm alias is already resolved (NEXTN -> EAGLE); + # write back to server_args, which is the object that actually holds the value. + from sglang.srt.utils import is_hip + + if ( + is_hip() + and not cfg.speculative_use_rejection_sampling + and cfg.speculative_algorithm in ("EAGLE", "EAGLE3") + and cfg.speculative_eagle_topk == 1 + and cfg.speculative_accept_threshold_single == 1.0 + and cfg.speculative_accept_threshold_acc == 1.0 + and not cfg.enable_deterministic_inference + ): + server_args.speculative_use_rejection_sampling = True + logger.info( + "ROCm requires rejection sampling for correct EAGLE spec-decode " + "sampling; enabling speculative_use_rejection_sampling by default." + ) + if cfg.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