diff --git a/tests/unit/test_cotiled_diag.py b/tests/unit/test_cotiled_diag.py
index 05bf0140..6fd274e0 100644
--- a/tests/unit/test_cotiled_diag.py
+++ b/tests/unit/test_cotiled_diag.py
@@ -47,12 +47,12 @@ def main():
     a_major = LayoutEnum.from_tensor(mQ).mma_major_mode()
     b_major = LayoutEnum.from_tensor(mK).mma_major_mode()
     
-    # v_fmha layout
-    v_fmha_layout = cute.make_layout(
-        (pv_n_tile, s_k, 1),
-        stride=(1, pv_n_tile, pv_n_tile * s_k),
-    )
-    v_major = LayoutEnum.ROW_MAJOR  # based on the v_fmha layout
+    # v_major: must match FmhaKernel's derivation
+    # In FmhaKernel, v_fmha is created with layout (pv_n_tile, s_k, 1) stride (1, pv_n_tile, ...)
+    # We create a temporary tensor to get the LayoutEnum
+    v_fmha_tensor = torch.randn(pv_n_tile, s_k, 1, dtype=torch.bfloat16, device='cuda')
+    mV_fmha = ct.from_dlpack(v_fmha_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_fmha_tensor))
+    v_major = LayoutEnum.from_tensor(mV_fmha).mma_major_mode()
 
     qk_mma = utils.sm100.make_trivial_tiled_mma(
         q_dtype, q_dtype, a_major, b_major, Float32,