Block-scaled GEMM using CUTLASS Operator API

This notebook walks through using CUTLASS Operator API to run block-scaled GEMMs -- the operations behind matrix-multiplication for narrow-precision types (like MXFP8).

It builds on the basic GEMM tutorial, which introduced the flow for a plain GEMM.

Key takeaway: the same flow works here. A block-scaled GEMM is still logically out = A @ B, only the operands are stored differently. Most of this notebook explains how to construct an Operand describing block-scaling. After that, expressing the GEMM, finding a supported kernel, compiling, and running it is completely identical to the basic GEMM.

import torch

import cutlass.operators as ops

if not (status := ops.utils.device.device_or_env_supports("100f")):
    print(f"This notebook requires a Blackwell GPU (sm_100f family).\n{status.error}")
    import sys

    sys.exit(0)

What is a block-scaled GEMM?

In a plain GEMM, we compute out = A @ B and the operands carry their own values in a single tensor. For narrower formats -- such as 8-bit/4-bit floating point -- a single tensor can no longer represent both the small magnitudes and the large outliers of a typical activation or weight tensor. The block-scaled formulation quantizes the bulk of the data into a narrow-precision tensor, and recovers dynamic range by multiplying it with a smaller separate tensor of scale factors. Each scale factor scales a small, contiguous block of the quantized tensor: out = (scale_A * quantized_A) @ (scale_B * quantized_B)

A block of V elements in quantized_A share the same scale_A multiplier factor (likewise for B). So, the scale factor's size is 1/V of the quantized tensor's.

Blackwell Tensor Cores can apply the scale-and-matmul in a single MMA instruction in hardware, rather than as separate steps. The format we demonstrate here is MXFP8: FP8 E4M3FN quantized values (8-bit, 4 exponent / 3 mantissa) paired with E8M0 scales -- one scale per 32 K-elements -- following the OCP Microscaling Format spec.

Operands: same A @ B, different storage

CUTLASS Operator API draws a clean line between a logical operand of an operation (A in out = A @ B), and the physical tensor(s) that store its values (A = scale_A * quantized_A).

The shared abstraction is Operand. It represents the logical operand, and encapsulates its constituent physical tensor(s). DenseTensor and ScaledOperand are 2 types of Operands:

• A plain GEMM computes out = A @ B and the Operands in it are DenseTensor. The logical value of the operand is trivially stored in the tensor.
• A block-scaled GEMM still logically computes out = A @ B, and the Operands are ScaledOperands. The logical values of A & B need to be reconstructed by multiplying the scale factors over the quantized tensors, upon which we then perform matrix multiplication as usual.

GemmArguments expect A & B to be Operands, so both forms work with it -- either way, GemmArguments express a logical GEMM over A and B.

Once the operation is expressed, the discovery / compile / run flow is exactly what you saw in the basic GEMM tutorial.

From this point in the notebook, the bulk of the work is constructing the ScaledOperand correctly.

Constructing ScaledOperands

A ScaledOperand has the following fields:

• quantized: the narrow-precision tensor holding the quantized values.
• scale: a tensor of scale factors, one per block of quantized.
• mode: the block shape over which each scale factor is broadcast, given as an (L, M, K) tuple. ops.ScaleMode.Blockwise1x32 is (1, 1, 32), meaning "per-element along L and M, every 32 elements along K".
• swizzle: how the scale tensor is laid out in memory. SwizzleNone is the natural shape (L, M, K // V), one scale per mode block. Swizzle32x4x4 is a special layout based on Blackwell MMA requirements, see ScaleSwizzleMode for details.

The latter two enums map directly to torch.nn.functional.scaled_mm, for easy interop:

• ScaleMode is the analog of torch.nn.functional.ScalingType (with members like BlockWise1x32)
• ScaleSwizzleMode is the analog of torch.nn.functional.SwizzleType

Like everything passed to GemmArguments, the fields above describe the operation and operands we have -- including their dtypes, block-scaling granularity, and swizzle layout -- independent of which operator will consume it. get_operators(args) then finds operators that support consuming operands described that way.

Step 1: Choose a format and a problem size

We pick MXFP8 format, which requires:

• quantized tensor of E4M3FN type
• scale factors of E8M0 type, shared by a block of 1x32 elements

This fixes the dtypes of the quantized & scale tensors, and ScaleMode.

The swizzle tag is determined by Blackwell's tcgen05.mma block-scaled MMA instruction (see documentation from cuDNN and cuBLAS for details). We tag our scale tensors as Swizzle32x4x4 to declare that we are providing data arranged in this layout.

M, N, K, L = 256, 512, 1024, 1

ab_dtype = torch.float8_e4m3fn      # narrow quantized type (MXFP8)
scale_dtype = torch.float8_e8m0fnu  # E8M0 scale type per the OCP MX spec
out_dtype = torch.float32

mode = ops.ScaleMode.Blockwise1x32            # one scale per 32 K-elements
swizzle = ops.ScaleSwizzleMode.Swizzle32x4x4  # tcgen05 block-scaled layout

Step 2: Build the quantized tensors

The quantized tensors quantized_A and quantized_B are no different than a DenseTensor from a plain GEMM. For a (L, M, K) @ (L, K, N) -> (L, M, N) GEMM, the Blackwell tcgen05 MMA requires both quantized_A and quantized_B to be K-major -- that is, quantized_A (with K columns) must be row-major, and quantized_B (with K rows) must be column-major. We construct those tensors here.

quantized_A = torch.randint(-1, 2, (L, M, K), device="cuda").to(ab_dtype)
quantized_B = torch.randint(-1, 2, (L, N, K), device="cuda").to(ab_dtype).transpose(1, 2)
out = torch.empty((L, M, N), device="cuda", dtype=out_dtype)

Step 3: Build the scale tensors

Logically, each scale factor in scale_A covers a (1, 1, 32) block of quantized_A, so scale_A carries one value per row of A and one per group of 32 K-elements -- i.e. a logical shape of (L, M, ceil(K / 32)).

Swizzle32x4x4 packs those values into 128 × 4 tiles. The tile boundaries impose two alignment requirements on the logical shape:

• the outer dimension (M for scale_A) is padded to a multiple of 128 -- the leading 32 × 4 in Swizzle32x4x4
• the inner dimension (K-blocks) is padded to a multiple of 4 -- the trailing 4 in the name

Thus the required number of elements for scale_A is L * round_up(M, 128) * round_up(ceil(K / 32), 4).

scale_B is also handled similarly, with M replaced by N.

The result of this calculation is also exposed as a helper, ScaledOperand.numel_scale(quantized_shape, mode, swizzle), which we use below.

scale_A_numel = ops.ScaledOperand.numel_scale((L, M, K), mode, swizzle)
scale_B_numel = ops.ScaledOperand.numel_scale((L, N, K), mode, swizzle)

scale_A = torch.rand(scale_A_numel, device="cuda").to(scale_dtype)
scale_B = torch.rand(scale_B_numel, device="cuda").to(scale_dtype)

API contract. Operator API only requires that the scale tensor is contiguous and has numel_scale(...) elements. The Python shape is otherwise free -- the 1D buffers above are valid input as-is. If a structured shape is more convenient, any reshape preserving the element count and contiguity also works:

from cutlass.operators.utils.common import round_up

scale_A = scale_A.view(L, round_up(M, 128), -1)
print(f"scale_A: shape={tuple(scale_A.shape)}, contiguous={scale_A.is_contiguous()}")
print(f"scale_B: shape={tuple(scale_B.shape)}, contiguous={scale_B.is_contiguous()}")

Most kernels for block-scaled GEMM expect that the scale storage is already laid out in the swizzled layout used by hardware. When we set swizzle=Swizzle32x4x4, we declare that we are providing input in this layout. In a real workflow, that layout is produced by a quantization step or an earlier kernel that writes the scale values directly into it.

We don't explicitly perform that rearrangement in this tutorial. We'll use torch._scaled_mm for checking our output, which also expects data in the same swizzled layout. Since kernels interpret the storage the same way, we can cross-check our result without an explicit swizzle step.

Step 4: Pair the quantized tensor with its scale via ScaledOperand

Now we package each quantized tensor with its scale tensor, mode, and swizzle into a ScaledOperand. Once paired, GemmArguments looks the same as in the plain GEMM tutorial -- the only difference is that A and B are ScaledOperands rather than plain tensors.

args = ops.GemmArguments(
    A=ops.ScaledOperand(quantized_A, scale_A, mode, swizzle),
    B=ops.ScaledOperand(quantized_B, scale_B, mode, swizzle),
    out=out, # same as DenseTensor(out)
    accumulator_type=torch.float32,
)

Step 5: Discover, compile, and run

From this point on, everything is identical to any other operation!

Despite the different set of tensors passed to the underlying kernel, the interface with the Operator remains the same -- there is no separate BlockScaledGemmArguments, block_scaled_get_operators, or run_block_scaled API. We hand args to get_operators(...), which walks the registered Operator catalog and finds the ones compatible with the operation we described -- same as any other operation.

We describe what operation to perform (via GemmArguments); the API finds operators that know how to do it.

We then pick the first one and run it.

target_sm = ops.utils.device.device_or_env_target_sm()

operators = ops.get_operators(args, target_sm=target_sm)
assert operators, "No block-scaled GEMM operator matched the given args"
operator = operators[0]
print(f"Picked: {operator.metadata.operator_name}")

operator.run(args)

Verifying against torch._scaled_mm

To confirm the result, we cross-check out against PyTorch's reference implementation, torch._scaled_mm.

def reference_scaled_mm(A, B, scale_A, scale_B, out_dtype):
    L, M, N = A.shape[0], A.shape[1], B.shape[2]
    scale_A = scale_A.reshape(L, -1)
    scale_B = scale_B.reshape(L, -1)
    out = torch.empty((L, M, N), device=A.device, dtype=out_dtype)
    for L_idx in range(L):
        out[L_idx] = torch._scaled_mm(
            A[L_idx], B[L_idx], scale_a=scale_A[L_idx], scale_b=scale_B[L_idx], out_dtype=out_dtype
        )
    return out


reference = reference_scaled_mm(quantized_A, quantized_B, scale_A, scale_B, out_dtype)
torch.testing.assert_close(out, reference)

Recap

You've run an MXFP8 block-scaled GEMM end-to-end through the same Operator API entry points used in the basic GEMM tutorial:

1. Describe the operation with GemmArguments. Each slot is an Operand -- a logical role like "the matrix A" -- and accepts whatever storage form the operand needs: a DenseTensor (any DLPack-compatible tensor) or a ScaledOperand when the value is reconstructed from a quantized tensor + scale.
2. Discover with get_operators(args). The Operator catalog filters down to the kernels that match the arguments you described.
3. Compile and run with operator.compile(args) / operator.run(args). Identical to the basic GEMM flow.

What changed from the basic to the block-scaled tutorial is only how the operands are described: ScaledOperand instead of plain tensors. We express the operation we want to perform, and Operator API maps it to a supported kernel. The same pattern -- describe GemmArguments -> get_operators -> compile / run -- works despite different kernel implementations.