From 0e00721886fdb0554b112230a99e62d6cc37744d Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 5 May 2026 13:06:15 -0700 Subject: [PATCH 1/4] some more markers and the view cache --- .../inference/contexts/dynamic_context.py | 88 ++++++++++++------- .../text_generation_controller.py | 9 +- 2 files changed, 63 insertions(+), 34 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index a0ed0b80a7d..affdcb07ae1 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -5,7 +5,7 @@ import operator import warnings from contextlib import nullcontext -from typing import List, Optional, Sequence, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import torch # type: ignore import torch.nn.functional as F # type: ignore @@ -1182,6 +1182,13 @@ def initialize_all_tensors(self) -> None: max_mamba_chunks=self._max_mamba_chunks, ) + # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. The slice + # + unsqueeze chain in current_input_and_position_ids constructs new + # TensorImpls each call (~30-60 us on the host); the underlying storage + # is fixed so views are reusable across steps. Bounded by the small set + # of token counts that recur (graph sizes + the eager batch size). + self._input_position_views: Dict[int, Tuple[Tensor, Tensor]] = {} + # Bind the shared MHA GPU views to both graph and non-graph metadata; # only one is active per step, so sharing storage is safe. self.graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view) @@ -2281,52 +2288,64 @@ def transfer_bookkeeping_to_gpu(self) -> None: tensors immediately before the H2D (GPU reads them at `[:n_active]` while CPU bookkeeping keeps them at `[paused_count:total_count)`). """ - n_active = self.total_request_count - self.paused_request_count - active_slice = slice(self.paused_request_count, self.total_request_count) - padded_active = max(n_active, self.padded_active_request_count) + with torch.cuda.nvtx.range("xfer_bk.compute_slices"): + n_active = self.total_request_count - self.paused_request_count + active_slice = slice(self.paused_request_count, self.total_request_count) + padded_active = max(n_active, self.padded_active_request_count) # Refresh request-level staging slots from the persistent CPU source. # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6 # 4-byte fields at max_requests=624). Negligible vs. the launch overhead # we save by merging the H2D memcpys into 1. - self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[ - active_slice - ] - self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice] - self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ - active_slice - ] + with torch.cuda.nvtx.range("xfer_bk.staging_request_fields"): + self._staging_request_in_prefill_status[:n_active] = ( + self.request_in_prefill_status_tensor[active_slice] + ) + self._staging_request_query_lengths[:n_active] = self.request_query_lengths[ + active_slice + ] + self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ + active_slice + ] # Sampling-parameter staging slots: read from `active_request_metadata`, # which `build_active_slices` + `pad_active_slices` already populated for # `[:padded_active]` (active values + neutral padding defaults). - self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][ - :padded_active - ] - self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active] - self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active] + with torch.cuda.nvtx.range("xfer_bk.staging_sampling_params"): + self._staging_temperature[:padded_active] = self.active_request_metadata[ + "temperature" + ][:padded_active] + self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][ + :padded_active + ] + self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][ + :padded_active + ] # Full-iteration CUDA graphs may have captured GPU consumers with the # padded graph request count. Keep those padded staging rows bounded so # graph replay never builds indices from stale request lengths. if n_active < padded_active: - self._staging_request_in_prefill_status[n_active:padded_active] = 0 - self._staging_request_query_lengths[n_active:padded_active] = 0 - self._staging_request_kv_length_offsets[n_active:padded_active] = 0 + with torch.cuda.nvtx.range("xfer_bk.zero_pad_rows"): + self._staging_request_in_prefill_status[n_active:padded_active] = 0 + self._staging_request_query_lengths[n_active:padded_active] = 0 + self._staging_request_kv_length_offsets[n_active:padded_active] = 0 # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer. # Copying the whole (max_tokens + max_requests)-sized buffer including # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves # 8 redundant launch overheads vs. the prior per-field copies. - self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) + with torch.cuda.nvtx.range("xfer_bk.h2d_copy"): + self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying # bytes. Nothing else to do here for MHA. # Mamba metadata: copy pre-computed CPU tensors to GPU buffers. - if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None: - self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) - self._pending_mamba_transfer = None + with torch.cuda.nvtx.range("xfer_bk.mamba_transfer"): + if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None: + self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) + self._pending_mamba_transfer = None def reset_tensors(self) -> None: """Fill all bookkeeping tensors with sentinel values.""" @@ -2428,14 +2447,21 @@ def current_input_and_position_ids( Return: (Tuple[Tensor, Tensor]) Flattened active input and position IDs. """ - num_tokens = num_warmup_tokens or self.padded_active_token_count - assert num_tokens >= self.padded_batch_dimensions.decode_req_count * ( - self.num_speculative_tokens + 1 - ) - return ( - self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0), - self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0), - ) + with torch.cuda.nvtx.range("cur_input.resolve_count"): + num_tokens = num_warmup_tokens or self.padded_active_token_count + with torch.cuda.nvtx.range("cur_input.assert_count"): + assert num_tokens >= self.padded_batch_dimensions.decode_req_count * ( + self.num_speculative_tokens + 1 + ) + cached = self._input_position_views.get(num_tokens) + if cached is not None: + return cached + with torch.cuda.nvtx.range("cur_input.build_views"): + input_ids = self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0) + pos_ids = self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0) + cached = (input_ids, pos_ids) + self._input_position_views[num_tokens] = cached + return cached def speculative_required_logit_indices(self) -> Tensor: """Token-level indices needed for speculative decode verification. 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 87edddea566..0c9888006f7 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -540,11 +540,14 @@ def _dynamic_step_context_init( input_ids (Tensor): The active input IDs. position_ids (Tensor): The active position IDs. """ - context = self.inference_wrapped_model.inference_context + with torch.cuda.nvtx.range("ctx_init.get_context"): + context = self.inference_wrapped_model.inference_context # Remove Float16Module wrapper if it exists - unwrapped_model = unwrap_model(self.inference_wrapped_model.model) - model_config = get_model_config(unwrapped_model) + with torch.cuda.nvtx.range("ctx_init.unwrap_model"): + unwrapped_model = unwrap_model(self.inference_wrapped_model.model) + with torch.cuda.nvtx.range("ctx_init.get_model_config"): + model_config = get_model_config(unwrapped_model) # Initialize attention state (100% CPU computation). range_push("initialize_attention_state") From 7930872d9c93e7db9f89fccd3effdb287235227c Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 5 May 2026 13:19:59 -0700 Subject: [PATCH 2/4] revert markers which were causing overhead --- .../inference/contexts/dynamic_context.py | 75 ++++++++----------- .../text_generation_controller.py | 9 +-- 2 files changed, 33 insertions(+), 51 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index affdcb07ae1..c08d0b7f757 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2288,64 +2288,52 @@ def transfer_bookkeeping_to_gpu(self) -> None: tensors immediately before the H2D (GPU reads them at `[:n_active]` while CPU bookkeeping keeps them at `[paused_count:total_count)`). """ - with torch.cuda.nvtx.range("xfer_bk.compute_slices"): - n_active = self.total_request_count - self.paused_request_count - active_slice = slice(self.paused_request_count, self.total_request_count) - padded_active = max(n_active, self.padded_active_request_count) + n_active = self.total_request_count - self.paused_request_count + active_slice = slice(self.paused_request_count, self.total_request_count) + padded_active = max(n_active, self.padded_active_request_count) # Refresh request-level staging slots from the persistent CPU source. # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6 # 4-byte fields at max_requests=624). Negligible vs. the launch overhead # we save by merging the H2D memcpys into 1. - with torch.cuda.nvtx.range("xfer_bk.staging_request_fields"): - self._staging_request_in_prefill_status[:n_active] = ( - self.request_in_prefill_status_tensor[active_slice] - ) - self._staging_request_query_lengths[:n_active] = self.request_query_lengths[ - active_slice - ] - self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ - active_slice - ] + self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[ + active_slice + ] + self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice] + self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ + active_slice + ] # Sampling-parameter staging slots: read from `active_request_metadata`, # which `build_active_slices` + `pad_active_slices` already populated for # `[:padded_active]` (active values + neutral padding defaults). - with torch.cuda.nvtx.range("xfer_bk.staging_sampling_params"): - self._staging_temperature[:padded_active] = self.active_request_metadata[ - "temperature" - ][:padded_active] - self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][ - :padded_active - ] - self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][ - :padded_active - ] + self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][ + :padded_active + ] + self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active] + self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active] # Full-iteration CUDA graphs may have captured GPU consumers with the # padded graph request count. Keep those padded staging rows bounded so # graph replay never builds indices from stale request lengths. if n_active < padded_active: - with torch.cuda.nvtx.range("xfer_bk.zero_pad_rows"): - self._staging_request_in_prefill_status[n_active:padded_active] = 0 - self._staging_request_query_lengths[n_active:padded_active] = 0 - self._staging_request_kv_length_offsets[n_active:padded_active] = 0 + self._staging_request_in_prefill_status[n_active:padded_active] = 0 + self._staging_request_query_lengths[n_active:padded_active] = 0 + self._staging_request_kv_length_offsets[n_active:padded_active] = 0 # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer. # Copying the whole (max_tokens + max_requests)-sized buffer including # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves # 8 redundant launch overheads vs. the prior per-field copies. - with torch.cuda.nvtx.range("xfer_bk.h2d_copy"): - self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) + self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True) # MHA metadata GPU views were already bound to state_data in # initialize_attention_state(); the H2D above populates the underlying # bytes. Nothing else to do here for MHA. # Mamba metadata: copy pre-computed CPU tensors to GPU buffers. - with torch.cuda.nvtx.range("xfer_bk.mamba_transfer"): - if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None: - self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) - self._pending_mamba_transfer = None + if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None: + self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer) + self._pending_mamba_transfer = None def reset_tensors(self) -> None: """Fill all bookkeeping tensors with sentinel values.""" @@ -2447,20 +2435,17 @@ def current_input_and_position_ids( Return: (Tuple[Tensor, Tensor]) Flattened active input and position IDs. """ - with torch.cuda.nvtx.range("cur_input.resolve_count"): - num_tokens = num_warmup_tokens or self.padded_active_token_count - with torch.cuda.nvtx.range("cur_input.assert_count"): - assert num_tokens >= self.padded_batch_dimensions.decode_req_count * ( - self.num_speculative_tokens + 1 - ) + num_tokens = num_warmup_tokens or self.padded_active_token_count + assert num_tokens >= self.padded_batch_dimensions.decode_req_count * ( + self.num_speculative_tokens + 1 + ) cached = self._input_position_views.get(num_tokens) if cached is not None: return cached - with torch.cuda.nvtx.range("cur_input.build_views"): - input_ids = self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0) - pos_ids = self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0) - cached = (input_ids, pos_ids) - self._input_position_views[num_tokens] = cached + input_ids = self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0) + pos_ids = self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0) + cached = (input_ids, pos_ids) + self._input_position_views[num_tokens] = cached return cached def speculative_required_logit_indices(self) -> Tensor: 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 0c9888006f7..87edddea566 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -540,14 +540,11 @@ def _dynamic_step_context_init( input_ids (Tensor): The active input IDs. position_ids (Tensor): The active position IDs. """ - with torch.cuda.nvtx.range("ctx_init.get_context"): - context = self.inference_wrapped_model.inference_context + context = self.inference_wrapped_model.inference_context # Remove Float16Module wrapper if it exists - with torch.cuda.nvtx.range("ctx_init.unwrap_model"): - unwrapped_model = unwrap_model(self.inference_wrapped_model.model) - with torch.cuda.nvtx.range("ctx_init.get_model_config"): - model_config = get_model_config(unwrapped_model) + unwrapped_model = unwrap_model(self.inference_wrapped_model.model) + model_config = get_model_config(unwrapped_model) # Initialize attention state (100% CPU computation). range_push("initialize_attention_state") From 1a69cd776b8b17765b1196f3cdbd96ac2f26b50c Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 5 May 2026 13:25:20 -0700 Subject: [PATCH 3/4] clarify --- megatron/core/inference/contexts/dynamic_context.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c08d0b7f757..0bd53b59192 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1182,10 +1182,9 @@ def initialize_all_tensors(self) -> None: max_mamba_chunks=self._max_mamba_chunks, ) - # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. The slice - # + unsqueeze chain in current_input_and_position_ids constructs new - # TensorImpls each call (~30-60 us on the host); the underlying storage - # is fixed so views are reusable across steps. Bounded by the small set + # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and + # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), + # we fix the underlying storage so views are reusable across steps. Bounded by the small set # of token counts that recur (graph sizes + the eager batch size). self._input_position_views: Dict[int, Tuple[Tensor, Tensor]] = {} From 394ead3e9aab19dd3f9d8099d8bf4e55a38394b7 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Thu, 7 May 2026 08:13:32 -0700 Subject: [PATCH 4/4] add unit test, fix comment --- .../inference/contexts/dynamic_context.py | 5 ++- .../contexts/test_dynamic_context.py | 40 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0bd53b59192..f9695f6c9a5 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1184,8 +1184,9 @@ def initialize_all_tensors(self) -> None: # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us), - # we fix the underlying storage so views are reusable across steps. Bounded by the small set - # of token counts that recur (graph sizes + the eager batch size). + # we fix the underlying storage so views are reusable across steps. The number of entries + # is bounded by the graph sizes plus eager-mode token counts, which are rounded up to + # multiples of TOKEN_ROUNDER and capped at max_tokens / TOKEN_ROUNDER distinct values. self._input_position_views: Dict[int, Tuple[Tensor, Tensor]] = {} # Bind the shared MHA GPU views to both graph and non-graph metadata; diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 0e307f600f1..6a4838f2c46 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -256,6 +256,46 @@ def test_token_overflow_error(self, is_hybrid_model: bool): ) ) # Exceeding max token count + @pytest.mark.internal + @rounder_override(64) + def test_current_input_and_position_ids_view_cache(self): + dynamic_context = self._get_dynamic_context( + params_dtype=torch.float32, + num_layers=2, + kv_channels=64, + num_attention_heads=8, + max_sequence_length=128, + buffer_size_gb=0.1, + block_size_tokens=128, + max_tokens=None, + ) + + num_tokens = 64 + dynamic_context.padded_active_token_count = num_tokens + + # First call: cache miss, populates entry. + assert num_tokens not in dynamic_context._input_position_views + input_ids_view, pos_ids_view = dynamic_context.current_input_and_position_ids() + assert num_tokens in dynamic_context._input_position_views + + # Second call: cache hit returns the same tensor objects. + cached_input_ids, cached_pos_ids = dynamic_context.current_input_and_position_ids() + assert cached_input_ids is input_ids_view + assert cached_pos_ids is pos_ids_view + + # Writing new values into the underlying storage must be reflected by the cached views. + device = dynamic_context.gpu_view.token_to_input_ids.device + new_input_ids = torch.arange(num_tokens, dtype=torch.long, device=device) + new_pos_ids = torch.arange(num_tokens, 2 * num_tokens, dtype=torch.long, device=device) + dynamic_context.gpu_view.token_to_input_ids[:num_tokens] = new_input_ids + dynamic_context.gpu_view.token_to_pos_ids[:num_tokens] = new_pos_ids + + refreshed_input_ids, refreshed_pos_ids = dynamic_context.current_input_and_position_ids() + assert refreshed_input_ids is input_ids_view + assert refreshed_pos_ids is pos_ids_view + assert torch.equal(refreshed_input_ids.squeeze(0), new_input_ids) + assert torch.equal(refreshed_pos_ids.squeeze(0), new_pos_ids) + @pytest.mark.internal @rounder_override(64) @pytest.mark.parametrize("is_hybrid_model", [False, True])