Module tp 

GLM-5.3 DSA tensor-parallel shard plan — MLA heads sharded, indexer replicated.

Shapes measured from LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3 layer 3. Same pure-plan approach as crate::layers::glm5next_kda::tp: the GPU copy is tp_shard reuse, but a wrong head range yields a running model with mixed-up heads, so the row arithmetic is data and gets proven without a GPU.

🔴 Why the indexer is REPLICATED, not sharded

index_n_heads = 32 divides cleanly at TP=2, so sharding looks free. It is not.

The indexer emits a token selection, not a partial sum. index_scores accumulates weights[h] * relu(scale · dot) as a plain sum over heads, so a head-sharded indexer gives each rank only a partial score. If each rank then takes its own top-k, the two ranks attend to different tokens — no crash, no shape error, just a wrong answer. Sharding is therefore only correct with an all-reduce of the score tensor before the top-k.

That reduce is the problem. Scores are [q_rows, n_pools] with n_pools = seq / kpool, so at a 262 144-token context one decode token costs 65 536 × 4 B = 256 KB per layer — about 2.8 MB/token across the 11 DSA layers, on a critical path the DS4F performance review measured as latency-bound (86 × 8 KB collectives/token). The weight saved by sharding the indexer is a small fraction of one 249.8 MB layer, once. The trade loses.

A second hazard argues the same way: the reference pins a deterministic tiebreak (higher score, then smaller pool index) precisely because torch.topk’s tie order is undefined. An all-reduced score changes the summation order and can flip a tie that straddles the cutoff.

So: indexer replicated, MLA heads sharded, o_proj row-parallel with the single all-reduce qwen3_attention already performs.

🪤 Traps this module encodes

• kv_a_proj_with_mqa is REPLICATED. It produces the shared latent KV that every head decompresses from — it is the MQA part of MLA and has no head axis. Sharding it starves each rank of half the latent.
• q_a_proj / kv_a_layernorm / q_a_layernorm are replicated — low-rank down-projections and their norms, no head structure. Same failure mode as KDA’s f_a/g_a.
• q_b_proj and kv_b_proj shard by head, at qk_head_dim and nope + v_dim per head respectively. The two strides differ; using one for the other silently mixes heads.
• o_proj is row-parallel on heads * v_head_dim. Column-slicing gives a plausible, wrong output.

Structs

DsaTensorPlan

DsaTpPlan

Per-rank shard plan for one DSA block.

Enums

DsaShard

How one DSA tensor maps onto TP ranks.