sglang 0.5.2rc2__py3-none-any.whl → 0.5.3.post1__py3-none-any.whl

sglang/srt/disaggregation/prefill.py

@@ -21,9 +21,10 @@ from __future__ import annotations
 
 import logging
 import threading
+import time
 from collections import deque
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Optional, Type
 
 import torch
 
@@ -42,7 +43,12 @@ from sglang.srt.disaggregation.utils import (
     poll_and_all_reduce,
     prepare_abort,
 )
-from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
+from sglang.srt.managers.schedule_batch import (
+    FINISH_LENGTH,
+    Req,
+    RequestStage,
+    ScheduleBatch,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -140,8 +146,10 @@ class PrefillBootstrapQueue:
         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
         kv_args.gpu_id = self.scheduler.gpu_id
 
-        kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
-        kv_manager = kv_manager_class(
+        kv_manager_class: Type[BaseKVManager] = get_kv_class(
+            self.transfer_backend, KVClassType.MANAGER
+        )
+        kv_manager: BaseKVManager = kv_manager_class(
             kv_args,
             DisaggregationMode.PREFILL,
             self.scheduler.server_args,
@@ -168,6 +176,7 @@ class PrefillBootstrapQueue:
             pp_rank=self.pp_rank,
         )
         self._process_req(req)
+        req.add_latency(RequestStage.PREFILL_PREPARE)
         self.queue.append(req)
 
     def extend(self, reqs: List[Req], num_kv_heads: int) -> None:
@@ -254,8 +263,11 @@ class PrefillBootstrapQueue:
 
             num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size)
             req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
+
             bootstrapped_reqs.append(req)
             indices_to_remove.add(i)
+            req.time_stats.wait_queue_entry_time = time.perf_counter()
+            req.add_latency(RequestStage.PREFILL_BOOTSTRAP)
 
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
@@ -309,6 +321,8 @@ class SchedulerDisaggregationPrefillMixin:
         self.result_queue = deque()
 
         while True:
+            self.launch_last_batch_sample_if_needed()
+
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
             self.waiting_queue.extend(
@@ -324,21 +338,8 @@ class SchedulerDisaggregationPrefillMixin:
                 result = self.run_batch(batch)
                 self.result_queue.append((batch.copy(), result))
 
-                if self.last_batch is None:
-                    # Create a dummy first batch to start the pipeline for overlap schedule.
-                    # It is now used for triggering the sampling_info_done event.
-                    tmp_batch = ScheduleBatch(
-                        reqs=None,
-                        forward_mode=ForwardMode.DUMMY_FIRST,
-                        next_batch_sampling_info=self.tp_worker.cur_sampling_info,
-                    )
-                    self.set_next_batch_sampling_info_done(tmp_batch)
-
             if self.last_batch:
                 tmp_batch, tmp_result = self.result_queue.popleft()
-                tmp_batch.next_batch_sampling_info = (
-                    self.tp_worker.cur_sampling_info if batch else None
-                )
                 self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
 
             if len(self.disagg_prefill_inflight_queue) > 0:
@@ -356,7 +357,6 @@ class SchedulerDisaggregationPrefillMixin:
         self: Scheduler,
         batch: ScheduleBatch,
         result: GenerationBatchResult,
-        launch_done: Optional[threading.Event] = None,
     ) -> None:
         """
         Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
@@ -367,41 +367,40 @@ class SchedulerDisaggregationPrefillMixin:
             next_token_ids,
             extend_input_len_per_req,
             extend_logprob_start_len_per_req,
+            copy_done,
         ) = (
             result.logits_output,
             result.next_token_ids,
             result.extend_input_len_per_req,
             result.extend_logprob_start_len_per_req,
+            result.copy_done,
         )
 
+        if copy_done is not None:
+            copy_done.synchronize()
+
         logprob_pt = 0
         # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
-        if self.enable_overlap:
-            # wait
-            logits_output, next_token_ids, _ = self.tp_worker.resolve_last_batch_result(
-                launch_done
-            )
-        else:
-            next_token_ids = result.next_token_ids.tolist()
-            if batch.return_logprob:
-                if logits_output.next_token_logprobs is not None:
-                    logits_output.next_token_logprobs = (
-                        logits_output.next_token_logprobs.tolist()
-                    )
-                if logits_output.input_token_logprobs is not None:
-                    logits_output.input_token_logprobs = tuple(
-                        logits_output.input_token_logprobs.tolist()
-                    )
+        next_token_ids = result.next_token_ids.tolist()
+        if batch.return_logprob:
+            if logits_output.next_token_logprobs is not None:
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.tolist()
+                )
+            if logits_output.input_token_logprobs is not None:
+                logits_output.input_token_logprobs = tuple(
+                    logits_output.input_token_logprobs.tolist()
+                )
 
         hidden_state_offset = 0
         for i, (req, next_token_id) in enumerate(
             zip(batch.reqs, next_token_ids, strict=True)
         ):
-            req: Req
             if req.is_chunked <= 0:
                 # There is no output_ids for prefill
                 req.output_ids.append(next_token_id)
                 self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
+                req.add_latency(RequestStage.PREFILL_FORWARD)
                 self.disagg_prefill_inflight_queue.append(req)
                 if (
                     logits_output is not None
@@ -410,9 +409,16 @@ class SchedulerDisaggregationPrefillMixin:
                     last_hidden_index = (
                         hidden_state_offset + extend_input_len_per_req[i] - 1
                     )
-                    req.hidden_states_tensor = (
-                        logits_output.hidden_states[last_hidden_index].cpu().clone()
-                    )
+                    req.output_topk_p = batch.spec_info.topk_p[i]
+                    req.output_topk_index = batch.spec_info.topk_index[i]
+                    if self.spec_algorithm.is_eagle3():
+                        req.hidden_states_tensor = (
+                            batch.spec_info.hidden_states[i].cpu().clone()
+                        )
+                    else:
+                        req.hidden_states_tensor = (
+                            logits_output.hidden_states[last_hidden_index].cpu().clone()
+                        )
                     hidden_state_offset += extend_input_len_per_req[i]
                 else:
                     req.hidden_states_tensor = None
@@ -432,6 +438,7 @@ class SchedulerDisaggregationPrefillMixin:
                     )
                     logprob_pt += num_input_logprobs
                 self.send_kv_chunk(req, last_chunk=True)
+                req.time_stats.prefill_transfer_queue_entry_time = time.perf_counter()
 
                 if req.grammar is not None:
                     # FIXME: this try-except block is for handling unexpected xgrammar issue.
@@ -471,8 +478,6 @@ class SchedulerDisaggregationPrefillMixin:
                 if self.enable_overlap:
                     self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
 
-        # We need to remove the sync in the following function for overlap schedule.
-        self.set_next_batch_sampling_info_done(batch)
         self.maybe_send_health_check_signal()
 
     def process_disagg_prefill_inflight_queue(
@@ -529,6 +534,9 @@ class SchedulerDisaggregationPrefillMixin:
             else:
                 assert False, f"Unexpected polling state {poll=}"
 
+        for req in done_reqs:
+            req.time_stats.completion_time = time.perf_counter()
+
         # Stream requests which have finished transfer
         self.stream_output(
             done_reqs,
@@ -537,6 +545,7 @@ class SchedulerDisaggregationPrefillMixin:
         )
         for req in done_reqs:
             req: Req
+            req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
             self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
             req.metadata_buffer_index = -1
 
@@ -665,7 +674,6 @@ class SchedulerDisaggregationPrefillMixin:
         self.running_mbs = [
             ScheduleBatch(reqs=[], batch_is_full=False) for _ in range(self.pp_size)
         ]
-        bids = [None] * self.pp_size
         pp_outputs: Optional[PPProxyTensors] = None
 
         # Either success or failed
@@ -737,10 +745,7 @@ class SchedulerDisaggregationPrefillMixin:
                 # send the outputs to the next step
                 if self.pp_group.is_last_rank:
                     if self.cur_batch:
-                        next_token_ids, bids[mb_id] = (
-                            result.next_token_ids,
-                            result.bid,
-                        )
+                        next_token_ids = result.next_token_ids
                         pp_outputs = PPProxyTensors(
                             {
                                 "next_token_ids": next_token_ids,
@@ -777,7 +782,6 @@ class SchedulerDisaggregationPrefillMixin:
                         next_token_ids=next_pp_outputs["next_token_ids"],
                         extend_input_len_per_req=None,
                         extend_logprob_start_len_per_req=None,
-                        bid=bids[next_mb_id],
                         can_run_cuda_graph=result.can_run_cuda_graph,
                     )
                     self.process_batch_result_disagg_prefill(
@@ -794,8 +798,6 @@ class SchedulerDisaggregationPrefillMixin:
 
                 # carry the outputs to the next stage
                 if not self.pp_group.is_last_rank:
-                    if self.cur_batch:
-                        bids[mb_id] = result.bid
                     if pp_outputs:
                         # send the outputs from the last round to let the next stage worker run post processing
                         self.pp_group.send_tensor_dict(
@@ -814,8 +816,10 @@ class SchedulerDisaggregationPrefillMixin:
 
                     # send out proxy tensors to the next stage
                     if self.cur_batch:
+                        # FIXME(lsyin): remove this assert
+                        assert result.pp_hidden_states_proxy_tensors.tensors is not None
                         self.pp_group.send_tensor_dict(
-                            result.pp_hidden_states_proxy_tensors,
+                            result.pp_hidden_states_proxy_tensors.tensors,
                             all_gather_group=self.attn_tp_group,
                         )
 

sglang/srt/disaggregation/utils.py

@@ -1,21 +1,17 @@
 from __future__ import annotations
 
-import dataclasses
 import os
 import random
-import threading
-import warnings
 from collections import deque
 from contextlib import nullcontext
 from enum import Enum
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional, Type
 
 import numpy as np
-import requests
 import torch
 import torch.distributed as dist
 
-from sglang.srt.utils import get_ip, is_npu
+from sglang.srt.utils import is_npu
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -89,7 +85,7 @@ class MetadataBuffers:
         self,
         size: int,
         hidden_size: int,
-        dtype: torch.dtype,
+        hidden_states_dtype: torch.dtype,
         max_top_logprobs_num: int = 128,
         custom_mem_pool: torch.cuda.MemPool = None,
     ):
@@ -111,7 +107,9 @@ class MetadataBuffers:
             # We transfer the metadata of first output token to decode
             # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
             self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device)
-
+            self.cached_tokens = torch.zeros(
+                (size, 16), dtype=torch.int32, device=device
+            )
             self.output_token_logprobs_val = torch.zeros(
                 (size, 16), dtype=torch.float32, device=device
             )
@@ -124,33 +122,49 @@ class MetadataBuffers:
             self.output_top_logprobs_idx = torch.zeros(
                 (size, max_top_logprobs_num), dtype=torch.int32, device=device
             )
+            # For PD + spec decode
+            self.output_topk_p = torch.zeros(
+                (size, 16), dtype=torch.float32, device=device
+            )
+            self.output_topk_index = torch.zeros(
+                (size, 16), dtype=torch.int64, device=device
+            )
             self.output_hidden_states = torch.zeros(
-                (size, hidden_size), dtype=dtype, device=device
+                (size, hidden_size), dtype=hidden_states_dtype, device=device
             )
 
     def get_buf_infos(self):
         ptrs = [
             self.output_ids.data_ptr(),
+            self.cached_tokens.data_ptr(),
             self.output_token_logprobs_val.data_ptr(),
             self.output_token_logprobs_idx.data_ptr(),
             self.output_top_logprobs_val.data_ptr(),
             self.output_top_logprobs_idx.data_ptr(),
+            self.output_topk_p.data_ptr(),
+            self.output_topk_index.data_ptr(),
             self.output_hidden_states.data_ptr(),
         ]
         data_lens = [
             self.output_ids.nbytes,
+            self.cached_tokens.nbytes,
             self.output_token_logprobs_val.nbytes,
             self.output_token_logprobs_idx.nbytes,
             self.output_top_logprobs_val.nbytes,
             self.output_top_logprobs_idx.nbytes,
+            self.output_topk_p.nbytes,
+            self.output_topk_index.nbytes,
             self.output_hidden_states.nbytes,
         ]
         item_lens = [
             self.output_ids[0].nbytes,
+            self.cached_tokens[0].nbytes,
             self.output_token_logprobs_val[0].nbytes,
             self.output_token_logprobs_idx[0].nbytes,
             self.output_top_logprobs_val[0].nbytes,
             self.output_top_logprobs_idx[0].nbytes,
+            self.output_topk_p[0].nbytes,
+            self.output_topk_index[0].nbytes,
             self.output_hidden_states[0].nbytes,
         ]
         return ptrs, data_lens, item_lens
@@ -158,16 +172,20 @@ class MetadataBuffers:
     def get_buf(self, idx: int):
         return (
             self.output_ids[idx],
+            self.cached_tokens[idx],
             self.output_token_logprobs_val[idx],
             self.output_token_logprobs_idx[idx],
             self.output_top_logprobs_val[idx],
             self.output_top_logprobs_idx[idx],
+            self.output_topk_p[idx],
+            self.output_topk_index[idx],
             self.output_hidden_states[idx],
         )
 
     def set_buf(self, req: Req):
 
         self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
+        self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
         if req.return_logprob:
             if req.output_token_logprobs_val:  # not none or empty list
                 self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
@@ -190,8 +208,17 @@ class MetadataBuffers:
                 ] = torch.tensor(
                     req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu"
                 )
-        # for PD + spec decode
+        # For PD + spec decode
         if req.hidden_states_tensor is not None:
+            # speculative_eagle_topk should not be greater than 16 currently
+            topk = req.output_topk_p.size(0)
+
+            self.output_topk_p[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_p
+            )
+            self.output_topk_index[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_index
+            )
             self.output_hidden_states[req.metadata_buffer_index].copy_(
                 req.hidden_states_tensor
             )
@@ -217,7 +244,9 @@ class KVClassType(Enum):
     BOOTSTRAP_SERVER = "bootstrap_server"
 
 
-def get_kv_class(transfer_backend: TransferBackend, class_type: KVClassType):
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: KVClassType
+) -> Optional[Type]:
     from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender
 
     if transfer_backend == TransferBackend.MOONCAKE:
@@ -305,49 +334,6 @@ def kv_to_page_num(num_kv_indices: int, page_size: int):
     return (num_kv_indices + page_size - 1) // page_size
 
 
-#########################
-# PDLB Registry
-#########################
-
-
-@dataclasses.dataclass
-class PDRegistryRequest:
-    """A request to register a machine itself to the LB."""
-
-    mode: str
-    registry_url: str
-    bootstrap_port: Optional[int] = None
-
-    def __post_init__(self):
-        if self.mode == "prefill" and self.bootstrap_port is None:
-            raise ValueError("Bootstrap port must be set in PREFILL mode.")
-        elif self.mode == "decode" and self.bootstrap_port is not None:
-            raise ValueError("Bootstrap port must not be set in DECODE mode.")
-        elif self.mode not in ["prefill", "decode"]:
-            raise ValueError(
-                f"Invalid mode: {self.mode}. Must be 'prefill' or 'decode'."
-            )
-
-
-def register_disaggregation_server(
-    mode: str, server_port: int, bootstrap_port: int, pdlb_url: str
-):
-    boostrap_port = bootstrap_port if mode == "prefill" else None
-    registry_request = PDRegistryRequest(
-        mode=mode,
-        registry_url=f"http://{get_ip()}:{server_port}",
-        bootstrap_port=boostrap_port,
-    )
-    res = requests.post(
-        f"{pdlb_url}/register",
-        json=dataclasses.asdict(registry_request),
-    )
-    if res.status_code != 200:
-        warnings.warn(
-            f"Failed to register disaggregation server: {res.status_code} {res.text}"
-        )
-
-
 #########################
 # Misc
 #########################

sglang/srt/distributed/device_communicators/all_reduce_utils.py

@@ -0,0 +1,16 @@
+MiB = 1024 * 1024
+
+SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
+    9: {
+        2: 64 * MiB,  # 64 MB
+        4: 32 * MiB,  # 32 MB
+        6: 64 * MiB,  # 64 MB
+        8: 64 * MiB,  # 64 MB
+    },
+    10: {
+        2: 64 * MiB,  # 64 MB
+        4: 32 * MiB,  # 32 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+}

sglang/srt/distributed/device_communicators/shm_broadcast.py

@@ -18,7 +18,7 @@ from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context  # type: ignore
 
 from sglang.srt.utils import (
     format_tcp_address,
-    get_ip,
+    get_local_ip_auto,
     get_open_port,
     is_valid_ipv6_address,
 )
@@ -191,7 +191,9 @@ class MessageQueue:
         self.n_remote_reader = n_remote_reader
 
         if connect_ip is None:
-            connect_ip = get_ip() if n_remote_reader > 0 else "127.0.0.1"
+            connect_ip = (
+                get_local_ip_auto("0.0.0.0") if n_remote_reader > 0 else "127.0.0.1"
+            )
 
         context = Context()
 

sglang/srt/distributed/device_communicators/symm_mem.py

@@ -0,0 +1,164 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/bf214ca22625e311a2c4c0dfbf7af19128f4919c/vllm/distributed/device_communicators/symm_mem.py
+import logging
+from typing import Optional, Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.distributed.device_communicators.all_reduce_utils import (
+    SYMM_MEM_ALL_REDUCE_MAX_SIZES,
+)
+from sglang.srt.utils import get_device_capability, is_cuda, is_hip
+
+try:
+    import torch.distributed._symmetric_memory as torch_symm_mem
+
+    symm_mem_available = True
+except ImportError:
+    symm_mem_available = False
+
+
+logger = logging.getLogger(__name__)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+symm_mem_is_available = False
+if _is_hip:
+    symm_mem_is_available = False
+if _is_cuda:
+    symm_mem_is_available = True
+
+
+class SymmMemCommunicator:
+    """
+    Thin wrapper around symmetric-memory collectives.
+
+    This communicator:
+      - Validates device capability and world size.
+      - Allocates a shared symmetric buffer.
+      - Chooses between 'multimem' and 'two-shot' all-reduce kernels.
+      - Exposes a fast-path all_reduce() compatible with bfloat16 inputs.
+
+    If any prerequisite is not met, the instance remains disabled and will
+    decline to perform symmetric-memory all-reduce.
+    """
+
+    # Mapping: compute capability major -> supported world sizes for multimem
+    # If the current (cc_major, world_size) is not listed, we fall back
+    # to the two-shot path.
+    _WORLD_SIZES_MULTIMEM = {
+        9: [4, 6, 8],
+        10: [6, 8],
+    }
+
+    def __init__(self, group: ProcessGroup, device: Union[int, str, torch.device]):
+        """
+        Args:
+            group: Torch process group used for rendezvous and naming.
+            device: Target CUDA device (index, 'cuda:X', or torch.device).
+        """
+
+        self.disabled = True
+
+        if not symm_mem_available:
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        torch.cuda.set_device(device)
+        self.dtype = torch.bfloat16
+        self.device = device
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+        self.device_capability = torch.cuda.get_device_capability(device)[0]
+        if self.device_capability < 9:
+            logger.warning(
+                "SymmMemCommunicator: Device capability %s not supported, "
+                "communicator is not available.",
+                self.device_capability,
+            )
+            return
+        if self.world_size not in SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability]:
+            logger.warning(
+                "SymmMemCommunicator: World size %d not supported, "
+                "communicator is not available.",
+                self.world_size,
+            )
+            return
+        self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability][
+            self.world_size
+        ]
+        self.buffer = torch_symm_mem.empty(
+            self.max_size // self.dtype.itemsize,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        handle = torch_symm_mem.rendezvous(self.buffer, self.group.group_name)
+        if handle.multicast_ptr == 0:
+            logger.warning(
+                "SymmMemCommunicator: symmetric memory "
+                "multicast operations are not supported."
+            )
+            self.buffer = None
+            self.disabled = True
+            return
+        self.disabled = False
+
+    def should_symm_mem_allreduce(self, inp: torch.Tensor):
+        """
+        Fast-path eligibility check for a given tensor.
+
+        Conditions:
+          - Communicator must be enabled.
+          - dtype must be bfloat16 (matches kernel + buffer dtype).
+          - Total byte size must be 4-byte aligned (hardware requirement).
+          - Payload must be smaller than the symmetric-memory max size.
+
+        Returns:
+            True if the symmetric-memory path can handle this tensor.
+        """
+        if self.disabled:
+            return False
+        if inp.dtype != self.dtype:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # enforce 4-byte alignment
+        if inp_size % 4 != 0:
+            return False
+        return inp_size < self.max_size
+
+    def all_reduce(
+        self, inp: torch.Tensor, *, out: Optional[torch.Tensor] = None
+    ) -> Optional[torch.Tensor]:
+        """
+        Perform an in-place sum all-reduce via symmetric memory.
+
+        Args:
+            inp: Input tensor on the target CUDA device (bfloat16).
+            out: Optional output tensor; if omitted, a new tensor is allocated.
+
+        Returns:
+            The reduced tensor (same shape as inp), or None if disabled.
+
+        Implementation details:
+            - Stages 'inp' into the symmetric buffer.
+            - Selects 'multimem' or 'two_shot' kernel based on topology.
+            - Writes the result into 'out' and returns it.
+        """
+        if out is None:
+            out = torch.empty_like(inp)
+        self.buffer[: inp.numel()].copy_(inp.view(-1))
+        if self.world_size in self._WORLD_SIZES_MULTIMEM[self.device_capability]:
+            torch.ops.symm_mem.multimem_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        else:
+            torch.ops.symm_mem.two_shot_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        out.copy_(self.buffer[: inp.numel()].view(out.shape))
+        return out