fix: avoid cuda_bf16.h, use inline PTX for BF16 conversion

This commit is contained in:
2026-05-28 05:12:08 +00:00
parent 1734d13f60
commit a64e55665b

View File

@@ -30,9 +30,25 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>
// NOTE: cuda_bf16.h is NOT included due to CUDA 13.2 C++17 compatibility bug
// (in_place_from error). __nv_bfloat16 is a built-in CUDA type and works
// without the header. We define helper functions manually.
// BF16 conversion helpers
__device__ __forceinline__ __nv_bfloat16 float_to_bfloat16(float f) {
__nv_bfloat16 h;
asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f));
return h;
}
__device__ __forceinline__ float bfloat16_to_float(__nv_bfloat16 h) {
float f;
asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h));
return f;
}
// CUTLASS C++ includes (CUDA device code only)
#if defined(__CUDA_ARCH__)
#include <cutlass/cutlass.h>
@@ -559,7 +575,7 @@ fmha_decode_kernel(
// Simplified: directly compute and write
// (Proper TMEM load + normalize below)
float o_val = 0.0f; // placeholder
o_head[j] = __float2bfloat16(o_val / row_sum);
o_head[j] = float_to_bfloat16(o_val / row_sum);
}
}
}