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

Move sC_flat_staged creation before const_expr guard

Browse Source
This commit is contained in:
biondizzle
2026-05-27 05:38:39 +00:00
parent 3c9dbc0c5d
commit 1caa737b09
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
dsv4/kernels/attention/fmha_smem_acc.py
View File

@@ -614,7 +614,8 @@ class FmhaKernel:
c_pipe.producer_tail()
else:
# Path 2: write sO_acc (FP32) -> sC_flat (BF16) -> TMA store to GMEM
# sC_flat has epi_s layout (same as what tma_c was created from)
# Reinterpret sC_flat as staged layout for TMA compatibility
sC_flat_staged = cute.make_tensor(sC_flat.iterator, cute.make_layout((128, self.pv_n_tile // self.num_c_stage, self.num_c_stage), stride=(self.pv_n_tile, self.num_c_stage, 1)))
# Step 1: Cast sO_acc -> sC_flat (BF16) with staged layout
for row in cutlass.range(0, 128, unroll=1):