glm5next: keep rejected DFlash draft keys out of the C4 tail ring

[shieldstar]

What

glm_kpool.py:_decode_update_kernel wrote tail[slot][pos % 4] = key for every row of a speculative request, including DFlash draft rows the verifier then rejected. The 4-slot C4 tail ring therefore carried keys of never-committed tokens into the next completed pool of the FP8 index. Glm5NextModelState already stages selector_num_accepted_tokens, but pooled_indexer.py never consumed it.

This PR adds per-slot DecodeTailRing bookkeeping: each decode step snapshots the ring, and the next step restores that snapshot and replays only the accepted prefix of the previously saved rows before writing its own.

Details

• Accepted length is position-derived: this step's first position minus the previous step's first position. selector_num_accepted_tokens is consumed as a cross-check; on disagreement positions win, so resume and replay re-base the request correctly instead of trusting a stale count.
• _prefill_tail_kernel zeroes the slot's saved count, so a fresh, resumed or short-tail prefill chunk never restores a stale ring.
• update_decode_pools takes tail_ring= and num_accepted_tokens= (ones as fallback); the pooled indexer allocates the ring sized 1 + num_speculative_tokens and passes main_metadata.selector_num_accepted_tokens[:num_reqs].
• Memory: about 100 KB at max_num_seqs=16, K=7.
• Only the DFlash path of GLM-5.3 is affected. The MTP snapshot/restore path is untouched.

Found while auditing the mixed prefill/decode path for the corruption described in #616; it is a separate defect (that corruption reproduces with speculative decoding off), reported here on its own.

Testing

Kernel-level (GB10, Triton 3.7.1, the dev/jovian-judgement container): _decode_update_kernel, _prefill_pool_kernel and _prefill_tail_kernel JIT-compile and run. A two-step scenario (max_num_seqs=16, K=7, two 8-row speculative requests at positions 100 and 200 plus a 300-row prefill; step 2 restarts them at 103 and 205 with num_accepted_tokens=[3, 5]) produces tails, ring state and completed pools bit-identical to a from-scratch run that only ever saw the accepted rows; the prefill slot's saved count is 0. The check script is small and can be attached if wanted.

An earlier revision of this patch failed to compile at engine warm-up (Loop carried variable saved changed type, a tl.static_range index reused as a runtime loop variable); that is fixed and is why the kernels are now exercised on the GPU rather than only py_compiled.

Serving-level, four-node DGX Spark TP4, image built from dev/jovian-judgement@2e67b303 plus this branch and #618, GLM-5.3-Flash-NVFP4-Spark with the DFlash2 draft (k=7), fp8 KV, --max-num-batched-tokens 1024:

vllm serve local-inference-lab/GLM-5.3-Flash-NVFP4-Spark -tp 4 --decode-context-parallel-size 1 \
  --block-size 256 --attention-backend B12X --moe-backend b12x --linear-backend b12x \
  --mamba-cache-mode align --enable-prefix-caching --enable-chunked-prefill --kv-cache-dtype fp8 \
  -O2 --compilation-config '{"cudagraph_mode": "FULL"}' --max-cudagraph-capture-size 128 \
  --speculative-config '{"method": "dflash", "model": "local-inference-lab/GLM-5.3-Flash-DFlash2", "num_speculative_tokens": 7, "kv_cache_dtype": "auto", "attention_backend": "FLASH_ATTN"}' \
  --max-num-batched-tokens 1024 --max-model-len 1000000 --max-num-seqs 16 --gpu-memory-utilization 0.82

• Engine start, FULL cudagraph capture and the DFlash2 draft warm-up all pass with the new kernels (3.5 min boot, KV pool 5,455,114 tokens at 1M context).

• llm-inference-bench coding-peak c1 smoke: 85.0 tok/s median, 91.0 max, 0 CJK runs; the same as the unpatched build (83-86), and since that number is acceptance-bound it shows DFlash acceptance on coding prompts is unchanged.

• llm-inference-bench estonia profile, c4 x 30, temperature 1.0: 26/30 (2 wall-clock timeouts mid-reasoning, one "United Kingdom", one correct answer written in Chinese). The unpatched build on the same chunk size scored 24-28/30 across three runs; the residual misses are the concurrent-prefill issue in GLM-5.3-Flash TP4 on 4x DGX Spark: decode corrupted by concurrent chunked prefill; rate scales with prefill tokens per step (4096 -> 1024 chunks: estonia 12/30 -> 28/30) #616 and are not affected by this change.

• Spec-decode counters over the estonia run: 193,491 accepted tokens over 80,796 draft steps (2.39 per step), no acceptance collapse.

• Four concurrent 90k-token needle prompts, 3 rounds: 10/12 clean, 1 runaway, 1 wrong (all of the GLM-5.3-Flash TP4 on 4x DGX Spark: decode corrupted by concurrent chunked prefill; rate scales with prefill tokens per step (4096 -> 1024 chunks: estonia 12/30 -> 28/30) #616 class).

  Corrected 2026-09-03: this line first read 11/12. Our needle harness scored a stream clean if the code appeared anywhere in the output, so a degenerate run that quoted the code mid-babble counted as a pass. Re-scored strictly (the visible answer is the code and the stream terminated on its own) the same log is 10/12. The comparison in this PR is unaffected, since the baseline was measured the same way, but the absolute number was wrong and is fixed here. Details in GLM-5.3-Flash TP4 on 4x DGX Spark: decode corrupted by concurrent chunked prefill; rate scales with prefill tokens per step (4096 -> 1024 chunks: estonia 12/30 -> 28/30) #616.

Not measured: a direct before/after on index-pool fidelity for generated tokens, which is what this fix changes; the semantics check above is the evidence for that part.

Duplicate check

gh pr list --state open --search "tail ring OR glm_kpool OR C4 tail OR rejected draft" returns nothing related; #505 (parallel C4 prefill pool writes) touches the prefill pool path, not the decode tail ring, and does not consume selector_num_accepted_tokens either.

Conventions

ruff check and ruff format --check (0.14.0, repo pyproject) pass on both files; every touched function has a Google-style docstring; the commit is DCO signed.

AI assistance

The audit that found the defect, the kernel change and the compile/semantics check were produced with AI assistance (Claude); the submitter reviewed every changed line and ran the tests listed above on our cluster.

Assumptions to check in review

• selector_num_accepted_tokens equals bonus + accepted rows of the previous verify step; it is only used as a cross-check, so a mismatch degrades to the position-derived value rather than corrupting state.
• DFlash rows for a request are position-ordered and the next step starts at prev_first + accepted.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}