Grouped GEMM with contiguous tensors via the CUTLASS Operator API#

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

[1]:

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:

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:

+----------------------------------+         ^
|                                  |  |      |
|               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.

[2]:

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).

[3]:

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:

[4]:

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

[5]:

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)