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

Gated RMS norm (norm_before_gate=False, per-group): output = rms_norm_per_group(input * silu(gate), weight, group_size)

Kernel: gated_rms_norm(input, gate, weight, output, hidden_size, eps, gate_stride, group_size) Grid: (num_tokens, 1, 1) Block: (min(hidden_size, 1024), 1, 1)