nvfp4-megamoe-kernel/memory/2026-05-29-tma-async.md

Session: 2026-05-29 04:33:00 UTC

TMA Async Load — Stage D

Started work on TMA async loads for FMHA kernel. Goal: replace scalar GMEM reads with TMA bulk async copies.

Key Discoveries

1. CUDA 13 cuTensorMapEncodeTiled requires byte strides (not element strides)

   • Old (CUDA 12): globalStrides[] = {1, cols} — element strides
   • New (CUDA 13): globalStrides[] = {cols*2, cols*2*rows} — byte strides
   • This was the root cause of ALL 2D descriptor creation failures

2. CUDA 13 cuTensorMapEncodeTiled requires rank >= 2 (2D, 3D, 4D, or 5D)

   • 1D descriptors still work but are limited
   • 2D descriptors work with byte strides
   • 3D descriptors (degenerate dim=1) also work

3. TMA load kernel HANGS — descriptor creates OK but cp.async.bulk.tensor.{2d,3d} never completes

   • Both 2D and 3D descriptors create successfully
   • The cp.async.bulk.tensor.2d / .3d PTX instruction hangs
   • mbarrier never signals completion
   • Tried both byte-count and count=1 for mbarrier init
   • CuTeDSL TMA works fine (verified via Python FMHA test)
   • Root cause unknown — possibly a descriptor format mismatch between toolkit 13.2 and driver 13.0

Current Status

• fmha_tma.cuh: TMA descriptor helper (3D, byte strides, BFLOAT16)
• fmha_6warp_tma.cuh: TMA-integrated multirow kernel
• test_fmha_tma.cu: Test harness
• BLOCKED: TMA load hangs on B200

Next Steps

• Need to figure out why cp.async.bulk.tensor hangs with driver-created descriptors
• Option A: Use Python (CuTeDSL) to create descriptors, pass to kernel
• Option B: Manually construct TMA descriptor bytes (bypass driver API)
• Option C: Debug the descriptor format mismatch