From 1caa737b091208a058b20868325b0c5176a341cb Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 05:38:39 +0000 Subject: [PATCH] Move sC_flat_staged creation before const_expr guard --- dsv4/kernels/attention/fmha_smem_acc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py index dd2c1095..85d1d215 100644 --- a/dsv4/kernels/attention/fmha_smem_acc.py +++ b/dsv4/kernels/attention/fmha_smem_acc.py @@ -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):