From aac1b25442e96d9b2dd4a314db3f62eea594f8b3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 29 May 2026 19:29:35 +0000 Subject: [PATCH] =?UTF-8?q?test:=20TMA=20QK=20diagnostic=20=E2=80=94=203?= =?UTF-8?q?=20variants=20to=20isolate=20failure?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_tma_qk_diag.cu | 367 +++++++++++++++++++++++++++++++++ 1 file changed, 367 insertions(+) create mode 100644 tests/unit/test_tma_qk_diag.cu diff --git a/tests/unit/test_tma_qk_diag.cu b/tests/unit/test_tma_qk_diag.cu new file mode 100644 index 00000000..d930bfc9 --- /dev/null +++ b/tests/unit/test_tma_qk_diag.cu @@ -0,0 +1,367 @@ +/** + * TMA + QK — test if the mbarrier wait in multi-warp is the issue. + * Approach: only thread 0 issues TMA AND waits, then __syncthreads(). + * All other threads skip the mbarrier entirely. + */ + +#include +#include +#include +#include +#include +#include + +#ifndef HD_VAL +#define HD_VAL 64 +#endif + +typedef unsigned short bf16_t; +static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); } +static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; } + +constexpr int HD = HD_VAL; +constexpr int SK = 128; +constexpr int MMA_K = 16; +constexpr int BLOCK_MN = 128; +constexpr int TILE_SZ = 128 * MMA_K; +constexpr int TMEM_N = (HD <= 128) ? 128 : 256; +constexpr int NKT = HD / MMA_K; +constexpr int CORES_MN = 16; +constexpr int NUM_READS = SK / 8; + +__device__ __forceinline__ void tma_mbarrier_init(uint32_t smem_mbar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(smem_mbar), "r"(count)); +} +__device__ __forceinline__ void tma_mbarrier_arrive_expect_tx(uint32_t smem_mbar, uint32_t tx_bytes) { + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(smem_mbar), "r"(tx_bytes) : "memory"); +} +__device__ __forceinline__ void tma_load_2d(uint32_t dst, uint64_t desc, uint32_t mbar, int cx, int cy) { + asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes " + "[%0], [%1, {%3, %4}], [%2];" :: "r"(dst), "l"(desc), "r"(mbar), "r"(cx), "r"(cy) : "memory"); +} +__device__ __forceinline__ void tma_mbarrier_wait(uint32_t smem_mbar, int phase) { + asm volatile("{\n\t.reg .pred P1;\n\tLAB_WAIT:mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t@P1 bra.uni DONE;\n\tbra.uni LAB_WAIT;\n\tDONE:\n\t}" :: "r"(smem_mbar), "r"(phase), "r"(0x989680) : "memory"); +} + +__device__ __forceinline__ bf16_t f32_to_bf16(float f) { bf16_t h; asm("cvt.rn.bf16.f32 %0, %1;" : "=h"(h) : "f"(f)); return h; } +__device__ __forceinline__ float bf16_to_f32(bf16_t h) { float f; asm("cvt.f32.bf16 %0, %1;" : "=f"(f) : "h"(h)); return f; } + +__device__ void tmem_alloc(uint32_t smem_ptr, int n) { + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem_ptr), "r"(n)); +} +__device__ void tmem_dealloc(uint32_t tmem_ptr, int n) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_ptr), "r"(n)); +} + +__device__ __forceinline__ uint64_t desc_encode(uint64_t b) { return b >> 4; } +__device__ __forceinline__ uint64_t make_umma_desc_kmajor_none(uint32_t smem_addr, int block_mn) { + uint64_t d = 0; + d |= desc_encode(smem_addr) & 0x3FFF; + d |= (desc_encode(block_mn * 16) & 0x3FFF) << 16; + d |= (desc_encode(128) & 0x3FFF) << 32; + d |= 1ULL << 46; + return d; +} +__device__ __forceinline__ uint32_t make_idesc(int bm, int bn) { + return (1U<<4)|(1U<<7)|(1U<<10)|((uint32_t)(bn>>3)<<17)|((uint32_t)(bm>>4)<<24); +} +__device__ void umma_ss_f16(uint32_t tc, uint64_t da, uint64_t db, uint32_t idesc, bool acc) { + uint32_t sb = acc ? 0x3F800000u : 0u; + asm volatile("{\n\t.reg .pred p;\n\tsetp.ne.b32 p, %4, 0;\n\ttcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, {%5, %6, %7, %8}, p;\n\t}" + :: "r"(tc), "l"(da), "l"(db), "r"(idesc), "r"(sb), "r"(0), "r"(0), "r"(0), "r"(0)); +} + +// Test A: TMA + mbarrier, ALL threads wait (original failing pattern) +__global__ void __launch_bounds__(128) +test_a_mbarrier_allwait( + float* out_s, const bf16_t* q, CUtensorMap* tma_k, int s_k +) { + const int tid = threadIdx.x; + const int wid = tid / 32, lane = tid % 32; + + 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; + bf16_t* sQ0 = (bf16_t*)(sbuf + off); off += TILE_SZ * 2; + bf16_t* sK0 = (bf16_t*)(sbuf + off); off += TILE_SZ * 2; + bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TILE_SZ * 2; + off = (off + 15) & ~(size_t)15; + uint64_t* sMbar = (uint64_t*)(sbuf + off); off += 8; + + if (wid == 1) tmem_alloc(__cvta_generic_to_shared(sTmemBase), TMEM_N); + if (tid == 0) { tma_mbarrier_init((uint32_t)__cvta_generic_to_shared(sMbar), 1); asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); } + __syncthreads(); + uint32_t tb = *sTmemBase; + uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar); + int phase = 0; + + uint32_t idesc = make_idesc(BLOCK_MN, BLOCK_MN); + for (int kt = 0; kt < NKT; kt++) { + // Q: direct load + for (int i = tid; i < TILE_SZ; i += 128) sQ0[i] = 0; + for (int d = tid; d < MMA_K; d += 128) { + int fd = kt * MMA_K + d; + if (fd < HD) { int ck = d/8, lc = d%8; sQ0[ck*CORES_MN*64+lc] = q[fd]; } + } + __syncthreads(); + + // K: TMA, ALL threads wait + if (tid == 0) { + tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_k, mbar_addr, kt*MMA_K, 0); + tma_mbarrier_arrive_expect_tx(mbar_addr, TILE_SZ * sizeof(bf16_t)); + } + tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; + __syncthreads(); + + // canonical write + for (int i = tid; i < TILE_SZ; i += 128) sK0[i] = 0; + for (int i = tid; i < s_k * MMA_K; i += 128) { + int r = i / MMA_K, c = i % MMA_K; + sK0[(c/8)*CORES_MN*64 + (r/8)*64 + (r%8)*8 + c%8] = sTmaBuf[i]; + } + __syncthreads(); + + // MMA + if (tid == 0) { + umma_ss_f16(tb, make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ0), BLOCK_MN), + make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK0), BLOCK_MN), idesc, kt > 0); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + } + __syncthreads(); + } + asm volatile("fence.sc.gpu;" ::: "memory"); __syncthreads(); + + // Read S + if (wid == 0) { + for (int n = 0; n < NUM_READS; n++) { + 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 + n * 8)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (lane == 0) for (int c=0;c<8;c++) { int col=n*8+c; if(col 0); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + } + __syncthreads(); + } + asm volatile("fence.sc.gpu;" ::: "memory"); __syncthreads(); + + if (wid == 0) { + for (int n = 0; n < NUM_READS; n++) { + 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 + n * 8)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (lane == 0) for (int c=0;c<8;c++) { int col=n*8+c; if(col 0); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + } + __syncthreads(); + } + asm volatile("fence.sc.gpu;" ::: "memory"); __syncthreads(); + + if (wid == 0) { + for (int n = 0; n < NUM_READS; n++) { + 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 + n * 8)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + if (lane == 0) for (int c=0;c<8;c++) { int col=n*8+c; if(col(ptr), gd, gs, td, ts, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (r!=CUDA_SUCCESS) { fprintf(stderr,"TMA fail %d\n",(int)r); return false; } + int dv=0; cudaDriverGetVersion(&dv); + if (dv<=13010 && rows*cols*2<131072) reinterpret_cast(out)[1] &= ~(1ULL<<21); + return true; +} + +int main() { + printf("=== TMA QK diagnostic (HD=%d) ===\n", HD); + bf16_t* h_q = (bf16_t*)calloc(HD, sizeof(bf16_t)); + bf16_t* h_k = (bf16_t*)calloc(SK*HD, sizeof(bf16_t)); + srand(42); + for (int i=0;i1e-4f ? fabsf(h[j]-s_ref[j])/fabsf(s_ref[j]) : fabsf(h[j]-s_ref[j]); + if (rel>mr) mr=rel; + if (rel>0.01f) bad++; + } + printf(" %s: max_rel=%.8f bad=%d %s\n", name, mr, bad, bad==0?"PASS":"FAIL"); + free(h); + }; + + // Test A: mbarrier, all wait + printf("Test A: mbarrier all-wait\n"); + cudaMemset(d_out, 0, SK*4); + cudaFuncSetAttribute(test_a_mbarrier_allwait, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_a); + test_a_mbarrier_allwait<<<1,128,smem_a>>>(d_out, d_q, d_tma_k, SK); + check("A", cudaDeviceSynchronize()); + + // Test B: mbarrier, tid0 only wait + printf("Test B: mbarrier tid0-wait\n"); + cudaMemset(d_out, 0, SK*4); + cudaFuncSetAttribute(test_b_mbarrier_tid0wait, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_a); + test_b_mbarrier_tid0wait<<<1,128,smem_a>>>(d_out, d_q, d_tma_k, SK); + check("B", cudaDeviceSynchronize()); + + // Test C: direct load (baseline) + printf("Test C: direct load (no TMA)\n"); + cudaMemset(d_out, 0, SK*4); + cudaFuncSetAttribute(test_c_direct_load, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_c); + test_c_direct_load<<<1,128,smem_c>>>(d_out, d_q, d_k, SK); + check("C", cudaDeviceSynchronize()); + + cudaFree(d_q); cudaFree(d_k); cudaFree(d_out); cudaFree(d_tma_k); + free(h_q); free(h_k); + return 0; +}