Split request_metadata 2D tensor into 1D slices

Author: tdene

megatron/core/inference/contexts/dynamic_context.py

@@ -239,9 +239,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
@@ -270,7 +270,7 @@ def __init__(
         use_flashinfer_fused_rope: bool = False,
         unified_memory_level: Optional[int] = 1,
         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)
 
@@ -395,9 +395,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
@@ -554,11 +554,14 @@ 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() if on_gpu else torch.device("cpu"),
+            )
+            for label, dtype, on_gpu in self.request_metadata_types
+        }
 
         # Per-token state.
         self.token_to_input_ids = torch.full(
@@ -1170,7 +1173,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)
@@ -1322,14 +1328,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] = (
@@ -1395,7 +1404,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]
 
@@ -1404,6 +1412,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]
@@ -1416,14 +1427,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,7 +6,7 @@
 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
 
@@ -313,20 +313,34 @@ 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",
+    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, True),
+        #    ("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
+        # ]
+        return [
+            ("temperature", torch.float32, True),
+            ("top_k", torch.float32, True),  # CPU for torch sampling
+            ("top_p", torch.float32, True),  # CPU for torch sampling
+            ("termination_id", torch.float32, True),
+            ("return_log_probs", torch.float32, True),  # CPU for non-selective logprobs
+            ("skip_prompt_log_probs", torch.float32, True),  # CPU for non-selective logprobs
         ]
-        return {k: v for v, k in enumerate(ret)}
 
     def add_event(self, type: DynamicInferenceEventType, payload: Optional[Any] = None) -> None:
         """Add event."""

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -23,11 +23,7 @@
     MaxSequenceLengthOverflowError,
     WarmupEngineMode,
 )
-from megatron.core.inference.inference_request import (
-    DynamicInferenceRequest,
-    InferenceRequest,
-    Status,
-)
+from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
@@ -573,11 +569,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
         return logits
 
     def _dynamic_step_sample_bookkeeping(
-        self,
-        *,
-        backend: str = "torch",
-        request_metadata: Optional[Tensor] = None,
-        request_metadata_labels: Dict[str, int] = None,
+        self, *, backend: str = "torch", request_metadata: Optional[Dict[str, Tensor]] = None
     ):
         """Perform bookkeeping necessary to sample logits for dynamic batching.
 
@@ -586,48 +578,62 @@ def _dynamic_step_sample_bookkeeping(
 
         Args:
             backend (str): The sampling backend to use.
-            request_metadata (Optional[Tensor]): An override for the tensor that manages all
-                request metadata, such as sampling parameters. By default, this metadata is
-                retrieved from the context.
-            request_metadata_labels (Optional[Dict]): An override for the map of metadata labels
-                to their index in the request_metadata tensor. By default, this metadata is
-                retrieved from the request object.
+            request_metadata (Optional[Dict[str, Tensor]]): An override for the tensors
+                that manage request metadata, such as sampling parameters. By default, this
+                metadata is retrieved from the context.
         """
         assert backend in ["torch"]
         context = self.inference_wrapped_model.inference_context
+        active_request_count = context.total_request_count - context.paused_request_count
 
         if request_metadata is None:
-            request_metadata = context.request_metadata[
-                context.paused_request_count : context.total_request_count, :
-            ]
-        if request_metadata_labels is None:
-            request_metadata_labels = DynamicInferenceRequest.get_metadata_labels()
-        active_request_count = request_metadata.size(0)
-
-        # Shorthand these, because the torch backend needs them.
-        temp = request_metadata[:, request_metadata_labels["temperature"]]
-        top_k = request_metadata[:, request_metadata_labels["top_k"]]
-        top_p = request_metadata[:, request_metadata_labels["top_p"]]
+            request_metadata = {
+                label: tensor[context.paused_request_count : context.total_request_count]
+                for label, tensor in context.request_metadata.items()
+            }
 
         # Copy data into relevant tensors.
-        self.temperature_cuda[:active_request_count].copy_(temp, non_blocking=True)
-        self.top_k_cuda[:active_request_count] = top_k.to(
+        #        self.temperature_cuda[:active_request_count].copy_(temp, non_blocking=True)
+        #        self.top_k_cuda[:active_request_count] = top_k.to(
+        #            dtype=torch.int32, copy=True, non_blocking=True
+        #        )
+        #        self.top_p_cuda[:active_request_count].copy_(top_p, non_blocking=True)
+        #        self.termination_id_cuda[:active_request_count] = request_metadata[
+        #            :, request_metadata_labels["termination_id"]
+        #        ].to(dtype=torch.int64, copy=True, non_blocking=True)
+        #        self.return_log_probs_cuda[:active_request_count] = request_metadata[
+        #            :, request_metadata_labels["return_log_probs"]
+        #        ].to(dtype=torch.bool, copy=True, non_blocking=True)
+        #        self.skip_prompt_log_probs_cuda[:active_request_count] = request_metadata[
+        #            :, request_metadata_labels["skip_prompt_log_probs"]
+        #        ].to(dtype=torch.bool, copy=True, non_blocking=True)
+        self.temperature_cuda[:active_request_count].copy_(
+            request_metadata["temperature"], non_blocking=True
+        )
+        self.top_k_cuda[:active_request_count] = request_metadata["top_k"].to(
             dtype=torch.int32, copy=True, non_blocking=True
         )
-        self.top_p_cuda[:active_request_count].copy_(top_p, non_blocking=True)
-        self.termination_id_cuda[:active_request_count] = request_metadata[
-            :, request_metadata_labels["termination_id"]
-        ].to(dtype=torch.int64, copy=True, non_blocking=True)
-        self.return_log_probs_cuda[:active_request_count] = request_metadata[
-            :, request_metadata_labels["return_log_probs"]
-        ].to(dtype=torch.bool, copy=True, non_blocking=True)
+        self.top_p_cuda[:active_request_count].copy_(request_metadata["top_p"], non_blocking=True)
+        self.termination_id_cuda[:active_request_count] = request_metadata["termination_id"].to(
+            dtype=torch.int64, copy=True, non_blocking=True
+        )
+        self.return_log_probs_cuda[:active_request_count] = request_metadata["return_log_probs"].to(
+            dtype=torch.bool, copy=True, non_blocking=True
+        )
         self.skip_prompt_log_probs_cuda[:active_request_count] = request_metadata[
-            :, request_metadata_labels["skip_prompt_log_probs"]
+            "skip_prompt_log_probs"
         ].to(dtype=torch.bool, copy=True, non_blocking=True)
 
         if backend == "torch":
             # Bucketize the core sampling parameters.
-            core_params = torch.stack((temp, top_k, top_p), dim=1)
+            core_params = torch.stack(
+                (
+                    request_metadata["temperature"],
+                    request_metadata["top_k"],
+                    request_metadata["top_p"],
+                ),
+                dim=1,
+            )
             _, inv_indices, cnts = torch.unique(
                 core_params, dim=0, return_inverse=True, return_counts=True
             )

tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py

@@ -258,16 +258,11 @@ def test_sample_from_dynamic_logits(self, backend):
                 rev_sampling_dict[idx] = sampling_params
 
         # Prepare metadata for sample bookkeeping.
-        request_metadata_labels = DynamicInferenceRequest.get_metadata_labels()
-        request_metadata = torch.empty(
-            (batch_size, len(request_metadata_labels)), dtype=torch.float32
-        ).cuda()
-        top_k_values = torch.Tensor([s.top_k for s in rev_sampling_dict]).cuda()
-        request_metadata[:, request_metadata_labels["top_k"]] = top_k_values
-        top_p_values = torch.Tensor([s.top_p for s in rev_sampling_dict]).cuda()
-        request_metadata[:, request_metadata_labels["top_p"]] = top_p_values
-        temp_values = torch.Tensor([s.temperature for s in rev_sampling_dict]).cuda()
-        request_metadata[:, request_metadata_labels["temperature"]] = temp_values
+        request_metadata = {
+            "temperature": torch.Tensor([s.temperature for s in rev_sampling_dict]).cuda(),
+            "top_k": torch.Tensor([s.top_k for s in rev_sampling_dict]).cuda(),
+            "top_p": torch.Tensor([s.top_p for s in rev_sampling_dict]).cuda(),
+        }
 
         # Bookkeeping.
         self.text_generation_controller._dynamic_step_sample_bookkeeping(