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

feat: per-head TMA descriptors for multi-head FMHA

Browse Source
This commit is contained in:
biondizzle
2026-05-29 19:44:58 +00:00
parent 9eb193458e
commit 754c6a692c
2 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
dsv4/kernels/attention/fmha_6warp_tma_multirow.cuh
View File

@@ -34,11 +34,11 @@ namespace dsv4::kernels::attention {
struct FmhaTmaMultiRowParams {
const bf16_t* __restrict__ q;
CUtensorMap* __restrict__ tma_k; // K: (s_k, HD) with tile (128, 16)
CUtensorMap* __restrict__ tma_k; // Array of [n_h] TMA descriptors for K
const bf16_t* __restrict__ v; // V: direct GMEM (HD, s_k)
bf16_t* __restrict__ o;
float* __restrict__ lse;
int s_k, T;
int s_k, T, n_h;
float scale;
int head_dim;
int q_head_stride, q_batch_stride;
@@ -111,7 +111,7 @@ fmha_6warp_tma_multirow_kernel(FmhaTmaMultiRowParams params) {
const uint32_t mbar_addr = (uint32_t)__cvta_generic_to_shared(sMbar);
int phase = 0;
// Row assignment
CUtensorMap* __restrict__ my_tma_k = params.tma_k + head_idx;
const bool my_warp_active = (T <= 32) ? (wid == 0) : is_softmax_warp;
const int my_row = my_warp_active ? (wid * 32 + lane) : 0;
const bool my_row_active = my_warp_active && (my_row < T);
@@ -136,7 +136,7 @@ fmha_6warp_tma_multirow_kernel(FmhaTmaMultiRowParams params) {
// Load K: TMA async
if (is_load_warp && lane == 0) {
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)params.tma_k,
tma_load_2d((uint32_t)__cvta_generic_to_shared(sTmaBuf), (uint64_t)my_tma_k,
mbar_addr, kt * MMA_K_BF16, 0);
tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES);
}