[Lumen-RL] Enable Rollout Router Replay (R3)

atom/config.py

@@ -2012,6 +2012,43 @@ def indexer_cp_unsupported_reason(
     return None
 
 
+def check_return_routed_experts(config: "Config") -> None:
+    """Refuse topologies that cannot assemble a full [seq_len-1] route tensor."""
+    if not config.enable_return_routed_experts:
+        return
+    plugin = config.plugin_config
+    if (plugin is not None and plugin.is_plugin_mode) or is_plugin_mode():
+        raise ValueError(
+            "enable_return_routed_experts is not supported in plugin mode: "
+            "capture lives in ATOM's ModelRunner, which a host framework does "
+            "not run"
+        )
+    if (
+        config.decode_context_parallel_size != 1
+        or config.prefill_context_parallel_size != 1
+    ):
+        raise ValueError(
+            "enable_return_routed_experts requires decode_context_parallel_size "
+            "== 1 and prefill_context_parallel_size == 1"
+        )
+    if config.pipeline_parallel_size != 1:
+        raise ValueError(
+            "enable_return_routed_experts requires pipeline_parallel_size == 1"
+        )
+    if config.kv_transfer_config:
+        raise ValueError(
+            "enable_return_routed_experts does not support KV transfer "
+            "(prefill/decode disaggregation or KV offload); routed-expert "
+            "metadata is not moved with KV blocks"
+        )
+    if config.enable_rapidserve:
+        raise ValueError(
+            "enable_return_routed_experts does not support RapidServe "
+            "prefill/decode disaggregation; decode skips KV allocation and "
+            "never initializes the process-local capture buffer"
+        )
+
+
 @dataclass
 class Config:
     model: str
@@ -2100,6 +2137,11 @@ class Config:
     kv_cache_dtype: str = "bf16"
     index_cache_dtype: str | None = None
     enable_prefix_caching: bool = True
+    # Return per-request MoE routes from generate() as int16
+    # [seq_len - 1, num_layers, top_k]. Scatter is CUDA-graph-safe (device
+    # write inside fused MoE). Requires DCP=PCP=PP=1, no plugin mode, and
+    # no KV transfer / offload / RapidServe P/D; prefix cache stays on.
+    enable_return_routed_experts: bool = False
     enable_chunked_prefill: bool = True
     enable_log_stats: bool = True
     enable_request_context_metrics: bool = False
@@ -2622,6 +2664,8 @@ def __post_init__(self):
             self.offload_kv_encoding = snapshot_kv_encoding(self.hf_config)
         self._resolve_dynamic_chunking()
 
+        check_return_routed_experts(self)
+
         if self.speculative_config is not None:
             num_spec = self.speculative_config.num_speculative_tokens
             is_dspark = getattr(self.speculative_config, "use_dspark", lambda: False)()

atom/model_engine/arg_utils.py

@@ -131,6 +131,7 @@ class EngineArgs:
     data_parallel_base_port: int | None = None
     enforce_eager: bool = False
     enable_prefix_caching: bool = True
+    enable_return_routed_experts: bool = False
     port: int = 8006
     kv_cache_dtype: str = "bf16"
     index_cache_dtype: str | None = None
@@ -299,6 +300,15 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             help="Enable prefix caching (default: enabled). "
             "Use --no-enable_prefix_caching to disable.",
         )
+        parser.add_argument(
+            "--enable-return-routed-experts",
+            dest="enable_return_routed_experts",
+            action="store_true",
+            help="Return per-request MoE expert ids on generate() as "
+            "routed_experts int16 [seq_len-1, num_layers, top_k]. "
+            "Requires DCP=PCP=PP=1, no plugin mode, and no KV transfer, "
+            "offload, or RapidServe.",
+        )
         parser.add_argument(
             "--port",
             type=int,

atom/model_engine/llm_engine.py

@@ -1074,6 +1074,9 @@ def postprocess(self, reqs: list[Sequence]):
                 "ttft": ttft,  # Time to first token in seconds
                 "tpot": tpot,  # Time per output token in seconds
             }
+            routed_experts = req.routed_experts
+            if routed_experts is not None:
+                outputs[req.id]["routed_experts"] = routed_experts
         return outputs
 
     def has_pending_requests(self):