From c2f332dc2a29f92cd950b9ad1d09536d1523eb44 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 18 Dec 2025 04:56:10 -0600
Subject: [PATCH 01/19] Reorganize code to slice tensors in context class

---
 .../inference/contexts/dynamic_context.py     | 22 +++++++
 .../text_generation_controller.py             | 63 +++++++------------
 .../test_text_generation_controller.py        | 22 +++----
 3 files changed, 51 insertions(+), 56 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 5b264b36302..66f28dbe898 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -860,6 +860,17 @@ 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: Dict[str, Tensor] = {}
+        for label, _, on_gpu in self.request_metadata_types:
+            if on_gpu:
+                tensor = torch.empty_like(self.request_metadata[label])
+            else:
+                tensor = torch.empty_like(
+                    self.request_metadata[label], device="cpu", pin_memory=True
+                )
+            self.active_request_metadata[label] = tensor
+
         # NOTE: Need to build this outside the UVM / TMS context to avoid IMA.
         if self.is_hybrid_model:
             self.mamba_metadata = MambaMetadata(
@@ -1062,6 +1073,17 @@ 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):
+        """Build the active slices of specific tensors. This is run on every forward step."""
+        active_slice = slice(self.paused_request_count, self.total_request_count)
+        batch_size = self.total_request_count - self.paused_request_count
+
+        # Request metadata all needs to be sliced.
+        for label, _, _ in self.request_metadata_types:
+            self.active_request_metadata[label][:batch_size].copy_(
+                self.request_metadata[label][active_slice], non_blocking=True
+            )
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 993d05afbe0..c83b51695d5 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -148,16 +148,6 @@ def _init_dynamic_sampling_tensors(self):
         self._sampling_backend = "torch"
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
-        # Keep track of request metadata.
-        self._request_metadata: Dict[str, Tensor] = {}
-        for label, dtype, on_gpu in context.request_metadata_types:
-            tensor = context.request_metadata[label]
-            if not on_gpu:
-                # Create pinned tensors for request metadata that lives on CPU.
-                # This is metadata which requires D2H copies, such as top_k for torch sampling.
-                tensor = torch.empty_like(tensor, device="cpu", pin_memory=True)
-            self._request_metadata[label] = tensor
-
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
             self._torch_sampling_buckets: List[Tuple] = []
@@ -600,12 +590,14 @@ def _dynamic_step_context_init(
             position_ids (Tensor): The active position IDs.
         """
         context = self.inference_wrapped_model.inference_context
-        active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Remove Float16Module wrapper if it exists
         unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
         model_config = get_model_config(unwrapped_model)
 
+        # Build active slices of all relevant tensors.
+        context.build_active_slices()
+
         # Initialize attention state.
         context.initialize_attention_state(
             construct_graph_dimensions=construct_graph_dimensions,
@@ -660,14 +652,6 @@ def _dynamic_step_context_init(
                 # Turn off symmetric all reduces for prefill
                 unwrapped_model.set_symmetric_ar(None)
 
-        # Get request metadata for this step.
-        for label, dtype, on_gpu in context.request_metadata_types:
-            if not on_gpu:
-                # We need a D2H copy from the context to the pinned memory buffer.
-                self._request_metadata[label].copy_(
-                    context.request_metadata[label], non_blocking=True
-                )
-
         # Get flat tokens, position ids.
         # If we are running a dummy forward step we want to use the token count agreed upon
         # by all EP ranks rather than the minimum number of tokens.
@@ -729,7 +713,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
     def _dynamic_step_sample_bookkeeping(self):
         """Perform bookkeeping necessary to sample logits for dynamic batching."""
         context = self.inference_wrapped_model.inference_context
-        active_request_slice = slice(context.paused_request_count, context.total_request_count)
+        active_request_count = context.total_request_count - context.paused_request_count
 
         if self._sampling_backend == "torch":
             # Bucketize the core sampling parameters.
@@ -737,9 +721,9 @@ def _dynamic_step_sample_bookkeeping(self):
             bucket_map = defaultdict(list)
 
             # Shorthands for the dictionary comprehension.
-            temp = self._request_metadata["temperature"][active_request_slice].tolist()
-            top_k = self._request_metadata["top_k"][active_request_slice].tolist()
-            top_p = self._request_metadata["top_p"][active_request_slice].tolist()
+            temp = context.active_request_metadata["temperature"][:active_request_count].tolist()
+            top_k = context.active_request_metadata["top_k"][:active_request_count].tolist()
+            top_p = context.active_request_metadata["top_p"][:active_request_count].tolist()
 
             for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
                 sampling_params = (t, k, p)
@@ -1187,12 +1171,12 @@ def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
             return_log_probs (bool): Whether to return the sampled log_probs.
         """
         context = self.inference_wrapped_model.inference_context
-        active_request_slice = slice(context.paused_request_count, context.total_request_count)
-
-        return_log_probs = self._request_metadata["return_log_probs"][active_request_slice]
-        top_n_log_probs = self._request_metadata["top_n_logprobs"][active_request_slice] > 0
+        active_request_count = context.total_request_count - context.paused_request_count
 
-        return return_log_probs.any(), top_n_log_probs.any()
+        return (
+            (context.active_request_metadata["return_log_probs"][:active_request_count]).any(),
+            (context.active_request_metadata["top_n_logprobs"][:active_request_count] > 0).any(),
+        )
 
     def _router_record_bookkeeping(self) -> Optional[np.ndarray]:
         """Collect flat routing indices for MoE router recording.
@@ -1397,7 +1381,6 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
         """
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
-        active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
             context.paused_request_count : context.total_request_count
@@ -1417,7 +1400,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
                 num_decode_requests, self.num_speculative_tokens + 1, -1
             )
             accepted_counts = self._accepted_token_counts_per_request[:num_decode_requests]
-            top_n_per_request = self._request_metadata["top_n_logprobs"][active_request_slice][
+            top_n_per_request = context.active_request_metadata["top_n_logprobs"][
                 :num_decode_requests
             ]
             max_top_n = int(top_n_per_request.max().item())
@@ -1446,8 +1429,8 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
             prefill_log_probs = log_probs_tensor[decode_len:]
 
             # Batch metadata reads: single CPU transfer for all prefill requests.
-            prefill_top_n = self._request_metadata["top_n_logprobs"][active_request_slice][
-                num_decode_requests:
+            prefill_top_n = context.active_request_metadata["top_n_logprobs"][
+                num_decode_requests:active_request_count
             ].tolist()
             max_top_n_prefill = int(max(prefill_top_n)) if prefill_top_n else 0
 
@@ -1474,7 +1457,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
                     prefill_log_probs_per_request = prefill_log_probs.split(
                         prefill_query_lengths.tolist(), dim=0
                     )
-                    prefill_skip_prompt = self._request_metadata["skip_prompt_log_probs"][
+                    prefill_skip_prompt = context.active_request_metadata["skip_prompt_log_probs"][
                         num_decode_requests:active_request_count
                     ].tolist()
 
@@ -1533,9 +1516,7 @@ def _dynamic_step_calculate_top_n_logprobs(
 
             top_n_results = {}
             for req_idx in range(active_request_count):
-                top_n = int(
-                    self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item()
-                )
+                top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item())
                 if top_n > 0:
                     # Get top-n logprobs and indices for this request (single token)
                     top_n_logits = torch.topk(log_probs[req_idx], k=top_n)
@@ -1557,14 +1538,14 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         top_n_results = {}
         for req_idx in range(active_request_count):
-            top_n = int(
-                self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item()
-            )
+            top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item())
             if top_n > 0:
                 request_log_probs = log_probs_per_request[
                     req_idx
                 ]  # [num_tokens_for_request, vocab_size]
-                skip_prompt = bool(self._request_metadata["skip_prompt_log_probs"][req_idx].item())
+                skip_prompt = bool(
+                    context.active_request_metadata["skip_prompt_log_probs"][req_idx].item()
+                )
 
                 # If skip_prompt_log_probs is True, only compute for last token
                 if skip_prompt and request_log_probs.size(0) > 1:
@@ -1746,7 +1727,7 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         # Note: termination_id tensor has per-request termination IDs from mixed sampling
         active_request_mask = (
             self._sampled_tokens_cuda[:active_request_count]
-            != self._request_metadata["termination_id"][active_request_slice]
+            != context.active_request_metadata["termination_id"][:active_request_count]
         ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte()
 
         # Mark requests as finished if they hit stop words (detected in previous step's post_process_requests)
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index dd4764ee92d..a5e36a501dc 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -315,12 +315,9 @@ def test_sample_from_dynamic_logits(
         temp_values = torch.Tensor([s.temperature for s in rev_sampling_dict])
         top_k_values = torch.Tensor([s.top_k for s in rev_sampling_dict]).to(torch.int32)
         top_p_values = torch.Tensor([s.top_p for s in rev_sampling_dict])
-        request_metadata = {
-            "temperature": temp_values,
-            "top_k": top_k_values,
-            "top_p": top_p_values,
-        }
-        self.text_generation_controller._request_metadata = request_metadata
+        context.active_request_metadata["temperature"][:batch_size].copy_(temp_values)
+        context.active_request_metadata["top_k"][:batch_size].copy_(top_k_values)
+        context.active_request_metadata["top_p"][:batch_size].copy_(top_p_values)
         self.text_generation_controller._sampling_backend = backend
 
         context.padded_active_token_count = batch_size
@@ -857,15 +854,10 @@ def test_dynamic_top_n_logprobs_calculation(
 
         # Prepare sampling params
         top_n = 5
-        request_metadata = {
-            "top_n_logprobs": torch.full((batch_size,), top_n, dtype=torch.int32).cuda(),
-            "skip_prompt_log_probs": torch.full(
-                (batch_size,), float(skip_prompt_log_probs), dtype=torch.float32
-            ).cuda(),
-        }
-        self.text_generation_controller._request_metadata = request_metadata
-        self.text_generation_controller._active_request_count = batch_size
-        self.text_generation_controller._active_request_slice = slice(0, batch_size)
+        context.active_request_metadata["top_n_logprobs"][:batch_size].fill_(top_n)
+        context.active_request_metadata["skip_prompt_log_probs"][:batch_size].fill_(
+            skip_prompt_log_probs
+        )
 
         if materialize_only_last_token_logits:
             # Decode mode: logits for last tokens only

From 9087f5504ff388f9b7546b903c0c353ff499fa85 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 18 Dec 2025 05:28:42 -0600
Subject: [PATCH 02/19] Slice additional tensors in context class

---
 .../inference/contexts/dynamic_context.py     | 75 ++++++++++++-------
 .../text_generation_controller.py             | 41 +++++-----
 .../test_text_generation_controller.py        | 11 +++
 3 files changed, 78 insertions(+), 49 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 66f28dbe898..769cb047295 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -871,8 +871,19 @@ def initialize_all_tensors(self) -> None:
                 )
             self.active_request_metadata[label] = tensor
 
+        self.active_request_ids = torch.empty_like(self.request_ids, dtype=torch.int64)
+        self.active_request_query_lengths = torch.empty_like(self.request_query_lengths)
+        self.active_request_output_lengths = torch.empty_like(self.request_output_lengths)
+        self.active_request_kv_length_offsets = torch.empty_like(self.request_kv_length_offsets)
+        self.active_request_to_kv_block_ids = torch.empty_like(self.request_to_kv_block_ids)
+
+        self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
+
         # NOTE: Need to build this outside the UVM / TMS context to avoid IMA.
         if self.is_hybrid_model:
+            self.active_mamba_indices = torch.empty_like(
+                self.request_query_lengths, dtype=torch.int32
+            )
             self.mamba_metadata = MambaMetadata(
                 max_requests=self.max_requests,
                 max_tokens=self.max_tokens,
@@ -1059,24 +1070,13 @@ def cu_kv_lengths(self) -> Tuple[Tensor, Tensor, int]:
             self.active_attn_metadata["mha_metadata"].state_data["max_seqlen_k"],
         )
 
-    def get_active_sequence_lengths(self) -> Tensor:
-        """Total sequence length (query + key) for active requests."""
-        lengths = self.request_kv_length_offsets + self.request_query_lengths
-        lengths = lengths[self.paused_request_count : self.total_request_count]
-        return lengths
-
-    def get_max_sequence_lengths(self) -> Tensor:
-        """Maximum sequence length for active requests."""
-        return self.request_output_lengths[self.paused_request_count : self.total_request_count]
-
     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):
+    def build_active_slices(self, batch_size: int):
         """Build the active slices of specific tensors. This is run on every forward step."""
         active_slice = slice(self.paused_request_count, self.total_request_count)
-        batch_size = self.total_request_count - self.paused_request_count
 
         # Request metadata all needs to be sliced.
         for label, _, _ in self.request_metadata_types:
@@ -1084,6 +1084,33 @@ def build_active_slices(self):
                 self.request_metadata[label][active_slice], non_blocking=True
             )
 
+        # The following tensor slices are used in various kernels.
+        self.active_request_ids[:batch_size].copy_(self.request_ids[active_slice])
+        self.active_request_query_lengths[:batch_size].copy_(
+            self.request_query_lengths[active_slice]
+        )
+        self.active_request_output_lengths[:batch_size].copy_(
+            self.request_output_lengths[active_slice]
+        )
+        self.active_request_kv_length_offsets[:batch_size].copy_(
+            self.request_kv_length_offsets[active_slice]
+        )
+        self.active_request_to_kv_block_ids[:batch_size].copy_(
+            self.request_to_kv_block_ids[active_slice]
+        )
+
+        torch.cumsum(
+            self.active_request_query_lengths[:batch_size],
+            dim=0,
+            out=self.active_request_last_token_idxs[:batch_size],
+        )
+        self.active_request_last_token_idxs[:batch_size] -= 1
+
+        if self.is_hybrid_model:
+            self.active_mamba_indices[:batch_size].copy_(
+                self.mamba_metadata.request_to_mamba_state_idx[active_slice]
+            )
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -1741,11 +1768,9 @@ def initialize_attention_state(
             else self.non_graph_attn_metadata  # type: ignore[assignment]
         )
 
-        # Update cu_query_seq_lengths, max_seqlen_q.
-        active_slice = slice(self.paused_request_count, self.total_request_count)
-        query_lengths_view = self.request_query_lengths[active_slice]
-        request_kv_length_offsets_view = self.request_kv_length_offsets[active_slice]
-        request_to_kv_block_ids_view = self.request_to_kv_block_ids[active_slice]
+        # Build active slices of all relevant tensors.
+        batch_size = self.total_request_count - self.paused_request_count
+        self.build_active_slices(batch_size)
 
         attn_dimensions = batch_dimensions
         if self.using_cuda_graph_this_step():
@@ -1762,16 +1787,15 @@ def initialize_attention_state(
 
         assert self.active_attn_metadata is not None
         self.active_attn_metadata["mha_metadata"].update(
-            request_query_lengths=query_lengths_view,
-            request_kv_length_offsets=request_kv_length_offsets_view,
-            request_to_kv_block_ids=request_to_kv_block_ids_view,
+            request_query_lengths=self.active_request_query_lengths[:batch_size],
+            request_kv_length_offsets=self.active_request_kv_length_offsets[:batch_size],
+            request_to_kv_block_ids=self.active_request_to_kv_block_ids[:batch_size],
             batch_dimensions=attn_dimensions,
             padded_batch_dimensions=self.padded_batch_dimensions,
             num_speculative_tokens=self.num_speculative_tokens,
         )
 
         if self.is_hybrid_model:
-            active_mamba_indices_view = self.mamba_metadata.request_to_mamba_state_idx[active_slice]
             token_to_request_idx_view = self.token_to_request_idx[: self.active_token_count]
             cu_seqlens = self.active_attn_metadata["mha_metadata"].state_data[
                 "cu_query_seq_lengths"
@@ -1783,7 +1807,7 @@ def initialize_attention_state(
                     self.mamba_slot_allocator.get_intermediate_gpu_data()
                 )
             self.mamba_metadata.update(
-                active_mamba_indices_view,
+                self.active_mamba_indices[:batch_size],
                 token_to_request_idx_view,
                 cu_seqlens,
                 batch_dimensions=attn_dimensions,
@@ -1961,11 +1985,8 @@ def last_token_logits(self, logits: Tensor) -> Tensor:
             selected = self.speculative_required_logit_indices(logits.device)
             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
-        return logits_2d[last_token_idxs, :]
+        active_request_count = self.total_request_count - self.paused_request_count
+        return logits_2d[self.active_request_last_token_idxs[:active_request_count], :]
 
     def _compute_prefix_match(
         self, req: DynamicInferenceRequest, prefill_chunk_length: int
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index c83b51695d5..04d2007640a 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -595,9 +595,6 @@ def _dynamic_step_context_init(
         unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
         model_config = get_model_config(unwrapped_model)
 
-        # Build active slices of all relevant tensors.
-        context.build_active_slices()
-
         # Initialize attention state.
         context.initialize_attention_state(
             construct_graph_dimensions=construct_graph_dimensions,
@@ -1276,9 +1273,7 @@ def _dynamic_step_calculate_log_probs_speculative(
         request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
             context.paused_request_count : context.total_request_count
         ]
-        request_query_lengths = context.request_query_lengths[
-            context.paused_request_count : context.total_request_count
-        ]
+        request_query_lengths = context.active_request_query_lengths[:active_request_count]
 
         num_prefill_requests = request_in_prefill_status_tensor.sum().item()
         num_decode_requests = active_request_count - num_prefill_requests
@@ -1385,9 +1380,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
         request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
             context.paused_request_count : context.total_request_count
         ]
-        request_query_lengths = context.request_query_lengths[
-            context.paused_request_count : context.total_request_count
-        ]
+        request_query_lengths = context.active_request_query_lengths[:active_request_count]
 
         num_prefill_requests = request_in_prefill_status_tensor.sum().item()
         num_decode_requests = active_request_count - num_prefill_requests
@@ -1506,7 +1499,6 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
-        active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Handle decode-only mode (only last token)
         if context.config.materialize_only_last_token_logits or context.is_decode_only():
@@ -1530,7 +1522,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         # Note: logits may be padded, so we only take the first active_token_count tokens
         log_probs = log_probs_tensor[: context.active_token_count]
 
-        active_query_lengths = context.request_query_lengths[active_request_slice]
+        active_query_lengths = context.active_request_query_lengths[:active_request_count]
 
         # Split log_probs across request boundaries
         # log_probs has shape [active_token_count, vocab_size]
@@ -1713,24 +1705,29 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Active sequence lengths.
-        active_request_ids = context.request_ids[active_request_slice].long()
-        active_sequence_lengths = context.get_active_sequence_lengths()
-
-        # After the forward pass and KV-cache rewind, get_active_sequence_lengths()
-        # returns kv_offsets + query_lengths which already includes all accepted
-        # speculative tokens (they were part of the query and survived the rewind).
-        # Only the newly sampled base token is not yet in the KV cache, so add 1.
-        active_sequence_lengths += 1
-        max_sequence_lengths = context.get_max_sequence_lengths()
+        # After the forward pass and KV-cache rewind, kv_offsets + query_lengths
+        # already includes all accepted speculative tokens (they were part of the
+        # query and survived the rewind).  Only the newly sampled base token is not
+        # yet in the KV cache, so add 1.
+        active_sequence_lengths = (
+            context.request_kv_length_offsets[active_request_slice]
+            + context.request_query_lengths[active_request_slice]
+            + 1
+        )
 
         # Request finished if termination_id or length >= max_sequence_length.
         # Note: termination_id tensor has per-request termination IDs from mixed sampling
         active_request_mask = (
             self._sampled_tokens_cuda[:active_request_count]
             != context.active_request_metadata["termination_id"][:active_request_count]
-        ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte()
+        ).byte() & torch.less(
+            active_sequence_lengths, context.active_request_output_lengths[:active_request_count]
+        ).byte()
+
+        active_request_ids = context.active_request_ids[:active_request_count]
 
-        # Mark requests as finished if they hit stop words (detected in previous step's post_process_requests)
+        # Mark requests as finished if they hit stop words
+        # (detected in previous step's post_process_requests)
         if self._get_stop_word_finished_ids_callback is not None:
             request_ids_list = active_request_ids.tolist()
             stop_word_finished_ids = self._get_stop_word_finished_ids_callback(request_ids_list)
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index a5e36a501dc..92b9f013403 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -322,6 +322,14 @@ def test_sample_from_dynamic_logits(
 
         context.padded_active_token_count = batch_size
         context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32)
+        context.active_request_query_lengths[:batch_size].fill_(1)
+        context.active_request_last_token_idxs[:batch_size].copy_(
+            torch.arange(
+                batch_size,
+                dtype=context.active_request_last_token_idxs.dtype,
+                device=context.active_request_last_token_idxs.device,
+            )
+        )
         context.paused_request_count = 0
         context.total_request_count = batch_size
 
@@ -908,6 +916,9 @@ def test_dynamic_top_n_logprobs_calculation(
             context.request_query_lengths = torch.tensor(
                 [0] * context.paused_request_count + query_lengths, dtype=torch.int32, device='cuda'
             )
+            context.active_request_query_lengths[:batch_size].copy_(
+                torch.tensor(query_lengths, dtype=context.active_request_query_lengths.dtype)
+            )
 
             # Create logits for all tokens
             logits = torch.randn(1, total_tokens, self.vocab_size).cuda()

From 86ae997976a60e5d4e52dfdd20ebbf05bc3b7749 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 18 Dec 2025 05:44:39 -0600
Subject: [PATCH 03/19] Slice tensors by padded_active_request_count

---
 .../inference/contexts/dynamic_context.py     | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 769cb047295..7f66a2a7ffd 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1075,28 +1075,31 @@ def get_active_request_count(self):
         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."""
-        active_slice = slice(self.paused_request_count, self.total_request_count)
+        """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_types:
             self.active_request_metadata[label][:batch_size].copy_(
-                self.request_metadata[label][active_slice], non_blocking=True
+                self.request_metadata[label][padded_slice], non_blocking=True
             )
 
         # The following tensor slices are used in various kernels.
-        self.active_request_ids[:batch_size].copy_(self.request_ids[active_slice])
+        self.active_request_ids[:batch_size].copy_(self.request_ids[padded_slice])
         self.active_request_query_lengths[:batch_size].copy_(
-            self.request_query_lengths[active_slice]
+            self.request_query_lengths[padded_slice]
         )
         self.active_request_output_lengths[:batch_size].copy_(
-            self.request_output_lengths[active_slice]
+            self.request_output_lengths[padded_slice]
         )
         self.active_request_kv_length_offsets[:batch_size].copy_(
-            self.request_kv_length_offsets[active_slice]
+            self.request_kv_length_offsets[padded_slice]
         )
         self.active_request_to_kv_block_ids[:batch_size].copy_(
-            self.request_to_kv_block_ids[active_slice]
+            self.request_to_kv_block_ids[padded_slice]
         )
 
         torch.cumsum(
@@ -1108,7 +1111,7 @@ def build_active_slices(self, batch_size: int):
 
         if self.is_hybrid_model:
             self.active_mamba_indices[:batch_size].copy_(
-                self.mamba_metadata.request_to_mamba_state_idx[active_slice]
+                self.mamba_metadata.request_to_mamba_state_idx[padded_slice]
             )
 
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
@@ -1751,6 +1754,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)
+        batch_size = self.total_request_count - self.paused_request_count
+
         # 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
@@ -1768,10 +1774,6 @@ def initialize_attention_state(
             else self.non_graph_attn_metadata  # type: ignore[assignment]
         )
 
-        # Build active slices of all relevant tensors.
-        batch_size = self.total_request_count - self.paused_request_count
-        self.build_active_slices(batch_size)
-
         attn_dimensions = batch_dimensions
         if self.using_cuda_graph_this_step():
             # Treat some decode requests as prefill requests to fit the cuda graph batch dimension.

From 7f28de293e492bbcd0653d4a92e9725e00084e7d Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 25 Dec 2025 17:40:57 -0600
Subject: [PATCH 04/19] Move context tensor padding into dedicated method

---
 .../inference/contexts/dynamic_context.py     | 34 +++++++++----------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 7f66a2a7ffd..c917050cff3 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1114,6 +1114,15 @@ def build_active_slices(self, batch_size: int):
                 self.mamba_metadata.request_to_mamba_state_idx[padded_slice]
             )
 
+    def pad_active_slices(self):
+        """Pad the active slices of specific tensors."""
+        # Some tensors need to be padded at the token level.
+        padding_token_slice = slice(self.active_token_count, self.padded_active_token_count)
+
+        self.token_to_block_idx[padding_token_slice] = self.kv_block_allocator.dummy_block_idx
+        self.token_to_local_position_within_kv_block[padding_token_slice] = 0
+        self.token_to_position_in_request[padding_token_slice] = 0
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -1750,23 +1759,6 @@ def initialize_attention_state(
                 prefill_req_count=padded_prefill_req_count,
                 decode_req_count=padded_decode_req_count,
             )
-        self.padded_active_token_count = self.padded_batch_dimensions.token_count
-        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)
-        batch_size = self.total_request_count - self.paused_request_count
-
-        # 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
-        )
-        self.token_to_local_position_within_kv_block[
-            self.active_token_count : self.padded_active_token_count
-        ] = 0
-        self.token_to_position_in_request[
-            self.active_token_count : self.padded_active_token_count
-        ] = 0
 
         self.active_attn_metadata = (
             self.graph_attn_metadata  # type: ignore[assignment]
@@ -1787,6 +1779,14 @@ def initialize_attention_state(
                     decode_req_count=adjusted_decode_req_count,
                 )
 
+        self.padded_active_token_count = self.padded_batch_dimensions.token_count
+        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()
+
+        batch_size = self.total_request_count - self.paused_request_count
         assert self.active_attn_metadata is not None
         self.active_attn_metadata["mha_metadata"].update(
             request_query_lengths=self.active_request_query_lengths[:batch_size],

From aba74cdaea80529c4a1c16672fef68971641d368 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Sun, 30 Nov 2025 13:33:11 -0800
Subject: [PATCH 05/19] Store logit output in static tensor

---
 .../text_generation_controller.py             | 82 ++++++++++++-------
 .../test_text_generation_controller.py        | 17 ++--
 2 files changed, 62 insertions(+), 37 deletions(-)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 04d2007640a..6e03af947e6 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -138,6 +138,10 @@ def _init_dynamic_sampling_tensors(self):
         """Initialize tensors needed for dynamic sampling."""
         context = self.inference_wrapped_model.inference_context
         max_requests = context.max_requests
+        if context.config.materialize_only_last_token_logits:
+            max_logits = max_requests
+        else:
+            max_logits = context.max_tokens
 
         # Callback to get request IDs that should be marked as finished due to stop words
         self._get_stop_word_finished_ids_callback = None
@@ -146,6 +150,15 @@ def _init_dynamic_sampling_tensors(self):
         logits_dtype = self.inference_wrapped_model.config.params_dtype
 
         self._sampling_backend = "torch"
+        self._enable_cuda_graph = False
+
+        # Initialize bookkeeping tensors.
+        if self._enable_cuda_graph:
+            self._all_logits_cuda = torch.empty(
+                (1, max_logits, vocab_size), dtype=logits_dtype, device=device
+            )
+        else:
+            self._all_logits_cuda = None
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
         # Used for inefficient torch sampling.
@@ -659,7 +672,7 @@ def _dynamic_step_context_init(
         else:
             return context.current_input_and_position_ids()
 
-    def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) -> Tensor:
+    def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
         """Forward step the model to get logits for dynamic batching.
 
         This also handles logits-broadcasting for pipeline parallelism.
@@ -670,6 +683,11 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
         """
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
+        logits_seq_len = (
+            active_request_count
+            if context.config.materialize_only_last_token_logits
+            else context.padded_active_token_count
+        )
 
         with torch.inference_mode():
             logits = self.inference_wrapped_model.run_one_forward_step(
@@ -677,6 +695,12 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
             )
             # logits shape: [1, seq_len, vocab_size]
 
+        assert logits_seq_len == (
+            active_request_count
+            if context.config.materialize_only_last_token_logits
+            else input_ids.shape[1]
+        )
+
         # Note: When speculative decoding is active (num_speculative_tokens > 0),
         # the model skips MTP computation during the forward pass. MTP logits
         # will be computed serially after verification to ensure they are
@@ -705,7 +729,11 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
                 pp_group=self.pp_group,
             )
 
-        return logits
+        # Copy logits to contiguous buffer.
+        if self._enable_cuda_graph:
+            self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits)
+        else:
+            self._all_logits_cuda = logits
 
     def _dynamic_step_sample_bookkeeping(self):
         """Perform bookkeeping necessary to sample logits for dynamic batching."""
@@ -1025,7 +1053,7 @@ def _verify_speculative_tokens(
             num_speculative_tokens=self.num_speculative_tokens,
         )
 
-    def _dynamic_step_sample_logits_and_verify_tokens(self, logits: Tensor, input_ids: Tensor):
+    def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         """
         Sample tokens from logits for dynamic batching with speculative tokens and verify the tokens.
         """
@@ -1040,6 +1068,7 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, logits: Tensor, input_id
         # These indices are always needed for input_ids slicing and tracking
         # accepted sequence positions, even when logits are pre-sliced.
         nvtx_range_push("mtp-spec-decoding/verify/logit-indices")
+        logits = self._all_logits_cuda
         required_logit_indices = context.speculative_required_logit_indices(logits.device)
 
         if context.config.materialize_only_last_token_logits:
@@ -1120,24 +1149,21 @@ def _prepare_speculative_tokens_for_next_forward_pass(
         # Expose the active slice so downstream code sees the right length.
         self._last_accepted_seq_indices = self._last_accepted_seq_indices_buf[:active_request_count]
 
-    def _dynamic_step_sample_logits(self, logits: Tensor):
-        """Sample tokens from logits for dynamic batching.
-
-        Args:
-            logits (Tensor): The logits from the forward pass.
-        """
+    def _dynamic_step_sample_logits(self):
+        """Sample tokens from logits for dynamic batching."""
         # TODO(ksanthanam): Evaluate whether it makes more sense to sample on 1 rank
         # and then broadcast the sampled tokens rather than broadcasting the raw logits.
 
         # Last token logits.
         context = self.inference_wrapped_model.inference_context
+        active_request_count = context.total_request_count - context.paused_request_count
+
         if context.config.materialize_only_last_token_logits:
             # When materialize_only_last_token_logits is true, last_token_logits is
             # already called in the forward pass of GPT.
-            required_token_logits = logits.squeeze(0)
+            required_token_logits = self._all_logits_cuda.squeeze(0)[:active_request_count, :]
         else:
-            # todo : Should do verification here and get approrpiate las token logits
-            required_token_logits = context.last_token_logits(logits)
+            required_token_logits = context.last_token_logits(self._all_logits_cuda)
 
         if self._sampling_backend == "torch":
             # Concatenate the outputs once to prevent repeated small writes.
@@ -1234,20 +1260,23 @@ def _router_record_bookkeeping(self) -> Optional[np.ndarray]:
         _ri_dtype = np.int16 if (config.num_moe_experts or 0) <= 32768 else np.int32
         return stacked_routing[:active_token_count].cpu().numpy().astype(_ri_dtype)
 
-    def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
+    def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]:
         """Calculate log probs from logits."""
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
+        logits_seq_len = (
+            active_request_count
+            if context.config.materialize_only_last_token_logits
+            else context.padded_active_token_count
+        )
 
         return context.calculate_log_probs(
-            logits,
+            self._all_logits_cuda[:, :logits_seq_len, :],
             self._sampled_tokens_cuda[:active_request_count],
             only_last_token_logits=context.config.materialize_only_last_token_logits,
         )
 
-    def _dynamic_step_calculate_log_probs_speculative(
-        self, logits: Tensor
-    ) -> Tuple[List[List[float]], Tensor]:
+    def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float]], Tensor]:
         """Calculate log probs from logits for speculative decoding.
 
         For decode requests, computes log probs for each accepted speculative token
@@ -1258,9 +1287,6 @@ def _dynamic_step_calculate_log_probs_speculative(
         - log_prob(accepted_token[j]) comes from logits at position j
         - log_prob(newly_sampled_token) comes from logits at position accepted_count
 
-        Args:
-            logits (Tensor): The main model logits [1, seq_len, vocab_size].
-
         Returns:
             Tuple of (log_probs_list, log_probs_tensor):
                 log_probs_list: List of lists, one per active request, containing
@@ -1279,6 +1305,7 @@ def _dynamic_step_calculate_log_probs_speculative(
         num_decode_requests = active_request_count - num_prefill_requests
 
         only_last = context.config.materialize_only_last_token_logits
+        logits = self._all_logits_cuda
         logits_squeezed = logits.squeeze(0).float()
         if only_last:
             log_probs_tensor = F.log_softmax(logits_squeezed, dim=-1)
@@ -1478,12 +1505,11 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
         return top_n_results if top_n_results else None
 
     def _dynamic_step_calculate_top_n_logprobs(
-        self, logits: Tensor, log_probs_tensor: Optional[Tensor] = None
+        self, log_probs_tensor: Optional[Tensor] = None
     ) -> Optional[Dict[int, List[Tuple[Tensor, Tensor]]]]:
         """Calculate top-n log probs from logits for dynamic batching.
 
         Args:
-            logits (Tensor): The logits to compute top-n log probs from.
             log_probs_tensor (Optional[Tensor]): Pre-computed log probabilities tensor.
                 If provided, avoids recomputing log_softmax. Should be the tensor
                 returned by calculate_log_probs.
@@ -1813,7 +1839,7 @@ async def async_generate_output_tokens_dynamic_batch(
 
             # Forward pass produces only base logits. When speculative decoding is
             # active, MTP logits are computed serially after verification.
-            logits = self._dynamic_step_forward_logits(input_ids, position_ids)
+            self._dynamic_step_forward_logits(input_ids, position_ids)
 
             # Commit Mamba intermediate states before update_requests, which
             # may swap request indices. The Python lists tracking EOS block IDs
@@ -1844,7 +1870,7 @@ async def async_generate_output_tokens_dynamic_batch(
             if self.num_speculative_tokens > 0:
                 # Phase 1: Verify speculative tokens using base logits only.
                 nvtx_range_push("mtp-spec-decoding/verify")
-                self._dynamic_step_sample_logits_and_verify_tokens(logits, input_ids)
+                self._dynamic_step_sample_logits_and_verify_tokens(input_ids)
                 nvtx_range_pop("mtp-spec-decoding/verify")
                 # Phase 2: Rewind KV cache for rejected tokens.
                 nvtx_range_push("mtp-spec-decoding/rewind-kv-cache")
@@ -1866,24 +1892,24 @@ async def async_generate_output_tokens_dynamic_batch(
                 # data-dependent boolean-mask sync overlaps with MTP GPU work.
                 context.kv_block_allocator.release_memory_blocks(blocks_to_release[remove_mask])
             else:
-                self._dynamic_step_sample_logits(logits)
+                self._dynamic_step_sample_logits()
 
             log_probs = None
             top_n_logprobs = None
             if return_log_probs or return_top_n_logprobs:
                 if self.num_speculative_tokens > 0:
                     log_probs, log_probs_tensor = (
-                        self._dynamic_step_calculate_log_probs_speculative(logits)
+                        self._dynamic_step_calculate_log_probs_speculative()
                     )
                     if return_top_n_logprobs:
                         top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs_speculative(
                             log_probs_tensor
                         )
                 else:
-                    log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs(logits)
+                    log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs()
                     if return_top_n_logprobs:
                         top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs(
-                            logits, log_probs_tensor
+                            log_probs_tensor
                         )
 
             if skip_bookkeeping:
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 92b9f013403..01ec988a7ba 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -338,7 +338,8 @@ def test_sample_from_dynamic_logits(
 
         # Sampling.
         logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda()
-        self.text_generation_controller._dynamic_step_sample_logits(logits)
+        self.text_generation_controller._all_logits_cuda = logits
+        self.text_generation_controller._dynamic_step_sample_logits()
         sampled_logits = self.text_generation_controller._sampled_tokens_cuda[:batch_size]
         vocab_indices = torch.arange(self.vocab_size).cuda()
 
@@ -881,7 +882,7 @@ def test_dynamic_top_n_logprobs_calculation(
 
             # Calculate top-n logprobs
             top_n_results = self.text_generation_controller._dynamic_step_calculate_top_n_logprobs(
-                logits, log_probs_tensor
+                log_probs_tensor
             )
 
             # Validate results
@@ -929,7 +930,7 @@ def test_dynamic_top_n_logprobs_calculation(
 
             # Calculate top-n logprobs
             top_n_results = self.text_generation_controller._dynamic_step_calculate_top_n_logprobs(
-                logits, log_probs_tensor
+                log_probs_tensor
             )
 
             # Validate results
@@ -1013,10 +1014,9 @@ def mock_sampling_func(logits, *args, **kwargs):
 
         # Mock logits matching input shape
         logits = torch.randn(1, 6, self.vocab_size, device='cuda')
+        self.text_generation_controller._all_logits_cuda = logits
 
-        self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(
-            logits, input_ids
-        )
+        self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
 
         # Verify acceptance counts
         accepted_counts = self.text_generation_controller._accepted_token_counts_per_request[:2]
@@ -1245,10 +1245,9 @@ def test_speculative_multinomial_sampling(self):
         # Since we are actually testing the internal math of `_torch_sampling_func` handling the shapes,
         # we DO NOT mock `_torch_sampling_func` here. We want it to run natively to prove it doesn't crash.
 
+        self.text_generation_controller._all_logits_cuda = logits
         try:
-            self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(
-                logits, input_ids
-            )
+            self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
         except RuntimeError as e:
             if "prob_dist must be 1 or 2 dim" in str(e):
                 pytest.fail("MTP logits were not flattened before calling multinomial sampling.")

From 24fd8bc5bdad2fae34b4958650be93b2e1b95e97 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 23 Dec 2025 10:11:48 -0600
Subject: [PATCH 06/19] Syntactic sugar for CG and awaiting

---
 megatron/core/inference/utils.py | 94 ++++++++++++++++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index 0914b81f005..5b54adc40f6 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -236,6 +236,100 @@ def tensor_swap(x, src_idxs, dst_idxs):
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
 
 
+class CUDAGraphCache:
+    """Syntactic sugar for capturing and replaying graphs over in-line code blocks.
+
+    In Python, neither context managers nor function calls allow the flexibility of this sugar.
+    This class allows for the following type of usage:
+
+    ```
+        cache_for_my_op = CUDAGraphCache()
+        token_count = some_token_count_value
+        request_count = some_request_count_value
+
+        for _ in cache_for_my_op(token_count, request_count):
+            logits = extract_logits(n).float()
+            probs = torch.softmax(logits, dim=-1)
+            output.copy_(sample(probs))
+    ```
+
+    The above example handles both capture and replay.
+    During capture, the for-loop yields twice; during replay, the for-loop body is skipped.
+    """
+
+    def __init__(self):
+        self._graphs: dict = {}
+
+    def __call__(self, *keys, pool=None, skip_warmup=False, eager=False):
+        """Default use: captures if the keys are new, replays if the keys are cached.
+
+        Args:
+            *keys: Arguments to key the graph on.
+            pool (Optional): CUDA memory pool to build the graphs within.
+            skip_warmup (Optional): If True, skip the warmup run before capture.
+            eager (Optional): If True, run the body once without capturing or replaying.
+        """
+        return self.capture(
+            *keys, pool=pool, replay_on_hit=True, skip_warmup=skip_warmup, eager=eager
+        )
+
+    def __contains__(self, key) -> bool:
+        return key in self._graphs
+
+    def __len__(self) -> int:
+        return len(self._graphs)
+
+    def __delitem__(self, key) -> None:
+        del self._graphs[key]
+
+    def clear(self) -> None:
+        """Remove all cached graphs."""
+        self._graphs.clear()
+
+    def capture(self, *keys, pool=None, replay_on_hit=False, skip_warmup=False, eager=False):
+        """Inline capture: captures if the keys are new, does nothing if the keys are cached.
+
+        Args:
+            *keys: Cache key components (stored as a tuple).
+            pool (Optional): CUDA memory pool to build the graphs within.
+            replay_on_hit (Optional): If True, replay the graph on cache hit.
+            skip_warmup (Optional): If True, skip the warmup run before capture.
+            eager (Optional): If True, run the body once without capturing or replaying.
+        """
+        if eager:
+            yield
+            return
+        key = keys if len(keys) != 1 else keys[0]
+        if key in self._graphs:
+            if replay_on_hit:
+                self._graphs[key].replay()
+            return
+        if not skip_warmup:
+            # Suggested best practice: run the body once before recording to ensure
+            # the kernels are JIT-compiled and the allocator state is stable.
+            yield
+        # Now perform the actual graph capture.
+        g = torch.cuda.CUDAGraph()
+        kwargs = {"pool": pool} if pool is not None else {}
+        with torch.cuda.graph(g, **kwargs):
+            yield
+        self._graphs[key] = g
+
+    def replay(self, *keys) -> None:
+        """Replay a previously captured graph."""
+        key = keys if len(keys) != 1 else keys[0]
+        self._graphs[key].replay()
+
+
+async def torch_awaitable(stream: torch.cuda.Stream | None = None):
+    """Syntactic sugar for returning an awaitable handle for non-distributed torch."""
+    if stream is None:
+        stream = torch.cuda.current_stream()
+    event = stream.record_event()
+    while not event.query():
+        await asyncio.sleep(0)
+
+
 async def await_process_call(call, process: multiprocessing.Process, timeout: float = 1.0):
     """Repeatedly wait for a multiprocessing callable to resolve, aborting upon process failure.
 

From b1970ba6326694ab164da491f89f2cd241c41879 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Sun, 12 Apr 2026 15:45:20 -0500
Subject: [PATCH 07/19] Introduce plumbing for more graphing

---
 .../core/inference/engines/dynamic_engine.py  | 15 +++++++--
 .../text_generation_controller.py             | 32 ++++++++++++++++---
 .../test_text_generation_controller.py        |  2 ++
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index f96235db0c0..bedebacfcd6 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -400,12 +400,21 @@ 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)
+            # Capture all relevant graphs in the pipeline.
+            # Note that some steps of the pipeline may capture multiple different variant graphs.
+            for setup_variant in controller.graph_capture_variants():
+                setup_variant(context)
+                controller._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
+                # Launch bookkeeping on a side stream so it overlaps with forward.
+                with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
+                    controller._pre_forward_bookkeeping_event.record()
+
+                controller._dynamic_step_forward_logits(input_ids, position_ids)
+
+                controller._pre_forward_bookkeeping_event.synchronize()
 
             # MTP CUDA graph warmup for this batch dimension.
             if mtp_warmup_enabled:
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 6e03af947e6..5d9a2d36de6 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -5,7 +5,7 @@
 import copy
 import functools
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, OrderedDict, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, OrderedDict, Tuple, Union
 
 import numpy as np
 import torch
@@ -150,17 +150,22 @@ def _init_dynamic_sampling_tensors(self):
         logits_dtype = self.inference_wrapped_model.config.params_dtype
 
         self._sampling_backend = "torch"
-        self._enable_cuda_graph = False
+        self._enable_cuda_graph = self.model_config.cuda_graph_impl == "local"
 
         # Initialize bookkeeping tensors.
         if self._enable_cuda_graph:
-            self._all_logits_cuda = torch.empty(
-                (1, max_logits, vocab_size), dtype=logits_dtype, device=device
+            self._all_logits_cuda = torch.zeros(
+                (1, max_logits, self.vocab_size), dtype=logits_dtype, device=device
             )
         else:
             self._all_logits_cuda = None
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
+        # Side stream for pre-forward bookkeeping. Work issued here runs concurrently
+        # with the forward pass; post-forward consumers synchronize on the event.
+        self._pre_forward_bookkeeping_stream = torch.cuda.Stream(device=device)
+        self._pre_forward_bookkeeping_event = torch.cuda.Event()
+
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
             self._torch_sampling_buckets: List[Tuple] = []
@@ -1584,6 +1589,18 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         return top_n_results if top_n_results else None
 
+    def graph_capture_variants(self) -> Generator[Callable, None, None]:
+        """Yield context-setup callables for each graph-capture variant.
+
+        During graph warmup, the engine runs the full step pipeline once per yielded callable.
+        Each callable exercises a different kernel path.
+        """
+        if not self._enable_cuda_graph:
+            yield lambda context: None
+            return
+
+        yield lambda context: None
+
     def dummy_forward(self):
         """Perform a dummy forward pass. This is used in expert model parallelism
         on ranks that do not have any real requests. It may run in eager mode."""
@@ -1837,6 +1854,11 @@ async def async_generate_output_tokens_dynamic_batch(
             if config.moe_enable_routing_replay:
                 RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
 
+            # Launch bookkeeping on a side stream so it overlaps with forward.
+            self._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(self._pre_forward_bookkeeping_stream):
+                self._pre_forward_bookkeeping_event.record()
+
             # Forward pass produces only base logits. When speculative decoding is
             # active, MTP logits are computed serially after verification.
             self._dynamic_step_forward_logits(input_ids, position_ids)
@@ -1862,9 +1884,9 @@ async def async_generate_output_tokens_dynamic_batch(
         # NOTE [TDE]: This will be moved once CPU and GPU methods are separated.
         await asyncio.sleep(0)
 
+        self._pre_forward_bookkeeping_event.synchronize()
         with torch.inference_mode():
             return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
-
             self._dynamic_step_sample_bookkeeping()
 
             if self.num_speculative_tokens > 0:
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 01ec988a7ba..fa01776754a 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -67,6 +67,7 @@ def setup_model(
         expert_model_parallel_size: int = 1,
         num_moe_experts: int = None,
         hybrid_layer_pattern: str = None,
+        cuda_graph_impl: str = 'none',
     ):
         if use_training_random_init:
             # This is necessary to induce the training behavior which permutes the random seed
@@ -102,6 +103,7 @@ def setup_model(
                 if hybrid_layer_pattern
                 else {}
             ),
+            cuda_graph_impl=cuda_graph_impl,
         )
         if dtype == torch.bfloat16:
             transformer_config.bf16 = True

From 989663425054bc20e0b26cbf2767259fe7898d4f Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 17 Apr 2026 09:50:21 -0500
Subject: [PATCH 08/19] Address reviewer comment: all metadata on GPU

---
 megatron/core/inference/config.py             |  4 ++--
 .../inference/contexts/dynamic_context.py     | 20 ++++++----------
 megatron/core/inference/inference_request.py  | 23 +++++++++----------
 3 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index e1a36ff1563..c4b092309c2 100644
--- a/megatron/core/inference/config.py
+++ b/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
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index c917050cff3..28fd9f55a90 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/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.
@@ -861,15 +861,9 @@ def initialize_all_tensors(self) -> None:
         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: Dict[str, Tensor] = {}
-        for label, _, on_gpu in self.request_metadata_types:
-            if on_gpu:
-                tensor = torch.empty_like(self.request_metadata[label])
-            else:
-                tensor = torch.empty_like(
-                    self.request_metadata[label], device="cpu", pin_memory=True
-                )
-            self.active_request_metadata[label] = tensor
+        self.active_request_metadata = {
+            label: torch.empty_like(tensor) for label, tensor in self.request_metadata.items()
+        }
 
         self.active_request_ids = torch.empty_like(self.request_ids, dtype=torch.int64)
         self.active_request_query_lengths = torch.empty_like(self.request_query_lengths)
@@ -1082,7 +1076,7 @@ def build_active_slices(self, batch_size: int):
         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_types:
+        for label in self.request_metadata:
             self.active_request_metadata[label][:batch_size].copy_(
                 self.request_metadata[label][padded_slice], non_blocking=True
             )
@@ -1452,7 +1446,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()
             )
@@ -2228,7 +2222,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,
diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py
index 3917bbba720..33fbcdf6518 100644
--- a/megatron/core/inference/inference_request.py
+++ b/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(

From 72d50ad6078499b49fe0d0bd8078226f52aa5113 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 17 Apr 2026 09:54:31 -0500
Subject: [PATCH 09/19] Address reviewer comment: fix padding bug

---
 megatron/core/inference/contexts/dynamic_context.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 28fd9f55a90..5148b5c9285 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1110,13 +1110,21 @@ def build_active_slices(self, batch_size: int):
 
     def pad_active_slices(self):
         """Pad the active slices of specific tensors."""
-        # Some tensors need to be padded at the token level.
-        padding_token_slice = slice(self.active_token_count, self.padded_active_token_count)
 
+        # Token-level padding.
+        padding_token_slice = slice(self.active_token_count, self.padded_active_token_count)
         self.token_to_block_idx[padding_token_slice] = self.kv_block_allocator.dummy_block_idx
         self.token_to_local_position_within_kv_block[padding_token_slice] = 0
         self.token_to_position_in_request[padding_token_slice] = 0
 
+        # Request-level padding.
+        active_request_count = self.total_request_count - self.paused_request_count
+        padding_request_slice = slice(active_request_count, self.padded_active_request_count)
+        self.active_request_query_lengths[padding_request_slice].fill_(
+            self.num_speculative_tokens + 1
+        )
+        self.active_request_last_token_idxs[padding_request_slice].fill_(0)
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 

From c5e411df0d5c23f4579ce9e40af7a68f434ae53e Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 22 Apr 2026 10:16:43 -0500
Subject: [PATCH 10/19] Remove CUDAGraphCache in favor of CudaGraphManager

---
 .../core/inference/engines/dynamic_engine.py  | 69 +++++++--------
 megatron/core/inference/utils.py              | 85 -------------------
 2 files changed, 35 insertions(+), 119 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index bedebacfcd6..8cbe5e304fb 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -405,41 +405,42 @@ def create_cuda_graphs(self, reset_context: bool = True):
 
             # Capture all relevant graphs in the pipeline.
             # Note that some steps of the pipeline may capture multiple different variant graphs.
-            for setup_variant in controller.graph_capture_variants():
-                setup_variant(context)
-                controller._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
-                # Launch bookkeeping on a side stream so it overlaps with forward.
-                with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
-                    controller._pre_forward_bookkeeping_event.record()
-
-                controller._dynamic_step_forward_logits(input_ids, position_ids)
-
-                controller._pre_forward_bookkeeping_event.synchronize()
-
-            # MTP CUDA graph warmup for this batch dimension.
-            if mtp_warmup_enabled:
-                n = cuda_graph_batch_dimension.req_count
-                if sp_enabled:
-                    n = round_up_to_nearest_multiple(n, tp_size)
-                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:
-                        with graph_capture():
-                            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,
-                            )
+            with torch.inference_mode():
+                for setup_variant in controller.graph_capture_variants():
+                    setup_variant(context)
+                    controller._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
+                    # Launch bookkeeping on a side stream so it overlaps with forward.
+                    with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
+                        controller._pre_forward_bookkeeping_event.record()
+
+                    controller._dynamic_step_forward_logits(input_ids, position_ids)
+
+                    controller._pre_forward_bookkeeping_event.synchronize()
+
+                # MTP CUDA graph warmup for this batch dimension.
+                if mtp_warmup_enabled:
+                    n = cuda_graph_batch_dimension.req_count
+                    if sp_enabled:
+                        n = round_up_to_nearest_multiple(n, tp_size)
+                    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:
+                            with graph_capture():
+                                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,
+                                )
 
-            context.reset()
+                context.reset()
 
         # Disable inference dispatcher after graph capture
         if is_inference_optimized_ep:
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index 5b54adc40f6..a6667a974d8 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -236,91 +236,6 @@ def tensor_swap(x, src_idxs, dst_idxs):
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
 
 
-class CUDAGraphCache:
-    """Syntactic sugar for capturing and replaying graphs over in-line code blocks.
-
-    In Python, neither context managers nor function calls allow the flexibility of this sugar.
-    This class allows for the following type of usage:
-
-    ```
-        cache_for_my_op = CUDAGraphCache()
-        token_count = some_token_count_value
-        request_count = some_request_count_value
-
-        for _ in cache_for_my_op(token_count, request_count):
-            logits = extract_logits(n).float()
-            probs = torch.softmax(logits, dim=-1)
-            output.copy_(sample(probs))
-    ```
-
-    The above example handles both capture and replay.
-    During capture, the for-loop yields twice; during replay, the for-loop body is skipped.
-    """
-
-    def __init__(self):
-        self._graphs: dict = {}
-
-    def __call__(self, *keys, pool=None, skip_warmup=False, eager=False):
-        """Default use: captures if the keys are new, replays if the keys are cached.
-
-        Args:
-            *keys: Arguments to key the graph on.
-            pool (Optional): CUDA memory pool to build the graphs within.
-            skip_warmup (Optional): If True, skip the warmup run before capture.
-            eager (Optional): If True, run the body once without capturing or replaying.
-        """
-        return self.capture(
-            *keys, pool=pool, replay_on_hit=True, skip_warmup=skip_warmup, eager=eager
-        )
-
-    def __contains__(self, key) -> bool:
-        return key in self._graphs
-
-    def __len__(self) -> int:
-        return len(self._graphs)
-
-    def __delitem__(self, key) -> None:
-        del self._graphs[key]
-
-    def clear(self) -> None:
-        """Remove all cached graphs."""
-        self._graphs.clear()
-
-    def capture(self, *keys, pool=None, replay_on_hit=False, skip_warmup=False, eager=False):
-        """Inline capture: captures if the keys are new, does nothing if the keys are cached.
-
-        Args:
-            *keys: Cache key components (stored as a tuple).
-            pool (Optional): CUDA memory pool to build the graphs within.
-            replay_on_hit (Optional): If True, replay the graph on cache hit.
-            skip_warmup (Optional): If True, skip the warmup run before capture.
-            eager (Optional): If True, run the body once without capturing or replaying.
-        """
-        if eager:
-            yield
-            return
-        key = keys if len(keys) != 1 else keys[0]
-        if key in self._graphs:
-            if replay_on_hit:
-                self._graphs[key].replay()
-            return
-        if not skip_warmup:
-            # Suggested best practice: run the body once before recording to ensure
-            # the kernels are JIT-compiled and the allocator state is stable.
-            yield
-        # Now perform the actual graph capture.
-        g = torch.cuda.CUDAGraph()
-        kwargs = {"pool": pool} if pool is not None else {}
-        with torch.cuda.graph(g, **kwargs):
-            yield
-        self._graphs[key] = g
-
-    def replay(self, *keys) -> None:
-        """Replay a previously captured graph."""
-        key = keys if len(keys) != 1 else keys[0]
-        self._graphs[key].replay()
-
-
 async def torch_awaitable(stream: torch.cuda.Stream | None = None):
     """Syntactic sugar for returning an awaitable handle for non-distributed torch."""
     if stream is None:

From 60fe5747de4770eff7315846a065606a630c2759 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 27 Apr 2026 14:06:22 -0500
Subject: [PATCH 11/19] Revert change being addressed by #4295

---
 megatron/core/inference/utils.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index a6667a974d8..0914b81f005 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -236,15 +236,6 @@ def tensor_swap(x, src_idxs, dst_idxs):
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
 
 
-async def torch_awaitable(stream: torch.cuda.Stream | None = None):
-    """Syntactic sugar for returning an awaitable handle for non-distributed torch."""
-    if stream is None:
-        stream = torch.cuda.current_stream()
-    event = stream.record_event()
-    while not event.query():
-        await asyncio.sleep(0)
-
-
 async def await_process_call(call, process: multiprocessing.Process, timeout: float = 1.0):
     """Repeatedly wait for a multiprocessing callable to resolve, aborting upon process failure.
 

From 3fc7bd584b5aca9712413d826d883f9ee901678a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 08:42:09 -0500
Subject: [PATCH 12/19] lint

---
 megatron/core/inference/engines/dynamic_engine.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 8cbe5e304fb..9b316dbb652 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -408,7 +408,9 @@ def create_cuda_graphs(self, reset_context: bool = True):
             with torch.inference_mode():
                 for setup_variant in controller.graph_capture_variants():
                     setup_variant(context)
-                    controller._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
+                    controller._pre_forward_bookkeeping_stream.wait_stream(
+                        torch.cuda.current_stream()
+                    )
                     # Launch bookkeeping on a side stream so it overlaps with forward.
                     with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
                         controller._pre_forward_bookkeeping_event.record()
@@ -435,8 +437,12 @@ def create_cuda_graphs(self, reset_context: bool = True):
                                         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),
+                                    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,
                                 )
 

From 94c05252e38140af077e9895febba31f152a5fcc Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 12:35:42 -0500
Subject: [PATCH 13/19] Defer `active_slice` logic to later PRs

---
 .../inference/contexts/dynamic_context.py     | 110 +++++++-----------
 .../text_generation_controller.py             |  35 +++---
 .../test_text_generation_controller.py        |  11 --
 3 files changed, 62 insertions(+), 94 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 5148b5c9285..f12159ad87a 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -865,19 +865,8 @@ def initialize_all_tensors(self) -> None:
             label: torch.empty_like(tensor) for label, tensor in self.request_metadata.items()
         }
 
-        self.active_request_ids = torch.empty_like(self.request_ids, dtype=torch.int64)
-        self.active_request_query_lengths = torch.empty_like(self.request_query_lengths)
-        self.active_request_output_lengths = torch.empty_like(self.request_output_lengths)
-        self.active_request_kv_length_offsets = torch.empty_like(self.request_kv_length_offsets)
-        self.active_request_to_kv_block_ids = torch.empty_like(self.request_to_kv_block_ids)
-
-        self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
-
         # NOTE: Need to build this outside the UVM / TMS context to avoid IMA.
         if self.is_hybrid_model:
-            self.active_mamba_indices = torch.empty_like(
-                self.request_query_lengths, dtype=torch.int32
-            )
             self.mamba_metadata = MambaMetadata(
                 max_requests=self.max_requests,
                 max_tokens=self.max_tokens,
@@ -1064,6 +1053,16 @@ def cu_kv_lengths(self) -> Tuple[Tensor, Tensor, int]:
             self.active_attn_metadata["mha_metadata"].state_data["max_seqlen_k"],
         )
 
+    def get_active_sequence_lengths(self) -> Tensor:
+        """Total sequence length (query + key) for active requests."""
+        lengths = self.request_kv_length_offsets + self.request_query_lengths
+        lengths = lengths[self.paused_request_count : self.total_request_count]
+        return lengths
+
+    def get_max_sequence_lengths(self) -> Tensor:
+        """Maximum sequence length for active requests."""
+        return self.request_output_lengths[self.paused_request_count : self.total_request_count]
+
     def get_active_request_count(self):
         """Returns the current number of active requests."""
         return self.total_request_count - self.paused_request_count
@@ -1081,49 +1080,9 @@ def build_active_slices(self, batch_size: int):
                 self.request_metadata[label][padded_slice], non_blocking=True
             )
 
-        # The following tensor slices are used in various kernels.
-        self.active_request_ids[:batch_size].copy_(self.request_ids[padded_slice])
-        self.active_request_query_lengths[:batch_size].copy_(
-            self.request_query_lengths[padded_slice]
-        )
-        self.active_request_output_lengths[:batch_size].copy_(
-            self.request_output_lengths[padded_slice]
-        )
-        self.active_request_kv_length_offsets[:batch_size].copy_(
-            self.request_kv_length_offsets[padded_slice]
-        )
-        self.active_request_to_kv_block_ids[:batch_size].copy_(
-            self.request_to_kv_block_ids[padded_slice]
-        )
-
-        torch.cumsum(
-            self.active_request_query_lengths[:batch_size],
-            dim=0,
-            out=self.active_request_last_token_idxs[:batch_size],
-        )
-        self.active_request_last_token_idxs[:batch_size] -= 1
-
-        if self.is_hybrid_model:
-            self.active_mamba_indices[:batch_size].copy_(
-                self.mamba_metadata.request_to_mamba_state_idx[padded_slice]
-            )
-
     def pad_active_slices(self):
         """Pad the active slices of specific tensors."""
-
-        # Token-level padding.
-        padding_token_slice = slice(self.active_token_count, self.padded_active_token_count)
-        self.token_to_block_idx[padding_token_slice] = self.kv_block_allocator.dummy_block_idx
-        self.token_to_local_position_within_kv_block[padding_token_slice] = 0
-        self.token_to_position_in_request[padding_token_slice] = 0
-
-        # Request-level padding.
-        active_request_count = self.total_request_count - self.paused_request_count
-        padding_request_slice = slice(active_request_count, self.padded_active_request_count)
-        self.active_request_query_lengths[padding_request_slice].fill_(
-            self.num_speculative_tokens + 1
-        )
-        self.active_request_last_token_idxs[padding_request_slice].fill_(0)
+        pass
 
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
@@ -1761,6 +1720,23 @@ def initialize_attention_state(
                 prefill_req_count=padded_prefill_req_count,
                 decode_req_count=padded_decode_req_count,
             )
+        self.padded_active_token_count = self.padded_batch_dimensions.token_count
+        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
+        )
+        self.token_to_local_position_within_kv_block[
+            self.active_token_count : self.padded_active_token_count
+        ] = 0
+        self.token_to_position_in_request[
+            self.active_token_count : self.padded_active_token_count
+        ] = 0
 
         self.active_attn_metadata = (
             self.graph_attn_metadata  # type: ignore[assignment]
@@ -1768,6 +1744,12 @@ def initialize_attention_state(
             else self.non_graph_attn_metadata  # type: ignore[assignment]
         )
 
+        # Update cu_query_seq_lengths, max_seqlen_q.
+        active_slice = slice(self.paused_request_count, self.total_request_count)
+        query_lengths_view = self.request_query_lengths[active_slice]
+        request_kv_length_offsets_view = self.request_kv_length_offsets[active_slice]
+        request_to_kv_block_ids_view = self.request_to_kv_block_ids[active_slice]
+
         attn_dimensions = batch_dimensions
         if self.using_cuda_graph_this_step():
             # Treat some decode requests as prefill requests to fit the cuda graph batch dimension.
@@ -1781,25 +1763,18 @@ def initialize_attention_state(
                     decode_req_count=adjusted_decode_req_count,
                 )
 
-        self.padded_active_token_count = self.padded_batch_dimensions.token_count
-        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()
-
-        batch_size = self.total_request_count - self.paused_request_count
         assert self.active_attn_metadata is not None
         self.active_attn_metadata["mha_metadata"].update(
-            request_query_lengths=self.active_request_query_lengths[:batch_size],
-            request_kv_length_offsets=self.active_request_kv_length_offsets[:batch_size],
-            request_to_kv_block_ids=self.active_request_to_kv_block_ids[:batch_size],
+            request_query_lengths=query_lengths_view,
+            request_kv_length_offsets=request_kv_length_offsets_view,
+            request_to_kv_block_ids=request_to_kv_block_ids_view,
             batch_dimensions=attn_dimensions,
             padded_batch_dimensions=self.padded_batch_dimensions,
             num_speculative_tokens=self.num_speculative_tokens,
         )
 
         if self.is_hybrid_model:
+            active_mamba_indices_view = self.mamba_metadata.request_to_mamba_state_idx[active_slice]
             token_to_request_idx_view = self.token_to_request_idx[: self.active_token_count]
             cu_seqlens = self.active_attn_metadata["mha_metadata"].state_data[
                 "cu_query_seq_lengths"
@@ -1811,7 +1786,7 @@ def initialize_attention_state(
                     self.mamba_slot_allocator.get_intermediate_gpu_data()
                 )
             self.mamba_metadata.update(
-                self.active_mamba_indices[:batch_size],
+                active_mamba_indices_view,
                 token_to_request_idx_view,
                 cu_seqlens,
                 batch_dimensions=attn_dimensions,
@@ -1989,8 +1964,11 @@ def last_token_logits(self, logits: Tensor) -> Tensor:
             selected = self.speculative_required_logit_indices(logits.device)
             return logits_2d[selected, :]
 
-        active_request_count = self.total_request_count - self.paused_request_count
-        return logits_2d[self.active_request_last_token_idxs[:active_request_count], :]
+        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
+        return logits_2d[last_token_idxs, :]
 
     def _compute_prefix_match(
         self, req: DynamicInferenceRequest, prefill_chunk_length: int
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 5d9a2d36de6..d5d494669cc 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1304,7 +1304,9 @@ def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float
         request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
             context.paused_request_count : context.total_request_count
         ]
-        request_query_lengths = context.active_request_query_lengths[:active_request_count]
+        request_query_lengths = context.request_query_lengths[
+            context.paused_request_count : context.total_request_count
+        ]
 
         num_prefill_requests = request_in_prefill_status_tensor.sum().item()
         num_decode_requests = active_request_count - num_prefill_requests
@@ -1412,7 +1414,9 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
         request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
             context.paused_request_count : context.total_request_count
         ]
-        request_query_lengths = context.active_request_query_lengths[:active_request_count]
+        request_query_lengths = context.request_query_lengths[
+            context.paused_request_count : context.total_request_count
+        ]
 
         num_prefill_requests = request_in_prefill_status_tensor.sum().item()
         num_decode_requests = active_request_count - num_prefill_requests
@@ -1530,6 +1534,7 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
+        active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Handle decode-only mode (only last token)
         if context.config.materialize_only_last_token_logits or context.is_decode_only():
@@ -1553,7 +1558,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         # Note: logits may be padded, so we only take the first active_token_count tokens
         log_probs = log_probs_tensor[: context.active_token_count]
 
-        active_query_lengths = context.active_request_query_lengths[:active_request_count]
+        active_query_lengths = context.request_query_lengths[active_request_slice]
 
         # Split log_probs across request boundaries
         # log_probs has shape [active_token_count, vocab_size]
@@ -1748,26 +1753,22 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Active sequence lengths.
-        # After the forward pass and KV-cache rewind, kv_offsets + query_lengths
-        # already includes all accepted speculative tokens (they were part of the
-        # query and survived the rewind).  Only the newly sampled base token is not
-        # yet in the KV cache, so add 1.
-        active_sequence_lengths = (
-            context.request_kv_length_offsets[active_request_slice]
-            + context.request_query_lengths[active_request_slice]
-            + 1
-        )
+        active_request_ids = context.request_ids[active_request_slice].long()
+        active_sequence_lengths = context.get_active_sequence_lengths()
+
+        # After the forward pass and KV-cache rewind, get_active_sequence_lengths()
+        # returns kv_offsets + query_lengths which already includes all accepted
+        # speculative tokens (they were part of the query and survived the rewind).
+        # Only the newly sampled base token is not yet in the KV cache, so add 1.
+        active_sequence_lengths += 1
+        max_sequence_lengths = context.get_max_sequence_lengths()
 
         # Request finished if termination_id or length >= max_sequence_length.
         # Note: termination_id tensor has per-request termination IDs from mixed sampling
         active_request_mask = (
             self._sampled_tokens_cuda[:active_request_count]
             != context.active_request_metadata["termination_id"][:active_request_count]
-        ).byte() & torch.less(
-            active_sequence_lengths, context.active_request_output_lengths[:active_request_count]
-        ).byte()
-
-        active_request_ids = context.active_request_ids[:active_request_count]
+        ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte()
 
         # Mark requests as finished if they hit stop words
         # (detected in previous step's post_process_requests)
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index fa01776754a..280bd566241 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -324,14 +324,6 @@ def test_sample_from_dynamic_logits(
 
         context.padded_active_token_count = batch_size
         context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32)
-        context.active_request_query_lengths[:batch_size].fill_(1)
-        context.active_request_last_token_idxs[:batch_size].copy_(
-            torch.arange(
-                batch_size,
-                dtype=context.active_request_last_token_idxs.dtype,
-                device=context.active_request_last_token_idxs.device,
-            )
-        )
         context.paused_request_count = 0
         context.total_request_count = batch_size
 
@@ -919,9 +911,6 @@ def test_dynamic_top_n_logprobs_calculation(
             context.request_query_lengths = torch.tensor(
                 [0] * context.paused_request_count + query_lengths, dtype=torch.int32, device='cuda'
             )
-            context.active_request_query_lengths[:batch_size].copy_(
-                torch.tensor(query_lengths, dtype=context.active_request_query_lengths.dtype)
-            )
 
             # Create logits for all tokens
             logits = torch.randn(1, total_tokens, self.vocab_size).cuda()

From c208c07cab92cb79110a774cf6c4af1d9f3fe99e Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 18:51:13 -0500
Subject: [PATCH 14/19] Defer graph_capture_variants and _pre_forward_*

---
 .../core/inference/engines/dynamic_engine.py  | 16 ++----------
 .../text_generation_controller.py             | 26 ++-----------------
 2 files changed, 4 insertions(+), 38 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 9b316dbb652..7bf651ae2a8 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -403,21 +403,9 @@ def create_cuda_graphs(self, reset_context: bool = True):
             if model_config.moe_enable_routing_replay:
                 RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
 
-            # Capture all relevant graphs in the pipeline.
-            # Note that some steps of the pipeline may capture multiple different variant graphs.
+            # Forward pass -> logits.
             with torch.inference_mode():
-                for setup_variant in controller.graph_capture_variants():
-                    setup_variant(context)
-                    controller._pre_forward_bookkeeping_stream.wait_stream(
-                        torch.cuda.current_stream()
-                    )
-                    # Launch bookkeeping on a side stream so it overlaps with forward.
-                    with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
-                        controller._pre_forward_bookkeeping_event.record()
-
-                    controller._dynamic_step_forward_logits(input_ids, position_ids)
-
-                    controller._pre_forward_bookkeeping_event.synchronize()
+                controller._dynamic_step_forward_logits(input_ids, position_ids)
 
                 # MTP CUDA graph warmup for this batch dimension.
                 if mtp_warmup_enabled:
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index d5d494669cc..1bf810d2c5a 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -5,7 +5,7 @@
 import copy
 import functools
 from collections import defaultdict
-from typing import Any, Callable, Dict, Generator, List, Optional, OrderedDict, Tuple, Union
+from typing import Any, Dict, List, Optional, OrderedDict, Tuple, Union
 
 import numpy as np
 import torch
@@ -161,11 +161,6 @@ def _init_dynamic_sampling_tensors(self):
             self._all_logits_cuda = None
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
-        # Side stream for pre-forward bookkeeping. Work issued here runs concurrently
-        # with the forward pass; post-forward consumers synchronize on the event.
-        self._pre_forward_bookkeeping_stream = torch.cuda.Stream(device=device)
-        self._pre_forward_bookkeeping_event = torch.cuda.Event()
-
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
             self._torch_sampling_buckets: List[Tuple] = []
@@ -1594,18 +1589,6 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         return top_n_results if top_n_results else None
 
-    def graph_capture_variants(self) -> Generator[Callable, None, None]:
-        """Yield context-setup callables for each graph-capture variant.
-
-        During graph warmup, the engine runs the full step pipeline once per yielded callable.
-        Each callable exercises a different kernel path.
-        """
-        if not self._enable_cuda_graph:
-            yield lambda context: None
-            return
-
-        yield lambda context: None
-
     def dummy_forward(self):
         """Perform a dummy forward pass. This is used in expert model parallelism
         on ranks that do not have any real requests. It may run in eager mode."""
@@ -1855,11 +1838,6 @@ async def async_generate_output_tokens_dynamic_batch(
             if config.moe_enable_routing_replay:
                 RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
 
-            # Launch bookkeeping on a side stream so it overlaps with forward.
-            self._pre_forward_bookkeeping_stream.wait_stream(torch.cuda.current_stream())
-            with torch.cuda.stream(self._pre_forward_bookkeeping_stream):
-                self._pre_forward_bookkeeping_event.record()
-
             # Forward pass produces only base logits. When speculative decoding is
             # active, MTP logits are computed serially after verification.
             self._dynamic_step_forward_logits(input_ids, position_ids)
@@ -1885,9 +1863,9 @@ async def async_generate_output_tokens_dynamic_batch(
         # NOTE [TDE]: This will be moved once CPU and GPU methods are separated.
         await asyncio.sleep(0)
 
-        self._pre_forward_bookkeeping_event.synchronize()
         with torch.inference_mode():
             return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
+
             self._dynamic_step_sample_bookkeeping()
 
             if self.num_speculative_tokens > 0:

From 7cb7c62b4b1561250f750107bed02a88b678cd5a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 18:52:35 -0500
Subject: [PATCH 15/19] Add comment

---
 .../text_generation_controllers/text_generation_controller.py   | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 1bf810d2c5a..dcb27dac1d3 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1068,6 +1068,7 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         # These indices are always needed for input_ids slicing and tracking
         # accepted sequence positions, even when logits are pre-sliced.
         nvtx_range_push("mtp-spec-decoding/verify/logit-indices")
+        # Use pre-allocated buffer for CUDA graph compatibility.
         logits = self._all_logits_cuda
         required_logit_indices = context.speculative_required_logit_indices(logits.device)
 
@@ -1307,6 +1308,7 @@ def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float
         num_decode_requests = active_request_count - num_prefill_requests
 
         only_last = context.config.materialize_only_last_token_logits
+        # Use pre-allocated buffer for CUDA graph compatibility.
         logits = self._all_logits_cuda
         logits_squeezed = logits.squeeze(0).float()
         if only_last:

From d4ae0e15d48b120b8c8da64dd420e281cc8a7db6 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 19:13:18 -0500
Subject: [PATCH 16/19] Address reviewer comment: fix for spec decoding

---
 .../text_generation_controller.py             | 27 +++++++++++--------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index dcb27dac1d3..f651cf8ec70 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -139,7 +139,8 @@ def _init_dynamic_sampling_tensors(self):
         context = self.inference_wrapped_model.inference_context
         max_requests = context.max_requests
         if context.config.materialize_only_last_token_logits:
-            max_logits = max_requests
+            # Under MTP, each decode request emits (num_speculative_tokens + 1) logit rows
+            max_logits = max_requests * (self.num_speculative_tokens + 1)
         else:
             max_logits = context.max_tokens
 
@@ -683,11 +684,17 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
         """
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
-        logits_seq_len = (
-            active_request_count
-            if context.config.materialize_only_last_token_logits
-            else context.padded_active_token_count
-        )
+        if context.config.materialize_only_last_token_logits:
+            if self.num_speculative_tokens > 0:
+                # Under MTP, each decode request emits (num_speculative_tokens + 1) logit rows.
+                logits_seq_len = (
+                    context.num_decode_requests * (self.num_speculative_tokens + 1)
+                    + context.num_prefill_requests
+                )
+            else:
+                logits_seq_len = active_request_count
+        else:
+            logits_seq_len = context.padded_active_token_count
 
         with torch.inference_mode():
             logits = self.inference_wrapped_model.run_one_forward_step(
@@ -695,11 +702,8 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
             )
             # logits shape: [1, seq_len, vocab_size]
 
-        assert logits_seq_len == (
-            active_request_count
-            if context.config.materialize_only_last_token_logits
-            else input_ids.shape[1]
-        )
+        if not context.config.materialize_only_last_token_logits:
+            assert logits_seq_len == input_ids.shape[1]
 
         # Note: When speculative decoding is active (num_speculative_tokens > 0),
         # the model skips MTP computation during the forward pass. MTP logits
@@ -1265,6 +1269,7 @@ def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]:
         """Calculate log probs from logits."""
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
+        # This code cannot be reached when we are using speculative decode.
         logits_seq_len = (
             active_request_count
             if context.config.materialize_only_last_token_logits

From f58657bc5e14d679a4210e28edd0571c261057e3 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 21:49:34 -0500
Subject: [PATCH 17/19] Fix CI

---
 .../text_generation_controller.py                           | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index f651cf8ec70..3a6276e6ce0 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -735,7 +735,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
 
         # Copy logits to contiguous buffer.
         if self._enable_cuda_graph:
-            self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits)
+            self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits[:, :logits_seq_len, :])
         else:
             self._all_logits_cuda = logits
 
@@ -1168,7 +1168,9 @@ def _dynamic_step_sample_logits(self):
             # already called in the forward pass of GPT.
             required_token_logits = self._all_logits_cuda.squeeze(0)[:active_request_count, :]
         else:
-            required_token_logits = context.last_token_logits(self._all_logits_cuda)
+            required_token_logits = context.last_token_logits(
+                self._all_logits_cuda[:, : context.padded_active_token_count, :]
+            )
 
         if self._sampling_backend == "torch":
             # Concatenate the outputs once to prevent repeated small writes.

From 31c51623aaa77a7e628090c5f865b757d59155fa Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 29 Apr 2026 00:34:00 -0500
Subject: [PATCH 18/19] Address reviewer comments

---
 .../inference/contexts/dynamic_context.py     | 18 +++++++++++++++++-
 .../text_generation_controller.py             | 19 +++----------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index f12159ad87a..2a6e4fb5342 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1937,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.
 
@@ -1950,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"
@@ -1962,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(
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 3a6276e6ce0..9ad51bd1d12 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -683,16 +683,8 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
             position_ids (Tensor): The position IDs.
         """
         context = self.inference_wrapped_model.inference_context
-        active_request_count = context.total_request_count - context.paused_request_count
         if context.config.materialize_only_last_token_logits:
-            if self.num_speculative_tokens > 0:
-                # Under MTP, each decode request emits (num_speculative_tokens + 1) logit rows.
-                logits_seq_len = (
-                    context.num_decode_requests * (self.num_speculative_tokens + 1)
-                    + context.num_prefill_requests
-                )
-            else:
-                logits_seq_len = active_request_count
+            logits_seq_len = context.num_last_token_logits
         else:
             logits_seq_len = context.padded_active_token_count
 
@@ -712,13 +704,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
 
         if self.model_is_pipeline_parallel:
             if context.config.materialize_only_last_token_logits:
-                if self.num_speculative_tokens > 0:
-                    logits_seq_len = (
-                        context.num_decode_requests * (self.num_speculative_tokens + 1)
-                        + context.num_prefill_requests
-                    )
-                else:
-                    logits_seq_len = active_request_count
+                logits_seq_len = context.num_last_token_logits
             else:
                 logits_seq_len = input_ids.shape[1]
             logits_shape = [1, logits_seq_len, self.vocab_size]
@@ -1272,6 +1258,7 @@ def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]:
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
         # This code cannot be reached when we are using speculative decode.
+        assert self.num_speculative_tokens == 0
         logits_seq_len = (
             active_request_count
             if context.config.materialize_only_last_token_logits

From 80c00a8fcb46ebb2e85c789a11495c286d27fa94 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 29 Apr 2026 16:25:37 -0500
Subject: [PATCH 19/19] lint

---
 megatron/core/inference/engines/dynamic_engine.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 7b76ecfc8d1..1c5a4c4cfb1 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -426,12 +426,8 @@ def create_cuda_graphs(self, reset_context: bool = True):
                                     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
-                                ),
+                                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),
                             )