absorb_o

Function absorb_o 

Source
pub fn absorb_o(
    cfg: &Glm5NextDsaConfig,
    o_local: &[f32],
    kv_b_local: &[f32],
    local_heads: usize,
) -> Result<Vec<f32>>
Expand description

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

The output-side twin of absorb_q. Absorbed MLA is a PAIR of transforms — Q into the latent space on the way in, the output projection back out of it on the way out — and shipping only the first leaves the decode kernel’s latent output being read by a V-space weight.

🪤 Unlike absorb_q, this one is done AFTER sharding, and that is not an optimisation that happens to be safe — it is exact. absorb_q pairs q_b head h with kv_b head h, so slicing before pairing would cross heads. Here both operands are indexed by the SAME h and the head axis is a plain outer sum, so this rank’s heads never touch another rank’s rows. Doing it on full heads would double the load-time cost for an identical result.

🪤 kv_b’s V half starts at nope, not 0. Using the K half compiles, runs, and gives a well-formed wrong answer — the same trap absorb_q documents from the other side.