[haosdent]

Overview

cooperative_topk now chooses between its single-pass and two-pass paths from each row's length instead of the
logits row stride. Rows that fit the single-pass staging no longer stream through TMA twice just because the
buffer is wide (max_model_len).

Claims

• InferenceX AgentX, GLM-5.3-NVFP4 on 4x GB200: together with [Perf] Emit the DSA indexer weights in BF16 on SM100 #62170 and [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172, decode time per turn
  drops 2.8%, 2.5%, 1.5% and 1.5% at concurrency 1, 4, 8 and 16 (main vs main: within 0.7%). This PR is
  about half of that.
• The decode top-k is up to 8% faster on B300 for rows that fit the single-pass staging when the stride is larger
  than the per-block staging, for example max_model_len 1M at 1-8 rows, or about 200K at 9-64 rows.
• Selection is unchanged: both paths compute the same exact radix top-k.

Validation

• tests/kernels/test_top_k_per_row.py: 265 passed, 37 skipped on GB200 (SM100); after rebasing onto [Bugfix] Make cooperative TopK exact on candidate overflow #59645's
  overflow recovery, 402 passed, 44 skipped on B300 (SM103). The new test_cooperative_topk_rows_shorter_than_stride
  covers every cluster size (1, 4, 8, 33 and 64 rows) with strides of 64K and 1M and rows of 300 to 300K tokens, so
  one launch mixes trivial, register, single-pass and streaming rows. It compares the selected values with
  torch.topk.

• InferenceX AgentX, GLM-5.3-NVFP4 on 4x GB200, main + [Perf] Emit the DSA indexer weights in BF16 on SM100 #62170, [Perf][Kernel] Size cooperative_topk's single pass by row length, not stride #62171 and [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172 vs main (median per-turn change
  on the same turns, 95% CI; setup below):

  concurrency	1	4	8	16
  decode time	-2.8% [-3.1, -2.6]	-2.5% [-3.1, -1.7]	-1.5% [-2.2, -1.1]	-1.5% [-2.4, +0.3]
  e2e latency	-3.3% [-3.9, -3.0]	-1.8% [-2.7, -0.9]	-1.2% [-2.0, -0.5]	-0.6% [-1.9, +2.4]
  KV cache tokens	+0.2%	+0.7%	+1.4%	0
  main vs main, decode	-0.1%	-0.7%	+0.3%	-0.3%

• cooperative_topk microbenchmark, k=2048, rows = requests x 4 (MTP), CUDA graph replay, B300:

  rows	context	stride	before (us)	after (us)
  4	40K	1M	10.99	10.27
  4	150K	1M	12.98	12.11
  16	40K	1M	12.18	11.24
  16	40K	64K	11.27	11.09

• Target verify forward (vllm bench forward, GLM-5.2-NVFP4, TP4, B300, MTP verify shapes, spread < 0.1%), on top
  of [Perf] Emit the DSA indexer weights in BF16 on SM100 #62170: -0.58% to -1.36% at 1, 3 and 7 requests, 25K and 94K context.

• Exactness: with real MTP verification on 32 GSM8K questions (greedy), the three changes together give the same
  outputs and the same acceptance length (3.4904) as main.

• MRCR (8 needles, 4K to 256K, 594 samples), GLM-5.3-NVFP4 on 4x GB200, greedy: main 0.903, with [Perf] Emit the DSA indexer weights in BF16 on SM100 #62170,
  [Perf][Kernel] Size cooperative_topk's single pass by row length, not stride #62171 and [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172 0.890; paired difference -0.013 [-0.038, +0.012], and every bin's CI includes zero. The
  256K bin, run three times each: -0.019 [-0.069, +0.031].

InferenceX AgentX setup and results for all GLM-5.3 PRs (4x GB200)

InferenceX c801c9f2d AgentX scenario (pinned AIPerf), 20-minute fast mode, GLM-5.3-NVFP4 rev e3b8f9f2, TP4, MTP 3
with synthetic acceptance 2.99, native KV offload 135 GB per rank, Rust frontend, max-num-seqs = 2 x
concurrency. vLLM main dcd0796 with its default backends and CUDA graph config, plus the PRs cumulatively in
the order below (FlashInfer with flashinfer-ai/flashinfer#6630, flashinfer-ai/flashinfer#6631 and
flashinfer-ai/flashinfer#6632 from #62174 on). One FlashInfer MoE autotune per concurrency, shared by all
variants. A fresh server per point, with the page cache dropped before each, and main run again at the end of
each concurrency as an A/A check. SGLang a916d91cd3 with IMEX, TRT-LLM
all-reduce fusion and --mem-fraction-static 0.90.

Each cell is the median, over the turns both runs finished (150 to 520), of the per-turn ratio to the previous
variant, with a 95% paired bootstrap CI. Decode is request latency minus TTFT. The bootstrap does not include
run-to-run variance; the main-vs-main row does.

decode time per turn	c=1	c=4	c=8	c=16
+ #62170, #62171, #62172	-2.8% [-3.1, -2.6]	-2.5% [-3.1, -1.7]	-1.5% [-2.2, -1.1]	-1.5% [-2.4, +0.3]
+ #62174	-1.6% [-1.8, -1.5]	-1.8% [-2.5, -0.9]	-3.3% [-3.9, -2.9]	-2.1% [-3.8, -0.9]
+ #62173	+0.5% (no-op here)	-0.3% [-0.5, -0.1]	-0.1% [-0.4, +0.3]	-0.9% [-2.0, +0.2]
all PRs vs main	-3.7% [-3.9, -3.5]	-4.3% [-4.8, -4.0]	-5.3% [-6.4, -4.8]	-4.8% [-6.5, -3.8]
main vs main (noise)	-0.1%	-0.7%	+0.3%	-0.3%

e2e latency per turn	c=1	c=4	c=8	c=16
+ #62170, #62171, #62172	-3.3%	-1.8%	-1.2%	-0.6%
+ #62174	-1.4%	-1.4%	-2.7%	-1.6%
+ #62173	+0.9%	-0.3%	-0.7%	-1.1%
all PRs vs main	-4.1%	-3.4%	-4.7%	-3.8%
main vs main (noise)	-0.6%	-1.2%	+0.4%	-0.4%

KV cache tokens per GPU (c=1 / 4 / 8 / 16): main 1,176,832 / 1,162,880 / 1,143,872 / 1,105,984; with all PRs
1,178,496 / 1,168,512 / 1,157,824 / 1,103,936. Window throughput (tok/s/GPU) moves up to 2.4% between runs in fast
mode, since it counts whichever long-context turns finish inside the 20 minutes; with all PRs it is +3.0%, -1.4%,
+4.8% and +1.2% at c=1/4/8/16.

vLLM vs SGLang, decode time per turn on the same turns (positive = vLLM slower):

	c=1	c=4	c=8	c=16
main	+10.4%	+12.7%	+13.7%	-2.9%
all PRs	+6.2%	+7.5%	+5.7%	-6.9%

#62170, #62171 and #62172 were measured as one group. vllm bench forward (target verify forward, spread < 0.1%)
splits their gain about 1/6, 1/2 and 1/3. #62173 changes only decode steps of 3, 5 or 7 requests; steps with 3 or
more requests are 0%, 4%, 13% and 37% of all steps at c=1/4/8/16. At c=1 both runs capture the same graphs, so
its +0.5% there is noise. At c=16, three more runs each of + #62174 and + #62173 (4 runs each, compared as
runs): decode time -0.54% [-0.93, -0.16], with every + #62173 run faster than every + #62174 run (permutation
p = 2/70); e2e latency -0.29% [-1.41, +0.82].

Details

The DSA indexer allocates its decode logits with max_model_len columns, and cooperative_topk gets that as
stride. The kernel sized per_block (elements per CTA of the cluster) from the stride. But large_topk splits
the row by its actual length (seq_len, in 16-byte units across the cluster ranks). With a long max_model_len,
every row took the two-pass streaming path even when its context fit the all-stages-resident single pass. For
example, 16 rows use 4-CTA clusters, and a 256K stride gives 64K elements per block, above the 48K single-pass
staging, while a 40K-token row needs only 10K per block.

per_block is now ceil(seq_len / CS), so each row picks its path from its own length. large_topk rounds each
block to 16-byte units, but the single-pass limit and the stage size are multiples of that unit, so this gives the
same path and stage count as the exact split. Rows in one launch can take different paths. Each row is its own
cluster, and the barrier setup already follows the chosen path.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[haosdent]

Could be

(static_cast<uint32_t>(sl) + CS - 1) / CS

[LopezCastroRoberto]

LGTM, thanks!

[github-actions]

✅ @haosdent, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[LopezCastroRoberto]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96880 for commit 789fb34874f4.