From 5932e928a879ec77cf5ddb9dad24ddcec9a5ab26 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 10:46:14 +0000 Subject: [PATCH] cleanup: remove debug test files (P4, P5) --- tests/unit/test_p4_tma_bit21_fix.cu | 117 ----------- tests/unit/test_p4_tma_descriptor_diff.py | 130 ------------ tests/unit/test_p4_tma_descriptor_dump.cu | 87 -------- tests/unit/test_p4_tma_load_test.cu | 148 -------------- tests/unit/test_p5_multitile.cu | 229 ---------------------- tests/unit/test_p5_python_minimal.py | 35 ---- tests/unit/test_p5_tma_multitile.cu | 121 ------------ 7 files changed, 867 deletions(-) delete mode 100644 tests/unit/test_p4_tma_bit21_fix.cu delete mode 100644 tests/unit/test_p4_tma_descriptor_diff.py delete mode 100644 tests/unit/test_p4_tma_descriptor_dump.cu delete mode 100644 tests/unit/test_p4_tma_load_test.cu delete mode 100644 tests/unit/test_p5_multitile.cu delete mode 100644 tests/unit/test_p5_python_minimal.py delete mode 100644 tests/unit/test_p5_tma_multitile.cu diff --git a/tests/unit/test_p4_tma_bit21_fix.cu b/tests/unit/test_p4_tma_bit21_fix.cu deleted file mode 100644 index 219d1261..00000000 --- a/tests/unit/test_p4_tma_bit21_fix.cu +++ /dev/null @@ -1,117 +0,0 @@ -/** - * P4: Test TMA with bit-21 workaround. - * Uses outermost-first dims, 2 strides, 128B-aligned GMEM. - */ -#include -#include -#include -#include -#include - -__global__ void tma_load_kernel(const void* tma_desc_ptr, int* result) { - __shared__ uint64_t mbar; - __shared__ uint16_t smem_out[16 * 16] __attribute__((aligned(128))); - - if (threadIdx.x == 0) { - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - asm volatile("mbarrier.init.shared.b64 [%0], 1;" :: "r"(mbar_addr)); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - __syncthreads(); - - if (threadIdx.x == 0) { - uint32_t smem_addr = __cvta_generic_to_shared(smem_out); - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes " - "[%0], [%1, {%3, %4}], [%2];" - :: "r"(smem_addr), "l"(tma_desc_ptr), "r"(mbar_addr), - "r"(0), "r"(0) - ); - } - __syncthreads(); - - if (threadIdx.x == 0) { - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - for (int w = 0; w < 1000000; w++) { - uint32_t state; - asm volatile("{\n\t.reg .pred p;\n\t" - "mbarrier.try_wait.parity.shared.b64 p, [%0], 0;\n\t" - "selp.b32 %1, 1, 0, p;\n\t}" : "=r"(state) : "r"(mbar_addr)); - if (state) { *result = 1; return; } - } - *result = -1; - } -} - -int main() { - const int ROWS = 128, COLS = 16; - const size_t DATA_SIZE = ROWS * COLS * 2; - - // Allocate 128B-aligned GMEM - void* d_data; - cudaMalloc(&d_data, DATA_SIZE + 128); - // Align to 128 bytes - uintptr_t aligned = ((uintptr_t)d_data + 127) & ~127ULL; - void* d_aligned = (void*)aligned; - cudaMemset(d_aligned, 1, DATA_SIZE); - - int* d_result; - cudaMalloc(&d_result, sizeof(int)); - - // Create descriptor: outermost-first, byte strides - auto make_desc = [&](void* ptr, bool fix_bit21) -> CUtensorMap { - CUtensorMap desc; - cuuint64_t tensorDims[] = {(cuuint64_t)ROWS, (cuuint64_t)COLS}; - cuuint64_t globalStrides[] = {(cuuint64_t)(COLS * 2), (cuuint64_t)2}; - cuuint32_t boxDims[] = {16, 16}; - cuuint32_t elementStrides[] = {1, 1}; - CUresult res = cuTensorMapEncodeTiled(&desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - ptr, tensorDims, globalStrides, boxDims, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - if (res != CUDA_SUCCESS) printf(" cuTensorMapEncodeTiled FAILED: %d\n", res); - if (fix_bit21) { - uint64_t* w = reinterpret_cast(&desc); - printf(" Before bit21 fix: word[1] = 0x%016lx\n", w[1]); - w[1] &= ~(1ULL << 21); - printf(" After bit21 fix: word[1] = 0x%016lx\n", w[1]); - } - return desc; - }; - - auto test_desc = [&](const char* name, CUtensorMap& desc) { - printf("=== %s ===\n", name); - auto* b = reinterpret_cast(&desc); - printf(" [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n"); - printf(" [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n"); - - void* d_desc; - cudaMalloc(&d_desc, sizeof(CUtensorMap)); - cudaMemcpy(d_desc, &desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - cudaMemset(d_result, 0, sizeof(int)); - - tma_load_kernel<<<1, 32>>>(d_desc, d_result); - cudaError_t err = cudaDeviceSynchronize(); - int r; cudaMemcpy(&r, d_result, sizeof(int), cudaMemcpyDeviceToHost); - - if (err != CUDA_SUCCESS) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), r); - else if (r == 1) printf(" SUCCESS\n"); - else if (r == -1) printf(" HANG (mbarrier timeout)\n"); - else printf(" result=%d\n", r); - - cudaFree(d_desc); - }; - - CUtensorMap desc1 = make_desc(d_aligned, false); - test_desc("Original (no fix)", desc1); - - CUtensorMap desc2 = make_desc(d_aligned, true); - test_desc("Bit-21 cleared", desc2); - - cudaFree(d_data); - cudaFree(d_result); - printf("\nPASSED\n"); - return 0; -} diff --git a/tests/unit/test_p4_tma_descriptor_diff.py b/tests/unit/test_p4_tma_descriptor_diff.py deleted file mode 100644 index 84493873..00000000 --- a/tests/unit/test_p4_tma_descriptor_diff.py +++ /dev/null @@ -1,130 +0,0 @@ -""" -P4: Dump TMA descriptor bytes from both CuTeDSL and cuTensorMapEncodeTiled. - -1. CuTeDSL: create a TMA descriptor for a (128,16) BF16 tensor via cute.compile - and dump the 128 bytes. -2. Driver API: use cuTensorMapEncodeTiled for the same tensor and dump 128 bytes. -3. memcmp and print differences. - -The CuTeDSL path already works (it's used in the existing FMHA kernel). -The raw Driver API path hangs when used with cp.async.bulk.tensor.2d. -By comparing descriptors byte-by-byte, we can identify the field that differs. -""" -import torch -import sys -import os -import struct -import numpy as np - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - - -def dump_driver_api_descriptor(): - """Create TMA descriptor using Driver API (cuTensorMapEncodeTiled).""" - from cuda.bindings import driver, runtime - - # Initialize CUDA - runtime.cudaFree(0) # Force context creation - - # Create a (128, 16) BF16 tensor on GPU - rows, cols = 128, 16 - data = torch.zeros(rows, cols, dtype=torch.bfloat16, device='cuda') - - # cuTensorMap descriptor: 128 bytes - tensor_map = driver.CUtensorMap() - - # cuTensorMapEncodeTiled args: - # - tensorMap: output - # - tensorRank: 2 - # - cudaDataType: CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 (6) - # - deviceAddress: data pointer - # - tensorDims: [128, 16] - # - globalStrides: [16*2, 2] (byte strides: row_stride=16*2 bytes, col_stride=2 bytes) - # - boxDims: [16, 16] (TMA tile size) - # - elementStrides: [1, 1] - # - interleave: CU_TENSOR_MAP_INTERLEAVE_NONE (0) - # - swizzle: CU_TENSOR_MAP_SWIZZLE_NONE (0) - # - l2Promotion: CU_TENSOR_MAP_L2_PROMOTION_NONE (0) - # - oobFill: CU_TENSOR_MAP_OOB_FILL_NONE (0) - - globalStrides = (ctypes.c_uint64 * 2)() - globalStrides[0] = cols * 2 # stride from row 0 to row 1 = 16 * 2 = 32 bytes - globalStrides[1] = 2 # stride from col 0 to col 1 = 2 bytes - - import ctypes - tensorDims = (ctypes.c_uint32 * 2)(rows, cols) - boxDims = (ctypes.c_uint32 * 2)(16, 16) - elementStrides = (ctypes.c_uint32 * 2)(1, 1) - - # Actually, let me use the cuda.bindings API directly - # cuTensorMapEncodeTiled is in cuda.bindings.driver - - result = driver.cuTensorMapEncodeTiled( - tensor_map, - 2, # tensorRank - driver.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - int(data.data_ptr()), - (rows, cols), - (cols * 2, 2), # globalStrides in bytes - (16, 16), # boxDims - (1, 1), # elementStrides - driver.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, - driver.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE, - driver.CUtensorMapL2Promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE, - driver.CUtensorMapOOBFill.CU_TENSOR_MAP_OOB_FILL_NONE, - ) - - if result != driver.CUresult.CUDA_SUCCESS: - print(f"cuTensorMapEncodeTiled failed: {result}") - return None - - # The descriptor is 128 bytes. Access it via the opaque field. - # CUtensorMap has an opaque byte array - desc_bytes = bytes(tensor_map) - return desc_bytes - - -def dump_cutedsl_descriptor(): - """Create TMA descriptor using CuTeDSL and dump bytes. - - CuTeDSL creates descriptors internally when you call cute.make_tensor - with a TMA layout. We need to intercept the descriptor bytes. - - Actually, CuTeDSL's TMA descriptors are created at JIT compile time - and stored in the kernel's parameter struct. We can't easily dump them - from Python. - - Alternative: use CuTe's TMA descriptor creation API directly. - cute.arch.make_tma_copy can create a descriptor that we can inspect. - """ - # This is harder than I thought. CuTeDSL hides the descriptor creation. - # Let me use a different approach: create a small CuTeDSL kernel that - # does a TMA load (which works), and use Nsight to capture the descriptor. - # Or: use the CUTLASS Python API directly. - - # Actually, the simplest approach: use the CUTLASS Python bindings - # that CuTeDSL uses internally. The TMA descriptor is a Python object - # before being passed to the kernel. - pass - - -def main(): - print("P4: TMA Descriptor Comparison") - print("=" * 60) - - # Step 1: Driver API descriptor - print("\n1. Driver API (cuTensorMapEncodeTiled) descriptor:") - desc_driver = dump_driver_api_descriptor() - if desc_driver is not None: - for i in range(0, 128, 16): - hex_str = ' '.join(f'{b:02x}' for b in desc_driver[i:i+16]) - print(f" [{i:3d}-{i+15:3d}]: {hex_str}") - else: - print(" FAILED to create descriptor") - - print("\nNote: CuTeDSL descriptor dump requires running inside a JIT kernel.") - print("Use the CUDA test (test_p4_tma_descriptor_diff.cu) for the full comparison.") - - -if __name__ == "__main__": - main() diff --git a/tests/unit/test_p4_tma_descriptor_dump.cu b/tests/unit/test_p4_tma_descriptor_dump.cu deleted file mode 100644 index f973ef0e..00000000 --- a/tests/unit/test_p4_tma_descriptor_dump.cu +++ /dev/null @@ -1,87 +0,0 @@ -/** - * P4: Dump TMA descriptor bytes for comparison. - * CUDA 13.2 compatible. - * - * Signature: cuTensorMapEncodeTiled( - * CUtensorMap*, CUtensorMapDataType, cuuint32_t tensorRank, - * void*, cuuint64_t* globalDim, cuuint64_t* globalStrides, - * cuuint32_t* boxDim, cuuint32_t* elementStrides, - * CUtensorMapInterleave, CUtensorMapSwizzle, - * CUtensorMapL2promotion, CUtensorMapFloatOOBfill) - */ -#include -#include -#include -#include - -int main() { - const int ROWS = 128; - const int COLS = 16; - const size_t SIZE = ROWS * COLS * 2; - - void* d_ptr; - cudaMalloc(&d_ptr, SIZE); - cudaMemset(d_ptr, 0, SIZE); - - // globalDim: tensor dimensions (ROWS, COLS) in elements - cuuint64_t globalDim[] = {(cuuint64_t)ROWS, (cuuint64_t)COLS}; - // globalStrides: byte strides between rows and between elements - cuuint64_t globalStrides[] = {(cuuint64_t)(COLS * 2), (cuuint64_t)2}; - // boxDim: TMA tile dimensions (16, 16) - cuuint32_t boxDim[] = {16, 16}; - // elementStrides: (1, 1) = contiguous - cuuint32_t elementStrides[] = {1, 1}; - - CUtensorMap tma_desc; - CUresult res; - - auto dump = [](const char* label, const CUtensorMap& desc) { - printf("=== %s ===\n", label); - auto* b = reinterpret_cast(&desc); - for (int i = 0; i < 128; i += 16) { - printf("[%3d-%3d]: ", i, i+15); - for (int j = 0; j < 16; j++) printf("%02x ", b[i+j]); - printf("\n"); - } - }; - - // 1: NO swizzle, OOB_NONE - res = cuTensorMapEncodeTiled(&tma_desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - if (res == CUDA_SUCCESS) dump("NO swizzle, OOB_NONE", tma_desc); - else printf("=== NO swizzle, OOB_NONE: FAILED (%d) ===\n", res); - - // 2: SWIZZLE_128B, OOB_NONE - res = cuTensorMapEncodeTiled(&tma_desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - if (res == CUDA_SUCCESS) dump("SWIZZLE_128B, OOB_NONE", tma_desc); - else printf("=== SWIZZLE_128B, OOB_NONE: FAILED (%d) ===\n", res); - - // 3: NO swizzle, OOB_FILL_ZERO - res = cuTensorMapEncodeTiled(&tma_desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA); - if (res == CUDA_SUCCESS) dump("NO swizzle, OOB_FILL_ZERO", tma_desc); - else printf("=== NO swizzle, OOB_FILL_ZERO: FAILED (%d) ===\n", res); - - // 4: SWIZZLE_128B, OOB_FILL_ZERO - res = cuTensorMapEncodeTiled(&tma_desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA); - if (res == CUDA_SUCCESS) dump("SWIZZLE_128B, OOB_FILL_ZERO", tma_desc); - else printf("=== SWIZZLE_128B, OOB_FILL_ZERO: FAILED (%d) ===\n", res); - - cudaFree(d_ptr); - printf("\nPASSED\n"); - return 0; -} diff --git a/tests/unit/test_p4_tma_load_test.cu b/tests/unit/test_p4_tma_load_test.cu deleted file mode 100644 index 935775fd..00000000 --- a/tests/unit/test_p4_tma_load_test.cu +++ /dev/null @@ -1,148 +0,0 @@ -/** - * P4: Test actual TMA loads with different descriptors. - * - * Creates TMA descriptors with various swizzle/OOB configs, - * launches a kernel that does cp.async.bulk.tensor.2d with each, - * and checks if the load completes (mbarrier signals) or hangs. - */ -#include -#include -#include -#include -#include - -#define MAX_WAIT 1000000 - -__global__ void tma_load_test_kernel( - const void* tma_desc_ptr, - int* result -) { - // SMEM: mbarrier (8 bytes aligned) + output buffer (512 bytes) - __shared__ uint64_t mbar; - __shared__ uint16_t smem_out[16 * 16]; // 256 BF16 values - - if (threadIdx.x == 0) { - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - asm volatile("mbarrier.init.shared.b64 [%0], 1;" :: "r"(mbar_addr)); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - __syncthreads(); - - if (threadIdx.x == 0) { - uint32_t smem_addr = __cvta_generic_to_shared(smem_out); - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - - // Issue TMA load: 16x16 BF16 = 512 bytes - asm volatile( - "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes " - "[%0], [%1, {%3, %4}], [%2];" - :: "r"(smem_addr), - "l"(tma_desc_ptr), - "r"(mbar_addr), - "r"(0), // coord row = 0 - "r"(0) // coord col = 0 - ); - } - __syncthreads(); - - // Wait for mbarrier - if (threadIdx.x == 0) { - uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - int waited = 0; - int arrived = 0; - while (waited < MAX_WAIT) { - uint32_t state; - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "mbarrier.try_wait.parity.shared.b64 p, [%0], 0;\n\t" - "selp.b32 %1, 1, 0, p;\n\t" - "}" - : "=r"(state) - : "r"(mbar_addr) - ); - if (state) { arrived = 1; break; } - waited++; - } - *result = arrived ? 1 : -1; - } -} - - -CUtensorMap create_descriptor(void* d_ptr, int swizzle, int oob_fill) { - CUtensorMap desc; - cuuint64_t globalDim[] = {128, 16}; - cuuint64_t globalStrides[] = {16 * 2, 2}; - cuuint32_t boxDim[] = {16, 16}; - cuuint32_t elementStrides[] = {1, 1}; - - CUtensorMapSwizzle sw = (swizzle == 0) ? CU_TENSOR_MAP_SWIZZLE_NONE : CU_TENSOR_MAP_SWIZZLE_128B; - CUtensorMapFloatOOBfill oob = (oob_fill == 0) ? CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE - : CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA; - - CUresult res = cuTensorMapEncodeTiled(&desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, sw, - CU_TENSOR_MAP_L2_PROMOTION_NONE, oob); - - if (res != CUDA_SUCCESS) { - printf(" Descriptor creation FAILED: %d\n", res); - } - return desc; -} - - -int main() { - const int ROWS = 128; - const int COLS = 16; - const size_t DATA_SIZE = ROWS * COLS * 2; - - void* d_data; - cudaMalloc(&d_data, DATA_SIZE); - cudaMemset(d_data, 1, DATA_SIZE); - - int* d_result; - cudaMalloc(&d_result, sizeof(int)); - - struct { const char* name; int swizzle; int oob; } configs[] = { - {"NO swizzle, OOB_NONE", 0, 0}, - {"SWIZZLE_128B, OOB_NONE", 1, 0}, - {"NO swizzle, OOB_FILL_ZERO", 0, 1}, - {"SWIZZLE_128B, OOB_FILL_ZERO", 1, 1}, - }; - - for (int i = 0; i < 4; i++) { - printf("Testing: %s\n", configs[i].name); - - CUtensorMap desc = create_descriptor(d_data, configs[i].swizzle, configs[i].oob); - void* d_desc; - cudaMalloc(&d_desc, sizeof(CUtensorMap)); - cudaMemcpy(d_desc, &desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - - cudaMemset(d_result, 0, sizeof(int)); - - tma_load_test_kernel<<<1, 32>>>(d_desc, d_result); - - cudaError_t err = cudaDeviceSynchronize(); - int h_result; - cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost); - - if (err != cudaSuccess) { - printf(" HANG or ERROR: %s (result=%d)\n", cudaGetErrorString(err), h_result); - } else if (h_result == 1) { - printf(" SUCCESS: TMA load completed\n"); - } else if (h_result == -1) { - printf(" HANG: mbarrier never signaled\n"); - } else { - printf(" UNKNOWN: result=%d\n", h_result); - } - - cudaFree(d_desc); - } - - cudaFree(d_data); - cudaFree(d_result); - printf("\nPASSED\n"); - return 0; -} diff --git a/tests/unit/test_p5_multitile.cu b/tests/unit/test_p5_multitile.cu deleted file mode 100644 index f8755f0e..00000000 --- a/tests/unit/test_p5_multitile.cu +++ /dev/null @@ -1,229 +0,0 @@ -/** - * P5: Minimal multi-tile FMHA test. - * Tests the multi-tile path (n_kv_tiles=2, N=256) against a CPU reference. - */ -#include "dsv4/kernels/attention/fmha_common.cuh" -#include "dsv4/kernels/attention/fmha_umma_desc.cuh" -#include "dsv4/kernels/attention/fmha_6warp_multihead.cuh" - -#include -#include -#include - -using namespace dsv4::kernels::attention; - -// Host-side BF16 helpers -static float hbf16_to_f32(uint16_t h) { - uint32_t u = ((uint32_t)h) << 16; - float f; memcpy(&f, &u, 4); return f; -} -static uint16_t hf32_to_bf16(float f) { - uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16); -} - -// CPU reference attention for single head -void reference_attention( - const bf16_t* q, const bf16_t* k, const bf16_t* v, - float* o_ref, float* lse_ref, - int hd, int s_k, float scale -) { - for (int d = 0; d < hd; d++) o_ref[d] = 0.0f; - float row_max = -INFINITY; - // Compute max - for (int j = 0; j < s_k; j++) { - float dot = 0.0f; - for (int d = 0; d < hd; d++) dot += hbf16_to_f32(q[d]) * hbf16_to_f32(k[j * hd + d]); - dot *= scale; - if (dot > row_max) row_max = dot; - } - // Compute exp and weighted sum - float row_sum = 0.0f; - for (int j = 0; j < s_k; j++) { - float dot = 0.0f; - for (int d = 0; d < hd; d++) dot += hbf16_to_f32(q[d]) * hbf16_to_f32(k[j * hd + d]); - dot *= scale; - float p = expf(dot - row_max); - row_sum += p; - for (int d = 0; d < hd; d++) o_ref[d] += p * hbf16_to_f32(v[d * s_k + j]); - } - for (int d = 0; d < hd; d++) o_ref[d] /= row_sum; - *lse_ref = logf(row_sum) + row_max; -} - -int main() { - constexpr int HD = 64; - constexpr int SK = 256; // 2 KV tiles - const float SCALE = 1.0f / sqrtf((float)HD); - - // Allocate host data - bf16_t h_q[HD], h_k[SK * HD], h_v[HD * SK]; - float h_o_ref[HD], h_lse_ref; - - srand(42); - for (int d = 0; d < HD; d++) h_q[d] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - for (int j = 0; j < SK * HD; j++) h_k[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - for (int j = 0; j < HD * SK; j++) h_v[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - - // CPU reference - reference_attention(h_q, h_k, h_v, h_o_ref, &h_lse_ref, HD, SK, SCALE); - - // GPU test - bf16_t *d_q, *d_k, *d_v, *d_o; - float *d_lse; - cudaMalloc(&d_q, HD * 2); - cudaMalloc(&d_k, SK * HD * 2); - cudaMalloc(&d_v, HD * SK * 2); - cudaMalloc(&d_o, HD * 2); - cudaMalloc(&d_lse, 4); - - cudaMemcpy(d_q, h_q, HD * 2, cudaMemcpyHostToDevice); - cudaMemcpy(d_k, h_k, SK * HD * 2, cudaMemcpyHostToDevice); - cudaMemcpy(d_v, h_v, HD * SK * 2, cudaMemcpyHostToDevice); - cudaMemset(d_o, 0, HD * 2); - cudaMemset(d_lse, 0, 4); - - FmhaParams params; - params.q = d_q; - params.k = d_k; - params.v = d_v; - params.o = d_o; - params.lse = d_lse; - params.s_k = SK; - params.scale = SCALE; - params.head_dim = HD; - params.n_kv_tiles = 2; // Multi-tile! - params.q_head_stride = 0; - params.q_batch_stride = 0; - params.k_head_stride = 0; - params.k_batch_stride = 0; - params.v_head_stride = 0; - params.v_batch_stride = 0; - params.o_head_stride = 0; - params.o_batch_stride = 0; - params.lse_head_stride = 0; - params.lse_batch_stride = 0; - - // SMEM size - constexpr int TILE_SZ = 128 * MMA_K_BF16; - constexpr int V_SUB_SZ = 256; - int smem = 4 + 4 + 4 + 16 + TILE_SZ*2 + TILE_SZ*2 + TILE_SZ*2 + V_SUB_SZ*2 + 128*4 + HD*4 + 256 + 127; - smem &= ~127; - - dim3 grid(1, 1, 1); - dim3 block(NTHREADS); - - if (smem > 48 * 1024) { - cudaFuncSetAttribute(fmha_6warp_multihead_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - } - - fmha_6warp_multihead_kernel<<>>(params); - cudaError_t err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { - printf("Kernel launch FAILED: %s\n", cudaGetErrorString(err)); - return 1; - } - - // Compare - bf16_t h_o[HD]; - float h_lse; - cudaMemcpy(h_o, d_o, HD * 2, cudaMemcpyDeviceToHost); - cudaMemcpy(&h_lse, d_lse, 4, cudaMemcpyDeviceToHost); - - float cos = 0, norm_a = 0, norm_b = 0; - for (int d = 0; d < HD; d++) { - float a = h_o_ref[d]; - float b = hbf16_to_f32(h_o[d]); - cos += a * b; - norm_a += a * a; - norm_b += b * b; - } - cos /= sqrtf(norm_a * norm_b + 1e-30f); - - printf("Multi-tile FMHA (HD=%d, SK=%d, n_kv_tiles=2):\n", HD, SK); - printf(" LSE: kernel=%.4f ref=%.4f\n", h_lse, h_lse_ref); - printf(" Cosine similarity: %.6f\n", cos); - printf(" Kernel O[0:5]:"); - for (int d = 0; d < 5; d++) printf(" %.4f", hbf16_to_f32(h_o[d])); - printf("\n Ref O[0:5]:"); - for (int d = 0; d < 5; d++) printf(" %.4f", h_o_ref[d]); - printf("\n"); - - // Also test: 2 separate single-tile launches + Python merge - bf16_t *d_k1, *d_v1, *d_k2, *d_v2; - bf16_t *d_o1, *d_o2; - float *d_lse1, *d_lse2; - cudaMalloc(&d_k1, 128 * HD * 2); - cudaMalloc(&d_k2, 128 * HD * 2); - cudaMalloc(&d_v1, HD * 128 * 2); - cudaMalloc(&d_v2, HD * 128 * 2); - cudaMalloc(&d_o1, HD * 2); - cudaMalloc(&d_o2, HD * 2); - cudaMalloc(&d_lse1, 4); - cudaMalloc(&d_lse2, 4); - - cudaMemcpy(d_k1, d_k, 128 * HD * 2, cudaMemcpyDeviceToDevice); - cudaMemcpy(d_k2, d_k + 128 * HD, 128 * HD * 2, cudaMemcpyDeviceToDevice); - cudaMemcpy(d_v1, d_v, HD * 128 * 2, cudaMemcpyDeviceToDevice); - cudaMemcpy(d_v2, d_v + HD * 128, HD * 128 * 2, cudaMemcpyDeviceToDevice); - - // Run single-tile on first half - FmhaParams p1 = params; - p1.k = d_k1; p1.v = d_v1; p1.o = d_o1; p1.lse = d_lse1; - p1.s_k = 128; p1.n_kv_tiles = 1; - fmha_6warp_multihead_kernel<<>>(p1); - - // Run single-tile on second half - FmhaParams p2 = params; - p2.k = d_k2; p2.v = d_v2; p2.o = d_o2; p2.lse = d_lse2; - p2.s_k = 128; p2.n_kv_tiles = 1; - fmha_6warp_multihead_kernel<<>>(p2); - - cudaDeviceSynchronize(); - - // Read single-tile results - bf16_t h_o1[HD], h_o2[HD]; - float h_lse1, h_lse2; - cudaMemcpy(h_o1, d_o1, HD * 2, cudaMemcpyDeviceToHost); - cudaMemcpy(h_o2, d_o2, HD * 2, cudaMemcpyDeviceToHost); - cudaMemcpy(&h_lse1, d_lse1, 4, cudaMemcpyDeviceToHost); - cudaMemcpy(&h_lse2, d_lse2, 4, cudaMemcpyDeviceToHost); - - // Python merge: O = (exp(lse1)*O1 + exp(lse2)*O2) / (exp(lse1) + exp(lse2)) - float e1 = expf(h_lse1), e2 = expf(h_lse2); - float h_o_merge[HD]; - for (int d = 0; d < HD; d++) { - float o1 = hbf16_to_f32(h_o1[d]) * e1; - float o2 = hbf16_to_f32(h_o2[d]) * e2; - h_o_merge[d] = (o1 + o2) / (e1 + e2); - } - - // Compare merge vs reference - float cos_merge = 0, nm_a = 0, nm_b = 0; - for (int d = 0; d < HD; d++) { - cos_merge += h_o_ref[d] * h_o_merge[d]; - nm_a += h_o_ref[d] * h_o_ref[d]; - nm_b += h_o_merge[d] * h_o_merge[d]; - } - cos_merge /= sqrtf(nm_a * nm_b + 1e-30f); - - printf(" Single-tile merge: cos=%.6f lse1=%.4f lse2=%.4f\n", cos_merge, h_lse1, h_lse2); - - // Compare in-kernel multi-tile vs Python merge - float cos_vs_merge = 0; nm_a = 0; nm_b = 0; - for (int d = 0; d < HD; d++) { - float a = h_o_merge[d]; - float b = hbf16_to_f32(h_o[d]); - cos_vs_merge += a * b; - nm_a += a * a; - nm_b += b * b; - } - cos_vs_merge /= sqrtf(nm_a * nm_b + 1e-30f); - printf(" In-kernel vs Python merge: cos=%.6f\n", cos_vs_merge); - - cudaFree(d_k1); cudaFree(d_k2); cudaFree(d_v1); cudaFree(d_v2); - cudaFree(d_o1); cudaFree(d_o2); cudaFree(d_lse1); cudaFree(d_lse2); - - if (cos >= 0.999990) { printf("PASS\n"); return 0; } - else { printf("FAIL (cos < 0.999990)\n"); return 1; } -} diff --git a/tests/unit/test_p5_python_minimal.py b/tests/unit/test_p5_python_minimal.py deleted file mode 100644 index 4f00895a..00000000 --- a/tests/unit/test_p5_python_minimal.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Minimal multi-tile test via Python.""" -import torch, sys, math -sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel') - -from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw - -torch.manual_seed(42) -hd = 64 -N = 256 -scale = 1.0 / math.sqrt(hd) - -q = torch.randn(1, 1, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() -k = torch.randn(1, 1, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() -v = torch.randn(1, 1, hd, N, dtype=torch.bfloat16, device='cuda').contiguous() - -print(f'q align: {q.data_ptr() % 128}, k align: {k.data_ptr() % 128}, v align: {v.data_ptr() % 128}') -print(f'q shape: {q.shape}, k shape: {k.shape}, v shape: {v.shape}') - -try: - o, lse = fmha_multitile_decode_raw(q, k, v, scale) - print(f'Output[0,0,0,:5]: {o[0,0,0,:5].float()}') - print(f'LSE: {lse[0,0,0].item():.4f}') - - # Reference - q_r = q[0,0].float() # (1, hd) - k_r = k[0,0].float() # (N, hd) - v_r = v[0,0].float().T # (N, hd) — V is (hd, N), transpose for reference - s = torch.matmul(q_r, k_r.T) * scale - s = torch.softmax(s, dim=-1) - o_ref = torch.matmul(s, v_r) - cos = torch.nn.functional.cosine_similarity(o[0,0].float().flatten().unsqueeze(0), o_ref.flatten().unsqueeze(0)).item() - print(f'Cosine vs reference: {cos:.6f}') - print(f'{"PASS" if cos >= 0.999990 else "FAIL"}') -except Exception as e: - print(f'FAILED: {e}') diff --git a/tests/unit/test_p5_tma_multitile.cu b/tests/unit/test_p5_tma_multitile.cu deleted file mode 100644 index af9c22a0..00000000 --- a/tests/unit/test_p5_tma_multitile.cu +++ /dev/null @@ -1,121 +0,0 @@ -/** - * P5: Test multi-tile TMA FMHA kernel with proper alignment. - */ -#include -#include -#include "dsv4/kernels/attention/fmha_common.cuh" -#include "dsv4/kernels/attention/fmha_umma_desc.cuh" -#include "dsv4/kernels/attention/fmha_tma.cuh" -#include "dsv4/kernels/attention/fmha_6warp_tma_multirow_multitile.cuh" -#include -#include -#include - -using namespace dsv4::kernels::attention; - -static float hbf16_to_f32(uint16_t h) { uint32_t u = ((uint32_t)h) << 16; float f; memcpy(&f, &u, 4); return f; } -static uint16_t hf32_to_bf16(float f) { uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16); } - -int main() { - constexpr int HD = 64; - constexpr int SK = 256; - const float SCALE = 1.0f / sqrtf((float)HD); - - bf16_t *d_q_raw, *d_k_raw, *d_v_raw, *d_o_raw; - float *d_lse_raw; - cudaMalloc(&d_q_raw, HD * 2 + 128); - cudaMalloc(&d_k_raw, SK * HD * 2 + 128); - cudaMalloc(&d_v_raw, HD * SK * 2 + 128); - cudaMalloc(&d_o_raw, HD * 2 + 128); - cudaMalloc(&d_lse_raw, 4 + 128); - - bf16_t *d_q = (bf16_t*)(((uintptr_t)d_q_raw + 127) & ~(uintptr_t)127); - bf16_t *d_k = (bf16_t*)(((uintptr_t)d_k_raw + 127) & ~(uintptr_t)127); - bf16_t *d_v = (bf16_t*)(((uintptr_t)d_v_raw + 127) & ~(uintptr_t)127); - bf16_t *d_o = (bf16_t*)(((uintptr_t)d_o_raw + 127) & ~(uintptr_t)127); - float *d_lse = (float*)(((uintptr_t)d_lse_raw + 127) & ~(uintptr_t)127); - - srand(42); - bf16_t h_q[HD], h_k[SK * HD], h_v[HD * SK]; - for (int d = 0; d < HD; d++) h_q[d] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - for (int j = 0; j < SK * HD; j++) h_k[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - for (int j = 0; j < HD * SK; j++) h_v[j] = hf32_to_bf16((float)(rand() % 100) / 100.0f); - - cudaMemcpy(d_q, h_q, HD * 2, cudaMemcpyHostToDevice); - cudaMemcpy(d_k, h_k, SK * HD * 2, cudaMemcpyHostToDevice); - cudaMemcpy(d_v, h_v, HD * SK * 2, cudaMemcpyHostToDevice); - - CUtensorMap h_tma_k, h_tma_v; - if (!create_tma_desc_2d_bf16(&h_tma_k, d_k, SK, HD, 128, 16)) { printf("K TMA FAILED\n"); return 1; } - if (!create_tma_desc_2d_bf16(&h_tma_v, d_v, HD, SK, 16, 16)) { printf("V TMA FAILED\n"); return 1; } - - CUtensorMap *d_tma_k, *d_tma_v; - cudaMalloc(&d_tma_k, sizeof(CUtensorMap)); - cudaMalloc(&d_tma_v, sizeof(CUtensorMap)); - cudaMemcpy(d_tma_k, &h_tma_k, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - cudaMemcpy(d_tma_v, &h_tma_v, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - - FmhaTmaMultiRowMultiTileParams params; - params.q = d_q; params.tma_k = d_tma_k; params.tma_v = d_tma_v; - params.o = d_o; params.lse = d_lse; - params.s_k = SK; params.T = 1; params.n_h = 1; params.scale = SCALE; - params.q_head_stride = 0; params.q_batch_stride = 0; - params.o_head_stride = 0; params.o_batch_stride = 0; - params.lse_head_stride = 0; params.lse_batch_stride = 0; - - // Compute SMEM (match kernel layout) - constexpr int HD_CHUNK = 256; - constexpr int TILE_SZ = 128 * 16; - constexpr int V_SUB_SZ = 16 * 16; - int hc = (HD <= 256) ? HD : HD_CHUNK; - size_t off = 0; - off += 4; off = (off+127)&~(size_t)127; // tmembase - off += 16; off = (off+127)&~(size_t)127; // mbar - off += TILE_SZ*2; off = (off+127)&~(size_t)127; // tmabuf - off += TILE_SZ*2; off = (off+127)&~(size_t)127; // q0 - off += TILE_SZ*2; off = (off+127)&~(size_t)127; // k0 - off += TILE_SZ*2; off = (off+127)&~(size_t)127; // pk - off += V_SUB_SZ*2; off = (off+127)&~(size_t)127; // v - off += 128*hc*4; off += 128*4; off += 128*4; off += 128*4; off += 128*4; off += 256; - int smem = (int)((off + 127) & ~(size_t)127); - - dim3 grid(1, 1, 1); - dim3 block(NTHREADS); - cudaFuncSetAttribute(fmha_6warp_tma_multirow_multitile_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - - fmha_6warp_tma_multirow_multitile_kernel<<>>(params); - cudaError_t err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { printf("Kernel FAILED: %s\n", cudaGetErrorString(err)); return 1; } - - // CPU reference - float h_o_ref[HD], row_max = -INFINITY; - for (int j = 0; j < SK; j++) { - float dot = 0; - for (int d = 0; d < HD; d++) dot += hbf16_to_f32(h_q[d]) * hbf16_to_f32(h_k[j*HD+d]); - dot *= SCALE; - if (dot > row_max) row_max = dot; - } - float row_sum = 0; - for (int d = 0; d < HD; d++) h_o_ref[d] = 0; - for (int j = 0; j < SK; j++) { - float dot = 0; - for (int d = 0; d < HD; d++) dot += hbf16_to_f32(h_q[d]) * hbf16_to_f32(h_k[j*HD+d]); - float p = expf(dot * SCALE - row_max); - row_sum += p; - for (int d = 0; d < HD; d++) h_o_ref[d] += p * hbf16_to_f32(h_v[d*SK+j]); - } - for (int d = 0; d < HD; d++) h_o_ref[d] /= row_sum; - - bf16_t h_o[HD]; - cudaMemcpy(h_o, d_o, HD * 2, cudaMemcpyDeviceToHost); - - float cos = 0, na = 0, nb = 0; - for (int d = 0; d < HD; d++) { float a = h_o_ref[d], b = hbf16_to_f32(h_o[d]); cos += a*b; na += a*a; nb += b*b; } - cos /= sqrtf(na * nb + 1e-30f); - - printf("Multi-tile TMA FMHA (HD=%d, SK=%d): cos=%.6f %s\n", HD, SK, cos, cos >= 0.999990 ? "PASS" : "FAIL"); - cudaFree(d_q_raw); cudaFree(d_k_raw); cudaFree(d_v_raw); cudaFree(d_o_raw); cudaFree(d_lse_raw); - cudaFree(d_tma_k); cudaFree(d_tma_v); - return (cos >= 0.999990) ? 0 : 1; -}