feat(moe): R3 rollout routing replay — slot-indexed expert-routing pool (scaffold)

[qywu]

What & why

R3 (arXiv:2510.11370) stabilizes RL training of MoE models: the training forward replays the exact expert selection the rollout made, instead of letting the inference and training routers pick experts independently. Independent routers + tiny numerical drift flip a token's top-k choice → train/inference KL blows up. R3 records per-token, per-layer top-k expert ids at rollout and replays them in training (KL cut ~50%; integrates with slime/veRL).

Design — store routing in the KV memory pool

The distinctive choice vs. vllm-ascend's transient RoutedExpertsCapturer: persist routing in a pool indexed by the KV slot, alongside KV.

On a prefix-cache hit the shared prefix tokens are dropped from the forward entirely — their MoE layers never run (out_cache_loc covers only new tokens). A transient capturer therefore returns nothing for the prefix; you'd have to force a re-forward. A slot-indexed pool makes routing follow the KV: a prefix hit that reuses slots also reuses the routing captured when those slots were first written — zero recompute, guaranteed consistent with the reused KV. Cost ≈ num_moe_layers * top_k * 4 B/token (~1-5% of KV).

This builds on infra the repo already has (moe/distribution_recorder.py already captures per-token per-layer topk_ids for offline analysis + tracks the current layer); R3 adds slot-indexed persistence and a return path.

What this PR lands (tested engine core)

runtime/cache/routed_experts_pool.py:

• RoutedExpertsPool — [size+1, num_moe_layers, top_k] int32, slot-indexed on dim 0 (row 0 reserved, mirrors the KV pools). store_layer() scatters by slot like set_kv_buffer; gather() is the prefix-hit retrieval.
• RoutedExpertsCapturer — per-forward begin_forward/capture/commit; no-op when inactive; skips + counts TP/EP-misaligned rows rather than corrupt the pool.
• process-global accessor (mirrors the distribution recorder).
• GenerateReqInput.return_routed_experts opt-in flag (propagated via __getitem__).

18 CPU unit tests: store/gather, reserved slot, sizing, the prefix-hit retrieval property, capture/commit, inactive no-op, TP-misalignment skip, detach-on-capture, flag propagation.

Remaining wiring (needs a running model) — documented follow-ups

In docs/guides/routing-replay-r3.md: allocate the pool via a ServerArgs flag; capture hook at topk.py::select_experts (beside the existing on_select_experts); thread out_cache_loc into the MoE block; TP/EP all-gather alignment; BYPASSED fused-routing capture; prefix-hit retrieval + return path across the gRPC/gateway boundary (proto + servicer + Rust gateway, same boundary as the return_token_ids follow-up).

Docs

Adds docs/guides/routing-replay-r3.md + sidebar entry.

Local note: this venv can't import the full kernel stack on main (unrelated triton_kernels/flashinfer version drift). The pool/capturer tests are kernel-free and run directly (18 passed); the one lazy-imported flag test also passes.

🤖 Draft opened for design review.

[qywu]

Update: forward wiring landed (commit 239ac75), guarded behind --enable-routing-replay (off by default → every hook is a no-op, default behavior byte-identical).

• server_args: --enable-routing-replay
• event_loop: build + install the pool right after the KV pool (sized to max_total_num_tokens; None for dense models)
• model_runner.forward: begin_forward(out_cache_loc) / commit() around model.forward
• topk.select_experts: capture_in_order(topk_ids) beside the existing on_select_experts; layer index by per-forward invocation order (no per-model edits); BYPASSED fused routing naturally skipped
• capture_in_order skips under is_current_stream_capturing() so it can never corrupt a CUDA graph

Now 25 CPU tests; all edited modules verified to import cleanly (no circular imports).

Not yet GPU-validated end to end (this venv can't run a model). Remaining, documented in the guide: CUDA-graph decode capture, TP/EP row alignment (safe-skipped today; TP=1 works), and the prefix-hit retrieval + return path across the gRPC/gateway boundary.