vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm

MXFP8 GEMM on 32x32 block-scaled weights for gfx950 (tl.dot_scaled).

y = x @ w.T with an MXFP8 activation (e4m3 values, one E8M0 scale per row and 32 K, [M, K / 32]) and an e4m3 weight whose E8M0 scales stay in the checkpoint's 32x32 blocks, [N / 32, K / 32], instead of being expanded to every row. Each weight scale byte is then read once per 32 output rows, and small-M shapes can use the packed kernel below.

Two kernels, picked per shape from a table tuned on MI355X:

• a tiled kernel for larger M;
• a packed kernel for small M, which fills the MFMA's rows with K panels instead of tokens and keeps only the matching-panel (block-diagonal) products, so a handful of tokens still streams the weight at full width.

Either can split K. The partials are normally summed in the same launch by the last program of each output tile to finish, in split order, so the result does not depend on scheduling. Otherwise a second launch reduces them.

The packed kernel and the in-launch split-K reduction on one XCD (_split_tile, _sum_splits, _split_counters, _k_partition) are adapted from the group32 GEMM in ROCm/aiter#5750.

Functions:

• rocm_mxfp8_block32_gemm –

  x @ weight.T for MXFP8 x and a 32x32 block-scaled MXFP8 weight.

• rocm_mxfp8_block32_gemm_quant –

  x @ weight.T for a raw bf16/fp16 x [M, K] and a 32x32

_default_config(M, N, K)

Untuned shapes: the tiers most tuned shapes settle on.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _default_config(M: int, N: int, K: int) -> _Config:
    """Untuned shapes: the tiers most tuned shapes settle on."""
    if M <= 8:
        return _packed(8, 16, 256, 2, 2, 3, 16, 0)
    if M <= 32:
        return _packed(32, 32, 256, 1, 2, 3, 16, 0)
    if M <= 256:
        return _tiled(64, 64, 256, 1, True, 4, 2, 16, 0)
    return _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)

_fused_quant_config(M, N, K)

The stock tier active at M with its tiling/scheduling retuned for the in-kernel quant prologue (see _TUNED_FUSED_QUANT); the stock config where the shape has no fused-tier entry. K-structure axes are always the stock tier's own, so the output stays bit-exact at every M.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _fused_quant_config(M: int, N: int, K: int) -> _Config:
    """The stock tier active at M with its tiling/scheduling retuned for the
    in-kernel quant prologue (see _TUNED_FUSED_QUANT); the stock config
    where the shape has no fused-tier entry. K-structure axes are always
    the stock tier's own, so the output stays bit-exact at every M."""
    base = _config(M, N, K)
    for max_m, (bm, bn, warps, stages, waves) in _TUNED_FUSED_QUANT.get((N, K), ()):
        if max_m >= M:
            return base._replace(
                block_m=bm,
                block_n=bn,
                num_warps=warps,
                num_stages=stages,
                waves_per_eu=waves,
            )
    return base

_k_partition(K, splits, step)

K per split, rounded up to whole steps, and the number of non-empty splits.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _k_partition(K: int, splits: int, step: int) -> tuple[int, int]:
    """K per split, rounded up to whole steps, and the number of non-empty
    splits."""
    size = triton.cdiv(triton.cdiv(K, splits), step) * step
    return size, triton.cdiv(K, size)

_launch(x, x_scale, weight, weight_scale, out, cfg, quant_in_kernel=False)

Launch cfg's kernel. quant_in_kernel=True (packed cfgs only) makes x the raw bf16/fp16 activation and x_scale None: the kernel quantizes x per 32-block in registers (see _block32_packed_quant_kernel).

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _launch(
    x: torch.Tensor,
    x_scale: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    out: torch.Tensor,
    cfg: _Config,
    quant_in_kernel: bool = False,
) -> None:
    """Launch cfg's kernel. quant_in_kernel=True (packed cfgs only) makes
    x the raw bf16/fp16 activation and x_scale None: the kernel quantizes
    x per 32-block in registers (see _block32_packed_quant_kernel)."""
    M, K = x.shape
    N = weight.shape[0]
    opts = dict(
        num_warps=cfg.num_warps,
        num_stages=cfg.num_stages,
        matrix_instr_nonkdim=cfg.matrix_instr_nonkdim,
        waves_per_eu=cfg.waves_per_eu,
    )
    grid_m, grid_n = triton.cdiv(M, cfg.block_m), triton.cdiv(N, cfg.block_n)
    tiles = grid_m * grid_n
    step = cfg.block_k * cfg.k_split if cfg.packed else cfg.block_k
    splits = cfg.fused_splits if cfg.fused_splits > 1 or cfg.packed else cfg.k_split
    k_per_split, splits = _k_partition(K, splits, step)
    counters = None
    if cfg.fused_splits > 1 and splits > 1 and tiles <= _SPLIT_COUNTERS:
        counters = _split_counters(out.device)
    if counters is not None:
        fused = splits
        target = out
        slots = torch.empty(
            tiles * fused * cfg.block_m * cfg.block_n,
            dtype=torch.float32,
            device=out.device,
        )
        grid: tuple[int, ...] = (triton.cdiv(tiles, _XCD_STRIDE) * _XCD_STRIDE * fused,)
    else:
        # Without counters, split K through a second launch.
        fused = 1
        target = (
            out
            if splits == 1
            else torch.empty((splits, M, N), dtype=torch.float32, device=x.device)
        )
        slots = counters = target  # Unused by the kernels.
        if cfg.packed or not cfg.n_first:
            grid = (grid_m, grid_n, splits)
        else:
            grid = (grid_n, grid_m, splits)
    args: tuple
    if quant_in_kernel:
        if not cfg.packed:
            raise ValueError("quant_in_kernel needs a packed tile config")
        args = (x, weight, weight_scale, target, slots, counters, M, N, K)
    else:
        args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K)
    stride = out.stride(0) if splits == 1 or fused > 1 else N
    if cfg.packed:
        if quant_in_kernel:
            # Same output contract as the packed kernel below (partials at
            # pid_k * M * N when FUSED_SPLITS == 1 with multiple pid_k), so
            # the split-K reduce after this branch applies unchanged.
            _block32_packed_quant_kernel[grid](
                *args,
                stride,
                BLOCK_M=cfg.block_m,
                BLOCK_N=cfg.block_n,
                BLOCK_K=cfg.block_k,
                K_PACK=cfg.k_split,
                K_PER_SPLIT=k_per_split,
                EVEN_K=K % step == 0,
                FUSED_SPLITS=fused,
                FP8_MAX=_FP8_MAX,
                TINY=_FP32_TINY,
                **opts,
            )
        else:
            _block32_packed_kernel[grid](
                *args,
                stride,
                BLOCK_M=cfg.block_m,
                BLOCK_N=cfg.block_n,
                BLOCK_K=cfg.block_k,
                K_PACK=cfg.k_split,
                K_PER_SPLIT=k_per_split,
                EVEN_K=K % step == 0,
                FUSED_SPLITS=fused,
                **opts,
            )
    else:
        _block32_tiled_kernel[grid](
            *args,
            stride,
            BLOCK_M=cfg.block_m,
            BLOCK_N=cfg.block_n,
            BLOCK_K=cfg.block_k,
            K_PER_SPLIT=k_per_split,
            N_FIRST=cfg.n_first,
            EVEN_K=K % k_per_split == 0,
            FUSED_SPLITS=fused,
            **opts,
        )
    if fused == 1 and splits > 1:
        _block32_splitk_reduce_kernel[(triton.cdiv(M, 32), triton.cdiv(N, 64))](
            target,
            out,
            M,
            N,
            out.stride(0),
            SPLIT_K=splits,
            BLOCK_M=32,
            BLOCK_N=64,
        )

_rocm_mxfp8_block32_gemm_quant_impl(x, weight, weight_scale, out_dtype)

x @ weight.T for raw bf16/fp16 x and a 32x32 block-scaled MXFP8 weight, quantizing x inside the GEMM at tiny M.

M <= 8 with a packed tile config (every tuned M <= 8 tier is packed) runs one launch whose K loop computes the per-32-block E8M0 scale and the e4m3 cast exactly as _mxfp8_quant_triton_kernel and feeds tl.dot_scaled from registers -- bit-identical output to mxfp8_e4m3_quantize + rocm_mxfp8_block32_gemm with one launch fewer on the decode critical path. Everything else (larger M, non-packed tiles, or non-contiguous or non-bf16/fp16 x) runs the stock two-launch path with identical results. The branch is inside this opaque op so a compiled caller cannot freeze the token-count decision.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _rocm_mxfp8_block32_gemm_quant_impl(
    x: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """``x @ weight.T`` for raw bf16/fp16 ``x`` and a 32x32 block-scaled
    MXFP8 weight, quantizing ``x`` inside the GEMM at tiny M.

    M <= 8 with a packed tile config (every tuned M <= 8 tier is packed)
    runs one launch whose K loop computes the per-32-block E8M0 scale and
    the e4m3 cast exactly as ``_mxfp8_quant_triton_kernel`` and feeds
    ``tl.dot_scaled`` from registers -- bit-identical output to
    ``mxfp8_e4m3_quantize`` + ``rocm_mxfp8_block32_gemm`` with one launch
    fewer on the decode critical path. Everything else (larger M,
    non-packed tiles, or non-contiguous or non-bf16/fp16 x) runs the stock
    two-launch path with identical results. The branch is inside this
    opaque op so a compiled caller cannot freeze the token-count decision.
    """
    M, K = x.shape
    out = torch.empty((M, weight.shape[0]), dtype=out_dtype, device=x.device)
    if M <= 0:
        return out
    cfg = _config(M, weight.shape[0], K)
    eligible = (
        cfg.packed
        and M <= _TINY_M_FUSED_QUANT_MAX_M
        and x.dtype in (torch.bfloat16, torch.float16)
        and x.is_contiguous()
    )
    if eligible:
        # Retuned tiling for the in-kernel quant prologue (same K-structure
        # as the stock tier at M, so bit-exactness holds); the stock path
        # below keeps the stock config untouched.
        fused_cfg = _fused_quant_config(M, weight.shape[0], K)
        shape = (M, weight.shape[0], K)
        if shape not in _FUSED_QUANT_HIT_SHAPES:
            _FUSED_QUANT_HIT_SHAPES.add(shape)
            # warning_once keeps the engagement log one-shot per shape;
            # WARNING, not INFO, so it reaches the server log.
            logger.warning_once(
                "tiny-M fused MXFP8 quant+GEMM HIT (M=%d, N=%d, K=%d, packed "
                "tile %dx%dx%d k_pack=%d%s)",
                M,
                weight.shape[0],
                K,
                fused_cfg.block_m,
                fused_cfg.block_n,
                fused_cfg.block_k,
                fused_cfg.k_split,
                (
                    f", fused_splits={fused_cfg.fused_splits}"
                    if fused_cfg.fused_splits > 1
                    else ""
                ),
            )
        _launch(x, None, weight, weight_scale, out, fused_cfg, quant_in_kernel=True)
    else:
        x_q, x_scale = mxfp8_e4m3_quantize(x)
        _launch(x_q, x_scale, weight, weight_scale, out, cfg)
    return out

_split_counters(device)

Zeroed arrival counters, one set per stream so concurrent launches never share a tile's counter. Each launch leaves the counters it used at zero.

None when a stream first needs them inside a CUDA graph capture: that allocation would come from the graph's pool, where it can take the address of an intermediate that every replay overwrites. vLLM warms up on its capture streams first, so this only affects captures without a warmup.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def _split_counters(device: torch.device) -> torch.Tensor | None:
    """Zeroed arrival counters, one set per stream so concurrent launches never
    share a tile's counter. Each launch leaves the counters it used at zero.

    None when a stream first needs them inside a CUDA graph capture: that
    allocation would come from the graph's pool, where it can take the address
    of an intermediate that every replay overwrites. vLLM warms up on its
    capture streams first, so this only affects captures without a warmup.
    """
    key = (device.index, torch.accelerator.current_stream(device).stream_id)
    counters = _split_counters_by_stream.get(key)
    if counters is None:
        if torch.cuda.is_current_stream_capturing():
            return None
        counters = torch.zeros(_SPLIT_COUNTERS, dtype=torch.int32, device=device)
        _split_counters_by_stream[key] = counters
    return counters

_split_tile(N, BLOCK_N, SPLITS)

(pid_m, pid_n, tile, split) of a 1D grid that keeps a tile's splits on one XCD, so they meet in that XCD's L2.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

@triton.jit
def _split_tile(N: tl.constexpr, BLOCK_N: tl.constexpr, SPLITS: tl.constexpr):
    """(pid_m, pid_n, tile, split) of a 1D grid that keeps a tile's splits on
    one XCD, so they meet in that XCD's L2."""
    pid = tl.program_id(0)
    tile = pid // (8 * SPLITS) * 8 + pid % 8
    grid_n: tl.constexpr = (N + BLOCK_N - 1) // BLOCK_N
    return tile // grid_n, tile % grid_n, tile, pid // 8 % SPLITS

_sum_splits(acc, out_ptrs, out_mask, slot_ptr, count_ptr, split, SPLITS)

Publish this split's partial; the tile's last arrival sums all of them in split order and re-arms the counter.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

@triton.jit
def _sum_splits(
    acc, out_ptrs, out_mask, slot_ptr, count_ptr, split, SPLITS: tl.constexpr
):
    """Publish this split's partial; the tile's last arrival sums all of them
    in split order and re-arms the counter."""
    BM: tl.constexpr = acc.shape[0]
    BN: tl.constexpr = acc.shape[1]
    local = tl.arange(0, BM)[:, None] * BN + tl.arange(0, BN)[None, :]
    tl.store(slot_ptr + split * BM * BN + local, acc)
    # The partial must be in L2 before the arrival is counted. The splits
    # share an XCD, so no device-scope release (an L2 writeback) is needed.
    tl.inline_asm_elementwise(
        "s_waitcnt vmcnt(0)", "=v,v", [split], dtype=tl.int32, is_pure=False, pack=1
    )
    tl.debug_barrier()
    if tl.atomic_add(count_ptr, 1, sem="acq_rel", scope="cta") == SPLITS - 1:
        total = tl.zeros((BM, BN), dtype=tl.float32)
        for s in tl.static_range(SPLITS):
            total += tl.load(slot_ptr + s * BM * BN + local, cache_modifier=".cv")
        tl.store(out_ptrs, total.to(out_ptrs.dtype.element_ty), mask=out_mask)
        tl.store(count_ptr, 0)

rocm_mxfp8_block32_gemm(x, x_scale, weight, weight_scale, out_dtype)

x @ weight.T for MXFP8 x and a 32x32 block-scaled MXFP8 weight.

Parameters:

• x

  (Tensor) –

  [M, K] e4m3 activation, contiguous.

• x_scale

  (Tensor) –

  [M, K / 32] E8M0 (uint8) activation scales.

• weight

  (Tensor) –

  [N, K] e4m3 weight, contiguous.

• weight_scale

  (Tensor) –

  [ceil(N / 32), K / 32] E8M0 (uint8) weight block scales.

• out_dtype

  (dtype) –

  Output dtype.

Returns:

• Tensor –

  The [M, N] product.

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def rocm_mxfp8_block32_gemm(
    x: torch.Tensor,
    x_scale: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """``x @ weight.T`` for MXFP8 ``x`` and a 32x32 block-scaled MXFP8 weight.

    Args:
        x: [M, K] e4m3 activation, contiguous.
        x_scale: [M, K / 32] E8M0 (uint8) activation scales.
        weight: [N, K] e4m3 weight, contiguous.
        weight_scale: [ceil(N / 32), K / 32] E8M0 (uint8) weight block scales.
        out_dtype: Output dtype.

    Returns:
        The [M, N] product.

    """
    return torch.ops.vllm.rocm_mxfp8_block32_gemm(
        x, x_scale, weight, weight_scale, out_dtype
    )

rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, out_dtype)

x @ weight.T for a raw bf16/fp16 x [M, K] and a 32x32 block-scaled MXFP8 weight [N, K] with weight_scale [N / 32, K / 32]. Tiny-M batches quantize inside the packed GEMM (see _rocm_mxfp8_block32_gemm_quant_impl).

Source code in vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py

def rocm_mxfp8_block32_gemm_quant(
    x: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    out_dtype: torch.dtype,
) -> torch.Tensor:
    """``x @ weight.T`` for a raw bf16/fp16 ``x`` [M, K] and a 32x32
    block-scaled MXFP8 ``weight`` [N, K] with ``weight_scale``
    [N / 32, K / 32]. Tiny-M batches quantize inside the packed GEMM (see
    ``_rocm_mxfp8_block32_gemm_quant_impl``).
    """
    return torch.ops.vllm.rocm_mxfp8_block32_gemm_quant(
        x, weight, weight_scale, out_dtype
    )