Module tp 

GLM-5.3 KDA tensor-parallel shard plan — head-parallel, one all-reduce.

Scoped to LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3. Shapes below are measured from the checkpoint’s safetensors headers, not inferred.

Why a pure plan

The GPU copy is three lines of crate::tp_shard reuse. The part that is easy to get silently wrong is which rows belong to this rank — and a wrong head range produces a running model with quietly mixed-up heads, not a crash. So the row arithmetic lives here as pure data that can be proven without a GPU, exactly like the EP residency proof.

The pattern this follows

Qwen3.5 GDN HeadParallel (weight_loader/qwen35/load_layers/linear_attn_arms.rs, helpers in tp_shard/gdn.rs): each rank owns a contiguous head range, out_proj is row-parallel, and one all-reduce follows it. KDA differs from GDN in two ways that matter here:

• KDA’s q/k/v_proj and q/k/v_conv1d are separate tensors on disk — GDN fuses them into in_proj_qkv / conv1d. So KDA needs no segmented slice: each tensor is sliced independently and the 3-segment trap (crate::tp_shard::gdn::segment_copy_plan) simply does not arise.
• KDA has no Z tensor. The output gate is low-rank g_a/g_b.

🪤 Traps this module encodes

• a_log is per-HEAD [64]; dt_bias is per-CHANNEL [8192]. The KDA module doc already calls this “the highest-risk line”. Under TP they shard at different granularity — heads vs heads×head_dim. Slicing dt_bias by head count silently keeps 1/128th of the right data.
• o_norm is [head_dim], NOT [heads*head_dim] — it is per-channel-within- a-head and therefore replicated, never sharded. It is 256 B; a “shard everything that looks per-head” rule corrupts it.
• f_a/g_a are down-projections [rank, hidden] — replicated. Only the _b up-projections carry head structure. Sharding an _a splits the low-rank bottleneck and every head reads a truncated gate.
• o_proj is row-parallel: [hidden, heads*head_dim] sliced on its INPUT dim. Each rank produces a partial [hidden] that is only correct after the all-reduce. Column-slicing it instead yields a plausible, wrong output.

Structs

KdaTensorPlan

One tensor’s placement. row_elems is the width of a row in elements; for KdaShard::ChannelCols the roles invert and row_elems is the sharded axis.

KdaTpPlan

The complete per-rank shard plan for one KDA block.

Enums

KdaShard

How one KDA tensor maps onto TP ranks.