Answer

Implementation and tests are in place and passing. Here is the verified solution.

Sparse MoE Layer in NumPy — Token Drop, Exact Dispatch/Combine, Aux Loss & Backward

1. Root-cause analysis (the traps this task targets)

There was no pre-existing file; the task is to produce a correct layer. The "root causes" are the subtle correctness traps that make naive implementations fail exactness/gradient checks:

1. Nondeterministic tie-breaking — np.argsort(-p) is not stable. Equal router scores must resolve by ascending expert index. Fix: np.argsort(-p, axis=-1, kind="stable") on the negated probabilities.
2. Wrong capacity — must be ceil(T*k/E * capacity_factor), and it is invariant to expert-parallel world_size (batch and expert set shard by the same factor).
3. Drop not position-prioritised / not counted — greedy in token-position order, gate-descending within a token; every rejection increments that expert's overflow.
4. Dropped assignment ≠ dropped token — dropped contributions must leave the residual untouched: y = x + Σ_kept gate·Expert(x).
5. Dispatch/combine lose top-k structure — exact form is dispatch_mask:(T,E,C) bool, combine_weights:(T,E,C) float, both zero on dropped slots.
6. Aux loss from the wrong quantity — must use router probability mass summed per expert: P_e=(1/T)Σ_t p[t,e], f_e=count/(T·k), aux=E·Σ f_e·P_e, pre-drop.
7. Backward misses the softmax→top-k→renormalise chain — needs quotient rule + softmax Jacobian; discrete drops are constants (straight-through).

2. Exact fix — implementation

Full code saved at ~/moe.py (reproduced in SOLUTION.md). Core:

def top_k_gating(router_logits, k):
    E = router_logits.shape[1]; k = int(min(k, E))
    p = softmax(router_logits, axis=-1)
    order = np.argsort(-p, axis=-1, kind="stable")   # ties -> ascending index
    idx = order[:, :k].astype(np.int64)
    selected = np.take_along_axis(p, idx, axis=-1)
    gate = selected / selected.sum(-1, keepdims=True)
    return p, idx, gate

def position_priority_dispatch(idx, gate, num_experts, capacity):
    T, k = idx.shape
    used = np.zeros(num_experts, dtype=np.int64)
    slot = np.full((T, k), -1, dtype=np.int64)
    kept = np.zeros((T, k), dtype=bool)
    for t in range(T):                      # earlier positions have priority
        for j in range(k):                  # higher gate first
            e = int(idx[t, j])
            if used[e] < capacity:
                slot[t, j] = used[e]; used[e] += 1; kept[t, j] = True
    attempts = np.zeros(num_experts, dtype=np.int64)
    np.add.at(attempts, idx.ravel(), 1)
    return slot, kept, attempts - used, (~kept).sum(axis=-1)

moe_forward builds M = combine.sum(axis=2), computes xe = einsum("td,edf->tef", x, W), out = einsum("tef,te->tf", xe, M), and returns y = x + out, plus dispatch/combine tensors, overflow, drops_per_token, and the aux loss. moe_backward uses:

dM      = np.einsum("tf,tef->te", dy, xe)
dW      = np.einsum("te,ta,tb->eab", M, x, dy)
dx      = (dy if residual else 0) + np.einsum("te,tb,eab->ta", M, dy, W)
dcombine= np.where(dispatch_mask, np.einsum("tf,tef->te", dy, xe)[:, :, None], 0.0)
# then gate_to_logits_grad: quotient rule + softmax Jacobian

3. Verification

Run python3 ~/test_moe.py. An independent dense reference loops over all experts and is used for forward comparison and central finite differences of the backward.

ok  tie-break ascending expert index
ok  capacity / overflow / position priority
ok  sparse forward == dense reference
ok  backward == finite difference (max err: dx=3.21e-09, dlogits=3.60e-09, dW=4.12e-09)
ok  aux loss & gradient (max err=2.51e-10, aux_uniform=1.0000)
ok  dropped tokens pass through residual unmodified
ok  capacity invariant to expert-parallel world size

ALL TESTS PASSED

• Backward vs. dense reference max error ≈4e-9 < 1e-6 ✅
• Equal router scores resolve to ascending expert indices ✅
• Edge checks: k==E, capacity==0 (output exactly x), single token/expert, use_residual=False (err <1e-9) ✅

4. Usage

y, cache, diag = moe_forward(x, logits, W, k=2, capacity_factor=1.0)
diag["dispatch_mask"], diag["combine_weights"]   # (T,E,C)
diag["overflow"], diag["drops_per_token"], diag["aux_loss"]
grads = moe_backward(dy, cache)   # dx, dlogits, dW, dgate, dcombine

Files created: ~/moe.py, ~/test_moe.py, ~/SOLUTION.md.