Skip to main content

rms_norm_backward

teeny_kernels::nn::norm::rmsnorm

Function rms_norm_backward 

Source 
pub fn rms_norm_backward<T: Triton, D: Float, const BLOCK_N: i32>(
    dy_ptr: T::Pointer<D>,
    x_ptr: T::Pointer<D>,
    dx_ptr: T::Pointer<D>,
    weight_ptr: T::Pointer<D>,
    dweight_ptr: T::Pointer<D>,
    rrms_ptr: T::Pointer<D>,
    _M: i32,
    N: i32,
)
where
    T::I32Tensor: Tensor<i32, 1> + Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
Expand description

RMSNorm backward pass.

dx[m,n] = rrms[m] * γ[n] * (dy[m,n] - x[m,n] * rrms[m]² * Σ_n dy[m,n]*γ[n]*x[m,n] / N)
dweight[n] = Σ_m dy[m,n] * x[m,n] * rrms[m]

Grid: [M] — one CTA per row.