ViperEkura
057c0d33df
- split quantize into quantize.cuh, templated on input type (bf16/fp16/fp32) - rename pybind entry quantize_bf16 to quantize; validate the fmt enum - fix fp8x2 packing: one 32-bit word packs two pairs (halves were dropped) - drop the dead OutFp8 template param; GEMM output is always bf16 - fp8_state.reset() restores recipe/format defaults too (test state leak) - rewrite tests for the two-primitive API with fp32-domain amax references
309 lines
12 KiB
Plaintext
309 lines
12 KiB
Plaintext
/*
|
|
FP8 family tests: single-warp MMA demo + full GEMM correctness.
|
|
|
|
Part 1 exercises one bf16 -> fp8 -> mma.sync m16n8k32 instruction pair
|
|
(sanity for astrai::mma_sync + the fragment layout contract).
|
|
Part 2 checks launch_fp8_gemm across all four operand layouts, both K
|
|
tiles, and ragged shapes against an fp32 CPU reference.
|
|
|
|
nvcc -I csrc -arch=sm_89 -std=c++17 -O3 csrc/tests/fp8_test.cu -o /tmp/fp8_test \
|
|
&& /tmp/fp8_test
|
|
*/
|
|
|
|
#include "test_utils.cuh"
|
|
|
|
#include <cuda_fp8.h>
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cuda_runtime.h>
|
|
#include <type_traits>
|
|
#include <vector>
|
|
|
|
#include "../kernels/common/mma.cuh"
|
|
#include "../kernels/fp8/gemm.cuh"
|
|
|
|
using namespace astrai::fp8;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Part 1: single-kernel BF16 -> FP8 MMA -> BF16 demo (m16n8k32)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
namespace {
|
|
|
|
constexpr int kMmaM = 16;
|
|
constexpr int kMmaN = 8;
|
|
constexpr int kMmaK = 32;
|
|
|
|
__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2,
|
|
float x3) {
|
|
__nv_fp8_e4m3 q0(x0);
|
|
__nv_fp8_e4m3 q1(x1);
|
|
__nv_fp8_e4m3 q2(x2);
|
|
__nv_fp8_e4m3 q3(x3);
|
|
return static_cast<unsigned>(q0.__x) |
|
|
(static_cast<unsigned>(q1.__x) << 8) |
|
|
(static_cast<unsigned>(q2.__x) << 16) |
|
|
(static_cast<unsigned>(q3.__x) << 24);
|
|
}
|
|
|
|
__device__ __forceinline__ unsigned load_quantize_fp8x4(
|
|
const bf16* src, float scale_inv) {
|
|
return pack_fp8x4(__bfloat162float(src[0]) * scale_inv,
|
|
__bfloat162float(src[1]) * scale_inv,
|
|
__bfloat162float(src[2]) * scale_inv,
|
|
__bfloat162float(src[3]) * scale_inv);
|
|
}
|
|
|
|
__global__ void fused_bf16_fp8_mma_kernel(
|
|
const bf16* __restrict__ a, const bf16* __restrict__ b,
|
|
bf16* __restrict__ out, float scale_a, float scale_b) {
|
|
const int lane = threadIdx.x;
|
|
const int group = lane >> 2;
|
|
const int thread_in_group = lane & 3;
|
|
const int k0 = thread_in_group * 4;
|
|
|
|
// PTX m16n8k32 A fragment: two rows, two 16-column K partitions.
|
|
unsigned a_frag[4];
|
|
a_frag[0] = load_quantize_fp8x4(&a[group * kMmaK + k0], 1.0f / scale_a);
|
|
a_frag[1] =
|
|
load_quantize_fp8x4(&a[(group + 8) * kMmaK + k0], 1.0f / scale_a);
|
|
a_frag[2] =
|
|
load_quantize_fp8x4(&a[group * kMmaK + k0 + 16], 1.0f / scale_a);
|
|
a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * kMmaK + k0 + 16],
|
|
1.0f / scale_a);
|
|
|
|
// B is supplied as row-major [N,K], equivalent to the col-major [K,N]
|
|
// operand required by the MMA instruction.
|
|
unsigned b_frag[2];
|
|
b_frag[0] = load_quantize_fp8x4(&b[group * kMmaK + k0], 1.0f / scale_b);
|
|
b_frag[1] =
|
|
load_quantize_fp8x4(&b[group * kMmaK + k0 + 16], 1.0f / scale_b);
|
|
|
|
float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
|
astrai::mma_sync<__nv_fp8_e4m3>(acc, a_frag, b_frag, acc);
|
|
|
|
const int col = thread_in_group * 2;
|
|
const float output_scale = scale_a * scale_b;
|
|
*reinterpret_cast<__nv_bfloat162*>(&out[group * kMmaN + col]) =
|
|
__floats2bfloat162_rn(acc[0] * output_scale, acc[1] * output_scale);
|
|
*reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * kMmaN + col]) =
|
|
__floats2bfloat162_rn(acc[2] * output_scale, acc[3] * output_scale);
|
|
}
|
|
|
|
static float quantize_e4m3(float value) {
|
|
return static_cast<float>(__nv_fp8_e4m3(value));
|
|
}
|
|
|
|
static bool test_single_mma() {
|
|
srand(0);
|
|
std::vector<float> a(kMmaM * kMmaK), b(kMmaN * kMmaK),
|
|
reference(kMmaM * kMmaN, 0.0f);
|
|
std::vector<bf16> a_bf16(kMmaM * kMmaK), b_bf16(kMmaN * kMmaK),
|
|
output(kMmaM * kMmaN);
|
|
for (float& value : a) value = randf() * 4.0f;
|
|
for (float& value : b) value = randf() * 4.0f;
|
|
for (int i = 0; i < kMmaM * kMmaK; ++i) {
|
|
a_bf16[i] = f2bf(a[i]);
|
|
a[i] = bf2f(a_bf16[i]);
|
|
}
|
|
for (int i = 0; i < kMmaN * kMmaK; ++i) {
|
|
b_bf16[i] = f2bf(b[i]);
|
|
b[i] = bf2f(b_bf16[i]);
|
|
}
|
|
|
|
const float amax = *std::max_element(
|
|
a.begin(), a.end(),
|
|
[](float x, float y) { return fabsf(x) < fabsf(y); });
|
|
const float bmax = *std::max_element(
|
|
b.begin(), b.end(),
|
|
[](float x, float y) { return fabsf(x) < fabsf(y); });
|
|
const float scale_a = fabsf(amax) / 448.0f;
|
|
const float scale_b = fabsf(bmax) / 448.0f;
|
|
|
|
for (int row = 0; row < kMmaM; ++row) {
|
|
for (int col = 0; col < kMmaN; ++col) {
|
|
float sum = 0.0f;
|
|
for (int k = 0; k < kMmaK; ++k) {
|
|
float qa = quantize_e4m3(a[row * kMmaK + k] / scale_a);
|
|
float qb = quantize_e4m3(b[col * kMmaK + k] / scale_b);
|
|
sum = fmaf(qa, qb, sum);
|
|
}
|
|
reference[row * kMmaN + col] = sum * scale_a * scale_b;
|
|
}
|
|
}
|
|
|
|
bf16 *d_a, *d_b, *d_out;
|
|
CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16)));
|
|
CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16)));
|
|
CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16)));
|
|
CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16),
|
|
cudaMemcpyHostToDevice));
|
|
CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16),
|
|
cudaMemcpyHostToDevice));
|
|
|
|
fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b);
|
|
CUDA_CHECK(cudaDeviceSynchronize());
|
|
CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16),
|
|
cudaMemcpyDeviceToHost));
|
|
|
|
float max_abs_error = 0.0f;
|
|
float max_rel_error = 0.0f;
|
|
for (int i = 0; i < kMmaM * kMmaN; ++i) {
|
|
float error = fabsf(bf2f(output[i]) - reference[i]);
|
|
max_abs_error = fmaxf(max_abs_error, error);
|
|
max_rel_error = fmaxf(
|
|
max_rel_error, error / fmaxf(fabsf(reference[i]), 1e-4f));
|
|
}
|
|
const bool pass = max_abs_error < 0.05f;
|
|
print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error,
|
|
max_rel_error, pass);
|
|
|
|
cudaFree(d_a);
|
|
cudaFree(d_b);
|
|
cudaFree(d_out);
|
|
return pass;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Part 2: GEMM correctness — layouts x K-tiles vs fp32 CPU reference
|
|
// ---------------------------------------------------------------------------
|
|
|
|
template <typename LA, typename LB, int kK, int Stages>
|
|
static bool run_gemm_case(const float* ha, const float* hb, int m, int n,
|
|
int k, int a_ld, int b_ld) {
|
|
__nv_fp8_e4m3 *da, *db;
|
|
__nv_bfloat16* dout;
|
|
float* dscale;
|
|
cudaMalloc(&da, (size_t)m * k);
|
|
cudaMalloc(&db, (size_t)n * k);
|
|
cudaMalloc(&dout, (size_t)m * n * 2);
|
|
cudaMalloc(&dscale, 4);
|
|
float one = 1.0f;
|
|
cudaMemcpy(dscale, &one, 4, cudaMemcpyHostToDevice);
|
|
// quantize inputs to e4m3 on host and upload byte-by-byte
|
|
std::vector<unsigned char> qa(m * k), qb(n * k);
|
|
for (int i = 0; i < m * k; ++i) {
|
|
__nv_fp8_e4m3 q(ha[i]);
|
|
qa[i] = *(unsigned char*)&q;
|
|
}
|
|
for (int i = 0; i < n * k; ++i) {
|
|
__nv_fp8_e4m3 q(hb[i]);
|
|
qb[i] = *(unsigned char*)&q;
|
|
}
|
|
cudaMemcpy(da, qa.data(), qa.size(), cudaMemcpyHostToDevice);
|
|
cudaMemcpy(db, qb.data(), qb.size(), cudaMemcpyHostToDevice);
|
|
|
|
FP8Params p = {};
|
|
p.a_ptr = da;
|
|
p.b_ptr = db;
|
|
p.out_ptr = dout;
|
|
p.scale = dscale;
|
|
p.m = m;
|
|
p.n = n;
|
|
p.k = k;
|
|
p.a_ld = a_ld;
|
|
p.b_ld = b_ld;
|
|
launch_fp8_gemm<FP8Format::E4M3, LA, LB, kK, Stages>(p, 0);
|
|
cudaError_t e = cudaDeviceSynchronize();
|
|
if (e != cudaSuccess) {
|
|
printf(" CUDA err: %s\n", cudaGetErrorString(e));
|
|
return false;
|
|
}
|
|
std::vector<unsigned short> hb16(m * n);
|
|
cudaMemcpy(hb16.data(), dout, (size_t)m * n * 2, cudaMemcpyDeviceToHost);
|
|
const float tol = 0.06f;
|
|
double max_rel = 0;
|
|
bool ok = true;
|
|
for (int i = 0; i < m && ok; ++i) {
|
|
for (int j = 0; j < n && ok; ++j) {
|
|
float ref = 0;
|
|
for (int kk = 0; kk < k; ++kk) {
|
|
// A reference reads the actual uploaded buffer: LA ColMajor
|
|
// means the buffer is [K][M] (ha_t), else [M][K].
|
|
float av = std::is_same_v<LA, ColMajor>
|
|
? (float)__nv_fp8_e4m3(ha[kk * m + i])
|
|
: (float)__nv_fp8_e4m3(ha[i * k + kk]);
|
|
float bv;
|
|
if (std::is_same_v<LB, ColMajor>)
|
|
bv = (float)__nv_fp8_e4m3(hb[j * k + kk]);
|
|
else
|
|
bv = (float)__nv_fp8_e4m3(hb[kk * n + j]);
|
|
ref += av * bv;
|
|
}
|
|
float got =
|
|
__bfloat162float(__ushort_as_bfloat16(hb16[i * n + j]));
|
|
float err = fabsf(got - ref);
|
|
float rel = err / fmaxf(fabsf(ref), 0.5f);
|
|
if (rel > max_rel) max_rel = rel;
|
|
if (err > tol * fmaxf(fabsf(ref), 1.0f)) ok = false;
|
|
}
|
|
}
|
|
printf(" max_rel=%.4f %s\n", max_rel, ok ? "PASS" : "FAIL");
|
|
cudaFree(da);
|
|
cudaFree(db);
|
|
cudaFree(dout);
|
|
cudaFree(dscale);
|
|
return ok;
|
|
}
|
|
|
|
static bool test_gemm() {
|
|
struct {
|
|
int m, n, k;
|
|
} cfgs[] = {
|
|
{128, 128, 128}, {256, 128, 256}, {128, 256, 64},
|
|
{100, 130, 96}, {64, 64, 160}, {300, 200, 320},
|
|
};
|
|
bool all = true;
|
|
for (auto& c : cfgs) {
|
|
float* ha = new float[c.m * c.k];
|
|
float* hb_rowmajor = new float[c.k * c.n]; // [K][N] for B RowMajor
|
|
float* hb_colmajor = new float[c.n * c.k]; // [N][K] for B ColMajor
|
|
for (int i = 0; i < c.m * c.k; ++i) ha[i] = randf();
|
|
for (int i = 0; i < c.k * c.n; ++i) hb_rowmajor[i] = randf();
|
|
for (int i = 0; i < c.k * c.n; ++i)
|
|
hb_colmajor[i / c.k * c.k + i % c.k] = hb_rowmajor[i];
|
|
float* ha_t = new float[c.k * c.m]; // [K][M] for A ColMajor
|
|
for (int i = 0; i < c.m; ++i)
|
|
for (int p = 0; p < c.k; ++p) ha_t[p * c.m + i] = ha[i * c.k + p];
|
|
printf("%dx%dx%d:\n", c.m, c.n, c.k);
|
|
printf(" NT K32:");
|
|
all &= run_gemm_case<RowMajor, ColMajor, 32, 3>(ha, hb_colmajor, c.m,
|
|
c.n, c.k, c.k, c.k);
|
|
printf(" NT K64:");
|
|
all &= run_gemm_case<RowMajor, ColMajor, 64, 2>(ha, hb_colmajor, c.m,
|
|
c.n, c.k, c.k, c.k);
|
|
printf(" NN K32:");
|
|
all &= run_gemm_case<RowMajor, RowMajor, 32, 3>(ha, hb_rowmajor, c.m,
|
|
c.n, c.k, c.k, c.n);
|
|
printf(" NN K64:");
|
|
all &= run_gemm_case<RowMajor, RowMajor, 64, 2>(ha, hb_rowmajor, c.m,
|
|
c.n, c.k, c.k, c.n);
|
|
printf(" TN K32:");
|
|
all &= run_gemm_case<ColMajor, ColMajor, 32, 3>(ha_t, hb_colmajor, c.m,
|
|
c.n, c.k, c.m, c.k);
|
|
printf(" TN K64:");
|
|
all &= run_gemm_case<ColMajor, ColMajor, 64, 2>(ha_t, hb_colmajor, c.m,
|
|
c.n, c.k, c.m, c.k);
|
|
printf(" TT K64:");
|
|
all &= run_gemm_case<ColMajor, RowMajor, 64, 2>(ha_t, hb_rowmajor, c.m,
|
|
c.n, c.k, c.m, c.n);
|
|
delete[] ha;
|
|
delete[] hb_rowmajor;
|
|
delete[] hb_colmajor;
|
|
delete[] ha_t;
|
|
}
|
|
return all;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
print_test_header();
|
|
bool ok = test_single_mma();
|
|
ok &= test_gemm();
|
|
printf(ok ? "All PASS\n" : "FAILURES\n");
|
|
return ok ? 0 : 1;
|
|
}
|