[Bug] Transformer Engine SM120 (Blackwell) Compatibility Issue

TE Version: 2.17.0
Affected: NVIDIA Blackwell architecture GPUs (SM120, compute capability 12.0)

Problem Description

On NVIDIA Blackwell architecture GPUs (SM120, compute capability 12.0), transformer_engine 2.17.0's LayerNormMLP module crashes when using fp32 dtype, with the error:

RuntimeError: /TransformerEngine/transformer_engine/common/activation/./../cast/dispatch/../fp8/gated_fp8.cuh:341
in function cast_gated_tma: CUDA Error: invalid argument

bf16/fp16 dtypes are not affected — only fp32 triggers this issue.

Environment

Component	Version/Model
GPU	NVIDIA Blackwell (SM120)
Compute Capability	12.0 (SM120)
Shared Memory per Block	100 KB
Transformer Engine	2.17.0+2e559f06
PyTorch	2.12.0+cu132
CUDA	13.2
Python	3.12

Minimal Reproduction Code

import torch
from transformer_engine.pytorch import LayerNormMLP

# Create fp32 LayerNormMLP module
module = LayerNormMLP(
    hidden_size=2560,
    ffn_hidden_size=10240,
    activation="swiglu",
    params_dtype=torch.float32,
).cuda()

# Create fp32 input
x = torch.randn(2, 16, 2560, dtype=torch.float32, device="cuda")

# Execute forward - crashes on SM120 GPUs
with torch.no_grad():
    out = module(x)

Expected: Forward completes successfully, output shape (2, 16, 2560)

Actual: Raises RuntimeError: CUDA Error: invalid argument

Control test: Changing params_dtype and input dtype to torch.bfloat16 allows the same code to run successfully on SM120.

Root Cause Analysis

TMA Kernel Shared Memory Exceeds Limit

The issue lies in TE's TMA (Tensor Memory Accelerator) kernel dispatch logic:

1. TE's dispatch condition (transformer_engine/common/cast/dispatch/gated.cuh:49):

   const bool use_tma_kernels = (cols % 32 == 0) && is_supported_by_CC_100();
   

2. is_supported_by_CC_100() implementation (transformer_engine/common/common.cu:231):

   bool is_supported_by_CC_100() {
     int deviceComputeCapability = cuda::sm_arch(cuda::current_device());
     return deviceComputeCapability >= 100;   // CC >= 10.0 = Blackwell and newer
   }
   

3. The problem: This function enables TMA for all GPUs with CC ≥ 10.0 (Blackwell architecture), but SM120's shared memory per block (100 KB) is insufficient for fp32 TMA kernels.

4. TMA kernel shared memory requirements (proportional to dtype):

   • bf16/fp16: ~64 KB ✅ (within SM120's 100 KB limit)
   • fp32: ~128 KB ❌ (exceeds SM120's 100 KB limit)

5. Trigger conditions (all must be met):

   • GPU is SM120 (Blackwell) or same architecture
   • Tensor dtype is fp32
   • cols % 32 == 0
   • Using te.LayerNormMLP or te.LayerNormLinear

Shared Memory Calculation

TMA kernel cast_gated_tma shared memory formula:

Parameter	Value
SHMEM_DIM_Y × SHMEM_DIM_X	32 × 128
BUFFERS_NUM	2
buff_elems_total	2 × 32 × 128 = 8192
TMA_SHMEM_ALIGNMENT	128 B

bf16 (2 byte/elem):

• Single buffer: DIVUP(8192 × 2, 128) = 16384 B
• 4 buffers (in_act, in_gate, out_act, out_gate): 4 × 16384 = 65536 B
• • alignment: 65664 B ≈ 64 KB ✅

fp32 (4 byte/elem):

• Single buffer: DIVUP(8192 × 4, 128) = 32768 B
• 4 buffers: 4 × 32768 = 131072 B
• • alignment: 131200 B ≈ 128 KB ❌

Reproduction Verification

Run the minimal reproduction code on an SM120 GPU to confirm:

• ✅ fp32 + TE 2.17.0 → crashes (reproduces the issue)
• ✅ bf16 + TE 2.17.0 → works (control test)

---

Relevant code locations:

• TMA kernel: transformer_engine/common/cast/fp8/gated_fp8.cuh:341
• Dispatch logic: transformer_engine/common/cast/dispatch/gated.cuh:49
• CC check: transformer_engine/common/common.cu:231

Reproduce the fp32 and bf16 LayerNormMLP cases on an SM120 GPU using the provided example. Read transformer_engine/common/cast/dispatch/gated.cuh:49, transformer_engine/common/common.cu:231, and transformer_engine/common/cast/fp8/gated_fp8.cuh:341 to trace the TMA dispatch and shared-memory requirements. Done means the fp32 forward completes with the expected shape while the bf16 control case remains working.