NVIDIA / NVIDIA/cutlass

[BUG] Wrong TV Layout C in MMA Atom

Open
#3,254 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug CuTe DSL inactive-30d inactive-90d
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.