leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.4 tag release update. (#3032)

Browse Source
This commit is contained in:
Junkai-Wu
2026-02-14 12:27:58 +08:00
committed by GitHub
parent 01687cfba1
commit d4bbf728ca
140 changed files with 41624 additions and 3691 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py
View File

@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
@@ -9,13 +9,11 @@
# its affiliates is strictly prohibited.
import argparse
import torch
from typing import Tuple
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
@@ -32,9 +30,8 @@ with optimizations for challenges that may arise with other problem sizes.
To run this example:
.. code-block:: bash
python examples/blackwell/tutorial_fp16_gemm_0.py \
--mnk 8192,8192,8192 \
--tolerance 1e-01
python examples/blackwell/tutorial_gemm/fp16_gemm_0.py \
--mnk 8192,8192,8192
Constraints for this example:
* The problem size of m and n must be divisible by the tile size m & n (128, 256)
@@ -128,7 +125,8 @@ def kernel(
num_stages=acc_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, threads_per_cta
pipeline.Agent.Thread,
threads_per_cta,
),
barrier_storage=storage.acc_mbar_ptr.data_ptr(),
).make_participants()
@@ -141,15 +139,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(0)
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -188,7 +186,7 @@ def kernel(
# (EpiTile, NumTiles)
gC_epi = cute.zipped_divide(tCgC, epi_tiler)
# Every thread loads 32x128 bits
# Every thread loads 64 x fp32
tmem_atom = cute.make_copy_atom(
tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
cutlass.Float32,
@@ -273,11 +271,7 @@ def kernel(
@cute.jit
def host_function(
a: cute.Tensor,
b: cute.Tensor,
c: cute.Tensor,
):
def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
# Construct tiled MMA
op = tcgen05.MmaF16BF16Op(
io_dtype,
@@ -354,6 +348,10 @@ def run_dense_gemm(
mnk: Tuple[int, int, int],
tolerance: float,
):
global torch, cutlass_torch
import torch
import cutlass.torch as cutlass_torch
print("===================================================================")
print("Running Blackwell fp16 GEMM example 0 with:")
print(f" mnk: {mnk}")
@@ -393,12 +391,7 @@ def run_dense_gemm(
)
# Entry point to the host JIT function
host_function(
a_tensor,
b_tensor,
c_tensor,
no_cache=True,
)
host_function(a_tensor, b_tensor, c_tensor, no_cache=True)
# Compute reference result and verify
ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu()
@@ -418,7 +411,11 @@ if __name__ == "__main__":
"Invalid format. Expected comma-separated integers."
)
if not torch.cuda.is_available():
from cuda.bindings import driver as cu_driver
cu_driver.cuInit(0)
err, device_count = cu_driver.cuDeviceGetCount()
if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
raise RuntimeError("A GPU is required to run this example")
parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 0")

34
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py
View File

@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
@@ -13,15 +13,14 @@
import argparse
import torch
from typing import Tuple
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
"""
@@ -78,7 +77,6 @@ acc_stage = 1
@cute.struct
class SharedStorage:
# each stage has 2 kinds of barrier, i.e. empty & full
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
@@ -174,15 +172,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -256,9 +254,9 @@ def kernel(
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor_like(tDgC[None, None, 0], acc_dtype)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor_like(tDgC[None, None, 0], io_dtype)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
@@ -356,13 +354,13 @@ def host_function(
tiled_mma = cute.make_tiled_mma(op)
# Construct SMEM layouts for A and B
a_smem_layout = utils.sm100.make_smem_layout_a(
a_smem_layout = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a.element_type,
ab_stages,
)
b_smem_layout = utils.sm100.make_smem_layout_b(
b_smem_layout = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b.element_type,
@@ -383,7 +381,7 @@ def host_function(
a_smem_layout_one_stage,
mma_tiler_mnk,
tiled_mma,
cta_layout_vmnk.shape,
cta_layout_vmnk.shape, # take the layout and extract the shape internally
)
b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
op,
@@ -438,6 +436,10 @@ def run_dense_gemm(
mnk: Tuple[int, int, int],
tolerance: float,
):
global torch, cutlass_torch
import torch
import cutlass.torch as cutlass_torch
print("===================================================================")
print("Running Blackwell fp16 GEMM example 1 with:")
print(f" mnk: {mnk}")
@@ -501,7 +503,11 @@ if __name__ == "__main__":
"Invalid format. Expected comma-separated integers."
)
if not torch.cuda.is_available():
from cuda.bindings import driver as cu_driver
cu_driver.cuInit(0)
err, device_count = cu_driver.cuDeviceGetCount()
if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
raise RuntimeError("A GPU is required to run this example")
parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 1")

778
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py Normal file
View File

@@ -0,0 +1,778 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import os
import sys
from typing import Type, Tuple
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import make_ptr
if __name__ == "__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
examples_dir = os.path.join(current_dir, "..", "..")
if examples_dir not in sys.path:
sys.path.insert(0, examples_dir)
from blackwell.tutorial_gemm.utils import create_parser, run
mma_tiler_mn = (128, 256)
mma_inst_shape_k = 64
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16
"""
The first tutorial NVFP4 block-scaled batched GEMM demonstrating a simple kernel implementation in CuTeDSL
This example demonstrates the kernel implementation of block-scaled batched GEMM with NVFP4 data type.
With large tile sizes (128x256x256), it can achieve very high performance on 8k×8k×8k problem sizes.
It can serve as a starting point to help users quickly experiment with optimizations for
challenges that may arise with other problem sizes.
To run this example:
.. code-block:: bash
python examples/blackwell/tutorial_gemm/nvfp4_gemm_0.py \
--mnkl 8192,8192,8192,1 --do_benchmark
Constraints for this example:
* The problem size of m, n and k must be divisible by the tile size m&n&k (128,256,256)
* The scaling factor vector size is 16.
* The A/B matrices have data contiguous on the k dimension.
* The C matrix has data contiguous on the n dimension.
* The A/B matrix data type is Float4E2M1FN.
* The SFA/SFB matrix data type is Float8E4M3FN.
"""
class Sm100BlockScaledDenseGemmKernel:
def __init__(self):
self.threads_per_cta = 128
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
self.num_tmem_alloc_cols = 512
# set stages for ab_pipeline and acc_pipeline
self.num_acc_stage = 1
self.num_ab_stage = 4
@cute.jit
def __call__(
self,
a_ptr: cute.Pointer,
b_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb_ptr: cute.Pointer,
c_ptr: cute.Pointer,
problem_size: tuple,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
):
# setup static attributes before smem/grid/tma computation
self.c_layout = utils.LayoutEnum.ROW_MAJOR
m, n, k, l = problem_size
# Setup attributes that depend on gemm inputs
mma_inst_tile_k = 4
self.mma_tiler = (
mma_tiler_mn[0],
mma_tiler_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0],
self.mma_tiler[1],
self.mma_tiler[2],
)
a_tensor = cute.make_tensor(
a_ptr,
cute.make_layout(
(m, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
),
)
b_tensor = cute.make_tensor(
b_ptr,
cute.make_layout(
(n, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
),
)
# make address offset of c_tensor 256bit aligned,
# so that epilogue could use vectorized store with larger vector size.
c_tensor = cute.make_tensor(
c_ptr,
cute.make_layout(
(cute.assume(m, 32), cute.assume(n, 16), l),
stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)),
),
)
# Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
# ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
a_tensor.shape, sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
b_tensor.shape, sf_vec_size
)
sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
mma_op = tcgen05.MmaMXF4NVF4Op(
sf_dtype,
(*mma_tiler_mn, mma_inst_shape_k),
tcgen05.CtaGroup.ONE,
tcgen05.OperandSource.SMEM,
)
tiled_mma = cute.make_tiled_mma(mma_op)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((1, 1, 1)),
(tiled_mma.thr_id.shape,),
)
# Compute A/B/SFA/SFB/C shared memory layout
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# TMA load for A
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
a_tensor,
a_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# TMA load for B
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
b_tensor,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# TMA load for SFA
sfa_smem_layout = cute.slice_(
self.sfa_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
sfa_tensor,
sfa_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
internal_type=cutlass.Int16,
)
# TMA load for SFB
sfb_smem_layout = cute.slice_(
self.sfb_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
sfb_tensor,
sfb_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
internal_type=cutlass.Int16,
)
# Compute TMA load bytes
a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
self.num_tma_load_bytes = (
a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
) * atom_thr_size
# Compute grid size
grid = (
cute.ceil_div(c_tensor.shape[0], self.cta_tile_shape_mnk[0]),
cute.ceil_div(c_tensor.shape[1], self.cta_tile_shape_mnk[1]),
c_tensor.shape[2],
)
# Launch the kernel synchronously
self.kernel(
tiled_mma,
tma_atom_a,
tma_tensor_a,
tma_atom_b,
tma_tensor_b,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb,
tma_tensor_sfb,
c_tensor,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
self.sfa_smem_layout_staged,
self.sfb_smem_layout_staged,
epilogue_op,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(1, 1, 1),
stream=stream,
)
return
# GPU device kernel
@cute.kernel
def kernel(
self,
tiled_mma: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b: cute.CopyAtom,
mB_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb: cute.CopyAtom,
mSFB_nkl: cute.Tensor,
mC_mnl: cute.Tensor,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
epilogue_op: cutlass.Constexpr,
):
"""
GPU device kernel performing the batched GEMM computation.
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
tidx, _, _ = cute.arch.thread_idx()
#
# Setup cta/thread coordinates
#
# Coords inside cluster
bidx, bidy, bidz = cute.arch.block_idx()
# Coords outside cluster
cta_coord = (bidx, bidy, bidz)
mma_tile_coord_mnl = (
cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
cta_coord[1],
cta_coord[2],
)
#
# Define shared storage for kernel
#
@cute.struct
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
# (MMA, MMA_M, MMA_K, STAGE)
sA = smem.allocate_tensor(
element_type=ab_dtype,
layout=a_smem_layout_staged.outer,
byte_alignment=128,
swizzle=a_smem_layout_staged.inner,
)
# (MMA, MMA_N, MMA_K, STAGE)
sB = smem.allocate_tensor(
element_type=ab_dtype,
layout=b_smem_layout_staged.outer,
byte_alignment=128,
swizzle=b_smem_layout_staged.inner,
)
# (MMA, MMA_M, MMA_K, STAGE)
sSFA = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfa_smem_layout_staged,
byte_alignment=128,
)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfb_smem_layout_staged,
byte_alignment=128,
)
#
# Initialize mainloop ab_pipeline, acc_pipeline and their states
#
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
).make_participants()
acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread,
self.threads_per_cta,
),
).make_participants()
#
# Local_tile partition global tensors
#
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gSFB_nkl = cute.local_tile(
mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/SFA/SFB/C
#
# (MMA, MMA_M, MMA_K, RestK)
thr_mma = tiled_mma.get_slice(0)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB = thr_mma.partition_B(gSFB_nkl)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
# Partition global/shared tensor for TMA load A/B/SFA/SFB
#
# TMA load A partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
0,
cute.make_layout(1),
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
# TMA load B partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
0,
cute.make_layout(1),
cute.group_modes(sB, 0, 3),
cute.group_modes(tCgB, 0, 3),
)
# TMA load partition for SFA tensor
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsSFA, tAgSFA = cpasync.tma_partition(
tma_atom_sfa,
0,
cute.make_layout(1),
cute.group_modes(sSFA, 0, 3),
cute.group_modes(tCgSFA, 0, 3),
)
tAsSFA = cute.filter_zeros(tAsSFA)
tAgSFA = cute.filter_zeros(tAgSFA)
# TMA load partition for SFB tensor
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsSFB, tBgSFB = cpasync.tma_partition(
tma_atom_sfb,
0,
cute.make_layout(1),
cute.group_modes(sSFB, 0, 3),
cute.group_modes(tCgSFB, 0, 3),
)
tBsSFB = cute.filter_zeros(tBsSFB)
tBgSFB = cute.filter_zeros(tBgSFB)
#
# Partition shared/tensor memory tensor for TiledMMA_A/B/C
#
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
# (MMA, MMA_M, MMA_N)
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
#
# Alloc tensor memory buffer
#
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=tmem_alloc_barrier,
)
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
#
# Make SFA/SFB tmem tensor
#
# Get SFA tmem ptr
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
dtype=sf_dtype,
)
# (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
# Get SFB tmem ptr
sfb_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA),
dtype=sf_dtype,
)
# (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
)
tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
#
# Partition for S2T copy of SFA/SFB
#
# Make S2T CopyAtom
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
sf_dtype,
)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFA_compact = cute.filter_zeros(sSFA)
# (MMA, MMA_MN, MMA_K)
tCtSFA_compact = cute.filter_zeros(tCtSFA)
tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFB_compact = cute.filter_zeros(sSFB)
# (MMA, MMA_MN, MMA_K)
tCtSFB_compact = cute.filter_zeros(tCtSFB)
tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), RestK)
tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
#
# Execute Data copy and Math computation in the k_tile loop
#
if warp_idx == 0:
# Wait for accumulator buffer empty
acc_empty = acc_producer.acquire_and_advance()
# Set ACCUMULATE field to False for the first k_tile iteration
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
# Execute k_tile loop
for k_tile in cutlass.range(
k_tile_cnt, prefetch_stages=self.num_ab_stage - 2
):
# Wait for AB buffer empty
ab_empty = ab_producer.acquire_and_advance()
# TMA load for A/B/SFA/SFB
cute.copy(
tma_atom_a,
tAgA[(None, ab_empty.count)],
tAsA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_b,
tBgB[(None, ab_empty.count)],
tBsB[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_sfa,
tAgSFA[(None, ab_empty.count)],
tAsSFA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_sfb,
tBgSFB[(None, ab_empty.count)],
tBsSFB[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
# Wait for AB buffer full
ab_full = ab_consumer.wait_and_advance()
# Copy SFA/SFB to tmem
s2t_stage_coord = (None, None, None, None, ab_full.index)
tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB_compact_s2t_staged,
tCtSFB_compact_s2t,
)
# tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kblock_idx,
ab_full.index,
)
# Set SFA/SFB tensor to tiled_mma
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(
tcgen05.Field.SFA,
tCtSFA[sf_kblock_coord].iterator,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_full.release()
acc_empty.commit()
#
# Epilogue
# Partition for epilogue
#
op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc)
# (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(tCgC)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rC = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype
)
# STG Atom
simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype)
tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
# Release TMEM allocation lock
tmem.relinquish_alloc_permit()
# Wait for accumulator buffer full
acc_full = acc_consumer.wait_and_advance()
# Copy accumulator to register
cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype))
tTR_rC.store(acc_vec)
# Store C to global memory
cute.copy(simt_atom, tTR_rC, tTR_gC)
acc_full.release()
# Deallocate TMEM
cute.arch.barrier()
tmem.free(acc_tmem_ptr)
return
def run_nvfp4_gemm(
mnkl: Tuple[int, int, int, int],
tolerance: float,
do_benchmark: bool = False,
warmup_iterations: int = 10,
iterations: int = 100,
use_cold_l2: bool = True,
):
run(
gemm_class=Sm100BlockScaledDenseGemmKernel,
ab_dtype=ab_dtype,
sf_dtype=sf_dtype,
c_dtype=c_dtype,
sf_vec_size=sf_vec_size,
mma_tiler_mn=mma_tiler_mn,
cluster_shape_mnk=(1, 1, 1),
mnkl=mnkl,
tolerance=tolerance,
do_benchmark=do_benchmark,
warmup_iterations=warmup_iterations,
iterations=iterations,
use_cold_l2=use_cold_l2,
)
if __name__ == "__main__":
parser = create_parser()
args = parser.parse_args()
if len(args.mnkl) != 4:
parser.error("--mnkl must contain exactly 4 values")
m, n, k, _ = args.mnkl
if m % mma_tiler_mn[0] != 0:
parser.error("m must be multiples of mma_tiler_mn[0] (got m={})".format(m))
if n % mma_tiler_mn[1] != 0:
parser.error("n must be multiples of mma_tiler_mn[1] (got n={})".format(n))
if k % 256 != 0:
parser.error("k must be a multiple of 256 (got k={})".format(k))
run_nvfp4_gemm(
args.mnkl,
args.tolerance,
args.do_benchmark,
)
print("PASS")

934
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py Normal file
View File

@@ -0,0 +1,934 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# This is the second tutorial nvfp4 GEMM. It builds on the first tutorial by adding 2CTA MMA
# instructions with a 2x1 cluster.
import argparse
import os
import sys
from typing import Type, Tuple
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import from_dlpack, make_ptr
if __name__ == "__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
examples_dir = os.path.join(current_dir, "..", "..")
if examples_dir not in sys.path:
sys.path.insert(0, examples_dir)
from blackwell.tutorial_gemm.utils import create_parser, run
mma_tiler_mn = (256, 256)
mma_inst_shape_k = 64
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16
cluster_shape_mnk = (2, 1, 1)
"""
The second tutorial further improves the performance of NVFP4 block-scaled batched GEMM
by adding 2CTA instructions and TMA multicast optimizations.
(1) The 2 CTA instructions could reduce the smem size requirement for B tensor,
increased num_ab_stage and improves the latency hiding capability.
For both 1CTA and 2CTA, the shared memory (smem) size per stage for the A, sfA, and sfB tensors is the same:
- For the A tensor, each stage requires 128 x 256 x sizeof(float4) = 16KB.
- For the sfA tensor, each stage requires 128 x (256 / 16) x sizeof(float8) = 2KB.
- For the sfB tensor, each stage requires 256 x (256 / 16) x sizeof(float8) = 4KB.
The situation is different for the B tensor:
- In the 1CTA case, each stage for the B tensor requires 256 x 256 x sizeof(float4) = 32KB.
- In the 2CTA case, only half this size is needed, i.e., 128 x 256 x sizeof(float4) = 16KB.
Therefore, the maximum number of AB stages is:
- For 1CTA: 227 // (16 + 32 + 2 + 4) = 4
- For 2CTA: 227 // (16 + 16 + 2 + 4) = 5
The latency hiding capability is:
- 1CTA: 512 * (4 - 1) = 1.5K cycles
- 2CTA: 512 * (5 - 1) = 2K cycles
(2) TMA multicast can help reduce L2 cache traffic.
Without TMA multicast, the L2 traffic per tile is typically 16KB + 32KB = 48KB (possibly less in practice, depending on hardware optimizations).
With TMA multicast in a cluster of shape (m, n), the L2 traffic per tile is reduced to 16KB / n + 32KB / m.
For example:
- In a 2x1 cluster: 16KB / 1 + 32KB / 2 = 24KB per tile
- In a 4x4 cluster: 16KB / 4 + 32KB / 4 = 12KB per tile
The first approach offers substantial capacity for hiding latency, whereas the second reduces the time required for data to become ready.
Both could be tried when the workload is latency-bound or limited by memory throughput.
To run this example:
.. code-block:: bash
python examples/blackwell/tutorial_gemm/nvfp4_gemm_1.py \
--mnkl 8192,8192,8192,1 --do_benchmark
Constraints for this example:
* The problem size of m, n and k must be divisible by the tile size m&n&k (256, 256, 256)
* The scaling factor vector size is 16.
* The A/B matrices have data contiguous on the k dimension.
* The C matrix has data contiguous on the n dimension.
* The A/B matrix data type is Float4E2M1FN.
* The SFA/SFB matrix data type is Float8E4M3FN.
"""
class Sm100BlockScaledDenseGemmKernel:
def __init__(self):
self.threads_per_cta = 128
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
self.num_tmem_alloc_cols = 512
# set stages for ab_pipeline and acc_pipeline
self.num_acc_stage = 1
self.num_ab_stage = 5
@cute.jit
def __call__(
self,
a_ptr: cute.Pointer,
b_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb_ptr: cute.Pointer,
c_ptr: cute.Pointer,
problem_size: tuple,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
):
# setup static attributes before smem/grid/tma computation
self.c_layout = utils.LayoutEnum.ROW_MAJOR
m, n, k, l = problem_size
self.use_2cta_instrs = False if mma_tiler_mn[0] == 128 else True
# Setup attributes that depend on gemm inputs
mma_inst_tile_k = 4
self.mma_tiler = (
mma_tiler_mn[0],
mma_tiler_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.mma_inst_shape_sfb = (
mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1),
mma_tiler_mn[1],
mma_inst_shape_k,
)
self.mma_tiler_sfb = (
self.mma_inst_shape_sfb[0],
self.mma_inst_shape_sfb[1],
mma_inst_shape_k * mma_inst_tile_k,
)
a_tensor = cute.make_tensor(
a_ptr,
cute.make_layout(
(m, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
),
)
b_tensor = cute.make_tensor(
b_ptr,
cute.make_layout(
(n, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
),
)
# 256bit aligned. row_major
c_tensor = cute.make_tensor(
c_ptr,
cute.make_layout(
(cute.assume(m, 32), cute.assume(n, 16), l),
stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)),
),
)
# Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
# ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
a_tensor.shape, sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
b_tensor.shape, sf_vec_size
)
sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
mma_op = tcgen05.MmaMXF4NVF4Op(
sf_dtype,
(*mma_tiler_mn, mma_inst_shape_k),
tcgen05.CtaGroup.ONE if not self.use_2cta_instrs else tcgen05.CtaGroup.TWO,
tcgen05.OperandSource.SMEM,
)
tiled_mma = cute.make_tiled_mma(mma_op)
# (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
# Note sfB don't support share among 2ctas
sfb_mma_op = tcgen05.MmaMXF4NVF4Op(
sf_dtype,
self.mma_inst_shape_sfb,
tcgen05.CtaGroup.ONE,
tcgen05.OperandSource.SMEM,
)
tiled_mma_sfb = cute.make_tiled_mma(sfb_mma_op)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // (2 if self.use_2cta_instrs else 1),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cta_tile_shape_mnk_sfb = (
self.mma_tiler_sfb[0] // (2 if self.use_2cta_instrs else 1),
self.mma_tiler_sfb[1],
self.mma_tiler_sfb[2],
)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout(cluster_shape_mnk),
(tiled_mma.thr_id.shape,),
)
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout(cluster_shape_mnk),
(tiled_mma_sfb.thr_id.shape,),
)
# Compute number of multicast CTAs for A/B
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
self.is_a_mcast = self.num_mcast_ctas_a > 1
self.is_b_mcast = self.num_mcast_ctas_b > 1
self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
# Compute A/B/SFA/SFB/C shared memory layout
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
cluster_shape_mnk[:2], tiled_mma.thr_id
)
# TMA load for A
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op,
a_tensor,
a_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# TMA load for B
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
cluster_shape_mnk[:2], tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b_tensor,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# TMA load for SFA
sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
cluster_shape_mnk[:2], tiled_mma.thr_id
)
sfa_smem_layout = cute.slice_(
self.sfa_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
sfa_op,
sfa_tensor,
sfa_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
internal_type=cutlass.Int16,
)
# TMA load for SFB
sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
cluster_shape_mnk[:2], tiled_mma.thr_id
)
sfb_smem_layout = cute.slice_(
self.sfb_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb_tensor,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
# Compute TMA load bytes
a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
self.num_tma_load_bytes = (
a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
) * atom_thr_size
# Compute grid size
grid = cute.round_up(
cute.ceil_div(
(c_tensor.layout.shape),
(self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], 1),
),
cluster_shape_mnk,
)
# Launch the kernel
self.kernel(
tiled_mma,
tiled_mma_sfb,
tma_atom_a,
tma_tensor_a,
tma_atom_b,
tma_tensor_b,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb,
tma_tensor_sfb,
c_tensor,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
self.sfa_smem_layout_staged,
self.sfb_smem_layout_staged,
self.cluster_layout_vmnk,
self.cluster_layout_sfb_vmnk,
epilogue_op,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=cluster_shape_mnk,
stream=stream,
)
return
# GPU device kernel
@cute.kernel
def kernel(
self,
tiled_mma: cute.TiledMma,
tiled_mma_sfb: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b: cute.CopyAtom,
mB_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb: cute.CopyAtom,
mSFB_nkl: cute.Tensor,
mC_mnl: cute.Tensor,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
cta_layout_vmnk: cute.Layout,
cta_layout_sfb_vmnk: cute.Layout,
epilogue_op: cutlass.Constexpr,
):
"""
GPU device kernel performing the batched GEMM computation.
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
tidx, _, _ = cute.arch.thread_idx()
#
# Setup cta/thread coordinates
#
# Coords inside cluster
bidx, bidy, bidz = cute.arch.block_idx()
cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
cta_in_cluster_coord_sfb_vmnk = cta_layout_sfb_vmnk.get_flat_coord(
cta_rank_in_cluster
)
# Coords outside cluster
mma_tile_coord_vmnk = (
bidx % cute.size(cta_layout_vmnk, mode=[0]),
bidx // cute.size(cta_layout_vmnk, mode=[0]),
bidy,
bidz,
)
mma_tile_coord_mnl = mma_tile_coord_vmnk[1:]
is_leader_cta = mma_tile_coord_vmnk[0] == 0
#
# Define shared storage for kernel
#
@cute.struct
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
# (MMA, MMA_M, MMA_K, STAGE)
sA = smem.allocate_tensor(
element_type=ab_dtype,
layout=a_smem_layout_staged.outer,
byte_alignment=128,
swizzle=a_smem_layout_staged.inner,
)
# (MMA, MMA_N, MMA_K, STAGE)
sB = smem.allocate_tensor(
element_type=ab_dtype,
layout=b_smem_layout_staged.outer,
byte_alignment=128,
swizzle=b_smem_layout_staged.inner,
)
# (MMA, MMA_M, MMA_K, STAGE)
sSFA = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfa_smem_layout_staged,
byte_alignment=128,
)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfb_smem_layout_staged,
byte_alignment=128,
)
#
# Compute multicast mask for A/B/SFA/SFB buffer full
#
a_full_mcast_mask = None
b_full_mcast_mask = None
sfa_full_mcast_mask = None
sfb_full_mcast_mask = None
if cutlass.const_expr(
self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs
):
a_full_mcast_mask = cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
)
b_full_mcast_mask = cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
)
sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
)
sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
cta_layout_sfb_vmnk, cta_in_cluster_coord_sfb_vmnk, mcast_mode=1
)
#
# Initialize mainloop ab_pipeline, acc_pipeline and their states
#
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
ab_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_tma_producer
)
ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cta_layout_vmnk,
).make_participants()
acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread,
self.threads_per_cta * (2 if self.use_2cta_instrs else 1),
),
cta_layout_vmnk=cta_layout_vmnk,
).make_participants()
#
# Local_tile partition global tensors
#
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gSFB_nkl = cute.local_tile(
mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/SFA/SFB/C
#
# (MMA, MMA_M, MMA_K, RestK)
thr_mma = tiled_mma.get_slice(mma_tile_coord_vmnk[0])
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_vmnk[0])
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
# tCgSFB = thr_mma.partition_B(gSFB_nkl)
tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
# Partition global/shared tensor for TMA load A/B/SFA/SFB
#
# TMA load A partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
# 0,
# cute.make_layout(1),
cta_in_cluster_coord_vmnk[2],
cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
# TMA load B partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
# 0,
# cute.make_layout(1),
cta_in_cluster_coord_vmnk[1],
cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
cute.group_modes(sB, 0, 3),
cute.group_modes(tCgB, 0, 3),
)
# TMA load SFA partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsSFA, tAgSFA = cpasync.tma_partition(
tma_atom_sfa,
# 0,
# cute.make_layout(1),
cta_in_cluster_coord_vmnk[2],
cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
cute.group_modes(sSFA, 0, 3),
cute.group_modes(tCgSFA, 0, 3),
)
tAsSFA = cute.filter_zeros(tAsSFA)
tAgSFA = cute.filter_zeros(tAgSFA)
# TMA load SFB partition_S/D
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
sfb_cta_layout = cute.make_layout(
cute.slice_(cta_layout_sfb_vmnk, (0, None, 0, 0)).shape
)
tBsSFB, tBgSFB = cpasync.tma_partition(
tma_atom_sfb,
cta_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB, 0, 3),
cute.group_modes(tCgSFB, 0, 3),
)
tBsSFB = cute.filter_zeros(tBsSFB)
tBgSFB = cute.filter_zeros(tBgSFB)
#
# Partition shared/tensor memory tensor for TiledMMA_A/B/C
#
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
# (MMA, MMA_M, MMA_N)
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
#
# Alloc tensor memory buffer
#
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
)
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
#
# Make SFA/SFB tmem tensor
#
# Get SFA tmem ptr
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
dtype=sf_dtype,
)
# (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
# Get SFB tmem ptr
sfb_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA),
dtype=sf_dtype,
)
# (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
)
tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
#
# Partition for S2T copy of SFA/SFB
#
# Make S2T CopyAtom
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(
tcgen05.CtaGroup.ONE
if not self.use_2cta_instrs
else tcgen05.CtaGroup.TWO
),
sf_dtype,
)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFA_compact = cute.filter_zeros(sSFA)
# (MMA, MMA_MN, MMA_K)
tCtSFA_compact = cute.filter_zeros(tCtSFA)
tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFB_compact = cute.filter_zeros(sSFB)
# (MMA, MMA_MN, MMA_K)
tCtSFB_compact = cute.filter_zeros(tCtSFB)
tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), RestK)
tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
#
# Execute Data copy and Math computation in the k_tile loop
#
if warp_idx == 0:
# Wait for accumulator buffer empty
if is_leader_cta:
acc_producer.acquire_and_advance()
# Set ACCUMULATE field to False for the first k_tile iteration
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
# Execute k_tile loop
for k_tile in cutlass.range(
k_tile_cnt, prefetch_stages=self.num_ab_stage - 2
):
# Wait for AB buffer empty
ab_empty = ab_producer.acquire_and_advance()
# TMA load A/B/SFA/SFB
cute.copy(
tma_atom_a,
tAgA[(None, ab_empty.count)],
tAsA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
mcast_mask=a_full_mcast_mask,
)
cute.copy(
tma_atom_b,
tBgB[(None, ab_empty.count)],
tBsB[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_sfa,
tAgSFA[(None, ab_empty.count)],
tAsSFA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
mcast_mask=sfa_full_mcast_mask,
)
cute.copy(
tma_atom_sfb,
tBgSFB[(None, ab_empty.count)],
tBsSFB[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
mcast_mask=sfb_full_mcast_mask,
)
if is_leader_cta:
# Wait for AB buffer full
ab_full = ab_consumer.wait_and_advance()
# Copy SFA/SFB to tmem
s2t_stage_coord = (None, None, None, None, ab_full.index)
tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB_compact_s2t_staged,
tCtSFB_compact_s2t,
)
# tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kblock_idx,
ab_full.index,
)
# Set SFA/SFB tensor to tiled_mma
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(
tcgen05.Field.SFA,
tCtSFA[sf_kblock_coord].iterator,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_full.release()
if is_leader_cta:
acc_producer.commit()
#
# Epilogue
# Partition for epilogue
#
# x32 or x128 all is ok.
op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc)
# (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(tCgC)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rC = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype
)
# STG Atom
simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype)
tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
# Wait for accumulator buffer full
acc_full = acc_consumer.wait_and_advance()
# Copy accumulator to register
cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype))
tTR_rC.store(acc_vec)
# Store C to global memory
cute.copy(simt_atom, tTR_rC, tTR_gC)
acc_full.release()
# Ensure used buffers are properly synchronized before producer exit.
# This could avoid the invalid dsmem access due to early leading CTA exit.
if warp_idx == 0:
ab_producer.tail()
if is_leader_cta:
acc_producer.tail()
# Deallocate TMEM
cute.arch.barrier()
tmem.free(acc_tmem_ptr)
return
def run_nvfp4_gemm(
mnkl: Tuple[int, int, int, int],
tolerance: float,
warmup_iterations: int = 10,
iterations: int = 100,
use_cold_l2: bool = True,
do_benchmark: bool = False,
):
run(
gemm_class=Sm100BlockScaledDenseGemmKernel,
ab_dtype=ab_dtype,
sf_dtype=sf_dtype,
c_dtype=c_dtype,
sf_vec_size=sf_vec_size,
mma_tiler_mn=mma_tiler_mn,
cluster_shape_mnk=cluster_shape_mnk,
mnkl=mnkl,
tolerance=tolerance,
do_benchmark=do_benchmark,
warmup_iterations=warmup_iterations,
iterations=iterations,
use_cold_l2=use_cold_l2,
)
if __name__ == "__main__":
parser = create_parser()
args = parser.parse_args()
if len(args.mnkl) != 4:
parser.error("--mnkl must contain exactly 4 values")
m, n, k, _ = args.mnkl
if m % mma_tiler_mn[0] != 0:
parser.error("M must be multiples of mma_tiler_mn[0] (got m={})".format(m))
if n % mma_tiler_mn[1] != 0:
parser.error("N must be multiples of mma_tiler_mn[1] (got n={})".format(n))
if k % 256 != 0:
parser.error("k must be a multiple of 256 (got k={})".format(k))
run_nvfp4_gemm(
args.mnkl,
args.tolerance,
do_benchmark=args.do_benchmark,
)
print("PASS")

366
examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py Normal file
View File

@@ -0,0 +1,366 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
from typing import Tuple
import torch
import cutlass
import cutlass.cute as cute
import cutlass.torch as cutlass_torch
from cutlass.cute.runtime import make_ptr
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
try:
return tuple(int(x.strip()) for x in s.split(","))
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
def create_parser():
parser = argparse.ArgumentParser(
description="Example of Sm100 Dense BlockScaled GEMM."
)
parser.add_argument(
"--mnkl",
type=parse_comma_separated_ints,
default=(8192, 8192, 8192, 8),
help="mnkl dimensions (comma-separated)",
)
parser.add_argument(
"--tolerance", type=float, default=1e-01, help="Tolerance for validation"
)
parser.add_argument(
"--do_benchmark", action="store_true", default=False, help="Do benchmark test"
)
return parser
def ceil_div(a, b):
return (a + b - 1) // b
# Helper function to create scale factor tensor SFA/SFB
# for 1x16 block scaled wise use case and follow the layout requirement
# defined in https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout
@cute.jit
def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
sf_ref_ptr: cute.Pointer,
sf_mma_ptr: cute.Pointer,
mn: int,
sf_k: int,
l: int,
mma_shape: tuple,
):
mma_permute_order = (3, 4, 1, 5, 2, 0)
permuted_shape = tuple(mma_shape[i] for i in mma_permute_order)
cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5))
sf_ref_tensor = cute.make_tensor(
sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k))
)
sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout)
sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
for i in cutlass.range(cute.size(sf_ref_tensor)):
mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
pass
def to_blocked(input_matrix):
rows, cols = input_matrix.shape
# Please ensure rows and cols are multiples of 128 and 4 respectively
n_row_blocks = ceil_div(rows, 128)
n_col_blocks = ceil_div(cols, 4)
padded = input_matrix
blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
return rearranged.flatten()
def run(
gemm_class,
ab_dtype,
sf_dtype,
c_dtype,
sf_vec_size,
mma_tiler_mn,
cluster_shape_mnk,
mnkl: Tuple[int, int, int, int],
tolerance: float,
warmup_iterations: int = 10,
iterations: int = 100,
use_cold_l2: bool = True,
do_benchmark: bool = False,
):
"""
Prepare A/B/SFA/SFB/C tensors, launch GPU kernel, and reference checking.
"""
print("=" * 60)
print("Launching Blackwell Dense BlockScaled GEMM Test")
print("-" * 60)
print(f"Input dimensions (m, n, k, l): {mnkl}")
print(f" m (rows): {mnkl[0]}")
print(f" n (cols): {mnkl[1]}")
print(f" k (inner): {mnkl[2]}")
print(f" l (batch): {mnkl[3]}")
print(f"Data Types & Precision:")
print(f" Input matrices (A, B): {ab_dtype}")
print(f" Scale factors (SFA, SFB): {sf_dtype}")
print(f" Output matrix (C): {c_dtype}")
print(f" Scale factor vector size: {sf_vec_size}")
print("Tile and cluster configuration:")
print(f" MMA tiler (M, N, K): {mma_tiler_mn}")
print(f" Cluster shape (M, N, K): {cluster_shape_mnk}")
print(f"Validation tolerance: {tolerance}")
print(f"Do benchmark: {do_benchmark}")
print("=" * 60)
# Unpack parameters
m, n, k, l = mnkl
if not torch.cuda.is_available():
raise RuntimeError("GPU is required to run this example!")
torch.manual_seed(1111)
# Create tensor A/B/C
a_ref = torch.randint(
0, 2, (l, m, k // 2), dtype=torch.uint8, device="cuda"
).permute(1, 2, 0)
b_ref = torch.randint(
0, 2, (l, n, k // 2), dtype=torch.uint8, device="cuda"
).permute(1, 2, 0)
# a_ref = torch.ones((l, m, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0)
# b_ref = torch.ones((l, n, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0)
a_ref_f4 = a_ref.view(torch.float4_e2m1fn_x2)
b_ref_f4 = b_ref.view(torch.float4_e2m1fn_x2)
c_tensor = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute(
1, 2, 0
)
# Create a torch tensor for scale factor tensor of A and B
def create_ref_scale_factor_tensor(l, mn, sf_k):
"""
Create the reference scale factor tensor on CPU.
Returns the reshaped/pruned tensor ready for ref computation and its original permuted form.
"""
ref_shape = (l, mn, sf_k)
ref_permute_order = (1, 2, 0)
ref_f8_random_int = torch.randint(1, 3, ref_shape, dtype=torch.int8)
ref_f8_torch_tensor_cpu = ref_f8_random_int.to(dtype=torch.float8_e4m3fn)
# permute to match ref_permute_order
ref_f8_torch_tensor_cpu_permuted = ref_f8_torch_tensor_cpu.permute(
*ref_permute_order
)
return ref_f8_torch_tensor_cpu_permuted
# Copy the reference scale factor tensor to the CUTE-format scale factor tensor
def create_cute_scale_factor_tensor(l, mn, sf_k, ref_f8_torch_tensor_cpu_permuted):
"""
Create the CUTE-format scale factor tensor on CUDA based on the reference tensor.
"""
atom_m = (32, 4)
atom_k = 4
mma_shape = (
l, # batch size
ceil_div(mn, atom_m[0] * atom_m[1]),
ceil_div(sf_k, atom_k),
atom_m[0],
atom_m[1],
atom_k,
)
mma_permute_order = (3, 4, 1, 5, 2, 0)
# Generate a random int8 tensor, then convert to float8_e4m3fn
rand_int_tensor = torch.randint(0, 2, mma_shape, dtype=torch.int8)
cute_f8_torch_tensor_cpu = rand_int_tensor.to(dtype=torch.float8_e4m3fn)
# Permute according to mma_permute_order
cute_f8_torch_tensor_cpu = cute_f8_torch_tensor_cpu.permute(*mma_permute_order)
# Call the helper function to do layout conversion
cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
make_ptr(
cutlass.Float8E4M3FN,
ref_f8_torch_tensor_cpu_permuted.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=32,
),
make_ptr(
cutlass.Float8E4M3FN,
cute_f8_torch_tensor_cpu.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=32,
),
mn,
sf_k,
l,
mma_shape,
)
return cute_f8_torch_tensor_cpu.cuda()
sf_k = ceil_div(k, sf_vec_size)
sfa_ref = create_ref_scale_factor_tensor(l, m, sf_k)
sfb_ref = create_ref_scale_factor_tensor(l, n, sf_k)
# sfa_ref.fill_(1)
# sfb_ref.fill_(1)
sfa_tensor = create_cute_scale_factor_tensor(l, m, sf_k, sfa_ref)
sfb_tensor = create_cute_scale_factor_tensor(l, n, sf_k, sfb_ref)
# Configure gemm kernel
gemm = gemm_class()
# Initialize Stream
current_stream = cutlass_torch.default_stream()
a_ptr = make_ptr(
ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b_ptr = make_ptr(
ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfa_ptr = make_ptr(
sf_dtype, sfa_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfb_ptr = make_ptr(
sf_dtype, sfb_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
# Compile gemm kernel
compiled_gemm = cute.compile(
gemm,
a_ptr,
b_ptr,
sfa_ptr,
sfb_ptr,
c_ptr,
(m, n, k, l),
current_stream,
)
# Launch GPU kernel
compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream)
# For batch l, do (m, k, l) @ (n, k, l).T along k for each batch.
# Result: (m, n, l)
# Allocate ref as (l, m, n) with n-contiguous layout, then permute to (m, n, l)
ref = torch.empty(
(l, m, n),
dtype=torch.float16,
device="cuda",
).permute(1, 2, 0)
for l_idx in range(l):
# Convert the scale factor tensor to blocked format
scale_a = to_blocked(sfa_ref[:, :, l_idx])
scale_b = to_blocked(sfb_ref[:, :, l_idx])
# (m, k) @ (n, k).T -> (m, n)
res = torch._scaled_mm(
a_ref_f4[:, :, l_idx],
b_ref_f4[:, :, l_idx].transpose(0, 1),
scale_a.cuda(),
scale_b.cuda(),
bias=None,
out_dtype=torch.float16,
)
ref[:, :, l_idx] = res
torch.testing.assert_close(c_tensor, ref, atol=tolerance, rtol=1e-02)
if do_benchmark:
def generate_tensors():
a_ptr = make_ptr(
ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b_ptr = make_ptr(
ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfa_ptr = make_ptr(
sf_dtype,
sfa_tensor.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=32,
)
sfb_ptr = make_ptr(
sf_dtype,
sfb_tensor.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=32,
)
args = cute.testing.JitArguments(
a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream
)
args.add_to_scope([a_ref_f4, b_ref_f4, sfa_tensor, sfb_tensor, c_tensor])
return args
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a_ref_f4.numel() * a_ref_f4.element_size()
+ b_ref_f4.numel() * b_ref_f4.element_size()
+ sfa_tensor.numel() * sfa_tensor.element_size()
+ sfb_tensor.numel() * sfb_tensor.element_size()
+ c_tensor.numel() * c_tensor.element_size()
)
workspace_count = cute.testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
# Return execution time in microseconds
time = cute.testing.benchmark(
compiled_gemm,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
stream=current_stream,
warmup_iterations=warmup_iterations,
iterations=iterations,
)
print(f"Execution time: {time} us")
peta_flops = (4 * m * n * k * l) / (time * 1e-6) / 1e9 / 1000000
print(f"FLOPS: {peta_flops} PFLOPS")
bytes_transfer = (
2 * m * k / 2 * l * a_ref_f4.element_size()
+ 2 * n * k / 2 * l * b_ref_f4.element_size()
+ 2 * m * n * l * c_tensor.element_size()
+ 2 * m * sf_k * l * sfa_tensor.element_size()
+ 2 * n * sf_k * l * sfb_tensor.element_size()
)
print(f"Bytes: {bytes_transfer} Bytes")
bandwidth = bytes_transfer / time * 1e-3
print(f"BW: {bandwidth} GB/s")