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

[JohnQinAMD]

Motivation

On ROCm/HIP, EAGLE speculative decoding always commits argmax (greedy) in verify. It ignores
temperature and top_p:

if sampling_info.is_all_greedy or _is_cpu or _is_npu or _is_hip or _is_xpu:
    target_predict = torch.argmax(...)   # greedy

_is_hip is in that list because the sampling-verify ops
(tree_speculative_sampling_target_only, top_k_renorm_prob, top_p_renorm_prob) are only imported
under is_cuda()/is_musa(). So any EAGLE spec-decode run at temperature>0 on ROCm decodes
greedy. On long reasoning output that turns into repetition loops. This is not model specific.

Modifications

1. eagle_utils.py — move the check into a small _verify_uses_greedy(...) helper, which returns
   exactly the old answer on every non-HIP platform. Add a Triton top-p renorm (pivot binary search)
   and a HIP wrapper to stand in for the CUDA-only sgl_kernel ops, and guard the CUDA-only imports.
   HIP then goes through the pure-Triton chain_speculative_sampling_triton.
2. arg_groups/speculative_hook.py — turn on speculative_use_rejection_sampling by default on HIP
   for EAGLE/EAGLE3 with topk==1, default accept thresholds, and non-deterministic mode, so users do
   not have to know about a flag. These conditions are a strict subset of the checks that raise just
   below, so turning it on can never break a config that used to work.

Accuracy Tests

Both sides were run on this branch with the same settings. Only the patch differs. GLM-5.2-FP8,
8x MI355X, NEXTN spec decode (num_steps=1, eagle_topk=1, draft_tokens=2), GPQA-Diamond 198
questions, temperature=1.0 / top_p=0.95, max_tokens=98304.

	old code	this PR
GPQA-Diamond	0.702 ± 0.033	0.904 ± 0.021
answers that used up the whole token budget	78/198 (39.4%)	5/198 (2.5%)
mean response length	134,921 chars	7,860 chars
accuracy on answers that did finish	0.9333	0.9223

+20.2 points, about 5.2 sigma. The cause is not model quality. On the questions where the old
code finishes, it scores 0.9333, which matches the patched build. It loses because 39.4% of its
answers never finish: greedy verify at temperature>0 falls into repetition loops and burns the
whole budget without giving an answer. You can also see it in wall-clock time (about 2x) and in a
first attempt that died because requests went past the default 1800s per-request timeout.

A probe on a non-greedy batch confirms which branch runs: old code takes GREEDY(argmax), this PR
takes SAMPLING. The same run shows the speculative_hook.py half working, with the server logging
enabling speculative_use_rejection_sampling by default and no error.

Speed Tests and Profiling

This is a correctness fix; the sampling path costs the same. Throughput under NEXTN spec decode on GSM8K n=200: 740-771 tok/s before, 764 tok/s after.

GSM8K cannot show this bug because it runs greedy, so argmax is already the right answer there.
It is only a no-regression check: 0.945-0.960 old vs 0.940 patched at n=200, throughput 740-771 vs
764 tok/s.

Known gaps (pre-existing, not regressions)

On HIP, EAGLE topk==1 with --enable-deterministic-inference or non-default accept thresholds is
left alone on purpose, because rejection sampling really is incompatible there. HIP tree drafting
(topk>1) still uses greedy, because the Triton chain sampler only handles chains.

_renorm_top_k_top_p_hip decides whether to run the top-k sort with bool((top_ks < vocab).any()),
which is a device sync on every verify step. Passing sampling_info.need_top_k_sampling and
need_top_p_sampling (what the CUDA branch already uses) would remove it. That measured well under
0.5% of a decode step, so I left it out to keep this PR a pure bug fix whose numbers match the code.

How to reproduce

Image lmsysorg/sglang-rocm:v0.5.17-rocm720-mi35x-20260817, /data/huggingface/GLM-5.2-FP8. Use
FP8, not MXFP4: on that image quark.py imports a name its overrides.py does not have, so the
server dies while loading weights.

# Unit tests. Pass the GPU devices even for CPU-only tests: importing sglang pulls in aiter,
# which runs rocminfo at import time.
docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video \
  --security-opt seccomp=unconfined --shm-size 16g -v "$PWD:/wt" -w /wt -e PYTHONPATH=/wt/python \
  lmsysorg/sglang-rocm:v0.5.17-rocm720-mi35x-20260817 \
  bash -lc 'python3 -m pytest test/registered/unit/spec/test_eagle_gate_routing.py \
            test/registered/spec/eagle/test_rocm_eagle_renorm.py -q'

# Serve with spec decode
python3 -m sglang.launch_server --model-path /data/huggingface/GLM-5.2-FP8 --trust-remote-code \
  --tp-size 8 --page-size 64 --kv-cache-dtype fp8_e4m3 \
  --dsa-prefill-backend tilelang --dsa-decode-backend tilelang \
  --speculative-algorithm NEXTN --speculative-num-steps 1 \
  --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \
  --reasoning-parser glm45 --context-length 163840 --mem-fraction-static 0.75 --port 30000

# GPQA-Diamond. lm-eval pinned at b315ef3b05176acc9732bb7fdec116abe1ecc476,
# task gpqa_diamond_cot_n_shot.
python3 -m lm_eval --model local-chat-completions --apply_chat_template \
  --tasks gpqa_diamond.yaml --output_path "$RESULTS" --log_samples \
  --model_args "model=$MODEL,base_url=http://0.0.0.0:30000/v1/chat/completions,api_key=EMPTY,\
eos_string=</s>,max_retries=5,num_concurrent=64,timeout=7200,tokenized_requests=False,max_length=163840" \
  --limit 198 --gen_kwargs "max_tokens=98304,temperature=1.0,top_p=0.95"

--limit 198 takes one permutation per question (process_docs uses n_repeats=2). You need
HF_TOKEN, the dataset is gated. lm-eval has to fall back to reasoning_content when content is
empty, or every graded answer from a thinking model is an empty string.

Two things to watch when probing the branch. Gate the probe on not sampling_info.is_all_greedy:
server warmup sends a greedy request, so a plain one-shot probe tells you nothing. And do not use
output diversity as a test: the old code still returns 8 out of 8 different completions at
temperature=1.0, because the bonus token goes through the normal sampler.

AI assistance

This PR was developed with AI assistance (Claude). The AI helped with the rebase and conflict
resolution, wrote the unit tests, and ran the measurements reported above. All code and results were
reviewed and verified by the author on real hardware; the commits carry a Co-authored-by trailer.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #33038846825
Latest PR Test (Extra): ❌ Run #33038846496
Latest PR Test (AMD ROCm 7.2): ❌ Run #33038846762

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ErenAta16]

Checked the pivot-search kernel and the routing logic separately.

_top_p_renorm_kernel: binary search for the largest tau with sum(p>=tau) >= top_p is the standard trick for sort-free top-p. 30 iterations against a probability in [0,1] gets tau to roughly 1e-9 precision, comfortably inside the <1e-6 tolerance claimed against the sort-based reference. The top_p >= 1.0 passthrough keeps that path bit-exact with CUDA's existing no-op behavior instead of running the approximation unnecessarily.

_verify_uses_greedy / the speculative_hook.py default flip: the new auto-enable conditions (topk==1, both accept thresholds at 1.0, not deterministic-inference) mirror the pre-existing validation block below it, so the flip can't produce a config that block would've rejected anyway. Lazy-importing sgl_kernel under if not _is_hip is correct: the gate guarantees the HIP branch always goes through rejection sampling, so the CUDA-only import is genuinely dead on that path, not just usually-true.

The new test compares the kernel against an independent torch nucleus reference (its own sort/cumsum implementation, not a restatement of the kernel's math) rather than a golden file, which is the right way to catch a broken pivot search instead of just pinning current output. Good fix for a real correctness gap, ROCm was silently downgrading EAGLE verify to greedy regardless of temperature/top_p.

[kkHuang-amd]

Could you add the e2e gsm8 testing result for greedy and non-greedy cases?

[JohnQinAMD]

Could you add the e2e gsm8 testing result for greedy and non-greedy cases?

Thanks @kkHuang-amd @ErenAta16 for the review

e2e GSM8K on GLM-5.2-MXFP4 (EAGLE topk==1 chain + rejection sampling, MI350X/ROCm), full test set (N=1319):

decode	GSM8K
greedy (temp=0)	96.59% (1274/1319)
sampling (temp=0.6)	96.74% (1276/1319)

Sampling ≈ greedy — temp>0 now samples correctly on HIP instead of the pre-fix silent-greedy fallback.

[ErenAta16]

The e2e gsm8k ask is reasonable, flagging that I don't have ROCm hardware to generate that myself, so it'll need to come from someone with access to the target device. The kernel-level fix stands on its own regardless: the GPU-free unit test isolates the actual bug (pivot-search binary search validated against an independent torch reference, plus the auto-enable/rejection-sampling flip checked against the existing validation block), so the e2e run would be confirming end-to-end behavior on top of an already-verified kernel, not the only evidence for correctness.

[JohnQinAMD]

Closing in favour of #37134, which carries the same fix adapted to the current arg_groups API (get_platform().is_hip and the resolving view instead of is_hip() and raw server_args). This one was written against the July API, and has since gone conflicting against main for that reason.