Remove debug test files, keep production B1/B2 unit tests
This commit is contained in:
@@ -1,211 +0,0 @@
|
||||
/**
|
||||
* 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: warp 1 reads rows 32-63 from the SAME TMEM address
|
||||
// Per P7 docs: different warps see different row slices from the same address
|
||||
// So we DON'T need a TMEM offset for rows 32-63 — warp 1 just reads from tb + col_base
|
||||
// This test uses warp 0 only, so we can only verify rows 0-31.
|
||||
// For rows 32-63, warp 1 must read the same address.
|
||||
// (Skipping rows 32-63 in this single-warp test.)
|
||||
}
|
||||
}
|
||||
__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");
|
||||
}
|
||||
@@ -1,151 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B1 FMHA debug: isolate the cosine failure to noPE vs RoPE path.
|
||||
|
||||
Strategy:
|
||||
1. Run mixed FP8 kernel with RoPE=0 (all noPE) → compare vs BF16
|
||||
2. Run mixed FP8 kernel with noPE=0 (all RoPE) → compare vs BF16
|
||||
3. Run with full split → see which part is broken
|
||||
4. Print per-dimension residual to find where the error lives
|
||||
"""
|
||||
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 cosine(a, b):
|
||||
return F.cosine_similarity(a.flatten().float(), b.flatten().float(), dim=0).item()
|
||||
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
HD = 512; NOPE = 448; ROPE = 64
|
||||
H = 4; B = 1; T = 1; N = 128 # small for debugging
|
||||
scale = 1.0 / math.sqrt(HD)
|
||||
|
||||
print(f"=== B1 FMHA Debug: N={N} H={H} HD={HD} NOPE={NOPE} ROPE={ROPE} ===\n")
|
||||
|
||||
# Generate Q and KV
|
||||
q_fp32 = torch.randn(B, H, T, HD, dtype=torch.float32) * 0.5
|
||||
k_fp32 = torch.randn(N, HD, dtype=torch.float32) * 0.5
|
||||
q_bf16 = q_fp32.bfloat16().cuda()
|
||||
|
||||
# Split KV
|
||||
k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE])
|
||||
k_rope_bf16 = k_fp32[:, NOPE:].bfloat16()
|
||||
k_nope_fp8 = k_nope_fp8.cuda()
|
||||
k_nope_scale = k_nope_scale.cuda()
|
||||
k_rope_bf16 = k_rope_bf16.cuda()
|
||||
|
||||
# --- FP32 Reference ---
|
||||
k_nope_dequant = dequantize_fp8_e4m3(k_nope_fp8.view(torch.uint8).cpu(), k_nope_scale.cpu())
|
||||
k_full = torch.cat([k_nope_dequant, k_fp32[:, NOPE:]], dim=-1) # (N, HD) FP32
|
||||
v_full = k_full.clone()
|
||||
|
||||
q_f = q_fp32.cuda() # (B, H, 1, HD)
|
||||
k_f = k_full.unsqueeze(0).unsqueeze(0).expand(B, -1, -1, -1).cuda()
|
||||
v_f = v_full.unsqueeze(0).unsqueeze(0).expand(B, -1, -1, -1).cuda()
|
||||
o_ref = F.scaled_dot_product_attention(q_f, k_f, v_f, scale=scale) # (B, H, 1, HD)
|
||||
|
||||
print(f"Reference output: |o|={o_ref.abs().max().item():.6f}")
|
||||
print(f" head 0: {o_ref[0,0,0,:8].tolist()}")
|
||||
print(f" head 0 noPE part: {o_ref[0,0,0,:8].tolist()}")
|
||||
print(f" head 0 RoPE part: {o_ref[0,0,0,448:456].tolist()}")
|
||||
|
||||
# --- Mixed FP8 kernel ---
|
||||
from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw
|
||||
o_mixed, lse = fmha_mixed_fp8_decode_raw(
|
||||
q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE)
|
||||
|
||||
print(f"\nMixed FP8 output: |o|={o_mixed.abs().max().item():.6f}")
|
||||
print(f" head 0: {o_mixed[0,0,0,:8].tolist()}")
|
||||
print(f" head 0 RoPE part: {o_mixed[0,0,0,448:456].tolist()}")
|
||||
|
||||
# Global cosine
|
||||
cos = cosine(o_mixed, o_ref.bfloat16())
|
||||
print(f"\nGlobal cosine: {cos:.6f}")
|
||||
|
||||
# Per-head cosine
|
||||
o_mixed_h = o_mixed.float().squeeze(2) # (B, H, HD)
|
||||
o_ref_h = o_ref.bfloat16().float().squeeze(2)
|
||||
per_head = F.cosine_similarity(o_mixed_h, o_ref_h, dim=-1) # (B, H)
|
||||
print(f"Per-head cosine: {per_head[0].tolist()}")
|
||||
print(f" min={per_head.min().item():.6f} mean={per_head.mean().item():.6f}")
|
||||
|
||||
# --- Per-dimension analysis ---
|
||||
# Compare noPE vs RoPE portions separately
|
||||
o_mixed_nope = o_mixed[0, 0, 0, :NOPE].float()
|
||||
o_ref_nope = o_ref[0, 0, 0, :NOPE].float()
|
||||
o_mixed_rope = o_mixed[0, 0, 0, NOPE:].float()
|
||||
o_ref_rope = o_ref[0, 0, 0, NOPE:].float()
|
||||
|
||||
cos_nope = F.cosine_similarity(o_mixed_nope.unsqueeze(0), o_ref_nope.unsqueeze(0), dim=1).item()
|
||||
cos_rope = F.cosine_similarity(o_mixed_rope.unsqueeze(0), o_ref_rope.unsqueeze(0), dim=1).item()
|
||||
|
||||
print(f"\nPer-dim cosine (head 0):")
|
||||
print(f" noPE (0..447): cos={cos_nope:.6f} |mixed|={o_mixed_nope.abs().max():.6f} |ref|={o_ref_nope.abs().max():.6f}")
|
||||
print(f" RoPE (448..511): cos={cos_rope:.6f} |mixed|={o_mixed_rope.abs().max():.6f} |ref|={o_ref_rope.abs().max():.6f}")
|
||||
|
||||
# Residual
|
||||
residual = (o_mixed[0,0,0,:] - o_ref[0,0,0,:].bfloat16()).float()
|
||||
print(f"\nResidual: |res|={residual.abs().max().item():.6f} mean={residual.mean().item():.6f}")
|
||||
print(f" noPE residual: |res|={residual[:NOPE].abs().max().item():.6f}")
|
||||
print(f" RoPE residual: |res|={residual[NOPE:].abs().max().item():.6f}")
|
||||
|
||||
# --- Per-head breakdown ---
|
||||
print(f"\nPer-head noPE/RoPE cosines:")
|
||||
for h in range(H):
|
||||
mn = o_mixed[0,h,0,:NOPE].float()
|
||||
rn = o_ref[0,h,0,:NOPE].float()
|
||||
mr = o_mixed[0,h,0,NOPE:].float()
|
||||
rr = o_ref[0,h,0,NOPE:].float()
|
||||
cn = F.cosine_similarity(mn.unsqueeze(0), rn.unsqueeze(0)).item()
|
||||
cr = F.cosine_similarity(mr.unsqueeze(0), rr.unsqueeze(0)).item()
|
||||
print(f" H{h}: noPE_cos={cn:.4f} rope_cos={cr:.4f} total_cos={per_head[0,h].item():.4f}")
|
||||
|
||||
# --- Score comparison ---
|
||||
# Compute reference scores manually
|
||||
q_h0 = q_fp32[0, 0, 0, :].cuda().float() # (HD,)
|
||||
k_all = k_full.cuda().float() # (N, HD)
|
||||
scores_ref = torch.matmul(q_h0, k_all.T) * scale # (N,)
|
||||
|
||||
# Compute noPE and RoPE scores separately
|
||||
scores_nope_ref = torch.matmul(q_h0[:NOPE], k_all[:, :NOPE].T) * scale
|
||||
scores_rope_ref = torch.matmul(q_h0[NOPE:], k_all[:, NOPE:].T) * scale
|
||||
|
||||
print(f"\nScore analysis (head 0):")
|
||||
print(f" noPE: [{scores_nope_ref.min().item():.4f}, {scores_nope_ref.max().item():.4f}]")
|
||||
print(f" RoPE: [{scores_rope_ref.min().item():.4f}, {scores_rope_ref.max().item():.4f}]")
|
||||
print(f" Total: [{scores_ref.min().item():.4f}, {scores_ref.max().item():.4f}]")
|
||||
|
||||
# Check: are the noPE and RoPE scores the right order of magnitude?
|
||||
nope_range = scores_nope_ref.max().item() - scores_nope_ref.min().item()
|
||||
rope_range = scores_rope_ref.max().item() - scores_rope_ref.min().item()
|
||||
print(f" noPE range: {nope_range:.4f} RoPE range: {rope_range:.4f}")
|
||||
print(f" noPE/RoPE ratio: {nope_range/rope_range:.2f}" if rope_range > 0 else " RoPE range is zero!")
|
||||
|
||||
# --- Check if the kernel is producing V=K correctly ---
|
||||
# In MQA self-attention, V=K. Check if the output magnitude matches
|
||||
# the expected: o = softmax(QK^T/sqrt(d)) @ K
|
||||
# With K=V and N=128, the output should be a weighted average of K rows
|
||||
print(f"\n K (row 0): {k_full[0,:8].tolist()}")
|
||||
print(f" o (head 0): {o_mixed[0,0,0,:8].float().tolist()}")
|
||||
print(f" o_ref (head 0): {o_ref[0,0,0,:8].tolist()}")
|
||||
|
||||
sys.exit(0 if cos >= 0.999 else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,171 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B1 FMHA isolate: test QK scores separately from PV.
|
||||
|
||||
Strategy:
|
||||
1. Compute QK scores with the kernel and with reference
|
||||
2. If QK is wrong, the bug is in QK. If QK is right, bug is in PV.
|
||||
3. Also test: single-head, single N=128 to minimize moving parts.
|
||||
"""
|
||||
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)
|
||||
HD = 512; NOPE = 448; ROPE = 64
|
||||
H = 1; B = 1; T = 1; N = 128
|
||||
scale = 1.0 / math.sqrt(HD)
|
||||
|
||||
print(f"=== B1 FMHA Isolate: QK vs PV (N={N} H={H}) ===\n")
|
||||
|
||||
# Generate Q and KV
|
||||
q_fp32 = torch.randn(B, H, T, HD, dtype=torch.float32) * 0.5
|
||||
k_fp32 = torch.randn(N, HD, dtype=torch.float32) * 0.5
|
||||
q_bf16 = q_fp32.bfloat16().cuda()
|
||||
|
||||
k_nope_fp8, k_nope_scale = quantize_fp8_e4m3(k_fp32[:, :NOPE])
|
||||
k_rope_bf16 = k_fp32[:, NOPE:].bfloat16()
|
||||
k_nope_fp8 = k_nope_fp8.cuda(); k_nope_scale = k_nope_scale.cuda()
|
||||
k_rope_bf16 = k_rope_bf16.cuda()
|
||||
|
||||
# --- Reference QK scores (head 0) ---
|
||||
k_nope_dequant = dequantize_fp8_e4m3(k_nope_fp8.view(torch.uint8).cpu(), k_nope_scale.cpu()).cuda()
|
||||
k_full = torch.cat([k_nope_dequant, k_fp32[:, NOPE:].cuda()], dim=-1) # (N, HD)
|
||||
|
||||
q_h0 = q_fp32[0, 0, 0, :].cuda().float()
|
||||
scores_ref = torch.matmul(q_h0, k_full.T) * scale # (N,)
|
||||
|
||||
# Separate noPE and RoPE scores
|
||||
scores_nope_ref = torch.matmul(q_h0[:NOPE], k_full[:, :NOPE].T) * scale
|
||||
scores_rope_ref = torch.matmul(q_h0[NOPE:], k_full[:, NOPE:].T) * scale
|
||||
|
||||
print(f"Reference scores (head 0):")
|
||||
print(f" Total: [{scores_ref.min():.4f}, {scores_ref.max():.4f}]")
|
||||
print(f" noPE: [{scores_nope_ref.min():.4f}, {scores_nope_ref.max():.4f}]")
|
||||
print(f" RoPE: [{scores_rope_ref.min():.4f}, {scores_rope_ref.max():.4f}]")
|
||||
print(f" First 8 total scores: {scores_ref[:8].tolist()}")
|
||||
|
||||
# --- Run kernel and extract LSE ---
|
||||
from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw
|
||||
o_mixed, lse = fmha_mixed_fp8_decode_raw(
|
||||
q_bf16, k_nope_fp8, k_nope_scale, k_rope_bf16, scale, rope_dim=ROPE)
|
||||
|
||||
# LSE should equal logsumexp of scores
|
||||
ref_lse = torch.logsumexp(scores_ref, dim=0)
|
||||
print(f"\nLSE comparison (head 0):")
|
||||
print(f" Kernel LSE: {lse[0,0,0].item():.4f}")
|
||||
print(f" Reference LSE: {ref_lse.item():.4f}")
|
||||
print(f" Diff: {abs(lse[0,0,0].item() - ref_lse.item()):.4f}")
|
||||
|
||||
# If LSE is close but output is wrong, bug is in PV.
|
||||
# If LSE is far off, bug is in QK.
|
||||
lse_close = abs(lse[0,0,0].item() - ref_lse.item()) < 0.1
|
||||
|
||||
# --- Check softmax probabilities ---
|
||||
# From reference scores
|
||||
probs_ref = F.softmax(scores_ref, dim=0)
|
||||
print(f"\nReference softmax probs: [{probs_ref.min():.6f}, {probs_ref.max():.6f}]")
|
||||
print(f" First 8 probs: {probs_ref[:8].tolist()}")
|
||||
|
||||
# --- Check output = P @ V ---
|
||||
# Reference: o = probs @ K (since V = K)
|
||||
o_ref = torch.matmul(probs_ref.unsqueeze(0), k_full).squeeze(0) # (HD,)
|
||||
|
||||
o_mixed_h0 = o_mixed[0, 0, 0, :].float()
|
||||
cos = F.cosine_similarity(o_mixed_h0.unsqueeze(0), o_ref.unsqueeze(0)).item()
|
||||
|
||||
print(f"\nOutput comparison (head 0):")
|
||||
print(f" cos(mixed, ref_P@V) = {cos:.6f}")
|
||||
print(f" |mixed| = {o_mixed_h0.norm():.6f}")
|
||||
print(f" |ref_P@V| = {o_ref.norm():.6f}")
|
||||
|
||||
# --- Check if PV is computing P @ V correctly ---
|
||||
# Compute P @ V step by step
|
||||
# The kernel does PV by splitting V into (SK_TILE, 16) sub-tiles
|
||||
# For N=128, HD=512: 32 sub-tiles of 16 dims each
|
||||
# P is (1, 128), V is (128, 512)
|
||||
# Expected: (1, 512)
|
||||
|
||||
# Verify that ref P@V matches the simple attention output
|
||||
o_ref_full = F.scaled_dot_product_attention(
|
||||
q_fp32.cuda()[:, :1, :, :], # (1, 1, 1, HD)
|
||||
k_full.unsqueeze(0).unsqueeze(0), # (1, 1, N, HD)
|
||||
k_full.unsqueeze(0).unsqueeze(0), # V=K
|
||||
scale=scale
|
||||
)
|
||||
cos_ref = F.cosine_similarity(o_ref.unsqueeze(0), o_ref_full[0,0,0,:].unsqueeze(0)).item()
|
||||
print(f" cos(ref_P@V, ref_SDPA) = {cos_ref:.6f}")
|
||||
|
||||
# --- Analyze the PV sub-tile structure ---
|
||||
# The kernel computes PV as:
|
||||
# For n_sub = 0..31:
|
||||
# MMA(P[128x128], V[128x16]) → TMEM[128x16] at offset n_sub*16
|
||||
# Then reads TMEM and accumulates
|
||||
#
|
||||
# The V matrix construction: V[row, d_base+dd] where d_base = n_sub*16
|
||||
# For noPE V: dequantized from FP8
|
||||
# For RoPE V: directly from k_rope_bf16
|
||||
|
||||
# Check that the V matrix indexing matches
|
||||
# V should be K, so V[row, d] = K[row, d]
|
||||
print(f"\nV matrix sanity (row 0):")
|
||||
# noPE part
|
||||
v_nope_ref = k_nope_dequant[0, :8] # first 8 noPE dims
|
||||
print(f" K_nope[0,:8] (dequant): {v_nope_ref.tolist()}")
|
||||
print(f" K_orig[0,:8]: {k_fp32[0, :8].tolist()}")
|
||||
cos_v = F.cosine_similarity(v_nope_ref.unsqueeze(0), k_fp32[0, :8].unsqueeze(0)).item()
|
||||
print(f" cos(dequant, orig) = {cos_v:.6f}")
|
||||
|
||||
# Diagnose
|
||||
if lse_close:
|
||||
print(f"\n*** DIAGNOSIS: LSE is close ({abs(lse[0,0,0].item() - ref_lse.item()):.4f}) but output cos is {cos:.6f}")
|
||||
print(f"*** BUG IS IN PV (probability-value multiply), NOT IN QK")
|
||||
else:
|
||||
print(f"\n*** DIAGNOSIS: LSE is far off ({abs(lse[0,0,0].item() - ref_lse.item()):.4f})")
|
||||
print(f"*** BUG IS IN QK (query-key scoring)")
|
||||
|
||||
# --- Extra: compare noPE-only output ---
|
||||
# Zero out RoPE dims in Q and K, run kernel, compare
|
||||
print(f"\n--- noPE-only test (RoPE zeroed) ---")
|
||||
q_nope_only = q_bf16.clone()
|
||||
q_nope_only[:, :, :, NOPE:] = 0 # zero RoPE in Q
|
||||
k_rope_zero = torch.zeros(N, ROPE, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
try:
|
||||
o_nope, lse_nope = fmha_mixed_fp8_decode_raw(
|
||||
q_nope_only, k_nope_fp8, k_nope_scale, k_rope_zero, scale, rope_dim=ROPE)
|
||||
|
||||
# Reference with zeroed RoPE
|
||||
q_nz = q_fp32.clone().cuda()
|
||||
q_nz[:, :, :, NOPE:] = 0
|
||||
k_nz = k_full.clone()
|
||||
k_nz[:, NOPE:] = 0
|
||||
o_ref_nope = F.scaled_dot_product_attention(
|
||||
q_nz, k_nz.unsqueeze(0).unsqueeze(0), k_nz.unsqueeze(0).unsqueeze(0), scale=scale)
|
||||
|
||||
cos_nope = F.cosine_similarity(
|
||||
o_nope[0,0,0,:].float().unsqueeze(0),
|
||||
o_ref_nope[0,0,0,:].float().unsqueeze(0)).item()
|
||||
print(f" noPE-only cos = {cos_nope:.6f}")
|
||||
except Exception as e:
|
||||
print(f" noPE-only test failed: {e}")
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,117 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B2 debug: read raw TMEM logits and compare with FP32 QK scores.
|
||||
|
||||
Bypass the weighted ReLU and top-k — just check if the FP8 GEMM
|
||||
produces correct logits after dequantization.
|
||||
"""
|
||||
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 = 8 # small for manual inspection
|
||||
TOP_K = 8
|
||||
|
||||
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
|
||||
w_h = torch.ones(N_IH, dtype=torch.bfloat16).cuda() # all positive weights for simplicity
|
||||
k_fp8, k_scale = quantize_fp8_e4m3(k_fp32)
|
||||
k_fp8 = k_fp8.cuda(); k_scale = k_scale.cuda()
|
||||
|
||||
# --- FP32 reference scores (full matrix) ---
|
||||
k_dequant = dequantize_fp8_e4m3(k_fp8.view(torch.uint8).cpu(), k_scale.cpu()).cuda()
|
||||
scores_ref = torch.einsum('nd,cd->nc', q_idx.float(), k_dequant.float()) # (64, 8)
|
||||
|
||||
print("=== FP32 Reference Scores (h, c) ===")
|
||||
for h in range(4):
|
||||
print(f" H{h}: {scores_ref[h, :4].tolist()}")
|
||||
|
||||
# --- Q per-row quantization (same as kernel) ---
|
||||
q_fp32 = q_idx.float()
|
||||
q_amax = q_fp32.abs().amax(dim=-1)
|
||||
q_scales = (q_amax / 448.0).clamp(min=1e-8)
|
||||
q_fp8_raw = (q_fp32 / q_scales.unsqueeze(-1)).clamp(-448, 448)
|
||||
q_fp8 = q_fp8_raw.to(torch.float8_e4m3fn).view(torch.uint8)
|
||||
|
||||
# FP8 dequant round-trip
|
||||
q_dequant = dequantize_fp8_e4m3(q_fp8, q_scales)
|
||||
cos_q = F.cosine_similarity(q_fp32.flatten().unsqueeze(0), q_dequant.flatten().unsqueeze(0)).item()
|
||||
print(f"\nQ FP8 round-trip cosine: {cos_q:.6f}")
|
||||
|
||||
# FP8 GEMM in Python: (q_dequant / q_scales) . (k_dequant / k_scales)
|
||||
# = q_dequant @ k_dequant^T / (q_scales * k_scales)
|
||||
# But wait — the MMA computes (q_fp8 @ k_fp8^T), and the result is
|
||||
# raw logits that need to be dequantized: logit * q_scale[h] * k_scale[c]
|
||||
# q_fp8 = q / q_scale, k_fp8 = k / k_scale
|
||||
# MMA output = q_fp8 @ k_fp8^T = (q/q_scale) @ (k/k_scale)^T
|
||||
# Dequantized = MMA_output * q_scale * k_scale = q @ k^T
|
||||
|
||||
# Simulate: compute the raw MMA output as (q_fp8 values @ k_fp8 values)
|
||||
q8_dequant_for_mma = q_dequant / q_scales.unsqueeze(-1) # "q_fp8" values (but in FP32)
|
||||
k8_dequant_for_mma = k_dequant / k_scale.unsqueeze(-1) # "k_fp8" values (but in FP32)
|
||||
mma_output = torch.einsum('nd,cd->nc', q8_dequant_for_mma, k8_dequant_for_mma) # (64, 8)
|
||||
|
||||
# Dequantized logits = mma_output * q_scale * k_scale
|
||||
dequant_logits = mma_output * q_scales.unsqueeze(-1) * k_scale.unsqueeze(0)
|
||||
|
||||
print(f"\nSimulated dequant logits:")
|
||||
for h in range(4):
|
||||
print(f" H{h}: {dequant_logits[h, :4].tolist()}")
|
||||
print(f" Ref: {scores_ref[h, :4].tolist()}")
|
||||
|
||||
cos_logits = F.cosine_similarity(dequant_logits.flatten().unsqueeze(0),
|
||||
scores_ref.flatten().unsqueeze(0)).item()
|
||||
print(f" Cosine vs ref: {cos_logits:.6f}")
|
||||
|
||||
# Now check: the kernel's MMA should produce similar raw logits.
|
||||
# If we run the kernel and it gives different results, the MMA itself is wrong.
|
||||
|
||||
# --- Run B2 kernel ---
|
||||
from dsv4.kernels.cuda.loader import get_cuda_module
|
||||
mod = get_cuda_module("indexer_fp8_score_topk", ["indexer_fp8_score_topk.cu"],
|
||||
extra_cuda_cflags=[
|
||||
"-gencode=arch=compute_100a,code=sm_100a",
|
||||
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
|
||||
])
|
||||
|
||||
topk_indices = torch.empty(TOP_K, dtype=torch.int32, device='cuda')
|
||||
mod.indexer_fp8_score_topk(
|
||||
q_idx, k_fp8, k_scale, w_h, topk_indices,
|
||||
N_IH, IHD, TOP_K)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
print(f"\nKernel top-k (w_h=1): {topk_indices.tolist()}")
|
||||
|
||||
# Reference top-k with w_h=1
|
||||
scores_relu = F.relu(scores_ref)
|
||||
ref_total = scores_relu.sum(0) # sum over heads
|
||||
ref_topk = ref_total.topk(TOP_K).indices
|
||||
print(f"Reference top-k (w_h=1): {ref_topk.tolist()}")
|
||||
print(f"Reference scores (w_h=1): {ref_total[ref_topk].tolist()}")
|
||||
|
||||
# Check overlap
|
||||
kernel_set = set(topk_indices.cpu().tolist()) - {-1}
|
||||
ref_set = set(ref_topk.tolist())
|
||||
overlap = len(kernel_set & ref_set)
|
||||
print(f"Overlap: {overlap}/{TOP_K}")
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,97 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B2 debug: compare kernel scores with FP32 reference.
|
||||
|
||||
Prints actual score values to find where the kernel diverges.
|
||||
"""
|
||||
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; TOP_K = 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
|
||||
w_h = torch.randn(N_IH, dtype=torch.bfloat16).cuda() * 0.3
|
||||
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()
|
||||
scores_full = torch.einsum('nd,cd->nc', q_idx.float(), k_dequant.float()) # (64, 128)
|
||||
scores_relu = F.relu(scores_full)
|
||||
ref_scores = (scores_relu * w_h.unsqueeze(-1).float()).sum(0) # (128,)
|
||||
|
||||
print(f"Reference scores: [{ref_scores.min():.4f}, {ref_scores.max():.4f}] mean={ref_scores.mean():.4f}")
|
||||
print(f" Positive scores: {(ref_scores > 0).sum().item()}/{N_COMP}")
|
||||
print(f" First 8: {ref_scores[:8].tolist()}")
|
||||
|
||||
# Also check per-head scores
|
||||
print(f"\nPer-head score stats:")
|
||||
for h in [0, 1, 32, 63]:
|
||||
head_scores = scores_full[h] # (128,)
|
||||
print(f" H{h}: [{head_scores.min():.4f}, {head_scores.max():.4f}] "
|
||||
f"relu_sum={scores_relu[h].sum():.4f} weighted_relu={(scores_relu[h]*w_h[h].float()).sum():.4f}")
|
||||
|
||||
# --- Run B2 kernel ---
|
||||
from dsv4.kernels.cuda.loader import get_cuda_module
|
||||
mod = get_cuda_module("indexer_fp8_score_topk", ["indexer_fp8_score_topk.cu"],
|
||||
extra_cuda_cflags=[
|
||||
"-gencode=arch=compute_100a,code=sm_100a",
|
||||
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
|
||||
])
|
||||
|
||||
topk_indices = torch.empty(TOP_K, dtype=torch.int32, device='cuda')
|
||||
mod.indexer_fp8_score_topk(
|
||||
q_idx, k_fp8, k_scale, w_h, topk_indices,
|
||||
N_IH, IHD, TOP_K)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
print(f"\nKernel top-k indices: {topk_indices[:16].tolist()}")
|
||||
print(f" Valid (>=0): {(topk_indices >= 0).sum().item()}")
|
||||
print(f" Unique: {len(set(topk_indices.cpu().tolist()))}")
|
||||
|
||||
# Compare: kernel-selected scores vs reference
|
||||
kernel_selected = topk_indices[topk_indices >= 0].cpu()
|
||||
if len(kernel_selected) > 0:
|
||||
ref_selected_scores = ref_scores[kernel_selected]
|
||||
print(f" Kernel-selected ref scores: {ref_selected_scores[:8].tolist()}")
|
||||
|
||||
# Reference top-k
|
||||
ref_topk = ref_scores.topk(TOP_K).indices
|
||||
print(f" Reference top-k: {ref_topk[:8].tolist()}")
|
||||
|
||||
# Overlap
|
||||
kernel_set = set(topk_indices.cpu().tolist()) - {-1}
|
||||
ref_set = set(ref_topk.tolist())
|
||||
overlap = len(kernel_set & ref_set)
|
||||
print(f" Overlap: {overlap}/{len(ref_set)}")
|
||||
|
||||
# Check: are the FP8 Q scales correct?
|
||||
# The kernel quantizes Q internally. Let's compute what it should use.
|
||||
q_fp32 = q_idx.float()
|
||||
q_amax = q_fp32.abs().amax(dim=-1) # (64,)
|
||||
q_scales = q_amax / 448.0
|
||||
q_scales = q_scales.clamp(min=1e-8)
|
||||
print(f"\nQ per-row scales: [{q_scales.min():.6f}, {q_scales.max():.6f}] mean={q_scales.mean():.6f}")
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,73 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""B2 minimal debug: test if the FP8 indexer kernel even completes.
|
||||
|
||||
Start with n_comp=1 (trivial case) and work up.
|
||||
Also test: does the kernel return at all? Or does it hang?
|
||||
"""
|
||||
import sys
|
||||
import math
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import time
|
||||
|
||||
|
||||
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 main():
|
||||
N_IH = 64; IHD = 128; TOP_K = 1
|
||||
|
||||
from dsv4.kernels.cuda.loader import get_cuda_module
|
||||
mod = get_cuda_module("indexer_fp8_score_topk", ["indexer_fp8_score_topk.cu"],
|
||||
extra_cuda_cflags=[
|
||||
"-gencode=arch=compute_100a,code=sm_100a",
|
||||
"-O3", "--use_fast_math", "--expt-relaxed-constexpr",
|
||||
])
|
||||
|
||||
# Test with increasing n_comp
|
||||
for n_comp in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
|
||||
print(f"\n--- n_comp={n_comp} ---", flush=True)
|
||||
torch.manual_seed(42)
|
||||
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
|
||||
w_h = torch.randn(N_IH, dtype=torch.bfloat16).cuda() * 0.3
|
||||
k_fp8, k_scale = quantize_fp8_e4m3(k_fp32)
|
||||
k_fp8 = k_fp8.cuda(); k_scale = k_scale.cuda()
|
||||
|
||||
tk = min(TOP_K, n_comp)
|
||||
topk_indices = torch.empty(tk, dtype=torch.int32, device='cuda')
|
||||
|
||||
# Add a CUDA event to time the kernel
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
try:
|
||||
mod.indexer_fp8_score_topk(
|
||||
q_idx, k_fp8, k_scale, w_h, topk_indices,
|
||||
N_IH, IHD, tk)
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
ms = start.elapsed_time(end)
|
||||
print(f" OK: {ms:.2f}ms indices={topk_indices.cpu().tolist()}", flush=True)
|
||||
except Exception as e:
|
||||
torch.cuda.synchronize()
|
||||
print(f" FAILED: {e}", flush=True)
|
||||
break
|
||||
|
||||
# Check for CUDA errors
|
||||
err = torch.cuda.current_device()
|
||||
if n_comp >= 128:
|
||||
# After a certain size, try timing it
|
||||
# Also check if the GPU is responsive
|
||||
print(f" GPU responsive: yes", flush=True)
|
||||
|
||||
print("\nDone.")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,65 +0,0 @@
|
||||
#!/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