[jianc99]

This is component 3 of the DFlash adaptive-verification series. Original implementation and component-only review: jianc99#3. Integration and full-system results: #46182.

Updated onto upstream main (6a319e7). All applicable pre-commit checks passed on this branch. The GPU tests and benchmark results below were run previously on the original source branch; they have not been rerun after this upstream update. Submission preparation used AI assistance.

---

Problem

DFLASH (selector and plain drafts), DSPARK, UNO dense verification and sampling-mask capture build their sampled-verify target distribution in build_speculative_verify_target_probs (srt/speculative/dflash_utils.py). EAGLE-family, standalone and n-gram verify build their own and are unchanged.

The function divided the whole [rows, vocab] logits matrix by the temperature, then ran torch.topk over the result. With Qwen3.8's 248,320-token vocabulary on GB300 this costs:

• 3.8 ms per verify step at 1024 rows (128 requests x 8-token blocks).
• 14.2 ms at 4096 rows (512 requests). Every tensor-parallel rank pays this.
• An extra [rows, vocab] fp32 temporary.

Fix

• Add sglang.kernels.ops.sampling.topk.radix_topk. It is resolved once per platform, like the renorm ops:
  • When FlashInfer is available, it uses FlashInfer's radix top-k with sorted output in deterministic mode.
  • Otherwise, or for CPU tensors, it uses torch.topk.
  • It accepts any k up to the vocabulary size.
• build_speculative_verify_target_probs selects the top-k on the raw logits and divides only the k selected values by the temperature. Dividing by a positive temperature keeps the order of the logits, so the probabilities stay the same.
• The DFlash2 selector in srt/models/dflash.py had a private copy of the same FlashInfer call. It now uses the shared op.
• Unchanged: the per-row top-k mask, softmax, top-p renormalization and dense scatter. The path without top-k computes the same values as before.

Ties: when several logits tie at a row's top-k boundary, which of them keep the probability is deterministic for a given batch, but it may differ from torch.topk's choice, and it can change with the number of rows, because the row count selects FlashInfer's kernel path. The number of kept tokens and the probability values are the same. As with torch.topk, ranks that see the same logits select the same ids.

Effect

Microbenchmark, GB300: vocab 248,320, top_k 20, top_p 0.95, temperature 1.

rows	before	after	peak temporary memory, before to after
512	2.07 ms	0.67 ms	0.95 to 0.47 GiB
1024	3.77 ms	1.04 ms	1.89 to 0.95 GiB
2048	7.17 ms	1.76 ms	3.79 to 1.90 GiB
4096	14.2 ms	3.21 ms	7.58 to 3.79 GiB

Below 128 rows the build is launch-bound, at about 0.5 ms both before and after.

End-to-end setup: Qwen3.8-27B with a DFlash2 draft (block 8), static DFLASH, TP1, GB300, --max-running-requests 128, SPEED-Bench first turns, T=1 / top_p 0.95 / top_k 20, up to 4096 output tokens. The before and after runs were interleaved across two GPUs. Decode tok/s = per-request decode rate (output tokens after the first / time after the first token), times concurrency.

concurrency	before, 3 runs (decode tok/s)	after, 4 runs (decode tok/s)	change
64	5863 (5811-5958)	6099 (5955-6197)	+4.0%
128	6908 (6858-6936)	7189 (7163-7239)	+4.1%

Acceptance length is unchanged: 3.58-3.61 on both sides at c128. At c128 the running cap equals the concurrency, so about one request waits in the queue; #46182 shows that this setup reads 7-15% higher at that level than a larger cap, on both sides, so compare c128 only within this table. At 512 requests (TP4) the per-rank saving is about 11 ms per verify step, from the microbenchmark above.

Tests

• New test/registered/kernels/ops/sampling/test_verify_target_probs.py, which checks:
  • Bitwise equality with the previous full-vocabulary implementation, for top-k 20, 2048, 2049 and 4096, in uniform and mixed batches including greedy rows, with and without top-p.
  • Rows masked to fewer finite logits than k.
  • Ties at each row's own top-k boundary in mixed batches: exact kept count, a valid top-k, the same probability values, and the same result on every call with the same batch.
  • Ties created by the temperature division.
  • The torch.topk fallback.
  • radix_topk against torch.topk for fp32 and bf16, contiguous and strided input, k up to 4096.
• Existing tests still pass: DFlash logits and selector, DSpark kernel parity and accept, speculative sampling, UNO sparse sampling, block verification, and import without FlashInfer.
• With about 1 s of GPU work queued, the call returns to the host in about 1 ms. It does not synchronize, so overlap scheduling is unaffected.

Related upstream work

• [Fix] sampling: Make top-k/top-p filtering consistent across sampling paths #41916 (semantic conflict). It changes the sparse top-k/top-p path of the same function to apply top-p against the full-vocabulary distribution. For that it reads scaled_logits, the full [rows, vocab] temperature-scaled matrix that this PR removes: its CUDA helper sums over scaled_logits, and its seeded/non-CUDA path takes a logsumexp over it. Beyond the overlapping lines, whichever lands second should compute that sum from the raw logits with the temperature applied inside the reduction (its helper already takes a temperatures argument for this) rather than bring the scaled copy back. With top-p set, that reduction is one more full-vocabulary pass per verify step.
• [Fix] dspark: Skip top-p renormalization on compact verify padding rows #43374 edits the dense top-p branch of the same function and adds a padding_verify_lens argument, one hunk away from this change; whichever lands second rebases.
• [Feat] speculative: Online acceptance estimator for DSpark confidence-scheduled verification #44866 fills compact padding rows with a point mass (0 in column 0, -1e9 elsewhere) instead of zeros. Those rows still tie at the top-k boundary, so the padding cost noted below was not measured with that fill.

Notes

• Compact (ragged) DSpark verification still builds probabilities for all bs x W rows of its strided buffer. Its zero-logit padding rows are all ties and cost the radix kernel about 4x a normal row: 1024 rows with half padding take 2.0 ms, against 1.0 ms with none (3.8 ms before this change). Restricting the work to verified rows needs the logit adjustments to run on packed rows, so it is left for a follow-up.
• UNO keeps its own FlashInfer top-k call. It can move onto this op separately.

---

CI States

Latest PR Test (Base): ❌ Run #40044008844
Latest PR Test (Extra): ❌ Run #40044008811
Latest PR Test (AMD ROCm 10): ❌ Run #40044008869