[BUG][CuTeDSL] Blackwell dense GEMM returns wrong results silently for tf32 + 2-CTA + MN-major B with MMA tiler N 32 / 96 / 160 / 224: get_smem_layout_atom_ab picks MN_SW64 for a 32-bit MN-major operand

Open
  • Is something wrong?
#3,792 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 4 days

@yunweili3 is already working on this.

Since Oct 27, 2026.

  • #3793 by @yunweili3 — open

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
25/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Stale
Tech stack
cpp, python
Domain
machine-learning, performance

Research direction

Start by reading cutlass.utils.blackwell_helpers.get_smem_layout_atom_ab and the sm100_smem_selector static assertion in include/cutlass/gemm/collective/builders/sm100_common.inl. Reproduce with the provided dense GEMM command and compare the listed tilers and operand types. Done means unsupported 32-bit MN-major layouts fail clearly, or the reported configurations produce correct results; linked pull request #3793 is already open.

Written by the indexing model from the issue text.

Description

Which component has the bug?

CuTe DSL

Bug Report

Describe the bug

On SM100 the Blackwell dense GEMM examples produce wrong results, silently, for a tf32 (Float32 input, TFloat32 MMA) MN-major B operand with use_2cta_instrs and an MMA tiler N of 32, 96, 160 or 224. can_implement() accepts the configuration, the kernel compiles and runs, and only the example's reference check notices (98.8% of the output mismatches). A K-major B, 1-CTA MMA, bf16/f16/fp8, or N 64 / 128 / 192 are all fine.

Root cause: cutlass.utils.blackwell_helpers.get_smem_layout_atom_ab returns MN_SW128_32B for a 32-bit MN-major operand only when the major extent is a multiple of 32 elements, and otherwise falls through to MN_SW64 / MN_SW32 / MN_INTER like it does for 16- and 8-bit types. For a 32-bit MN-major operand those atoms are not consumable by UMMA; the C++ builder says so explicitly (sm100_smem_selector in include/cutlass/gemm/collective/builders/sm100_common.inl: static_assert(..., "for mn-major tf32 operands, SW128_32B is the only available smem layout")). With 2 CTAs the B tile per CTA has N/2 rows, so N ∈ {32, 96, 160, 224} gives 16 / 48 / 80 / 112 rows → MN_SW64 → garbage.

Steps/Code to reproduce bug

B200, CUDA 13, nvidia-cutlass-dsl 4.8.0 and 4.8.1, main 91843b71, examples/python/CuTeDSL:

python cute/blackwell/kernel/dense_gemm/dense_gemm.py --ab_dtype TFloat32 --c_dtype Float16 --acc_dtype Float32 \
  --mma_tiler_mn 256,96 --cluster_shape_mn 2,1 --mnkl 1024,960,256,2 --use_tma_store --use_2cta_instrs --a_major k --b_major n

AssertionError: Tensor-likes are not close!
Mismatched elements: 1941548 / 1966080 (98.8%)
Greatest absolute difference: 151.0 at index (0, 546, 546) (up to 0.1 allowed)
Greatest relative difference: 1.0 at index (0, 0, 0) (up to 0.001 allowed)

Same command with other tilers / types (all --b_major n, 2-CTA unless noted):

--mma_tiler_mn B smem tile per CTA (n, k) atom picked result
256,32 and 128,32 (16, 32) MN_SW64 wrong
256,64 (32, 32) MN_SW128_32B pass
256,96 and 128,96 (48, 32) MN_SW64 wrong
256,160 and 128,160 (80, 32) MN_SW64 wrong
256,192 (96, 32) MN_SW128_32B pass
256,224 (112, 32) MN_SW64 wrong
1-CTA 128,32 / 128,64 / 128,96 (32..96, 32) MN_SW128_32B pass
bf16 256,64 / 256,96 (32, 64) / (48, 64) MN_SW64 / MN_SW32 pass
fp8 e4m3 256,96 (48, 128) MN_SW32 pass

Independent of --use_tma_store and of A's majorness; dense_gemm_persistent_dynamic.py behaves the same.

Expected behavior

Either a correct result or a clear error. Since UMMA has no other MN-major layout for 32-bit operands, I think the helper should mirror the C++ static_assert and raise when the major extent is not a multiple of 32 elements, so that the examples (and user kernels built on the helper) fail at compile time instead of returning garbage. With that change (PR to follow) the command above fails with

ValueError: MN-major TFloat32 operands need a major mode extent that is a multiple of 32 elements (SW128_32B is the only available SMEM layout atom), got 48

and the passing configurations above are unchanged. The dense GEMM examples' can_implement() could additionally reject tf32 + 2-CTA + MN-major B with N % 64 != 0 up front, if you prefer a False over an exception there.

Environment details

B200 (sm_100a), CUDA 13.1, Python 3.10, nvidia-cutlass-dsl 4.8.0 / 4.8.1, torch 2.14.0+cu130, main at 91843b71.

Dominant language
C++
Stars
10.7k
Forks
2.2k
Avg merge
3d 20h
Merged PRs (30d)
4

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

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.

More from NVIDIA/cutlass

All issues in NVIDIA/cutlass

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.