[ormandj]

Motivation

Compact verify scatters verified logits into a strided tensor and fills each request's unverified rows with zero logits. The resulting [bs * verify_num_draft_tokens, vocab] logits are passed to acceptance. When a request samples with top-p, the accept path builds the target probabilities with softmax and FlashInfer top_p_renorm_probs over every row, padding included.

Without logit adjustments, a padding row is exactly uniform after softmax. In FlashInfer's radix (AIR) top-p, every element of a uniform row falls into the same histogram bucket on every pass, and in the kernel microbenchmark below these rows make the call 3.5x to 12.5x slower than a copy of a verified row. The renorm also has no effect on such a row: FlashInfer keeps every probability >= the threshold it finds, which for a uniform row keeps every token, then divides by the row sum (1 up to rounding).

Modifications

• build_speculative_verify_target_probs gets an optional padding_verify_lens. In the dense top-p branch (top-p only, no top-k) on CUDA, padding rows are not renormalized. A Triton kernel copies each request's last verified row into its padding rows in place, so the zero-logit padding rows no longer reach the top-p kernel, and a second pass writes the padding rows' saved uniform softmax value back into the output. Rows of requests without padding are not touched, and no extra [rows, vocab] buffer is allocated.
• The DSpark accept path passes the compact-verify verify_lens only in compact mode, when the batch has no grammar and verify_logits_adjustments_are_noop holds (no logit bias, penalties or custom logit processor). Host-side lengths disable substitution for full-width batches; device-only layouts (built when the verify token count is rounded up to a captured CUDA-graph size) conservatively enable it. For a full-width device-only layout, the two kernels launch but leave every row untouched. In all other cases, and for all other callers, the code path is unchanged.
• The strided logits themselves are unchanged.

Accuracy Tests

New CUDA test test/registered/spec/dspark/test_dspark_compact_top_p.py. It uses a compact batch with verify lengths 6/3/1/4, per-request temperatures and top-p values, and verified rows on which top-p keeps two tokens and drops a third:

• target probs at vocab 32,000 and 129,280: verified rows are bit-identical to the full-renorm path; padding rows match it within rtol 1e-5;
• accept outputs (correct_len, bonus, cap_trim_lens) are bit-identical to the full-renorm path over 4 seeds, and match the expected values for a rejection on an inner verified row, a rejection on the last verified row, and acceptance through the last verified row into the padding rows;
• with a device-only full-width layout (lengths 6/6/6/6, vocab 129,280), target probs are bit-identical to the full-renorm path and peak CUDA allocation is within 1 MiB of it;
• with a logit bias added to the logits, under which only a top-p renorm rejects the drafts on padding rows, accept outputs are bit-identical to the full-renorm path.

python -m pytest test/registered/spec/dspark/test_dspark_compact_top_p.py on one RTX PRO 6000 Blackwell (SM120): 4 passed, 2 subtests passed. test/registered/unit/kernels/test_kernels_namespace.py: 67 passed. On 2026-10-29 this was rebased onto main 974831643f; the new head keeps the full renorm for batches whose logits watermarking forced, the CPU-runnable test/registered/unit/kernels/test_kernels_namespace.py passes on it (apart from one failure main has too), and the SM120 GPU results above are from the previous head and will be re-run.

Speed Tests and Profiling

Kernel microbenchmark (RTX PRO 6000 Blackwell, SM120): FlashInfer top_p_renorm_probs on [rows, 129280] FP32 probabilities, top_p 0.95, half the rows padding, timed in a CUDA graph. "Copy of last row" is the kernel input this change produces.

rows	zero-logit padding rows	copy of last verified row
6	494 us	142 us
48	2352 us	188 us

This times the top-p kernel only. The two added padding-row kernels, and end-to-end serving with this version, were not measured.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #40961599764
Latest PR Test (Extra): ❌ Run #40961599687
Latest PR Test (AMD ROCm 10): ❌ Run #40961599658