gated_rms_norm

Function gated_rms_norm 

Source
pub fn gated_rms_norm(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    input: DevicePtr,
    gate: DevicePtr,
    weight: &DenseWeight,
    output: DevicePtr,
    num_tokens: u32,
    hidden_size: u32,
    gate_stride: u32,
    eps: f32,
    group_size: u32,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: Gated RMS norm of num_tokens rows, one block per row; the gate row of token t starts at t * gate_stride. The order of norm and gate is the resolved kernel’s. gated_rms_norm and gated_rms_norm_f32_input in kernels/gb10/common/rms_norm.cu compute rms_norm(input) * weight * silu(gate) and ignore group_size. The per-group shadows, e.g. kernels/gb10/minimax-m2-229b/nvfp4/rms_norm.cu, gate first and normalise each group_size slice.