vllm.models.qwen3_5.nvidia.ops.fused_norm_quant

Gemma (1 + w) add-RMSNorm fused with the consumer linear's input quant.

Replaces gemma_fused_add_rmsnorm followed by a standalone static FP8 (scaled_fp8_quant) or NVFP4 (scaled_fp4_quant) kernel. The quantized values are computed from the bf16-rounded normed activation with the same scale arithmetic as those kernels, so codes match the unfused chain.

Functions:

• gemma_add_rms_norm_quant –

  Gemma RMSNorm of x (+ residual) quantized for the next linear.

_cvt_e2m1x2(lo, hi)

Pack fp32 (lo, hi) pairs into E2M1 bytes (lo in the low nibble).

Source code in vllm/models/qwen3_5/nvidia/ops/fused_norm_quant.py

@triton.jit
def _cvt_e2m1x2(lo, hi):
    """Pack fp32 (lo, hi) pairs into E2M1 bytes (lo in the low nibble)."""
    return tl.inline_asm_elementwise(
        """
        {
        .reg .b8 b0, b1, b2, b3;
        cvt.rn.satfinite.e2m1x2.f32 b0, $5, $1;
        cvt.rn.satfinite.e2m1x2.f32 b1, $6, $2;
        cvt.rn.satfinite.e2m1x2.f32 b2, $7, $3;
        cvt.rn.satfinite.e2m1x2.f32 b3, $8, $4;
        mov.b32 $0, {b0, b1, b2, b3};
        }
        """,
        "=r,r,r,r,r,r,r,r,r",
        [lo, hi],
        dtype=tl.uint8,
        is_pure=True,
        pack=4,
    )

gemma_add_rms_norm_quant(x, residual, weight, eps, scale, nvfp4, write_y=False)

Gemma RMSNorm of x (+ residual) quantized for the next linear.

Parameters:

• x

  (Tensor) –

  (num_tokens, hidden) bf16/fp16 activation, contiguous. Read only.

• residual

  (Tensor | None) –

  same shape as x, or None. Updated in place to x + residual.

• weight

  (Tensor) –

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

• eps

  (float) –

  RMS norm epsilon.

• scale

  (Tensor) –

  FP8: the static per-tensor input scale. NVFP4: the global scale (input_global_scale_inv).

• nvfp4

  (bool) –

  quantize to NVFP4 with swizzled 128x4 block scales instead of static FP8.

• write_y

  (bool, default: False ) –

  also return the bf16 normed activation.

Returns:

• Tensor –

  (q, block_scale, y): quantized data (fp8, or packed fp4 uint8),

• Tensor –

  NVFP4 block scales (fp8, (round_up(M, 128), round_up(H // 16, 4));

• Tensor | None –

  rows >= M are left unwritten) or scale for FP8, and the normed

• tuple[Tensor, Tensor, Tensor | None] –

  activation (or None).

Source code in vllm/models/qwen3_5/nvidia/ops/fused_norm_quant.py

def gemma_add_rms_norm_quant(
    x: torch.Tensor,
    residual: torch.Tensor | None,
    weight: torch.Tensor,
    eps: float,
    scale: torch.Tensor,
    nvfp4: bool,
    write_y: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Gemma RMSNorm of ``x (+ residual)`` quantized for the next linear.

    Args:
        x: (num_tokens, hidden) bf16/fp16 activation, contiguous. Read only.
        residual: same shape as x, or None. Updated in place to x + residual.
        weight: (hidden,) Gemma norm weight; the norm scales by (1 + weight).
        eps: RMS norm epsilon.
        scale: FP8: the static per-tensor input scale. NVFP4: the global
            scale (``input_global_scale_inv``).
        nvfp4: quantize to NVFP4 with swizzled 128x4 block scales instead of
            static FP8.
        write_y: also return the bf16 normed activation.

    Returns:
        ``(q, block_scale, y)``: quantized data (fp8, or packed fp4 uint8),
        NVFP4 block scales (fp8, ``(round_up(M, 128), round_up(H // 16, 4))``;
        rows >= M are left unwritten) or ``scale`` for FP8, and the normed
        activation (or None).

    """
    assert x.dim() == 2 and x.is_contiguous()
    assert x.dtype in (torch.bfloat16, torch.float16)
    assert residual is None or (residual.shape == x.shape and residual.is_contiguous())
    assert weight.shape == (x.shape[1],) and weight.is_contiguous()
    assert scale.numel() == 1 and scale.dtype == torch.float32
    M, H = x.shape
    if nvfp4:
        # Whole 4-column scale tiles: padded scale columns are not written.
        assert H % 64 == 0
        q = torch.empty((M, H // 2), dtype=torch.uint8, device=x.device)
        block_scale = torch.empty(
            (round_up(M, 128), round_up(H // 16, 4) // 4),
            dtype=torch.int32,
            device=x.device,
        ).view(FP8_DTYPE)
    else:
        q = torch.empty((M, H), dtype=FP8_DTYPE, device=x.device)
        block_scale = scale
    y = torch.empty_like(x) if write_y else None
    if M == 0:
        return q, block_scale, y

    # Swept on GB300 at H=5120, M=1..8192: within 3% of the best of 2/4/8/16
    # warps at every swept M except NVFP4 M=129 (4%).
    max_m_16_warps = 128 if nvfp4 else 384
    _gemma_add_rms_norm_quant_kernel[(M,)](
        x,
        residual if residual is not None else x,
        weight,
        scale,
        q,
        block_scale,
        y if y is not None else x,
        eps,
        H=H,
        HAS_RESIDUAL=residual is not None,
        QUANT_NVFP4=nvfp4,
        WRITE_Y=write_y,
        num_warps=16 if max_m_16_warps >= M else 4,
        launch_pdl=current_platform.is_arch_support_pdl(),
        # Overlaps the consumer's weight load with this kernel. Swept on
        # GB300: ~2% win at M <= 128; the prelaunched consumer grid steals
        # SMs at M >= 192 (fp8 +4-7%).
        TRIGGER_EARLY=M <= 128,
    )
    return q, block_scale, y