Skip to content

vllm.models.minimax_m3.nvidia.ops.norm_mxfp8

Gemma RMSNorm (optionally fused with the residual add) in one Triton pass.

With an sf scale buffer the output is E4M3 data plus UE8M0 scales in the 128x4 swizzled layout, replacing the BF16 norm followed by a separate MXFP8 quantize (which re-reads the full normed activation). Without it, the output is the plain BF16 norm. The residual is bitwise the BF16 norm's; the quantized data and scales match the reference quantize within one E4M3 step (the rstd reduction order differs).

Functions:

gemma_rmsnorm_mxfp8(x, weight, eps, residual=None, mxfp8=True)

GemmaRMSNorm of x (plus residual in place when given).

Parameters:

  • x

    (Tensor) –

    (num_tokens, hidden) BF16 activation, contiguous along the row (any row stride). Read only.

  • weight

    (Tensor) –

    (hidden,) norm weight; the norm scales by (1 + weight).

  • eps

    (float) –

    RMS norm epsilon.

  • residual

    (Tensor | None, default: None ) –

    same shape/dtype/stride(1) as x; updated in place to x + residual. None for a plain norm.

  • mxfp8

    (bool, default: True ) –

    return MXFP8 (E4M3 data + UE8M0 scales) or the plain BF16 normed activation.

Returns:

  • Tensor –

    (data, scale): with mxfp8=True, E4M3 data (num_tokens,

  • Tensor | None –

    hidden) and UE8M0 scales in the 128x4 swizzled layout with zeroed

  • tuple[Tensor, Tensor | None] –

    padding rows; with mxfp8=False, the BF16 normed activation and

  • tuple[Tensor, Tensor | None] –

    None.

Source code in vllm/models/minimax_m3/nvidia/ops/norm_mxfp8.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def gemma_rmsnorm_mxfp8(
    x: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    residual: torch.Tensor | None = None,
    mxfp8: bool = True,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """GemmaRMSNorm of ``x`` (plus ``residual`` in place when given).

    Args:
        x: (num_tokens, hidden) BF16 activation, contiguous along the row
            (any row stride). Read only.
        weight: (hidden,) norm weight; the norm scales by (1 + weight).
        eps: RMS norm epsilon.
        residual: same shape/dtype/stride(1) as x; updated in place to
            x + residual. None for a plain norm.
        mxfp8: return MXFP8 (E4M3 data + UE8M0 scales) or the plain BF16
            normed activation.

    Returns:
        ``(data, scale)``: with ``mxfp8=True``, E4M3 data (num_tokens,
        hidden) and UE8M0 scales in the 128x4 swizzled layout with zeroed
        padding rows; with ``mxfp8=False``, the BF16 normed activation and
        None.

    """
    assert x.dim() == 2 and x.stride(1) == 1, x.shape
    assert weight.is_contiguous()
    assert residual is None or (residual.shape == x.shape and residual.stride(1) == 1)
    M, H = x.shape
    sf = None
    if mxfp8:
        padded_m = round_up(M, SF_ROW_TILE)
        out = torch.empty((M, H), dtype=torch.uint8, device=x.device)
        sf = torch.empty(
            padded_m * round_up(H // MXFP8_BLOCK, 4), dtype=torch.uint8, device=x.device
        )
        if M == 0:
            return out.view(torch.float8_e4m3fn), sf
    else:
        padded_m = M
        out = torch.empty((M, H), dtype=torch.bfloat16, device=x.device)
        if M == 0:
            return out, None
    _gemma_rmsnorm_mxfp8_kernel[(padded_m,)](
        x,
        residual,
        weight,
        out,
        sf,
        M,
        x.stride(0),
        residual.stride(0) if residual is not None else x.stride(0),
        eps,
        H=H,
        launch_pdl=current_platform.is_arch_support_pdl(),
        num_warps=4,
    )
    if mxfp8:
        return out.view(torch.float8_e4m3fn), sf
    return out, None