Complete multi-PV-tile fixes: pv_n_tile, v_fmha layout, MMA construction, n_corr_tiles

2026-05-23 09:02:26 +00:00
parent eedcfd7d21
commit eb5c538c9b
1 changed files with 2 additions and 0 deletions
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -19,6 +19,8 @@ class FmhaKernel:
        self.head_dim = head_dim
        self.s_k = s_k
        self.n_kv_tiles = s_k // 128
+        self.pv_n_tile = min(head_dim, 256)  # tcgen05 MMA max N=256
+        self.n_pv_tiles = head_dim // self.pv_n_tile

        self.pv_n_tile = min(head_dim, 256)  # tcgen05 MMA max N=256
        self.n_pv_tiles = head_dim // self.pv_n_tile