diff --git a/tests/test_pv64_with_softmax.py b/tests/test_pv64_with_softmax.py
index de9e749b..a92680fb 100644
--- a/tests/test_pv64_with_softmax.py
+++ b/tests/test_pv64_with_softmax.py
@@ -45,7 +45,11 @@ class Pv64WithSoftmax:
         tOtO = pv_thr.make_fragment_C(pv_as)
         self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.o_dtype.width
         self.tmem_s0_offset = 0; self.tmem_p0_offset = 32
-        self.tmem_o0_offset = find_tmem_tensor_col_offset(tOtO)
+        p_cols_fp32 = self.pv_mma_tiler[2] * self.q_dtype.width // self.qk_acc_dtype.width
+        p_end = self.tmem_p0_offset + p_cols_fp32
+        s_cols = self.qk_mma_tiler[1]
+        o_after = max(s_cols, p_end)
+        self.tmem_o0_offset = ((o_after + 31) // 32) * 32
         tCS = qk_mma.make_fragment_C(cute.append(qk_as, self.num_acc_stage))
         tCO = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage))
         self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols([tCS, tCO], arch="sm_100")
@@ -58,7 +62,14 @@ class Pv64WithSoftmax:
         self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
         self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
         self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
-        self.v_major = LayoutEnum.from_tensor(v).mma_major_mode()
+        v_fmha = cute.make_tensor(
+            v.iterator,
+            cute.make_layout(
+                (HEAD_DIM, 128, 1),
+                stride=(1, HEAD_DIM, HEAD_DIM * 128),
+            ),
+        )
+        self.v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode()
         self.c_layout = LayoutEnum.from_tensor(c)
         qk_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, self.a_major, self.b_major, self.qk_acc_dtype, self.cta_group, (128,128), tcgen05.OperandSource.SMEM)
         pv_mma = utils.sm100.make_trivial_tiled_mma(self.q_dtype, self.q_dtype, cute.nvgpu.OperandMajorMode.K, self.v_major, self.qk_acc_dtype, self.cta_group, (128,HEAD_DIM), tcgen05.OperandSource.TMEM)
@@ -67,7 +78,7 @@ class Pv64WithSoftmax:
         v_s = cute.slice_(self.v_smem_s,(None,None,None,0))
         tma_q,mQ = cute.nvgpu.make_tiled_tma_atom_A(utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn,qk_mma.thr_id),q,q_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
         tma_k,mK = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,qk_mma.thr_id),k,k_s,self.mma_tiler,qk_mma,self.cluster_layout_vmnk.shape)
-        tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
+        tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape)
         epi_s = cute.select(self.c_smem_s,mode=[0,1])
         tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
         self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.a_smem_s,self.b_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream)
@@ -223,26 +234,24 @@ def test():
     q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
     k = torch.randn(n, hd, 1, dtype=torch.bfloat16, device='cuda')
     v = torch.ones(n, hd, dtype=torch.bfloat16, device='cuda')
-    v = v.as_strided((n, hd), (1, n)).unsqueeze(-1)
+    v_kernel = v.unsqueeze(-1)
     c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
     qf = q[:,:,0].float(); kf = k[:,:,0].float()
-    ref = (qf @ kf.T).bfloat16().float() @ v[:,:,0].float()
+    ref = (qf @ kf.T).bfloat16().float() @ v.float()
     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).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v))
+    mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
     mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
     stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
     kernel = Pv64WithSoftmax()
-    print('Compiling...', flush=True)
+    print("Compiling...", flush=True)
     compiled = cute.compile(kernel, mQ, mK, mV, mC, stream)
-    print(f'tilePlikeFP32={kernel.tilePlikeFP32} s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset}', flush=True)
+    print(f"tilePlikeFP32={kernel.tilePlikeFP32} s0={kernel.tmem_s0_offset} p0={kernel.tmem_p0_offset} o0={kernel.tmem_o0_offset}", flush=True)
     compiled(mQ, mK, mV, mC, stream)
     torch.cuda.synchronize()
     out = c[:,:,0].float()
     cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
-    print(f'PV64 with softmax: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
-    if cos < 0.99:
-        print(f'  out[0,:4]={out[0,:4].tolist()} ref[0,:4]={ref[0,:4].tolist()}')
+    print(f"PV64 with softmax: cosine {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}")
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test()