debug: add printf to TMEM kernel to find hang point

This commit is contained in:
2026-05-28 07:39:53 +00:00
parent 44fb04fa1f
commit 0ddcc6bafd
2 changed files with 9 additions and 14 deletions

View File

@@ -117,7 +117,7 @@ int main() {
printf("=== FMHA SM100 Decode Kernel Test Suite ===\n\n");
int all_pass = 1;
int head_dims[] = {64, 128};
int head_dims[] = {64};
int s_ks[] = {128};
for (int t = 0; t < 2; t++) {
@@ -154,8 +154,8 @@ int main() {
cudaMemcpy(dk,hkb,B*sk*HD*2,cudaMemcpyHostToDevice);
cudaMemcpy(dv,hvb,B*HD*sk*2,cudaMemcpyHostToDevice);
all_pass &= test_kernel("reference", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
// all_pass &= test_kernel("tmem_epilogue", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
// all_pass &= test_kernel("reference", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
all_pass &= test_kernel("tmem_epilogue", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
cudaFree(dq);cudaFree(dk);cudaFree(dv);cudaFree(do_);cudaFree(d_lse);
free(hq);free(hk);free(hv);free(ho_ref);free(hqb);free(hkb);free(hvb);