fix: use expf for softmax (not exp2f with scale)

This commit is contained in:
2026-05-28 05:34:03 +00:00
parent 3cb339129b
commit 8eb735618f
2 changed files with 52 additions and 4 deletions

View File

@@ -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:
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:

View File

@@ -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;
}
}