dockyard_rl.models.jax.moe.experts

Grouped-GEMM routed-expert compute for MoE (SwiGLU), Flax NNX.

JAX mirror of models/dtensor/moe/experts.py. The routed-expert FFN runs as a single grouped matmul over the local experts via jax.lax.ragged_dot (the JAX equivalent of torch _grouped_mm) instead of a per-expert loop. Unlike the torch kernel (CUDA-only, HV-8), ragged_dot lowers on CPU, so this is CPU-parity-validatable against a per-expert reference.

Expert weights are 3D (num_experts, *, *) and shard Shard(0) over the EP mesh (J8b); under EP, num_experts here is the LOCAL expert count. __call__ orchestrates dispatch -> grouped compute -> combine via an injected dispatcher.

Shape suffixes: D=model dim, F=hidden dim, E=experts, R=routed tokens.

Module Contents

Classes

GroupedExperts

SwiGLU routed experts computed as a grouped matmul.

Data

Array

API

dockyard_rl.models.jax.moe.experts.Array

None

class dockyard_rl.models.jax.moe.experts.GroupedExperts(dim: int, hidden_dim: int, num_experts: int, dispatcher: dockyard_rl.models.jax.moe.dispatch.LocalTokenDispatcher, *, rngs: flax.nnx.Rngs, param_dtype: jax.numpy.dtype = jnp.float32, compute_dtype: Optional[jax.numpy.dtype] = None)

Bases: flax.nnx.Module

SwiGLU routed experts computed as a grouped matmul.

Args: dim: Model dimension (D). hidden_dim: FFN intermediate dimension (F). num_experts: Local expert count (E / ep under expert parallelism). dispatcher: Token dispatcher providing dispatch/combine. compute_dtype: Optional matmul dtype (the GPU path uses bfloat16 for the grouped GEMM); None keeps the input dtype (fp32 CPU parity).

Initialization