Module tp 

2026-09-25: GLM-5.3-Flash DSA tensor-parallel shard plan: MLA heads sharded, the indexer replicated.

Owner: model-arch (GLM-5.3 DSA). Invariants:

• DsaTpPlan::new refuses tp_rank >= tp_size and a config that fails Glm5NextDsaConfig::validate.
• Every indexer tensor, kv_a_proj_with_mqa, and the low-rank down-projections and their norms are replicated; q_b_proj and kv_b_proj shard by head; o_proj is row-parallel.

Pure data, as in crate::glm5next_kda::tp, so the per-rank arithmetic is tested without a GPU.

Why the indexer is replicated

The indexer emits a token selection, not a partial sum. index_scores adds weights[h] * relu(scale * dot) over the indexer heads, so a head-sharded indexer gives each rank a partial score, and a per-rank top-k over partial scores can select different tokens on different ranks. Sharding is correct only with an all-reduce of the scores before the top-k. Scores are [q_rows, n_pools] with n_pools = seq / kpool: at a 262,144-token context (max_position_embeddings in MODEL.toml) that is 65,536 x 4 B = 256 KiB per decode token per DSA layer, against a replicated indexer of about 15 MB per layer (the indexer.* sizes in tp/tests.rs). An all-reduce would also change the order of the per-head sum, which can reorder pools whose scores tie at the cutoff.

• kv_a_proj_with_mqa produces the shared latent every head decompresses from; it has no head axis.
• q_b_proj shards at qk_head_dim rows per head and kv_b_proj at qk_nope_head_dim + v_head_dim; the two strides differ.
• o_proj is row-parallel on heads * v_head_dim.

Structs

DsaTensorPlan

DsaTpPlan

2026-09-25: Per-rank shard plan for one DSA block.

Enums

DsaShard

2026-09-25: How one DSA tensor maps onto TP ranks.