fix: use expf for softmax (not exp2f with scale)
This commit is contained in:
@@ -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:
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user