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

atom/config.py

@@ -1728,6 +1728,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 DP-attention, 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
     # Seconds between engine-status lines. Validated > 0 by EngineStats.
@@ -2142,6 +2147,24 @@ def __post_init__(self):
 
                 self.kv_transfer_config = ast.literal_eval(self.kv_transfer_config)
 
+        if self.enable_return_routed_experts:
+            from atom.model_ops.fused_moe.routed_experts_capturer import (
+                check_return_routed_experts,
+            )
+
+            check_return_routed_experts(
+                self.decode_context_parallel_size,
+                self.prefill_context_parallel_size,
+                self.pipeline_parallel_size,
+                kv_transfer_config=self.kv_transfer_config,
+                enable_rapidserve=(
+                    self.enable_rapidserve
+                    or self.disagg_is_decode
+                    or "RapidServeModelRunner" in self.runner_qualname
+                ),
+                enable_dp_attention=self.enable_dp_attention,
+            )
+
         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

@@ -49,6 +49,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
@@ -207,6 +208,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 DP-attention, and no KV transfer, "
+            "offload, or RapidServe.",
+        )
         parser.add_argument(
             "--port",
             type=int,

atom/model_engine/llm_engine.py

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

atom/model_engine/model_runner.py

@@ -61,6 +61,7 @@
     Pool,
     PoolPlan,
     SubPoolSpec,
+    page_pool,
     plan_pools,
 )
 from atom.model_ops.decode_input_ids import (
@@ -1446,6 +1447,18 @@ def _sub_pool_specs(self) -> list[SubPoolSpec]:
         specs = list(self.attn_metadata_builder.sub_pool_specs())
         if hasattr(self, "draft_kv_builder"):
             specs += self.draft_kv_builder.sub_pool_specs()
+        geo = self._routed_experts_geometry()
+        if geo is not None:
+            from atom.model_ops.fused_moe.routed_experts_capturer import (
+                capture_bytes_per_kv_block,
+            )
+
+            num_layers, top_k = geo
+            specs.append(
+                page_pool(
+                    capture_bytes_per_kv_block(self.block_size, num_layers, top_k)
+                )
+            )
         return specs
 
     def _state_pool_names(self) -> list[str]:
@@ -1594,6 +1607,16 @@ def get_num_blocks(self) -> dict[str, object]:
 
         torch.set_default_device("cpu")
 
+        geo = self._routed_experts_geometry()
+        if geo is not None:
+            from atom.model_ops.fused_moe.routed_experts_capturer import (
+                capture_pad_row_bytes,
+            )
+
+            available_for_kv = max(
+                0, int(available_for_kv) - capture_pad_row_bytes(*geo)
+            )
+
         specs = self._sub_pool_specs()
 
         # Sub-pool sizing is pure arithmetic over the byte budget — see
@@ -1972,6 +1995,11 @@ def allocate_kv_cache(self, num_kvcache_blocks):
             num_blocks=num_kvcache_blocks,
         )
 
+        # Capture buffer is part of the PAGE budget (per-slot int32 routes).
+        # Allocate it before the expected-vs-actual check so the measurement
+        # includes those bytes rather than warning on a false mismatch.
+        self._maybe_init_routed_experts_capturer(num_kvcache_blocks)
+
         # Cross-validate: compare estimated vs actual KV cache allocation.
         # `actual_kv_bytes` includes BOTH the unified pool tensors (counted by
         # `block_bytes × num_blocks`) AND the per-request cache tensors (state
@@ -2012,6 +2040,122 @@ def allocate_kv_cache(self, num_kvcache_blocks):
             torch.distributed.barrier()
         return True
 
+    def _routed_experts_geometry(self) -> tuple[int, int] | None:
+        """``(num_layers, top_k)`` for the capture buffer, or None if unused.
+
+        Assigns ``moe_capture_layer_id`` on each FusedMoE. Cached so pool
+        planning and capturer init share one walk.
+        """
+        if hasattr(self, "_routed_experts_geo"):
+            return self._routed_experts_geo
+        if not getattr(self.config, "enable_return_routed_experts", False):
+            self._routed_experts_geo = None
+            return None
+        if not hasattr(self, "model"):
+            return None
+        import re
+
+        from atom.model_ops.fused_moe.routed_experts_capturer import fused_moe_modules
+
+        moes: list = fused_moe_modules(self.model)
+        if not moes:
+            raise ValueError(
+                "enable_return_routed_experts requires at least one FusedMoE "
+                "module; this model has none (dense checkpoints cannot "
+                "return routed_experts)"
+            )
+        layer_ids: list[int] = []
+        for ordinal, moe in enumerate(moes):
+            prefix = getattr(moe, "prefix", "") or ""
+            match = re.search(r"layers\.(\d+)", prefix)
+            if match:
+                idx = int(match.group(1))
+            elif getattr(moe, "layer_id", None) is not None:
+                idx = int(moe.layer_id)
+            else:
+                idx = ordinal
+            moe.moe_capture_layer_id = idx
+            layer_ids.append(idx)
+        self._routed_experts_geo = (
+            max(layer_ids) + 1,
+            max(int(m.top_k) for m in moes),
+        )
+        return self._routed_experts_geo
+
+    def _maybe_init_routed_experts_capturer(self, num_kvcache_blocks: int) -> None:
+        geo = self._routed_experts_geometry()
+        if geo is None:
+            return
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            RoutedExpertsCapturer,
+        )
+
+        num_layers, top_k = geo
+        num_slots = int(num_kvcache_blocks) * int(self.block_size)
+        RoutedExpertsCapturer.init(
+            num_slots=num_slots,
+            num_layers=num_layers,
+            top_k=top_k,
+            device=self.device,
+        )
+
+    def _store_routed_experts_step(
+        self,
+        batch: ScheduledBatch,
+        *,
+        stream: torch.cuda.Stream | None = None,
+        wait_event: torch.cuda.Event | None = None,
+        slot_mapping: torch.Tensor | None = None,
+    ) -> bool:
+        if not getattr(self.config, "enable_return_routed_experts", False):
+            return False
+        if getattr(batch, "is_dummy_run", False) or self.rank != 0:
+            return False
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            RoutedExpertsCapturer,
+        )
+
+        capturer = RoutedExpertsCapturer.get()
+        if capturer is None:
+            return False
+        return capturer.store_step(slot_mapping, stream=stream, wait_event=wait_event)
+
+    def _commit_routed_experts(self, *, keep_last: bool = False) -> None:
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            RoutedExpertsCapturer,
+        )
+
+        capturer = RoutedExpertsCapturer.get()
+        if capturer is not None:
+            capturer.commit_pending(keep_last=keep_last)
+
+    def _attach_routed_experts(
+        self, batch: ScheduledBatch, output: ScheduledBatchOutput
+    ) -> None:
+        """CPU gather in this runner process. The capturer is not visible to
+        the EngineCore scheduler (runner is an ``AsyncIOProcManager`` child).
+        """
+        if not getattr(self.config, "enable_return_routed_experts", False):
+            return
+        if getattr(batch, "is_dummy_run", False) or self.rank != 0:
+            return
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            RoutedExpertsCapturer,
+        )
+
+        capturer = RoutedExpertsCapturer.get()
+        if capturer is None:
+            return
+        block_tables = getattr(batch, "block_tables", None)
+        if not block_tables or len(block_tables) != len(batch.req_ids):
+            return
+        output.routed_experts = capturer.export_batch(
+            batch.req_ids,
+            block_tables,
+            batch.context_lens,
+            self.block_size,
+        )
+
     def get_dp_padding(self, num_tokens: int) -> tuple[int, torch.Tensor | None]:
         dp_size = self.config.parallel_config.data_parallel_size
         dp_rank = self.config.parallel_config.data_parallel_rank
@@ -3080,11 +3224,21 @@ def postprocess(
                 sampled_logprobs = get_tp_group().broadcast(sampled_logprobs, src=0)
 
         self.forward_done_event.record()
-        # Capture before prepare_sampled_ids(), which advances self.prev_batch to current batch.
+        deferred = self.tokenID_processor.is_deferred_out
+        # Capture before prepare_sampled_ids(), which advances self.prev_batch.
         prev_batch = self.tokenID_processor.prev_batch
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            _current_slot_mapping,
+        )
+
+        slots = _current_slot_mapping()
         token_id_dict, logprobs_map = self.tokenID_processor.prepare_sampled_ids(
             batch, sampled_tokens, self.forward_done_event, sampled_logprobs
         )
+        # Blocking D2H after recv: previous copy_done already waited, so this
+        # memcpy does not absorb the current decode. cpu_buffer is complete
+        # before the next step attaches this batch as prev_batch.
+        self._store_routed_experts_step(batch, slot_mapping=slots)
         # Extract req_ids and token_ids from dict (key -1 is the is_deferred_out flag)
         req_ids_out = [k for k in token_id_dict if k != -1]
         token_ids_out = [token_id_dict[k] for k in req_ids_out]
@@ -3151,7 +3305,7 @@ def postprocess(
         if verify_scheduler is not None:
             dspark_ell = verify_scheduler.ell_nonblocking()
 
-        return ScheduledBatchOutput(
+        output = ScheduledBatchOutput(
             req_ids=req_ids_out,
             token_ids=token_ids_out,
             draft_token_ids=draft_token_ids,
@@ -3161,6 +3315,10 @@ def postprocess(
             logprobs=logprobs_map,
             dspark_ell=dspark_ell,
         )
+        attach_batch = prev_batch if deferred and prev_batch is not None else batch
+        if not (deferred and prev_batch is None):
+            self._attach_routed_experts(attach_batch, output)
+        return output
 
     def _record_kv_cache_ready(self, batch: ScheduledBatch) -> None:
         """Publish a GPU event for final prefill chunks to transfer connectors."""
@@ -3250,17 +3408,21 @@ def forward(self, batch: ScheduledBatch) -> ScheduledBatchOutput:
                     None,  # nothing verified -> segment's last row
                     align_only=True,
                 )
+            self.forward_done_event.record()
+            # No token copy_done on this path; blocking store is fine.
+            self._store_routed_experts_step(batch)
             reset_forward_context()
             # Mark this slot's GPU work (attention consumed its metadata) done.
             self._record_forward_vars_event()
             self._record_kv_cache_ready(batch)
-            return ScheduledBatchOutput(
+            output = ScheduledBatchOutput(
                 req_ids=list(batch.req_ids),
                 token_ids=[],
                 num_rejected=None,
                 num_bonus=None,
                 draft_token_ids=None,
             )
+            return output
 
         fwd_output = self.postprocess(
             batch,
@@ -4213,9 +4375,29 @@ def _kv_budget_extra_reserve(self, total_bytes: int) -> int:
         safety_margin = int(total_bytes * 0.02)
         return 4 * safety_margin
 
+    def _refuse_routed_experts_capture(self) -> None:
+        if not getattr(self.config, "enable_return_routed_experts", False):
+            return
+        from atom.model_ops.fused_moe.routed_experts_capturer import (
+            check_return_routed_experts,
+        )
+
+        check_return_routed_experts(
+            getattr(self.config, "decode_context_parallel_size", 1),
+            getattr(self.config, "prefill_context_parallel_size", 1),
+            getattr(self.config, "pipeline_parallel_size", 1),
+            kv_transfer_config=getattr(self.config, "kv_transfer_config", None),
+            enable_rapidserve=True,
+            enable_dp_attention=bool(
+                getattr(self.config, "enable_dp_attention", False)
+            ),
+        )
+
     def get_num_blocks(self) -> dict[str, object]:
         # Decode in disagg mode owns no GPU memory — kvcache is imported from
-        # prefill.
+        # prefill. That also skips `_maybe_init_routed_experts_capturer`, so
+        # capture cannot be served from this process.
+        self._refuse_routed_experts_capture()
         if self.config.disagg_is_decode:
             transfer = self.attn_metadata_builder.state_transfer()
             if transfer.copies:
@@ -4231,6 +4413,7 @@ def get_num_blocks(self) -> dict[str, object]:
 
     def allocate_kv_cache(self, num_kvcache_blocks):
         # Decode in disagg mode: kvcache is imported from prefill, not allocated.
+        self._refuse_routed_experts_capture()
         if self.config.disagg_is_decode:
             logger.info("decode skipping kv cache allocation")
             return True

atom/model_engine/request.py

@@ -2,16 +2,17 @@
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 
 @dataclass
 class RequestOutput:
     """Output structure passed to stream callback."""
 
     request_id: int
-    output_tokens: List[int]
+    output_tokens: list[int]
     finished: bool
-    finish_reason: Optional[str] = None
-    kv_transfer_params_output: Optional[Dict[str, Any]] = None
+    finish_reason: str | None = None
+    kv_transfer_params_output: dict[str, Any] | None = None
     num_cached_tokens: int = 0
+    routed_experts: Any | None = None