From a75b4ac483166189a45290783cb0a18af5ff0ea5 Mon Sep 17 00:00:00 2001 From: Ali Hassani <68103095+alihassanijr@users.noreply.github.com> Date: Tue, 5 Dec 2023 15:35:41 -0500 Subject: [PATCH] Fix Stream-K reduce bug in epilogue with broadcast (#1224) Co-authored-by: Ali Hassani --- .../epilogue/threadblock/epilogue_with_broadcast.h | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h b/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h index 7a97e0cb32..6b7351cfd1 100644 --- a/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h +++ b/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h @@ -946,13 +946,13 @@ class EpilogueWithBroadcast< // if (OutputOp::kStoreZ) { + destination_iterator += reduce_fragment_idx; destination_iterator.store(frag_Z); - ++destination_iterator; } if (OutputOp::kStoreT) { + tensor_iterator += reduce_fragment_idx; tensor_iterator.store(frag_T); - ++tensor_iterator; } } }; @@ -1698,13 +1698,13 @@ class EpilogueWithBroadcast< // if (OutputOp::kStoreZ) { + destination_iterator += reduce_fragment_idx; destination_iterator.store(frag_Z); - ++destination_iterator; } if (OutputOp::kStoreT) { + tensor_iterator += reduce_fragment_idx; tensor_iterator.store(frag_T); - ++tensor_iterator; } } };