[ROCm] EAGLE: stop draft_probs being recycled out from under the verify sampler

[JohnQinAMD]

Two commits against one bug: the first removes its cause, the second stops the sampler turning a malformed input into an invalid token id. Supersedes #57, which carried the second commit alone.

Problem

With EAGLE and rejection sampling, a request at temperature > 0 comes back with stray ! and repeated </think> / <tool_call> fragments — a few percent of generated tokens on GLM-5.2. temperature = 0 is clean, because a greedy batch takes the argmax path and never reaches the sampler.

The chain:

#
1	draft_probs is the one draft-graph output the target forward does not consume — it is read afterwards, by eagle_sample. The tensor the replay returns points into the graph's private memory pool.
2	The pool recycles that block in between. What lands there is the DSA top-k mask, which writes -inf into positions past the current sequence.
3	In the sampler, q_val == q_val rejects NaN but not -inf, so p - q goes to +inf and norm_sum with it.
4	total_cumsum > coin * norm_sum is then never true, and the kernel commits its fallback, VOCAB_SIZE - 1.
5	VOCAB_SIZE is the embedding width, 154880 for GLM-5.2, against a tokenizer vocabulary of 154856. The committed id is an untrained row. It decodes to an empty string, so nothing is visible in the output, but it enters the context and the model degenerates from there.

Three properties of the corruption identified step 2: each run of -inf ends on a 64-aligned column (page_size = 64), the runs advance with the decode step at roughly the accept length (the mask's valid region tracks the sequence), and they land at the same draft step every time (the pool's allocation order is fixed at capture). --disable-cuda-graph removes them entirely, which is what ruled out the producer.

Fix

Copy draft_probs out at the graph boundary. Only rejection sampling produces it, so this is a no-op everywhere else.

Test q against the probability range rather than for NaN, in the two residual passes:

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

Reject rather than accept when q is not a probability, in the accept test. That site is where a corrupt row does the most damage: coin * q < p is -inf < p for an -inf q and 0 < p for a zero one, so the step accepts unconditionally and commits the draft head's token instead of the target's. Zero also passes the range test above, so widening that guard does not cover this site. X was sampled from q, so q(X) has to be strictly positive:

q_is_prob = (q > 0.0) & (q <= 1.0)
if q_is_prob & (coin * q < p):

The failure is invisible from the outside — acceptance climbing toward 1.0 reads as a good draft model, not as a broken test.

The second is worth keeping on its own: a sampling kernel should not be able to emit an invalid token id because an input was malformed, whatever wrote that input. It also covers the VOCAB_SIZE - 1 fallback, which stays wrong in principle for a padded vocabulary even though this makes it unreachable in the case observed.

Usage

No flag. The path is reached with EAGLE and rejection sampling, which on ROCm means:

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

Send a request with temperature > 0 and count ! in the reply.

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
non-finite entries in draft_probs	5 to 59 per batch	0
sampler fallback emissions	70 / 1050 committed tokens	0 / 3437
stray tokens per 1200 generated	25 to 81	0

Cost, bench_serving random 1024/512, concurrency 1, both arms on this commit's base so the copy is the only difference:

	without	with
median ITL	2.56, 2.56, 2.56, 2.56	2.57, 2.57, 2.57
output tok/s	290.8, 305.2, 297.5, 282.7	287.5, 301.8, 305.2

Median ITL is stable to the hundredth of a millisecond in both arms and moves by one of them, 0.4%, which is the display resolution of the metric. Throughput varies by more than that within each arm, so it carries no signal either way. The copy is one (bs, steps, vocab) fp32 tensor per decode step, about 3 MB at these shapes.

GSM8K 200q 5-shot at temperature=0: 0.940 accuracy, 0.000 invalid, accept len 3.035. Unchanged, as expected for a path this cannot reach — three repeats on one unchanged server span 0.920 to 0.945.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34646795717
Latest PR Test (Extra): ❌ Run #34646795214
Latest PR Test (AMD ROCm 10): ❌ Run #34646795657