Implementation:Turboderp org Exllamav2 ExLlamaV2HeadNorm

Knowledge Sources	Turboderp_org_Exllamav2
Domains	Normalization, Attention
Last Updated	2026-02-15 00:00 GMT

Overview

Per-attention-head normalization module that applies either RMSNorm or LayerNorm independently to each attention head's output, with CUDA-accelerated and pure-PyTorch forward paths.

Description

ExLlamaV2HeadNorm is a subclass of ExLlamaV2Module that performs normalization at the granularity of individual attention heads rather than across the full hidden dimension. This is used by architectures that require per-head normalization (e.g., certain Cohere or custom model variants).

Key components:

__init__(model, key, num_heads, head_dim, archparams) -- Initialises the module with the number of attention heads and per-head dimension. Sets variance_epsilon from the model config's norm_eps.
load() -- Loads the normalization weight (and optional bias) tensor. If the loaded weight has shape (head_dim,) (a single head's worth), it is repeated across num_heads to produce shape (num_heads, head_dim). Optionally adds archparams.norm_constant_bias to the weight. Creates an nn.LayerNorm instance for the PyTorch fallback path.
unload() -- Releases the layernorm, weight, and bias tensors.
forward(hidden_states, ...) -- The primary forward path using ext_c.head_norm() CUDA kernel. Passes the weight, optional bias, variance epsilon, and a boolean indicating whether to use RMSNorm mode (when archparams.headnorm == "rmsnorm") versus standard LayerNorm.
forward_torch(hidden_states, ...) -- Pure PyTorch fallback. In RMSNorm mode, computes variance and applies reciprocal square root scaling without mean subtraction. In LayerNorm mode, subtracts the mean before computing variance and scaling.

Both forward paths accept an intermediates flag; when True, they return a dictionary with a hidden_states key instead of a bare tensor.

Usage

Use ExLlamaV2HeadNorm as a layer within transformer attention blocks that require per-head normalization. It is instantiated automatically by the model architecture loader when the model's config specifies head-level normalization. Users do not typically create this module directly.

Code Reference

Source Location

Repository: Turboderp_org_Exllamav2
File: exllamav2/headnorm.py
Lines: L1-177

Signature

class ExLlamaV2HeadNorm(ExLlamaV2Module):

    name: str = "LayerNorm"
    layernorm: nn.LayerNorm | None
    weight: nn.Parameter | None
    bias: nn.Parameter | None
    variance_epsilon: float
    head_dim: int
    num_heads: int

    def __init__(
        self,
        model: ExLlamaV2,
        key: str,
        num_heads: int,
        head_dim: int,
        archparams=None
    ):
        ...

    def load(self) -> None:
        ...

    def unload(self) -> None:
        ...

    def forward(
        self,
        hidden_states: torch.Tensor,
        cache=None,
        attn_params=None,
        past_len=None,
        intermediates: bool = False,
        loras=None,
        **kwargs
    ) -> torch.Tensor | dict[str, torch.Tensor]:
        ...

    def forward_torch(
        self,
        hidden_states: torch.Tensor,
        cache=None,
        attn_params=None,
        past_len=None,
        intermediates: bool = False,
        loras=None,
        **kwargs
    ) -> torch.Tensor | dict[str, torch.Tensor]:
        ...

Import

from exllamav2.headnorm import ExLlamaV2HeadNorm

I/O Contract

Inputs

Name	Type	Required	Description
model	ExLlamaV2	Yes	The parent model instance
key	str	Yes	The tensor key path for loading weights from safetensors
num_heads	int	Yes	Number of attention heads
head_dim	int	Yes	Dimension per attention head
archparams	object	No	Architecture parameters; defaults to model.config.arch.lm
hidden_states	torch.Tensor	Yes (forward)	Input tensor to normalize, shape (batch, seq_len, num_heads * head_dim)
intermediates	bool	No (default False)	If True, return a dict with "hidden_states" key instead of bare tensor

Outputs

Name	Type	Description
hidden_states	torch.Tensor	Normalized tensor of the same shape as input, with per-head normalization applied
intermediates dict	dict[str, torch.Tensor]	When intermediates=True: dictionary containing "hidden_states" key

Usage Examples

Module in an Attention Block (Internal Usage)

from exllamav2.headnorm import ExLlamaV2HeadNorm

# Typically created internally during model construction
head_norm = ExLlamaV2HeadNorm(
    model=model,
    key="model.layers.0.self_attn.q_norm",
    num_heads=32,
    head_dim=128,
)
head_norm.load()

# Forward pass
normalized = head_norm.forward(hidden_states)

PyTorch Fallback for Debugging

# Use the pure-PyTorch path (no CUDA kernel)
normalized = head_norm.forward_torch(hidden_states, intermediates=True)
print(normalized["hidden_states"].shape)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment