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

fix: SMEM layout and printf in PV-only test

Browse Source
This commit is contained in:
biondizzle
2026-05-29 19:08:39 +00:00
parent 106f103c83
commit 5eff53c145
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/unit/test_pv_only.cu
View File

@@ -47,10 +47,11 @@ test_pv_only_kernel(
size_t off = 0;
uint32_t* sTmemBase = (uint32_t*)(sbuf + off); off = 4;
off = (off + 15) & ~(size_t)15;
float* s_p_vals = (float*)(sbuf + off); off += SK * sizeof(float);
off = (off + 15) & ~(size_t)15;
bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t);
off = (off + 127) & ~(size_t)127;
bf16_t* sV = (bf16_t*)(sbuf + off); off += V_SUB_SZ * sizeof(bf16_t);
float* s_p_vals = (float*)(sbuf + off); off += SK * sizeof(float);
// Load P values into SMEM
for (int j = tid; j < SK; j += 128) s_p_vals[j] = p_input[j];
@@ -112,7 +113,7 @@ test_pv_only_kernel(
}
int main() {
printf("PV-Only Test (HD=%d, SK=%d)\n", HD);
printf("PV-Only Test (HD=%d, SK=%d)\n", (int)HD, (int)SK);
bf16_t* h_v = (bf16_t*)malloc(HD * SK * sizeof(bf16_t));
float* h_p = (float*)malloc(SK * sizeof(float));
@@ -128,7 +129,7 @@ int main() {
cudaMemcpy(d_v, h_v, HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_p, h_p, SK * sizeof(float), cudaMemcpyHostToDevice);
int smem = 4 + 16 + TILE_SZ + V_SUB_SZ*2 + SK*4 + 256;
size_t smem = 4 + 16 + SK*4 + 16 + TILE_SZ*2 + 128 + V_SUB_SZ*2 + 256;
cudaFuncSetAttribute(test_pv_only_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
test_pv_only_kernel<<<1, 128, smem>>>(d_out, d_v, d_p);
cudaError_t err = cudaDeviceSynchronize();