NVIDIA / NVIDIA/cutlass

[BUG] Cutlass and Cute-DSL generate suboptimal code for UMMA that use more registers than necessary

Open
#2,408 24 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug CuTe DSL
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Describe the bug
The SASS code for UMMA uses more registers than necessary: the registers holding tmem address and idesc keep changing between instructions.
This is important as it affects Blackwell FMHA perf, e.g. causing a lot of register spill in Blackwell FMHA (example 77 and cute-dsl blackwell's fmha.py).

SASS code from example 71 (71_blackwell_gemm_with_collective_builder), we see the tmem (4th argument to UTCHMMA) and idesc registers keep changing between 4 instructions:

UTCHMMA gdesc[UR38], gdesc[UR40], tmem[UR17], tmem[UR26], idesc[UR27], UP2 ; 
UPLOP3.LUT UP2, UPT, UPT, UPT, UPT, 0x80, 0x8 ;   
UMOV UR35, 0x40004040 ;                                                            
UMOV UR33, 0x40004040 ;                                                            
UTCHMMA gdesc[UR34], gdesc[UR36], tmem[UR17], tmem[UR24], idesc[UR25], UPT ;       
UMOV UR31, 0x40004040 ;                                                            
UMOV UR29, 0x40004040 ;                                                            
UMOV UR43, 0x40004040 ;                                                            
UTCHMMA gdesc[UR30], gdesc[UR32], tmem[UR17], tmem[UR22], idesc[UR23], UPT ;      
UTCHMMA gdesc[UR42], gdesc[UR28], tmem[UR17], tmem[UR20], idesc[UR21], UPT ;      

SASS code from cute-dsl blackwell dense_gemm_persistent.py, similar:

Image

Steps/Code to reproduce bug
Compile example 71 or cute-dsl blackwell dense_gemm_persistent.py and view the SASS

Expected behavior
The register holding tmem address and idesc should stay the same between instructions.
This is the SASS from a gemm kernel in cublas (nvjet), where the tmem address (4th argument to UTCHMMA) and idesc use the same registers throughout instructions:
Image

Environment details (please complete the following information):
Newest Pytorch docker image: nvcr.io/nvidia/pytorch:25.05-py3
nvcc release 12.9, V12.9.41, Build cuda_12.9.r12.9/compiler.35813241_0

cc: @thakkarV

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

Compile example 71 (71_blackwell_gemm_with_collective_builder) and cute-dsl blackwell's dense_gemm_persistent.py, then inspect the generated SASS for UMMA instructions. Compare the tmem address and idesc registers across instructions; done means they remain stable and the excessive register use or spill is reduced.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.