vllm.models.flint.ops.nvfp4_gemv

Small-M BF16 x NVFP4 (W4A16) and BF16 x BF16 GEMM for SM100 in CuTe DSL.

out[M, N] = (x[M, K] @ dequant(W)[N, K]^T) * global_scale with the weight in the ModelOpt layout used by the FlashInfer CuTe-DSL W4A16 kernel: weight uint8 [N, K / 2] (two E2M1 codes per byte, low nibble first) and weight_scale the 128x4-swizzled E4M3 block scales (one per 16 values), so both kernels share the same tensors.

Numerics: every weight is dequantized exactly to BF16 (e2m1 * e4m3 via cvt.rn.bf16x2.e2m1x2 and mul.rn.bf16x2), multiplied with the BF16 activations on tensor cores (mma.sync m16n8k16, FP32 accumulate), and the FP32 global scale is applied to the FP32 accumulator before the single output rounding to BF16.

Each CTA owns 16 output features (the MMA M dimension, swap-AB) and 8 * nt tokens (nt = 1, 2 or 4, i.e. up to 32); its warps (and, with split > 1, the CTAs of a cluster) split K and reduce through shared memory / DSMEM. All weight loads are issued before griddepcontrol.wait so they overlap the previous kernel under PDL.

Functions:

• bf16_gemv –

  Small-M BF16 GEMM x @ weight.T (FP32 accumulate); same kernel and

• dequantize_rows –

  Exact BF16 e2m1 * e4m3 of all weight rows (global scale not

• nvfp4_gemv –

  Small-M W4A16 GEMM: x @ dequant(weight).T * global_scale.

• pick_nt –

  Smallest tuned token tile count covering m, or None if the GEMV

• tuned –

  Whether the GEMV is the fastest choice for this shape at some M.

• warmup –

  Compile every variant the GEMV may pick for this shape.

bf16_gemv(x, weight, silu_mul=False, pdl=True)

Small-M BF16 GEMM x @ weight.T (FP32 accumulate); same kernel and arguments as nvfp4_gemv with a row-major BF16 [N, K] weight.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def bf16_gemv(
    x: torch.Tensor,
    weight: torch.Tensor,
    silu_mul: bool = False,
    pdl: bool = True,
) -> torch.Tensor:
    """Small-M BF16 GEMM ``x @ weight.T`` (FP32 accumulate); same kernel and
    arguments as ``nvfp4_gemv`` with a row-major BF16 ``[N, K]`` weight."""
    assert weight.dtype == torch.bfloat16 and weight.is_contiguous()
    assert weight.shape[1] == x.shape[1]
    return _launch(
        x,
        weight.view(torch.uint8).view(-1),
        weight.view(torch.uint8).view(-1),
        _one(x.device),
        weight.shape[0],
        silu_mul,
        pdl,
        bf16=True,
    )

dequantize_rows(weight, weight_scale)

Exact BF16 e2m1 * e4m3 of all weight rows (global scale not applied); weight_scale is 128x4-swizzled.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def dequantize_rows(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor:
    """Exact BF16 ``e2m1 * e4m3`` of all weight rows (global scale not
    applied); ``weight_scale`` is 128x4-swizzled."""
    global _dequant
    if _dequant is None:
        _dequant = _dequant_kernel()
    rows, k = weight.shape[0], weight.shape[1] * 2
    assert rows % 16 == 0 and k % 256 == 0
    out = torch.empty(rows, k, dtype=torch.bfloat16, device=weight.device)
    _dequant[(rows // 16, k // 256)](
        weight,
        weight_scale.view(torch.float8_e4m3fn).view(-1),
        out,
        k,
        16,
        256,
        num_warps=4,
    )
    return out

nvfp4_gemv(x, weight, weight_scale, global_scale, silu_mul=False, pdl=True)

Small-M W4A16 GEMM: x @ dequant(weight).T * global_scale.

Parameters:

• x

  (Tensor) –

  BF16 [M, K] activations with unit inner stride; (N, K) and M must be covered by the tuned tables (pick_nt).

• weight

  (Tensor) –

  uint8 [N, K / 2] packed E2M1 codes (ModelOpt layout).

• weight_scale

  (Tensor) –

  128x4-swizzled E4M3 block scales (swizzle_blockscale).

• global_scale

  (Tensor) –

  FP32 scalar tensor (one element).

• silu_mul

  (bool, default: False ) –

  Treat the output as [gate | up] halves and return the BF16 silu(gate) * up ([M, N / 2]) of the BF16 gate / up outputs (exact SiLU, one BF16 rounding).

• pdl

  (bool, default: True ) –

  Launch with programmatic dependent launch: weight loads start while the previous kernel drains. Disable after kernels that need the whole GPU to make progress (e.g. attention), whose remaining CTAs the early, waiting GEMV CTAs would starve.

Returns:

• Tensor –

  The BF16 [M, N] ([M, N / 2] with silu_mul) output.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def nvfp4_gemv(
    x: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    global_scale: torch.Tensor,
    silu_mul: bool = False,
    pdl: bool = True,
) -> torch.Tensor:
    """Small-M W4A16 GEMM: ``x @ dequant(weight).T * global_scale``.

    Args:
        x: BF16 ``[M, K]`` activations with unit inner stride; ``(N, K)``
            and M must be covered by the tuned tables (``pick_nt``).
        weight: uint8 ``[N, K / 2]`` packed E2M1 codes (ModelOpt layout).
        weight_scale: 128x4-swizzled E4M3 block scales (``swizzle_blockscale``).
        global_scale: FP32 scalar tensor (one element).
        silu_mul: Treat the output as ``[gate | up]`` halves and return the
            BF16 ``silu(gate) * up`` (``[M, N / 2]``) of the BF16 gate / up
            outputs (exact SiLU, one BF16 rounding).
        pdl: Launch with programmatic dependent launch: weight loads start
            while the previous kernel drains. Disable after kernels that need
            the whole GPU to make progress (e.g. attention), whose remaining
            CTAs the early, waiting GEMV CTAs would starve.

    Returns:
        The BF16 ``[M, N]`` (``[M, N / 2]`` with ``silu_mul``) output.

    """
    assert weight.shape[1] * 2 == x.shape[1]
    return _launch(
        x,
        weight.view(-1),
        weight_scale.view(torch.uint8).view(-1),
        global_scale.reshape(1),
        weight.shape[0],
        silu_mul,
        pdl,
        bf16=False,
    )

pick_nt(n, k, m, silu_mul=False, bf16=False)

Smallest tuned token tile count covering m, or None if the GEMV is not the fastest choice for this shape and M.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def pick_nt(
    n: int, k: int, m: int, silu_mul: bool = False, bf16: bool = False
) -> int | None:
    """Smallest tuned token tile count covering ``m``, or None if the GEMV
    is not the fastest choice for this shape and M."""
    nts = [nt for nt in _table(silu_mul, bf16).get((n, k), ()) if 8 * nt >= m]
    return min(nts) if nts else None

tuned(n, k, silu_mul=False, bf16=False)

Whether the GEMV is the fastest choice for this shape at some M.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def tuned(n: int, k: int, silu_mul: bool = False, bf16: bool = False) -> bool:
    """Whether the GEMV is the fastest choice for this shape at some M."""
    return (n, k) in _table(silu_mul, bf16)

warmup(n, k, silu_mul=False, pdl=True, bf16=False)

Compile every variant the GEMV may pick for this shape.

Source code in vllm/models/flint/ops/nvfp4_gemv.py

def warmup(
    n: int,
    k: int,
    silu_mul: bool = False,
    pdl: bool = True,
    bf16: bool = False,
) -> None:
    """Compile every variant the GEMV may pick for this shape."""
    for nt in _configs(n, k, silu_mul, bf16):
        _get_compiled(n, k, nt, silu_mul, pdl, bf16)