vllm.models.beam.amd.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/amd/layernorm.py

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

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

    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.

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

def beam_norm_residual_norm(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_weight: torch.Tensor | None,
    post_eps: float | None,
    next_weight: torch.Tensor | None,
    next_eps: float | None,
    scales: tuple[float, float] | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Fuse post-norm, scaled residual add and next pre-norm, keeping each cast."""
    assert x.ndim == residual.ndim == 2 and x.shape == residual.shape
    assert x.stride(-1) == residual.stride(-1) == 1
    hidden = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    normalized = torch.empty_like(hidden) if next_eps 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 or 0.0,
            next_eps or 0.0,
            *(scales if scales is not None else (1.0, 1.0)),
            post_eps is not None,
            post_weight is not None,
            next_eps is not None,
            next_weight is not None,
            scales 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/amd/layernorm.py

def beam_rms_norm(
    x: torch.Tensor, weight: torch.Tensor | None, 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,
        weight is not None,
        triton.next_power_of_2(n),
        num_warps=4 if n <= 1024 else 8,
        enable_fp_fusion=False,
    )
    return out