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

feat: double-buffer TMA pipeline in multi-row kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-30 03:20:49 +00:00
parent 4a9c850e9c
commit 762f054d6d
2 changed files with 101 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/unit/test_fmha_6warp_tma_multirow.cu
View File

@@ -34,13 +34,14 @@ static size_t compute_smem() {
size_t off = 0;
off += 4; off = (off+127)&~(size_t)127;
off += 16; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += 16 * MY_MMA_K * 2;
off += 128 * 4;
off += 128 * 4;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sTmaBuf
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sQ0
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sK0
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sK1 (double buffer)
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sPk
off += 16 * MY_MMA_K * 2; // sV
off += 128 * 4; // sRowMax
off += 128 * 4; // sRowSum
return off;
}