[BUG] Wrong TV Layout C in MMA Atom
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
Which component has the problem?
CuTe DSL
Bug Report
Describe the bug
When N equals 8, the SM90 WGMMA Atom has an incorrect TV Layout C.
Steps/Code to reproduce bug
import cutlass
import cutlass.cute as cute
import cutlass.utils.hopper_helpers as sm90_utils
@cute.jit
def make_mma(N: cutlass.Constexpr):
tiled_mma = sm90_utils.make_trivial_tiled_mma(
cutlass.Float16,
cutlass.Float16,
cute.nvgpu.OperandMajorMode.K,
cute.nvgpu.OperandMajorMode.K,
cutlass.Float32,
(1, 1, 1),
tiler_mn=(64, N),
)
print(tiled_mma)
if __name__ == '__main__':
make_mma(8)
make_mma(16)
make_mma(32)
make_mma(64)
Expected behavior
Wrong Result:
Tiled MMA
Thr Layout VMNK: (128,1,1,1):(1,0,0,0)
Permutation MNK: (_,_,_)
MMA Atom
ThrID: 128:1
Shape MNK: (64,8,16)
TV Layout A: (128,(64,16)):(0,(1,64))
TV Layout B: (128,(8,16)):(0,(1,8))
TV Layout C: ((4,8,4),(2,2)):((128,1,16),(64,8))
Right Result:
Tiled MMA
Thr Layout VMNK: (128,1,1,1):(1,0,0,0)
Permutation MNK: (_,_,_)
MMA Atom
ThrID: 128:1
Shape MNK: (64,8,16)
TV Layout A: (128,(64,16)):(0,(1,64))
TV Layout B: (128,(8,16)):(0,(1,8))
TV Layout C: ((4,8,4),(2,2,1)):((128,1,16),(64,8,512))
Reference Code: https://github.com/NVIDIA/cutlass/blob/982cb9e718bcd4d7ac546b1795702a08326dfe4b/include/cute/atom/mma_traits_sm90_gmma.hpp#L432-L435
Environment details (please complete the following information):
- nvidia-cutlass-dsl 4.5.1
- nvidia-cutlass-dsl-libs-base 4.5.1
- nvidia-cutlass-dsl-libs-cu13 4.5.1
Additional context
None.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by running the provided CuTe DSL reproduction for N=8, 16, 32, and 64. Read include/cute/atom/mma_traits_sm90_gmma.hpp at the referenced lines and compare the generated TV Layout C with the expected output. Done means the N=8 layout matches the right result without regressing the other cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- hpc
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 62/100