perf: optimize logits allgather and parallelize eagle3 input projection

python/tokenspeed/runtime/distributed/comm_backend/auto.py

@@ -50,7 +50,7 @@ def __init__(self):
         self._custom_ar = CustomAllReduceBackend(fallback=self._nccl)
         self._trtllm_ar = TrtllmAllReduceBackend(fallback=self._nccl)
         self._triton_ar = TritonAllReduceBackend(fallback=self._nccl)
-        self._rsag = TritonRSAGBackend()
+        self._rsag = TritonRSAGBackend(fallback=self._nccl)
 
     @property
     def nccl(self) -> NcclBackend:
@@ -75,22 +75,18 @@ def configure(
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
-        return self._rsag.token_all_gather(tensor, rank, group, scattered_num_tokens)
+        return self._rsag.token_all_gather(tensor, group, scattered_num_tokens)
 
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
-        return self._rsag.token_reduce_scatter(
-            tensor, rank, group, scattered_num_tokens
-        )
+        return self._rsag.token_reduce_scatter(tensor, group, scattered_num_tokens)
 
     # ---- Public CommBackend interface ----
 
@@ -106,6 +102,9 @@ def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tenso
     def all_gather(
         self, tensor: torch.Tensor, group: Group, dim: int = 0
     ) -> torch.Tensor:
+        if tensor.dim() == 2 and dim in (-1, tensor.dim() - 1):
+            return self._rsag.all_gather(tensor, group, dim)
+
         return self._nccl.all_gather(tensor, group, dim)
 
     def all_gather_into_tensor(

python/tokenspeed/runtime/distributed/comm_backend/base.py

@@ -60,7 +60,6 @@ def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: ..
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor: ...
@@ -69,7 +68,6 @@ def token_all_gather(
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor: ...

python/tokenspeed/runtime/distributed/comm_backend/custom_allreduce.py

@@ -120,7 +120,6 @@ def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor:
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
@@ -129,7 +128,6 @@ def token_all_gather(
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:

python/tokenspeed/runtime/distributed/comm_backend/nccl.py

@@ -179,7 +179,6 @@ def recv(
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
@@ -216,7 +215,6 @@ def token_all_gather(
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
@@ -245,4 +243,5 @@ def token_reduce_scatter(
         torch.distributed.reduce_scatter_tensor(
             output, padded_input.contiguous(), group=res["device_group"]
         )
+        rank = group.index(torch.distributed.get_rank())
         return output[: scattered_num_tokens[rank]].contiguous()

python/tokenspeed/runtime/distributed/comm_backend/triton_allreduce.py

@@ -96,7 +96,6 @@ def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor:
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
@@ -105,7 +104,6 @@ def token_all_gather(
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:

python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py

@@ -25,13 +25,16 @@
 """
 
 import torch
+import torch.distributed as dist
 from tokenspeed_kernel.ops.communication.triton import (
     all_gather,
+    all_gather_inner,
     create_state,
     reduce_scatter,
 )
+from tokenspeed_kernel.platform import current_platform
 
-from tokenspeed.runtime.distributed.comm_backend.base import Group
+from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group
 from tokenspeed.runtime.distributed.process_group_manager import (
     process_group_manager as pg_manager,
 )
@@ -47,44 +50,76 @@ class TritonRSAGBackend:
     (group, hidden_size) pair because RSAG pre-allocates buffers.
     """
 
-    def __init__(self):
+    def __init__(self, fallback: CommBackend):
+        self._fallback = fallback
         # (group_tuple, hidden_size) -> Triton RS/AG state
         self._instances = {}
 
-    def _get_or_create(self, rank: int, group: Group, hidden_size: int):
+    def _get_or_create(self, group: Group, hidden_size: int):
         key = (group, hidden_size)
         if key in self._instances:
             return self._instances[key]
 
         max_num_tokens = self._get_max_num_gathered_tokens()
-        rsag = create_state(
+        state = create_state(
             group=pg_manager.get_process_group("nccl", group),
-            rank_in_group=rank,
+            rank_in_group=group.index(dist.get_rank()),
             max_tokens=max_num_tokens,
             hidden_size=hidden_size,
         )
-        self._instances[key] = rsag
-        return rsag
+        self._instances[key] = state
+        return state
+
+    def all_gather(
+        self,
+        tensor: torch.Tensor,
+        group: Group,
+        dim: int = 0,
+    ) -> torch.Tensor:
+        if tensor.dim() != 2:
+            return self._fallback.all_gather(tensor, group=group, dim=dim)
+
+        if dim == 0:
+            return self.token_all_gather(
+                tensor,
+                group=group,
+                scattered_num_tokens=[tensor.size(0)] * len(group),
+            )
+
+        if (
+            current_platform().is_nvidia
+            and dim in (-1, tensor.dim() - 1)
+            and tensor.dtype == torch.bfloat16
+        ):
+            hidden_size = tensor.size(-1) * len(group)
+            state = self._get_or_create(group, hidden_size)
+            return all_gather_inner(
+                state,
+                tensor,
+                tp_hidden_dim=hidden_size,
+                skip_entry_sync=False,
+                safe=False,
+            )
+
+        return self._fallback.all_gather(tensor, group=group, dim=dim)
 
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
-        rsag = self._get_or_create(rank, group, tensor.size(-1))
-        return all_gather(rsag, tensor, token_list_in_group=scattered_num_tokens)
+        state = self._get_or_create(group, tensor.size(-1))
+        return all_gather(state, tensor, token_list_in_group=scattered_num_tokens)
 
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
-        rsag = self._get_or_create(rank, group, tensor.size(-1))
-        return reduce_scatter(rsag, tensor, token_list_in_group=scattered_num_tokens)
+        state = self._get_or_create(group, tensor.size(-1))
+        return reduce_scatter(state, tensor, token_list_in_group=scattered_num_tokens)
 
     def _get_max_num_gathered_tokens(self):
         """Compute max buffer size for TritonRSAG.

python/tokenspeed/runtime/distributed/comm_backend/trtllm_allreduce.py

@@ -194,7 +194,6 @@ def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor:
     def token_all_gather(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:
@@ -203,7 +202,6 @@ def token_all_gather(
     def token_reduce_scatter(
         self,
         tensor: torch.Tensor,
-        rank: int,
         group: Group,
         scattered_num_tokens: list[int],
     ) -> torch.Tensor:

python/tokenspeed/runtime/distributed/comm_manager.py

@@ -120,7 +120,6 @@ def pre_attn_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext):
 
         return token_all_gather(
             hidden_states,
-            rank=self.mapping.attn.tp_rank,
             group=self.mapping.attn.tp_group,
             scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx),
         )
@@ -132,24 +131,20 @@ def post_attn_comm(
             return hidden_states, residual
 
         if self.use_all_reduce(self.is_moe):
-            hidden_states = all_reduce(
-                hidden_states, self.mapping.attn.tp_rank, self.mapping.attn.tp_group
-            )
+            hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group)
             # The output residual is expected to have attn_tp_num_tokens.
             # For first layer, the input residual has attn_tp_num_tokens.
             # Otherwise, if this layer experiences a RSAG -> AR switch, residual needs allgather.
             if self.layer_id > 0 and not self.use_all_reduce(self.prev_is_moe):
                 residual = token_all_gather(
                     residual,
-                    rank=self.mapping.attn.tp_rank,
                     group=self.mapping.attn.tp_group,
                     scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx),
                 )
         else:
             token_list = self.attn_tp_group_scattered_num_tokens(ctx)
             hidden_states = token_reduce_scatter(
                 hidden_states,
-                rank=self.mapping.attn.tp_rank,
                 group=self.mapping.attn.tp_group,
                 scattered_num_tokens=token_list,
             )
@@ -177,7 +172,6 @@ def pre_dense_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext):
 
         return token_all_gather(
             hidden_states,
-            rank=self.mapping.dense.tp_rank,
             group=self.mapping.dense.tp_group,
             scattered_num_tokens=self.dense_tp_group_scattered_num_tokens(ctx),
         )
@@ -191,7 +185,6 @@ def pre_moe_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext):
 
         return token_all_gather(
             hidden_states,
-            rank=self.mapping.moe.tp_ep_rank,
             group=self.mapping.moe.tp_ep_group,
             scattered_num_tokens=self.moe_tp_ep_group_scattered_num_tokens(ctx),
         )
@@ -211,13 +204,10 @@ def post_dense_comm(
             return hidden_states, residual
 
         if self.use_all_reduce(is_moe=False):
-            hidden_states = all_reduce(
-                hidden_states, self.mapping.dense.tp_rank, self.mapping.dense.tp_group
-            )
+            hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group)
             return hidden_states, residual
         hidden_states = token_reduce_scatter(
             hidden_states,
-            rank=self.mapping.dense.tp_rank,
             group=self.mapping.dense.tp_group,
             scattered_num_tokens=self.dense_tp_group_scattered_num_tokens(ctx),
         )
@@ -230,13 +220,10 @@ def post_moe_comm(
             return hidden_states, residual
 
         if self.use_all_reduce(is_moe=True):
-            hidden_states = all_reduce(
-                hidden_states, self.mapping.moe.tp_ep_rank, self.mapping.moe.tp_ep_group
-            )
+            hidden_states = all_reduce(hidden_states, self.mapping.moe.tp_ep_group)
             return hidden_states, residual
         hidden_states = token_reduce_scatter(
             hidden_states,
-            rank=self.mapping.moe.tp_ep_rank,
             group=self.mapping.moe.tp_ep_group,
             scattered_num_tokens=self.moe_tp_ep_group_scattered_num_tokens(ctx),
         )
@@ -251,7 +238,6 @@ def post_final_norm_comm(
             return hidden_states, residual
         hidden_states = token_all_gather(
             hidden_states,
-            rank=self.mapping.attn.tp_rank,
             group=self.mapping.attn.tp_group,
             scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx),
         )