[mickaelseznec]

Overview

The DCP sparse-indexer top-k merge (_merge_dcp_topk_global) all-gathers each rank's packed (score, global_id) candidates along dim 1. Since NCCL gathers along dim 0, GroupCoordinator.all_gather then reorders the result into [rows, world_size * topk, 2]: a full extra copy of every candidate. This PR gathers along dim 0 and lets the CuTeDSL selector read the rank-major [world_size, rows, topk, 2] buffer in place. The selected set is unchanged.

With W ranks, R rows and K = topk candidates per rank:

# Before
packed = pack(local_topk)                    # [R, K, 2]  (score, global_id)
buf = nccl_all_gather(packed)                # [W*R, K, 2]  rank-major
gathered = buf.view(W, R, K, 2)              #              (free view)
           .movedim(0, 1)                    # [R, W, K, 2]
           .reshape(R, W*K, 2)               # [R, W*K, 2]  <-- full copy
selector(gathered)                           # one CTA per row r:
    for key_idx in range(W*K // 512):
        col = tid + key_idx * 512
        key = gathered[r, col]

# After
packed = pack(local_topk)                    # [R, K, 2]
buf = nccl_all_gather(packed)                # [W*R, K, 2]  rank-major
gathered = buf.view(W, R, K, 2)              # free view, read in place
selector(gathered)                           # one CTA per row r:
    for key_idx in range(W*K // 512):        # same loop, same keys
        rank = (key_idx * 512) // K          # compile-time constant (K % 512 == 0)
        cand = tid + (key_idx * 512) % K     # compile-time constant + tid
        key = gathered[rank, r, cand]

Claims

• Removes the reorder copy from every DCP top-k merge (prefill and decode) of DSA models with a sparse indexer under DCP (DeepSeek-V3.2, GLM-5.x).
• Kernel level on B200 (6710 rows, DCP8, top-2048): the merge's gather-side work drops from 1438 µs (reorder copy + selector) to 839 µs (selector on the rank-major input; 838 µs on the row-major input, so reading in place costs nothing).
• End to end, in a downstream fork (internal DSA model, 8×B200, TP8 × DCP8, top-2048): prefill throughput +5–6% (TTFT −5 to −6%) over 5k–100k-token prompts.

Validation

• tests/v1/attention/test_indexer_dcp_localize.py: 114 passed on B200. The selector test now feeds the rank-major layout and compares it with an fp64 reference over all candidates. The faked-gather merge tests follow the dim-0 gather.

• Downstream end-to-end benchmark (same container for both runs, 64 generated tokens, 3 iterations):

  prompt	batch	prefill tok/s before	after	ratio
  5k	1	12,746	13,411	1.05×
  5k	16	12,677	13,312	1.05×
  20k	1	14,356	15,218	1.06×
  20k	16	14,153	15,011	1.06×
  50k	4	13,554	14,334	1.06×
  100k	16	13,391	14,148	1.06×

  Decode was unchanged there, but that fork serves decode-sized batches through a separate symmetric-memory exchange. Upstream's decode path goes through this merge too, and is not measured here.

• No public-model end-to-end run yet (e.g. DeepSeek-V3.2 or GLM-5.x with DCP).

Details

• The selector's only caller is _merge_dcp_topk_global, which already requires topk ∈ (512, 1024, 2048). So each rank contributes topk candidates, a multiple of the 512-thread block, and every 512-wide slice of a row's candidates falls inside one rank. Each thread key's (rank, candidate) is then a compile-time constant plus the thread index: no runtime division. The CompileKey(topk, num_candidates) and the selector's radix passes are unchanged.
• stable_topk_from_gathered_candidates_cutedsl now takes only the rank-major layout, asserting its shape.
• get_warmup_keys skips a pooled top-k (kpool indexers) that isn't a multiple of 512. The merge rejects such a top-k at runtime anyway, so warmup no longer compiles a kernel that can never run.
• Related open PRs touching the same code:
  • [Bugfix][Attention] Reserve the sparse indexer's DCP top-k merge scratch through the workspace manager #59367 gathers into a rank-major buffer, then copies it to row-major for the selector. With this PR the selector could read that buffer directly and drop the copy.
  • [Perf][DCP] Sparse indexer: key only significant bits in the DCP top-k merge #64129 changes the selector's key bits in the same kernel and merge function; whichever lands second will need a small rebase.

---

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.

[GirasoleY]

Looks good. Thanks for the improvements!

[github-actions]

✅ @mickaelseznec, 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.

[GirasoleY]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99330 for commit b23ca0ce11f8.