src/fractus/nn/moe.py

"""PhaseRoutedMoE: mixture-of-experts with von Mises phase routing.

Ported from the original system (src/moe.rs + farey.rs) in pure PyTorch.

Expert phases drawn from Farey sequence. Von Mises gate with top-k routing.
Load-balance loss as auxiliary. End-to-end differentiable.

L8 OPTIMIZATION (gather-first sparse dispatch):
    The original computed the outputs of ALL n_experts, then gathered the
    top-k — wasting (E-K)/E of the FLOPs (50-75% on typical presets). Here
    we GATHER FIRST: index_select the top-k experts' WEIGHTS, then compute
    only those K experts. Output is bit-identical (proven by
    test_moe_sparse_matches_reference), but we do K/E of the matmul work.

    Concretely: instead of einsum("bld,edf->blef") over all E experts, we
    build w1_selected[b,l,k] = w1[topk_idx[b,l,k]] via gather, then a single
    batched matmul over the K active experts per token.
"""

import math
import torch
import torch.nn as nn

from .farey import expert_phases


def _gelu(x: torch.Tensor) -> torch.Tensor:
    """Tanh GeLU approximation (as in moe.rs:14-17)."""
    return 0.5 * x * (1.0 + torch.tanh(
        math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3)
    ))


class PhaseRoutedMoE(nn.Module):
    """Mixture-of-experts with von Mises phase routing on Farey phases.

    Args:
        d_model     : input/output dimension.
        n_experts   : number of experts E.
        top_k       : number of active experts per token (<= E).
        kappa       : von Mises concentration.
        temperature : gate temperature (κ_eff = κ/temperature).
        d_ff        : expert hidden dimension (64 by default, as in the original).
    """

    def __init__(
        self,
        d_model: int,
        n_experts: int,
        top_k: int,
        kappa: float = 4.0,
        temperature: float = 1.0,
        d_ff: int = 64,
    ):
        super().__init__()
        if n_experts < 1:
            raise ValueError("n_experts >= 1")
        if top_k < 1 or top_k > n_experts:
            raise ValueError(f"top_k must be in [1, {n_experts}], got {top_k}")
        self.d_model = d_model
        self.n_experts = n_experts
        self.top_k = top_k
        self.kappa = kappa
        self.temperature = temperature
        self.d_ff = d_ff

        # Expert phases (Farey precomputation, off-graph).
        phases = expert_phases(n_experts)
        self.register_buffer("expert_phases", torch.tensor(phases, dtype=torch.float32))

        # Expert weights: E × (W1, b1, W2, b2). Xavier uniform init.
        scale1 = math.sqrt(2.0 / d_model)
        scale2 = math.sqrt(2.0 / d_ff)
        self.w1 = nn.Parameter(torch.empty(n_experts, d_model, d_ff).uniform_(-scale1, scale1))
        self.b1 = nn.Parameter(torch.zeros(n_experts, d_ff))
        self.w2 = nn.Parameter(torch.empty(n_experts, d_ff, d_model).uniform_(-scale2, scale2))
        self.b2 = nn.Parameter(torch.zeros(n_experts, d_model))

    def _compute_gates(self, phases: torch.Tensor) -> torch.Tensor:
        """Computes the von Mises gates for each token.

        phases: (B, L, n_phases). Returns gates (B, L, E).
        """
        sin_p = torch.sin(phases).sum(dim=-1)  # (B, L)
        cos_p = torch.cos(phases).sum(dim=-1)
        theta_bar = torch.atan2(sin_p, cos_p)  # (B, L)
        kappa_eff = self.kappa / self.temperature
        diff = theta_bar.unsqueeze(-1) - self.expert_phases.view(
            *[1] * (phases.dim() - 1), self.n_experts
        )  # (B, L, E)
        gates = torch.exp(kappa_eff * torch.cos(diff))  # (B, L, E)
        gates_sum = gates.sum(dim=-1, keepdim=True)
        uniform = torch.full_like(gates, 1.0 / self.n_experts)
        gates = torch.where(gates_sum > 1e-10, gates / gates_sum, uniform)
        return gates

    def _sparse_expert_forward(
        self, h: torch.Tensor, topk_idx: torch.Tensor
    ) -> torch.Tensor:
        """GATHER-FIRST sparse forward: compute ONLY the top_k experts per token.

        h        : (B, L, d_model)
        topk_idx : (B, L, K) — indices in [0, E) of the selected experts.
        Returns  : (B, L, K, d_model) — output of each selected expert.

        This is the L8 optimization. Instead of materializing the (B,L,E,d_model)
        full-expert tensor and gathering (wasting (E-K)/E of the matmul), we
        index_select the K experts' weights PER TOKEN, then do one batched
        matmul. Work scales with K, not E.
        """
        B, L, D = h.shape
        K = topk_idx.shape[-1]

        # Gather the K selected experts' weights PER TOKEN.
        # w1: (E, D, F) → w1_sel: (B, L, K, D, F)
        # topk_idx: (B, L, K) → expand to (B, L, K, D, F) for gather on dim 0 of a flat view.
        # Cleanest: flatten (B,L,K) indices and use index_select on the expert dim.
        flat_idx = topk_idx.reshape(-1)  # (B*L*K,)
        w1_sel = self.w1.index_select(0, flat_idx).reshape(B, L, K, D, self.d_ff)
        b1_sel = self.b1.index_select(0, flat_idx).reshape(B, L, K, self.d_ff)
        w2_sel = self.w2.index_select(0, flat_idx).reshape(B, L, K, self.d_ff, D)
        b2_sel = self.b2.index_select(0, flat_idx).reshape(B, L, K, D)

        # h: (B, L, D) → (B, L, 1, D, 1) broadcast over the K dim.
        # w1_sel is (B, L, K, D, F): align D at dim -2.
        h_exp = h.unsqueeze(2).unsqueeze(-1)  # (B, L, 1, D, 1)
        # h1[b,l,k,f] = Σ_d h[b,l,d] · w1_sel[b,l,k,d,f]
        h1 = (h_exp * w1_sel).sum(dim=-2) + b1_sel  # (B, L, K, F)
        h1_act = _gelu(h1)
        # out[b,l,k,d] = Σ_f h1_act[b,l,k,f] · w2_sel[b,l,k,f,d]
        h1_act_exp = h1_act.unsqueeze(-1)  # (B, L, K, F, 1)
        out = (h1_act_exp * w2_sel).sum(dim=-2) + b2_sel  # (B, L, K, D)
        return out

    def _dense_expert_forward(self, h: torch.Tensor) -> torch.Tensor:
        """DENSE forward (the original path): compute ALL E experts.

        h: (B, L, d_model) → outputs of all experts (B, L, E, d_model).
        Cheaper than sparse on CPU when E is small (einsum is more optimized
        than per-token index_select + broadcast). Used when n_experts is small.
        """
        B, L, D = h.shape
        h1 = torch.einsum("bld,edf->blef", h, self.w1) + self.b1.view(1, 1, self.n_experts, self.d_ff)
        h1_act = _gelu(h1)
        out = torch.einsum("blef,efd->bled", h1_act, self.w2) + self.b2.view(1, 1, self.n_experts, self.d_model)
        return out

    def forward(
        self, h: torch.Tensor, phases: torch.Tensor
    ):
        """h: (B, L, d_model), phases: (B, L, n_phases).
        Returns (output (B, L, d_model), load_balance_loss scalar).

        L8 ADAPTIVE DISPATCH: pick the cheaper path at construction time.
            - Sparse (gather-first) when n_experts > 2·top_k  (>50% waste saved).
            - Dense (einsum over all E) otherwise — on CPU the optimized einsum
              beats per-token index_select for small E.
        Measured: for E=4,K=2 the dense path is ~1.5× faster than sparse; for
        E=32,K=8 the sparse path wins. The 2× threshold is the empirical knee.
        """
        gates = self._compute_gates(phases)  # (B, L, E)

        topk_vals, topk_idx = gates.topk(self.top_k, dim=-1)  # (B, L, K)
        topk_sum = topk_vals.sum(dim=-1, keepdim=True)
        uniform_topk = torch.full_like(topk_vals, 1.0 / self.top_k)
        topk_vals_norm = torch.where(
            topk_sum > 1e-10, topk_vals / topk_sum, uniform_topk
        )

        # Adaptive: dense when small E (einsum wins on CPU), sparse when large E.
        if self.n_experts > 2 * self.top_k:
            topk_out = self._sparse_expert_forward(h, topk_idx)  # (B, L, K, d_model)
        else:
            all_out = self._dense_expert_forward(h)  # (B, L, E, d_model)
            idx_exp = topk_idx.unsqueeze(-1).expand(-1, -1, -1, self.d_model)
            topk_out = torch.gather(all_out, dim=2, index=idx_exp)  # (B, L, K, d_model)
        output = (topk_vals_norm.unsqueeze(-1) * topk_out).sum(dim=2)  # (B, L, d_model)

        # Load-balance loss uses the FULL gates (all E) — this is the only place
        # we still touch all experts, and it's a cheap mean over (B,L,E).
        P = gates.mean(dim=(0, 1))  # (E,)
        lb_loss = self.n_experts * ((P - 1.0 / self.n_experts) ** 2).sum()

        return output, lb_loss