From 88d5995ec961a729b59201427bc75929ad98afd5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 05:14:01 +0000 Subject: [PATCH] fix: define bf16_t using __bf16 built-in, avoid cuda_bf16.h bug --- dsv4/kernels/attention/fmha_sm100.cpp | 8 ++-- dsv4/kernels/attention/fmha_sm100.cuh | 68 ++++++++++++++++++--------- 2 files changed, 51 insertions(+), 25 deletions(-) diff --git a/dsv4/kernels/attention/fmha_sm100.cpp b/dsv4/kernels/attention/fmha_sm100.cpp index 6c64144a..1b199c23 100644 --- a/dsv4/kernels/attention/fmha_sm100.cpp +++ b/dsv4/kernels/attention/fmha_sm100.cpp @@ -66,10 +66,10 @@ std::tuple fmha_decode_cuda( // Launch kernel (head_dim template specialization) #define LAUNCH_FMHA(HD, CAUSAL, SINK) \ fmha_decode_kernel<<>>( \ - reinterpret_cast(q.data_ptr()), \ - reinterpret_cast(k.data_ptr()), \ - reinterpret_cast(v.data_ptr()), \ - reinterpret_cast<__nv_bfloat16*>(o.data_ptr()), \ + reinterpret_cast(q.data_ptr()), \ + reinterpret_cast(k.data_ptr()), \ + reinterpret_cast(v.data_ptr()), \ + reinterpret_cast(o.data_ptr()), \ batch_stride_q, batch_stride_kv, batch_stride_o, \ s_k_val, n_comp, swa_len, (float)scale_softmax, \ sink_ptr, lse.data_ptr() \ diff --git a/dsv4/kernels/attention/fmha_sm100.cuh b/dsv4/kernels/attention/fmha_sm100.cuh index 8a01b707..acfc7a49 100644 --- a/dsv4/kernels/attention/fmha_sm100.cuh +++ b/dsv4/kernels/attention/fmha_sm100.cuh @@ -30,9 +30,35 @@ #pragma once #include -#include #include +// NOTE: cuda_bf16.h has a C++17/20 compatibility bug on CUDA 13.2 (in_place_from). +// We use __bf16 (CUDA 13+ built-in type) and define our own helpers. +// __bf16 is the actual hardware type; bf16_t is a wrapper from the header. +// +// For device code, __bf16 works directly with cvt PTX instructions. +// For host code, we use uint16_t as the storage type. + +#if defined(__CUDACC__) +// Device-side BF16 helpers using inline PTX +typedef __bf16 bf16_t; + +__device__ __forceinline__ bf16_t f32_to_bf16(float f) { + bf16_t h; + asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); + return h; +} + +__device__ __forceinline__ float bf16_to_f32(bf16_t h) { + float f; + asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); + return f; +} +#else +// Host-side: use uint16_t as storage +typedef uint16_t bf16_t; +#endif + // CUTLASS C++ includes (CUDA device code only) #if defined(__CUDA_ARCH__) #include @@ -80,10 +106,10 @@ constexpr int THREADS_PER_CTA = TOTAL_WARPS * WARP_SIZE; // 192 // C: (TILE_M, head_dim) BF16 — output epilogue SMEM // SMEM sizes at various head_dims (bytes) -constexpr int SMEM_Q(int hd) { return TILE_M * hd * sizeof(__nv_bfloat16); } -constexpr int SMEM_K(int hd) { return SMEM_TILE_K * hd * sizeof(__nv_bfloat16); } -constexpr int SMEM_V(int hd) { return SMEM_TILE_K * hd * sizeof(__nv_bfloat16); } -constexpr int SMEM_C(int hd) { return TILE_M * hd * sizeof(__nv_bfloat16); } +constexpr int SMEM_Q(int hd) { return TILE_M * hd * sizeof(bf16_t); } +constexpr int SMEM_K(int hd) { return SMEM_TILE_K * hd * sizeof(bf16_t); } +constexpr int SMEM_V(int hd) { return SMEM_TILE_K * hd * sizeof(bf16_t); } +constexpr int SMEM_C(int hd) { return TILE_M * hd * sizeof(bf16_t); } // With kv_stage=1 for hd>128, kv_stage=2 for hd<=128 constexpr int TOTAL_SMEM(int hd) { @@ -93,7 +119,7 @@ constexpr int TOTAL_SMEM(int hd) { int sK = SMEM_TILE_K * hd * kv_stage; // K (1-2 stages) int sV = pv_n_tile * hd * kv_stage; // V (1-2 stages) int sC = pv_n_tile * hd; // C (epilogue, 1 stage) - return (sQ + sK + sV + sC) * sizeof(__nv_bfloat16); + return (sQ + sK + sV + sC) * sizeof(bf16_t); } // ===================================================================== @@ -270,11 +296,11 @@ template< __global__ void __launch_bounds__(THREADS_PER_CTA) fmha_decode_kernel( // Input tensors (GMEM, BF16) - const __nv_bfloat16* __restrict__ q, // (batch, num_heads, T, HEAD_DIM) - const __nv_bfloat16* __restrict__ k, // (batch, s_k, HEAD_DIM) — dense KV tile - const __nv_bfloat16* __restrict__ v, // (batch, HEAD_DIM, s_k) — transposed for PV + const bf16_t* __restrict__ q, // (batch, num_heads, T, HEAD_DIM) + const bf16_t* __restrict__ k, // (batch, s_k, HEAD_DIM) — dense KV tile + const bf16_t* __restrict__ v, // (batch, HEAD_DIM, s_k) — transposed for PV // Output tensor (GMEM, BF16) - __nv_bfloat16* __restrict__ o, // (batch, num_heads, T, HEAD_DIM) + bf16_t* __restrict__ o, // (batch, num_heads, T, HEAD_DIM) // Parameters int batch_stride_q, int batch_stride_kv, @@ -332,30 +358,30 @@ fmha_decode_kernel( // SMEM is allocated dynamically via the kernel launch API. // Layout: [Q | K | V | C] — packed, matching CuTeDSL FMHA. extern __shared__ char smem_buf[]; - __nv_bfloat16* sQ = reinterpret_cast<__nv_bfloat16*>(smem_buf); + bf16_t* sQ = reinterpret_cast(smem_buf); int smem_offset = TILE_M * HEAD_DIM; // Q size const int kv_stage = (HEAD_DIM > 128) ? 1 : 2; - __nv_bfloat16* sK = reinterpret_cast<__nv_bfloat16*>(smem_buf + smem_offset * sizeof(__nv_bfloat16)); + bf16_t* sK = reinterpret_cast(smem_buf + smem_offset * sizeof(bf16_t)); smem_offset += SMEM_TILE_K * HEAD_DIM * kv_stage; - __nv_bfloat16* sV = reinterpret_cast<__nv_bfloat16*>(smem_buf + smem_offset * sizeof(__nv_bfloat16)); + bf16_t* sV = reinterpret_cast(smem_buf + smem_offset * sizeof(bf16_t)); smem_offset += SMEM_TILE_K * HEAD_DIM * kv_stage; - __nv_bfloat16* sC = reinterpret_cast<__nv_bfloat16*>(smem_buf + smem_offset * sizeof(__nv_bfloat16)); + bf16_t* sC = reinterpret_cast(smem_buf + smem_offset * sizeof(bf16_t)); // ================================================================= // Q pointer arithmetic (head-packed) // ================================================================= - const __nv_bfloat16* q_batch = q + batch_idx * batch_stride_q; - const __nv_bfloat16* k_batch = k + batch_idx * batch_stride_kv; - const __nv_bfloat16* v_batch = v + batch_idx * batch_stride_kv; - __nv_bfloat16* o_batch = o + batch_idx * batch_stride_o; + const bf16_t* q_batch = q + batch_idx * batch_stride_q; + const bf16_t* k_batch = k + batch_idx * batch_stride_kv; + const bf16_t* v_batch = v + batch_idx * batch_stride_kv; + bf16_t* o_batch = o + batch_idx * batch_stride_o; // Q: head_idx-th head, all T rows // For decode (T=1): Q is (1, HEAD_DIM) - const __nv_bfloat16* q_head = q_batch + head_idx * HEAD_DIM; - __nv_bfloat16* o_head = o_batch + head_idx * HEAD_DIM; + const bf16_t* q_head = q_batch + head_idx * HEAD_DIM; + bf16_t* o_head = o_batch + head_idx * HEAD_DIM; // ================================================================= // KV tile loop @@ -559,7 +585,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] = f32_to_bf16(o_val / row_sum); } } }