Module build 

Loading one DSA block: TP sharding plus the three load-time transforms the runtime cannot do per token.

Takes a load closure yielding an uploaded BF16 tensor by layer-relative name, rather than a WeightStore, so the transforms are testable and the loader wiring stays one call site.

The three transforms, and why each is here rather than in decode

1. q_absorb — q_b_proj pre-multiplied by kv_b_proj’s K half, so Q arrives in the 512-dim latent space the decode kernel dots against. Doing it per token would be a second GEMM on the critical path for a weight that never changes.
2. weights_proj scaled by index_heads^-0.5 — dsa_index_scores does not apply the factor. Folding it into the weight is exact (a positive scalar) and free.
3. o_absorb — o_proj pre-multiplied by kv_b_proj’s V half, the counterpart of q_absorb and the half that was missing. The decode kernel consumes the latent KV directly, so its output is [local_heads, kv_lora_rank] in LATENT space; the raw checkpoint o_proj expects [local_heads, v_head_dim] in V space. Feeding one to the other is not a numeric drift — the GEMM reads kv_lora_rank / v_head_dim = 2x past the end of every weight row. Measured 2026-08-28: CUDA_ERROR_ILLEGAL_ADDRESS at layer 3, grid=[256,1,1] block=[16,16,1], on the first forward that reached it.
4. ape upconverted BF16 → F32 — the kernel’s parameter is const float* while the checkpoint stores BF16. This is the #341/#347 dtype-mismatch class: reading it at the wrong width is silent.

Functions

absorb_o

o_absorb[i][h*kvl + c] = Σ_r o_proj[i][h*vd + r] · kv_b[h*(nope+vd) + nope + r][c].

absorb_q

q_absorb[h*kvl + c][k] = Σ_r kv_b[h*(nope+vd) + r][c] · q_b[h*nope + r][k].

build_dsa_weights

Bind one DSA block for this rank.

shard_host

Apply one tensor’s shard plan to full host values.

Type Aliases

LoadFn

A tensor as the checkpoint holds it: full (unsharded) BF16 values on the host.