From 6fb3d54c02fac588c10624664da78230456e9fe8 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 07:54:01 +0000 Subject: [PATCH] docs: update here-docs with CuTeDSL rationale for NVIDIA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updated fmha_common.cuh, fmha_sm100.cuh, fmha_epilogue_sm100.cuh, and fmha_sm100_launch.cuh with comprehensive here-docs explaining: 1. The 4 CuTeDSL gaps that forced us to raw CUDA C++: - TMEM round-trip broken (Ld32x32bOp/St32x32bOp column mismatch) - Float→int impossible (arith.fptosi not lowerable) - epilogue_tma_store blocks multi-CTA - hd=512 MLIR optimizer hangs 2. TMEM lane mapping (verified on B200): - Lane i → positions i*4+0..3, 128 FP32 per column - Warp-collective: ALL 32 lanes must call ld/st or HANG - Column address = tmem_base + column_index 3. Key insight for NVIDIA: float→int gap is the single most impactful limitation, blocking ALL quantization-epilogue fusion --- dsv4/kernels/attention/fmha_common.cuh | 126 +++++++++++++----- .../kernels/attention/fmha_epilogue_sm100.cuh | 54 ++++++-- dsv4/kernels/attention/fmha_sm100.cuh | 14 +- dsv4/kernels/attention/fmha_sm100_launch.cu | 25 +++- 4 files changed, 163 insertions(+), 56 deletions(-) diff --git a/dsv4/kernels/attention/fmha_common.cuh b/dsv4/kernels/attention/fmha_common.cuh index 559125f0..4afd189a 100644 --- a/dsv4/kernels/attention/fmha_common.cuh +++ b/dsv4/kernels/attention/fmha_common.cuh @@ -1,48 +1,104 @@ /** - * DSV4 FMHA shared definitions — base header for raw CUDA kernels. + * DSV4 FMHA shared definitions — base header for raw CUDA C++ kernels. * * ================================================================== - * WHY THIS EXISTS + * WHY RAW CUDA C++ INSTEAD OF CUTEDSL * ================================================================== - * CuTeDSL (the Python DSL for CUTLASS) has fundamental limitations - * on Blackwell SM100 that make certain operations impossible: * - * 1. TMEM round-trip is BROKEN (Ld32x32bOp/St32x32bOp column mismatch) - * 2. Float-to-int conversion is IMPOSSIBLE (arith.fptosi not lowerable) - * 3. epilogue_tma_store BLOCKS multi-CTA (can't accept flat_divide coords) - * 4. hd=512 MLIR backend HANGS (>3hr optimizer time) + * CuTeDSL (the Python DSL for CUTLASS) has FOUR fundamental gaps on + * Blackwell SM100 that made it impossible to ship a production FMHA + * kernel. We spent weeks hitting these walls before switching to raw + * CUDA C++ with inline PTX. The reference kernel you see here was + * written and verified correct in ~2 hours. * - * This header provides the building blocks for writing FMHA in raw - * CUDA C++ with inline PTX, bypassing ALL of the above. + * GAP 1: TMEM round-trip is BROKEN + * --------------------------------- + * CuTeDSL's Ld32x32bOp and St32x32bOp atoms have different column + * mappings — they address DIFFERENT TMEM columns even when constructed + * with the same layout. A NO-OP round-trip (load → store unchanged) + * corrupts data by ~3% (cos ~0.97). This makes in-kernel O rescale + * (the D1.5 fix) impossible in CuTeDSL — you can't read the TMEM + * accumulator, multiply by acc_scale, and write it back. The MoE + * kernel avoids this because it only does a ONE-WAY TMEM → regs → + * SMEM → GMEM epilogue. FMHA needs a round-trip for per-KV-tile + * rescaling. Raw CUDA with tcgen05.ld/st PTX gives us full control + * over TMEM addressing. + * + * GAP 2: Float-to-int conversion is IMPOSSIBLE + * --------------------------------------------- + * The CuTeDSL MLIR pipeline CANNOT lower ANY float→int operation. + * arith.fptosi, nvvm.inline_ptx with cvt.rni.s32.f32, llvm.bitcast + * Float32→Int32 — ALL fail with "LLVM ERROR: unsupported operation." + * This makes quantize-in-epilogue fusion (NVFP4-1.1) impossible in + * CuTeDSL. Every quantization kernel needs f32→i32. Raw CUDA has + * __float2int_rn which works perfectly. + * + * GAP 3: epilogue_tma_store BLOCKS multi-CTA + * ------------------------------------------- + * The epilogue_tma_store helper cannot accept flat_divide-based GMEM + * coordinates, which are required for multi-CTA grid launches. This + * means FMHA is stuck at per-head Python launch (128 launches per Pro + * decode step) instead of a single GPU launch. The MoE kernel avoids + * this by using the explicit epilogue_tmem_copy + epilogue_smem_copy + + tma_partition pipeline. Raw CUDA gives us full control over TMA + * descriptors and GMEM writes. + * + * GAP 4: hd=512 MLIR backend HANGS + * --------------------------------- + * The CuTeDSL MLIR optimizer cannot process the hd=512 kernel in + * reasonable time. Tracer completes in 0.8s (kernel is structurally + * correct), but the MLIR optimizer runs for 3+ hours before we kill + * it. This is a CuTeDSL/MLIR toolchain limitation, not a kernel bug. + * Raw CUDA compiles hd=512 in seconds. * * ================================================================== - * WHAT WORKS (tested on B200) + * WHAT WORKS (tested on B200, CUDA 13.2, SM100) * ================================================================== * - BF16 conversion via inline PTX cvt.rn.bf16.f32 / cvt.f32.bf16 - * - Warp reductions (fmax, sum) - * - TMEM alloc/dealloc via tcgen05 PTX - * - TMEM load/store via tcgen05.ld/st (uint32_t b32 registers) - * - TMEM fence via tcgen05.fence + * - Warp reductions (fmax, sum) via shfl_xor + * - TMEM alloc/dealloc via tcgen05.alloc/dealloc PTX + * - TMEM load/store via tcgen05.ld/st PTX (uint32_t b32 registers) + * - TMEM fence via tcgen05.wait::st/ld.sync.aligned + * - One-way TMEM epilogue: TMEM → regs → normalize → BF16 → GMEM + * (verified cos 0.999999 at hd=64, cos 0.999998 at hd=128) * * ================================================================== - * WHAT'S BROKEN / NEEDS WORK + * TMEM LANE MAPPING (verified on B200 via test_tmem_lane_mapping.cu) * ================================================================== - * - TMEM load/store column addressing: the exact column offset - * calculation for row groups (8 row-groups per column) needs - * verification. The kernel using these ops hangs on B200. - * - tcgen05.mma (QK/PV GEMM): UMMA SMEM descriptor construction - * is placeholder only. The descriptor bitfield format is known - * (see cute/arch/mma_sm100_desc.hpp SmemDescriptor) but the - * exact values for our Q/K layouts haven't been validated. + * tcgen05.st/ld 16x256b.x1.b32 are warp-collective operations: + * - ALL 32 lanes in a warp MUST execute them (or HANG) + * - Each lane i reads/writes positions i*4+0..i*4+3 within the column + * - 32 lanes × 4 FP32 = 128 FP32 per column + * - For row 0: lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127 + * - HD values need ceil(HD/128) TMEM columns + * - Column address = tmem_base + column_index + * - tmem_base is WRITTEN to SMEM by tcgen05.alloc (read it back after alloc) + * + * CRITICAL: If fewer than 32 lanes call tmem_store/tmem_load, the + * warp is divergent on a collective operation and the GPU HANGS. + * Always loop over enough columns that all 32 lanes participate, + * even if some write don't-care data to unused columns. + * + * ================================================================== + * TMEM ALLOC/DEALLOC + * ================================================================== + * - tcgen05.alloc: pass SMEM pointer (via __cvta_generic_to_shared) + * and num_columns (power of 2, minimum 32). Must be called by an + * ENTIRE fully active warp (all 32 lanes). The alloc WRITES the + * tmem_base value to the SMEM location. + * - tcgen05.dealloc: pass tmem_base (the VALUE from SMEM, not the + * SMEM pointer) and num_columns. Must also be warp-collective. * * ================================================================== * KEY INSIGHT FOR NVIDIA * ================================================================== - * CuTeDSL's inability to lower float→int is a fundamental gap. - * Every quantization kernel needs f32→i32. The fact that nvvm.inline_ptx - * also fails suggests the CuTeDSL MLIR pipeline simply doesn't have a - * lowering path for ANY float→integer type conversion. This makes - * quantize-in-epilogue fusion impossible in CuTeDSL. + * The CuTeDSL float→int gap is the single most impactful limitation. + * It blocks ALL quantization-epilogue fusion (NVFP4-1.1, NVFP4-1.2) + * and forces a separate kernel launch for every quantize step. The + * TMEM round-trip issue is the second most impactful — it blocks + * in-kernel O rescaling, forcing 5-9 Python kernel launches per + * decode step instead of 1 GPU launch. Both of these work trivially + * in raw CUDA C++ with inline PTX. */ #pragma once @@ -78,7 +134,13 @@ __device__ __forceinline__ float wsum(float v) { for(int o=16;o>0;o>>=1) v+=__shfl_xor_sync(0xFFFFFFFF,v,o); return v; } -// TMEM operations — using uint32_t registers per CUTLASS reference +// ================================================================== +// TMEM operations — inline PTX for Blackwell SM100 +// ================================================================== +// These wrap the tcgen05 PTX instructions for TMEM management. +// All TMEM load/store ops are WARP-COLLECTIVE: all 32 lanes must +// execute them. See the header comment for the lane mapping. +// ================================================================== __device__ void tmem_alloc(uint32_t smem_ptr, int num_cols) { asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" @@ -90,24 +152,26 @@ __device__ void tmem_dealloc(uint32_t tmem_ptr, int num_cols) { :: "r"(tmem_ptr), "r"(num_cols)); } -/** Load 16 rows × 256 bits from TMEM column. 4 uint32_t registers per thread. */ +/** Load 16 rows × 256 bits from TMEM column. Warp-collective. 4 uint32_t per lane. */ __device__ void tmem_load(uint32_t col_addr, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) { asm volatile("tcgen05.ld.sync.aligned.16x256b.x1.b32 {%0, %1, %2, %3}, [%4];" : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(col_addr)); } -/** Store 16 rows × 256 bits to TMEM column. 4 uint32_t registers per thread. */ +/** Store 16 rows × 256 bits to TMEM column. Warp-collective. 4 uint32_t per lane. */ __device__ void tmem_store(uint32_t col_addr, uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3) { asm volatile("tcgen05.st.sync.aligned.16x256b.x1.b32 [%0], {%1, %2, %3, %4};" :: "r"(col_addr), "r"(r0), "r"(r1), "r"(r2), "r"(r3)); } +/** Wait for prior TMEM stores to be visible. Call after tmem_store. */ __device__ void tmem_fence_store() { asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); } +/** Wait for prior TMEM loads to complete. Call after tmem_load. */ __device__ void tmem_fence_load() { asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); } diff --git a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh index 7de3c15a..1aa0e6cf 100644 --- a/dsv4/kernels/attention/fmha_epilogue_sm100.cuh +++ b/dsv4/kernels/attention/fmha_epilogue_sm100.cuh @@ -1,20 +1,50 @@ /** * DSV4 FMHA Phase 2 — TMEM accumulator + one-way correction epilogue. * - * STATUS: WORKING — TMEM pipeline functional (SMEM → TMEM → regs → normalize → GMEM) + * ================================================================== + * STATUS: WORKING — cos 0.999999 at hd=64, cos 0.999998 at hd=128 + * ================================================================== * - * This kernel proves the MoE-style one-way correction epilogue works for FMHA: - * 1. Compute attention in SMEM (same as reference) - * 2. Write accumulator to TMEM (warp-collective store) - * 3. Read from TMEM to registers (warp-collective load) - * 4. Normalize in registers (per-lane math) - * 5. Cast to BF16 and write to GMEM + * This kernel proves the MoE-style one-way correction epilogue works + * for FMHA on Blackwell SM100, using raw CUDA C++ with inline PTX + * (bypassing all CuTeDSL limitations — see fmha_common.cuh). * - * TMEM lane mapping (verified on B200 via test_tmem_lane_mapping.cu): - * tcgen05.st/ld 16x256b.x1.b32 is warp-collective. Each lane i - * writes/reads positions i*4+0..i*4+3 within the column. - * 32 lanes × 4 FP32 = 128 FP32 per column. - * For row 0: lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127. + * Pipeline: SMEM → TMEM (warp-collective store) → regs (warp-collective + * load) → normalize in regs → BF16 cast → GMEM. + * + * ================================================================== + * WHY THIS MATTERS (Priority 2 from ROADMAP) + * ================================================================== + * This is the one-way correction epilogue pattern that the MoE kernel + * uses successfully in CuTeDSL: + * TMEM → regs (tcgen05.ld) → [normalize/cast/pack] → GMEM + * + * In CuTeDSL, this pattern is done with epilogue_tmem_copy_and_partition + * + epilogue_smem_copy_and_partition (paired atoms, one-way only). + * FMHA needs this for the final epilogue (after all KV tiles), and it + * also UNBLOCKS: + * - D2 multi-CTA grid (128 Python launches → 1 GPU launch) + * - NVFP4-1.2 (register slot for FP4 amax + pack in epilogue) + * - In-kernel normalize (O / row_sum without TMEM round-trip) + * + * In raw CUDA, we implement this with tcgen05.ld/st PTX directly, + * giving us full control over TMEM addressing and avoiding the + * Ld32x32bOp/St32x32bOp column mismatch that plagues CuTeDSL. + * + * ================================================================== + * TMEM LANE MAPPING (verified on B200 via test_tmem_lane_mapping.cu) + * ================================================================== + * tcgen05.st/ld 16x256b.x1.b32 are warp-collective operations: + * - ALL 32 lanes in a warp MUST execute them (or GPU HANGS) + * - Each lane i writes/reads positions i*4+0..i*4+3 within the column + * - 32 lanes × 4 FP32 = 128 FP32 per column + * - For row 0: lane 0 = positions 0-3, lane 1 = 4-7, ..., lane 31 = 124-127 + * - HD values need ceil(HD/128) TMEM columns + * - Column address = tmem_base + column_index + * + * CRITICAL: If fewer than 32 lanes call tmem_store/tmem_load, the + * warp is divergent on a collective operation and the GPU HANGS. + * Always iterate over enough columns that all 32 lanes participate. */ #pragma once #include "fmha_common.cuh" diff --git a/dsv4/kernels/attention/fmha_sm100.cuh b/dsv4/kernels/attention/fmha_sm100.cuh index 6fc81d32..08979484 100644 --- a/dsv4/kernels/attention/fmha_sm100.cuh +++ b/dsv4/kernels/attention/fmha_sm100.cuh @@ -7,19 +7,19 @@ * * This is the CORRECT reference implementation. It proves that: * - The online softmax with O rescale approach is mathematically correct - * - D3 SWA masking works + * - D3 SWA masking works in raw CUDA * - Raw CUDA C++ compiles and runs on Blackwell SM100 without CuTeDSL + * - The kernel infrastructure (nvcc compilation, standalone test) works * * ================================================================== - * WHY RAW CUDA INSTEAD OF CUTEDSL + * WHY THIS EXISTS (see fmha_common.cuh for the full rationale) * ================================================================== * CuTeDSL hit 4 fundamental walls on Blackwell: - * 1. TMEM round-trip broken (D1.5) — Ld32x32bOp/St32x32bOp mismatch - * 2. Float→int impossible — arith.fptosi not lowerable to PTX + * 1. TMEM round-trip broken (Ld32x32bOp/St32x32bOp column mismatch) + * 2. Float→int impossible (arith.fptosi not lowerable to PTX) * 3. epilogue_tma_store blocks multi-CTA * 4. hd=512 MLIR optimizer hangs * - * Writing in raw CUDA gives us full PTX control and bypasses all of these. * This reference kernel took ~2 hours to get working. The equivalent * CuTeDSL kernel took weeks and still has the D1.5 blocker. * @@ -34,7 +34,9 @@ * These are all solvable incrementally. The critical milestone is: * CORRECT FMHA OUTPUT IN RAW CUDA ON BLACKWELL SM100. * - * Next phase: Parallelize across threads, add tcgen05.mma for QK/PV. + * Next phase: Add tcgen05.mma for QK/PV tensor core acceleration + * (see fmha_epilogue_sm100.cuh for the TMEM pipeline that's already + * working). */ #pragma once #include "fmha_common.cuh" diff --git a/dsv4/kernels/attention/fmha_sm100_launch.cu b/dsv4/kernels/attention/fmha_sm100_launch.cu index 652358aa..6df0bf3f 100644 --- a/dsv4/kernels/attention/fmha_sm100_launch.cu +++ b/dsv4/kernels/attention/fmha_sm100_launch.cu @@ -4,20 +4,31 @@ * ================================================================== * STATUS: COMPILES but doesn't run via torch.utils.cpp_extension * ================================================================== - * The kernel compiles cleanly with nvcc (see test_fmha_sm100.py), - * but torch JIT compilation fails due to __bf16 / bf16_t type - * conflicts with PyTorch's -D__CUDA_NO_BFLOAT16_CONVERSIONS__ flag. + * The kernel compiles cleanly with nvcc, but torch JIT compilation + * fails due to bf16_t (unsigned short) conflicting with PyTorch's + * -D__CUDA_NO_BFLOAT16_CONVERSIONS__ flag, which triggers an nvcc + * internal compiler error when __bf16 is used on SM100. * - * Workaround: Use the standalone test (test_fmha_sm100_standalone.cu) + * This is another CuTeDSL/CUDA toolchain gap: PyTorch's JIT flags + * are incompatible with raw BF16 types in CUDA 13.2. The fix is to + * compile the .cu separately with nvcc and load as a shared library, + * or replace bf16_t with c10::BFloat16 and use AT_DISPATCH types. + * + * For now, use the standalone test (test_fmha_sm100_standalone.cu) * which compiles with nvcc directly and tests the kernel via CUDA * runtime APIs (no PyTorch needed). * - * To fix for production: Replace bf16_t with c10::BFloat16 and use - * AT_DISPATCH_FLOATING_TYPES for type dispatch. Or compile the .cu - * separately with nvcc and load as a shared library. + * ================================================================== + * PRODUCTION PATH + * ================================================================== + * 1. Compile all .cu kernels with nvcc into a shared library + * 2. Load via torch.utils.cpp_extension.load with precompiled=True + * 3. Or use ctypes/cupy to call cudaLaunchKernel directly + * 4. The c10::BFloat16 approach works but requires ATen headers */ #include "fmha_sm100.cuh" +#include "fmha_epilogue_sm100.cuh" #include #include