perf: eliminate end_sync in custom allreduce by delaying input tensor release until next sync

aiter/dist/device_communicators/communicator_cuda.py

@@ -25,6 +25,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
     _ar_quant_no_prefill_max_bytes = int(
         os.environ.get("AITER_AR_QUANT_NO_PREFILL_MAX_BYTES", -1)
     )
+    _custom_ar_hold_inputs = (
+        os.environ.get("AITER_CUSTOM_AR_HOLD_INPUTS", "0") == "1"
+    )
 
     def __init__(
         self,
@@ -82,6 +85,7 @@ def __init__(
             self.ca_comm = CustomAllreduce(
                 group=self.cpu_group,
                 device=self.device,
+                hold_inputs=self._custom_ar_hold_inputs,
                 # symm_mem_enabled=(
                 #     self.symm_mem_comm is not None and not self.symm_mem_comm.disabled
                 # ),

aiter/dist/device_communicators/custom_all_reduce.py

@@ -351,6 +351,7 @@ def __init__(
         device: Union[int, str, torch.device],
         max_size=1024 * 1024 * 1024,  # 2GB bf16/half
         enable_register_for_capturing: bool = True,
+        hold_inputs: bool = False,
     ) -> None:
         """
         Args:
@@ -364,6 +365,8 @@ def __init__(
         """
         self._IS_CAPTURING = False
         self.disabled = True
+        self.hold_inputs = hold_inputs
+        self._inputs: List[torch.Tensor] = []
 
         if not custom_ar:
             # disable because of missing custom allreduce library
@@ -487,6 +490,24 @@ def __init__(
             offsets,
         )
 
+    def _record_collective_inputs(self, *inputs: torch.Tensor) -> None:
+        """Keep current inputs alive until the next collective is enqueued.
+
+        When final device-side end_sync is disabled, peer ranks may still
+        read a registered input after Python has returned from the launch. The
+        next custom collective's internal start_sync proves all ranks have
+        advanced past the prior collective, so replacing this list after that
+        next launch safely releases the previous inputs.
+        """
+        if not self.hold_inputs:
+            return
+        self._inputs = [
+            inp for inp in inputs if isinstance(inp, torch.Tensor)
+        ]
+
+    def flush_recorded_inputs(self) -> None:
+        self._inputs.clear()
+
     @contextmanager
     def capture(self):
         """
@@ -706,7 +727,9 @@ def reduce_scatter(
             split_dim,
             reg,
             reg_bytes,
+            end_sync=not self.hold_inputs,
         )
+        self._record_collective_inputs(inp)
 
     def custom_reduce_scatter(
         self, input: torch.Tensor, output: torch.Tensor, dim: int = 0
@@ -754,7 +777,9 @@ def all_gather_reg(self, inp: torch.Tensor, out: torch.Tensor = None, dim: int =
             inp,
             out,
             dim,
+            end_sync=not self.hold_inputs,
         )
+        self._record_collective_inputs(inp)
         return out
 
     def all_gather_unreg(
@@ -851,6 +876,7 @@ def fused_ar_rms(
                     reg_bytes,
                     use_1stage,
                     gemma_norm,
+                    end_sync=not self.hold_inputs,
                 )
             else:
                 ops.fused_allreduce_rmsnorm_pad(
@@ -865,7 +891,9 @@ def fused_ar_rms(
                     reg_bytes,
                     use_1stage,
                     gemma_norm,
+                    end_sync=not self.hold_inputs,
                 )
+            self._record_collective_inputs(inp)
             return out, res_out
         else:
             if out is None:
@@ -888,7 +916,9 @@ def fused_ar_rms(
                 reg_bytes,
                 use_1stage,
                 gemma_norm,
+                end_sync=not self.hold_inputs,
             )
+            self._record_collective_inputs(inp)
             return out, res_out, scale_out
 
     def custom_fused_ar_rms(
@@ -1108,7 +1138,9 @@ def fused_ar_rms_per_group_quant(
             use_1stage,
             bf16_ptr,
             transpose_scale,
+            end_sync=not self.hold_inputs,
         )
+        self._record_collective_inputs(inp)
         if emit_bf16:
             return out, res_out, scale_out, bf16_out
         return out, res_out, scale_out
@@ -1361,7 +1393,9 @@ def fused_ar_rms_mxfp4_quant(
             reg_bytes,
             use_1stage,
             bf16_ptr,
+            end_sync=not self.hold_inputs,
         )
+        self._record_collective_inputs(inp)
         if emit_bf16:
             return out, res_out, scale_out, bf16_out
         return out, res_out, scale_out
@@ -1479,6 +1513,7 @@ def custom_fused_ar_rms_mxfp4_quant(
         )
 
     def close(self):
+        self.flush_recorded_inputs()
         if not self.disabled and getattr(self, "_ptr", 0):
             ops.dispose(self._ptr)
             self._ptr = 0

aiter/ops/custom_all_reduce.py

@@ -46,6 +46,7 @@ def reduce_scatter(
     split_dim: int,
     reg_ptr: int,
     reg_bytes: int,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -55,6 +56,7 @@ def all_gather_reg(
     inp: torch.Tensor,
     out: torch.Tensor,
     dim: int,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -66,6 +68,7 @@ def all_gather_unreg(
     out: torch.Tensor,
     reg_bytes: int,
     dim: int,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -82,6 +85,7 @@ def fused_allreduce_rmsnorm(
     reg_bytes: int,
     use_1stage: bool,
     gemma_norm: bool = False,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -98,6 +102,7 @@ def fused_allreduce_rmsnorm_pad(
     reg_bytes: int,
     use_1stage: bool,
     gemma_norm: bool = False,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -115,6 +120,7 @@ def fused_allreduce_rmsnorm_quant(
     reg_bytes: int,
     use_1stage: bool,
     gemma_norm: bool = False,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -134,6 +140,7 @@ def fused_allreduce_rmsnorm_quant_per_group(
     use_1stage: bool,
     bf16_out_ptr: int = 0,
     transpose_scale: bool = False,
+    end_sync: bool = True,
 ) -> None: ...
 
 
@@ -151,6 +158,7 @@ def fused_allreduce_rmsnorm_mxfp4_quant(
     reg_bytes: int,
     use_1stage: bool,
     bf16_out_ptr: int = 0,
+    end_sync: bool = True,
 ) -> None: ...