diff --git a/tests/unit/test_cotiled_diag.py b/tests/unit/test_cotiled_diag.py
index 47bc8e47..1a6e8f92 100644
--- a/tests/unit/test_cotiled_diag.py
+++ b/tests/unit/test_cotiled_diag.py
@@ -39,12 +39,11 @@ def main():
 
     # Derive major modes exactly as FmhaKernel does
     from cutlass.utils import LayoutEnum
+    from cutlass.cute.nvgpu import OperandMajorMode
     a_major = LayoutEnum.from_tensor(mQ).mma_major_mode()
     b_major = LayoutEnum.from_tensor(mK).mma_major_mode()
-
-    # V FMHA layout: use explicit constants
-    # pv_n_tile=256, s_k=128
-    v_major = LayoutEnum.COL_MAJOR  # layout (256, 128, 1) stride (1, 256, 32768) = col-major
+    v_major = LayoutEnum.from_tensor(mV).mma_major_mode()
+    print(f"a_major: {a_major}, b_major: {b_major}, v_major: {v_major}")  # layout (256, 128, 1) stride (1, 256, 32768) = col-major
 
     c_layout = LayoutEnum.from_tensor(mC)