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<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;