ITADN

[BUG] `cute.gemm` rejects valid SM90 WGMMA m64n8k32 in SS mode — degenerate partition layout `(1,1):(0,0)`

#3290Closed16bit-ykiko 创建于 2026-06-02
bug? - Needs TriageCuTe DSL
1
16bit-ykikocommented
### Which component has the problem? CuTe DSL ### Bug Report ## Describe the bug `cute.gemm` fails at compile time with `"invalid layout of A/B/D"` when using SS mode (both operands from shared memory) with `tiler_mn=(64, 8)` for FP8 WGMMA on SM90. The error occurs because the thread partition of the smem descriptor collapses to `(1,1):(0,0)` when N=8 — exactly one atom wide. The hardware supports `wgmma.mma_async.sync.aligned.m64n8k32.f32.e4m3.e4m3` (confirmed in PTX ISA and CUTLASS's own `SM90_WGMMA_SHAPES_FP8_DENSE` which lists `(64, 8, 32)` as Level 2). The equivalent C++ CuTe code works correctly with N=8 in SS mode. **Error message:** ``` 'cute.gemm' op invalid layout of A/B/D. A: (1,1):(0,0), B: (1,1):(0,0), D:((2,2),1,1):((1,2),0,0) ``` N=16 and above compile successfully. ## Steps/Code to reproduce bug ```python import torch import cutlass, cutlass.cute as cute, cutlass.utils as utils import cutlass.utils.hopper_helpers as sm90_utils import cutlass.cute.nvgpu as nvgpu, cutlass.cute.nvgpu.warpgroup as warpgroup from cutlass.cute.runtime import from_dlpack class Repro: def __init__(self, N): self.N = N @cute.kernel def kernel(self, mma: cute.TiledMma, sA_l, sB_l): tidx, _, _ = cute.arch.thread_idx() smem = utils.SmemAllocator() st = smem.allocate(self._ss) sA = st.sA.get_tensor(sA_l.outer, swizzle=sA_l.inner)[(None, None, 0)] sB = st.sB.get_tensor(sB_l.outer, swizzle=sB_l.inner)[(None, None, 0)] t = mma.get_slice(tidx) rA = t.make_fragment_A(t.partition_A(sA)) rB = t.make_fragment_B(t.partition_B(sB)) acc = t.make_fragment_C(t.partition_shape_C((64, self.N))) warpgroup.fence() for k in range(cute.size(rA, mode=[2]), unroll_full=True): mma.set(warpgroup.Field.ACCUMULATE, k != 0) cute.gemm(mma, acc, rA[None, None, k], rB[None, None, k], acc) warpgroup.commit_group() warpgroup.wait_group(0) @cute.jit def __call__(self, dummy: cute.Tensor): mma = sm90_utils.make_trivial_tiled_mma( a_dtype=cute.Float8E4M3FN, b_dtype=cute.Float8E4M3FN, a_leading_mode=nvgpu.OperandMajorMode.K, b_leading_mode=nvgpu.OperandMajorMode.K, acc_dtype=cute.Float32, atom_layout_mnk=(1, 1, 1), tiler_mn=(64, self.N), a_source=warpgroup.OperandSource.SMEM, ) tile = (64, self.N, 128) lay_a = sm90_utils.make_smem_layout_a(utils.LayoutEnum.ROW_MAJOR, tile, cute.Float8E4M3FN, 1) lay_b = sm90_utils.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, tile, cute.Float8E4M3FN, 1) @cute.struct class SS: sA: cute.struct.Align[cute.struct.MemRange[cute.Float8E4M3FN, cute.cosize(lay_a)], 128] sB: cute.struct.Align[cute.struct.MemRange[cute.Float8E4M3FN, cute.cosize(lay_b)], 128] self._ss = SS self.kernel(mma, lay_a, lay_b).launch(grid=(1,1,1), block=(128,1,1), smem=SS.size_in_bytes()) dummy = from_dlpack(torch.zeros(16, dtype=torch.int8, device="cuda"), assumed_align=16) # N=8: FAILS cute.compile(Repro(8), dummy) # N=16, 32, 64, 128: works fine # cute.compile(Repro(16), dummy) ``` **Output:** ``` cute.gemm SS mode, M=64 K=128 FP8: N= 8: FAIL — 'cute.gemm' op invalid layout of A/B/D. A: (1,1):(0,0), B: (1,1):(0,0), D:((2,2),1,1):((1,2),0,0) N= 16: OK N= 32: OK N= 64: OK N=128: OK ``` ## Expected behavior `cute.gemm` should compile successfully for N=8 in SS mode. ## Environment details - **Environment**: Bare-metal, NVIDIA H20 (SM90) - **CUDA**: 12.8 - **Python**: 3.13 - **nvidia-cutlass-dsl**: 4.5.2 (pip) - **OS**: Linux x86-64
关闭于 2026-07-02 2 条评论