diff --git a/include/triton/Dialect/TritonGPU/Transforms/Utility.h b/include/triton/Dialect/TritonGPU/Transforms/Utility.h
index 41094258aa..bc87edc1a1 100644
--- a/include/triton/Dialect/TritonGPU/Transforms/Utility.h
+++ b/include/triton/Dialect/TritonGPU/Transforms/Utility.h
@@ -280,6 +280,8 @@ class PatternRewriterWithAsyncTaskIds {
   OpTy replaceOpWithNewOp(Operation *op, Args &&...args) {
     auto newOp =
         rewriter->replaceOpWithNewOp<OpTy>(op, std::forward<Args>(args)...);
+    if (!asyncTaskIds.empty())
+      setAsyncTaskIds(newOp, asyncTaskIds);
     return newOp;
   }