[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:
Johnsonms
2026-03-17 21:40:15 -07:00
committed by GitHub
parent 1b741cabaa
commit 982748aa73
5 changed files with 3024 additions and 3 deletions

View File

@@ -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
)