fix: avoid cuda_bf16.h, use inline PTX for BF16 conversion
This commit is contained in:
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user