diff --git a/tests/unit/test_fmha_6warp_multirow.cu b/tests/unit/test_fmha_6warp_multirow.cu index b6afcd38..60b73ffb 100644 --- a/tests/unit/test_fmha_6warp_multirow.cu +++ b/tests/unit/test_fmha_6warp_multirow.cu @@ -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; }