Add production-safe local Grug MoE backends

[dlwh]

Add local Grug MoE backends for scatter, Sonic-style XLA dispatch, interleaved W13 layout, and a production-safe custom-VJP down/gather boundary. MoEExpertMlp owns expert weight initialization and W13 layout so model code can select the backend without layout plumbing.

Part of #5328

[claude]

Claude finished @dlwh's task in 6m 22s —— View job

---

Code review

Reviewed for bugs and CLAUDE.md/AGENTS.md compliance.

One issue flagged inline:

• _gather_sum_reference in quack_moe.py:56-68 duplicates the identical helper in grug_moe.py:217-229. Suggested importing it instead — see inline comment for the suggested fix.

Other things considered and not flagged:

• MoeLocalImplementation typed as Literal[...] rather than StrEnum — follows the established pattern of the pre-existing MoeImplementation in the same module, so flagging would be inconsistent with current code.
• Boolean dispatch flags (interleaved_w13, quack_down, split_moe_w13_output(interleaved=...)) — these are localized dispatch keys, not the kind of compound state machine the AGENTS.md rule targets.
• Optional-dep guard in grug_moe.py:131-156 only catches ModuleNotFoundError; a transitive ImportError of a different subclass (or a partial install where find_spec("cuda") succeeds but cuda.bindings.driver import fails) would bypass the friendly message and bypass the test's skip check at test_grugformer_moe.py:257-279. Low-risk in the typical install path; mentioning here for awareness rather than as a blocking issue.
• custom_vjp residuals/cotangents arity in quack_interleaved_down_gather_sum, the interleaved/concat W13 layout symmetry, and the dispatch_positions construction in _prepare_moe_dispatch_indices_with_assignment_ids all look consistent across forward/backward.
  · Branch: codex/moe-quack-down-backend

[dlwh]

🤖 Specification for the updated >500 LOC backend PR.

Problem: moe_mlp only exposed the EP ring/ragged implementations on main, and model code had to own expert MLP weight layout. The earlier PR version also tried to make the fastest path depend on CuTe/QuACK kernels, but our GH200 lowering work showed the d-gated ragged TMA path is not production-safe yet. We still need an opt-in local path that captures the useful W13 interleaving and custom-VJP boundary without CUDA-only optional imports.

Approach: lib/levanter/src/levanter/grug/grug_moe.py now accepts ring, ragged_all_to_all, scatter, sonic_xla, sonic_xla_interleaved_w13, and sonic_xla_interleaved_w13_custom_vjp_down. EP choices stay EP-only; local choices route through separate straight-line functions. MoEExpertMlp owns expert weight initialization, sharding, backend selection, and W13 concat/interleaved layout. experiments/grug/moe/model.py delegates expert MLP execution to MoEExpertMlp. The custom-VJP backend lives in lib/levanter/src/levanter/grug/custom_vjp_moe.py and uses JAX ragged_dot plus an explicit VJP for the down/gather boundary; it adds no new package dependencies.

Key code:

@jax.custom_vjp
def custom_vjp_interleaved_down_gather_sum(
    w13_out_interleaved,
    combine_weights,
    w_down,
    token_ids_sort,
    sorted_assignment_ids,
    dispatch_positions,
    group_sizes,
):
    out, _ = _custom_vjp_interleaved_down_gather_sum_forward(
        w13_out_interleaved,
        combine_weights,
        w_down,
        token_ids_sort,
        sorted_assignment_ids,
        dispatch_positions,
        group_sizes,
    )
    return out

The backward saves the dispatch output from forward, replays only the SwiGLU activation pullback and down-projection pullback, and computes the combine-weight gradient as sum(dout_sorted * dispatch_output) before unsorting by assignment id.

Tests: lib/levanter/tests/grug/test_grugformer_moe.py covers scatter vs Sonic-style XLA value/gradient parity, concat vs interleaved W13 value/gradient parity, MoEExpertMlp layout-hiding parity, custom-VJP down/gather value/gradient parity against the interleaved XLA path, and EP abstract-mesh lowering. Local validation ran ./infra/pre-commit.py --all-files --fix and uv run --package marin --group dev pytest -m 'not slow' lib/levanter/tests/grug/test_grugformer_moe.py -q.

[claude]

Parallel implementation of _gather_sum_reference in grug_moe.py — the body is identical (only the type annotations differ: Float[Array, ...] vs plain jax.Array). Since quack_moe is in the same package and is itself imported lazily from grug_moe, you can import the helper directly without circular-import concerns.

Per root AGENTS.md: "If a suitable implementation exists, use it. Do not create parallel implementations."

Suggested change

	def _gather_sum_reference(
	dispatch_output: jax.Array,
	dispatch_positions: jax.Array,
	combine_weights: jax.Array,
	) -> jax.Array:
	acc = jnp.zeros((dispatch_positions.shape[0], dispatch_output.shape[1]), dtype=dispatch_output.dtype)
	weights = combine_weights.astype(dispatch_output.dtype)
	for topk_index in range(dispatch_positions.shape[1]):
	gathered = jnp.take(dispatch_output, dispatch_positions[:, topk_index], axis=0)
	acc = (acc + (gathered * weights[:, topk_index, None]).astype(dispatch_output.dtype)).astype(
	dispatch_output.dtype
	)
	return acc
	from levanter.grug.grug_moe import _gather_sum_reference