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)