docs: update here-docs with CuTeDSL rationale for NVIDIA
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
This commit is contained in:
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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 <ATen/ATen.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
|
||||
Reference in New Issue
Block a user