diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 2ed27d7e..e499fb11 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -366,16 +366,19 @@ class FmhaKernel:
                     cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP)
                     cute.arch.fence_view_async_tmem_store()
                 else:
-                    # SMEM-P: write P to sP using coordinate-indexed store.
-                    for j0 in range(32):
-                        for j1 in range(4):
-                            coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0]
-                            m_coord = coord[0]
-                            k_coord = coord[1]
-                            k0 = k_coord % 16
-                            k1 = (k_coord // 16) % 4
-                            k2 = k_coord // 64
-                            _sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
+                    # SMEM-P: write P to sP using make_tiled_copy_C + retile.
+                    # The retile() call transforms rP_bf16 from QK C-fragment layout
+                    # to the SMEM copy's source layout, matching partition_D(sP).
+                    _smem_p_store_atom = cute.make_copy_atom(
+                        cute.nvgpu.CopyUniversalOp(),
+                        self.q_dtype,
+                        num_bits_per_copy=16,
+                    )
+                    _tiled_smem_p = cute.make_tiled_copy_C(_smem_p_store_atom, qk_mma)
+                    _thr_smem_p = _tiled_smem_p.get_slice(sfw_idx)
+                    _tRS_sP = _thr_smem_p.partition_D(_sP_nostage)
+                    _tRS_rP = _tiled_smem_p.retile(rP_bf16)
+                    cute.copy(_tiled_smem_p, _tRS_rP, _tRS_sP)
                     cute.arch.fence_proxy("async.shared", space="cta")
                 if kt > 0:
                     for i in range(n_corr_tiles):
diff --git a/tests/unit/test_d1_layout_diag.py b/tests/unit/test_d1_layout_diag.py
index 981b82dd..500c761d 100644
--- a/tests/unit/test_d1_layout_diag.py
+++ b/tests/unit/test_d1_layout_diag.py
@@ -1,35 +1,41 @@
-"""D1: Print SMEM-P layout diagnostics for hd=128."""
-import torch, math, cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.torch as ct
+"""D1: Print SMEM-P layout diagnostics inside the kernel."""
+import torch, math
+import cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.torch as ct
 import cuda.bindings.driver as cuda
-from cutlass.cute.nvgpu import tcgen05
-from cutlass import Float32, BFloat16
-from cutlass.utils import LayoutEnum
+from dsv4.kernels.attention.fmha import FmhaKernel
 
-for hd in [64, 128, 256]:
-    q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
-    k = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
 
-    a_major = LayoutEnum.from_tensor(ct.from_dlpack(q)).mma_major_mode()
-    b_major = LayoutEnum.from_tensor(ct.from_dlpack(k)).mma_major_mode()
-    qk_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, a_major, b_major, Float32, tcgen05.CtaGroup.ONE, (128,128), tcgen05.OperandSource.SMEM)
-    pv_n_tile = min(hd, 256)
-    pv_a_major = a_major  # SMEM-P path uses a_major
-    pv_mma = utils.sm100.make_trivial_tiled_mma(BFloat16, BFloat16, pv_a_major, b_major, Float32, tcgen05.CtaGroup.ONE, (128,pv_n_tile), tcgen05.OperandSource.SMEM)
+def test_layout(hd):
+    m = 128; n_kv = 1  # minimal
+    torch.manual_seed(42)
+    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(n_kv, hd, 1, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda')
+    c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+    lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
 
-    qk_mma_tiler = (128, 128, cute.size(qk_mma.shape_mnk, mode=[2]) * 4)
-    pv_mma_tiler = (128, pv_n_tile, cute.size(pv_mma.shape_mnk, mode=[2]) * (128 // cute.size(pv_mma.shape_mnk, mode=[2])))
+    kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False)
+    pv_n_tile = kernel.pv_n_tile
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 
-    p_tmem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
-    p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
+    v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
+    c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
+    mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
+    mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
+    mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
+    mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
+    mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
 
-    # QK C-fragment layout
-    qk_thr = qk_mma.get_slice(0)
-    qk_as = qk_thr.partition_shape_C(qk_mma_tiler[:2])
-    tStS = qk_thr.make_fragment_C(qk_as)
+    # The _setup method prints are visible at compile time
+    print(f'--- hd={hd} ---', flush=True)
+    print(f'  pv_n_tile={pv_n_tile}, use_smem_p={kernel.use_smem_p}', flush=True)
+    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
+    print(f'  tmem_p0_offset={kernel.tmem_p0_offset}', flush=True)
+    print(f'  tmem_o0_offset={kernel.tmem_o0_offset}', flush=True)
+    print(f'  tOrP0_offset={kernel.tOrP0_offset}', flush=True)
+    print(f'  num_tmem_alloc_cols={kernel.num_tmem_alloc_cols}', flush=True)
 
-    print(f'--- hd={hd} ---')
-    print(f'  p_tmem_s.outer = {p_tmem_s.outer}')
-    print(f'  p_smem_s.outer = {p_smem_s.outer}')
-    print(f'  tStS.layout = {tStS.layout}')
-    print(f'  pv_mma_tiler = {pv_mma_tiler}')
-    print()
+
+if __name__ == '__main__':
+    for hd in [64, 128, 256]:
+        test_layout(hd)