diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index a0ed0b80a7d..f9695f6c9a5 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. 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. 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; # only one is active per step, so sharing storage is safe. self.graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view) @@ -2432,10 +2439,14 @@ def current_input_and_position_ids( 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), - ) + cached = self._input_position_views.get(num_tokens) + if cached is not None: + return 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: """Token-level indices needed for speculative decode verification. 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])