Skip to content

vllm.models.muse_spark.nvidia.ops.prepare_trtllm_nvfp4

Pre-expert RMSNorm fused with the TRT-LLM MoE's NVFP4 input quantization.

Matches MuseSpark1xRMSNorm (one BF16 rounding) followed by scaled_fp4_quant(..., is_sf_swizzled_layout=False) with the static global scale, and reads the pooler output either contiguous or in the rank-block layout [tp, M, H / tp] left by a column all-gather.

Functions:

_quantize_nvfp4_global_scale(values, global_scale, BLOCK_K)

scaled_fp4_quant's arithmetic (nvfp4_utils.cuh), linear scales.

Source code in vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
@triton.jit
def _quantize_nvfp4_global_scale(values, global_scale, BLOCK_K: tl.constexpr):
    """``scaled_fp4_quant``'s arithmetic (nvfp4_utils.cuh), linear scales."""
    groups = tl.reshape(values, [BLOCK_K // 16, 16])
    amax = tl.max(tl.abs(groups), axis=1)
    sf = (global_scale * (amax * _rcp_approx(6.0))).to(tl.float8e4nv)
    sf_f32 = sf.to(tl.float32)
    out_scale = tl.where(
        sf_f32 != 0.0, _rcp_approx(sf_f32 * _rcp_approx(global_scale)), 0.0
    )
    scaled = groups * out_scale[:, None]
    magnitude = tl.abs(scaled)
    # Round to nearest even on the E2M1 grid {0, .5, 1, 1.5, 2, 3, 4, 6},
    # saturating, as cvt.rn.satfinite.e2m1x2.f32 does.
    code = (
        (magnitude > 0.25).to(tl.uint8)
        + (magnitude >= 0.75).to(tl.uint8)
        + (magnitude > 1.25).to(tl.uint8)
        + (magnitude >= 1.75).to(tl.uint8)
        + (magnitude > 2.5).to(tl.uint8)
        + (magnitude >= 3.5).to(tl.uint8)
        + (magnitude > 5.0).to(tl.uint8)
    )
    code |= (scaled < 0).to(tl.uint8) << 3
    lo, hi = tl.split(tl.reshape(code, [BLOCK_K // 2, 2]))
    return (lo & 0x0F) | (hi << 4), sf

norm_quant_trtllm_nvfp4(pooled, num_tokens, norm_weight, norm_eps, gain_center, global_scale)

RMSNorm + NVFP4 quant of [M, H] pooled states, in any rank-block layout.

pooled is [M, H] or the [tp * M, H / tp] output of a dim-0 all-gather of column shards. Returns packed [M, H/2] uint8 values and linear [M, H/16] uint8 (E4M3) block scales.

Source code in vllm/models/muse_spark/nvidia/ops/prepare_trtllm_nvfp4.py 
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
def norm_quant_trtllm_nvfp4(
    pooled: torch.Tensor,
    num_tokens: int,
    norm_weight: torch.Tensor,
    norm_eps: float,
    gain_center: float,
    global_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """RMSNorm + NVFP4 quant of ``[M, H]`` pooled states, in any rank-block layout.

    ``pooled`` is ``[M, H]`` or the ``[tp * M, H / tp]`` output of a dim-0
    all-gather of column shards. Returns packed ``[M, H/2]`` uint8 values and
    linear ``[M, H/16]`` uint8 (E4M3) block scales.
    """
    hidden = norm_weight.shape[0]
    shard = pooled.shape[1]
    assert pooled.dtype == torch.bfloat16 and pooled.is_contiguous()
    assert pooled.shape[0] * shard == num_tokens * hidden and hidden % _BLOCK_K == 0
    packed = torch.empty(
        num_tokens, hidden // 2, dtype=torch.uint8, device=pooled.device
    )
    sf = torch.empty(num_tokens, hidden // 16, dtype=torch.uint8, device=pooled.device)
    if num_tokens == 0:
        return packed, sf
    enable_pdl = _enable_pdl()
    _norm_quant_trtllm_nvfp4_kernel[(num_tokens, hidden // _BLOCK_K)](
        pooled,
        norm_weight,
        global_scale,
        packed,
        sf,
        num_tokens,
        norm_eps,
        gain_center,
        HIDDEN=hidden,
        SHARD=shard,
        BLOCK_K=_BLOCK_K,
        USE_PDL=enable_pdl,
        num_warps=8,
        enable_fp_fusion=False,
        **({"launch_pdl": True} if enable_pdl else {}),
    )
    return packed, sf