Skip to content

vllm.model_executor.layers.quantization.utils.nvfp4_utils

Functions:

_csf32_blocks(rows, cols) cached

Origin, width and page of each 32x2 block in storage order, and the flat [rows, cols] position of each code.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@functools.cache
def _csf32_blocks(rows: int, cols: int):
    """Origin, width and page of each 32x2 block in storage order, and the flat
    ``[rows, cols]`` position of each code."""
    tile_r, tile_c, pair, block = np.meshgrid(
        np.arange(-(-rows // 128)),
        np.arange(-(-cols // 4)),
        np.arange(2),
        np.arange(4),
        indexing="ij",
    )
    r0 = (128 * tile_r + 32 * block).ravel()
    c0 = (4 * tile_c + 2 * pair).ravel()
    page = (tile_r * tile_c.shape[1] + tile_c).ravel() // 32
    keep = (r0 < rows) & (c0 < cols)
    r0, c0, page = r0[keep], c0[keep], page[keep]
    width = np.minimum(cols - c0, 2)
    i, j = np.meshgrid(np.arange(32), np.arange(2), indexing="ij")
    cells = (r0[:, None, None] + i) * cols + c0[:, None, None] + j
    return r0, c0, width, page, cells[j < width[:, None, None]]

csf_weights_transform(encodings)

Pair CSF scale streams back into canonical E4M3 scale tensors.

The returned transform is lazy over a (name, tensor) weight stream: tensors named <name>.<stream> for a declared encoding are buffered per scale and, once the encoding's stream set completes, replaced by the decoded <name> tensor; everything else passes through unchanged.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
def csf_weights_transform(
    encodings: Iterable[str],
) -> Callable[[Iterable[tuple[str, torch.Tensor]]], Iterable[tuple[str, torch.Tensor]]]:
    """Pair CSF scale streams back into canonical E4M3 scale tensors.

    The returned transform is lazy over a ``(name, tensor)`` weight stream:
    tensors named ``<name>.<stream>`` for a declared encoding are buffered per
    scale and, once the encoding's stream set completes, replaced by the
    decoded ``<name>`` tensor; everything else passes through unchanged.
    """
    suffixes = {
        f".{stream}": (encoding, stream)
        for encoding in encodings
        for stream in CSF_SCALE_ENCODINGS[encoding][0]
    }

    def transform(
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[tuple[str, torch.Tensor]]:
        pending: dict[str, tuple[str, dict[str, torch.Tensor]]] = {}
        decoded = 0
        for name, tensor in weights:
            for suffix, (encoding, stream) in suffixes.items():
                if name.endswith(suffix):
                    base = name[: -len(suffix)]
                    break
            else:
                yield name, tensor
                continue
            if base in pending and pending[base][0] != encoding:
                raise ValueError(f"Mixed CSF scale encodings for {base}")
            parts = pending.setdefault(base, (encoding, {}))[1]
            parts[stream] = tensor
            if len(parts) == len(CSF_SCALE_ENCODINGS[encoding][0]):
                del pending[base]
                yield base, decode_csf_scale(encoding, parts)
                decoded += 1
        if pending:
            raise ValueError(
                f"CSF scale streams without a partner: {sorted(pending)[:4]}"
            )
        if decoded:
            logger.info("Decoded %d CSF scale tensors at load", decoded)

    return transform

decode_csf_scale(encoding, streams)

Rebuild one [rows, cols] float8_e4m3fn block-scale matrix.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
269
270
271
272
273
def decode_csf_scale(encoding: str, streams: dict[str, torch.Tensor]) -> torch.Tensor:
    """Rebuild one ``[rows, cols]`` float8_e4m3fn block-scale matrix."""
    names, decode = CSF_SCALE_ENCODINGS[encoding]
    out = decode(*(streams[name].cpu() for name in names))
    return torch.from_numpy(out).view(torch.float8_e4m3fn)

nvfp4_weight_padding_bytes(layer)

Compute K-dim padding (packed bytes) of an NVFP4 weight, Aligned with pad_nvfp4_weight_for_cutlass above.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
104
105
106
107
108
def nvfp4_weight_padding_bytes(layer: torch.nn.Module) -> int:
    """Compute K-dim padding (packed bytes) of an NVFP4 weight,
    Aligned with `pad_nvfp4_weight_for_cutlass` above.
    """
    return layer.weight.shape[1] - layer.input_size_per_partition // 2

pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes)

Pad packed FP4 activations to match the K-dimension padding applied to weights. The padding is in bytes (tensor dimension), not FP4 elements.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
111
112
113
114
115
116
117
118
119
120
def pad_nvfp4_activation_for_cutlass(
    x_fp4: torch.Tensor,
    weights_padding_bytes: int,
) -> torch.Tensor:
    """Pad packed FP4 activations to match the K-dimension padding applied to weights.
    The padding is in bytes (tensor dimension), not FP4 elements.
    """
    if weights_padding_bytes > 0:
        return torch.nn.functional.pad(x_fp4, (0, weights_padding_bytes)).contiguous()
    return x_fp4

pad_nvfp4_weight_for_cutlass(weight, alignment=32)

Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.

CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible by 32 for aligned memory access and efficient tensor core operations.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def pad_nvfp4_weight_for_cutlass(
    weight: torch.Tensor,
    alignment: int = 32,
) -> tuple[torch.Tensor, int]:
    """Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
    the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.

    CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible
    by 32 for aligned memory access and efficient tensor core operations.
    """
    weight_current_rows = weight.shape[0]

    # Pad N dimension (rows) if not aligned
    if weight_current_rows % alignment != 0:
        total_rows = round_up(weight_current_rows, alignment)
        pad_rows = total_rows - weight_current_rows
        weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_rows)).contiguous()

    # Check K dimension alignment
    # 2 FP4 items are packed per byte in the input dimension
    weight_current_col_bytes = weight.shape[1]
    weight_current_col_elements = weight_current_col_bytes * 2

    weights_padding_bytes = 0
    if weight_current_col_elements % alignment != 0:
        total_cols = round_up(weight_current_col_elements, alignment)
        pad_cols = total_cols - weight_current_col_elements
        # Convert from FP4 element count to bytes (2 FP4 values per byte)
        # pad_cols is always even since alignment=32 and current elements are even
        pad_bytes = pad_cols // 2
        weight = torch.nn.functional.pad(weight, (0, pad_bytes, 0, 0)).contiguous()
        weights_padding_bytes = pad_bytes

    return weight, weights_padding_bytes

slice_nvfp4_output(out, output_size)

Slice the output tensor to remove padding in N dimension if weight was padded.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
123
124
125
126
127
128
129
130
def slice_nvfp4_output(
    out: torch.Tensor,
    output_size: int,
) -> torch.Tensor:
    """Slice the output tensor to remove padding in N dimension if weight was padded."""
    if out.shape[-1] != output_size:
        return out[..., :output_size].contiguous()
    return out

swizzle_blockscale(scale)

Pad and block-interleave the FP4 block-scales so that they match the data layout expected by the CUTLASS / FlashInfer kernels.

Parameters:

  • scale

    (Tensor) –

    The FP4 block-scale tensor to swizzle.

Returns:

  • Tensor –

    The swizzled tensor with the same logical shape as scale.

Source code in vllm/model_executor/layers/quantization/utils/nvfp4_utils.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
    """Pad and block-interleave the FP4 block-scales so that they match the data
    layout expected by the CUTLASS / FlashInfer kernels.

    Args:
        scale: The FP4 block-scale tensor to swizzle.

    Returns:
        The swizzled tensor with the same logical shape as `scale`.

    """
    assert scale.dtype == torch.float8_e4m3fn, (
        "swizzle_blockscale expects the input tensor to be in "
        "torch.float8_e4m3fn format."
    )

    scale_ndim = scale.ndim
    if scale_ndim == 2:
        scale = scale.unsqueeze(0)  # (1, M, K)
    assert scale.ndim == 3, "Expected a 2-D or 3-D tensor for block scales."

    B, M, K = scale.shape

    M_padded = round_up(M, 128)
    K_padded = round_up(K, 4)

    padded = torch.zeros(
        (B, M_padded, K_padded), dtype=scale.dtype, device=scale.device
    )
    padded[:B, :M, :K] = scale

    # Reshape / permute to the layout required by the kernel.
    padded = padded.reshape(B, M_padded // 128, 4, 32, K_padded // 4, 4)
    swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda()

    if scale_ndim == 2:
        return swizzled.reshape(M_padded, K_padded)
    return swizzled.reshape(B, M_padded, K_padded)