[BUG] NVFP4BlockScaling forward pass fails on RTX 5080 (sm_120) when K is a multiple of 128; sm_100/B100 TMEM+UMMA hardcoded in row_cast_col_hadamard_transform_cast_fusion.cu

**TE version:** 2.14.1 **Hardware:** NVIDIA GeForce RTX 5080 (consumer Blackwell, sm_120, 16 GB VRAM) **Stack:** torch 2.11.0+cu130, Python 3.12, WSL2 Ubuntu 24.04, CUDA 13.0 **NVIDIA Driver:** 596.36 (Studio); same behavior on 595.79 **Build:** TE 2.14.1 source-built with `TORCH_CUDA_ARCH_LIST=12.0` ## Symptom `is_nvfp4_available()` returns `True` on RTX 5080 sm_120, but actual NVFP4 forward kernels fail with `CUDA Error: invalid argument` at most production dims. Initial square-shape boundary probe: | dim (square M=N=K) | Result | |---|---| | 64 | ✅ OK | | 96 | ✅ OK | | 128 | ❌ FAIL | | 192 | ✅ OK | | 256, 384, 512, 768, 1024, 2048, 4096 | ❌ FAIL | After asymmetric probe (varying M/K/N independently — see "Why some shapes work" section below for full data), we found the failure is gated by the **K (input/inner) dim alone**, with a precise pattern: > **Kernel works iff `K % 32 == 0 AND K % 128 != 0`** All multiples of 128 fail; all non-multiples of 32 fail. M and N can be anything (tested up to 4096) and don't affect the outcome. Error stack: ``` RuntimeError: /TransformerEngine/transformer_engine/common/hadamard_transform/ row_cast_col_hadamard_transform_cast_fusion.cu:1200 in function row_col_rht_gemm_ntt_w_sfc: CUDA Error: invalid argument ``` ## Reproduction ```python import torch import transformer_engine.pytorch as te from transformer_engine.common.recipe import NVFP4BlockScaling print(torch.cuda.get_device_name(0)) # "NVIDIA GeForce RTX 5080" print(torch.cuda.get_device_capability()) # (12, 0) → sm_120 print(te.is_nvfp4_available()) # True # Works at dim=64, 96, 192: x = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") layer = te.Linear(64, 64, params_dtype=torch.bfloat16, device="cuda") with te.fp8_autocast(enabled=True, fp8_recipe=NVFP4BlockScaling()): out = layer(x) # OK # Fails at dim=128, 256+: x = torch.randn(1024, 4096, dtype=torch.bfloat16, device="cuda") layer = te.Linear(4096, 4096, params_dtype=torch.bfloat16, device="cuda") with te.fp8_autocast(enabled=True, fp8_recipe=NVFP4BlockScaling()): out = layer(x) # CUDA Error: invalid argument ``` Tested with `CUDA_LAUNCH_BLOCKING=1` for synchronous error reporting. All four `NVFP4BlockScaling` recipe variants tested (default RHT, RHT off, power_2_scale + RHT off, smaller shape) — all fail at production dims. ## Root cause analysis (source-level) The failing kernel (`transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu`) has hardcoded **sm_100/B100 hardware feature dependencies** that consumer Blackwell sm_120 does not have: ### 1. TMEM (Tensor Memory) — sm_100/B100 ONLY hardware Around line 1110: ```cpp static uint32_t constexpr Sm100TmemCapacityColumns = 512; static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns; static uint32_t constexpr AccumulatorPipelineStageCount = TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape)); ``` TMEM is exclusive to sm_100/B100/B200 silicon. Consumer Blackwell sm_120 (RTX 5080/5090) does not have TMEM hardware at all. ### 2. UMMA (Unified MMA) — sm_100 instruction class ```cpp cutlass::detail::CustomizedPipelineTmaUmmaAsync<...> cutlass::PipelineUmmaAsync<...> ``` UMMA is the sm_100 Tensor Core instruction set. sm_120 has the standard HMMA/QMMA Tensor Core instructions but not UMMA. ### 3. 232 KB shared memory budget — sm_100-specific Line 1131: ```cpp static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes ``` sm_120 has only ~99 KB opt-in dynamic shared memory per block (vs sm_100's 228 KB). This budget assumption is sm_100-specific. ### 4. Explicit sm_100 selectors ```cpp cutlass::gemm::collective::detail::sm100_smem_selector<...> make_tma_copy_A_sm100(...) make_tma_copy_B_sm100(...) ``` These select sm_100-specific layout / TMA descriptor / kernel parameters. There are no `sm_120_*` equivalents in the current codebase. ### `is_nvfp4_available()` is a false positive on sm_120 The helper checks compute-capability tier (`>= sm_100`-level Blackwell) without testing for TMEM/UMMA hardware availability. As a result it returns `True` on sm_120 even though no usable NVFP4 kernel exists for that target. ## Why some shapes work — precise K-axis pattern (mod 128) After the initial square-shape boundary probe, we ran an asymmetric probe varying M, K, N independently: ### Test 1: M scales (K=N=64 fixed) M=64, 192, 256, 1024, 4096 → **ALL OK**. M is NOT a bottleneck. ### Test 2: N scales (M=K=192 fixed) N=64, 192, 256, 1024, 4096 → **ALL OK**. N is NOT a bottleneck. ### Test 3: K scales (M=N=192 fixed) K=64, 192 → OK; K=256, 1024, 4096 → **FAIL**. **K IS the sole bottleneck.** ### Test 4: Production target (K=N=4096) varying M All FAIL regardless of M. Confirms K=4096 gates everything. ### Fine-grained K sweep (M=N=192, K from 16 to 3072, multiples of 16) ``` K=16 FAIL K=160 OK K=320 OK K=32 OK K=176 FAIL K=384 FAIL K=48 FAIL K=192 OK K=448 OK K=64 OK K=208 FAIL K=512 FAIL K=80 FAIL K=224 OK K=640 FAIL K=96 OK K=240 FAIL K=768 FAIL K=112 FAIL K=256 FAIL K=1024 FAIL K=128 FAIL K=272 FAIL K=1536 FAIL K=144 FAIL K=288 OK K=2048 FAIL K=320 OK K=3072 FAIL ``` **The pattern is exact:** the kernel works iff `K % 32 == 0 AND K % 128 != 0`. Working K (mod 128) ∈ {32, 64, 96}. Failing K (mod 128) ∈ {0, 16, 48, 80, 112}. The kernel's CUTLASS dispatch likely has a single-K-tile path for K-values in the {32, 64, 96}-mod-128 family that bypasses the TMEM+UMMA pipeline, versus a multi-K-tile path for K-values that are multiples of 128 (triggering the sm_100-only TMEM accumulator). The non-multiples-of-32 fail at an earlier gate (presumably an alignment requirement somewhere in the kernel). This is what makes the bug invisible to most users: tests at dim ∈ {64, 192} look fine; production dims (typical 128/256/512/1024/4096) ALL fall in the failing class because they're all multiples of 128. ## Implication for software workarounds We considered K-chunking the input to route production matmuls through working K values, but the math doesn't work: ``` Largest working K chunk: 448 Linear(K=4096) chunks needed: ceil(4096 / 448) = 10 Per-chunk time @ 60 TFLOPS: ~63 μs Launch overhead (10 calls): ~300 μs Total chunked NVFP4: ~930 μs bf16 native baseline: ~378 μs Net: 2.5× SLOWER than bf16 ``` To break even, we'd need K-chunk size ≥ 1024, but K=1024 is a multiple of 128 → fails. **No workaround is performant at production scale.** ## Driver verification Tested with NVIDIA Studio Driver 595.79 → 596.36 (full Windows reboot between tests). Got bit-identical boundary patterns — confirming this is a TE source-level issue, not a driver-level launch validation bug. ## Suggested fix (two non-exclusive options) ### Option A (preferred): add an sm_120-targeted NVFP4 Hadamard kernel A separate kernel (e.g. `transformer_engine/common/hadamard_transform/sm120_row_cast_col_hadamard_transform_cast_fusion.cu` or a templated dispatch on `__CUDA_ARCH__ == 1200`) that uses sm_120-available primitives: - Replace TMEM with shared memory + register tiles (or TMEM-equivalent spilling pattern) - Replace UMMA with HMMA/QMMA (sm_120 Tensor Core instructions) - Replace TMA descriptors with `cp.async` or `cute::CooperativeAtomTile` - Fit within sm_120's 99 KB opt-in dynamic shared memory budget This unlocks NVFP4 for the entire RTX 50 series consumer Blackwell line. ### Option B (interim): fix `is_nvfp4_available()` to return False on sm_120 Until the sm_120 kernel is ready, prevent the false positive at the helper level. Users will at least know they need to fall back to FP8 or BF16 without hitting the cryptic CUDA invalid-argument error mid-training. ```python def is_nvfp4_available(): cap = torch.cuda.get_device_capability() if cap == (10, 0): # sm_100, B100 return True if cap == (12, 0): # sm_120, consumer Blackwell — kernel not ready return False return False ``` (Or check for TMEM availability via a `cudaDeviceGetAttribute` query if such an attribute is exposed.) ## Parallel-path verification (cuBLASLt + PyTorch nightly also RED) To rule out NVFP4-specific issues in TE alone, we tested the parallel software paths to FP4 on RTX 5080 sm_120. All RED: ### cuBLASLt 13.4.1.1 (CUDA 13.0) ```bash # Symbol scan of libcublasLt.so.13 nm -D libcublasLt.so.13 | grep -iE 'fp4|float4|e2m1|nvfp4|mxfp4' # (zero matches) strings libcublasLt.so.13 | grep -iE 'fp4|float4|e2m1' # (zero matches) # Header check (2674 lines) grep -E 'FP4|FLOAT4|E2M1' /usr/local/cuda/include/cublasLt.h # (zero matches) ``` **FP4 is not in the cuBLASLt API at all on CUDA 13.0.** This is consistent with NVIDIA's FP4 path going through cutlass (TE's underlying templated library), not cuBLASLt. FP8 IS in cuBLASLt and works fine on sm_120 (we verified 1.43× speedup via TE's `Float8BlockScaling` recipe — see "Useful contrast" section below). ### PyTorch 2.13.0.dev20260502+cu130 (nightly, 2026-05-02) ```python import torch # FP4 dtype IS exposed: torch.float4_e2m1fn_x2 # ok, itemsize=1 torch.uint4 # ok, itemsize=1 # But cast op is NOT wired: x = torch.randn(64, 64, dtype=torch.bfloat16, device='cuda') y = x.to(torch.float4_e2m1fn_x2) # RuntimeError: copy_() does not support casting Float4_e2m1fn_x2 # to different types. Source dtype is BFloat16 target dtype is Float4_e2m1fn_x2 ``` **Same exact error as PyTorch 2.11.0+cu130** — two minor versions later, no kernel-layer progress on the FP4 path. PyTorch's FP4 is dtype-scaffolding only; the actual kernels haven't been implemented. Likely waiting on cuBLASLt or cutlass-sm_120 to ship working dispatch first. ## Useful contrast — FP8 works perfectly on the same hardware We benchmarked `Float8BlockScaling` on the SAME RTX 5080 sm_120 silicon with the same TE 2.14.1 install and same shape we used for NVFP4: | Recipe | TFLOPS | vs bf16 | |---|---|---| | torch.nn.Linear bf16 (baseline) | 96.89 | 1.00× | | **Float8BlockScaling** | **138.10** | **1.43×** ← works! | | MXFP8BlockScaling | FAIL | clean error: "MXFP8 (for all gemm layouts) is not supported on 12.0+ architectures yet" | | NVFP4BlockScaling | FAIL | cryptic CUDA invalid argument | **Two observations:** 1. **FP8 GEMM dispatches correctly to sm_120 tensor cores.** This proves the hardware works for fp-quantized matmul; the issue is specifically that NVFP4's kernel uses sm_100-only primitives (TMEM/UMMA), not a missing hardware capability. 2. **MXFP8 fails with a CLEAN explicit error message** — "not supported on 12.0+ architectures yet" — and skips kernel launch. This is the correct gating pattern. **NVFP4 should follow it** instead of returning True from `is_nvfp4_available()` and then crashing inside the kernel with a cryptic CUDA error. The fact that NVFP4 has the SAME hardware-feature gap as MXFP8 (both need sm_100 features absent on sm_120) but DIFFERENT user-facing behavior (clean error vs cryptic crash) suggests the gating logic exists in the codebase — it just needs to be applied to NVFP4 too. ## Project context We're a small ML/AI project building on RTX 5080 single-card local training. We attempted NVFP4 to accelerate 4-bit weight matmul during supervised fine-tuning (~3× target speedup vs the bnb NF4 LUT path that doesn't engage Tensor Cores at all). Blocking on this issue means we stay on bitsandbytes NF4 (CUDA-cores only) for now. We also confirmed via separate boundary probe + source inspection of bnb 0.49.2 that bnb's FP4 path is hardcoded LUT (`fp4_dequantization_lut[8]`) with no cuBLASLt/cutlass dispatch — so bnb FP4 also doesn't engage Tensor Cores on consumer Blackwell. We'll track and re-test when: - TE >= 2.15 ships with explicit consumer Blackwell support in changelog, OR - A separate sm_120 Hadamard kernel lands in the `transformer_engine/common/hadamard_transform/` directory, OR - `is_nvfp4_available()` is updated to return `False` on sm_120 (would unblock graceful FP8 fallback for downstream users) Thanks for the great library — looking forward to consumer Blackwell support landing!