Improve performance of request_metadata logic

megatron/core/inference/contexts/dynamic_context.py

@@ -238,9 +238,9 @@ class DynamicInferenceContext(BaseInferenceContext):
         use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation.
             If None, defaults to using flash-infer if available.
         metrics_writer (Optional['WandbModule']): Wandb module for writing metrics.
-        num_request_metadata (Optional[int]): Number of metadata fields to track per request.
-            These represent metadata that is needed by the text generation controller,
-            and that must be kept in sync with active requests through update_requests.
+        request_metadata_types (Optional[List[Tuple[str, torch.dtype, bool]]]): 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.
     """
 
     DEFAULT_MAX_TOKENS = 16384
@@ -271,7 +271,7 @@ def __init__(
         cuda_graph_max_tokens: Optional[int] = None,
         cuda_graph_mixed_prefill_count: Optional[int] = 16,
         metrics_writer: Optional['WandbModule'] = None,
-        num_request_metadata: Optional[int] = None,
+        request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None,
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
@@ -399,9 +399,9 @@ def __init__(
         )
 
         # Track request metadata.
-        if num_request_metadata is None:
-            num_request_metadata = len(DynamicInferenceRequest.get_metadata_labels())
-        self.num_request_metadata = num_request_metadata
+        if request_metadata_types is None:
+            request_metadata_types = DynamicInferenceRequest.get_metadata_types()
+        self.request_metadata_types = request_metadata_types
 
         # Initialize context state.
         self.params_dtype = params_dtype
@@ -537,11 +537,12 @@ def allocate_all_tensors(self, *, is_init: bool) -> None:
         )
 
         # Track request metadata.
-        self.request_metadata = torch.empty(
-            (self.max_total_requests, self.num_request_metadata),
-            dtype=torch.float32,
-            device=torch.cuda.current_device(),
-        )
+        self.request_metadata = {
+            label: torch.empty(
+                (self.max_total_requests,), dtype=dtype, device=torch.cuda.current_device()
+            )
+            for label, dtype, _ in self.request_metadata_types
+        }
 
         # Per-token state.
         self.token_to_input_ids = torch.full(
@@ -999,7 +1000,7 @@ def add_dummy_requests_parallel(
         num_tokens_to_generate: List[int] = []
         request_ids: List[int] = []
         prompt_tokens: List[Tensor] = []
-        metadata_rows: List[List[float]] = []
+        metadata_cols: List[List] = [[] for _ in self.request_metadata_types]
 
         for req in requests:
             assert isinstance(
@@ -1020,7 +1021,8 @@ def add_dummy_requests_parallel(
                     device=self.token_to_input_ids.device, dtype=self.token_to_input_ids.dtype
                 )
             )
-            metadata_rows.append(req.tracked_metadata)
+            for i, m in enumerate(req.tracked_metadata):
+                metadata_cols[i].append(m)
 
         total_new_tokens = sum(lengths)
         if self.active_token_count + total_new_tokens > self.max_tokens:
@@ -1034,9 +1036,6 @@ def add_dummy_requests_parallel(
             num_tokens_to_generate, dtype=self.request_query_lengths.dtype, device=device
         )
         request_ids_tensor = torch.tensor(request_ids, dtype=self.request_ids.dtype, device=device)
-        metadata_tensor = torch.tensor(
-            metadata_rows, dtype=self.request_metadata.dtype, device=self.request_metadata.device
-        )
 
         block_counts = torch.div(
             lengths_tensor + (self.block_size_tokens - 1),
@@ -1053,7 +1052,10 @@ 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
-        self.request_metadata[request_slice] = metadata_tensor
+        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()
+            )
 
         dummy_block_idx = self.block_allocator.dummy_block_idx
         self.request_last_kv_block_id[request_slice] = dummy_block_idx
@@ -1325,7 +1327,10 @@ def reset(self) -> None:
         self.request_last_kv_block_id.fill_(-1)
         self.request_last_kv_block_offset.fill_(0)
         self.request_to_kv_block_ids.fill_(-1)
-        self.request_metadata.fill_(0)
+
+        # Reset request metadata.
+        for metadata_tensor in self.request_metadata.values():
+            metadata_tensor.fill_(0)
 
         # Reset token indexes.
         self.token_to_input_ids.fill_(0)
@@ -1482,14 +1487,17 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int]
             raise TokenOverflowError(req.request_id)
 
         self.request_ids[current_id] = req.request_id
+
         # Handle request metadata.
-        metadata = req.tracked_metadata
         assert (
-            len(metadata) == self.num_request_metadata
-        ), "Request added to context with invalid metadata length"
-        self.request_metadata[current_id] = torch.tensor(
-            metadata, dtype=torch.float32, device=self.request_metadata.device
-        )
+            req.get_metadata_types() == self.request_metadata_types
+        ), "Request added to context with invalid metadata types"
+        metadata = req.tracked_metadata
+        metadata_types = req.get_metadata_types()
+        for m, m_type in zip(metadata, metadata_types):
+            label, _, _ = m_type
+            self.request_metadata[label][current_id] = m
+
         # Handle length and block assignments.
         self.request_query_lengths[current_id] = chunk_length
         self.request_output_lengths[current_id] = (
@@ -1556,7 +1564,6 @@ def _move_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens):
         self.request_kv_length_offsets[dst_idxs] = self.request_kv_length_offsets[src_idxs]
         self.request_query_lengths[dst_idxs] = self.request_query_lengths[src_idxs]
         self.request_output_lengths[dst_idxs] = self.request_output_lengths[src_idxs]
-        self.request_metadata[dst_idxs] = self.request_metadata[src_idxs]
         self.request_ids[dst_idxs] = self.request_ids[src_idxs]
         next_tokens[dst_idxs] = next_tokens[src_idxs]
 
@@ -1565,6 +1572,9 @@ def _move_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens):
         self.request_last_kv_block_id[dst_idxs] = self.request_last_kv_block_id[src_idxs]
         self.request_last_kv_block_offset[dst_idxs] = self.request_last_kv_block_offset[src_idxs]
 
+        for metadata_tensor in self.request_metadata.values():
+            metadata_tensor[dst_idxs] = metadata_tensor[src_idxs]
+
         if self.is_hybrid_model:
             self.mamba_metadata.request_to_mamba_state_idx[dst_idxs] = (
                 self.mamba_metadata.request_to_mamba_state_idx[src_idxs]
@@ -1577,14 +1587,16 @@ def _swap_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens):
         tensor_swap(self.request_kv_length_offsets, src_idxs, dst_idxs)
         tensor_swap(self.request_query_lengths, src_idxs, dst_idxs)
         tensor_swap(self.request_output_lengths, src_idxs, dst_idxs)
-        tensor_swap(self.request_metadata, src_idxs, dst_idxs)
         tensor_swap(self.request_ids, src_idxs, dst_idxs)
         tensor_swap(next_tokens, src_idxs, dst_idxs)
         tensor_swap(self.request_to_kv_block_ids, src_idxs, dst_idxs)
         tensor_swap(self.request_kv_block_counts, src_idxs, dst_idxs)
         tensor_swap(self.request_last_kv_block_id, src_idxs, dst_idxs)
         tensor_swap(self.request_last_kv_block_offset, src_idxs, dst_idxs)
 
+        for metadata_tensor in self.request_metadata.values():
+            tensor_swap(metadata_tensor, src_idxs, dst_idxs)
+
         if self.is_hybrid_model:
             tensor_swap(self.mamba_metadata.request_to_mamba_state_idx, src_idxs, dst_idxs)
 

megatron/core/inference/inference_request.py

@@ -6,12 +6,13 @@
 import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum, auto
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.tokenizers import MegatronTokenizer
+from megatron.core.utils import experimental_api
 
 
 def serialize_tensor(tensor: torch.Tensor) -> bytes:
@@ -228,6 +229,7 @@ def deserialize(cls, obj: dict) -> "DynamicInferenceEvent":
         return event
 
 
+@experimental_api
 @dataclass(kw_only=True)
 class DynamicInferenceRequest(InferenceRequest):
     """Class for one inference request
@@ -313,21 +315,27 @@ 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_labels().keys()]
+        return [getattr(sp, field) for field, _, _ in self.get_metadata_types()]
 
     @staticmethod
-    def get_metadata_labels() -> Dict[str, int]:
-        """Provides human-readable labels for the tracked metadata fields."""
-        ret = [
-            "temperature",
-            "top_k",
-            "top_p",
-            "termination_id",
-            "return_log_probs",
-            "skip_prompt_log_probs",
-            "top_n_logprobs",
+    def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
+        """Keeps track of all request metadata names, dtypes, and target device.
+
+        Returns:
+            List[Tuple[str, torch.dtype, bool]]: 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
         ]
-        return {k: v for v, k in enumerate(ret)}
 
     def add_event(self, type: DynamicInferenceEventType, payload: Optional[Any] = None) -> None:
         """Add event."""