D1.3: Use MLIR-compatible expression for tOrP0 offset (same as Stage C)

dsv4/kernels/attention/fmha.py

@@ -175,17 +175,15 @@ class FmhaKernel:
         tOrP = tOrP_base[(None,None,None,0)]
         tCrP = pv_mma.make_fragment_A(sP)
         # tOrP0: PV A-operand with TMEM column offset for P0 (TMEM-P path).
-        # The softmax warps store P at tmem_p0_offset FP32 columns. PV MMA's
-        # tOrP fragment uses BF16 elements. Offset in BF16 elements =
-        # tmem_p0_offset * (FP32_width / BF16_width) = offset * 2.
-        # For SMEM-P, offset is 0 (P not in TMEM).
-        # Must be defined unconditionally (CuTeDSL scoping).
-        _p0_bf16_offset = max(self.tmem_p0_offset, 0) * (32 // 16)  # Python int
-        # Must define tOrP0 before conditional (CuTeDSL scoping rule).
-        # Initialize to tOrP, then override with offset for TMEM-P.
-        tOrP0 = tOrP
-        if _p0_bf16_offset > 0:
-            tOrP0 = cute.make_tensor(tOrP.iterator + _p0_bf16_offset, tOrP.layout)
+        # Same pattern as Stage C kernel: uses MLIR-compatible arithmetic.
+        # tmem_p0_offset is in FP32 columns; tOrP uses BF16 elements.
+        # Offset = acc_width / q_width * tmem_p0_offset.
+        # For SMEM-P (tmem_p0_offset=-1), tOrP0 is unused by the MMA warp
+        # but must be valid for CuTeDSL compilation.
+        tOrP0 = cute.make_tensor(
+            tOrP.iterator + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p0_offset,
+            tOrP.layout,
+        )
 
         tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage))
         pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)