Implement FlashInfer sampling backend

Author: tdene

megatron/core/inference/config.py

@@ -2,10 +2,11 @@
 
 from dataclasses import InitVar, dataclass
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import List, Literal, Optional, Tuple
 
 import torch
 
+from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_attr_wrapped_model
@@ -297,6 +298,9 @@ class InferenceConfig:
     Defaults to 0, which means no logging.
     """
 
+    sampling_backend: Literal['torch', 'flashinfer'] = 'torch'
+    """Which sampling kernels to use during inference."""
+
     request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None
     """
     A list of the per-request metadata types to track. Each entry is a tuple
@@ -320,3 +324,16 @@ def __post_init__(self, verbose: bool):
                 f"prefix_caching_routing_alpha must be in [0, 1], "
                 f"got {self.prefix_caching_routing_alpha}"
             )
+
+        if self.sampling_backend == 'flashinfer':
+            try:
+                import flashinfer  # noqa: F401
+            except ImportError:
+                raise ValueError(
+                    "sampling_backend='flashinfer' requires flashinfer to be installed."
+                )
+
+        if self.request_metadata_types is None:
+            self.request_metadata_types = DynamicInferenceRequest.get_metadata_types(
+                sampling_backend=self.sampling_backend
+            )

megatron/core/inference/contexts/dynamic_context.py

@@ -1129,6 +1129,21 @@ def pad_active_slices(self):
         self.token_to_local_position_within_kv_block[padding_token_slice] = 0
         self.token_to_position_in_request[padding_token_slice] = 0
 
+        # Pad request-level sampling metadata for FlashInfer.  Padded slots
+        # must contain disabled-filter sentinels so the kernel doesn't trip on
+        # stale values.  The last-token gather also reads padded slots when
+        # materialize_only_last_token_logits is False; index 0 is always
+        # in-bounds for the logits buffer.
+        if self.config.sampling_backend == 'flashinfer':
+            n = self.total_request_count - self.paused_request_count
+            padded_n = self.padded_active_request_count
+            if padded_n > n:
+                self.active_request_metadata["temperature"][n:padded_n].fill_(1.0)
+                self.active_request_metadata["top_k"][n:padded_n].fill_(0)
+                self.active_request_metadata["top_p"][n:padded_n].fill_(0.0)
+                if not self.config.materialize_only_last_token_logits:
+                    self.active_request_last_token_idxs[n:padded_n].fill_(0)
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -2229,7 +2244,8 @@ def add_request(
 
         # Handle request metadata.
         assert (
-            req.get_metadata_types() == self.request_metadata_types
+            req.get_metadata_types(sampling_backend=self.config.sampling_backend)
+            == self.request_metadata_types
         ), "Request added to context with invalid metadata types"
         metadata = req.tracked_metadata
         metadata_types = req.get_metadata_types()

megatron/core/inference/engines/dynamic_engine.py

@@ -364,6 +364,8 @@ def create_cuda_graphs(self, reset_context: bool = True):
             unwrapped_model = controller.inference_wrapped_model.model
             set_inference_cuda_graphed_iteration_for_ep_inference(unwrapped_model)
 
+        model_config = controller.inference_wrapped_model.model.config
+
         tbar = enumerate(context.cuda_graph_batch_dimensions_list)
         if HAVE_TQDM:
             tbar = tqdm(tbar, total=len(context.cuda_graph_batch_dimensions_list))
@@ -392,13 +394,14 @@ def create_cuda_graphs(self, reset_context: bool = True):
                 controller._pre_forward_bookkeeping_stream.wait_stream(
                     torch.cuda.current_stream()
                 )
-                # Launch bookkeeping on a side stream so it overlaps with forward.
                 with torch.cuda.stream(controller._pre_forward_bookkeeping_stream):
+                    controller._dynamic_step_sample_bookkeeping()
                     controller._pre_forward_bookkeeping_event.record()
 
                 controller._dynamic_step_forward_logits(input_ids, position_ids)
 
                 controller._pre_forward_bookkeeping_event.synchronize()
+                controller._dynamic_step_sample_logits()
 
             context.reset()
 

megatron/core/inference/inference_request.py

@@ -472,24 +472,37 @@ def tracked_metadata(self) -> List[Any]:
         return [getattr(sp, field) for field, _, _ in self.get_metadata_types()]
 
     @staticmethod
-    def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
+    def get_metadata_types(
+        sampling_backend: str = 'torch',
+    ) -> List[Tuple[str, torch.dtype, bool]]:
         """Keeps track of all request metadata names, dtypes, and target device.
 
+        Args:
+            sampling_backend: Which sampling backend is in use.
+                FlashInfer needs sampling parameters directly on GPU.
+
         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
+        types = [
+            ("temperature", torch.float32, False),
+            ("top_k", torch.int32, False),
+            ("top_p", torch.float32, False),
             ("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_log_probs", torch.bool, False),
+            ("skip_prompt_log_probs", torch.bool, False),
+            ("top_n_logprobs", torch.int32, False),
         ]
+        if sampling_backend == 'flashinfer':
+            gpu_fields = {"temperature", "top_k", "top_p"}
+            types = [
+                (label, dtype, True if label in gpu_fields else on_gpu)
+                for label, dtype, on_gpu in types
+            ]
+        return types
 
     def add_event(
         self, type: DynamicInferenceEventType, payload: Optional[Any] = None