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

fix: warp-stride for TMA canonical writes — only load warp calls them

Browse Source 
write_smem_canonical used NTHREADS=192 as the stride, but in the TMA
kernel only the load warp (32 threads) calls it. With threadIdx.x in
[160,191] and stride 192, only 32 out of 2048 elements got written.
Fix: template STRIDE parameter, default 192, TMA kernel uses 32.
This commit is contained in:
biondizzle
2026-05-29 18:25:47 +00:00
parent 3e14a25bb0
commit 8e09fae3a1
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -162,7 +162,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
phase ^= 1;
__syncthreads();
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16>(sQ, sQ_tma);
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sQ, sQ_tma);
__syncthreads();
// --- TMA load K sub-tile ---
@@ -175,7 +175,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) {
phase ^= 1;
__syncthreads();
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16>(sK, sK_tma);
if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sK, sK_tma);
__syncthreads();
// MMA: sQ × sK → TMEM

5
dsv4/kernels/attention/fmha_umma_desc.cuh
View File

@@ -121,14 +121,15 @@ constexpr int MMA_K_BF16 = 16;
*
* ROWS must be a multiple of 8 (ideally 128 for full MMA tile).
* COLS must be a multiple of 8.
* STRIDE is the thread stride (default NTHREADS=192 for full CTA).
*/
template<int ROWS, int COLS>
template<int ROWS, int COLS, int STRIDE = NTHREADS>
__device__ void write_smem_canonical(bf16_t* __restrict__ dst,
const bf16_t* __restrict__ src) {
constexpr int CORES_MN = ROWS / 8;
constexpr int CORES_K = COLS / 8;
constexpr int TOTAL = ROWS * COLS;
for (int i = threadIdx.x; i < TOTAL; i += NTHREADS) {
for (int i = threadIdx.x; i < TOTAL; i += STRIDE) {
int r = i / COLS;
int c = i % COLS;
int core_mn = r / 8;