[ROCm] EAGLE: reject non-finite draft probabilities, not just NaN

[JohnQinAMD]

Problem

speculative_sampling_classic_kernel guards the draft probability q against NaN with q_val == q_val. Infinities satisfy that test, so a -inf in draft_probs reaches p - q as +inf, norm_sum becomes infinite, and the cumulative search total_cumsum > coin * norm_sum is never true. The kernel then takes its degenerate fallback.

That fallback is VOCAB_SIZE - 1, where VOCAB_SIZE is the padded tensor width — 154880 for GLM-5.2, against the tokenizer's 154856. The id it emits is an untrained padding embedding. It decodes to an empty string, so nothing is visible in the output, but it is committed to the context and the model degenerates from there into repeated </think> / <tool_call> fragments and stray ! (token 0).

This only fires where verify actually samples. A greedy batch takes the argmax path and never reaches the kernel, which is why temperature=0 workloads are unaffected.

Fix

Test q against the probability range rather than for NaN. A comparison with NaN is false, so the range test rejects NaN along with the infinities and with negatives.

q_val = tl.where((q_val >= 0.0) & (q_val <= 1.0), q_val, 0.0)

Usage

No flag. The kernel is reached when EAGLE verify samples, which on ROCm means --speculative-use-rejection-sampling:

python3 -m sglang.launch_server \
  --model-path /path/to/GLM-5.2-MXFP4 --trust-remote-code \
  --tp 4 --ep-size 4 --kv-cache-dtype fp8_e4m3 \
  --speculative-algorithm EAGLE --speculative-num-steps 5 \
  --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 \
  --speculative-use-rejection-sampling

Then send a request with temperature > 0. Before this patch the reply carries stray ! characters at a percent-level rate; after it, none.

Results

GLM-5.2-MXFP4, MI355X, TP4/EP4, EAGLE steps=5 topk=1 draft=6, rejection sampling on, temperature=1.0 top_p=0.95:

	before	after
fallback emissions	70 / 1050 committed tokens	0 / 3437
stray token 0 per 1200 generated	25 to 81	0

GSM8K 200q 5-shot at temperature=0: 0.935 accuracy, 0.000 invalid, accept len 3.03 — unchanged, as expected for a path this patch cannot reach.

Not addressed here

Two things this patch deliberately leaves alone.

Where the non-finite draft probabilities come from. draft_probs carries -inf entries, consistently at draft step index 2, 5 to 59 of them per batch. They are a CUDA-graph artefact, not a softmax output:

	CUDA graph on	--disable-cuda-graph
non-finite entries in draft_probs	5 to 59 per batch, always step 2	0 over 250+ verify calls
stray ! per 1200 generated tokens	25 to 81	0

Two further properties of those entries: each run of them ends on a 64-aligned column, and the run's position climbs monotonically with the decode step at roughly the accept length. A wave64 tile tail whose position tracks the sequence is the shape of an attention-mask fill, so the draft step's softmax output buffer appears to overlap something in the graph's memory pool that writes -inf. Pinning down the exact pair needs allocation-interval instrumentation of that pool, which is out of scope here.

This kernel has to tolerate the values either way, which is what the NaN guard already intended.

The fallback itself. final_token = VOCAB_SIZE - 1 is still an invalid id whenever the vocabulary dimension is padded, even though the guard above makes it unreachable in the case observed. Falling back to argmax(p) would need the pass-1 loop to carry a running max, which is more than this fix should take on.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34638460549
Latest PR Test (Extra): ❌ Run #34638460202
Latest PR Test (AMD ROCm 10): ❌ Run #34638460442

[JohnQinAMD]

Superseded by #58, which carries this commit unchanged plus the root cause it defends against: draft_probs is returned from the draft CUDA graph's private memory pool and that block is recycled before eagle_sample reads it, which is what puts the -inf there in the first place.

Closing in favour of the combined PR — the two halves are not independently useful. This one alone leaves the draft distribution corrupt (just no longer visible, since the guard zeroes those entries), and the root-cause fix alone leaves a sampler that can still emit VOCAB_SIZE - 1 on any future malformed input.