[BUG] Cutlass and Cute-DSL generate suboptimal code for UMMA that use more registers than necessary
Nobody has claimed this yet.
- 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:
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:
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
- 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
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