rms_norm_warp_row

Function rms_norm_warp_row 

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

Warp-per-row RMS norm for SHORT rows — one warp per row instead of one block, so the grid shrinks 8x and the reduction needs no shared memory or barrier. Profitable exactly for the Qwen3 per-head q_norm/k_norm during prefill (num_rows = heads * seq, hidden_size = head_dim), where the block-per-row kernel measured ~43x above its bandwidth floor.