Reorganize order of operations in inference context and text generation controller

megatron/core/inference/config.py

@@ -297,10 +297,10 @@ class InferenceConfig:
     Defaults to 0, which means no logging.
     """
 
-    request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None
+    request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None
     """
     A list of the per-request metadata types to track. Each entry is a tuple
-    consisting of the string label, the target dtype, and whether to store the data on GPU.
+    consisting of the string label and the target dtype.
     """
 
     use_synchronous_zmq_collectives: bool = False

megatron/core/inference/contexts/dynamic_context.py

@@ -844,7 +844,7 @@ def initialize_all_tensors(self) -> None:
             label: torch.empty(
                 (self.max_requests,), dtype=dtype, device=torch.cuda.current_device()
             )
-            for label, dtype, _ in self.request_metadata_types
+            for label, dtype in self.request_metadata_types
         }
 
         # Per-token state.
@@ -860,6 +860,11 @@ def initialize_all_tensors(self) -> None:
         self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids)
         self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids)
 
+        # Static tensor addresses of active slices to enable fast inference kernels.
+        self.active_request_metadata = {
+            label: torch.empty_like(tensor) for label, tensor in self.request_metadata.items()
+        }
+
         # NOTE: Need to build this outside the UVM / TMS context to avoid IMA.
         if self.is_hybrid_model:
             self.mamba_metadata = MambaMetadata(
@@ -1062,6 +1067,23 @@ def get_active_request_count(self):
         """Returns the current number of active requests."""
         return self.total_request_count - self.paused_request_count
 
+    def build_active_slices(self, batch_size: int):
+        """Build the active slices of specific tensors. This is run on every forward step.
+
+        If the context is reordered to active -> paused -> finished, this can be graphed.
+        """
+        padded_slice = slice(self.paused_request_count, self.paused_request_count + batch_size)
+
+        # Request metadata all needs to be sliced.
+        for label in self.request_metadata:
+            self.active_request_metadata[label][:batch_size].copy_(
+                self.request_metadata[label][padded_slice], non_blocking=True
+            )
+
+    def pad_active_slices(self):
+        """Pad the active slices of specific tensors."""
+        pass
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -1391,7 +1413,7 @@ def add_dummy_requests_parallel(
         self.request_output_lengths[request_slice] = lengths_tensor + tokens_to_generate_tensor
         self.request_kv_length_offsets[request_slice] = 0
         self.request_kv_block_counts[request_slice] = block_counts
-        for i, (label, dtype, _) in enumerate(self.request_metadata_types):
+        for i, (label, dtype) in enumerate(self.request_metadata_types):
             self.request_metadata[label][request_slice] = torch.tensor(
                 metadata_cols[i], dtype=dtype, device=torch.cuda.current_device()
             )
@@ -1702,6 +1724,9 @@ def initialize_attention_state(
         self.padded_active_request_count = self.padded_batch_dimensions.req_count
         self.padding_slice = slice(self.active_token_count, self.padded_active_token_count)
 
+        self.build_active_slices(self.padded_active_request_count)
+        self.pad_active_slices()
+
         # Update token position indexes.
         self.token_to_block_idx[self.active_token_count : self.padded_active_token_count] = (
             self.kv_block_allocator.dummy_block_idx
@@ -1912,6 +1937,20 @@ def speculative_required_logit_indices(self, device: torch.device) -> Tensor:
 
         return torch.cat([decode_indices, prefill_last_indices])
 
+    @property
+    def num_last_token_logits(self) -> int:
+        """Number of rows produced by `last_token_logits` for the current step.
+
+        Single source of truth for the bound: one row per request, with
+        `(num_speculative_tokens + 1)` rows per decode request when MTP is active.
+        """
+        if self.num_speculative_tokens > 0:
+            return (
+                self.num_decode_requests * (self.num_speculative_tokens + 1)
+                + self.num_prefill_requests
+            )
+        return self.total_request_count - self.paused_request_count
+
     def last_token_logits(self, logits: Tensor) -> Tensor:
         """Select the logit positions needed for token generation.
 
@@ -1925,7 +1964,7 @@ def last_token_logits(self, logits: Tensor) -> Tensor:
             logits (Tensor): Output logits of forward pass, shape [1, S, H].
 
         Return:
-            (Tensor) Selected logits, shape [N, H].
+            (Tensor) Selected logits, shape [N, H], where N == num_last_token_logits.
         """
         # todo: @lmcafee, remove these asserts?
         assert logits.size(0) == 1, f"logits.size(0) ({tuple(logits.shape)}) != 1"
@@ -1937,12 +1976,14 @@ def last_token_logits(self, logits: Tensor) -> Tensor:
 
         if self.num_speculative_tokens > 0:
             selected = self.speculative_required_logit_indices(logits.device)
+            assert selected.numel() == self.num_last_token_logits
             return logits_2d[selected, :]
 
         paused = self.paused_request_count
         total = self.total_request_count
         query_lengths = self.request_query_lengths[paused:total]
         last_token_idxs = torch.cumsum(query_lengths, dim=0) - 1
+        assert last_token_idxs.numel() == self.num_last_token_logits
         return logits_2d[last_token_idxs, :]
 
     def _compute_prefix_match(
@@ -2183,7 +2224,7 @@ def add_request(
         metadata = req.tracked_metadata
         metadata_types = req.get_metadata_types()
         for m, m_type in zip(metadata, metadata_types):
-            label, _, _ = m_type
+            label, _ = m_type
             if not isinstance(m, torch.Tensor):
                 m = torch.as_tensor(
                     m,

megatron/core/inference/engines/dynamic_engine.py

@@ -400,39 +400,39 @@ def create_cuda_graphs(self, reset_context: bool = True):
 
             # 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)
-
-            # MTP CUDA graph warmup for this batch dimension.
-            if mtp_warmup_enabled:
-                n = cuda_graph_batch_dimension.req_count
-                # pylint: disable-next=possibly-used-before-assignment
-                if sp_enabled:
-                    n = round_up_to_nearest_multiple(n, tp_size)
-                # pylint: disable-next=possibly-used-before-assignment
-                if n > 0 and n not in mtp_seen_batch_sizes:
-                    mtp_seen_batch_sizes.add(n)
-                    device = torch.cuda.current_device()
-                    batch_dim = n // tp_size if sp_enabled else n
-                    # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay.
-                    for depth in mtp_warmup_depths:
-                        unwrapped.compute_mtp_single_step(
-                            hidden_states=torch.zeros(
-                                (batch_dim, 1, model_config.hidden_size),
-                                device=device,
-                                dtype=model_config.params_dtype,
-                            ),
-                            next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long),
-                            position_ids=torch.zeros((1, n), device=device, dtype=torch.int64),
-                            depth=depth,
-                            cache_key=("mtp", n, depth),
-                        )
+            with torch.inference_mode():
+                controller._dynamic_step_forward_logits(input_ids, position_ids)
+
+                # MTP CUDA graph warmup for this batch dimension.
+                if mtp_warmup_enabled:
+                    n = cuda_graph_batch_dimension.req_count
+                    # pylint: disable-next=possibly-used-before-assignment
+                    if sp_enabled:
+                        n = round_up_to_nearest_multiple(n, tp_size)
+                    # pylint: disable-next=possibly-used-before-assignment
+                    if n > 0 and n not in mtp_seen_batch_sizes:
+                        mtp_seen_batch_sizes.add(n)
+                        device = torch.cuda.current_device()
+                        batch_dim = n // tp_size if sp_enabled else n
+                        # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay.
+                        for depth in mtp_warmup_depths:
+                            unwrapped.compute_mtp_single_step(
+                                hidden_states=torch.zeros(
+                                    (batch_dim, 1, model_config.hidden_size),
+                                    device=device,
+                                    dtype=model_config.params_dtype,
+                                ),
+                                next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long),
+                                position_ids=torch.zeros((1, n), device=device, dtype=torch.int64),
+                                depth=depth,
+                                cache_key=("mtp", n, depth),
+                            )
 
-            context.reset()
+                context.reset()
 
         # Disable inference dispatcher after graph capture
         if is_inference_optimized_ep:

megatron/core/inference/inference_request.py

@@ -476,26 +476,25 @@ def tracked_metadata(self) -> List[Any]:
                     "in its sampling_params. Defaulting to -1."
                 )
             sp.termination_id = -1
-        return [getattr(sp, field) for field, _, _ in self.get_metadata_types()]
+        return [getattr(sp, field) for field, _ in self.get_metadata_types()]
 
     @staticmethod
-    def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
-        """Keeps track of all request metadata names, dtypes, and target device.
+    def get_metadata_types() -> List[Tuple[str, torch.dtype]]:
+        """Keeps track of all request metadata names and dtypes.
 
         Returns:
-            List[Tuple[str, torch.dtype, bool]]: Mapping from metadata name to:
+            List[Tuple[str, torch.dtype]]: Mapping from metadata name to:
                 name (str) - The name of the metadata field.
                 dtype (torch.dtype) - The datatype of the metadata.
-                on_device (bool) - Whether the metadata lives on GPU (True) or CPU (False).
         """
         return [
-            ("temperature", torch.float32, False),  # CPU for torch sampling
-            ("top_k", torch.int32, False),  # CPU for torch sampling
-            ("top_p", torch.float32, False),  # CPU for torch sampling
-            ("termination_id", torch.int64, True),
-            ("return_log_probs", torch.bool, False),  # CPU for non-selective logprobs
-            ("skip_prompt_log_probs", torch.bool, False),  # CPU for non-selective logprobs
-            ("top_n_logprobs", torch.int32, False),  # CPU for torch sampling
+            ("temperature", torch.float32),
+            ("top_k", torch.int32),
+            ("top_p", torch.float32),
+            ("termination_id", torch.int64),
+            ("return_log_probs", torch.bool),
+            ("skip_prompt_log_probs", torch.bool),
+            ("top_n_logprobs", torch.int32),
         ]
 
     def add_event(