[avigyabb]

Stacked on #2413 (chunked k-pool indexer scoring): the shards walk that PR's chunk loop. Review only the top commit.

Summary

Under sequence parallelism every TP rank runs the k-pool indexer on the same gathered sequence with the same frozen, replicated indexer weights, so at TP8 all 8 ranks computed identical O(sq^2) pool scores (every query scores every pool). SKYRL_DSA_INDEXER_TP_SHARD=1 passes the TP group as query_shard_group. Each rank scores and selects pools only for its contiguous slice of query rows, then the int32 selections are all-gathered. Per query row nothing changes, so selections are bitwise identical. Opt-in; GLM-5.3-Flash only.

Full GLM-5.3-Flash, 64 B200 (TCP NCCL), 512k tokens/seq, TP8: warm fwd+bwd 1545 -> 1094 s, i.e. ~1/3 of the step was replicated indexer scoring (sharding 8 ways removes 7/8 of it).

Testing

• gpu_ci/patches/megatron/mcore_ext/test_dsa_kpool_tp_shard.py (two ranks, ragged last shard, many score chunks: identical to unsharded), together with A's test_dsa_kpool.py: 7 passed on this branch (B200)

#2418 (context parallelism) is stacked on this PR.

Split out of #2341.

🤖 Generated with Claude Code

---

Note

Medium Risk
Opt-in distributed path changes which tokens the sparse DSA indexer selects if sharding or gather is wrong; selections are intended to be bitwise-identical and are GPU-tested, but errors would not raise.

Overview
Adds an opt-in (SKYRL_DSA_INDEXER_TP_SHARD=1) optimization for GLM-5.3-Flash k-pool DSA: tensor-parallel ranks no longer each run the full O(sq×pools) indexer scoring on identical gathered inputs.

glm5_next/dsa.py passes the TP process group into fused_qk_topk_kpool as query_shard_group. In mcore_ext/dsa_kpool.py, each rank scores and top-k’s only its contiguous slice of query rows (the chunked path gains q_lo/q_hi), pads the last shard when lengths don’t divide evenly, all-gathers int32 pool ids, and reassembles a full [batch, sq, budget] tensor—per-row behavior stays the same as unsharded.

Docs in the Megatron patches README note the deviation and the new GPU test. test_dsa_kpool_tp_shard.py runs a 2-rank NCCL check (ragged shard + small score chunks) that sharded selection matches a single rank bitwise.

^{Reviewed by Cursor Bugbot for commit d0637d4. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces tensor-parallel query sharding for the k-pool indexer, allowing query rows to be split across the tensor-parallel group to reduce redundant scoring overhead. The changes include adding an opt-in environment variable SKYRL_DSA_INDEXER_TP_SHARD, updating the fused_qk_topk_kpool function to shard queries and gather selections, and adding a multi-GPU unit test to verify exact matching. A critical issue was identified in the all_gather_into_tensor implementation where a dimension mismatch between the input and output tensors will cause a runtime error; reshaping the gathered tensor as suggested is required to fix this.

[gemini-code-assist]

The all_gather_into_tensor API expects the output tensor to have the same number of dimensions as the input tensor, with the 0th dimension sized to world_size * input_tensor.size(0). Currently, local has 3 dimensions (batch, rows_per_rank, budget), but gathered is initialized with 4 dimensions (shard_size, batch, rows_per_rank, budget). This will cause a RuntimeError at runtime whenever batch > 1 (or even batch = 1 on strict PyTorch versions) due to dimension and size mismatches.\n\nTo fix this, initialize gathered with 3 dimensions (shard_size * batch, rows_per_rank, budget) and then use .view() to reshape it before permuting.

            gathered = torch.empty((shard_size * batch, rows_per_rank, budget), dtype=torch.int32, device=q.device)\n            torch.distributed.all_gather_into_tensor(gathered, local.contiguous(), group=query_shard_group)\n            local = gathered.view(shard_size, batch, rows_per_rank, budget).permute(1, 0, 2, 3).reshape(batch, shard_size * rows_per_rank, budget)[:, :sq]

[greptile-apps]

Confidence Score: 4/5

[High risk] Adds optional tensor-parallel sharding to the k-pool indexer query scoring.

The PR appears safe to merge, though the opt-in attention wiring should receive integration coverage.

Findings

1.  Production wiring remains untested ▶

_{Reviews (1) · Last reviewed commit: "[megatron] GLM-5.3-Flash DSA: shard the ..."}

[greptile-apps]

Production wiring remains untested. This test passes a process group directly to fused_qk_topk_kpool, but production must read SKYRL_DSA_INDEXER_TP_SHARD and pass Megatron's TP group through the attention path. The test would still pass if that wiring stopped working, leaving redundant scoring or a collective failure undetected. Please add coverage for the flag-enabled attention path.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!