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

sC_flat with simple (128, pv_n_tile) layout for full epi_tile coverage

Browse Source
This commit is contained in:
biondizzle
2026-05-27 05:27:51 +00:00
parent 6fb0e6a417
commit 7d14a2f764
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -226,7 +226,7 @@ class FmhaKernel:
sO_acc = smem.allocate_tensor(element_type=Float32, layout=sO_acc_layout, byte_alignment=128)
# sC_flat: BF16 SMEM buffer with epi_s layout (non-swizzled) for TMA store
# Used to cast sO_acc (FP32) -> BF16 and TMA store to GMEM
sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.select(self.c_smem_s, mode=[0, 1]).outer, byte_alignment=128)
sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((128, self.pv_n_tile), stride=(self.pv_n_tile, 1)), byte_alignment=128)
else:
sO_acc = smem.allocate_tensor(element_type=Float32, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)
sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)