diff --git a/NVFP4-1.1_RAW_CUDA_CPP_APPROACH.md b/NVFP4-1.1_RAW_CUDA_CPP_APPROACH.md index d4c0bf22..703782db 100644 --- a/NVFP4-1.1_RAW_CUDA_CPP_APPROACH.md +++ b/NVFP4-1.1_RAW_CUDA_CPP_APPROACH.md @@ -77,4 +77,52 @@ The kernel itself compiles and is valid SM100 PTX. The launch wrapper needs to b The fix for the JIT is to not use `__bf16` when `__CUDA_NO_BFLOAT16_CONVERSIONS__` is defined, and use `unsigned short` instead. But honestly, the JIT approach isn't the right path — we should compile the kernel directly with nvcc and load it as a shared library. -Let me update the test to compile with nvcc directly and test the kernel with a proper CUDA launch:nvcc uses `-Xcompiler -fPIC` instead of `-fPIC` directly: \ No newline at end of file +Let me update the test to compile with nvcc directly and test the kernel with a proper CUDA launch:nvcc uses `-Xcompiler -fPIC` instead of `-fPIC` directly:The `Internal Compiler Error (codegen): "unsupported float variant!"` is back! Even with direct nvcc compilation (no PyTorch JIT). This is coming from the `bf16_t` (`__bf16`) type when used in the kernel. The `o0 *= rescale` line is multiplying a float by a float — that shouldn't be the issue. Let me check if the `__bf16` type is causing the ICE when used as function parameters:The `__bf16` type is used as kernel parameters and SMEM buffer type. The ICE is from nvcc's codegen when it encounters `__bf16` in certain contexts on SM100. This is likely a CUDA 13.2 compiler bug. + +The fix: use `unsigned short` (which is the underlying storage for BF16) everywhere except when doing BF16 operations (conversions):Wait, that removes the BF16 type entirely. I need `uint16_t` for storage and inline asm for conversions. Let me rewrite the conversion functions to work with `uint16_t`:The sed command corrupted the typedef. `typedef __bf16 uint16_t` is wrong. Let me fix this properly:Also fix the sed damage to the launch wrapper:Those `uint16_t*` casts are correct since `bf16_t` is `unsigned short`. The `sizeof(bf16_t)` will be 2 which is correct. Let me also check for `uint16_to_f32` references that were incorrectly renamed:Excellent progress! The kernel compiles to PTX but the PTX assembler (`ptxas`) rejects the inline asm because I have the `tcgen05.mma` instruction syntax wrong. The arguments are mismatched. This is the hard part — getting the exact PTX syntax right for the UMMA operations. + +Let me look at the exact PTX syntax from the CUTLASS reference:Now I can see the EXACT PTX syntax: + +**QK (S→S, both SMEM):** +``` +tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, {%5, %6, %7, %8}, p; +`` +Args: tmem_c (r), desc_a (l), desc_b (l), idescE_hi32 (r), scaleC (r), mask[0-3] (r) + +**PV (T→S, TMEM A + SMEM B):** +``` +tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p; +``` +Args: tmem_c (r), tmem_a (r), desc_b (l), idescE_hi32 (r), scaleC (r), mask[0-3] (r) + +The key differences from my code: +1. `idescE` is passed as the upper 32 bits (`idescE>>32`), not the full 64-bit value +2. `scaleC` is a float, not an integer +3. There's a `mask` array of 4 uint32 values +4. The constraint string is different + +Let me fix the QK GEMM in the kernel:Multiple PTX syntax errors. The inline asm PTX I wrote doesn't match the actual PTX ISA syntax. I need to get the exact PTX instructions right. Let me study the CUTLASS C++ implementation more carefully to extract the EXACT PTX:Key insight: the TMEM alloc takes `[%0]` (output register) and `%1` (num_columns), but the output is the *pointer* in the `[%0]`. Let me check — actually no, the format is `"tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"` where `%0` is the SMEM pointer (input) and `%1` is the num columns (input). There's no output — TMEM is allocated at the specified SMEM address. + +Wait, looking more carefully at the CUTLASS code: +```cpp +uint32_t dst_intptr = cute::cast_smem_ptr_to_uint(dst_ptr); +asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(dst_intptr), "r"(num_columns)); +``` + +So `%0` is the SMEM pointer where TMEM will be mapped, and `%1` is the number of columns. No output register — TMEM is implicitly accessible. + +This is different from what I assumed. Let me study the MMA PTX more carefully. The complexity here is real — there are many PTX instructions with specific syntax. Let me take a step back and focus on getting the kernel to compile with the TMEM alloc/dealloc working, then add MMA later. + +Let me simplify the kernel to just TMEM alloc + dealloc + a simple SMEM→GMEM output (no MMA yet), verify it compiles and runs, then incrementally add the MMA:The launch wrapper still references `fmha_decode` but the kernel is now `fmha_decode_ref`. Let me update the launcher:Also the template params changed — no more CAUSAL/SINK template params for the ref version:Also fix the template macro — the ref kernel only takes HD, not (D,C,S):**✅ COMPILATION SUCCESSFUL!** The kernel compiled to a shared library. The runtime error is just missing PyTorch symbols for the pybind11 binding — the CUDA kernel itself compiled fine. + +This is real, tangible progress. The FMHA kernel compiles on Blackwell SM100. Now I need to write a proper test that doesn't need pybind11 — just launch the kernel directly using CUDA runtime APIs from Python via `cupy` or raw `ctypes` with `cudaLaunchKernel`. + +But the P@V accumulation in the reference kernel has a bug — multiple threads computing different KV columns but the same output element need a reduction. Let me fix the kernel's P@V logic properly, then write a standalone C test that doesn't need PyTorch bindings:Already committed. Let me now write a standalone CUDA test that launches the kernel directly, without PyTorch bindings:Let me just rewrite the test file:The include path needs to reference the full path from the repo root. The include is `"fmha_sm100.cuh"` but the file is at `dsv4/kernels/attention/fmha_sm100.cuh`. With `-I{REPO}`, I need to use `#include "dsv4/kernels/attention/fmha_sm100.cuh"`:**MASSIVE PROGRESS!** The kernel compiles, launches, and runs on the B200! The output is all zeros, which means the kernel logic has a bug — but the infrastructure works. The GPU kernel executes without crashing. + +The bug is clear: the P@V accumulation is broken (I noted this earlier — the thread-to-column mapping for the reduction is wrong). Let me fix the kernel logic. The issue is: + +1. Q is loaded by all 192 threads but only `tid < HD` loads happen +2. The KV loop processes columns in parallel across threads, but the P@V accumulation needs a proper parallel reduction +3. The row_max computation across warps is broken + +Let me fix the kernel with a simpler, correct approach: \ No newline at end of file diff --git a/dsv4/kernels/attention/fmha_sm100.cuh b/dsv4/kernels/attention/fmha_sm100.cuh index da3e3f2c..fd327228 100644 --- a/dsv4/kernels/attention/fmha_sm100.cuh +++ b/dsv4/kernels/attention/fmha_sm100.cuh @@ -126,13 +126,13 @@ fmha_decode_ref( // Online softmax with O rescale float new_max = fmaxf(row_max, s_val); if (new_max > row_max) { - float rescale = exp2f((row_max - new_max) * scale * 1.4426950408889634f); + float rescale = expf(row_max - new_max); for (int d = 0; d < HD; d++) o_acc[d] *= rescale; row_sum *= rescale; row_max = new_max; } - float p_val = exp2f((s_val - row_max) * scale * 1.4426950408889634f); + float p_val = expf(s_val - row_max); row_sum += p_val; for (int d = 0; d < HD; d++) { @@ -154,7 +154,7 @@ fmha_decode_ref( // LSE if (lse_out && tid == 0) { - lse_out[batch * gridDim.y + head] = logf(row_sum) + row_max * 1.4426950408889634f; + lse_out[batch * gridDim.y + head] = logf(row_sum) + row_max; } }