cleanup: remove debug test files (P4, P5)
This commit is contained in:
@@ -1,117 +0,0 @@
|
||||
/**
|
||||
* P4: Test TMA with bit-21 workaround.
|
||||
* Uses outermost-first dims, 2 strides, 128B-aligned GMEM.
|
||||
*/
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
__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<uint64_t*>(&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<uint8_t*>(&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;
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
|
||||
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<const uint8_t*>(&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;
|
||||
}
|
||||
@@ -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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
#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;
|
||||
}
|
||||
@@ -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 <cstdio>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
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<HD, 128>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
}
|
||||
|
||||
fmha_6warp_multihead_kernel<HD, 128><<<grid, block, smem>>>(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<HD, 128><<<grid, block, smem>>>(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<HD, 128><<<grid, block, smem>>>(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; }
|
||||
}
|
||||
@@ -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}')
|
||||
@@ -1,121 +0,0 @@
|
||||
/**
|
||||
* P5: Test multi-tile TMA FMHA kernel with proper alignment.
|
||||
*/
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#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 <cstdio>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
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<HD>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
|
||||
fmha_6warp_tma_multirow_multitile_kernel<HD><<<grid, block, smem>>>(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;
|
||||
}
|
||||
Reference in New Issue
Block a user