Graphed sampling

Author: tdene

megatron/core/inference/engines/dynamic_engine.py

@@ -298,6 +298,10 @@ def create_cuda_graphs(self, reset_context: bool = True):
             # Forward pass -> logits.
             controller._dynamic_step_forward_logits(input_ids, position_ids)
 
+            # Sampling.
+            if controller._sampling_backend not in ["torch"]:
+                controller._dynamic_step_sample_logits()
+
             context.reset()
 
         # Memory usage.

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -99,15 +99,17 @@ def _init_dynamic_sampling_tensors(self):
         max_requests = context.max_total_requests
 
         model_config = get_model_config(self.inference_wrapped_model.model)
-        sampling_backend = model_config.sampling_backend
+        self._sampling_backend = model_config.sampling_backend
+        self._enable_cuda_graph = model_config.enable_cuda_graph
+        if self._enable_cuda_graph:
+            self._recording_graph: bool = False
+            self._sampling_graph_map: Dict[int, torch.cuda.CUDAGraph] = {}
 
         device = torch.cuda.current_device()
         logits_dtype = self.inference_wrapped_model.inference_wrapper_config.params_dtype
         # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
         vocab_size = self.inference_wrapped_model.inference_wrapper_config.padded_vocab_size
 
-        self._sampling_backend = "torch"
-
         # Keep track of request metadata.
         self._active_request_count = None
         self._request_metadata: Dict[str, Tensor] = {}
@@ -590,10 +592,12 @@ def _dynamic_step_context_init(
 
         # Get flat tokens, position ids.
         if construct_graph_dimensions is not None:
+            self._recording_graph = True
             return context.current_input_and_position_ids(
                 num_warmup_tokens=construct_graph_dimensions.token_count
             )
         else:
+            self._recording_graph = False
             return context.current_input_and_position_ids()
 
     def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) -> Tensor:
@@ -679,9 +683,6 @@ def _dynamic_step_sample_logits(self):
         """Sample tokens from logits for dynamic batching."""
         # 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.
-        context = self.inference_wrapped_model.inference_context
-        active_request_count = context.total_request_count - context.paused_request_count
-
         if self._sampling_backend == "torch":
             # Concatenate the outputs once to prevent repeated small writes.
             token_list = []
@@ -699,7 +700,24 @@ def _dynamic_step_sample_logits(self):
             sampled_indices = torch.cat(indices_list, dim=0)
             self._sampled_tokens_cuda.index_copy_(0, sampled_indices, sampled_tokens)
         elif self._sampling_backend == "flashinfer":
-            self.flashinfer_sampling_func(active_request_count)
+            if self._enable_cuda_graph:
+                if self._recording_graph:
+                    g = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(g):
+                        self.flashinfer_sampling_func(self._active_request_count)
+                    self._sampling_graph_map[self._active_request_count] = g
+                else:
+                    # Find the appropriate graph to replay.
+                    graph_size = next(
+                        (
+                            size
+                            for size in self._sampling_graph_map.keys()
+                            if size >= self._active_request_count
+                        ),
+                        None,
+                    )
+                    assert graph_size is not None
+                    self._sampling_graph_map[graph_size].replay()
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.