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

fix: revert V to direct load (V TMA needs debugging), K TMA works

Browse Source
This commit is contained in:
biondizzle
2026-05-29 19:35:44 +00:00
parent cd2c028b39
commit fb971781aa
2 changed files with 15 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
dsv4/kernels/attention/fmha_6warp_tma.cuh
View File

@@ -49,7 +49,8 @@ __global__ void __launch_bounds__(192)
fmha_6warp_tma_kernel(
const bf16_t* __restrict__ q,
CUtensorMap* __restrict__ tma_k,
CUtensorMap* __restrict__ tma_v,
const bf16_t* __restrict__ v,
CUtensorMap* __restrict__ tma_v_unused,
bf16_t* __restrict__ o,
float* __restrict__ lse,
int s_k,
@@ -206,35 +207,25 @@ fmha_6warp_tma_kernel(
int d_base = n * 16;
for (int kt = 0; kt < NKT_PV; kt++) {
// ---- Warp 5: Fill sPk ----
// ---- Warp 5: Fill sPk and load V sub-tile (direct) ----
if (is_load_warp) {
// Fill sPk from s_p_vals
for (int i = lane; i < TILE_SZ; i += 32) sPk[i] = 0;
if (lane < 16) {
int c = lane;
int ck = c / 8, lc = c % 8;
sPk[ck * CORES_MN * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * MMA_K_BF16 + c]);
}
}
__syncthreads();
// ---- TMA load V sub-tile ----
// V is (HD, SK) in GMEM. Load (16, 16) tile at (d_base, kt*16)
// TMA coord: (kt*16, d_base) = (inner_dim_offset, outer_dim_offset)
if (is_load_warp && lane == 0) {
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)tma_v,
mbar_addr, kt * MMA_K_BF16, d_base);
tma_mbarrier_arrive_expect_tx(mbar_addr, V_SUB_SZ * sizeof(bf16_t));
}
tma_mbarrier_wait(mbar_addr, phase); phase ^= 1;
__syncthreads();
// Convert sTmaBuf (row-major 16×16) → sV (canonical 16×16)
for (int i = tid; i < V_SUB_SZ; i += 192) sV[i] = 0;
for (int i = tid; i < 16 * MMA_K_BF16; i += 192) {
int dd = i / MMA_K_BF16, lr = i % MMA_K_BF16; // row, col in (16,16) tile
int g_mn = dd / 8, g_k = lr / 8;
int llr = dd % 8, lc = lr % 8;
sV[g_k * 2 * 64 + g_mn * 64 + llr * 8 + lc] = sTmaBuf[i];
// Load V sub-tile: direct from GMEM
for (int i = lane; i < V_SUB_SZ; i += 32) sV[i] = 0;
for (int dd = lane; dd < 16; dd += 32) {
for (int lr = 0; lr < MMA_K_BF16; lr++) {
int r = kt * MMA_K_BF16 + lr;
int g_mn = dd / 8, g_k = lr / 8;
int llr = dd % 8, lc = lr % 8;
sV[g_k * 2 * 64 + g_mn * 64 + llr * 8 + lc] = v[(d_base + dd) * SK_TILE + r];
}
}
}
__syncthreads();

2
tests/unit/test_fmha_6warp_tma.cu
View File

@@ -122,7 +122,7 @@ int main() {
if (smem > 48 * 1024) {
cudaFuncSetAttribute(fmha_6warp_tma_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
}
fmha_6warp_tma_kernel<HD><<<1, 192, smem>>>(d_q, d_tma_k, d_tma_v, d_o, d_lse, SK, SCALE);
fmha_6warp_tma_kernel<HD><<<1, 192, smem>>>(d_q, d_tma_k, d_v, d_tma_v, d_o, d_lse, SK, SCALE);
cudaError_t launch_err = cudaGetLastError();
if (launch_err != cudaSuccess) { printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err)); return 1; }