mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-07-17 09:07:47 +00:00
247 lines
7.5 KiB
Plaintext
247 lines
7.5 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "91d43c2b",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Grouped GEMM with contiguous tensors via the CUTLASS Operator API\n",
|
|
"\n",
|
|
"Note: this notebook requires a GPU with compute capability 100:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "f671f602",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import cutlass.operators as ops\n",
|
|
"\n",
|
|
"if not (status := ops.utils.device.device_or_env_supports(\"100f\")):\n",
|
|
" print(f\"This notebook requires a Blackwell GPU (sm_100f family).\\n{status.error}\")\n",
|
|
" import sys\n",
|
|
"\n",
|
|
" sys.exit(0)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "bc4adf7d",
|
|
"metadata": {},
|
|
"source": [
|
|
"This notebook shows how to use the CUTLASS Operator API to discover, compile, and execute\n",
|
|
"operators supporting contiguous offset grouped GEMMs.\n",
|
|
"\n",
|
|
"In a \"contiguous offset\" grouped GEMM, `G` different problems are executed\n",
|
|
"in which problems differ only in the `M` mode. Their problem sizes are thus\n",
|
|
"represented as:\n",
|
|
"\n",
|
|
"```text\n",
|
|
"M0 x N x K\n",
|
|
"M1 x N x K\n",
|
|
"M2 x N x K\n",
|
|
"...\n",
|
|
"M(G-1) x N x K\n",
|
|
"```\n",
|
|
"\n",
|
|
"The grouped GEMM is referred to as \"contiguous\" because operands for different\n",
|
|
"problems in the group are contained within contiguous tensors.\n",
|
|
"\n",
|
|
"Rather than having `G` different tensors for each of operands `A` and `B`, tensors\n",
|
|
"for different problems in the group are packed together:\n",
|
|
"* `A` is of shape `(TotalM, K)`, where `TotalM` is the sum of all `M` mode sizes for problems in the group.\n",
|
|
"The `A` operands for each problem in the group are stacked along the `M` mode to form this input. More on this below.\n",
|
|
"* `B` is of shape `(G, K, N)`, where `B[i, :, :]` represents the GEMM `B` operand for the `i`th problem in the group.\n",
|
|
"\n",
|
|
"For example, with `G=3` (three problems in the group), with `M` mode sizes of M0, M1, and M2,\n",
|
|
"respectively, the tensor `A` would be laid out as follows:\n",
|
|
"\n",
|
|
"```text\n",
|
|
"\n",
|
|
" +----------------------------------+ ^ \n",
|
|
" | | | | \n",
|
|
" | A0 | M0 | \n",
|
|
" | | | | \n",
|
|
" |- - - - - - - - - - - -| | \n",
|
|
" | | | |\n",
|
|
" | | | TotalM \n",
|
|
" | A1 | M1 |\n",
|
|
" | | | |\n",
|
|
" | | | | \n",
|
|
" |- - - - - - - - - - - -| | \n",
|
|
" | A2 | M2 | \n",
|
|
" +----------------------------------+ v \n",
|
|
"```\n",
|
|
"\n",
|
|
"The extents of individual `A` operands packed within the overall contiguous offset `A` tensor\n",
|
|
"are provided by an auxiliary `offsets` vector of shape `(G,)`. `offsets[i]` indicates the ending\n",
|
|
"M coordinate (exclusive) for the `i`th `A` operand.\n",
|
|
"\n",
|
|
"Thus, for the example above, `offsets = [M0, M0 + M1, M0 + M1 + M2]`.\n",
|
|
"\n",
|
|
"The output of the operation is of shape `(TotalM, N)`. The `i`th output occupies `out[start:end, :]`,\n",
|
|
"where `start` and `end` are `offsets[i-1]` and `offsets[i]`, respectively (unless `i=0`, in which case\n",
|
|
"`start` is 0).\n",
|
|
"\n",
|
|
"The reference code below shows the computation of this Operator."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "6185f60a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import torch\n",
|
|
"\n",
|
|
"\n",
|
|
"def reference_contiguous_offset_grouped_gemm(A, B, offsets, out_dtype):\n",
|
|
" G, K, N = B.shape\n",
|
|
" TotalM = A.shape[0]\n",
|
|
"\n",
|
|
" out = torch.empty((TotalM, N), dtype=out_dtype, device=A.device)\n",
|
|
"\n",
|
|
" start = 0\n",
|
|
" for i in range(G):\n",
|
|
" end = offsets[i]\n",
|
|
" out[start:end, :] = A[start:end, :] @ B[i, :, :]\n",
|
|
" start = end\n",
|
|
"\n",
|
|
" return out"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d0bf2f91",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Contiguous offset grouped GEMM in PyTorch"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "4308a6a2",
|
|
"metadata": {},
|
|
"source": [
|
|
"The same operation is performed by `torch`'s `torch._grouped_mm` (torch < 2.10)\n",
|
|
"and `torch.nn.functional.grouped_mm` (torch >= 2.10)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "043906af",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"TotalM = 8192\n",
|
|
"G = 12\n",
|
|
"K = 1024\n",
|
|
"N = 2048\n",
|
|
"\n",
|
|
"offsets = torch.arange(\n",
|
|
" TotalM // G, TotalM, TotalM // G, device=\"cuda\", dtype=torch.int32\n",
|
|
")\n",
|
|
"offsets[-1] = TotalM\n",
|
|
"\n",
|
|
"A = torch.randint(-2, 3, (TotalM, K), device=\"cuda\", dtype=torch.bfloat16)\n",
|
|
"B = torch.randint(-2, 3, (G, N, K), device=\"cuda\", dtype=torch.bfloat16).permute(\n",
|
|
" 0, 2, 1\n",
|
|
")\n",
|
|
"\n",
|
|
"out_torch = torch._grouped_mm(A, B, offsets, out_dtype=torch.bfloat16)\n",
|
|
"reference = reference_contiguous_offset_grouped_gemm(\n",
|
|
" A, B, offsets, out_dtype=torch.bfloat16\n",
|
|
")\n",
|
|
"\n",
|
|
"torch.testing.assert_close(out_torch, reference)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0d0e9479",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Contiguous offset grouped GEMM in CUTLASS Operator API\n",
|
|
"\n",
|
|
"CUTLASS Operator API exposes this contiguous offset grouped GEMM via `GroupedGemmArguments`,\n",
|
|
"which are constructed similarly to `GemmArguments`, but take in an `offsets`\n",
|
|
"tensor as well:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"id": "ff8d3ef1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"out = torch.empty((TotalM, N), device=\"cuda\", dtype=torch.bfloat16)\n",
|
|
"\n",
|
|
"args = ops.GroupedGemmArguments(\n",
|
|
" A,\n",
|
|
" B,\n",
|
|
" out,\n",
|
|
" accumulator_type=torch.float32,\n",
|
|
" offsets=offsets,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0dc6d1cb",
|
|
"metadata": {},
|
|
"source": [
|
|
"One can then use the same APIs for finding, compiling, and executing a\n",
|
|
"operator supporting this operation"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "80213e1e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"operators = ops.get_operators(args)\n",
|
|
"\n",
|
|
"assert operators, \"No operators found\"\n",
|
|
"\n",
|
|
"# Select the first operator found for simplicity\n",
|
|
"operator = operators[0]\n",
|
|
"\n",
|
|
"compiled_artifact = operator.compile(args)\n",
|
|
"\n",
|
|
"# Execute the operator\n",
|
|
"operator.run(args, compiled_artifact=compiled_artifact)\n",
|
|
"\n",
|
|
"torch.testing.assert_close(out, reference)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.5"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|