Implementation:InternLM Lmdeploy RmsNorm

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, Normalization
Last Updated	2026-02-07 15:00 GMT

Overview

CUDA kernel declarations for RMS (Root Mean Square) normalization and fused residual-bias-RMSNorm operations.

Description

This header provides the API for RMSNorm operations used in LLaMA-family and similar transformer architectures. invokeRMSNorm() applies standard RMS normalization to a tensor using learned weights and an epsilon parameter. invokeRMSNormQK() applies RMSNorm in-place, designed for normalizing query and key projections. invokeBiasResidualRMSNorm() fuses bias addition, residual connection, and RMSNorm into a single kernel for efficiency. invokeResidualBiasRMSNorm() is a type-erased variant that accepts DataType at runtime. ApplyBias() applies bias with optional offset indexing and scaling.

Usage

Use these kernels in transformer layer normalization, particularly in models that use RMSNorm (e.g., LLaMA, InternLM). The fused variants reduce memory bandwidth by combining residual addition and normalization in a single pass.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/norm/rms_norm.h

Signature

void invokeRMSNorm(Tensor& out, const Tensor& x, const Tensor& w, float eps, cudaStream_t st);

void invokeRMSNormQK(Tensor& x, const Tensor& w, float eps, cudaStream_t st);

template<class T>
void invokeBiasResidualRMSNorm(
    T* residual, T* hidden_states, const T* weights, const T* bias,
    int dims, int num, float eps, cudaStream_t st);

void invokeResidualBiasRMSNorm(
    void* hidden_states, void* residual, const void* weights, const void* bias,
    DataType dtype, int dims, int num, float eps, cudaStream_t st);

void ApplyBias(Tensor& x, const Tensor& bias, cudaStream_t st);

Import

#include "src/turbomind/kernels/norm/rms_norm.h"

I/O Contract

Inputs

Name	Type	Required	Description
x	Tensor / T*	Yes	Input tensor or hidden states
w / weights	Tensor / const T*	Yes	Learned RMSNorm weight parameters
bias	const T*	No	Bias vector (for fused variants)
residual	T*	No	Residual connection tensor (fused variants)
eps	float	Yes	Epsilon for numerical stability
dims	int	Yes	Hidden dimension size
num	int	Yes	Number of tokens

Outputs

Name	Type	Description
out / hidden_states	Tensor / T*	Normalized output (may be in-place)
residual	T*	Updated residual (for fused variants)

Usage Examples

using namespace turbomind;

// Standard RMSNorm
invokeRMSNorm(output, input, weights, 1e-6f, stream);

// Fused bias + residual + RMSNorm
invokeBiasResidualRMSNorm(
    residual_ptr, hidden_ptr, weight_ptr, bias_ptr,
    hidden_dim, num_tokens, 1e-6f, stream);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime