[RFC] Replace routing replay with CUDA-graph-compatible device cache approach

Motivation

Routing replay (--enable-return-routed-experts) captures which MoE experts process each token during inference. This is needed by RL training pipelines (GRPO, RLHF) where the training step reconstructs expert routing decisions from the inference pass.

We've been running a fork with a production-grade routing replay replacement on our internal GPU clusters for large MoE models (120B and 400B+ parameter). The current upstream implementation has fundamental issues that prevent real-world use - it breaks under CUDA graphs, doesn't work multi-node, misses the monolithic kernel path entirely, and has no MTP or prefix caching support.

We'd like to upstream our replacement. The implementation is code complete and validated across 9 configurations with <5% throughput overhead and zero accuracy impact. A companion FlashInfer PR (adding routing_replay_out to MoE kernel launchers) is currently under review.

What's wrong with the current implementation

The current approach uses a capture() callback in the router combined with multiprocessing.SharedMemory and fcntl file locking. We found 8 issues:

1. CUDA graph incompatibility. The device buffer (routed_experts_capturer.py:133) is not marked with cudagraph_mark_tensor_static. During CUDA graph capture, the graph snapshots this buffer. On replay, it restores the snapshot before executing, overwriting clear_buffer() calls made outside the graph. Positions not rewritten by the current forward retain stale snapshot data instead of zeros. This took us 5 different approaches to solve correctly.

2. Monolithic kernels not captured. Routing is only captured via select_experts() in the router (base_router.py:280). The monolithic trtllm-gen path (apply_monolithic()) bypasses the router, so no routing data is captured for FP8/MXFP8 models using FlashInfer's fused kernels. This is the default serving path for quantized MoE models.

3. No multi-node support. SharedMemory is node-local (routed_experts_capturer.py:14,58-60). On multi-node TP setups (needed for 400B+ models), the scheduler on node 0 can't read shared memory from workers on other nodes. Our replacement removes shared memory entirely - routing data flows through ModelRunnerOutput via Ray DAG to the scheduler, same as other per-request outputs.

4. No MTP support. There's no handling of speculative tokens from MTP anywhere in the routing replay code path. With MTP, the model captures routing for all tokens including rejected speculative positions - the output needs trimming to match actual accepted tokens.

5. No prefix caching support. The host buffer is initialized with fill(0) (routed_experts_capturer.py:154). This means cached positions (prefix hits) are indistinguishable from "routed to expert 0" - 0 is a valid expert ID.

6. Synchronous D2H. save_captured_experts does .cpu().numpy() (routed_experts_capturer.py:230) which blocks the GPU, then takes a file lock (line 232) for shared memory access.

7. Buffer layout mismatch. Device buffer is (N, L, K) int32 (routed_experts_capturer.py:133-137). FlashInfer's routing_replay_out parameter needs contiguous (N, K) per layer, which requires (L, N, K) for zero-copy slicing.

8. Memory overhead. Uses int32 for expert IDs that fit in int16.

Proposed design

Device cache + async D2H pipeline replacing the shared-memory capturer/reader:

• Pre-allocated (L, N, K) int16 device buffer. buffer[layer_id] gives a contiguous (N, K) view per layer.
• Each FusedMoE layer gets a persistent module attribute module._routing_replay_out = buffer[layer_id]. torch.compile captures module attributes by reference, so CUDA graph replay always writes to the live buffer.
• cudagraph_mark_tensor_static on each per-layer view prevents snapshot/restore from zeroing the data.
• Async D2H via CUDA events + pinned memory scatter, only on TP rank 0.
• Per-request numpy host cache initialized with -1 sentinel.

Monolithic kernel integration: Thread routing_replay_out through the apply_monolithic() call chain so FlashInfer's fused kernels write expert IDs directly during routing. For the non-monolithic Triton path: write topk_ids.to(int16) after select_experts(). A _monolithic_writes_routing_replay flag distinguishes kernel capabilities.

MTP + prefix caching: Seqlen clamping using authoritative token count from request state. Output trim in output_processor to match actual accepted tokens. Host cache init with -1 sentinel instead of 0.

Multi-node: Device buffer on all TP ranks (symmetric CUDA graphs). Only rank 0 does D2H and host cache management. Data flows through ModelRunnerOutput -> Ray DAG -> scheduler, replacing the SharedMemory path entirely.

FlashInfer dependency

This depends on a companion FlashInfer PR that adds routing_replay_out as an optional parameter to all MoE kernel launchers (FP8 block, FP8 per-tensor, BF16, FP4, MXINT4) and routing kernels (Custom, DeepSeek, Llama4). When None (the default), there is zero overhead - the kernel skips the write entirely.

Files changed

Full rewrite: routed_experts_capturer.py

Modified (~12 files): gpu_model_runner.py (capturer lifecycle), layer.py (moe_layer_id), moe_runner_base.py (non-monolithic write), modelopt.py / modular_kernel.py / fp8.py / trtllm_fp8_moe.py (monolithic threading), output_processor.py (MTP trim), OpenAI API protocol + serving files (response field).

Removed: RoutedExpertsReader, SharedMemory usage, fcntl locking, capture() callback in router, KV cache slot_mapping retrieval.

Validation

Tested on GB200 GPUs with a 120B MoE model (BF16 dummy weights) and a 400B+ MoE model (MXFP8 and BF16, real weights). All configs use expert parallelism.

Functional tests (9 configs, all passing - 5 correctness checks + 1000-prompt scale test at 256 concurrency, ISL=1024/OSL=1024):

• FP8/MXFP8 monolithic kernel path (single-node and multi-node)
• BF16 non-monolithic Triton path (single-node and multi-node)
• Prefix caching enabled
• Data parallelism (TP=4, DP=2)
• MTP speculative decoding (5 draft tokens)
• MTP + prefix caching combined

Throughput (400B+ MoE, MXFP8, 2 nodes, TP=8, 1000 prompts, 256 concurrency):

Accuracy (GSM8K, 4 seeds, 1319 problems): pass@1 = 95.77% for both baseline and RR. Zero accuracy impact.

API compatibility

Fully preserved - same CLI flag (--enable-return-routed-experts), same output field (routed_experts), same shape ([seq_len, num_moe_layers, top_k]). Only internals change.

Plan

1. FlashInfer PR merges (under review)
2. vLLM pins FlashInfer version with routing_replay_out support
3. vLLM PR: 3 commits (core plumbing, MTP+PC support, tests)

Happy to adjust the PR structure or split differently based on feedback.