[ROCm] Fix EAGLE spec-decode verify silently sampling greedy on HIP

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

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

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()