Extend dynamic inference asynchronous scheduling support

examples/inference/advanced/gpt_dynamic_inference.py

@@ -128,7 +128,10 @@ def _process_step_result(result):
         """Process a single engine step result, updating bookkeeping state."""
         nonlocal total_output_tokens, num_requests_finished
 
-        is_decode_only = engine.is_decode_only
+        decode_only = engine.decode_only
+        is_decode_only = (
+            decode_only.launched if decode_only.launched is not None else decode_only.consumed
+        )
 
         # Record cuda_graph_request_count.
         cuda_graph_request_count = result["cuda_graph_request_count"]
@@ -228,7 +231,7 @@ def _process_step_result(result):
             add_times.append(get_curr_time(do_broadcast=False) - add_start)
 
             # Step inference engine (i.e., generate a token for each active request).
-            # Before step, we haven't done the scheduling, so we cannot know the is_decode_only
+            # The engine reports the consumed and launched decode-only states after scheduling.
             try:
                 result = engine.step_modern()
             except EngineSuspendedError as e:

megatron/core/inference/config.py

@@ -140,11 +140,8 @@ class AsyncScheduleMode(str, Enum):
     LEGACY = "legacy"
     """Resolve requests before preparing the next forward pass."""
 
-    SERIAL = "serial"
-    """Prepare and forward speculatively before resolving the sampled requests."""
-
-    OVERLAP = "overlap"
-    """Overlap async scheduling prepare/sample and forward/resolve phases."""
+    ASYNC = "async"
+    """Overlap asynchronous scheduling phases by reordering them to prepare-before-resolve."""
 
 
 @dataclass

megatron/core/inference/sampling/base.py

@@ -25,6 +25,7 @@ def sample_kernel(
         no_top_p: bool,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
+        output: Optional[Tensor] = None,
         eager: bool = False,
         cache_key: Any = None,
     ) -> Tensor:
@@ -41,6 +42,7 @@ def sample_kernel(
             gather_indices: If provided, only sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: Per-token request mapping; when set, sampling
                 parameters are gathered per-token instead of per-request.
+            output: Optional caller-owned destination tensor of shape `[n]`.
             eager, cache_key: Accepted for API symmetry; ignored (no CUDA graph).
 
         Returns:
@@ -63,12 +65,24 @@ def sample_speculative(
         """Sample tokens for the speculative-verify path.
 
         Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests contribute 1.
-        Builds the per-token request mapping and dispatches to `sample_kernel`.
-        The `sample_kernel` is forced eager so its own `CudaGraphManager` wrapper does not fire.
+        Builds the per-token request mapping and dispatches to the return-valued `sample_kernel`.
 
         When `gather_indices` is supplied, the kernel selects via `logits[gather_indices[:n], :]`.
         When `gather_indices` is None, `required_logits` is expected to be already pre-gathered to
         the layout described above (e.g. when `materialize_only_last_token_logits=True` upstream).
+
+        Args:
+            required_logits: Logits containing base and speculative rows.
+            num_decode: Number of decode requests.
+            num_prefill: Number of prefill requests.
+            num_speculative_tokens: Number of draft tokens per decode request.
+            context: The active DynamicInferenceContext.
+            gather_indices: Optional rows to gather from `required_logits`.
+            eager: Whether to bypass a wrapped CUDA graph.
+            cache_key: CUDA graph lookup key.
+
+        Returns:
+            Sampled token IDs for all required base and speculative rows.
         """
         # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
@@ -106,13 +120,15 @@ def sample_speculative(
 
     @abstractmethod
     def log_probs_kernel(
-        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+        self, logits: Tensor, context, *, token_to_request_index: Optional[Tensor] = None
     ) -> Tensor:
         """Per-row log-probs of the distribution this backend samples from.
 
         Args:
             logits: `[num_rows, vocab_size]` raw logits.
-            temperature, top_k, top_p: `[num_rows]` per-row sampling params.
+            context: The active DynamicInferenceContext.
+            token_to_request_index: Optional per-row request mapping. When
+                omitted, each logits row maps to the request at the same index.
 
         Returns:
             `[num_rows, vocab_size]` log-probs; filtered-out tokens are `-inf`.

megatron/core/inference/sampling/flashinfer_sampling.py

@@ -49,6 +49,7 @@ def sample_kernel(
         no_top_p: bool,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
+        output: Optional[Tensor] = None,
         eager: bool = False,
         cache_key: Any = None,
     ) -> Tensor:
@@ -65,10 +66,11 @@ def sample_kernel(
             gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: When set, sampling parameters are gathered
                 per-token rather than per-request (speculative decoding path).
+            output: Optional caller-owned destination tensor of shape `[n]`.
             eager, cache_key: Accepted for API symmetry; ignored (no CUDA graph).
 
         Returns:
-            Sampled token ids of shape `[n]`.
+            Sampled token IDs in `output`, or a newly allocated tensor when it is not provided.
         """
         del eager, cache_key
 
@@ -106,36 +108,63 @@ def sample_kernel(
             # multinomial forces a device-to-host sync, whereas sampling_from_probs
             # stays on-device and keeps the RNG's philox offset advancing per launch.
             probs = torch.softmax(scaled, dim=-1)
-            return flashinfer.sampling.sampling_from_probs(
+            sampled_tokens = flashinfer.sampling.sampling_from_probs(
                 probs, deterministic=True, generator=self._rng
             ).long()
         elif no_top_k:
             # Top-p only -> dedicated exact nucleus kernel.
             probs = torch.softmax(scaled, dim=-1)
             top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
-            return flashinfer.sampling.top_p_sampling_from_probs(
+            sampled_tokens = flashinfer.sampling.top_p_sampling_from_probs(
                 probs, top_p_safe, deterministic=True, generator=self._rng
             ).long()
         elif no_top_p:
             # Top-k only -> dedicated exact top-k kernel.
             probs = torch.softmax(scaled, dim=-1)
             top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
-            return flashinfer.sampling.top_k_sampling_from_probs(
+            sampled_tokens = flashinfer.sampling.top_k_sampling_from_probs(
                 probs, top_k_safe, deterministic=True, generator=self._rng
             ).long()
         else:
             # Mixed batch (some top-k, some top-p, or requests using both) -> joint
             # kernel, fed the temperature-scaled logits.
             top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
             top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
-            return flashinfer.sampling.top_k_top_p_sampling_from_logits(
+            sampled_tokens = flashinfer.sampling.top_k_top_p_sampling_from_logits(
                 scaled, top_k_safe, top_p_safe, deterministic=True, generator=self._rng
             ).long()
 
+        if output is None:
+            return sampled_tokens
+        output.copy_(sampled_tokens)
+        return output
+
     def log_probs_kernel(
-        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+        self, logits: Tensor, context, *, token_to_request_index: Optional[Tensor] = None
     ) -> Tensor:
-        """Per-row log-probs of the FlashInfer top-k / top-p sampling distribution."""
+        """Per-row log-probs of the FlashInfer top-k / top-p sampling distribution.
+
+        Args:
+            logits (Tensor): Raw logits with shape `[num_rows, vocab_size]`.
+            context: Active dynamic inference context providing GPU sampling metadata.
+            token_to_request_index (Optional[Tensor]): Optional mapping from each
+                logits row to its request index.
+
+        Returns:
+            Tensor: Per-row log probabilities for the processed distribution.
+        """
+        gpu_view = context.gpu_view
+        if token_to_request_index is None:
+            num_rows = logits.size(0)
+            temperature = gpu_view.temperature[:num_rows]
+            top_k = gpu_view.top_k[:num_rows]
+            top_p = gpu_view.top_p[:num_rows]
+        else:
+            token_to_request_index = token_to_request_index.to(logits.device, non_blocking=True)
+            temperature = gpu_view.temperature[token_to_request_index]
+            top_k = gpu_view.top_k[token_to_request_index]
+            top_p = gpu_view.top_p[token_to_request_index]
+
         temperature = temperature.clamp(min=1e-6)
         probs = torch.softmax(logits / temperature.unsqueeze(1), dim=-1)
 

megatron/core/inference/sampling/torch_sampling.py

@@ -114,14 +114,34 @@ def sample_from_logits(
         return sampled
 
     def log_probs_kernel(
-        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+        self, logits: Tensor, context, *, token_to_request_index: Optional[Tensor] = None
     ) -> Tensor:
         """Per-row log-probs of the temperature, top-k/top-p sampling distribution.
 
         Buckets rows by identical (temperature, top_k, top_p) and reuses `filter_logits`
-        (the same filter as `sample_from_logits`) so log-probs match how this backend
-        samples. `temperature`/`top_k`/`top_p` are per-row `[num_rows]` tensors.
+        (the same filter as `sample_from_logits`) so log-probs match how this backend samples.
+
+        Args:
+            logits (Tensor): Raw logits with shape `[num_rows, vocab_size]`.
+            context: Active dynamic inference context providing CPU sampling metadata.
+            token_to_request_index (Optional[Tensor]): Optional CPU mapping from
+                each logits row to its request index.
+
+        Returns:
+            Tensor: Per-row log probabilities for the processed distribution.
         """
+        active_request_count = context.total_request_count - context.paused_request_count
+        metadata = context.active_request_metadata
+        if token_to_request_index is None:
+            temperature = metadata["temperature"][:active_request_count]
+            top_k = metadata["top_k"][:active_request_count]
+            top_p = metadata["top_p"][:active_request_count]
+        else:
+            assert not token_to_request_index.is_cuda
+            temperature = metadata["temperature"][:active_request_count][token_to_request_index]
+            top_k = metadata["top_k"][:active_request_count][token_to_request_index]
+            top_p = metadata["top_p"][:active_request_count][token_to_request_index]
+
         temps = temperature.tolist()
         top_ks = top_k.tolist()
         top_ps = top_p.tolist()
@@ -148,10 +168,11 @@ def sample_kernel(
         no_top_p: bool,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
+        output: Optional[Tensor] = None,
         eager: bool = False,
         cache_key: Any = None,
     ) -> Tensor:
-        """Bucket active requests by `(temperature, top_k, top_p)` and sample each bucket.
+        """Bucket active requests by sampling parameters and sample each bucket.
 
         Args:
             logits: Logits tensor of shape `[>=n, vocab_size]`.
@@ -163,11 +184,12 @@ def sample_kernel(
             gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: When set, the loop dispatches per-token rather than
                 per-request (used by the speculative path).
+            output: Optional caller-owned destination tensor of shape `[n]`.
             eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).
             cache_key: Accepted for API symmetry; ignored.
 
         Returns:
-            Sampled token ids of shape `[n]`.
+            Sampled token IDs in `output`, or a newly allocated tensor when it is not provided.
         """
         del eager, cache_key, no_top_k, no_top_p
 
@@ -191,7 +213,8 @@ def sample_kernel(
         if gather_indices is not None:
             logits = logits[gather_indices[:n], :]
 
-        output = torch.empty(n, device=logits.device, dtype=torch.int64)
+        if output is None:
+            output = torch.empty(n, device=logits.device, dtype=torch.int64)
         token_list = []
         indices_list = []
         for idx_tensor, (_, temp, top_k, top_p) in zip(bucket_index_tensors, buckets):