kda_chunked_prenorm

Function kda_chunked_prenorm 

Source
pub fn kda_chunked_prenorm(
    qn: &[f32],
    kn: &[f32],
    v: &[f32],
    gate: &[f32],
    beta: &[f32],
    dims: KdaDims,
    chunk: usize,
    state: &mut [f32],
) -> Vec<f32>
Expand description

Same chunked formulation, but q/k are already L2-normalised — Atlas’s contract, where the conv path (fused on decode, l2_norm_bf16 on prefill) has already normalised them.

Split out for the same reason as kda_recurrent_prenorm: re-normalising an already-unit vector is nearly a no-op in fp32 but on a bf16-rounded vector it RESTORES the norm the rounding destroyed, which makes the reference disagree with the kernel and looks exactly like a kernel bug. Atlas’s prefill L2 writes bf16, so this path is always the bf16 case.