Further speed up torch sampling

Author: tdene

megatron/core/inference/inference_request.py

@@ -12,7 +12,7 @@
 
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.tokenizers import MegatronTokenizer
-from megatron.core.utils import internal_api
+from megatron.core.utils import experimental_api
 
 
 def serialize_tensor(tensor: torch.Tensor) -> bytes:
@@ -229,7 +229,7 @@ def deserialize(cls, obj: dict) -> "DynamicInferenceEvent":
         return event
 
 
-@internal_api
+@experimental_api
 @dataclass(kw_only=True)
 class DynamicInferenceRequest(InferenceRequest):
     """Class for one inference request

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -89,11 +89,6 @@ def _init_dynamic_sampling_tensors(self):
         vocab_size = self.inference_wrapped_model.inference_wrapper_config.padded_vocab_size
 
         self._sampling_backend = "torch"
-
-        # Initialize bookkeeping tensors.
-        self._sampling_logits_cuda = torch.empty(
-            max_requests, vocab_size, dtype=logits_dtype, device=device
-        )
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
         # Keep track of request metadata.
@@ -116,6 +111,11 @@ def _init_dynamic_sampling_tensors(self):
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
             self._torch_sampling_buckets: Iterator[Tuple] = []
+        elif self._sampling_backend == "flashinfer":
+            self._sampling_logits_cuda = torch.empty(
+                max_requests, vocab_size, dtype=logits_dtype, device=device
+            )
+            raise NotImplementedError
 
     def tokenize_prompt(self, prompt: str, add_BOS: bool = False) -> List[int]:
         """Utility to tokenize the input prompts.
@@ -593,23 +593,10 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
                 pp_group=self.pp_group,
             )
 
-        # Last token logits.
-        if self._materialize_only_last:
-            # When materialize_only_last_token_logits is true, last_token_logits is
-            # already called in the forward pass of GPT.
-            last_token_logits = logits.squeeze(0)
-        else:
-            last_token_logits = context.last_token_logits(logits)
-        # Copy last_token_logits to contiguous buffer.
-        self._sampling_logits_cuda[: self._active_request_count].copy_(
-            last_token_logits, non_blocking=True
-        )
-
         return logits
 
     def _dynamic_step_sample_bookkeeping(self):
         """Perform bookkeeping necessary to sample logits for dynamic batching."""
-
         if self._sampling_backend == "torch":
             # Bucketize the core sampling parameters.
             # Doing so via list comprehension is orders of magnitude faster than via torch.
@@ -633,27 +620,43 @@ def _dynamic_step_sample_bookkeeping(self):
                 (indices, temp[rep], top_k[rep], top_p[rep]) for indices, rep in bucket_map.values()
             )
 
-    def _dynamic_step_sample_logits(self):
-        """Sample tokens from logits for dynamic batching."""
+    def _dynamic_step_sample_logits(self, logits: Tensor):
+        """Sample tokens from logits for dynamic batching.
+
+        Args:
+            logits (Tensor): The logits from the forward pass.
+        """
         # 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.
+        if self._materialize_only_last:
+            # When materialize_only_last_token_logits is true, last_token_logits is
+            # already called in the forward pass of GPT.
+            last_token_logits = logits.squeeze(0)
+        else:
+            last_token_logits = context.last_token_logits(logits)
+
         if self._sampling_backend == "torch":
             # Concatenate the outputs once to prevent repeated small writes.
             token_list = []
             indices_list = []
 
             for indices, temp, top_k, top_p in self._torch_sampling_buckets:
                 token_list.append(
-                    self._torch_sampling_func(
-                        self._sampling_logits_cuda[indices, :], temp, top_k, top_p
-                    )
+                    self._torch_sampling_func(last_token_logits[indices, :], temp, top_k, top_p)
                 )
                 indices_list.append(torch.tensor(indices))
 
             # Single write to the output tensor.
             sampled_tokens = torch.cat(token_list, dim=0)
             sampled_indices = torch.cat(indices_list, dim=0)
             self._sampled_tokens_cuda[sampled_indices] = sampled_tokens
+        elif self._sampling_backend == "flashinfer":
+            # Copy last_token_logits to contiguous buffer.
+            self._sampling_logits_cuda[: self._active_request_count].copy_(
+                last_token_logits, non_blocking=True
+            )
+            raise NotImplementedError
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.
@@ -850,10 +853,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._dynamic_step_sample_bookkeeping()
-        self._dynamic_step_sample_logits()
-
         return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
+        self._dynamic_step_sample_bookkeeping()
+        self._dynamic_step_sample_logits(logits)
 
         log_probs = None
         top_n_logprobs = None