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

Staged sC_flat with (128, pv_n_tile//2, 2) to match TMA atom

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

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

@@ -616,22 +616,24 @@ class FmhaKernel:
# 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)
# Step 1: Cast sO_acc -> sC_flat (BF16)
# Step 1: Cast sO_acc -> sC_flat (BF16) with staged layout
for row in cutlass.range(0, 128, unroll=1):
for col in cutlass.range(0, self.pv_n_tile, unroll=1):
val = sO_acc[row, col]
if const_expr(self.normalize):
inv_row_sum = Float32(1.0) / row_sum
val = val * inv_row_sum
sC_flat[row, col] = val.to(self.o_dtype)
sC_flat_staged[row, col // self.num_c_stage, col % self.num_c_stage] = val.to(self.o_dtype)
cute.arch.fence_proxy("async.shared", space="cta")
# Step 2: TMA store sC_flat -> GMEM
# Use tCgC (already partitioned) for the GMEM side of TMA
tCgC_xfm = transform_partitioned_tensor_layout(tCgC)
tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile)
# sC_flat (128, pv_n_tile) -> split to match TMA stage: (128, pv_n_tile//2, 2)
sC_flat_staged = cute.logical_divide(sC_flat, cute.make_layout((128, self.pv_n_tile // 2, 2), stride=(self.pv_n_tile, 2, 1)))
# sC_flat (128, pv_n_tile, stride=(pv_n_tile, 1))
# Split pv_n_tile into (pv_n_tile//2, 2) to match TMA stage decomposition
# Result: (128, pv_n_tile//2, 2, stride=(pv_n_tile, 2, 1))
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)))
tOsC, tOgO = cpasync.tma_partition(
tma_c, 0, cute.make_layout(1),
sC_flat_staged,