Fix Stream-K reduce bug in epilogue with broadcast (#1224)

Co-authored-by: Ali Hassani <[email protected]>
NVIDIA · Dec 5, 2023 · a75b4ac · a75b4ac
1 parent e9e30c2
commit a75b4ac
Showing 1 changed file with 4 additions and 4 deletions.
diff --git 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;
       }
     }
 };