Support routing replay with asynchronous scheduling (#7310)

Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>

Author: lmcafee-nvidia

examples/inference/README.md

@@ -117,7 +117,7 @@ Inference can collect traces two ways.
 
 | Path | Enable with | Captures | CUDA graphs |
 |------|-------------|----------|------------|
-| **Sink** | `--moe-enable-routing-replay` + legacy scheduling | top-K indices only | on |
+| **Sink** | `--moe-enable-routing-replay` | top-K indices only | on |
 | **Hook** | no replay + `--cuda-graph-impl none` | indices **+ hidden states + router weights** | must be off |
 
 Only the hook path captures the hidden states and router weights that
@@ -140,14 +140,14 @@ Forward hooks do not fire during CUDA graph replay. MoE cudagraphs must be disab
 
 **Inference — sink** (routing indices only, graphs on):
 
-Routing replay requires legacy scheduling. Because async scheduling is enabled by
-default, explicitly select legacy mode when enabling the routing replay sink.
+The routing replay sink supports both legacy and async dynamic-batching schedules.
+The example below explicitly selects async scheduling.
 
 ```bash
 --moe-routing-trace-path /path/to/trace_dir
 --moe-routing-trace-max-inference-steps 200
 --moe-enable-routing-replay
---inference-dynamic-batching-async-sched-mode legacy
+--inference-dynamic-batching-async-sched-mode async
 ```
 
 **Inference — hook** (adds hidden states + weights for predictability):

examples/inference/advanced/gpt_dynamic_inference_with_coordinator.py

@@ -94,6 +94,16 @@ async def main(
         num_requests_added = 0
         next_suspend_at = args.suspend_resume_interval or 0
         cycles_done = 0
+        batch_boundaries = (
+            {int(boundary) for boundary in args.batch_boundaries.split(",")}
+            if args.drain_between_batches and args.batch_boundaries
+            else set()
+        )
+
+        async def wait_for_batch_boundary():
+            """Drain submitted requests before starting the next configured batch."""
+            if num_requests_added in batch_boundaries and futures:
+                await asyncio.gather(*futures)
 
         while True:
             current_time = time.time_ns() / 10**9
@@ -103,6 +113,7 @@ async def main(
                     num_requests_added < num_requests_total
                     and requests[num_requests_added].time_arrival <= current_time
                 ):
+                    await wait_for_batch_boundary()
                     request = requests[num_requests_added]
                     # These add-request calls will queue up the request on a zmq socket and return
                     # instantaneously. They will return an asyncio future which can be awaited for
@@ -123,6 +134,7 @@ async def main(
                 for i in range(
                     min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
                 ):
+                    await wait_for_batch_boundary()
                     # Change sampling parameters to force different generation lengths.
                     request = requests[num_requests_added]
                     n = request.sampling_params.num_tokens_to_generate

megatron/core/inference/contexts/dynamic_context.py

@@ -19,6 +19,7 @@
     InferenceBatchDimensions,
 )
 from megatron.core.inference.config import (
+    AsyncScheduleMode,
     InferenceConfig,
     KVCacheManagementMode,
     PrefixCachingEvictionPolicy,
@@ -2731,7 +2732,10 @@ def initialize_attention_state(
             )
 
         if self.moe_enable_routing_replay:
-            if self.using_cuda_graph_this_step():
+            if (
+                self.config.async_sched_mode == AsyncScheduleMode.ASYNC
+                or self.using_cuda_graph_this_step()
+            ):
                 self.moe_routing_metadata.enable_static_buffer_recording()
             else:
                 self.moe_routing_metadata.disable_static_buffer_recording()

megatron/core/inference/contexts/kv_block_allocator.py

@@ -641,42 +641,39 @@ def evict_lru_blocks(self, num_blocks_needed: int) -> bool:
     # Per-block routing storage methods (for MoE routing replay)
     # =========================================================================
 
-    def store_routing_per_block(self, flat_routing: Optional[np.ndarray]) -> None:
+    def store_routing_per_block(
+        self, flat_routing: np.ndarray, block_ids: np.ndarray, positions: np.ndarray
+    ) -> None:
         """Scatter flat routing indices into per-block storage.
 
-        Uses the context's token-to-block mapping to distribute each token's
-        routing data into the appropriate block. Matched (prefix-cached) blocks
-        already have routing from the original request and are not overwritten
-        here since their tokens are not in the active token layout.
+        Matched prefix-cache blocks already contain routing from the original
+        request and are absent from this forward's token layout.
 
         Args:
             flat_routing: ndarray of shape [active_token_count, num_layers, topk]
-                aligned with the context's active-token layout, or None.
+                aligned with the supplied token layout.
+            block_ids: KV-block ID for each routing row.
+            positions: Position within the KV block for each routing row.
         """
-        if flat_routing is None:
-            return
-
-        context = self.context
-        token_count = context.active_token_count
+        token_count = flat_routing.shape[0]
         if token_count == 0:
             return
 
-        assert (
-            flat_routing.shape[0] == token_count
-        ), f"Routing token count {flat_routing.shape[0]} != active token count {token_count}"
-
-        # Token-to-block mapping for all active tokens
-        block_ids_np = context.token_to_block_idx[:token_count].cpu().numpy()
-        positions_np = context.token_to_local_position_within_kv_block[:token_count].cpu().numpy()
+        assert block_ids.shape == (
+            token_count,
+        ), f"Routing block IDs shape {block_ids.shape} != ({token_count},)"
+        assert positions.shape == (
+            token_count,
+        ), f"Routing positions shape {positions.shape} != ({token_count},)"
 
         dummy = self.dummy_block_idx
 
         # Group tokens by block_id using sort for efficient scatter
         unique_blocks, inverse, counts = np.unique(
-            block_ids_np, return_inverse=True, return_counts=True
+            block_ids, return_inverse=True, return_counts=True
         )
         sorted_indices = np.argsort(inverse, kind='stable')
-        sorted_positions = positions_np[sorted_indices]
+        sorted_positions = positions[sorted_indices]
         sorted_routing = flat_routing[sorted_indices]
 
         offset = 0
@@ -707,19 +704,20 @@ def reconstruct_routing_from_blocks(
 
         Returns:
             ndarray [total_routing_tokens, num_layers, topk] or None if any
-            block is missing routing data.
+            required block is missing routing data.
         """
         block_size = self.context.block_size_tokens
         routing_parts = []
         tokens_collected = 0
 
         for bid in block_ids:
-            routing = self.get_block_routing(bid)
-            if routing is None:
-                return None  # Missing routing data for this block
+            # Speculative preparation can add a trailing block the result does not need.
             remaining = total_routing_tokens - tokens_collected
             if remaining <= 0:
                 break
+            routing = self.get_block_routing(bid)
+            if routing is None:
+                return None  # Missing routing data for this block
             take = min(block_size, remaining)
             routing_parts.append(routing[:take])
             tokens_collected += take

megatron/core/inference/contexts/routing_metadata.py

@@ -1,9 +1,11 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, cast
 
 import torch
 
+from megatron.core.inference.config import AsyncScheduleMode
+
 if TYPE_CHECKING:
     from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 
@@ -23,6 +25,12 @@ class RoutingMetadata:
     """
 
     def __init__(self, context: 'DynamicInferenceContext', moe_router_topk: int):
+        """Initialize lazy routing storage.
+
+        Args:
+            context (DynamicInferenceContext): Context owning the recorded token layout.
+            moe_router_topk (int): Number of selected experts per token.
+        """
         self.context = context
         self.max_tokens = context.max_tokens
         self.moe_router_topk = moe_router_topk
@@ -31,6 +39,8 @@ def __init__(self, context: 'DynamicInferenceContext', moe_router_topk: int):
         # 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.buffer_index_cuda: Optional[torch.Tensor] = None
+        self.layer_indices_cuda: Optional[list[torch.Tensor]] = None
         self.num_moe_layers: Optional[int] = None
 
     def _ensure_buffer_allocated(self) -> None:
@@ -41,29 +51,40 @@ def _ensure_buffer_allocated(self) -> None:
         if self.routing_indices_buffer is not None:
             return
 
-        self.num_moe_layers = len(RouterReplay.global_router_replay_instances)
+        self.num_moe_layers = len(RouterReplay.get_instances(is_mtp_layer=False))
 
         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]:
+        shape = (self.max_tokens, self.num_moe_layers, self.moe_router_topk)
+        if self.context.config.async_sched_mode == AsyncScheduleMode.ASYNC:
+            shape = (2, *shape)
+            self.buffer_index_cuda = torch.zeros(1, dtype=torch.int64, device=self.device)
+            # Separate scalars keep every compiled input aligned, unlike slices of arange.
+            self.layer_indices_cuda = [
+                torch.tensor([i], dtype=torch.int64, device=self.device)
+                for i in range(self.num_moe_layers)
+            ]
+        self.routing_indices_buffer = torch.empty(shape, dtype=torch.int32, device=self.device)
+
+    def get_routing_indices(self, buffer_index: int | None = None) -> Optional[torch.Tensor]:
         """Get the recorded routing indices.
 
-        Automatically uses the static buffer when CUDA graphs are active,
-        otherwise retrieves from RouterReplay utility.
+        Async scheduling reads the selected bank in both eager and graphed execution.
+        Legacy execution reads static graph storage or stacks eager router outputs.
+
+        Args:
+            buffer_index (int | None): Bank written by the async base forward, or
+                ``None`` for legacy recording.
 
         Returns:
             Tensor of shape [num_tokens, num_moe_layers, topk] or None if not available.
         """
-        if self.context.using_cuda_graph_this_step():
+        if buffer_index is not None:
+            if self.routing_indices_buffer is None:
+                return None
+            return self.routing_indices_buffer[buffer_index, : self.context.active_token_count]
+        elif 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
@@ -72,13 +93,13 @@ def get_routing_indices(self) -> Optional[torch.Tensor]:
             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()
+            recorded_data = RouterReplay.get_recorded_data(is_mtp_layer=False)
             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)
+            return torch.stack(cast(list[torch.Tensor], recorded_data), dim=1)
 
     def enable_static_buffer_recording(self) -> None:
         """Enable recording into the static buffer for CUDA graph compatibility.
@@ -89,8 +110,13 @@ def enable_static_buffer_recording(self) -> None:
         """
         self._ensure_buffer_allocated()
         if self.routing_indices_buffer is not None:
-            RouterReplay.set_global_static_buffers(self.routing_indices_buffer)
+            RouterReplay.set_global_static_buffers(
+                self.routing_indices_buffer,
+                is_mtp_layer=False,
+                buffer_index=self.buffer_index_cuda,
+                layer_indices=self.layer_indices_cuda,
+            )
 
     def disable_static_buffer_recording(self) -> None:
         """Disable static buffer recording, reverting to normal tensor assignment."""
-        RouterReplay.clear_global_static_buffers()
+        RouterReplay.clear_global_static_buffers(is_mtp_layer=False)