pub fn gated_rms_norm_prefill(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
input: DevicePtr,
gate: DevicePtr,
weight: &DenseWeight,
output: DevicePtr,
heads_per_token: u32,
head_dim: u32,
eps: f32,
num_actual_tokens: u32,
input_token_stride: u32,
gate_token_stride: u32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: Gated RMS norm of every (head, token) row of a prefill chunk in
one launch. input and output rows of token t start at
t * input_token_stride, and gate rows at t * gate_token_stride, both in
BF16 elements.