Module build 

2026-09-25: Loading one DSA block for this rank: tensor-parallel sharding plus four load-time weight transforms.

Owner: model-arch (GLM-5.3 DSA). Invariants: none beyond the types.

build_dsa_weights takes a load closure that returns a checkpoint tensor by layer-relative name, so the transforms are testable without a WeightStore.

The transforms, each done once at load

1. q_absorb (absorb_q): q_b_proj multiplied through kv_b_proj’s K half, so Q is in the kv_lora_rank-wide latent space the decode kernel reads.
2. weights_proj scaled by index_heads^-0.5: dsa_index_scores does not apply the factor, and a positive scale folds into the weight exactly.
3. o_absorb (absorb_o): o_proj multiplied through kv_b_proj’s V half. The decode kernel’s output is [local_heads, kv_lora_rank] in latent space, while the checkpoint o_proj takes [local_heads, v_head_dim] (256 per head against 512 on GLM-5.3).
4. ape uploaded as f32: the dsa_kpool_compress parameter is const float* and the checkpoint stores BF16.

Functions

absorb_o

2026-09-25: The output-side absorption: o_absorb[i][h*kvl + c] = Σ_r o_proj[i][h*vd + r] · kv_b[h*(nope+vd) + nope + r][c].

absorb_q

2026-09-25: q_absorb[h*kvl + c][k] = Σ_r kv_b[h*(nope+vd) + r][c] · q_b[h*nope + r][k] over all full_heads heads, computed on the host.

build_dsa_weights

2026-09-25: Load, transform, shard and upload one DSA block for this rank.

shard_host

2026-09-25: Apply one tensor’s shard plan to full host values.

Type Aliases

LoadFn

2026-09-25: Returns a checkpoint tensor by layer-relative name: the full (unsharded) BF16 values, widened to f32 on the host.