Grouped GEMM with contiguous tensors via the CUTLASS Operator API

Note: this notebook requires a GPU with compute capability 100:

python

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)

This notebook shows how to use the CUTLASS Operator API to discover, compile, and execute operators supporting contiguous offset grouped GEMMs.

In a "contiguous offset" grouped GEMM, G different problems are executed in which problems differ only in the M mode. Their problem sizes are thus represented as:

text

M0 x N x K
M1 x N x K
M2 x N x K
...
M(G-1) x N x K

The grouped GEMM is referred to as "contiguous" because operands for different problems in the group are contained within contiguous tensors.

Rather than having G different tensors for each of operands A and B, tensors for different problems in the group are packed together:

A is of shape (TotalM, K), where TotalM is the sum of all M mode sizes for problems in the group. The A operands for each problem in the group are stacked along the M mode to form this input. More on this below.
B is of shape (G, K, N), where B[i, :, :] represents the GEMM B operand for the ith problem in the group.

For example, with G=3 (three problems in the group), with M mode sizes of M0, M1, and M2, respectively, the tensor A would be laid out as follows:

text


    +----------------------------------+         ^  
    |                                  |  |      |  
    |               A0                 |  M0     |  
    |                                  |  |      |  
    |-  -  -  -  -  -  -  -  -  -  -  -|         |  
    |                                  |  |      |
    |                                  |  |    TotalM  
    |               A1                 |  M1     |
    |                                  |  |      |
    |                                  |  |      |  
    |-  -  -  -  -  -  -  -  -  -  -  -|         |  
    |               A2                 |  M2     |  
    +----------------------------------+         v

The extents of individual A operands packed within the overall contiguous offset A tensor are provided by an auxiliary offsets vector of shape (G,). offsets[i] indicates the ending M coordinate (exclusive) for the ith A operand.

Thus, for the example above, offsets = [M0, M0 + M1, M0 + M1 + M2].

The output of the operation is of shape (TotalM, N). The ith output occupies out[start:end, :], where start and end are offsets[i-1] and offsets[i], respectively (unless i=0, in which case start is 0).

The reference code below shows the computation of this Operator.

python

import torch


def reference_contiguous_offset_grouped_gemm(A, B, offsets, out_dtype):
    G, K, N = B.shape
    TotalM = A.shape[0]

    out = torch.empty((TotalM, N), dtype=out_dtype, device=A.device)

    start = 0
    for i in range(G):
        end = offsets[i]
        out[start:end, :] = A[start:end, :] @ B[i, :, :]
        start = end

    return out

Contiguous offset grouped GEMM in PyTorch

The same operation is performed by torch's torch._grouped_mm (torch < 2.10) and torch.nn.functional.grouped_mm (torch >= 2.10).

python

TotalM = 8192
G = 12
K = 1024
N = 2048

offsets = torch.arange(
    TotalM // G, TotalM, TotalM // G, device="cuda", dtype=torch.int32
)
offsets[-1] = TotalM

A = torch.randint(-2, 3, (TotalM, K), device="cuda", dtype=torch.bfloat16)
B = torch.randint(-2, 3, (G, N, K), device="cuda", dtype=torch.bfloat16).permute(
    0, 2, 1
)

out_torch = torch._grouped_mm(A, B, offsets, out_dtype=torch.bfloat16)
reference = reference_contiguous_offset_grouped_gemm(
    A, B, offsets, out_dtype=torch.bfloat16
)

torch.testing.assert_close(out_torch, reference)

Contiguous offset grouped GEMM in CUTLASS Operator API

CUTLASS Operator API exposes this contiguous offset grouped GEMM via GroupedGemmArguments, which are constructed similarly to GemmArguments, but take in an offsets tensor as well:

python

out = torch.empty((TotalM, N), device="cuda", dtype=torch.bfloat16)

args = ops.GroupedGemmArguments(
    A,
    B,
    out,
    accumulator_type=torch.float32,
    offsets=offsets,
)

One can then use the same APIs for finding, compiling, and executing a operator supporting this operation

python

operators = ops.get_operators(args)

assert operators, "No operators found"

# Select the first operator found for simplicity
operator = operators[0]

compiled_artifact = operator.compile(args)

# Execute the operator
operator.run(args, compiled_artifact=compiled_artifact)

torch.testing.assert_close(out, reference)