[BUG] SM120 NVF4 GEMM (example 79a): misaligned address crash

[m96-chan]

Which component has the problem?

CUTLASS C++

Bug Report

Summary

Example 79a (79a_blackwell_geforce_nvfp4_bf16_gemm.cu) crashes with "misaligned address" on RTX 5090 (SM120). This is related to #2905 but affects a different mainloop (sm120_blockscaled_mma_tma.hpp for NVF4 instead of sm120_mma_tma_blockwise_scaling.hpp for FP8).

Environment

• GPU: NVIDIA GeForce RTX 5090 (SM 12.0)
• CUDA Toolkit: 13.1
• CUTLASS: v4.3.4
• OS: Windows 11

Steps to Reproduce

1. Build example 79a:

cd examples/79_blackwell_geforce_gemm
nvcc -O2 -arch=sm_120a -std=c++17 --expt-relaxed-constexpr \
  -I ../../include -I ../../tools/util/include \
  -DCUTLASS_ARCH_MMA_SM120_SUPPORTED=1 \
  79a_blackwell_geforce_nvfp4_bf16_gemm.cu -o 79a_test

2. Run:

./79a_test --m=256 --n=256 --k=256

3. Result:

========= COMPUTE-SANITIZER
========= Misaligned shared or local address
=========     at ...MainloopSm120TmaWarpSpecializedBlockScaled...

Root Cause Analysis

Two alignment issues:

1. TMA descriptor alignment: prefetch.tensormap requires 64-byte alignment. Params structs lack alignas(64).
2. Scale factor smem alignment: ldmatrix.sync.aligned requires 16-byte alignment. Comparing 79a vs 87a mainloops:

// sm120_mma_tma_blockwise_scaling.hpp (87a) - uses array_aligned (16-byte default)
cute::array_aligned<ElementSF, ...> smem_scale_A;

// sm120_blockscaled_mma_tma.hpp (79a) - no alignment
cute::ArrayEngine<ElementSF, ...> smem_SFA;  // Missing alignment
cute::ArrayEngine<ElementSF, ...> smem_SFB;

Expected Behavior

Disposition: Passed

Affected Files

1. include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp
2. include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp

Appendix: Proposed Fix

sm120_blockscaled_mma_tma.hpp:

 struct TensorStorage : cute::aligned_struct<128, _0> {
   alignas(1024) cute::ArrayEngine<SmemAllocTypeA, ...> smem_A;
   alignas(1024) cute::ArrayEngine<SmemAllocTypeB, ...> smem_B;
-  cute::ArrayEngine<ElementSF, ...> smem_SFA;
-  cute::ArrayEngine<ElementSF, ...> smem_SFB;
+  alignas(128) cute::ArrayEngine<ElementSF, ...> smem_SFA;
+  alignas(128) cute::ArrayEngine<ElementSF, ...> smem_SFB;
 };

-struct Params {
+struct alignas(64) Params {
-  TMA_A tma_load_a;
-  TMA_B tma_load_b;
-  TMA_SFA tma_load_sfa;
-  TMA_SFB tma_load_sfb;
+  alignas(64) TMA_A tma_load_a;
+  alignas(64) TMA_B tma_load_b;
+  alignas(64) TMA_SFA tma_load_sfa;
+  alignas(64) TMA_SFB tma_load_sfb;
   // ...
 };

sm90_epilogue_tma_warpspecialized.hpp:

-struct Params {
+struct alignas(64) Params {
   // ...
-  TMA_C tma_load_c;
-  TMA_D tma_store_d;
+  alignas(64) TMA_C tma_load_c;
+  alignas(64) TMA_D tma_store_d;
 };

After fix: Passed, 3043 GFLOPS

Notes on Pull Request

The fix touches sm90_epilogue_tma_warpspecialized.hpp which is shared across SM90/SM100/SM120. I only have SM120 hardware, so I cannot verify impact on other architectures. Happy to submit a PR if the team would like.

Related Issues

• [BUG] SM120 blockwise FP8 GEMM: TMA descriptor misalignment in Params struct #2905 (FP8 blockwise GEMM - same TMA alignment issue, different mainloop)
• [QST][CUTEDSL] Address Misalignment in FP8 Gemm  #2902 (partition_S alignment drop discussion)

[hwu36]

@lsyyy666

[m96-chan]

I've created a workaround for SM120 (RTX 5090) misaligned address issues.

Fork with fix: https://github.com/m96-chan/cutlass/tree/fix/sm120-alignment
Commit: m96-chan@65e7e401

Root causes identified:

1. TMA descriptors require 64-byte alignment for prefetch.tensormap PTX
2. ldmatrix.sync.aligned requires 16-byte alignment for shared memory
3. partition_S() can drop alignment from 1024 to 8 bytes

Fixes applied:

• alignas(64) on Params struct and TMA descriptor members
• alignas(128) for smem_SFA/SFB scale factor storage
• Applied to all SM90/SM100/SM120 epilogue and mainloop collectives

Benchmark results (RTX 5090, SM120a):

NVF4 BF16 GEMM (Example 79a equivalent):

Matrix Size	TFLOPS
8192x8192	174
16384x16384	316

Pure NVF4 GEMM:

Matrix Size	TFLOPS
8192x8192	270
16384x16384	446

Files modified: 17 headers across cute/arch, cutlass/epilogue, cutlass/gemm collectives.

I don't currently have access to SM90 / SM100 hardware, so testing there would be very helpful.

[lsyyy666]

@m96-chan Thanks for the feedback! Quick question, why did you need to alignas(128) for smem_SFA/SFB scale factor storage while it only needs to be alignas(16) ?

[m96-chan]

@lsyyy666 Thank you for the review!

Why I originally used alignas(128):

When debugging the misaligned address crash, my focus was on identifying where alignment was being lost (traced to partition_S() dropping from 1024 to 8 bytes). For the fix value, I chose 128 empirically as a "safe" power-of-2 without deeply analyzing the exact PTX requirements for ldmatrix.sync.aligned.

Confirmed alignas(16) works:

I tested with alignas(16) on RTX 5090 (SM 120a):

• Build passes
• NVF4 GEMM/GEMV tests pass
• Benchmarks run correctly

You're right that 16-byte alignment is sufficient since ldmatrix loads 16 bytes per 4-thread group. The 128-byte choice was overly conservative. Thanks for the correction!

[lsyyy666]

@m96-chan Got it, thanks. To futher clarify, the 16-byte alignment requirement is probably due to TMA load. The scale factors are loaded by ld.shared instead of ldmatrix.sync.aligned.

[lsyyy666]

More on tensor descriptor - CUDA toolkit 13.1 is automatically aligning it to 128B which has known issues with MSVC.
CUDA toolkit 13.2 will revert the alignment of tensor descriptors back to 64B by default. Before 13.2, you can workaround this issue with alignas(64)

[m96-chan]

@lsyyy666 Thank you for the clarification - I was exploring empirically on this new architecture, so this is very helpful.

Since CUDA 13.1 support is likely needed on both sides, how about proceeding with this workaround (alignas(16) for scale factors + alignas(64) for TMA descriptors)?

If SM90/SM100 testing would help, I can run tests on RunPod - let me know.