Support async scheduling for prefill transitions

megatron/core/inference/contexts/dynamic_context.py

@@ -3587,26 +3587,27 @@ def commit_sampled_tokens(self, sampled_tokens_cpu: Tensor) -> None:
 
         self.token_to_input_ids[:active_request_count] = sampled_tokens_cpu
 
-    def resolve_requests(self, active_requests_mask: Tensor) -> Tensor:
+    def resolve_requests(self, active_requests_mask: Tensor) -> Tuple[Tensor, Tensor]:
         """Resolve finished requests after an async scheduling forward pass.
 
-        Async scheduling supports only request completion. The active request rows
-        and current decode-token rows are compacted in survivor order so any
-        following legacy or async scheduling step sees a consistent context.
+        Prefill requests transition to decode during resolution. Request rows use
+        the same hole-filling order as ``update_requests`` so seeded sampling stays
+        consistent with legacy scheduling. Decode token rows are moved in the same
+        order when prepare has already built the successor input; prefill token rows
+        are left untouched because prepare rebuilds them after resolution.
 
         Args:
             active_requests_mask (Tensor): 1D mask marking requests that remain active.
 
         Returns:
-            Tensor: Request IDs for requests that finished during resolution.
+            Tuple[Tensor, Tensor]: Request IDs that finished and source row indices
+                for surviving requests in their resolved destination order.
         """
         if active_requests_mask.is_cuda:
             active_requests_mask = active_requests_mask.cpu()
 
         if self.num_speculative_tokens != 0:
             raise RuntimeError("Async scheduling does not support speculative tokens.")
-        if self.num_prefill_requests != 0:
-            raise RuntimeError("Async scheduling only supports decode-only steps.")
         if self.paused_request_count != 0:
             raise RuntimeError("Async scheduling does not support paused requests.")
 
@@ -3617,22 +3618,36 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tensor:
                 f"got {active_requests_mask.numel()}."
             )
 
-        survivor_idxs = torch.nonzero(active_requests_mask == 1, as_tuple=True)[0]
+        had_prefill_requests = self.num_prefill_requests != 0
+        self.num_prefill_requests = 0
+        self.request_in_prefill_status_tensor[self.request_in_prefill_status_tensor == 1] = 0
+
         finished_idxs = torch.nonzero(active_requests_mask == 0, as_tuple=True)[0]
         finished_request_ids = self.request_ids[finished_idxs].clone()
 
+        active_request_count = int(active_requests_mask.sum().item())
+        survivor_idxs = torch.arange(active_request_count, device='cpu')
+        finished_idxs_on_left = torch.nonzero(
+            active_requests_mask[:active_request_count] == 0, as_tuple=True
+        )[0]
+        active_idxs_on_right = (
+            torch.nonzero(active_requests_mask[active_request_count:] == 1, as_tuple=True)[0]
+            + active_request_count
+        )
+        assert finished_idxs_on_left.numel() == active_idxs_on_right.numel()
+        survivor_idxs[finished_idxs_on_left] = active_idxs_on_right
+
         self.reset_attention_state()
 
         if finished_idxs.numel() > 0:
             self.release_memory_blocks_from_request_indexes(finished_idxs)
 
-        active_request_count = survivor_idxs.numel()
         if active_request_count == 0:
             self.request_to_kv_block_ids.fill_(-1)
             self.total_request_count = 0
             self.active_token_count = 0
             self.reset_mamba_state()
-            return finished_request_ids
+            return finished_request_ids, survivor_idxs
 
         dst_idxs = torch.arange(active_request_count, device='cpu')
         if not torch.equal(survivor_idxs, dst_idxs):
@@ -3652,22 +3667,24 @@ def resolve_requests(self, active_requests_mask: Tensor) -> Tensor:
             for metadata_tensor in self.request_metadata.values():
                 metadata_tensor[dst_idxs] = metadata_tensor[survivor_idxs]
 
-            self.token_to_input_ids[dst_idxs] = self.token_to_input_ids[survivor_idxs]
-            self.token_to_pos_ids[dst_idxs] = self.token_to_pos_ids[survivor_idxs]
-            self.token_to_block_idx[dst_idxs] = self.token_to_block_idx[survivor_idxs]
-            self.token_to_local_position_within_kv_block[dst_idxs] = (
-                self.token_to_local_position_within_kv_block[survivor_idxs]
-            )
-            self.token_to_position_in_request[dst_idxs] = self.token_to_position_in_request[
-                survivor_idxs
-            ]
+            if not had_prefill_requests:
+                self.token_to_input_ids[dst_idxs] = self.token_to_input_ids[survivor_idxs]
+                self.token_to_pos_ids[dst_idxs] = self.token_to_pos_ids[survivor_idxs]
+                self.token_to_block_idx[dst_idxs] = self.token_to_block_idx[survivor_idxs]
+                self.token_to_local_position_within_kv_block[dst_idxs] = (
+                    self.token_to_local_position_within_kv_block[survivor_idxs]
+                )
+                self.token_to_position_in_request[dst_idxs] = self.token_to_position_in_request[
+                    survivor_idxs
+                ]
 
-        self.token_to_request_idx[:active_request_count] = dst_idxs
+        if not had_prefill_requests:
+            self.token_to_request_idx[:active_request_count] = dst_idxs
         stale_slice = slice(active_request_count, old_active_request_count)
         self.request_to_kv_block_ids[stale_slice] = -1
         self.total_request_count = active_request_count
-        self.active_token_count = active_request_count
-        return finished_request_ids
+        self.active_token_count = 0 if had_prefill_requests else active_request_count
+        return finished_request_ids, survivor_idxs
 
     def update_requests(
         self,

megatron/core/inference/engines/dynamic_engine.py

@@ -1002,6 +1002,8 @@ def _validate_async_sched_support_for_config(self) -> None:
             return
 
         model_config = self.controller.inference_wrapped_model.model.config
+        if self.enable_chunked_prefill:
+            raise ValueError("Async scheduling does not support chunked prefill.")
         if self.num_speculative_tokens > 0:
             raise ValueError("Async scheduling does not support speculative tokens.")
         if self.context.is_hybrid_model:
@@ -1608,14 +1610,30 @@ def get_prefix_coordination_metrics(self) -> dict:
         """
         return {"waits": self._prefix_coordination_waits}
 
-    def schedule_waiting_requests(self):
-        """Tries to schedule any requests in the waiting pool."""
+    def _should_defer_async_sched_admission(self) -> bool:
+        """Return whether admission must wait for pending async logits.
+
+        Returns:
+            bool: Whether a ready request must remain queued for one drain step.
+        """
+        return (
+            self.context.config.async_sched_mode != AsyncScheduleMode.LEGACY
+            and self.controller.has_pending_async_forward()
+        )
+
+    def schedule_waiting_requests(self) -> bool:
+        """Try to schedule requests from the waiting pool.
+
+        Returns:
+            bool: Whether a ready request remained queued for one drain step.
+        """
         # Keep track of which requests get scheduled.
         waiting_before = set(self.waiting_request_ids)
         if self.enable_chunked_prefill:
             self.schedule_chunked_prefill()
+            admission_deferred = False
         else:
-            self.schedule_non_chunked_prefill()
+            admission_deferred = self.schedule_non_chunked_prefill()
         waiting_after = set(self.waiting_request_ids)
 
         # Re-stamp kv_cache_epoch on requests that were just scheduled.
@@ -1625,11 +1643,16 @@ def schedule_waiting_requests(self):
                 if req.kv_cache_epoch is None:
                     req.kv_cache_epoch = [(0, self._generation_epoch)]
 
-    def schedule_non_chunked_prefill(self):
-        """
-        Perform the same original scheduling logic for non-chunked runs
+        return admission_deferred
+
+    def schedule_non_chunked_prefill(self) -> bool:
+        """Schedule non-chunked prefill requests.
+
+        Returns:
+            bool: Whether a ready request remained queued for one drain step.
         """
         prefix_caching_enabled = self.context.enable_prefix_caching
+        admission_deferred = False
         if prefix_caching_enabled:
             pending_block_hashes = set()
             pending_request_ids = []
@@ -1666,6 +1689,10 @@ def schedule_non_chunked_prefill(self):
                     if not self._cg_admission_check(req, candidate):
                         break
 
+                if self._should_defer_async_sched_admission():
+                    admission_deferred = True
+                    break
+
                 # Add these hashes to pending.
                 if prefix_caching_enabled:
                     for block_hash in req.precomputed_block_hashes:
@@ -1685,6 +1712,8 @@ def schedule_non_chunked_prefill(self):
         if prefix_caching_enabled and pending_request_ids:
             self.waiting_request_ids.extendleft(reversed(pending_request_ids))
 
+        return admission_deferred
+
     def _cg_admission_gating_active(self) -> bool:
         """Cudagraph-aware admission gating is active when --inference-cuda-graph-all-prefills
         is set, the engine has prefill/mixed CGs, and the batch-dim list is populated.
@@ -1934,8 +1963,8 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
         if self.state in (EngineState.SUSPENDED, EngineState.SUSPENDING):
             raise EngineSuspendedError(self.context.step_count)
 
-        # schedule requests
-        self.schedule_waiting_requests()
+        # Schedule requests, or leave a ready admission queued for one drain step.
+        admission_deferred = self.schedule_waiting_requests()
 
         # The print block (async_bookkeep) and metrics block both fire on this
         # condition after step_count is incremented. Predict it up-front so we
@@ -1974,11 +2003,21 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
             self.step_start_event.record()
         while True:
             controller_result: DynamicBatchControllerStepResult = (
-                await self.controller.async_generate_output_tokens_dynamic_batch()
+                await self.controller.async_generate_output_tokens_dynamic_batch(
+                    drain_pending_forward=admission_deferred
+                )
             )
             if not controller_result.primer_only:
                 result = controller_result.output
                 break
+
+        if admission_deferred:
+            # Admit against the resolved batch, then leave its mixed forward pending.
+            assert not self.schedule_waiting_requests(), "Async admission remained deferred."
+            primer_result = await self.controller.async_generate_output_tokens_dynamic_batch()
+            assert (
+                primer_result.primer_only or primer_result.output is None
+            ), "Async admission may only launch a forward primer."
         if will_log_this_step:
             self.step_end_event.record()
             self.step_end_event.synchronize()