Skip to content

vllm.models.flint.ops.norm

Flint's RMSNorm, fused residual-add RMSNorm and SiLU-and-mul (Triton).

Numerics of the reference implementation's CUDA norm and activation kernels, everything in fp32 with one output rounding:

  • rms_norm: bf16(x * rsqrt(mean(x^2) + eps) * w).
  • fused_add_rms_norm: s = x + residual in fp32, residual_out = bf16(s), out = bf16(s * rsqrt(mean(s^2) + eps) * w) from the unrounded s.
  • silu_and_mul: bf16(g / (1 + exp(-g)) * u) (precise expf and division).

rsqrt is rsqrt.approx.ftz.f32 (rsqrtf under fast math). The fused Flint kernels (qkv_fused, the decode MoE block's norm and shared expert, and the GEMV SiLU epilogue) follow the same rules; the routed experts use the approximate ex2 / rcp SiLU of the TRTLLM MoE (see moe).

Classes:

  • RMSNorm –

    Holds a norm weight; rms_norm(x), or with a residual

Functions:

RMSNorm

Bases: Module

Holds a norm weight; rms_norm(x), or with a residual fused_add_rms_norm(x, residual).

Source code in vllm/models/flint/ops/norm.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class RMSNorm(nn.Module):
    """Holds a norm ``weight``; ``rms_norm(x)``, or with a residual
    ``fused_add_rms_norm(x, residual)``."""

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

    def forward(
        self, x: torch.Tensor, residual: torch.Tensor | None = None
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            return rms_norm(x, self.weight, self.variance_epsilon)
        return fused_add_rms_norm(x, residual, self.weight, self.variance_epsilon)

fused_add_rms_norm(x, residual, weight, eps)

In place: residual = bf16(x + residual), x = rms_norm(x + residual); returns (x, residual).

Source code in vllm/models/flint/ops/norm.py 
82
83
84
85
86
87
88
89
90
def fused_add_rms_norm(
    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> tuple[torch.Tensor, torch.Tensor]:
    """In place: ``residual = bf16(x + residual)``, ``x = rms_norm(x +
    residual)``; returns ``(x, residual)``."""
    assert residual.shape == x.shape and residual.is_contiguous()
    assert x.is_contiguous()
    _launch_norm(x, residual, weight, eps, x)
    return x, residual

silu_and_mul(x)

silu(x[:, :D]) * x[:, D:] for x [M, 2D].

Source code in vllm/models/flint/ops/norm.py 
103
104
105
106
107
108
109
110
111
112
113
def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
    """``silu(x[:, :D]) * x[:, D:]`` for ``x`` [M, 2D]."""
    m, d2 = x.shape
    d = d2 // 2
    assert x.dtype == torch.bfloat16 and x.stride(-1) == 1
    bd = min(d, 1024)
    assert d % bd == 0
    out = torch.empty(m, d, dtype=x.dtype, device=x.device)
    if m > 0:
        _silu_and_mul_kernel[(m, d // bd)](x, out, x.stride(0), D=d, BD=bd, num_warps=4)
    return out