feat(trainer): Composer 2-style plausibility filter for router replay

docs/inference.md

@@ -330,6 +330,8 @@ This however is not free, it adds a significant overhead to the HTTP requests as
 
 Router replay works with `inference.kv_cache_offload` (native and Mooncake). vLLM keeps routing rows in a local store keyed by block hash, and no connector moves them with the KV, so a loaded block has no rows on this instance. prime-rl therefore limits an offload load to the prompt tokens the client already has routing for: the blocks before `routed_experts_prompt_start`. Multi-turn rollouts keep their offload hits, because the reloaded earlier turns sit before that position. Offloaded prefixes past it (first turns, prompts shared across a group or across nodes) are recomputed. Loaded blocks then serve local GPU prefix hits like any other block, so a local hit is cut at the first loaded block that a request needs routing from (for example a first turn that hits a prefix reloaded by another rollout's later turn), and the rest is recomputed. Under P/D, decode-side loads are not limited: decode requests emit routing only from their first forward.
 
+Plain replay can force an expert that the trainer's router scores very low, so its recomputed gating weight is far from the one inference used. `trainer.router_replay_min_score_ratio` (α, off by default) filters these out, following Cursor's Composer 2 report ([arXiv:2603.24477](https://arxiv.org/abs/2603.24477), section 6.2). For each token and MoE layer, a replayed expert is kept only if its router score is at least α times the lowest score among the trainer router's own top-k. Dropped experts are replaced by the router's own top-k picks, best first, so each token still uses `top_k` distinct experts. Scores are the router's gating scores after the score function (softmax/sigmoid probabilities, before route normalization and without the load-balancing selection bias; full-softmax probabilities for `topk_softmax` routers). The router stays trainable because gating weights are still computed by the trainer. Smaller α keeps more of the replay; for routers without a selection bias, α = 1 falls back to the trainer's own routing. Cursor does not publish their threshold. On Qwen3-30B-A3B (bf16, FP8 inference, and off-policy rollouts), no α in 0.3–0.9 lowered mismatch KL or raised reward compared with plain replay, so leave it off unless your own run shows a gain.
+
 ### Sampling Replay
 
 Truncated sampling (`top_p < 1`, `top_k`) renormalizes the sampling distribution over a sampling mask of surviving token ids. The rollout logprobs reflect that (`logprobs_mode = "processed_logprobs"`), so the trainer must renormalize over the same mask — otherwise every importance ratio is biased and training collapses (DeepSeek V3.2's "Keep Sampling Mask", [arXiv:2512.02556](https://arxiv.org/abs/2512.02556) §3.1; Cognition's [SWE-1.7 post](https://cognition.com/blog/swe-1-7)). prime-rl handles this automatically: vLLM records the sampling mask at sampling time (`--return-sampling-mask`, native since vLLM 0.28) and the trainer renormalizes its logprobs over it.

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -715,6 +715,15 @@ class TrainerConfig(BaseTrainerConfig):
     enable_router_replay: bool = False
     """Return routed experts in the batch so the trainer can replay routing. Requires ``enable_return_routed_experts=true`` on the vLLM server (or ``--enable-return-routed-experts``) and is only supported for custom models."""
 
+    router_replay_min_score_ratio: float | None = Field(None, gt=0, le=1)
+    """Composer 2-style router replay filter. When set, a replayed expert is kept only if its router score is at least this fraction of the lowest score among the trainer router's own top-k; dropped experts are replaced by the router's own top-k picks. Requires ``enable_router_replay``."""
+
+    @model_validator(mode="after")
+    def router_replay_filter_requires_router_replay(self):
+        if self.router_replay_min_score_ratio is not None and not self.enable_router_replay:
+            raise ValueError("router_replay_min_score_ratio requires enable_router_replay")
+        return self
+
     @model_validator(mode="after")
     def validate_lora_broadcast(self):
         if self.model.lora is not None and self.weight_broadcast.type in ("nccl", "nixl"):

src/prime_rl/trainer/model.py

@@ -50,6 +50,7 @@
     supports_custom_impl,
 )
 from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Indexer
+from prime_rl.trainer.models.deepseek_v4.moe import DeepseekV4HashRouter
 from prime_rl.trainer.models.fusions import (
     apply_model_fusions,
     get_fsdp_shard_placement_fn,
@@ -201,6 +202,23 @@ def apply_force_balanced_routing(model: nn.Module) -> None:
     )
 
 
+def apply_router_replay_filter(model: nn.Module, min_score_ratio: float) -> None:
+    """Filter implausible replayed experts in every learned MoE router (see `filter_replayed_experts`)."""
+    num_routers = 0
+    for layer in get_language_model(model).layers:
+        mlp = layer.mlp if hasattr(layer, "mlp") else layer.feed_forward if hasattr(layer, "feed_forward") else None
+        # Hash layers pass their token-id selection in as `routed_experts`; it must not be filtered.
+        if isinstance(mlp, MoE) and not isinstance(mlp.router, DeepseekV4HashRouter):
+            mlp.router.replay_min_score_ratio = min_score_ratio
+            num_routers += 1
+
+    if num_routers == 0:
+        raise ValueError("No MoE routers found for router_replay_min_score_ratio. Is this a MoE model?")
+    get_logger().info(
+        f"Filtering replayed experts below {min_score_ratio} x the router's own top-k on {num_routers} MoE layers"
+    )
+
+
 def is_tt_moe_model(model: nn.Module) -> bool:
     config = getattr(model.config, "text_config", model.config)
     return hasattr(config, "num_experts") or hasattr(config, "n_routed_experts")

src/prime_rl/trainer/models/layers/moe.py

@@ -49,6 +49,32 @@ def _record_moe_routing_statistics_fake(
     return None
 
 
+def filter_replayed_experts(
+    scores: torch.Tensor,
+    routed_experts: torch.Tensor,
+    min_score_ratio: float,
+    selection_bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Drop implausible replayed experts, as in Composer 2's router replay (arXiv:2603.24477, sec. 6.2).
+
+    A replayed expert is kept if its score is at least ``min_score_ratio`` times the lowest score among
+    the router's own top-k. Dropped slots take the router's own top-k experts (best first) that are not
+    already kept, so the result has ``top_k`` distinct experts whenever the replayed ones are distinct.
+    """
+    with torch.no_grad():
+        top_k = routed_experts.shape[-1]
+        selection_scores = scores if selection_bias is None else scores + selection_bias
+        own = torch.topk(selection_scores, k=top_k, dim=1, sorted=True).indices
+        threshold = min_score_ratio * scores.gather(1, own).amin(dim=1, keepdim=True)
+        keep = scores.gather(1, routed_experts) >= threshold
+        kept_ids = routed_experts.masked_fill(~keep, -1)
+        taken = (own.unsqueeze(-1) == kept_ids.unsqueeze(1)).any(dim=-1)
+        candidates = own.gather(1, torch.argsort(taken.to(torch.int8), dim=1, stable=True))
+        drop_rank = ((~keep).cumsum(dim=1) - 1).clamp(min=0)
+        fill = candidates.gather(1, drop_rank).to(routed_experts.dtype)
+        return torch.where(keep, routed_experts, fill)
+
+
 @dataclass
 class MoEArgs:
     num_experts: int = 8
@@ -175,6 +201,8 @@ def __init__(
         self.route_scale = route_scale
         self.topk_sorted = topk_sorted
         self.force_balanced = False
+        # Set via trainer.router_replay_min_score_ratio, see `filter_replayed_experts`.
+        self.replay_min_score_ratio: float | None = None
 
     def forward(
         self,
@@ -218,6 +246,13 @@ def forward(
         # NOTE: selection biases are only used for routing. The gating value
         #       top_scores is still derived from the original scores/logits.
 
+        if routed_experts is not None and self.replay_min_score_ratio is not None:
+            # topk_softmax scores are logits; the threshold is a ratio of full-softmax probabilities.
+            filter_scores = F.softmax(scores.float(), dim=1) if self.score_func == "topk_softmax" else scores
+            routed_experts = filter_replayed_experts(
+                filter_scores, routed_experts, self.replay_min_score_ratio, self.selection_bias
+            )
+
         if routed_experts is not None:
             top_scores = scores.gather(dim=1, index=routed_experts)
             selected_experts_indices = routed_experts

src/prime_rl/trainer/rl/train.py

@@ -43,6 +43,7 @@
 from prime_rl.multimodal import get_multimodal_adapter
 from prime_rl.trainer.rl.annotations import AnnotationWriter
 from prime_rl.trainer.model import (
+    apply_router_replay_filter,
     forward,
     get_full_offload_dtype_policy,
     get_expert_load_stats,
@@ -145,6 +146,8 @@ def train(config: TrainerConfig):
     t0 = time.perf_counter()
     loading_from_ckpt_later = checkpoint_step is not None
     model = setup_model(config.model, parallel_dims, loading_from_ckpt_later)
+    if config.router_replay_min_score_ratio is not None:
+        apply_router_replay_filter(model, config.router_replay_min_score_ratio)
     logger.debug(f"Initialized model in {format_time(time.perf_counter() - t0)}")
 
     processor = None

tests/unit/train/models/test_moe.py

@@ -14,6 +14,8 @@
     GroupedExperts,
     MoE,
     MoEArgs,
+    TokenChoiceTopKRouter,
+    filter_replayed_experts,
 )
 from prime_rl.trainer.models.qwen3_5 import (
     Qwen3_5Config,
@@ -306,3 +308,32 @@ def test_expert_load_stats():
     assert stats.keys() == expected.keys()
     for name, value in expected.items():
         assert stats[name].item() == pytest.approx(value)
+
+
+def test_filter_replayed_experts():
+    scores = torch.tensor([[0.30, 0.25, 0.20, 0.15, 0.06, 0.04], [0.01, 0.02, 0.40, 0.30, 0.15, 0.12]])
+    routed = torch.tensor([[3, 4, 1], [0, 2, 1]], dtype=torch.int32)
+    # Thresholds are 0.5 * 0.20 and 0.5 * 0.15: dropped slots take the router's best experts not already kept.
+    filtered = filter_replayed_experts(scores, routed, min_score_ratio=0.5)
+    assert filtered.dtype == routed.dtype
+    assert filtered.tolist() == [[3, 0, 1], [3, 2, 4]]
+
+    # The bias steers the router's own top-k to [4, 0, 1]; the threshold uses unbiased scores (0.6 * 0.10).
+    scores = torch.tensor([[0.30, 0.25, 0.20, 0.15, 0.10, 0.05]])
+    bias = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.3, 0.0])
+    routed = torch.tensor([[5, 2, 3]])
+    assert filter_replayed_experts(scores, routed, 0.6, bias).tolist() == [[4, 2, 3]]
+
+    router = TokenChoiceTopKRouter(
+        dim=8, num_experts=6, top_k=3, score_func="softmax", route_norm=True, route_scale=1.0
+    )
+    router.init_weights(0.5)
+    x = torch.randn(16, 8)
+    routed = torch.stack([torch.randperm(6)[:3] for _ in range(16)])
+    router.replay_min_score_ratio = 1.0
+    top_scores, selected, _, _ = router(x, routed_experts=routed)
+    # With ratio 1 only the router's own top-k survive, as distinct experts.
+    own = torch.topk(router.gate(x), k=3, dim=1).indices
+    assert selected.sort(dim=1).values.tolist() == own.sort(dim=1).values.tolist()
+    top_scores.pow(2).sum().backward()
+    assert router.gate.weight.grad.abs().sum() > 0