[BUG] `cute.gemm` rejects valid SM90 WGMMA m64n8k32 in SS mode — degenerate partition layout `(1,1):(0,0)`
bug? - Needs TriageCuTe DSL
### 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 条评论