[faresobeid]

Adds an opt-in filter to trainer-side router replay, following Cursor's Composer 2 report (arXiv:2603.24477, sec. 6.2). Recommendation: keep it off. In GPU tests it did not beat plain replay (numbers below). It is offered as a default-off option so the result is reproducible and the option is available for models where it might help.

What changes

• New trainer.router_replay_min_score_ratio: float | None = None (0 < α ≤ 1). A validator requires trainer.enable_router_replay.
• filter_replayed_experts (trainer/models/layers/moe.py): per token, a replayed expert is kept only if its router score is ≥ α × the lowest score in the router's own top-k (selected with the selection bias, like normal routing). Dropped slots take the router's own top-k experts, best first, skipping kept ones, so each token keeps top_k distinct experts. Vectorized, under no_grad; gating weights are still recomputed by the trainer, so the router stays trainable.
• Scores are post-score-function gating scores (softmax/sigmoid/sqrtsoftplus, before route norm/scale, without selection bias); full-softmax probabilities for topk_softmax (gpt-oss).
• Lives in TokenChoiceTopKRouter.forward, so it covers every custom MoE model. apply_router_replay_filter (called from rl/train.py) sets α on all MoE routers except DeepSeek-V4 hash routers, whose routed_experts are a token-id selection that must not be filtered.
• Unset: code path unchanged, output bit-identical to main.
• Docs: a paragraph in docs/inference.md › Router Replay that explains α and says it did not beat plain replay.

No new metric: a replaced-fraction metric would need another per-layer buffer and custom-op mutation. Breaking changes: none.

GPU results (Qwen3-30B-A3B, 1 node, reverse-text)

4 trainer GPUs (EP 4) + 4 vLLM GPUs, batch 256, seed 42, 30 steps per arm, means over steps 5-29. Mismatch KL is token-level over completion tokens. Off-policy = max_staleness = 4. α sweep 0.3/0.5/0.7/0.9 was run for FP8 and off-policy; α = 0.5 shown.

regime	arm	mismatch KL mean	p99	replaced slots	reward
bf16 on-policy	replay	0.0053	0.071	0	0.710
	α = 0.5	0.0049	0.076	0.37%	0.710
FP8 inference	replay	0.067	1.23	0	0.719
	α = 0.5	0.068	1.21	0.37%	0.639
off-policy bf16	replay	0.065	1.17	0	0.625
	α = 0.5	0.140	2.86	0.37%	0.617

• Plain replay does the work: against no replay it cuts mean mismatch KL 2-7x and p99 2-9x in all three regimes.
• No α in 0.3-0.9 lowers mismatch or raises reward compared with plain replay. Off-policy, every α raises mismatch (mean +47% at α = 0.3, about 2x at α ≥ 0.5).
• The filter swaps a fixed share of near-tie experts (0.37% at α = 0.5) in every regime. That share does not grow with FP8 or staleness, so it does not target the extra mismatch.
• Cost: trainer forward+backward +0-6%. Step time is within noise (runs are inference-bound).
• Single seed per arm. Not run: multi-node GLM-4.5-Air, CP, gpt-oss.

Validation

• pytest tests/unit/train/models/{test_moe,test_deepseek_v4_cpu,test_qwen3_moe,test_glm4_moe,test_nemotron_h,test_moe_conversions}.py tests/unit/test_configs.py tests/unit/orchestrator -m "not gpu" on current main: 418 passed. New test_filter_replayed_experts checks kept/replaced ids, dtype, α = 1 → router's own top-k, and that the gradient reaches the gate.
• torch.compile(fullgraph=True) CPU smoke: matches eager exactly, ids distinct.
• ruff check ., ruff format --check ., uv lock --check.

---

Note

Medium Risk
When enabled, it changes MoE expert selection during RL forward passes on the replay path; default-off, but incorrect α could worsen trainer/inference mismatch for MoE training.

Overview
Adds an opt-in Composer 2–style plausibility filter for MoE router replay on the trainer. New trainer.router_replay_min_score_ratio (α, default off, requires enable_router_replay) drops replayed experts whose gating score is below α times the weakest score in the trainer router’s own top-k, and fills those slots from the trainer’s top-k so each token still routes to top_k distinct experts. Gating weights are still taken from the trainer for the final expert set.

Implementation: vectorized filter_replayed_experts runs in TokenChoiceTopKRouter.forward when replay ids are present (with a topk_softmax full-softmax score path); apply_router_replay_filter enables it on learned MoE routers at RL startup and skips DeepseekV4HashRouter. Docs in docs/inference.md describe α and note it should stay off unless your runs benefit. Unit tests cover replacement logic, selection bias, α=1 behavior, and router gradients.

^{Reviewed by Cursor Bugbot for commit 90ee189. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 90ee189. Configure here.}

[cursor]

Wrong own top-k after softmax

Low Severity

For topk_softmax routers, the filter adds selection_bias to full-softmax probabilities when choosing the trainer’s own top-k, but the normal routing path adds that bias to raw logits. Those rankings are not the same, so the keep threshold and replacements can diverge from the router’s real top-k whenever both topk_softmax and a selection bias are active.

Additional Locations (1)

• src/prime_rl/trainer/models/layers/moe.py#L65-L67

 

^{Reviewed by Cursor Bugbot for commit 90ee189. Configure here.}