From 8e09fae3a1d805fdb502744606e2e576a8e0a36c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 29 May 2026 18:25:47 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20warp-stride=20for=20TMA=20canonical=20wr?= =?UTF-8?q?ites=20=E2=80=94=20only=20load=20warp=20calls=20them?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- dsv4/kernels/attention/fmha_6warp_tma.cuh | 4 ++-- dsv4/kernels/attention/fmha_umma_desc.cuh | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh index 0e3f7a08..5579fa67 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma.cuh @@ -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 diff --git a/dsv4/kernels/attention/fmha_umma_desc.cuh b/dsv4/kernels/attention/fmha_umma_desc.cuh index a5465c09..0ecb2a0c 100644 --- a/dsv4/kernels/attention/fmha_umma_desc.cuh +++ b/dsv4/kernels/attention/fmha_umma_desc.cuh @@ -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 +template __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;