[avigyabb]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds an opt-in flag that splits KPool indexer scoring across the tensor-parallel group by query rows, so TP ranks stop recomputing the same O(queries x pools) scores; the selected indices are unchanged.

Follows up #7522 (by @HollowMan6), which added the KPool DSA indexer.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Mechanism

Under sequence parallelism, DSAIndexer.forward_before_topk gathers x and qr over the TP group before projecting, and the indexer projections are duplicated across TP ranks (average_gradients_across_tp_domain; without the indexer loss they are frozen). Without sequence parallelism the activations are already replicated across TP. Either way, every TP rank calls fused_qk_topk_kpool with the same q, k, weights, gate scores and mask/varlen bounds, and every rank computes the same full [batch, queries, pools] score matrix and top-k. At TP=8, seven of the eight scorings are redundant. That cost grows with the square of the sequence length, so it takes up more of the step as context grows.

Each query row's pool selection depends only on that row and the pooled keys. The rows can therefore be split across ranks and the results concatenated, with every row matching the unsharded computation.

Change

• fused_qk_topk_kpool(..., query_shard_group=None): when a group is passed (requires return_index_scores=False), rank r scores rows [r*ceil(sq/G), min(sq, (r+1)*ceil(sq/G))) through the existing chunked path _kpool_topk_from_score_chunks. That path gains query_start/query_end and keeps its _KPOOL_SCORE_CHUNK_BYTES (2 GiB) cap on the score buffer per chunk. The local pool IDs are cast to int32, padded to equal length, combined with one all_gather_into_tensor, and reassembled into the full [batch, sq, budget] tensor. Pool expansion, tail append and mask filtering then run unchanged on the full tensor, so every rank returns the same indices.
• New TransformerConfig.dsa_indexer_kpool_shard_queries_across_tp: bool = False (CLI --dsa-indexer-kpool-shard-queries-across-tp). When it is set and TP > 1, DSAttention passes the TP group for the KPool top-k. DSAIndexer.forward_with_scores passes it when return_index_scores=False, which is the path DSAIndexer.forward uses.
• DSAttention skips sharding when the attention query rows are already sequence-local per TP rank (sequence_parallel_query_is_local). In that case q/weights are already sliced to the local rows and differ across ranks.
• Default off. With the flag off, behavior is unchanged: query_start=0, query_end=sq, and no collective.

Communication per indexer call: one all-gather of batch * ceil(sq/TP) * (index_topk / kpool) int32 per rank. At 512k tokens, index_topk=2048, kpool=4 and TP=8, that is 128 MiB per rank.

Alternatives considered

• Keep the indexer replicated (status quo). This is simplest, but the redundant scoring grows with the square of the sequence length.
• Shard the keys/pools instead of the queries. This needs a distributed top-k (a local top-k plus a merge of TP x budget candidates per row) and changes tie-breaking. Splitting by query needs only a concatenation and keeps results identical.
• All-gather the expanded token indices instead of pool IDs. The payload would be kpoolx larger (and int64). Gathering the pool IDs and expanding after the gather moves the least data.
• Turn it on unconditionally. It adds one collective per indexer call. At short sequences the scoring is cheap and the collective may not pay for itself, so it is opt-in.

Test plan

• New tests/unit_tests/transformer/experimental_attention_variant/test_kpool_query_shard.py (43 cases, CPU-runnable, uses CUDA if available). It simulates every rank of a group in one process: a fake group provides size()/rank(), and torch.distributed.all_gather_into_tensor is replaced by one that returns the shards every simulated rank produced. Each simulated rank's output is compared bitwise to the unsharded call:
  • causal explicit mask, shard sizes 2/3/8, sq in {5, 24, 37} (ragged last shard, and empty shards when sq < TP), batch 1/2, with and without a small _KPOOL_SCORE_CHUNK_BYTES so shard edges fall inside the chunk loop;
  • packed THD (cu_seqlens_kv + varlen bounds), shard sizes 2/4;
  • each rank scores only ceil(sq/TP) rows;
  • the shard bounds cover every row once;
  • query_shard_group with return_index_scores=True raises.
• tests/unit_tests/models/test_hybrid_moe_model.py: added the new field to GOLDEN_CONFIG. Not run locally because it needs the distributed GPU harness.

Run locally (CPU only, CUDA_VISIBLE_DEVICES=""):

pytest tests/unit_tests/transformer/experimental_attention_variant/test_kpool_query_shard.py
43 passed

Formatting: black --skip-magic-trailing-comma --skip-string-normalization (26.3.0), isort (5.13.2), ruff check (0.9.x) and pylint (3.2.6, 10.00/10) on the changed files.

Not run here: the multi-GPU DSA unit tests and a real NCCL run of this exact port.

Validation (from SkyRL)

The same change, with the same contiguous query split and the same int32 all-gather, has been running in SkyRL on a vendored copy of the #7522 KPool code: NovaSky-AI/SkyRL#2422. The model is GLM-5.3-Flash, run on 64 B200 (8 nodes x 8) with TP8:

Measurement	Unsharded	Sharded (TP8)
KPool selection microbenchmark, 512k tokens, one B200, per rank	21.07 s	2.63 s
Full model, 256k tokens/seq, steady-state training step	218 s	102 s (2.1x)
Full model, 256k tokens/seq, peak allocated memory	98.5 GiB	98.4 GiB
Full model, 512k tokens/seq, warm fwd+bwd (#2422)	1545 s	1094 s

The 256k step times and the microbenchmark come from SkyRL's long-context benchmark runs. Those runs used SkyRL-specific settings, including a per-step torch.cuda.empty_cache(), so the absolute step times will differ elsewhere. The speedup depends on how much of the step the indexer takes, which grows with sequence length. In SkyRL's 2-rank NCCL test (test_dsa_kpool_tp_shard.py in #2422), sharded selections were bitwise identical to a single rank.

Risks / notes for reviewers

• The change relies on TP ranks having identical indexer inputs. This holds for the gathered SP path and for non-SP TP, and the sequence-local-query case is excluded. If ranks ever diverged numerically, the gather would make all ranks use the same selection, rather than each rank keeping its own.
• Only the KPool path (dsa_indexer_kpool > 1) is covered. The per-token indexer (fused_qk_topk_naive / fused kernels) is unchanged.

Issue tracking

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.