vllm.models.beam.nvidia.layernorm

Beam kernels preserving FP32 normalization and affine arithmetic.

Classes:

• RMSNorm –

  Parameters for normalization performed directly by the model kernels.

Functions:

• beam_norm_residual_norm –

  Fuse post-norm, scaled residual add and next pre-norm, keeping each cast.

• beam_rms_norm –

  Normalize strided token/head rows, rounding only after the affine multiply.

RMSNorm

Bases: Module

Parameters for normalization performed directly by the model kernels.

Source code in vllm/models/beam/nvidia/layernorm.py

class RMSNorm(nn.Module):
    """Parameters for normalization performed directly by the model kernels."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))

    def forward(self, *args, **kwargs):
        raise RuntimeError("Call the model's fused normalization kernels directly")

beam_norm_residual_norm(x, residual, post_weight, post_eps, next_weight, next_eps, scales)

Fuse post-norm, scaled residual add and next pre-norm, keeping each cast.

Returns the new residual and, when next_weight is given, its normalization for the next sublayer.

Source code in vllm/models/beam/nvidia/layernorm.py

def beam_norm_residual_norm(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_weight: torch.Tensor,
    post_eps: float,
    next_weight: torch.Tensor | None,
    next_eps: float | None,
    scales: tuple[float, float],
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Fuse post-norm, scaled residual add and next pre-norm, keeping each cast.

    Returns the new residual and, when ``next_weight`` is given, its
    normalization for the next sublayer.
    """
    assert x.ndim == residual.ndim == 2 and x.shape == residual.shape
    assert x.stride(-1) == residual.stride(-1) == 1
    assert (next_weight is None) == (next_eps is None)
    hidden = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    normalized = torch.empty_like(hidden) if next_weight is not None else None
    n = x.shape[-1]
    if hidden.numel():
        _norm_residual_norm_kernel[(x.shape[0],)](
            x,
            residual,
            post_weight,
            next_weight,
            hidden,
            normalized,
            n,
            x.stride(0),
            residual.stride(0),
            post_eps,
            next_eps or 0.0,
            *scales,
            next_weight is not None,
            triton.next_power_of_2(n),
            num_warps=4 if n <= 1024 else 8,
            enable_fp_fusion=False,
        )
    return hidden, normalized

beam_rms_norm(x, weight, eps)

Normalize strided token/head rows, rounding only after the affine multiply.

Source code in vllm/models/beam/nvidia/layernorm.py

def beam_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    """Normalize strided token/head rows, rounding only after the affine multiply."""
    assert x.ndim in (2, 3) and x.stride(-1) == 1
    n = x.shape[-1]
    heads = x.shape[1] if x.ndim == 3 else 1
    out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    if out.numel() == 0:
        return out
    _rms_norm_kernel[(x.numel() // n,)](
        x,
        weight,
        out,
        n,
        heads,
        x.stride(0),
        x.stride(-2) if x.ndim == 3 else n,
        eps,
        triton.next_power_of_2(n),
        num_warps=4 if n <= 1024 else 8,
        enable_fp_fusion=False,
    )
    return out