From e2ba21df46413a1078d3f597135ac291fbc333fc Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 21:05:12 +0000 Subject: [PATCH] D1.3: Use MLIR-compatible expression for tOrP0 offset (same as Stage C) --- dsv4/kernels/attention/fmha.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index fb3848e8..1c376ecd 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/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)