File size: 4,671 Bytes
88b8ef2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | // ============================================================================
// bf16-input vectorized fused FP4 quantize + CUTLASS SFA/SFB scales.
//
// Same per-block scale selection and e2m1 rounding as
// quantize_fp4_dynamic_sfa_fp16_vec, with bf16 source elements. Each
// thread quantizes one 16-element block: two 16-byte loads, one 8-byte
// packed store, one SFA byte at the tile-interleaved offset.
// ============================================================================
#include "quantize_fp4_sfa_bf16.cuh"
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
# include "cutlass/cutlass.h"
# include "cutlass/detail/sm100_blockscaled_layout.hpp"
# include "cute/tensor.hpp"
# define FV_HAVE_CUTLASS 1
#else
# define FV_HAVE_CUTLASS 0
#endif
namespace flash_rt {
namespace fp4 {
#if FV_HAVE_CUTLASS
namespace {
using CfgVecB = cutlass::detail::Sm1xxBlockScaledConfig<16>;
__device__ __forceinline__ uint8_t fp32_to_e2m1_bvec(float x) {
uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
float ax = fabsf(x);
uint8_t mant;
if (ax <= 0.25f) mant = 0u;
else if (ax <= 0.75f) mant = 1u;
else if (ax <= 1.25f) mant = 2u;
else if (ax <= 1.75f) mant = 3u;
else if (ax <= 2.5f) mant = 4u;
else if (ax <= 3.5f) mant = 5u;
else if (ax <= 5.0f) mant = 6u;
else mant = 7u;
return sign | mant;
}
template <class LayoutSF>
__global__ void kernel_quantize_fp4_sfa_bf16_vec(
const int4* __restrict__ src, // bf16 [N, D] as int4 (8 elements)
uint2* __restrict__ dst_packed, // [N, D/2] bytes as uint2 (1 block)
uint8_t* __restrict__ dst_sfa,
LayoutSF layout,
int N, int D8) { // D8 = D / 8 int4 chunks per row
const int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
const int row = blockIdx.y;
const int n_blocks = D8 >> 1; // 16 elements per block
if (row >= N || block_idx >= n_blocks) return;
const int4 raw0 = src[row * D8 + 2 * block_idx];
const int4 raw1 = src[row * D8 + 2 * block_idx + 1];
const __nv_bfloat16* h0 = reinterpret_cast<const __nv_bfloat16*>(&raw0);
const __nv_bfloat16* h1 = reinterpret_cast<const __nv_bfloat16*>(&raw1);
float vals[16];
float amax = 0.f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
vals[i] = __bfloat162float(h0[i]);
vals[8 + i] = __bfloat162float(h1[i]);
}
#pragma unroll
for (int i = 0; i < 16; ++i) {
const float a = fabsf(vals[i]);
if (a > amax) amax = a;
}
float desired = amax / 6.f;
if (desired < 1e-12f) desired = 1e-12f;
__nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
const float bs_dq = static_cast<float>(bs_q);
dst_sfa[layout(row, block_idx * 16, 0)] =
*reinterpret_cast<uint8_t*>(&bs_q);
const float inv_bs = 1.f / bs_dq;
uint2 out;
uint8_t* ob = reinterpret_cast<uint8_t*>(&out);
#pragma unroll
for (int p = 0; p < 8; ++p) {
const uint8_t lo = fp32_to_e2m1_bvec(vals[2 * p] * inv_bs);
const uint8_t hi = fp32_to_e2m1_bvec(vals[2 * p + 1] * inv_bs);
ob[p] = static_cast<uint8_t>(lo | (hi << 4));
}
dst_packed[row * n_blocks + block_idx] = out;
}
} // namespace
#endif // FV_HAVE_CUTLASS
int quantize_fp4_dynamic_sfa_bf16_vec(
const void* src_bf16, void* dst_packed, void* dst_sfa,
int N, int D, bool is_sfb, cudaStream_t stream) {
#if FV_HAVE_CUTLASS
if (D % 16 != 0) return -1;
if ((reinterpret_cast<uintptr_t>(src_bf16) & 15) ||
(reinterpret_cast<uintptr_t>(dst_packed) & 7)) return -1;
const int n_blocks = D / 16;
const int threads = 128;
dim3 grid((n_blocks + threads - 1) / threads, N);
auto shape = cute::make_shape(
is_sfb ? 1 : N,
is_sfb ? N : 1,
D, 1);
if (is_sfb) {
auto layout = CfgVecB::tile_atom_to_shape_SFB(shape);
kernel_quantize_fp4_sfa_bf16_vec<<<grid, threads, 0, stream>>>(
reinterpret_cast<const int4*>(src_bf16),
reinterpret_cast<uint2*>(dst_packed),
reinterpret_cast<uint8_t*>(dst_sfa),
layout, N, D >> 3);
} else {
auto layout = CfgVecB::tile_atom_to_shape_SFA(shape);
kernel_quantize_fp4_sfa_bf16_vec<<<grid, threads, 0, stream>>>(
reinterpret_cast<const int4*>(src_bf16),
reinterpret_cast<uint2*>(dst_packed),
reinterpret_cast<uint8_t*>(dst_sfa),
layout, N, D >> 3);
}
const cudaError_t e = cudaGetLastError();
return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
#else
(void)src_bf16; (void)dst_packed; (void)dst_sfa;
(void)N; (void)D; (void)is_sfb; (void)stream;
return -2;
#endif
}
} // namespace fp4
} // namespace flash_rt
|