[DRAFT] Paged Stashing for ring-of-experts MoE Activation Memory

[abhinavgoel95]

Problem

When use_ring_of_experts=True with expert parallelism EP=4, the ring-all-gather inflates the token buffer on each device to batch × EP × seq × top_k tokens. For DeepSeek V3 671B with batch=2, seq=4096, top_k=8, EP=4, this is 262,144 tokens/device.

The moe_mlpwo GMM output has shape (262144, 7168) = 3.5 GB per layer. With scan_layers=True over 60 MoE layers and host offloading, this totals ~210 GB of host memory — exceeding typical limits (180 GB).

Solution: Paged Stashing

Inspired by Megatron-LM PR #2690, this PR introduces a paged stash mechanism that decouples the compute buffer (static worst-case shape, required by XLA) from the storage buffer (actual tokens only).

Key idea

• A single shared static buffer is allocated once for all MoE layers: total_capacity = num_moe_layers × expected_tokens_per_layer + max_chunk_slack
• Each layer writes only its actual routed tokens into this buffer at a tracked offset using lax.dynamic_update_slice (static size, dynamic start index — XLA-compatible)
• On backward pass, each layer reads back its compact slice and expands to the full compute shape
• Expected actual tokens ≈ worst_case / EP = 262144 / 4 = 65536, so the shared buffer is ~4× smaller than per-layer worst-case storage

Memory comparison (DeepSeek V3 671B, EP=4, 60 MoE layers)

Approach	Storage for moe_mlpwo	Notes
Baseline (offload per layer)	~210 GB	262k × 7168 × 60 layers
Paged stash (this PR)	~52 GB	65k × 7168 × 60 layers

Implementation

New file: src/MaxText/layers/paged_stash.py

• make_stash_fns(max_chunk, hidden) — returns (stash_fn, restore_fn) pair with jax.custom_vjp so gradients flow correctly through the compact/expand operations
• stash_buffer_size(num_moe_layers, expected_per_layer, max_chunk) — buffer sizing helper
• expected_tokens_per_layer(batch, ep, seq, top_k) — computes expected load (= worst_case / EP)

Modified: src/MaxText/layers/moe.py

When config.ring_paged_stash=True, replaces the checkpoint_name("moe_mlpwo") call with stash_fn/restore_fn calls that pack actual tokens into the shared buffer.

Modified: src/maxtext/configs/types.py

Adds two new config fields:

• ring_paged_stash: bool = False
• ring_paged_stash_safety_margin: float = 1.5

TODOs / Open Questions

• Decoder scan carry wiring: The shared stash_buf and write_ptr need to be threaded through the lax.scan carry in decoders.py / deepseek.py. Currently the per-layer integration is sketched but the cross-layer buffer threading is not yet wired up. Looking for guidance on the best pattern here.
• restore_fn backward: The gradient accumulation in restore_fn_bwd needs review — specifically the scatter from compact → full shape.
• Tests: Unit tests for stash/restore round-trip and gradient correctness.
• moe_mlpwi_0/moe_mlpwi_1: Same technique could apply to wi outputs if they are offloaded.