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

test: debug — just QK+softmax+P read (no PV)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:08:06 +00:00
parent 5d75decd57
commit 0672373e51
1 changed files with 2 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/unit/test_fmha_hd16.cu
View File

@@ -86,11 +86,8 @@ test_fmha_hd16(const bf16_t* q, const bf16_t* k, const bf16_t* v,
}
__syncthreads();
// STEP 3: PV — register math (decode T=1)
// O[d] = Σ P[0,j] × V[d,j] for j=0..SK-1
// P is already in s_vals (warp 0, lane 0). Read V from GMEM.
// STEP 3: Just read P from TMEM and write to output (debug)
if (wid == 0) {
// Re-read P from TMEM (already written back)
float p_vals[SK];
for (int n = 0; n < SK / 8; n++) {
float tmp[8];
@@ -98,15 +95,7 @@ test_fmha_hd16(const bf16_t* q, const bf16_t* k, const bf16_t* v,
asm volatile("tcgen05.wait::ld.sync.aligned;");
if (lane == 0) for (int c=0;c<8;c++) p_vals[n*8+c] = tmp[c];
}
// Compute O[d] = Σ p[j] × V[d,j]
if (lane == 0) {
for (int d = 0; d < HD; d++) {
float ov = 0.0f;
for (int j = 0; j < SK; j++)
ov += p_vals[j] * bf16_to_f32(v[d * SK + j]);
o_out[d] = f32_to_bf16(ov);
}
}
if (lane == 0) for (int j=0;j<SK;j++) o_out[j] = f32_to_bf16(p_vals[j]);
}
__syncthreads();