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.