From 6f0475f0db1582d4a663ca210c376347824ec3a1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 24 May 2026 03:36:36 +0000
Subject: [PATCH] fix: guard tma_p creation with gP check

---
 dsv4/kernels/attention/fmha.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index ac6abcfe..ef015e40 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -115,16 +115,16 @@ class FmhaKernel:
         tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile)
 
         # SMEM-P: gP buffer and TMA for P (GMEM→SMEM via TMA)
-        # gP is created by the caller and passed as a parameter
-        # tma_p is created from the traced gP tensor
         if self.use_smem_p and gP is not None:
             p_s = cute.slice_(self.p_smem_s,(None,None,None,0))
             tma_p,gP = cute.nvgpu.make_tiled_tma_atom_A(
                 utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, pv_mma.thr_id),
                 gP, p_s, self.qk_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape
             )
+        elif not self.use_smem_p:
+            tma_p = tma_q  # dummy, dead code
         else:
-            tma_p = tma_q  # dummy
+            raise ValueError("use_smem_p=True but no gP provided")
         # Always create a valid mLSE tensor for the kernel.
         # CuTeDSL doesn't support None parameters in @cute.kernel.
         # For normalize=True, mLSE is unused (dead-code-eliminated by compiler).