Record moe routing decisions during inference.

examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py

@@ -145,6 +145,9 @@ async def main(
                     result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs
                 throughput = len(req.generated_tokens) / req.latency
                 throughputs.append(throughput)
+                if req.routing_indices is not None:
+                    result_dict["routing_indices"] = req.routing_indices.tolist()
+
                 json_results[req.request_id] = result_dict
             throughput_dict = {"throughput": throughputs}
             if args.throughput_check_only:

megatron/core/inference/contexts/dynamic_context.py

@@ -35,6 +35,7 @@
 from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
 from .base_context import BaseInferenceContext
 from .dynamic_block_allocator import BlockAllocator
+from .routing_metadata import RoutingMetadata
 
 try:
     from .fused_kv_append_kernel import triton_append_key_value_cache
@@ -469,6 +470,13 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             max_seqlen=self.max_sequence_length,
         )
 
+        self.moe_enable_routing_replay = model_config.moe_enable_routing_replay
+        if self.moe_enable_routing_replay:
+            assert (
+                model_config.num_moe_experts is not None
+            ), "Router recording/replay requested but no MoE experts specified!"
+            self.moe_routing_metadata = RoutingMetadata(self, model_config.moe_router_topk)
+
         # CUDA graph config list
         self.use_cuda_graphs_for_non_decode_steps = (
             inference_config.use_cuda_graphs_for_non_decode_steps
@@ -1294,6 +1302,12 @@ def initialize_attention_state(
                 padded_batch_dimensions=self.padded_batch_dimensions,
             )
 
+        if self.moe_enable_routing_replay:
+            if self.using_cuda_graph_this_step():
+                self.moe_routing_metadata.enable_static_buffer_recording()
+            else:
+                self.moe_routing_metadata.disable_static_buffer_recording()
+
     def reset(self) -> None:
         """Reset entire context.
 

megatron/core/inference/contexts/routing_metadata.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+
+from megatron.core.transformer.moe.router_replay import RouterReplay
+
+
+class RoutingMetadata:
+    """Manages routing indices metadata for MoE layers during inference.
+
+    This class provides static buffers for CUDA graph compatibility when
+    recording routing decisions. It holds a reference to the inference context
+    to automatically determine whether to use static buffers based on CUDA graph state.
+
+    Args:
+        context (DynamicInferenceContext): The inference context.
+        moe_router_topk (int): Number of experts selected per token.
+    """
+
+    def __init__(self, context: 'DynamicInferenceContext', moe_router_topk: int):
+        self.context = context
+        self.max_tokens = context.max_tokens
+        self.moe_router_topk = moe_router_topk
+        self.device = torch.cuda.current_device()
+
+        # Static buffer allocated lazily in _ensure_buffer_allocated().
+        # We defer allocation because RouterReplay instances don't exist yet at init time.
+        self.routing_indices_buffer: Optional[torch.Tensor] = None
+        self.num_moe_layers: Optional[int] = None
+
+    def _ensure_buffer_allocated(self) -> None:
+        """Allocate the static buffer if not already allocated.
+
+        Gets the actual number of MoE layers from RouterReplay instances.
+        """
+        if self.routing_indices_buffer is not None:
+            return
+
+        self.num_moe_layers = len(RouterReplay.global_router_replay_instances)
+
+        if self.num_moe_layers == 0:
+            return
+
+        # Static buffer for CUDA graph compatibility.
+        # Shape: [max_tokens, num_moe_layers, moe_router_topk]
+        self.routing_indices_buffer = torch.empty(
+            (self.max_tokens, self.num_moe_layers, self.moe_router_topk),
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+    def get_routing_indices(self) -> Optional[torch.Tensor]:
+        """Get the recorded routing indices.
+
+        Automatically uses the static buffer when CUDA graphs are active,
+        otherwise retrieves from RouterReplay utility.
+
+        Returns:
+            Tensor of shape [num_tokens, num_moe_layers, topk] or None if not available.
+        """
+        if self.context.using_cuda_graph_this_step():
+            # Return view of static buffer up to current token count.
+            if self.routing_indices_buffer is None:
+                return None
+            # Only return up to active token count, to skip entries
+            # for padding tokens.
+            return self.routing_indices_buffer[: self.context.active_token_count]
+        else:
+            # Get from RouterReplay and stack into [num_tokens, num_layers, topk].
+            recorded_data = RouterReplay.get_recorded_data()
+            if recorded_data is None or len(recorded_data) == 0:
+                return None
+            if recorded_data[0] is None:
+                return None
+            # Stack: list of [num_tokens, topk] -> [num_tokens, num_layers, topk]
+            return torch.stack(recorded_data, dim=1)
+
+    def enable_static_buffer_recording(self) -> None:
+        """Enable recording into the static buffer for CUDA graph compatibility.
+
+        This sets up RouterReplay instances to copy routing indices into our
+        pre-allocated static buffer instead of creating new tensors.
+        Allocates the buffer lazily on first call.
+        """
+        self._ensure_buffer_allocated()
+        if self.routing_indices_buffer is not None:
+            RouterReplay.set_global_static_buffers(self.routing_indices_buffer)
+
+    def disable_static_buffer_recording(self) -> None:
+        """Disable static buffer recording, reverting to normal tensor assignment."""
+        RouterReplay.clear_global_static_buffers()

megatron/core/inference/engines/dynamic_engine.py

@@ -41,6 +41,7 @@
 from megatron.core.inference.utils import Counter, await_process_call
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
 from megatron.core.utils import (
     deprecate_args,
     experimental_api,
@@ -297,6 +298,12 @@ def create_cuda_graphs(self, reset_context: bool = True):
                     f"{tbar_idx}/{len(context.cuda_graph_batch_dimensions_list)}. {tbar_str}"
                 )
 
+            # Enable routing recording during warmup if routing replay is enabled.
+            # This ensures the record_indices copy operation is captured in the CUDA graph.
+            model_config = controller.inference_wrapped_model.model.config
+            if model_config.moe_enable_routing_replay:
+                RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+
             # Forward pass -> logits.
             controller._dynamic_step_forward_logits(input_ids, position_ids)
 
@@ -837,6 +844,7 @@ def post_process_requests(
         sample: torch.Tensor,
         log_probs: torch.Tensor,
         top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
+        routing_indices_per_request: Optional[Dict[int, torch.Tensor]] = None,
     ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest]]:
         """
         Handles post-processing for requests after a step.
@@ -850,6 +858,9 @@ def post_process_requests(
             log_probs: (List): Log probs for each request
             top_n_logprobs: (Dict): Top-n log probs for each request. Maps request_idx to
                 list of (top_n_logprobs, top_n_indices) tuples.
+            routing_indices_per_request: (Dict[int, Tensor]): MoE routing indices
+                pre-mapped by request_id. Each value is a tensor of shape
+                [num_tokens_this_step, num_layers, topk].
 
         Returns:
             A list of active requests and completed requests as `DynamicInferenceRequest` objects
@@ -979,6 +990,23 @@ def post_process_requests(
                     else:
                         request.generated_top_n_logprobs.append(logit_dict)
 
+            # Process routing indices if available (keyed by request_id)
+            # Each step's routing is a tensor of shape [num_tokens_this_step, num_layers, topk]
+            # We concatenate along dim=0 to accumulate: [total_tokens, num_layers, topk]
+            if (
+                routing_indices_per_request is not None
+                and request_id in routing_indices_per_request
+            ):
+                step_routing = routing_indices_per_request[
+                    request_id
+                ]  # [num_tokens, num_layers, topk]
+                if request.routing_indices is None:
+                    request.routing_indices = step_routing.clone()
+                else:
+                    request.routing_indices = torch.cat(
+                        [request.routing_indices, step_routing], dim=0
+                    )
+
         # Handle evicted requests.
         if evict_request_ids is not None and evict_request_ids.numel() > 0:
 
@@ -1248,6 +1276,7 @@ async def async_bookkeep(
             sample = step_result["sample"]
             log_probs = step_result["log_probs"]
             top_n_logprobs = step_result.get("top_n_logprobs", None)
+            routing_indices_per_request = step_result.get("routing_indices_per_request", None)
             cuda_graph_request_count = step_result["cuda_graph_request_count"]
 
             # Add paused events.
@@ -1266,6 +1295,7 @@ async def async_bookkeep(
                 sample,
                 log_probs,
                 top_n_logprobs,
+                routing_indices_per_request,
             )
 
         else:

megatron/core/inference/inference_request.py

@@ -271,6 +271,9 @@ class DynamicInferenceRequest(InferenceRequest):
     # remaining prompt tokens are used for chunked prefill
     remaining_prompt_tokens: Optional[torch.Tensor] = None
     latency: Optional[float] = None
+    # routing_indices stores MoE routing decisions for all tokens generated so far.
+    # Shape: [total_tokens, num_layers, topk] - accumulated across all generation steps
+    routing_indices: Optional[torch.Tensor] = None
     finished_chunk_token_count: int = 0
     stop_word_ids: Optional[List[List[int]]] = None  # Tokenized stop words (populated internally)
 
@@ -309,6 +312,17 @@ def serialize(self) -> dict:
         torch.cuda.nvtx.range_push("DynamicInferenceRequest.serialize")
         obj = super().serialize()
         obj["events"] = [e.serialize() for e in self.events]
+
+        # Sanity check routing_indices: Tensor [total_tokens - 1, num_layers, topk]
+        if self.routing_indices is not None:
+            total_tokens = len(self.prompt_tokens) + len(self.generated_tokens)
+            # the last generated token does not undergo a forward pass
+            # hence we expect routing indices for total_tokens - 1
+            assert self.routing_indices.shape[0] == total_tokens - 1, (
+                f"routing_indices first dimension {self.routing_indices.shape[0]} does not match "
+                f"total tokens {total_tokens-1}."
+            )
+
         torch.cuda.nvtx.range_pop()
         return obj
 
@@ -499,6 +513,9 @@ def merge_lists(key):
 
         prompt_tokens = self.requests[0].prompt_tokens
         prompt_text = self.requests[0].prompt
+        routing_indices = None
+        if self.requests[0].routing_indices is not None:
+            routing_indices = torch.cat([r.routing_indices for r in self.requests])
         generated_tokens = merge_lists("generated_tokens")
         try:
             generated_text = "".join(r.generated_text for r in self.requests)
@@ -522,6 +539,7 @@ def merge_lists(key):
             status=self.requests[-1].status,
             latency=self.latency,
             events=merge_lists("events"),
+            routing_indices=routing_indices,
         )
 
         return request

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -27,10 +27,12 @@
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding
 from megatron.core.models.multimodal.llava_model import LLaVAModel
+from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
 from megatron.core.transformer.utils import set_model_to_sequence_parallel
-from megatron.core.utils import get_asyncio_loop, get_model_config, unwrap_model
+from megatron.core.utils import get_asyncio_loop, get_model_config, get_pg_size, unwrap_model
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
@@ -675,6 +677,66 @@ def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
 
         return return_log_probs.any(), top_n_log_probs.any()
 
+    def _router_record_bookkeeping(self) -> Optional[Dict[int, Tensor]]:
+        """Collect and map routing indices per request for MoE router recording.
+
+        This method retrieves recorded routing decisions and maps them to individual
+        requests using the context's request_ids and query_lengths. Uses the context's
+        routing_metadata when available (which handles CUDA graph static buffers automatically).
+        Must be called while context attributes are still valid (before request transitions).
+
+        Returns:
+            Optional[Dict[int, Tensor]]: A dictionary mapping request_id to a tensor of
+                shape [num_tokens, num_layers, topk]. Returns None if routing replay is
+                disabled or no routing data was recorded.
+        """
+        config = self.inference_wrapped_model.model.config
+        if not config.moe_enable_routing_replay:
+            return None
+
+        # Get routing indices - use routing_metadata if available (handles CUDA graph static buffers)
+        context = self.inference_wrapped_model.inference_context
+        if context.moe_routing_metadata is None:
+            return None
+
+        stacked_routing = context.moe_routing_metadata.get_routing_indices()
+
+        if stacked_routing is None:
+            return None
+
+        # Get active request info from context
+        active_request_slice = slice(context.paused_request_count, context.total_request_count)
+        active_request_ids = context.request_ids[active_request_slice].tolist()
+        active_query_lengths = context.request_query_lengths[active_request_slice].tolist()
+        active_token_count = context.active_token_count
+
+        # Get TP group for all-gather if using sequence parallelism
+        # With sequence parallelism, each TP rank only sees a portion of the tokens,
+        # so we need to gather routing indices across all TP ranks.
+        tp_group = self.inference_wrapped_model.tp_group
+        tp_size = get_pg_size(tp_group)
+
+        # All-gather across TP group if using sequence parallelism (tp_size > 1)
+        if tp_size > 1 and get_model_config(self.inference_wrapped_model.model).sequence_parallel:
+            # gather_from_sequence_parallel_region gathers along dim 0
+            # [local_token_count, num_layers, topk] -> [global_token_count, num_layers, topk]
+            stacked_routing = gather_from_sequence_parallel_region(stacked_routing, group=tp_group)
+
+        # Slice to real tokens (remove CUDA padding)
+        stacked_routing = stacked_routing[:active_token_count]
+
+        # Split by request along token dimension
+        # stacked_routing has shape [active_token_count, num_layers, topk]
+        routing_splits = stacked_routing.split(active_query_lengths, dim=0)
+
+        # Map to request IDs
+        routing_indices_per_request = {}
+        for req_id, routing_split in zip(active_request_ids, routing_splits):
+            # routing_split has shape [num_tokens_for_request, num_layers, topk]
+            routing_indices_per_request[req_id] = routing_split
+
+        return routing_indices_per_request
+
     def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
         """Calculate log probs from logits."""
         context = self.inference_wrapped_model.inference_context
@@ -889,8 +951,16 @@ async def async_generate_output_tokens_dynamic_batch(
             context.padded_active_request_count if context.is_decode_only() else None
         )
 
+        # Enable routing recording before forward pass if routing replay is enabled
+        config = self.inference_wrapped_model.model.config
+        if config.moe_enable_routing_replay:
+            RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+
         logits = self._dynamic_step_forward_logits(input_ids, position_ids)
 
+        # Collect routing indices per request (must be done before context transitions)
+        routing_indices_per_request = self._router_record_bookkeeping()
+
         # This is the best place to yield control back to event loop.
         # At this point we have enqueued FW pass GPU kernels asynchronously.
         # While they are running, we can do other useful CPU work.
@@ -922,6 +992,7 @@ async def async_generate_output_tokens_dynamic_batch(
             "sample": self._sampled_tokens_cuda[:active_request_count],
             "log_probs": log_probs,
             "top_n_logprobs": top_n_logprobs,
+            "routing_indices_per_request": routing_indices_per_request,
             "cuda_graph_request_count": cuda_graph_request_count,
         }
         ret.update(request_bookkeeping)