biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P4: 128B-aligned GMEM, proper SMEM alignment, bit21 test

Browse Source
This commit is contained in:
biondizzle
2026-05-30 08:41:15 +00:00
parent 90c806733f
commit 58c087416b
1 changed files with 55 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
tests/unit/test_p4_tma_bit21_fix.cu
View File

@@ -1,16 +1,6 @@
/**
* P4: Test TMA load with the bit-21 workaround from CUTLASS.
*
* Root cause of the TMA hang: driver 13.0 can't read descriptors
* created by toolkit 13.2's cuTensorMapEncodeTiled. CUTLASS clears
* bit 21 of desc[1] as a workaround for driver <= 13.1 with small tensors.
*
* This test:
* 1. Creates a 2D TMA descriptor with NO swizzle
* 2. Dumps the descriptor bytes
* 3. Clears bit 21 of word[1] (the 64-bit word at offset 8)
* 4. Dumps the modified descriptor
* 5. Tests the TMA load with both descriptors
* P4: Test TMA with bit-21 workaround.
* Uses outermost-first dims, 2 strides, 128B-aligned GMEM.
*/
#include <cuda.h>
#include <cuda_runtime.h>
@@ -18,12 +8,9 @@
#include <cstdint>
#include <cstring>
__global__ void tma_load_kernel(
const void* tma_desc_ptr,
int* result
) {
__global__ void tma_load_kernel(const void* tma_desc_ptr, int* result) {
__shared__ uint64_t mbar;
__shared__ uint16_t smem_out[16 * 16];
__shared__ uint16_t smem_out[16 * 16] __attribute__((aligned(128)));
if (threadIdx.x == 0) {
uint32_t mbar_addr = __cvta_generic_to_shared(&mbar);
@@ -38,131 +25,90 @@ __global__ void tma_load_kernel(
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)
:: "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);
int waited = 0;
while (waited < 1000000) {
for (int w = 0; w < 1000000; w++) {
uint32_t state;
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
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)
);
"selp.b32 %1, 1, 0, p;\n\t}" : "=r"(state) : "r"(mbar_addr));
if (state) { *result = 1; return; }
waited++;
}
*result = -1;
}
}
CUtensorMap create_descriptor(void* d_ptr, bool clear_bit21) {
CUtensorMap desc;
// (128, 16) BF16, row-major, NO swizzle
// CUDA 13: globalDim is innermost-first, globalStrides in bytes, rank-1 strides
cuuint64_t globalDim[] = {16, 128}; // (cols, rows) innermost-first
cuuint64_t globalStrides[] = {16 * 2}; // row stride in bytes (rank-1 strides!)
cuuint32_t boxDim[] = {16, 16};
cuuint32_t elementStrides[] = {1, 1};
CUresult res = cuTensorMapEncodeTiled(&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) {
printf(" cuTensorMapEncodeTiled FAILED: %d\n", res);
return desc;
}
// Apply bit-21 workaround if requested
if (clear_bit21) {
uint64_t* words = reinterpret_cast<uint64_t*>(&desc);
words[1] &= ~(1ULL << 21);
}
return desc;
}
int main() {
const size_t DATA_SIZE = 128 * 16 * 2; // (128, 16) BF16
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);
cudaMemset(d_data, 1, DATA_SIZE);
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));
// Test 1: Original descriptor (no bit-21 fix)
printf("=== Test 1: Original descriptor (no fix) ===\n");
{
CUtensorMap desc = create_descriptor(d_data, false);
// Dump first 16 bytes
auto* b = reinterpret_cast<const uint8_t*>(&desc);
printf(" Bytes [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n");
printf(" Bytes [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n");
// 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 h_result;
cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost);
int r; cudaMemcpy(&r, d_result, sizeof(int), cudaMemcpyDeviceToHost);
if (err != cudaSuccess) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), h_result);
else if (h_result == 1) printf(" SUCCESS\n");
else if (h_result == -1) printf(" HANG (mbarrier timeout)\n");
else printf(" UNKNOWN: result=%d\n", h_result);
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);
}
};
// Test 2: Bit-21 cleared (CUTLASS workaround)
printf("\n=== Test 2: Bit-21 cleared (CUTLASS workaround) ===\n");
{
CUtensorMap desc = create_descriptor(d_data, true);
auto* b = reinterpret_cast<const uint8_t*>(&desc);
printf(" Bytes [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n");
printf(" Bytes [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n");
CUtensorMap desc1 = make_desc(d_aligned, false);
test_desc("Original (no fix)", desc1);
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 h_result;
cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost);
if (err != cudaSuccess) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), h_result);
else if (h_result == 1) printf(" SUCCESS\n");
else if (h_result == -1) printf(" HANG (mbarrier timeout)\n");
else printf(" UNKNOWN: result=%d\n", h_result);
cudaFree(d_desc);
}
CUtensorMap desc2 = make_desc(d_aligned, true);
test_desc("Bit-21 cleared", desc2);
cudaFree(d_data);
cudaFree(d_result);