Function rms_norm 

pub fn rms_norm(
    x: &[f32],
    w: &[f32],
    tokens: usize,
    dim: usize,
    eps: f32,
) -> Vec<f32>

RMSNorm.forward: w * (x * rsqrt(mean(x^2) + eps)) in f32, back to bf16.