From 61cb8f002d8ec98f0a9062487f8e7d0536869724 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 21:07:10 +0000 Subject: [PATCH] D1.3: Pre-compute tOrP0_offset in _setup, use const_expr for compile-time selection --- dsv4/kernels/attention/fmha.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index ca0ad7c0..ca6f082d 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -75,6 +75,9 @@ class FmhaKernel: self.num_tmem_alloc_cols = 1 while self.num_tmem_alloc_cols < total: self.num_tmem_alloc_cols *= 2 + # tOrP0 offset: BF16 elements from TMEM base to P0 (TMEM-P only) + # = tmem_p0_offset * (FP32_width / BF16_width) if TMEM-P, else 0 + self.tOrP0_offset = max(self.tmem_p0_offset, 0) * 2 # Python int cta = cute.size(qk_mma.thr_id.shape) q_s = cute.slice_(self.q_smem_s,(None,None,None,0)) k_s = cute.slice_(self.k_smem_s,(None,None,None,0)) @@ -175,11 +178,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). - # Use const_expr to handle TMEM-P vs SMEM-P at compile time. - # For TMEM-P: apply P0 offset (BF16 elements = p0_offset * 2) - # For SMEM-P: no offset needed (tOrP0 unused by MMA warp) - _p0_bf16_offset = max(self.tmem_p0_offset, 0) * 2 # Python int, computed at JIT time - tOrP0 = const_expr(lambda: cute.make_tensor(tOrP.iterator + _p0_bf16_offset, tOrP.layout) if _p0_bf16_offset > 0 else tOrP) + # self.tOrP0_offset is pre-computed in _setup as a Python int. + # We must avoid tOrP.iterator + 0 (MLIR OpResult + int not supported). + # Use const_expr to select the right construction at compile time. + @const_expr + def _make_tOrP0(): + if self.tOrP0_offset > 0: + return cute.make_tensor(tOrP.iterator + self.tOrP0_offset, tOrP.layout) + return tOrP + tOrP0 = _make_tOrP0() tCtO_fake = pv_mma.make_fragment_C(cute.append(pv_as, self.num_acc_stage)) pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)