From eedcfd7d2144ff6fc759436ac88a26f3effef50c Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 23 May 2026 09:02:01 +0000
Subject: [PATCH] Fix v_fmha layout to use pv_n_tile instead of head_dim for
 multi-PV-tile support

---
 dsv4/kernels/attention/fmha.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 1ae02921..c7f3a3bc 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -92,7 +92,10 @@ class FmhaKernel:
         v_fmha = cute.make_tensor(
             v.iterator,
             cute.make_layout(
-                (self.head_dim, self.s_k, 1),
+                (self.pv_n_tile, self.s_k, 1),
+                stride=(1, self.pv_n_tile, self.pv_n_tile * self.s_k),
+            ),
+        )
                 stride=(1, self.head_dim, self.head_dim * self.s_k),
             ),
         )