[Hopper CuTeDSL] Add grouped GEMM persistent kernel and tests (#3091)
Implement grouped GEMM (C_g = A_g x B_g for g groups) on Hopper using CuTe DSL, extending the dense persistent GEMM with per-group TMA descriptor management. Kernel design (grouped_gemm.py): - Warp-specialized pipeline: DMA warp group handles TMA loads and per-group tensormap updates; MMA warp group runs WGMMA and stores C - StaticPersistentGroupTileScheduler for cross-group tile scheduling - Per-group TMA descriptor updates via GMEM or SMEM mode - Supports fp16, fp8 (E4M3FN/E5M2), int8 with mixed A/B dtypes - Configurable tile shapes (128x128, 128x256) and cluster shapes - Fix base TensorMapManager: hoist uniform_smem_ptrs outside predicated block to avoid illegal @P0 R2UR on sm_90a Tests (test/examples/CuTeDSL/hopper/test_grouped_gemm.py): - L0 compile and L1 correctness pytest suite covering tile shapes, dtypes, major modes, cluster shapes, group counts, and mixed sizes - Move to test/examples/CuTeDSL/hopper/ following sm_100a convention - Fix deprecated startdir arg in test_sharding.py pytest hook
This commit is contained in:
@@ -19,6 +19,8 @@ from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass import const_expr
|
||||
from cutlass.cute.core import AddressSpace as _CuteAddressSpace
|
||||
from cutlass.cute.core import make_ptr as _cute_make_ptr
|
||||
|
||||
|
||||
class TensorMapUpdateMode(Enum):
|
||||
@@ -138,11 +140,25 @@ class TensorMapManager:
|
||||
warp_idx = cute.arch.make_warp_uniform(
|
||||
cute.arch.warp_idx(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM):
|
||||
# Hoist SMEM pointer integer values into warp-uniform registers before
|
||||
# entering predicated blocks. This avoids predicated R2UR lowering on sm_90a.
|
||||
uniform_smem_ptrs = tuple(
|
||||
_cute_make_ptr(
|
||||
p.dtype,
|
||||
cute.arch.make_warp_uniform(p.toint(), loc=loc, ip=ip),
|
||||
mem_space=_CuteAddressSpace.smem,
|
||||
assumed_align=p.alignment,
|
||||
)
|
||||
for p in tensormap_smem_ptr
|
||||
)
|
||||
else:
|
||||
uniform_smem_ptrs = tensormap_smem_ptr
|
||||
# updates before touching tensormap in global memory
|
||||
if warp_idx == warp_id:
|
||||
if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM):
|
||||
for copy_atom, tensor, smem_ptr in zip(
|
||||
tma_copy_atom, tensor_gmem, tensormap_smem_ptr
|
||||
tma_copy_atom, tensor_gmem, uniform_smem_ptrs
|
||||
):
|
||||
cute.nvgpu.cpasync.update_tma_descriptor(
|
||||
copy_atom, tensor, smem_ptr, loc=loc, ip=ip
|
||||
@@ -154,7 +170,7 @@ class TensorMapManager:
|
||||
cute.arch.sync_warp(loc=loc, ip=ip)
|
||||
# updates to tensormap in global memory
|
||||
if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM):
|
||||
for gmem_ptr, smem_ptr in zip(tensormap_gmem_ptr, tensormap_smem_ptr):
|
||||
for gmem_ptr, smem_ptr in zip(tensormap_gmem_ptr, uniform_smem_ptrs):
|
||||
cute.nvgpu.cpasync.cp_fence_tma_desc_release(
|
||||
gmem_ptr, smem_ptr, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user