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

auto: pre-test commit

Browse Source
This commit is contained in:
biondizzle
2026-05-28 15:51:55 +00:00
parent 102174fade
commit 6e5401df3b
1 changed files with 13 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
tests/unit/test_tmem_all_lanes.cu
View File

@@ -173,7 +173,7 @@ test_tmem_all_lanes(const bf16_t* q, const bf16_t* k, const bf16_t* v,
// and figure out the mapping on the host.
if (wid == 0) {
for (int n = 0; n < 8; n++) { // 8 reads of 8 columns = 64 columns
for (int n = 0; n < 16; n++) { // 16 reads of 8 columns = 128 columns
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]),
@@ -183,10 +183,9 @@ test_tmem_all_lanes(const bf16_t* q, const bf16_t* k, const bf16_t* v,
// Each lane writes its 8 values to GMEM
// For column (n*8 + c), lane i's value goes to position lane in that column
// Layout: tmem_dump[col * 32 + lane] = lane's value for column col
// (32 lanes, 1 value per column per lane)
for (int c = 0; c < 8; c++) {
int col = n * 8 + c;
tmem_dump[col * 32 + lane] = tmp[c];
if (col < 128) tmem_dump[col * 32 + lane] = tmp[c];
}
}
}
@@ -213,9 +212,9 @@ int main() {
cudaMalloc(&d_q, HD*sizeof(bf16_t));
cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
cudaMalloc(&d_v, HD*SK*sizeof(bf16_t));
// 64 columns × 32 lanes = 2048 FP32
cudaMalloc(&d_tmem_dump, 64 * 32 * sizeof(float));
cudaMemset(d_tmem_dump, 0, 64 * 32 * sizeof(float));
// 128 columns × 32 lanes = 4096 FP32
cudaMalloc(&d_tmem_dump, 128 * 32 * sizeof(float));
cudaMemset(d_tmem_dump, 0, 128 * 32 * sizeof(float));
cudaMemcpy(d_q, h_q, HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_v, h_v, HD*SK*sizeof(bf16_t), cudaMemcpyHostToDevice);
@@ -227,8 +226,8 @@ int main() {
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
float* h_dump = (float*)malloc(64 * 32 * sizeof(float));
cudaMemcpy(h_dump, d_tmem_dump, 64 * 32 * sizeof(float), cudaMemcpyDeviceToHost);
float* h_dump = (float*)malloc(128 * 32 * sizeof(float));
cudaMemcpy(h_dump, d_tmem_dump, 128 * 32 * sizeof(float), cudaMemcpyDeviceToHost);
// Reference
float s[SK];
@@ -255,7 +254,7 @@ int main() {
float target = o_ref[d];
int best_col = -1, best_lane = -1;
float best_diff = 1e10f;
for (int col = 0; col < 64; col++) {
for (int col = 0; col < 128; col++) {
for (int ln = 0; ln < 32; ln++) {
float val = h_dump[col * 32 + ln];
float diff = fabsf(val - target);
@@ -272,11 +271,11 @@ int main() {
// Print column summary: for each column, which lanes have non-zero values?
printf("\n=== Non-zero lanes per column ===\n");
for (int col = 0; col < 64; col++) {
int nz = 0;
for (int ln = 0; ln < 32; ln++) if (fabsf(h_dump[col*32+ln]) > 1e-6f) nz++;
if (nz > 0) printf(" col %2d: %d non-zero lanes\n", col, nz);
}
for (int col = 0; col < 128; col++) {
int nz = 0;
for (int ln = 0; ln < 32; ln++) if (fabsf(h_dump[col*32+ln]) > 1e-6f) nz++;
if (nz > 0) printf(" col %3d: %d non-zero lanes, lane0=%10.6f\n", col, nz, h_dump[col*32+0]);
}
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_tmem_dump);
free(h_q); free(h_k); free(h_v); free(h_dump);