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

v4.4 release update v2. (#2999)

Browse Source
This commit is contained in:
Junkai-Wu
2026-02-03 20:48:31 -05:00
committed by GitHub
parent 1cfbb53a23
commit 6b3e607b85
91 changed files with 13242 additions and 1488 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py
+183 -101
View File
@@ -34,8 +34,8 @@ import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.runtime import from_dlpack
import cutlass.utils as utils
from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
@@ -1046,45 +1046,70 @@ class PersistentDenseGemmKernel:
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
#
# Persistent tile scheduling loop for epilogue
#
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
if cutlass.const_expr(self.use_tma_store):
assert tma_atom_c is not None and sC is not None
utils.gemm.sm100.epilogue_tma_store(
self,
tidx,
warp_idx,
acc_pipeline,
tiled_mma,
tma_atom_c,
tCtAcc_base,
sC,
tCgC,
epi_tile,
tile_sched,
epilogue_op,
clc_pipeline,
clc_consumer_state,
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilogue_warp_id),
)
else:
utils.gemm.sm100.epilogue(
self,
tidx,
acc_pipeline,
tiled_mma,
tCtAcc_base,
tCgC,
epi_tile,
tile_sched,
epilogue_op,
tmem_dealloc_barrier,
None,
None,
clc_pipeline,
clc_consumer_state,
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage, producer_group=c_producer_group
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
cur_tile_coord[1],
cur_tile_coord[2],
)
num_tiles_executed = tile_sched.num_tiles_executed
if cutlass.const_expr(self.use_tma_store):
acc_consumer_state = utils.gemm.sm100.epilogue_tma_store(
self,
tidx,
warp_idx,
tma_atom_c,
tCtAcc_base,
sC,
tCgC,
epi_tile,
num_tiles_executed,
epilogue_op,
mma_tile_coord_mnl,
acc_consumer_state,
acc_pipeline,
c_pipeline,
)
else:
acc_consumer_state = utils.gemm.sm100.epilogue(
self,
tidx,
tCtAcc_base,
tCgC,
epi_tile,
epilogue_op,
mma_tile_coord_mnl,
acc_consumer_state,
acc_pipeline,
)
#
# Advance to next tile
#
clc_pipeline.consumer_wait(clc_consumer_state)
work_tile = tile_sched.get_current_work()
clc_pipeline.consumer_release(clc_consumer_state)
clc_consumer_state.advance()
if cutlass.const_expr(self.use_tma_store):
# Wait for C store complete
c_pipeline.producer_tail()
else:
# Synchronize before TMEM dealloc (done by the caller)
tmem_dealloc_barrier.arrive_and_wait()
#
# Dealloc the tensor memory buffer
#
@@ -1150,8 +1175,11 @@ class PersistentDenseGemmKernel:
return num_tmem_alloc_cols
def check_supported_dtypes(
self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric]
) -> bool:
self,
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
):
"""
Check if the dtypes are valid
@@ -1173,8 +1201,10 @@ class PersistentDenseGemmKernel:
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
}
if ab_dtype not in valid_ab_dtypes:
raise testing.CantImplementError(f"Unsupported AB dtype: {ab_dtype}")
if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes:
raise testing.CantImplementError(
f"Unsupported AB dtype: {a_dtype} and {b_dtype}"
)
if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}:
raise testing.CantImplementError(
@@ -1198,8 +1228,13 @@ class PersistentDenseGemmKernel:
cutlass.Int32: {cutlass.Uint8, cutlass.Int8},
}
# Check compatibility between accumulator type and AB type
if ab_dtype not in acc_ab_compatibility[self.acc_dtype]:
return False
if (
a_dtype not in acc_ab_compatibility[self.acc_dtype]
or b_dtype not in acc_ab_compatibility[self.acc_dtype]
):
raise testing.CantImplementError(
f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}"
)
# Define compatibility mapping between accumulator type and C type
acc_c_compatibility = {
@@ -1228,11 +1263,11 @@ class PersistentDenseGemmKernel:
}
# Check compatibility between accumulator type and C type
if c_dtype not in acc_c_compatibility[self.acc_dtype]:
return False
raise testing.CantImplementError(
f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}"
)
return True
def check_mma_tiler_and_cluster_shape(self) -> bool:
def check_mma_tiler_and_cluster_shape(self):
"""Check if the mma tiler and cluster shape are valid.
:raises testing.CantImplementError: If the mma tiler and cluster shape are invalid
@@ -1273,12 +1308,13 @@ class PersistentDenseGemmKernel:
n: int,
k: int,
l: int,
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
) -> bool:
):
"""
Check if the tensor alignment is valid
@@ -1290,8 +1326,10 @@ class PersistentDenseGemmKernel:
:type k: int
:param l: The number of columns in the C tensor
:type l: int
:param ab_dtype: The data type of the A and B operands
:type ab_dtype: Type[cutlass.Numeric]
:param a_dtype: The data type of the A operand
:type a_dtype: Type[cutlass.Numeric]
:param b_dtype: The data type of the B operand
:type b_dtype: Type[cutlass.Numeric]
:param c_dtype: The data type of the output tensor
:type c_dtype: Type[cutlass.Numeric]
:param a_major: The major axis of the A tensor
@@ -1301,8 +1339,7 @@ class PersistentDenseGemmKernel:
:param c_major: The major axis of the C tensor
:type c_major: str
:return: True if the problem shape is valid, False otherwise
:rtype: bool
:raises testing.CantImplementError: If the tensor alignment is invalid
"""
# TODO: move to utils
@@ -1313,15 +1350,15 @@ class PersistentDenseGemmKernel:
return num_major_elements % num_contiguous_elements == 0
if (
not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l))
or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l))
or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l))
):
raise testing.CantImplementError(
f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {ab_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}"
f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}"
)
def check_epilog_store_option(self, m: int, n: int) -> bool:
def check_epilog_store_option(self, m: int, n: int):
"""
Check if the epilogue store option is valid
@@ -1346,7 +1383,8 @@ class PersistentDenseGemmKernel:
def can_implement(
self,
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
@@ -1357,8 +1395,10 @@ class PersistentDenseGemmKernel:
:param mnkl: Problem size as a tuple (M, N, K, L).
:type mnkl: Tuple[int, int, int, int]
:param ab_dtype: Data type for input tensors A and B.
:type ab_dtype: Type[cutlass.Numeric]
:param a_dtype: Data type for input tensors A.
:type a_dtype: Type[cutlass.Numeric]
:param b_dtype: Data type for input tensors B.
:type b_dtype: Type[cutlass.Numeric]
:param c_dtype: Data type for output tensor C.
:type c_dtype: Type[cutlass.Numeric]
:param a_major: Major dimension of the A tensor layout ("m" or "k").
@@ -1373,14 +1413,14 @@ class PersistentDenseGemmKernel:
try:
# Skip unsupported types
self.check_supported_dtypes(ab_dtype, c_dtype)
self.check_supported_dtypes(a_dtype, b_dtype, c_dtype)
# Skip invalid mma tile shape and cluster shape
self.check_mma_tiler_and_cluster_shape()
m, n, k, l = mnkl
self.check_tensor_alignment(
m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major
)
self.check_epilog_store_option(m, n)
except testing.CantImplementError:
@@ -1432,43 +1472,72 @@ def bmm(
@lru_cache(maxsize=1)
def prepare_tensors(
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
init_random: bool = True,
normal_mean: float = 0.0,
normal_std: float = 1.0,
):
"""Prepare tensors for GEMM.
Returns:
Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage):
- *_f32: Float32 tensors with the logical data (for reference and fp8 conversion)
- *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype)
"""
import torch
from cutlass.torch import dtype as torch_dtype
m, n, k, l = mnkl
if a_major == "k":
a = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
elif a_major == "m":
a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if b_major == "n":
b = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
elif b_major == "k":
b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1)
b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if c_major == "n":
c = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
elif c_major == "m":
c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if init_random:
a.random_(-2, 3)
b.random_(-2, 3)
c.random_(-2, 3)
# Uniform random initialization in range [-2, 3)
a_f32.random_(-2, 3)
b_f32.random_(-2, 3)
c_f32.random_(-2, 3)
else:
# Normal (Gaussian) initialization with user-specified mean and std
a_f32.normal_(mean=normal_mean, std=normal_std)
b_f32.normal_(mean=normal_mean, std=normal_std)
c_f32.normal_(mean=normal_mean, std=normal_std)
return (
a.to(dtype=torch_dtype(ab_dtype)),
b.to(dtype=torch_dtype(ab_dtype)),
c.to(dtype=torch_dtype(c_dtype)),
)
# For float8 types, use uint8 as storage type to avoid dlpack limitation
# (dlpack doesn't support float8 types)
# For other types, convert to the target dtype
a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype)
b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype)
c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype)
a_storage = a_f32.to(dtype=a_storage_dtype)
b_storage = b_f32.to(dtype=b_storage_dtype)
c_storage = c_f32.to(dtype=c_storage_dtype)
return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage)
@lru_cache(maxsize=1)
@@ -1499,7 +1568,7 @@ def compile_bmm(
)
# Check if configuration can be implemented
can_implement = gemm.can_implement(
mnkl, a.element_type, c.element_type, a_major, b_major, c_major
mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major
)
if not can_implement:
raise testing.CantImplementError(
@@ -1594,21 +1663,30 @@ def run(
)
# Run and verify BMM with torch
a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major)
a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors(
mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major
)
leading_dim_a = 2 if a_major == "k" else 1
leading_dim_b = 1 if b_major == "k" else 2
leading_dim_c = 2 if c_major == "n" else 1
a_ = from_dlpack(
a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(
b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(
c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_c)
# Create CuTe tensors, passing float32 source for fp8 conversion
a_ = create_cute_tensor_for_fp8(
a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32
)
b_ = create_cute_tensor_for_fp8(
b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32
)
c_ = create_cute_tensor_for_fp8(
c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32
)
print("Compile Blackwell Persistent Dense GEMM with:")
print(f"ab_dtype: {ab_dtype}, c_dtype: {c_dtype}, acc_dtype: {acc_dtype}")
print(f"a_major: {a_major}, b_major: {b_major}, c_major: {c_major}")
print(f"mma_tiler_mn: {mma_tiler_mn}, cluster_shape_mn: {cluster_shape_mn}")
print(f"use_2cta_instrs: {use_2cta_instrs}, use_tma_store: {use_tma_store}")
compiled_fn = compile_bmm(
mnkl,
@@ -1640,13 +1718,15 @@ def run(
compiled_fn(a_, b_, c_, current_stream)
# Manually quantize to be comparable
# Use float32 source data for reference calculation
ref = (
torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32))
torch.bmm(a_f32, b_f32)
.to(dtype=torch_dtype(c_dtype))
.to(dtype=torch.float32)
)
# Read back the result from CuTe tensor (c_storage was updated in-place)
torch.testing.assert_close(
c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
)
if not benchmark:
@@ -1654,32 +1734,33 @@ def run(
def generate_tensors():
init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8]
a, b, c = prepare_tensors(
a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors(
mnkl,
ab_dtype,
ab_dtype,
c_dtype,
a_major,
b_major,
c_major,
init_random=not init_normal,
)
a_ = from_dlpack(
a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(
b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(
c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_c)
a_ = create_cute_tensor_for_fp8(
a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32
)
b_ = create_cute_tensor_for_fp8(
b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32
)
c_ = create_cute_tensor_for_fp8(
c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32
)
return testing.JitArguments(a_, b_, c_, current_stream)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a.numel() * a.element_size()
+ b.numel() * b.element_size()
+ c.numel() * c.element_size()
a_storage.numel() * a_storage.element_size()
+ b_storage.numel() * b_storage.element_size()
+ c_storage.numel() * c_storage.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
@@ -1735,6 +1816,7 @@ def prepare_parser():
action="store_true",
help="Enable 2CTA MMA instructions feature",
)
parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k")
parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k")
parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n")
examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py
+2977
View File
File diff suppressed because it is too large Load Diff