d9d 0.1.0__py3-none-any.whl

d9d/module/block/moe/communications/base.py

@@ -0,0 +1,58 @@
+import abc
+
+import torch
+
+
+class ExpertCommunicationHandler(abc.ABC):
+    """Abstract base class for Mixture-of-Experts communication strategies."""
+
+    @abc.abstractmethod
+    def dispatch(
+            self,
+            hidden_states: torch.Tensor,
+            topk_ids: torch.Tensor,
+            topk_weights: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Prepares and routes local hidden states to their target experts (possibly on other workers).
+
+        This process involves:
+
+        1. All-to-All Communication: Transfers hidden states to workers containing the assigned experts. States
+        assigned to multiple experts are replicated.
+
+        2. Permutation: Sorts tokens by expert ID to prepare for Grouped GEMM.
+
+        Args:
+            hidden_states: Input tokens. Shape: `(num_tokens, hidden_size)`.
+            topk_ids: Indices of the top-k experts selected for each token. Shape: `(num_tokens, k)`.
+            topk_weights: Routing weights associated with the selected experts. Shape: `(num_tokens, k)`.
+
+        Returns:
+            A tuple containing:
+
+            - Permuted hidden states received by this rank. Shape: `(num_received_tokens, hidden_size)`.
+            - Permuted weights matching the hidden states order. Shape: `(num_received_tokens)`.
+            - Expert count tensor indicating how many tokens each local expert received. Shape: `(num_local_experts)`.
+        """
+
+        ...
+
+    @abc.abstractmethod
+    def combine(
+            self,
+            hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Restores hidden states to their original order and location.
+
+        Undoes the permutation and performs the reverse All-to-All communication
+        to return processed results to the workers that originated the requests.
+
+        Args:
+            hidden_states: The processed hidden states. Shape: `(num_received_tokens, hidden_size)`.
+
+        Returns:
+            The combined hidden states with the original shape and order. Shape: `(num_tokens, hidden_size)`.
+        """
+        ...

d9d/module/block/moe/communications/deepep.py

@@ -0,0 +1,300 @@
+from typing import Any
+
+import torch
+from deep_ep import Buffer, EventOverlap
+from torch.autograd.function import FunctionCtx
+
+from d9d.kernel.moe.indices_to_multihot import fused_indices_to_multihot
+from d9d.kernel.moe.permute_with_probs import moe_permute_with_probs, moe_unpermute_mask
+from d9d.module.block.moe.communications import ExpertCommunicationHandler
+
+# see https://github.com/deepseek-ai/DeepEP/blob/main/README.md for examples
+# TODO: implement computation/communication overlap for PP case
+
+_buffer: Buffer | None = None
+
+
+def get_hidden_state_bytes(x: torch.Tensor) -> int:
+    """
+    Calculates the byte size of a hidden state tensor row.
+
+    Args:
+        x: Input tensor. Shape: `(?, hidden_size)`.
+    """
+
+    return x.size(1) * max(x.element_size(), 2)
+
+
+def init_deepep_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
+    """
+    Initializes or expands the global DeepEP communication buffer.
+
+    Checks if the existing buffer is sufficient for the required hidden dimension
+    and process group size. If not, it allocates a new buffer.
+
+    Args:
+        group: The process group intended for communication.
+        hidden_bytes: Size of a single hidden state vector in bytes.
+    """
+
+    global _buffer  # noqa: PLW0603
+    num_nvl_bytes, num_rdma_bytes = 0, 0
+    for config in (
+            Buffer.get_dispatch_config(group.size()),
+            Buffer.get_combine_config(group.size()),
+    ):
+        num_nvl_bytes = max(
+            config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes
+        )
+        num_rdma_bytes = max(
+            config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
+        )
+
+    # Allocate buffer if not existed or not enough buffer
+    if (
+            _buffer is None
+            or _buffer.group != group
+            or _buffer.num_nvl_bytes < num_nvl_bytes
+            or _buffer.num_rdma_bytes < num_rdma_bytes
+    ):
+        _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)
+
+
+class DeepEpDispatch(torch.autograd.Function):
+    """Autograd function for the DeepEP Dispatch operation."""
+
+    @staticmethod
+    def forward(
+            ctx: FunctionCtx,
+            x: torch.Tensor,
+            topk_idx: torch.Tensor,
+            topk_weights: torch.Tensor,
+            num_experts: int
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        list,
+        tuple,
+        EventOverlap
+    ]:
+        previous_event = Buffer.capture()
+        (
+            num_tokens_per_rank,
+            num_tokens_per_rdma_rank,
+            num_tokens_per_expert,
+            is_token_in_rank,
+            previous_event
+        ) = _buffer.get_dispatch_layout(
+            topk_idx, num_experts,
+            previous_event=previous_event,
+            async_finish=True,
+            allocate_on_comm_stream=True
+        )
+
+        (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            num_recv_tokens_per_expert_list,
+            handle,
+            event
+        ) = _buffer.dispatch(
+            x,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            num_tokens_per_rank=num_tokens_per_rank,
+            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
+            is_token_in_rank=is_token_in_rank,
+            num_tokens_per_expert=num_tokens_per_expert,
+            previous_event=previous_event,
+            async_finish=True,
+            allocate_on_comm_stream=True
+        )
+
+        event.current_stream_wait()
+
+        num_recv_tokens_per_expert_list = torch.tensor(num_recv_tokens_per_expert_list)
+
+        ctx.handle = handle
+
+        return (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            num_recv_tokens_per_expert_list,
+            handle
+        )
+
+    @staticmethod
+    def backward(
+            ctx: FunctionCtx,
+            grad_recv_x: torch.Tensor,
+            grad_recv_topk_idx: torch.Tensor,
+            grad_recv_topk_weights: torch.Tensor,
+            grad_num_recv_tokens_per_expert_list: list,
+            grad_handle: Any
+    ) -> tuple[
+        torch.Tensor,
+        None,
+        torch.Tensor,
+        None
+    ]:
+        handle = ctx.handle
+
+        prev_event = Buffer.capture()
+
+        (
+            combined_grad_x,
+            combined_grad_recv_topk_weights,
+            event
+        ) = _buffer.combine(
+            grad_recv_x.contiguous(),
+            handle,
+            topk_weights=grad_recv_topk_weights,
+            async_finish=True,
+            previous_event=prev_event,
+            allocate_on_comm_stream=True
+        )
+
+        event.current_stream_wait()
+
+        return combined_grad_x, None, combined_grad_recv_topk_weights, None
+
+
+class DeepEpCombine(torch.autograd.Function):
+    """Autograd function for the DeepEP Combine operation."""
+
+    @staticmethod
+    def forward(
+            ctx: FunctionCtx,
+            x: torch.Tensor,
+            handle: Any
+    ) -> torch.Tensor:
+        previous_event = Buffer.capture()
+
+        combined_x, _, event = _buffer.combine(
+            x,
+            handle,
+            async_finish=True,
+            previous_event=previous_event,
+            allocate_on_comm_stream=True
+        )
+
+        event.current_stream_wait()
+
+        ctx.handle = handle
+
+        return combined_x
+
+    @staticmethod
+    def backward(ctx: FunctionCtx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
+        handle = ctx.handle
+
+        previous_event = Buffer.capture()
+
+        grad_x, _, _, _, _, event = _buffer.dispatch(
+            grad_output.contiguous(),
+            handle=handle,
+            async_finish=True,
+            previous_event=previous_event,
+            allocate_on_comm_stream=True
+        )
+
+        event.current_stream_wait()
+
+        return grad_x, None
+
+
+class DeepEpCommunicationHandler(ExpertCommunicationHandler):
+    """Handles MoE communication using the high-performance DeepEP library."""
+
+    def __init__(self, num_experts: int):
+        """Constructs the DeepEpCommunicationHandler."""
+
+        self._num_experts = num_experts
+        self._num_experts_per_shard = None  # late-initialization
+
+        # == fields saved for post-dispatch ==
+
+        self._handle = None
+        self._hidden_shape_before_permute = None
+        self._unpermute_mapping = None
+
+    def setup(self, group: torch.distributed.ProcessGroup, hidden_size: int, hidden_dtype: torch.dtype):
+        """
+        Initializes the backend buffer and calculates expert sharding.
+
+        Args:
+            group: The process group containing all experts.
+            hidden_size: Dimensionality of the hidden states.
+            hidden_dtype: Data type of the hidden states.
+        """
+
+        init_deepep_buffer(group, hidden_size * hidden_dtype.itemsize)
+
+        if self._num_experts % group.size() != 0:
+            raise ValueError("num_experts must be divisible by distributed group size")
+
+        self._num_experts_per_shard = self._num_experts // group.size()
+
+    def dispatch(
+            self,
+            hidden_states: torch.Tensor,
+            topk_ids: torch.Tensor,
+            topk_weights: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        (
+            hidden_states,
+            topk_ids,
+            topk_weights,
+            tokens_per_expert,
+            handle
+        ) = DeepEpDispatch.apply(
+            hidden_states,
+            topk_ids,
+            topk_weights,
+            self._num_experts
+        )
+
+        routing_map, routing_probs = fused_indices_to_multihot(
+            topk_ids, topk_weights, self._num_experts_per_shard
+        )
+
+        self._hidden_shape_before_permute = hidden_states.shape
+
+        hidden_states, routing_probs, reverse_permute_map = moe_permute_with_probs(
+            hidden_states,
+            routing_probs,
+            routing_map,
+            num_out_tokens=tokens_per_expert.sum().item()
+        )
+
+        self._handle = handle
+        self._unpermute_mapping = reverse_permute_map
+
+        return hidden_states, routing_probs, tokens_per_expert
+
+    def combine(
+            self,
+            hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        if self._handle is None:
+            raise ValueError("you fucked up moe communication order: you should dispatch first and after that combine")
+
+        hidden_states = moe_unpermute_mask(
+            hidden_states,
+            self._unpermute_mapping,
+            restore_shape=self._hidden_shape_before_permute,
+        )
+
+        hidden_states = DeepEpCombine.apply(
+            hidden_states,
+            self._handle
+        )
+
+        self._handle = None
+        self._unpermute_mapping = None
+        self._hidden_shape_before_permute = None
+
+        return hidden_states

d9d/module/block/moe/communications/naive.py

@@ -0,0 +1,68 @@
+from typing import cast
+
+import torch
+from torch import Size
+
+from d9d.kernel.moe import (
+    fused_indices_to_multihot,
+    moe_permute_with_probs,
+    moe_unpermute_mask,
+)
+from d9d.module.block.moe.communications import ExpertCommunicationHandler
+
+
+class NoCommunicationHandler(ExpertCommunicationHandler):
+    """
+    Handles MoE routing within a single device or when no cross-device routing is needed.
+
+    This handler does not perform network operations. It only permutes elements
+    mostly for local logical grouping or debugging.
+    """
+
+    def __init__(self, num_experts: int):
+        """Constructs the NoCommunicationHandler."""
+        self._num_experts = num_experts
+
+        self._hidden_shape_before_permute: Size | None = None
+        self._unpermute_mapping: torch.Tensor | None = None
+
+    def dispatch(
+            self,
+            hidden_states: torch.Tensor,
+            topk_ids: torch.Tensor,
+            topk_weights: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        with torch.no_grad():
+            tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=self._num_experts).cpu()
+
+        routing_map, routing_probs = fused_indices_to_multihot(
+            topk_ids, topk_weights, self._num_experts
+        )
+
+        self._hidden_shape_before_permute = hidden_states.shape
+
+        hidden_states, routing_probs, reverse_permute_map = moe_permute_with_probs(
+            hidden_states,
+            routing_probs,
+            routing_map,
+            num_out_tokens=cast(int, tokens_per_expert.sum().item())
+        )
+
+        self._unpermute_mapping = reverse_permute_map
+
+        return hidden_states, routing_probs, tokens_per_expert
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self._unpermute_mapping is None:
+            raise ValueError("Cannot run combine before running dispatch!")
+
+        hidden_states = moe_unpermute_mask(
+            hidden_states,
+            self._unpermute_mapping,
+            restore_shape=self._hidden_shape_before_permute,
+        )
+
+        self._unpermute_mapping = None
+        self._hidden_shape_before_permute = None
+
+        return hidden_states

d9d/module/block/moe/grouped_experts.py

@@ -0,0 +1,81 @@
+import torch
+from torch import nn
+
+from d9d.kernel.swiglu import silu_mul
+from d9d.module.base import ModuleLateInit
+
+from .grouped_linear import GroupedLinear
+
+
+class GroupedSwiGLU(nn.Module, ModuleLateInit):
+    """
+    Executes a collection of SwiGLU experts efficiently using Grouped GEMM.
+
+    This module implements the architectural pattern: `down_proj(SiLU(gate_proj(x)) * up_proj(x))`.
+    It applies this operation across multiple discrete experts in parallel without padding or masking.
+    """
+
+    def __init__(
+            self,
+            hidden_dim: int,
+            intermediate_dim: int,
+            num_experts: int
+    ):
+        """
+        Constructs the GroupedSwiGLU module.
+
+        Args:
+            hidden_dim: Dimensionality of the input and output hidden states.
+            intermediate_dim: Dimensionality of the intermediate projection.
+            num_experts: Total number of experts managed by this local instance.
+        """
+
+        super().__init__()
+        self._num_experts = num_experts
+
+        self.gate_proj = GroupedLinear(num_experts, hidden_dim, intermediate_dim)
+        self.up_proj = GroupedLinear(num_experts, hidden_dim, intermediate_dim)
+        self.down_proj = GroupedLinear(num_experts, intermediate_dim, hidden_dim)
+
+    def forward(
+            self,
+            permuted_x: torch.Tensor,
+            permuted_probs: torch.Tensor,
+            tokens_per_expert: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Computes expert outputs for sorted input tokens.
+
+        Args:
+            permuted_x: Input tokens sorted by their assigned expert.
+                Shape: `(total_tokens, hidden_dim)`.
+            permuted_probs: Routing weights/probabilities corresponding to the sorted tokens.
+                Shape: `(total_tokens)`.
+            tokens_per_expert: Number of tokens assigned to each consecutive expert. It is a CPU tensor.
+                Shape: `(num_experts)`.
+
+        Returns:
+            The computed and weighted output tokens (still permuted).
+            Shape: `(total_tokens, hidden_dim)`.
+        """
+
+        if permuted_x.numel() == 0:  # handle cases when there are no routed experts to this instance
+            return permuted_x
+
+        probs = permuted_probs[:, None].to(permuted_x.dtype)
+        values = self.down_proj(
+            silu_mul(
+                self.gate_proj(permuted_x, tokens_per_expert),
+                self.up_proj(permuted_x, tokens_per_expert)
+            ),
+            tokens_per_expert
+        )
+
+        return probs * values
+
+    def reset_parameters(self):
+        """Resets parameters for all internal linear projections."""
+
+        self.gate_proj.reset_parameters()
+        self.up_proj.reset_parameters()
+        self.down_proj.reset_parameters()

d9d/module/block/moe/grouped_linear.py

@@ -0,0 +1,78 @@
+import math
+
+import torch
+from torch import nn
+from torch.distributed.tensor import DTensor
+
+from d9d.core.autograd import GradDirection
+from d9d.kernel.gmm import gmm
+from d9d.module.base import ModuleLateInit
+
+
+class GroupedLinear(nn.Module, ModuleLateInit):
+    """
+    Applies a linear transformation using Grouped GEMM (Generalized Matrix Multiplication).
+
+    This module allows efficient execution of multiple linear layers (experts) in parallel, where each expert
+    processes a variable number of tokens.
+    It is the computational core of the Mixture-of-Experts layer.
+    """
+
+    def __init__(
+            self,
+            n_groups: int,
+            in_features: int,
+            out_features: int,
+            device: torch.device | str | None = None,
+            dtype: torch.dtype | None = None
+    ):
+        """
+        Constructs the GroupedLinear layer.
+
+        Args:
+            n_groups: Number of groups (experts).
+            in_features: Input hidden size.
+            out_features: Output hidden size.
+            device: Target device.
+            dtype: Target data type.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(n_groups, in_features, out_features,
+                                               device=device, dtype=dtype))
+
+        self.n_groups = n_groups
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor, x_groups: torch.Tensor) -> torch.Tensor:
+        """
+        Performs the grouped matrix multiplication.
+
+        Args:
+            x: Flattened input tensor containing tokens for all groups.
+                Shape: `(total_tokens, in_features)`.
+            x_groups: CPU Tensor indicating the number of tokens assigned to each group.
+                Must sum to `total_tokens`. Shape: `(n_groups,)`.
+
+        Returns:
+            The output tensor. Shape: `(total_tokens, out_features)`.
+        """
+
+        weight: torch.Tensor = self.weight
+
+        if isinstance(weight, DTensor):
+            weight = weight.to_local()
+
+        return gmm(
+            x,
+            weight,
+            x_groups,
+            a_grad_direction=GradDirection.inputs,
+            b_grad_direction=GradDirection.weight
+        )
+
+    def reset_parameters(self):
+        """Initializes weights using a uniform distribution based on input features."""
+        nn.init.uniform_(self.weight, -1 / math.sqrt(self.in_features), 1 / math.sqrt(self.in_features))

d9d/module/block/moe/layer.py

@@ -0,0 +1,122 @@
+import torch
+from torch import nn
+from torch.distributed import ProcessGroup
+
+from d9d.module.base import ModuleLateInit
+
+from .communications import (
+    DeepEpCommunicationHandler,
+    ExpertCommunicationHandler,
+    NoCommunicationHandler,
+)
+from .grouped_experts import GroupedSwiGLU
+from .router import TopKRouter
+
+# TODO: implement expert bias
+# TODO: shared experts
+
+
+class MoELayer(nn.Module, ModuleLateInit):
+    """
+    A complete Mixture-of-Experts (MoE) block comprising routing, communication, and computation.
+
+    This layer integrates:
+
+    1.  **Router**: Selects experts for each token.
+    2.  **Communicator**: Handles token dispatch to local or remote experts (EP).
+    3.  **Experts**: Performs parallelized computation (Grouped SwiGLU).
+    """
+
+    def __init__(
+            self,
+            hidden_dim: int,
+            intermediate_dim_grouped: int,
+            num_grouped_experts: int,
+            top_k: int,
+            router_renormalize_probabilities: bool
+    ):
+        """
+        Constructs the MoELayer.
+
+       Args:
+           hidden_dim: Hidden size.
+           intermediate_dim_grouped: Intermediate dimension for the Expert FFNs.
+           num_grouped_experts: Total number of experts.
+           top_k: Number of experts to route each token to.
+           router_renormalize_probabilities: Configures router probability normalization behavior.
+       """
+
+        super().__init__()
+        self.router = TopKRouter(
+            dim=hidden_dim, num_experts=num_grouped_experts, top_k=top_k,
+            renormalize_probabilities=router_renormalize_probabilities
+        )
+        self.grouped_experts = GroupedSwiGLU(
+            hidden_dim=hidden_dim,
+            intermediate_dim=intermediate_dim_grouped,
+            num_experts=num_grouped_experts
+        )
+        self._communicator: ExpertCommunicationHandler = NoCommunicationHandler(num_grouped_experts)
+
+        self._num_grouped_experts = num_grouped_experts
+        self._hidden_dim = hidden_dim
+
+        self.tokens_per_expert = nn.Buffer(torch.empty((num_grouped_experts,), dtype=torch.int64), persistent=False)
+
+    def enable_distributed_communicator(self, group: ProcessGroup):
+        """
+        Switches from local no-op communication to distributed DeepEP communication.
+
+        This should be called during model initialization if the model is running in a
+        distributed Expert Parallel environment.
+
+        Args:
+            group: The PyTorch process group spanning the expert parallel ranks.
+        """
+
+        communicator = DeepEpCommunicationHandler(num_experts=self._num_grouped_experts)
+        communicator.setup(group, self._hidden_dim, self.router.gate.weight.dtype)
+        self._communicator = communicator
+
+    @torch.no_grad()
+    def _update_tokens_per_expert(self, expert_indices: torch.Tensor):
+        self.tokens_per_expert.add_(expert_indices.view(-1).bincount(minlength=self._num_grouped_experts))
+
+    @torch.no_grad()
+    def reset_stats(self):
+        """Resets the expert load balancing counters."""
+        self.tokens_per_expert.zero_()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Routes tokens to experts, computes, and combines results.
+
+        Args:
+            hidden_states: Input tensor. Shape: `(batch_size, seq_len, hidden_dim)`.
+
+        Returns:
+            Output tensor combined from experts. Shape: `(batch_size, seq_len, hidden_dim)`.
+        """
+
+        old_shape = hidden_states.shape
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+        expert_indices, expert_scores = self.router(hidden_states)
+        self._update_tokens_per_expert(expert_indices)
+        hidden_states, expert_scores, expert_count = self._communicator.dispatch(
+            hidden_states, expert_indices, expert_scores
+        )
+        hidden_states = self.grouped_experts(hidden_states, expert_scores, expert_count)
+        hidden_states = self._communicator.combine(hidden_states)
+        hidden_states = hidden_states.reshape(*old_shape)
+
+        return hidden_states
+
+    def reset_parameters(self):
+        """Resets module parameters."""
+        self.router.reset_parameters()
+        self.grouped_experts.reset_parameters()
+
+        nn.init.zeros_(self.tokens_per_expert)