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

46
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
View File

@@ -30,13 +30,11 @@ import argparse
from typing import Type, Tuple, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
@@ -64,7 +62,7 @@ This GEMM kernel supports the following features:
This GEMM works as follows:
1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations.
2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
3. EPILOGUE warp:
- Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
@@ -1008,7 +1006,10 @@ class BlockwiseGemmKernel:
)
# fence view async shared
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.sched_sync_barrier.arrive_and_wait()
# commit tile info pipeline
tile_info_pipeline.producer_commit(tile_info_producer_state)
@@ -1123,7 +1124,10 @@ class BlockwiseGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1295,7 +1299,10 @@ class BlockwiseGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1450,7 +1457,10 @@ class BlockwiseGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1684,7 +1694,10 @@ class BlockwiseGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1851,7 +1864,10 @@ class BlockwiseGemmKernel:
tRS_sC[(None, None, None, c_buffer)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.epilog_sync_barrier.arrive_and_wait()
#
@@ -1881,7 +1897,10 @@ class BlockwiseGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -2553,6 +2572,9 @@ class BlockwiseGemmKernel:
def create_tensors(
l, m, n, k, a_major, b_major, cd_major, ab_dtype, c_dtype, scale_dtype
):
import torch
import cutlass.torch as cutlass_torch
torch.manual_seed(1111)
a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
@@ -2613,6 +2635,9 @@ def run(
use_cold_l2: bool = False,
**kwargs,
):
import torch
import cutlass.torch as cutlass_torch
"""
Prepare A/B/C tensors, launch GPU kernel, and reference checking.
"""
@@ -2688,6 +2713,7 @@ def run(
# try to check CUDA version to decide the opt level
try:
from cutlass import CUDA_VERSION
opt_level = (
3
if CUDA_VERSION.major < 13

47
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
View File

@@ -30,13 +30,11 @@ import argparse
from typing import Type, Tuple, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
@@ -80,7 +78,7 @@ This GEMM kernel supports the following features:
This GEMM works as follows:
1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations.
2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
3. EPILOGUE warp:
- Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
@@ -1034,7 +1032,10 @@ class BlockwiseContiguousGroupedGemmKernel:
)
# fence view async shared
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.sched_sync_barrier.arrive_and_wait()
# commit tile info pipeline
tile_info_pipeline.producer_commit(tile_info_producer_state)
@@ -1150,7 +1151,10 @@ class BlockwiseContiguousGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1322,7 +1326,10 @@ class BlockwiseContiguousGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1479,7 +1486,10 @@ class BlockwiseContiguousGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1715,7 +1725,10 @@ class BlockwiseContiguousGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1884,7 +1897,10 @@ class BlockwiseContiguousGroupedGemmKernel:
tRS_sC[(None, None, None, c_buffer)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.epilog_sync_barrier.arrive_and_wait()
#
@@ -1914,7 +1930,10 @@ class BlockwiseContiguousGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -2595,6 +2614,8 @@ class BlockwiseContiguousGroupedGemmKernel:
def create_mask(num_groups, expect_m, fixed_m=False, m_aligned=128):
import torch
valid_m = 0
group_m_list = []
gidx_mapping = []
@@ -2632,6 +2653,9 @@ def create_tensors(
scale_dtype,
fixed_m=False,
):
import torch
import cutlass.torch as cutlass_torch
torch.manual_seed(1111)
valid_m, group_m_list, _gidx_mapping = create_mask(l, m, fixed_m)
@@ -2702,6 +2726,9 @@ def run(
fixed_m: bool = False,
**kwargs,
):
import torch
import cutlass.torch as cutlass_torch
"""
Prepare A/B/C tensors, launch GPU kernel, and reference checking.
"""

47
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
View File

@@ -30,13 +30,11 @@ import argparse
from typing import Type, Tuple, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
@@ -79,7 +77,7 @@ Matrix A/C Memory Layout Diagrams:
This GEMM works as follows:
1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations.
2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations.
2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction.
3. EPILOGUE warp:
- Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld.
@@ -1041,7 +1039,10 @@ class BlockwiseMaskedGroupedGemmKernel:
)
# fence view async shared
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.sched_sync_barrier.arrive_and_wait()
# commit tile info pipeline
tile_info_pipeline.producer_commit(tile_info_producer_state)
@@ -1156,7 +1157,10 @@ class BlockwiseMaskedGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1328,7 +1332,10 @@ class BlockwiseMaskedGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1483,7 +1490,10 @@ class BlockwiseMaskedGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1717,7 +1727,10 @@ class BlockwiseMaskedGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -1884,7 +1897,10 @@ class BlockwiseMaskedGroupedGemmKernel:
tRS_sC[(None, None, None, c_buffer)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
self.epilog_sync_barrier.arrive_and_wait()
#
@@ -1914,7 +1930,10 @@ class BlockwiseMaskedGroupedGemmKernel:
for idx in cutlass.range(4, unroll_full=True):
tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)]
is_valid_tile = tile_info[3] == 1
cute.arch.fence_proxy("async.shared", space="cta")
cute.arch.fence_proxy(
"async.shared",
space="cta",
)
tile_info_pipeline.consumer_release(tile_info_consumer_state)
tile_info_consumer_state.advance()
@@ -2586,6 +2605,8 @@ class BlockwiseMaskedGroupedGemmKernel:
def create_mask(num_groups: int, m: int, fixed_m=False, tile_m=128):
import torch
# align with block_m (or block_n if swapAB)
masked_m_candidates = list(
filter(
@@ -2617,6 +2638,9 @@ def create_tensors(
scale_dtype,
fixed_m=False,
):
import torch
import cutlass.torch as cutlass_torch
torch.manual_seed(1111)
_gidx_mapping, masked_m = create_mask(l, m, fixed_m)
@@ -2684,6 +2708,9 @@ def run(
fixed_m: bool = False,
**kwargs,
):
import torch
import cutlass.torch as cutlass_torch
"""
Prepare A/B/C tensors, launch GPU kernel, and reference checking.
"""