[Feature]: Router Replay(R3): capture per-token MoE routing for train/inference alignment

🚀 The feature, motivation and pitch

[RFC] Return per-token MoE routing from the engine (Router Replay)

• Status: Proposal
• Method reference: Ma et al., "Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers", arXiv:2510.11370 — Rollout Routing Replay (R3): record routing from the inference engine and replay it during training to prevent RL collapse. This RFC adds the engine-side capture that R3 consumes.
• Prior art (other engines): vLLM #38079 (enable_return_routed_experts) · SGLang return_routed_experts (#20964, #31116) · Megatron-LM #4168 (trainer replay).

1. Summary

Add an opt-in capability to return, per token and per MoE layer, the pre-EPLB logical top-k expert IDs the router selected during generation, surfaced on CompletionOutput.routed_experts with contract (seq_len − 1, num_moe_layers, top_k). A downstream MoE-RL trainer replays this exact routing in its own forward pass, removing the train/inference routing mismatch. vLLM and SGLang already expose this; TensorRT-LLM does not.

2. Motivation

In MoE reinforcement learning the rollout engine and the trainer are separate implementations, so even at identical weights their routers disagree on a non-trivial fraction of tokens. This mismatch destabilizes and can collapse training. R3 (arXiv:2510.11370) fixes it by forcing the trainer to reuse the routing the rollout actually took — but that requires the inference engine to return its routing. TRT-LLM currently cannot, which blocks using it as an MoE-RL rollout engine that stays consistent with the trainer.

3. Goals / Non-goals

Goals

• Faithfully return the pre-EPLB logical top-k on the separated-routing backends (CUTLASS / DeepGemm) with no kernel changes.
• Zero overhead when disabled; opt-in per request.
• Off the forward critical path — no synchronization, locking, or shared-memory writes on the hot path (the naive approach deadlocks the event loop → EP NCCL timeout, cf. vLLM #38079).
• Transparent to engine features: correct with prefix caching, CUDA graphs, the overlap scheduler, and EPLB/EP all enabled — Router Replay is a passive recorder and must not require any of them to be turned off.

Non-goals

• Fully-fused backends that compute top-k inside the kernel (fail-closed).
• PP > 1, MTP, speculative decoding (fail-closed; interface reserved).
• The trainer-side replay algorithm (lives in the trainer; out of scope).

4. Proposed design

4.1 Capture

The captured value is the output of routing_method.apply() — the pre-EPLB logical top-k, before _load_balancer_route(). This point is backend-agnostic (CUTLASS / DeepGemm share it) and requires no kernel change; it is also EPLB-invariant (EPLB only remaps logical → physical after this point, so the captured routing is identical whether EPLB is on or off). Non-separated / fused backends fail-closed via assert_capturable.

A per-rank RouteCapturer pre-allocates a device buffer; a one-line capture(layer_id, topk_ids) right after top-k writes it inside the CUDA graph (a Python-side hook would not execute during graph replay and would silently miss all graphed decode steps). A RouteCopier moves the buffer to host with one non-blocking D2H per step on a dedicated stream — never synchronizing on the forward path.

All capture / staging / commit are keyed by the immutable iter_counter (step id) rather than current request state, so the overlap scheduler stays correct.

4.2 Storage & lifetime

Routes are held in a per-request route store (PerRequestRouteStore, authoritative), indexed by absolute position so that chunked prefill, preemption, and recompute are idempotent and never drop rows. A separate shared route cache (SharedRouteCache), keyed by the reusable-block identity (pool_id, block_id, offset) + block_generation, holds routes for cross-request prefix sharing.

The invariant is route validity == KV-block validity: routes ride with the KV. On a prefix-cache hit, the hit tokens are not re-forwarded, so their routes are read from the shared cache into the request's route store; when a refit flushes the KV cache, the corresponding routes are invalidated in the same step. This makes prefix caching, preemption, and weight-drifted async rollout fall out of one rule rather than needing special cases. (Deliberately no per-policy salt, which would defeat cross-version KV reuse.)

4.3 Output

On request finish the route store assembles to [seq_len − 1, num_moe_layers, top_k] (genuine gap → -1 sentinel, fail-closed) and is attached to CompletionOutput.routed_experts.

SamplingParams(..., return_routed_experts=True)
#  → RequestOutput.outputs[0].routed_experts : int [tokens, num_moe_layers, top_k]

4.4 Correctness considerations

Because capture rides the live inference path, the design is built around the cases that would otherwise corrupt or lose routes:

Case	How the design handles it
Continuous batching / chunked prefill	attribute rows by absolute position, mirroring the engine's token layout (a naive prompt_len layout misattributes every row)
CUDA graph (batched decode)	capture into the device buffer inside the graph (a Python-side hook never fires during graph replay)
Prefix-cache hit	hit tokens are not re-forwarded → reuse their routes from the shared route cache (SharedRouteCache)
Overlap scheduler	key everything on the immutable iter_counter, never on live request state
Prefill/decode boundary (off-by-one)	the S−1 contract drops the final position, aligned to the trainer (cf. SGLang #20964)
Expert-parallel ID space	capture global logical expert IDs, in the same space the trainer replays
Preemption / SWA eviction	the per-request route store is authoritative → no snapshot; only the shared-cache entry is invalidated, recompute overwrites idempotently
Async weight refit / version skew	route validity = KV validity; routes invalidate with the KV cache flush; no per-policy salt
Missing vs. padding	genuine gap → -1 sentinel (fail-closed); legal padding → arange(top_k) dummy — never conflated
Unsupported paths (fused / PP>1 / MTP / spec-decode)	fail-closed via assert_capturable — never silently return zeros
Concurrency / deadlock	no sync / lock / shared-memory write on the forward path → no EP NCCL timeout (cf. vLLM #38079)

5. Components to add

#	Name	Location	Responsibility
C1	RouteCapturer (new class, per-rank singleton)	tensorrt_llm/_torch/route_capture.py (new; beside expert_statistic.py)	device buffer; capture(layer_id, topk_ids) inside the CUDA graph; get_routed_experts()
C2	assert_capturable(backend) (new fn)	same file	fail-close for fused / PP>1 / non-dropless backends
C3	capture hook (1 line ×2 sites)	_torch/modules/fused_moe/moe_scheduler.py (after routing_method.apply())	RouteCapturer.capture(layer_idx, token_selected_experts)
C4	RouteCopier (new class)	_torch/pyexecutor/route_copier.py (new)	token→slot map; one non-blocking D2H per step on a side stream
C5	RouteManager (new class, BaseResourceManager)	_torch/pyexecutor/route_manager.py (new; registered in resource_manager.py)	authoritative PerRequestRouteStore + SharedRouteCache; commit by iter_counter; assemble on finish
C6	driver hooks (edits)	_torch/pyexecutor/{model_engine,py_executor}.py	create capturer/manager; stage per step; copy from shared cache on prefix hit; attach on finish
C7	opt-in gates	LlmArgs / SamplingParams	enable_return_routed_experts (engine) + return_routed_experts (per request)

6. Implementation plan

• Phase 1 — capture + route store + output: return routed_experts for the common path (separated-routing backend). Correctness first.
• Phase 2 — full lifetime: SharedRouteCache for prefix-cache reuse, preemption/refit invalidation, and async weight drift, at full throughput with graphs/overlap on.
• Phase 3 (optional): fail-closed coverage extended (or, orthogonally, deterministic MoE kernels for bit-exact logprobs — a separate numerical-determinism effort, not routing).

7. Alternatives / related

• File lock + shared memory (vLLM's first attempt): deadlocks the event loop under concurrency; rejected in favor of device buffer + non-blocking D2H (#38079).
• In-kernel capture for fused backends: covers all backends but needs kernel changes; deferred, fail-closed for now.
• slime #2262 captures routing on SGLang and adds deterministic kernels; confirms the approach but is SGLang/DeepEP-specific.

8. Open questions

• Device-buffer sizing against a dynamic max_num_tokens.
• Exact mapping of the SharedRouteCache block-identity key to reusable-block semantics under async refit.
• Interaction with disaggregated (prefill/decode-split) serving.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.