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

fix: V tile copy — V is (HD, SK_TOTAL) so tile columns are not contiguous

Browse Source
This commit is contained in:
biondizzle
2026-05-28 23:55:52 +00:00
parent 869460a932
commit c65baabcc9
1 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/unit/test_fmha_6warp_multirow.cu
View File

@@ -187,10 +187,19 @@ static int test_multitile_merge(int T) {
float* lse_per_tile = (float*)malloc(N_TILES * T * sizeof(float));
float* o_per_tile = (float*)malloc(N_TILES * T * HD * sizeof(float));
// We need separate V buffer per tile since V is (HD, SK_TOTAL) and
// the tile columns are not contiguous in the original array
bf16_t* h_v_tile = (bf16_t*)malloc(HD * SK * sizeof(bf16_t));
for (int tile = 0; tile < N_TILES; tile++) {
// K/V for this tile
// K for this tile: K is (SK_TOTAL, HD), rows are contiguous
cudaMemcpy(d_k, h_k + tile * SK * HD, SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_v, h_v + tile * HD * SK, HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
// V for this tile: V is (HD, SK_TOTAL), tile columns are NOT contiguous
// Copy V[d, tile*SK..(tile+1)*SK-1] into h_v_tile[d, 0..SK-1]
for (int d = 0; d < HD; d++)
memcpy(h_v_tile + d * SK, h_v + d * SK_TOTAL + tile * SK, SK * sizeof(bf16_t));
cudaMemcpy(d_v, h_v_tile, HD * SK * sizeof(bf16_t), cudaMemcpyHostToDevice);
FmhaMultiRowParams params;
params.q = d_q; params.k = d_k; params.v = d_v; params.o = d_o; params.lse = d_lse;
@@ -265,7 +274,7 @@ static int test_multitile_merge(int T) {
printf(" min_cos=%.8f %s\n", min_cos, failed==0?"PASSED":"FAILED");
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse);
free(h_q); free(h_k); free(h_v); free(h_o_merged);
free(h_q); free(h_k); free(h_v); free(h_o_merged); free(h_v_tile);
free(lse_per_tile); free(o_per_tile);
return failed == 0;
}