dsv3: fix pure-JAX MoE einsum (wi_0 / wi_1 axis labels)

[ultrons]

Summary

Fixes a latent correctness bug in DSv3's pure-JAX MoE backends
(--moe_backend=jax and --moe_backend=ragged_dot). Both paths raise
TypeError immediately due to misaligned einsum/ragged_dot axis labels
against the actual weight storage layout.

TypeError: dot_general requires contracting dimensions to have the
same shape, got (7168,) and (2048,).

Root cause

init_params (model.py around line 2929) stores wi_0 / wi_1 with
shape (cfg.E, cfg.D_moe, cfg.D) — D_moe in the middle axis, D as
the last. Production paths (gmm_v2 / megablox / kernel-agent)
consume this layout correctly.

Two pure-JAX consumers labeled the axes wrong:

• expert_mlp_jax line 889-890 (EP=1 branch): einsum string
  "td,edm->tem" assumes (E, D, D_moe) — places d (contracting)
  on the middle axis, but the middle axis is actually D_moe.
• _rd_block (expert_mlp_ragged_dot no-mesh path): ragged_dot
  expects rhs (G, K, N), but wi_0/wi_1 arrive as (E, D_moe, D)
  = (G, N, K) — same contracting mismatch.

The down-projection wo is unaffected: its contracting dim is D_moe,
which already aligns with the middle axis of storage.

Fix

# expert_mlp_jax (EP=1)
- gate_all = jax.nn.silu(jnp.einsum("td,edm->tem", flat_x, wi_0))
- up_all   = jnp.einsum("td,edm->tem", flat_x, wi_1)
+ gate_all = jax.nn.silu(jnp.einsum("td,emd->tem", flat_x, wi_0))
+ up_all   = jnp.einsum("td,emd->tem", flat_x, wi_1)

# _rd_block (ragged_dot path)
- gate = jax.nn.silu(jax.lax.ragged_dot(sorted_x, wi_0, group_sizes))
- up   = jax.lax.ragged_dot(sorted_x, wi_1, group_sizes)
+ gate = jax.nn.silu(jax.lax.ragged_dot(sorted_x, wi_0.swapaxes(-1, -2), group_sizes))
+ up   = jax.lax.ragged_dot(sorted_x, wi_1.swapaxes(-1, -2), group_sizes)

Plus a comment block explaining the axis convention so future drifts
get caught at code-review time.

Why this stayed latent

Production training (v304, v3xx) uses --moe_backend gmm_v2 or
--moe_backend megablox — never --moe_backend jax. No integration
test exercises the broken paths, so a weight-layout refactor (storing
D_moe in the middle, presumably for FSDP / tile-layout reasons) broke
the pure-JAX consumers silently.

Verification

JAX 0.10 / CPU.

Lowering (was broken; now works):

• DSv3 mini (L=2, full D=7168) forward lowers cleanly under
  moe_backend=jax: 3198 HLO lines, 26 dot ops, no errors.
• Same under moe_backend=ragged_dot: 3721 HLO lines, 25 dots.

Numerical (cross-check between the two pure-JAX paths on a nano
config, D=64, D_moe=32, E=4):

• fp32: max rel diff 2.0e-5, mean rel diff 3.0e-7 — bit-level
  agreement.
• bf16: max 8% / mean 0.2% — expected summation-order noise.

Why this matters now

Surfaced while building perfsim's HLO cross-compile validation harness
(perfsim PR ultrons/perfsim#88, task #110). perfsim needs the pure-JAX
path to lower on CPU since Pallas backends are TPU-only. Without this
fix the no-hardware HLO validation track is blocked.

Base branch

I targeted autoperf/dsv3_train_full since main doesn't carry the
dsv3 module. If a different base is preferred (e.g. a stable release
branch), happy to re-target.

Test plan

• DSv3 mini forward lowers under moe_backend=jax
• DSv3 mini forward lowers under moe_backend=ragged_dot
• Numerical cross-check (fp32 bit-level, bf16 noise-level)
• Integration test (suggest adding one — see commit message for
  a minimal pytest skeleton)