diff --git a/dsv4/kernels/cuda/test_fp8_gemm_tmem_read.cu b/dsv4/kernels/cuda/test_fp8_gemm_tmem_read.cu deleted file mode 100644 index 0da87921..00000000 --- a/dsv4/kernels/cuda/test_fp8_gemm_tmem_read.cu +++ /dev/null @@ -1,211 +0,0 @@ -/** - * B2 debug: FP8 GEMM + TMEM read test kernel. - * Produces raw dequantized logits for comparison with FP32 reference. - */ -#include -#include -#include -#include -#include -#include -#include -#include -#include - -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(&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(q_bf16.data_ptr()), - k8.data_ptr(), k_scale.data_ptr(), - logits_out.data_ptr(), 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"); -} diff --git a/tests/unit/test_b1_debug_cosine.py b/tests/unit/test_b1_debug_cosine.py deleted file mode 100644 index 7821e74c..00000000 --- a/tests/unit/test_b1_debug_cosine.py +++ /dev/null @@ -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() diff --git a/tests/unit/test_b1_isolate_qk_pv.py b/tests/unit/test_b1_isolate_qk_pv.py deleted file mode 100644 index 866f6c07..00000000 --- a/tests/unit/test_b1_isolate_qk_pv.py +++ /dev/null @@ -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() diff --git a/tests/unit/test_b2_debug_qk.py b/tests/unit/test_b2_debug_qk.py deleted file mode 100644 index 47944d73..00000000 --- a/tests/unit/test_b2_debug_qk.py +++ /dev/null @@ -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() diff --git a/tests/unit/test_b2_debug_scores.py b/tests/unit/test_b2_debug_scores.py deleted file mode 100644 index 79a31723..00000000 --- a/tests/unit/test_b2_debug_scores.py +++ /dev/null @@ -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() diff --git a/tests/unit/test_b2_minimal.py b/tests/unit/test_b2_minimal.py deleted file mode 100644 index 242dc1f2..00000000 --- a/tests/unit/test_b2_minimal.py +++ /dev/null @@ -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() diff --git a/tests/unit/test_b2_tmem_read.py b/tests/unit/test_b2_tmem_read.py deleted file mode 100644 index a343489d..00000000 --- a/tests/unit/test_b2_tmem_read.py +++ /dev/null @@ -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()