biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

D1.3: Pre-compute tOrP0_offset in _setup, use const_expr for compile-time selection

Browse Source
This commit is contained in:
biondizzle
2026-05-23 21:07:10 +00:00
parent 77035f8a54
commit 61cb8f002d
1 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
dsv4/kernels/attention/fmha.py
View File

@@ -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)