Note

This page is a rendered version of 006_block_scaled_gemm.ipynb available on GitHub.

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.

[1]:

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.

[2]:

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.

[3]:

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.

[6]:

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.

[7]:

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)

Picked: cutedsl.PersistentDenseBlockScaledGemmOperator_sm100_tnt_AFloat8E4M3FN_BFloat8E4M3FN_outFloat32_SFAFloat8E8M0FNU_SFBFloat8E8M0FNU_accFloat32_scaleScaleMode.Blockwise1x32_swizzleScaleSwizzleMode.Swizzle32x4x4_2cta_cluster2x1x1_tile128x64x128_tma_store

Verifying against torch._scaled_mm#

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

[8]:

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.