Add B2 TMEM read debug kernel and test
This commit is contained in:
219
dsv4/kernels/cuda/test_fp8_gemm_tmem_read.cu
Normal file
219
dsv4/kernels/cuda/test_fp8_gemm_tmem_read.cu
Normal file
@@ -0,0 +1,219 @@
|
||||
/**
|
||||
* B2 debug: FP8 GEMM + TMEM read test kernel.
|
||||
* Produces raw dequantized logits for comparison with FP32 reference.
|
||||
*/
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_fp8.hpp>
|
||||
#include <ATen/ATen.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/extension.h>
|
||||
#include <cstdint>
|
||||
#include <cfloat>
|
||||
|
||||
static constexpr float E4M3_MAX = 448.0f;
|
||||
typedef unsigned short bf16_t;
|
||||
|
||||
__device__ __forceinline__ float bf16_to_f32_ptx(bf16_t h) {
|
||||
float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f;
|
||||
}
|
||||
__device__ __forceinline__ uint8_t fp8_e4m3_from_f32(float x) {
|
||||
x = fminf(fmaxf(x, -E4M3_MAX), E4M3_MAX);
|
||||
__nv_fp8_e4m3 v(x);
|
||||
return *reinterpret_cast<uint8_t*>(&v);
|
||||
}
|
||||
__device__ __forceinline__ uint64_t desc_encode(uint64_t byte_val) { return byte_val >> 4; }
|
||||
__device__ __forceinline__ uint64_t make_umma_desc_kmajor_none(uint32_t smem_addr, int block_mn) {
|
||||
const uint64_t LBO = block_mn * 16;
|
||||
const uint64_t SBO = 128;
|
||||
uint64_t desc = 0;
|
||||
desc |= desc_encode(smem_addr) & 0x3FFF;
|
||||
desc |= (desc_encode(LBO) & 0x3FFF) << 16;
|
||||
desc |= (desc_encode(SBO) & 0x3FFF) << 32;
|
||||
desc |= 1ULL << 46;
|
||||
return desc;
|
||||
}
|
||||
__device__ __forceinline__ uint32_t make_idesc_f8_e4m3(int block_m, int block_n) {
|
||||
return (1U << 4) | ((uint32_t)(block_n >> 3) << 17) | ((uint32_t)(block_m >> 4) << 24);
|
||||
}
|
||||
__device__ void umma_ss_f8f6f4(uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b,
|
||||
uint32_t i_desc, bool accumulate) {
|
||||
uint32_t scaleC_bits = accumulate ? 0x3F800000u : 0u;
|
||||
asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::1.kind::f8f6f4 [%0], %1, %2, %3, p;\n\t}"
|
||||
:: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(i_desc), "r"(scaleC_bits));
|
||||
}
|
||||
__device__ void tmem_alloc(uint32_t smem_ptr, int num_cols) {
|
||||
asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
|
||||
:: "r"(smem_ptr), "r"(num_cols));
|
||||
}
|
||||
__device__ void tmem_dealloc(uint32_t tmem_ptr, int num_cols) {
|
||||
asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
|
||||
:: "r"(tmem_ptr), "r"(num_cols));
|
||||
}
|
||||
__device__ __forceinline__ int canon_idx_fp8_128x32(int r, int c) {
|
||||
int core_mn = r >> 3; int core_k = c >> 4;
|
||||
int local_r = r & 7; int local_c = c & 15;
|
||||
return core_k * 16 * 128 + core_mn * 128 + local_r * 16 + local_c;
|
||||
}
|
||||
|
||||
__global__ void __launch_bounds__(192)
|
||||
test_fp8_gemm_tmem_read_kernel(
|
||||
const bf16_t* q_bf16, const uint8_t* k_fp8, const float* k_scale,
|
||||
float* logits_out, int n_comp, int n_ih, int ihd
|
||||
) {
|
||||
constexpr int MMA_K_F8 = 32;
|
||||
constexpr int NKT = 4;
|
||||
constexpr int SK_TILE = 128;
|
||||
constexpr int TILE_F8 = 128 * 32;
|
||||
constexpr int TMEM_COLS = 512;
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int wid = tid >> 5;
|
||||
const int lane = tid & 31;
|
||||
const bool is_mma_warp = (wid == 4);
|
||||
|
||||
extern __shared__ __align__(128) char sbuf[];
|
||||
size_t off = 0;
|
||||
uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off += 4;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
uint8_t* sQ8 = (uint8_t*)(sbuf + off); off += TILE_F8;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
uint8_t* sK8 = (uint8_t*)(sbuf + off); off += TILE_F8;
|
||||
off = (off + 127) & ~(size_t)127;
|
||||
float* sQ_scale = (float*)(sbuf + off); off += 128 * sizeof(float);
|
||||
|
||||
for (int h = 0; h < n_ih; h++) {
|
||||
float local_max = 0.0f;
|
||||
for (int d = tid; d < ihd; d += 192) {
|
||||
float val = fabsf(bf16_to_f32_ptx(q_bf16[h * ihd + d]));
|
||||
local_max = fmaxf(local_max, val);
|
||||
}
|
||||
for (int o = 16; o > 0; o >>= 1)
|
||||
local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, o));
|
||||
__shared__ float _q_amax[6];
|
||||
if ((tid & 31) == 0) _q_amax[tid >> 5] = local_max;
|
||||
__syncthreads();
|
||||
float amax = 0.0f;
|
||||
if (tid < 32) {
|
||||
amax = (tid < 6) ? _q_amax[tid] : 0.0f;
|
||||
for (int o = 16; o > 0; o >>= 1)
|
||||
amax = fmaxf(amax, __shfl_down_sync(0xffffffff, amax, o));
|
||||
}
|
||||
amax = __shfl_sync(0xffffffff, amax, 0);
|
||||
float scale = amax / E4M3_MAX;
|
||||
if (scale < 1e-8f) scale = 1e-8f;
|
||||
if (tid == 0) sQ_scale[h] = scale;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (is_mma_warp) tmem_alloc((uint32_t)__cvta_generic_to_shared(sTmemBase), TMEM_COLS);
|
||||
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
|
||||
__syncthreads();
|
||||
uint32_t tb = *sTmemBase;
|
||||
|
||||
const int n_k_tiles = (n_comp + SK_TILE - 1) / SK_TILE;
|
||||
const uint32_t idesc_f8 = make_idesc_f8_e4m3(128, 128);
|
||||
|
||||
for (int kv_tile = 0; kv_tile < n_k_tiles; kv_tile++) {
|
||||
const int kv_start = kv_tile * SK_TILE;
|
||||
const int kv_len = min(SK_TILE, n_comp - kv_start);
|
||||
|
||||
for (int kt = 0; kt < NKT; kt++) {
|
||||
for (int i = tid; i < TILE_F8; i += 192) { sQ8[i] = 0; sK8[i] = 0; }
|
||||
__syncthreads();
|
||||
for (int i = tid; i < n_ih * MMA_K_F8; i += 192) {
|
||||
int row = i / MMA_K_F8, col = i % MMA_K_F8;
|
||||
int d = kt * MMA_K_F8 + col;
|
||||
if (d < ihd) {
|
||||
float val = bf16_to_f32_ptx(q_bf16[row * ihd + d]);
|
||||
float inv_scale = 1.0f / sQ_scale[row];
|
||||
sQ8[canon_idx_fp8_128x32(row, col)] = fp8_e4m3_from_f32(val * inv_scale);
|
||||
}
|
||||
}
|
||||
for (int i = tid; i < kv_len * MMA_K_F8; i += 192) {
|
||||
int row = i / MMA_K_F8, col = i % MMA_K_F8;
|
||||
int d = kt * MMA_K_F8 + col;
|
||||
int g_row = kv_start + row;
|
||||
sK8[canon_idx_fp8_128x32(row, col)] = k_fp8[(int64_t)g_row * ihd + d];
|
||||
}
|
||||
__syncthreads();
|
||||
if (is_mma_warp && lane == 0) {
|
||||
uint64_t dq = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sQ8), 128);
|
||||
uint64_t dk = make_umma_desc_kmajor_none((uint32_t)__cvta_generic_to_shared(sK8), 128);
|
||||
umma_ss_f8f6f4(tb, dq, dk, idesc_f8, kt > 0);
|
||||
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
asm volatile("fence.sc.gpu;" ::: "memory");
|
||||
__syncthreads();
|
||||
|
||||
// Warp 0 reads TMEM and stores logits
|
||||
if (wid == 0) {
|
||||
for (int n = 0; n < SK_TILE / 8; n++) {
|
||||
int col_base = n * 8;
|
||||
if (col_base >= kv_len) break;
|
||||
int cols_valid = min(8, kv_len - col_base);
|
||||
|
||||
// Row group 0-31
|
||||
float tmp[8] = {};
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),
|
||||
"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7])
|
||||
: "r"(tb + col_base));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory");
|
||||
if (lane < n_ih && lane < 32) {
|
||||
for (int j = 0; j < cols_valid; j++) {
|
||||
float k_s = k_scale[kv_start + col_base + j];
|
||||
logits_out[(int64_t)lane * n_comp + kv_start + col_base + j] = tmp[j] * sQ_scale[lane] * k_s;
|
||||
}
|
||||
}
|
||||
|
||||
// Row group 32-63
|
||||
float tmp2[8] = {};
|
||||
asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];"
|
||||
: "=f"(tmp2[0]),"=f"(tmp2[1]),"=f"(tmp2[2]),"=f"(tmp2[3]),
|
||||
"=f"(tmp2[4]),"=f"(tmp2[5]),"=f"(tmp2[6]),"=f"(tmp2[7])
|
||||
: "r"(tb + SK_TILE + col_base));
|
||||
asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory");
|
||||
if (lane < n_ih - 32 && lane < 32) {
|
||||
int h = lane + 32;
|
||||
for (int j = 0; j < cols_valid; j++) {
|
||||
float k_s = k_scale[kv_start + col_base + j];
|
||||
logits_out[(int64_t)h * n_comp + kv_start + col_base + j] = tmp2[j] * sQ_scale[h] * k_s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (is_mma_warp) tmem_dealloc(tb, TMEM_COLS);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
void test_fp8_gemm_tmem_read_cuda(
|
||||
torch::Tensor q_bf16, torch::Tensor k_fp8, torch::Tensor k_scale,
|
||||
torch::Tensor logits_out, int64_t n_ih, int64_t ihd
|
||||
) {
|
||||
int n_comp = k_fp8.size(0);
|
||||
auto k8 = k_fp8.dtype() == torch::kUInt8 ? k_fp8 : k_fp8.view(torch::kUInt8);
|
||||
size_t smem = 0;
|
||||
smem += 4; smem = (smem + 127) & ~127;
|
||||
smem += 128 * 32; smem = (smem + 127) & ~127;
|
||||
smem += 128 * 32; smem = (smem + 127) & ~127;
|
||||
smem += 128 * 4;
|
||||
cudaFuncSetAttribute(test_fp8_gemm_tmem_read_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
test_fp8_gemm_tmem_read_kernel<<<1, 192, smem, c10::cuda::getCurrentCUDAStream()>>>(
|
||||
reinterpret_cast<const bf16_t*>(q_bf16.data_ptr<at::BFloat16>()),
|
||||
k8.data_ptr<uint8_t>(), k_scale.data_ptr<float>(),
|
||||
logits_out.data_ptr<float>(), n_comp, (int)n_ih, (int)ihd);
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("test_fp8_gemm_tmem_read", &test_fp8_gemm_tmem_read_cuda,
|
||||
"B2 debug: FP8 GEMM + TMEM read test");
|
||||
}
|
||||
65
tests/unit/test_b2_tmem_read.py
Normal file
65
tests/unit/test_b2_tmem_read.py
Normal file
@@ -0,0 +1,65 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B2 TMEM read verification: compare raw FP8 GEMM logits with FP32 reference."""
|
||||
import sys
|
||||
import math
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def quantize_fp8_e4m3(x_fp32):
|
||||
amax = x_fp32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
|
||||
scale = amax / 448.0
|
||||
fp8 = (x_fp32 / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
|
||||
return fp8.view(torch.uint8), scale.squeeze(-1)
|
||||
|
||||
|
||||
def dequantize_fp8_e4m3(fp8_uint8, scale):
|
||||
fp8 = fp8_uint8.view(torch.float8_e4m3fn)
|
||||
return fp8.float() * scale.unsqueeze(-1).float()
|
||||
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
N_IH = 64; IHD = 128; N_COMP = 128
|
||||
|
||||
q_idx = torch.randn(N_IH, IHD, dtype=torch.bfloat16).cuda() * 0.5
|
||||
k_fp32 = torch.randn(N_COMP, IHD, dtype=torch.float32) * 0.5
|
||||
k_fp8, k_scale = quantize_fp8_e4m3(k_fp32)
|
||||
k_fp8 = k_fp8.cuda(); k_scale = k_scale.cuda()
|
||||
|
||||
# FP32 reference
|
||||
k_dequant = dequantize_fp8_e4m3(k_fp8.view(torch.uint8).cpu(), k_scale.cpu()).cuda()
|
||||
ref_logits = torch.einsum('nd,cd->nc', q_idx.float(), k_dequant.float())
|
||||
|
||||
# Run test kernel
|
||||
from dsv4.kernels.cuda.loader import get_cuda_module
|
||||
mod = get_cuda_module("test_fp8_gemm_tmem_read", ["test_fp8_gemm_tmem_read.cu"],
|
||||
extra_cuda_cflags=[
|
||||
"-gencode=arch=compute_100a,code=sm_100a",
|
||||
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
|
||||
])
|
||||
|
||||
logits = torch.zeros(N_IH, N_COMP, dtype=torch.float32, device='cuda')
|
||||
mod.test_fp8_gemm_tmem_read(q_idx, k_fp8, k_scale, logits, N_IH, IHD)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
cos = F.cosine_similarity(logits.flatten().unsqueeze(0), ref_logits.flatten().unsqueeze(0)).item()
|
||||
print(f"Global cosine: {cos:.6f}")
|
||||
|
||||
per_head_cos = F.cosine_similarity(logits, ref_logits, dim=-1)
|
||||
print(f"Per-head cos: min={per_head_cos.min():.6f} mean={per_head_cos.mean():.6f}")
|
||||
|
||||
for h in [0, 1, 31, 32, 33, 63]:
|
||||
hc = per_head_cos[h].item()
|
||||
print(f" H{h}: cos={hc:.6f} |kernel|={logits[h].abs().max():.4f} |ref|={ref_logits[h].abs().max():.4f}")
|
||||
|
||||
# Column-wise
|
||||
for c in [0, 32, 64, 96, 127]:
|
||||
col_cos = F.cosine_similarity(logits[:, c].unsqueeze(0), ref_logits[:, c].unsqueeze(0)).item()
|
||||
print(f" Col {c}: cos={col_cos:.6f}")
|
||||
|
||||
sys.exit(0 if cos >= 0.999 else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user