66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
#!/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()
|