biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: define bf16_t using __bf16 built-in, avoid cuda_bf16.h bug

Browse Source
This commit is contained in:
biondizzle
2026-05-28 05:14:01 +00:00
parent f0660d0bd7
commit 88d5995ec9
2 changed files with 51 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
dsv4/kernels/attention/fmha_sm100.cpp
View File

@@ -66,10 +66,10 @@ std::tuple<torch::Tensor, torch::Tensor> fmha_decode_cuda(
// Launch kernel (head_dim template specialization)
#define LAUNCH_FMHA(HD, CAUSAL, SINK) \
fmha_decode_kernel<HD, 1, CAUSAL, SINK><<<grid, block, smem_bytes>>>( \
reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>()), \
reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>()), \
reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>()), \
reinterpret_cast<__nv_bfloat16*>(o.data_ptr<at::BFloat16>()), \
reinterpret_cast<const bf16_t*>(q.data_ptr<at::BFloat16>()), \
reinterpret_cast<const bf16_t*>(k.data_ptr<at::BFloat16>()), \
reinterpret_cast<const bf16_t*>(v.data_ptr<at::BFloat16>()), \
reinterpret_cast<bf16_t*>(o.data_ptr<at::BFloat16>()), \
batch_stride_q, batch_stride_kv, batch_stride_o, \
s_k_val, n_comp, swa_len, (float)scale_softmax, \
sink_ptr, lse.data_ptr<float>() \

68
dsv4/kernels/attention/fmha_sm100.cuh
View File

@@ -30,9 +30,35 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>
// 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 <cutlass/cutlass.h>
@@ -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<bf16_t*>(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<bf16_t*>(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<bf16_t*>(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<bf16_t*>(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);
}
}
}