Mixture of Experts

"""Sparse mixture of experts, transcribed from the objective.

No ML library. One router, E expert MLPs, top-k selection per token, and
the auxiliary load-balancing loss written out so the two terms of

    L = L_task + alpha * E * sum_e f_e * P_e

can be read against each other.

The thing to notice: `+ '`' + `f_e` is a hard count of tokens routed to expert
e, so it carries no gradient. `+ '`' + `P_e` is the mean router probability
and is soft. All the gradient flows through P_e, with f_e acting only as
a per-expert weight. That asymmetry is how a discrete routing decision
gets a differentiable penalty, and it is the whole trick.
"""

import math


def softmax(values):
    top = max(values)
    exps = [math.exp(value - top) for value in values]
    total = sum(exps)
    return [value / total for value in exps]


def matvec(matrix, vector):
    return [
        sum(row[idx] * vector[idx] for idx in range(len(vector)))
        for row in matrix
    ]


def relu(values):
    return [value if value > 0.0 else 0.0 for value in values]


def expert_forward(token, w_in, w_out):
    """One expert: a plain two-layer MLP. Nothing special about it.

    The experts are ordinary. Everything interesting is in the router.
    """
    return matvec(w_out, relu(matvec(w_in, token)))


def route(token, router_weights, n_experts, top_k):
    """Score every expert, keep the k best, renormalize over those.

    Renormalizing over the kept experts matters: without it the output
    magnitude depends on how much probability mass happened to land
    outside the top-k, which varies per token and injects noise the
    residual stream has to absorb.
    """
    logits = matvec(router_weights, token)
    probabilities = softmax(logits)

    ranked = sorted(range(n_experts), key=lambda expert: -probabilities[expert])
    chosen = ranked[:top_k]

    kept_mass = sum(probabilities[expert] for expert in chosen)
    weights = {expert: probabilities[expert] / kept_mass for expert in chosen}
    return weights, probabilities


def moe_layer(tokens, router_weights, experts, top_k, capacity_factor):
    """Forward pass with capacity, so token dropping is visible.

    Capacity is the ugly part and it is not an implementation detail.
    Batched expert computation needs a fixed buffer per expert, so an
    uneven batch overflows some buffers and the overflow tokens are
    DROPPED -- they pass through with no expert applied at all.

    Consequence worth stating plainly: the output for a token depends on
    which other tokens share its batch. In training that is noise. At
    serving time it means the same input can produce different outputs,
    which surprises people who expect a deterministic model.
    """
    n_tokens = len(tokens)
    n_experts = len(experts)
    width = len(tokens[0])

    capacity = max(1, int(n_tokens * top_k / n_experts * capacity_factor))

    outputs = [[0.0] * width for _ in range(n_tokens)]
    assignments = [[] for _ in range(n_experts)]
    dropped = 0

    router_probabilities = []
    for position, token in enumerate(tokens):
        weights, probabilities = route(token, router_weights, n_experts, top_k)
        router_probabilities.append(probabilities)
        for expert, weight in weights.items():
            if len(assignments[expert]) >= capacity:
                dropped += 1
                continue
            assignments[expert].append((position, weight))

    for expert, assigned in enumerate(assignments):
        w_in, w_out = experts[expert]
        for position, weight in assigned:
            expert_output = expert_forward(tokens[position], w_in, w_out)
            for idx in range(width):
                # The router weight multiplies the expert output. This is
                # what makes the routing decision differentiable at all:
                # top-k has no derivative, but the weight of the chosen
                # expert does, so the router learns by having its
                # confidence scaled against how useful the choice was.
                outputs[position][idx] += weight * expert_output[idx]

    return outputs, assignments, router_probabilities, dropped


def load_balancing_loss(assignments, router_probabilities, n_tokens):
    """The auxiliary term: E * sum_e f_e * P_e.

    Minimized when every expert gets an equal share. That is not a goal
    anybody has -- it is a defence against a failure the TASK loss
    causes. Early in training one expert is marginally better by chance,
    the router reduces loss fastest by sending everything there, those
    tokens make it better still, and the rest stagnate for want of
    gradient. The end state is the memory cost of many experts and the
    capacity of one.

    The loss curve looks entirely healthy throughout. Only the routing
    histogram shows it, which is why collapse is usually discovered late.
    """
    n_experts = len(assignments)

    fraction = [len(assigned) / n_tokens for assigned in assignments]
    mean_probability = [
        sum(probabilities[expert] for probabilities in router_probabilities) / n_tokens
        for expert in range(n_experts)
    ]

    # fraction[] is a hard count: no gradient. mean_probability[] is soft
    # and is the only path a gradient can take.
    return n_experts * sum(
        fraction[expert] * mean_probability[expert] for expert in range(n_experts)
    )


def routing_histogram(assignments, n_tokens):
    """The diagnostic that actually matters. Watch this, not the loss."""
    shares = [len(assigned) / n_tokens for assigned in assignments]
    total = sum(shares)
    if total <= 0.0:
        return {"shares": shares, "entropy": 0.0, "max_share": 0.0}

    normalized = [share / total for share in shares]
    entropy = -sum(
        share * math.log(share) for share in normalized if share > 0.0
    )
    return {
        "shares": shares,
        # Uniform routing gives log(E); collapse drives this toward 0.
        "entropy": entropy,
        "max_share": max(shares),
    }

"""Mixture of experts with routing health as a returned value.

The design point: router collapse is invisible in the loss and visible
only in the routing histogram, so the histogram is part of the layer's
output rather than something a caller reassembles. A diagnostic that has
to be reconstructed by hand does not get watched, and this particular
failure is expensive to discover late.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Final, Sequence

import numpy as np
from numpy.typing import NDArray

# Uniform routing over E experts gives entropy log(E). Below this fraction
# of it, the router has effectively collapsed.
COLLAPSE_ENTROPY_RATIO: Final[float] = 0.5


class RouterCollapsed(RuntimeError):
    """Routing entropy has fallen far enough that experts are starving.

    Not a numerical error -- the primary objective PREFERS this solution.
    Sending every token to the currently-best expert lowers the task loss
    fastest, so collapse is an attractor rather than a bug, and the loss
    curve looks healthy the whole way down.
    """


class CapacityTooTight(ValueError):
    """A capacity factor below 1.0 guarantees dropping even under perfect
    routing, which is almost never intended."""


class InvalidTopK(ValueError):
    """More experts requested per token than exist."""


class RouterPrecisionWarning(UserWarning):
    """Router logits computed in reduced precision.

    Nothing in the task loss penalizes router confidence, so logits can
    grow without bound; in bfloat16 that reaches instability quickly. The
    standard fix is float32 for the router alone, which is cheap because
    the router is the smallest weight in the layer.
    """


@dataclass(frozen=True, slots=True)
class MoEConfig:
    """Frozen because these are a training contract, not runtime knobs."""

    n_experts: int
    top_k: int = 2
    capacity_factor: float = 1.25
    balancing_weight: float = 0.01
    router_jitter: float = 0.0

    def __post_init__(self) -> None:
        if self.n_experts < 2:
            raise ValueError(f"need at least 2 experts, got {self.n_experts}")
        if not 1 <= self.top_k <= self.n_experts:
            raise InvalidTopK(
                f"top_k {self.top_k} outside [1, {self.n_experts}]"
            )
        if self.capacity_factor < 1.0:
            raise CapacityTooTight(
                f"capacity_factor {self.capacity_factor} drops tokens even "
                "under perfectly uniform routing"
            )
        if self.balancing_weight < 0.0:
            raise ValueError("balancing_weight must be non-negative")

    def capacity_for(self, n_tokens: int) -> int:
        """Per-expert buffer size. Per BATCH, which is the whole problem."""
        exact = n_tokens * self.top_k / self.n_experts * self.capacity_factor
        return max(1, int(exact))


@dataclass(frozen=True, slots=True)
class RoutingHealth:
    """What a caller must monitor. Returned, not reconstructed."""

    shares: NDArray[np.float64]
    entropy: float
    max_entropy: float
    dropped_fraction: float

    @property
    def entropy_ratio(self) -> float:
        """1.0 is uniform routing; toward 0.0 is collapse."""
        if self.max_entropy <= 0.0:
            return 0.0
        return self.entropy / self.max_entropy

    @property
    def has_collapsed(self) -> bool:
        return self.entropy_ratio < COLLAPSE_ENTROPY_RATIO

    def raise_if_collapsed(self) -> None:
        """Guard clause for a training loop that should stop early."""
        if not self.has_collapsed:
            return
        busiest = int(np.argmax(self.shares))
        raise RouterCollapsed(
            f"entropy ratio {self.entropy_ratio:.3f} below "
            f"{COLLAPSE_ENTROPY_RATIO}; expert {busiest} holds "
            f"{self.shares[busiest]:.1%} of tokens"
        )


@dataclass(slots=True)
class MoEOutput:
    """The layer output alongside everything needed to judge it."""

    activations: NDArray[np.float32]
    balancing_loss: float
    health: RoutingHealth
    assignments: list[NDArray[np.int32]] = field(default_factory=list)


class Router:
    """Top-k routing with the histogram as an output.

    Deliberately separate from the experts: the experts are ordinary MLPs
    and the router is where every interesting failure lives, so they do
    not belong in one class.
    """

    def __init__(self, weights: NDArray[np.float32], config: MoEConfig) -> None:
        if weights.shape[0] != config.n_experts:
            raise ValueError(
                f"router has {weights.shape[0]} rows for "
                f"{config.n_experts} experts"
            )
        # float64 here to make the precision point explicit. In a real
        # bfloat16 model the router alone runs in float32 -- cheap,
        # because it is the smallest weight in the layer.
        self._weights = weights.astype(np.float64)
        self._config = config

    def route(
        self,
        tokens: NDArray[np.float32],
        rng: np.random.Generator | None = None,
    ) -> tuple[NDArray[np.int32], NDArray[np.float64], NDArray[np.float64]]:
        """Return chosen experts, renormalized weights, and full probabilities.

        The full probability matrix is returned because the balancing loss
        needs the soft distribution over ALL experts, not just the chosen
        ones -- the gradient path runs through the unchosen entries too.
        """
        logits = tokens.astype(np.float64) @ self._weights.T

        if self._config.router_jitter > 0.0:
            # No mutable default; a shared Generator across calls would
            # make one caller's draws depend on another's.
            generator = rng if rng is not None else np.random.default_rng()
            logits = logits + generator.normal(
                scale=self._config.router_jitter, size=logits.shape
            )

        shifted = logits - logits.max(axis=-1, keepdims=True)
        exps = np.exp(shifted)
        probabilities = exps / exps.sum(axis=-1, keepdims=True)

        top_k = self._config.top_k
        chosen = np.argpartition(-probabilities, top_k - 1, axis=-1)[:, :top_k]
        chosen_probabilities = np.take_along_axis(probabilities, chosen, axis=-1)

        # Renormalize over the kept experts. Without this the output
        # magnitude depends on how much mass fell outside the top-k,
        # which varies per token and injects noise downstream.
        kept_mass = chosen_probabilities.sum(axis=-1, keepdims=True)
        weights = chosen_probabilities / np.maximum(kept_mass, 1e-12)

        return chosen.astype(np.int32), weights, probabilities


class MoELayer:
    """Experts plus a router, with capacity enforced and reported."""

    def __init__(
        self,
        router: Router,
        expert_w_in: NDArray[np.float32],
        expert_w_out: NDArray[np.float32],
        config: MoEConfig,
    ) -> None:
        if expert_w_in.shape[0] != config.n_experts:
            raise ValueError(
                f"{expert_w_in.shape[0]} expert weight blocks for "
                f"{config.n_experts} experts"
            )
        self._router = router
        self._w_in = expert_w_in
        self._w_out = expert_w_out
        self._config = config

    def forward(
        self,
        tokens: NDArray[np.float32],
        rng: np.random.Generator | None = None,
    ) -> MoEOutput:
        if tokens.ndim != 2:
            raise ValueError(f"expected (tokens, width), got {tokens.shape}")

        n_tokens, width = tokens.shape
        chosen, weights, probabilities = self._router.route(tokens, rng)
        capacity = self._config.capacity_for(n_tokens)

        activations = np.zeros((n_tokens, width), dtype=np.float32)
        assignments: list[NDArray[np.int32]] = []
        dropped = 0

        for expert in range(self._config.n_experts):
            slot_positions, slot_ranks = np.nonzero(chosen == expert)
            if slot_positions.size == 0:
                assignments.append(np.empty(0, dtype=np.int32))
                continue

            # Capacity truncation. Order is arrival order within the
            # batch, so which tokens survive depends on batch composition.
            if slot_positions.size > capacity:
                dropped += int(slot_positions.size - capacity)
                slot_positions = slot_positions[:capacity]
                slot_ranks = slot_ranks[:capacity]

            assignments.append(slot_positions.astype(np.int32))

            hidden = tokens[slot_positions] @ self._w_in[expert].T
            np.maximum(hidden, 0.0, out=hidden)
            expert_output = hidden @ self._w_out[expert].T

            # Router weight scales the expert output; this is the only
            # differentiable path to the routing decision.
            scale = weights[slot_positions, slot_ranks].astype(np.float32)
            activations[slot_positions] += scale[:, None] * expert_output

        health = self._health(assignments, n_tokens, dropped)
        balancing = self._balancing_loss(assignments, probabilities, n_tokens)

        return MoEOutput(
            activations=activations,
            balancing_loss=self._config.balancing_weight * balancing,
            health=health,
            assignments=assignments,
        )

    def _balancing_loss(
        self,
        assignments: Sequence[NDArray[np.int32]],
        probabilities: NDArray[np.float64],
        n_tokens: int,
    ) -> float:
        """E * sum_e f_e * P_e, with the gradient asymmetry preserved.

        f_e is a hard count and carries no gradient. P_e is soft. The
        product penalizes high router confidence on already-overloaded
        experts, which is the only differentiable handle available on a
        discrete decision.
        """
        n_experts = self._config.n_experts
        fraction = np.array(
            [assigned.size / n_tokens for assigned in assignments], dtype=np.float64
        )
        mean_probability = probabilities.mean(axis=0)
        return float(n_experts * np.dot(fraction, mean_probability))

    def _health(
        self,
        assignments: Sequence[NDArray[np.int32]],
        n_tokens: int,
        dropped: int,
    ) -> RoutingHealth:
        shares = np.array(
            [assigned.size / n_tokens for assigned in assignments], dtype=np.float64
        )
        total = shares.sum()
        if total <= 0.0:
            return RoutingHealth(shares, 0.0, 0.0, 1.0)

        normalized = shares / total
        nonzero = normalized[normalized > 0.0]
        entropy = float(-np.sum(nonzero * np.log(nonzero)))

        return RoutingHealth(
            shares=shares,
            entropy=entropy,
            max_entropy=float(np.log(self._config.n_experts)),
            dropped_fraction=dropped / (n_tokens * self._config.top_k),
        )

"""Mixture of experts with tokens sorted into contiguous expert blocks.

The optimization that matters: a naive implementation gathers scattered
rows per expert, which strides across memory and leaves BLAS idle.
Sorting the batch by expert assignment makes each expert's inputs one
CONTIGUOUS slice, so the layer becomes a pass over E contiguous segments
with two GEMMs each. The inverse permutation puts the outputs back.

What this does NOT change: the capacity mechanism, and therefore the
batch-composition dependence of the output. That is inherent to batched
expert computation, not an artefact of any implementation -- the same
input really can produce different outputs depending on what shares its
batch, and no amount of optimization removes it.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Final

import numpy as np
from numpy.typing import NDArray

COLLAPSE_ENTROPY_RATIO: Final[float] = 0.5


@dataclass(frozen=True, slots=True)
class SortedRouting:
    """A permutation that makes each expert's tokens contiguous.

    order          -- token indices sorted by assigned expert
    expert_offsets -- E+1 boundaries, so expert e owns order[off[e]:off[e+1]]
    weights        -- router weight per (permuted) slot
    probabilities  -- full soft distribution, kept for the balancing loss
    """

    order: NDArray[np.int64]
    expert_offsets: NDArray[np.int64]
    weights: NDArray[np.float32]
    probabilities: NDArray[np.float32]
    dropped: int


def route_and_sort(
    tokens: NDArray[np.float32],
    router_weights: NDArray[np.float32],
    n_experts: int,
    top_k: int,
    capacity_factor: float,
) -> SortedRouting:
    """One matmul for routing, one argsort for locality.

    Router logits in float32 even if `+ '`' + `tokens` is lower precision: nothing
    in the task loss penalizes router confidence, so logits grow without
    bound, and in bfloat16 that reaches instability quickly. The router is
    the smallest weight in the layer, so the upcast is nearly free.
    """
    n_tokens = tokens.shape[0]

    logits = tokens.astype(np.float32, copy=False) @ router_weights.T
    logits -= logits.max(axis=-1, keepdims=True)
    np.exp(logits, out=logits)
    logits /= logits.sum(axis=-1, keepdims=True)
    probabilities = logits  # now a distribution; reused below, not recomputed

    # Partial partition rather than a full sort: only the top-k boundary
    # matters, and k is 1 or 2 against E of 8 to 128.
    chosen = np.argpartition(-probabilities, top_k - 1, axis=-1)[:, :top_k]
    chosen_probabilities = np.take_along_axis(probabilities, chosen, axis=-1)
    kept_mass = chosen_probabilities.sum(axis=-1, keepdims=True)
    slot_weights = chosen_probabilities / np.maximum(kept_mass, 1e-12)

    flat_experts = chosen.reshape(-1)
    flat_tokens = np.repeat(np.arange(n_tokens, dtype=np.int64), top_k)
    flat_weights = slot_weights.reshape(-1).astype(np.float32)

    # Stable sort by expert: within an expert, arrival order is preserved,
    # which is what makes capacity truncation deterministic FOR A FIXED
    # BATCH -- and undefined across different batches.
    order = np.argsort(flat_experts, kind="stable")
    sorted_experts = flat_experts[order]

    counts = np.bincount(sorted_experts, minlength=n_experts)
    capacity = max(1, int(n_tokens * top_k / n_experts * capacity_factor))

    keep = np.ones(order.size, dtype=bool)
    offsets = np.zeros(n_experts + 1, dtype=np.int64)
    cursor = 0
    dropped = 0
    for expert in range(n_experts):
        count = int(counts[expert])
        if count > capacity:
            keep[cursor + capacity : cursor + count] = False
            dropped += count - capacity
        cursor += count

    kept_order = order[keep]
    kept_experts = flat_experts[kept_order]
    kept_counts = np.bincount(kept_experts, minlength=n_experts)
    offsets[1:] = np.cumsum(kept_counts)

    return SortedRouting(
        order=flat_tokens[kept_order],
        expert_offsets=offsets,
        weights=flat_weights[kept_order],
        probabilities=probabilities,
        dropped=dropped,
    )


def moe_forward(
    tokens: NDArray[np.float32],
    routing: SortedRouting,
    expert_w_in: NDArray[np.float32],
    expert_w_out: NDArray[np.float32],
    out: NDArray[np.float32] | None = None,
) -> NDArray[np.float32]:
    """Two GEMMs per expert over contiguous slices.

    The permutation is the whole optimization. `+ '`' + `permuted` is a copy of
    the batch in expert order, so `+ '`' + `permuted[lo:hi]` is a contiguous block
    that BLAS can stream. The cost is a full extra copy of the activations,
    which roughly doubles peak memory for this layer -- worth it because
    the alternative is a strided gather inside every expert matmul.
    """
    n_tokens, width = tokens.shape
    n_experts = expert_w_in.shape[0]
    n_slots = routing.order.size

    permuted = np.ascontiguousarray(tokens[routing.order])
    # Preallocated and written in place, so no per-expert intermediate.
    permuted_out = np.empty((n_slots, width), dtype=np.float32)

    for expert in range(n_experts):
        lo = int(routing.expert_offsets[expert])
        hi = int(routing.expert_offsets[expert + 1])
        if hi <= lo:
            continue

        block = permuted[lo:hi]
        hidden = block @ expert_w_in[expert].T
        np.maximum(hidden, 0.0, out=hidden)
        np.matmul(hidden, expert_w_out[expert].T, out=permuted_out[lo:hi])

    # Router weight scales the output before scatter-add. Still the only
    # differentiable path to the routing decision.
    permuted_out *= routing.weights[:, None]

    activations = out if out is not None else np.zeros(
        (n_tokens, width), dtype=np.float32
    )
    if out is not None:
        activations.fill(0.0)

    # np.add.at is the correct scatter-add: a token routed to k experts
    # appears k times in `+ '`' + `order` and its contributions must sum.
    np.add.at(activations, routing.order, permuted_out)
    return activations


def balancing_loss(
    routing: SortedRouting, n_tokens: int, n_experts: int, weight: float
) -> float:
    """E * sum_e f_e * P_e, reusing the probability matrix already computed.

    Recomputing the softmax here would double the router cost for no
    reason; the forward pass already materialized it.
    """
    counts = np.diff(routing.expert_offsets).astype(np.float32)
    fraction = counts / np.float32(n_tokens)
    mean_probability = routing.probabilities.mean(axis=0)
    return float(weight * n_experts * np.dot(fraction, mean_probability))


def routing_health(
    routing: SortedRouting, n_tokens: int, n_experts: int, top_k: int
) -> dict[str, float]:
    """The metric to watch. Cheap enough to compute every step.

    A model whose task loss looks perfect while entropy_ratio sits near
    zero has the memory footprint of E experts and the capacity of one.
    No other signal reports that.
    """
    counts = np.diff(routing.expert_offsets).astype(np.float64)
    total = counts.sum()
    if total <= 0.0:
        return {"entropy_ratio": 0.0, "max_share": 0.0, "dropped_fraction": 1.0}

    shares = counts / total
    nonzero = shares[shares > 0.0]
    entropy = float(-np.sum(nonzero * np.log(nonzero)))
    max_entropy = float(np.log(n_experts))

    return {
        "entropy_ratio": entropy / max_entropy if max_entropy > 0.0 else 0.0,
        "max_share": float(shares.max()),
        "dropped_fraction": routing.dropped / max(1, n_tokens * top_k),
        "collapsed": float(
            (entropy / max_entropy if max_entropy > 0.0 else 0.0)
            < COLLAPSE_ENTROPY_RATIO
        ),
    }