From f9bb2c18e362f3eb71b6a59784888008c659eff0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 01:41:36 +0000 Subject: [PATCH] fix: all epilogue warps do TMA store, no dynamic if inside method --- tests/unit/test_fmha_v3_stage_c.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index c868ec98..2aad1970 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -190,11 +190,11 @@ class FmhaV3StageCMulti: ) epilog_sync_bar.arrive_and_wait() + # All warps in the epilogue group do the TMA store together c_buffer = 0 - if warp_idx == self.epilogue_warp_id[0]: - cute.copy(tma_c, tCsC[(None, c_buffer)], tCgC_tma[(None, 0, 0, 0, 0, 0, 0)]) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.copy(tma_c, tCsC[(None, c_buffer)], tCgC_tma[(None, 0, 0, 0, 0, 0, 0)]) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) epilog_sync_bar.arrive_and_wait() @cute.kernel