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,12 @@ 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.request_metadata_types is None:
+            self.request_metadata_types = DynamicInferenceRequest.get_metadata_types()
+            # FlashInfer reads sampling params directly on GPU.
+            if self.sampling_backend == 'flashinfer':
+                self.request_metadata_types = [
+                    (label, dtype, True if label in {"temperature", "top_k", "top_p"} else on_gpu)
+                    for label, dtype, on_gpu in self.request_metadata_types
+                ]

megatron/core/inference/contexts/dynamic_context.py

@@ -1069,6 +1069,28 @@ def get_active_request_count(self):
         """Returns the current number of active requests."""
         return self.total_request_count - self.paused_request_count
 
+    def _encode_sampling_metadata(self, label: str, value: Tensor) -> Tensor:
+        """Remap disabled-filter sentinels to FlashInfer passthrough values.
+
+        Megatron's user-facing convention uses ``top_k == 0`` and ``top_p == 0.0``
+        to mean "no filter." FlashInfer's sampling kernels instead treat a large
+        ``top_k`` and ``top_p == 1.0`` as the identity operation. This helper
+        translates at the data-entry boundary so the sampling path never has to
+        reason about two encodings, and ``build_active_slices`` can be optimized
+        without silently corrupting sampling semantics.
+        """
+        if self.config.sampling_backend != 'flashinfer':
+            return value
+        if label == 'top_k':
+            return torch.where(
+                value == 0,
+                torch.full_like(value, torch.iinfo(value.dtype).max),
+                value,
+            )
+        if label == 'top_p':
+            return torch.where(value == 0.0, torch.ones_like(value), value)
+        return value
+
     def build_active_slices(self, batch_size: int):
         """Build the active slices of specific tensors. This is run on every forward step.
 
@@ -1459,9 +1481,12 @@ def add_dummy_requests_parallel(
         self.request_kv_length_offsets[request_slice] = 0
         self.request_kv_block_counts[request_slice] = block_counts
         for i, (label, dtype, _) in enumerate(self.request_metadata_types):
-            self.request_metadata[label][request_slice] = torch.tensor(
+            value = torch.tensor(
                 metadata_cols[i], dtype=dtype, device=torch.cuda.current_device()
             )
+            self.request_metadata[label][request_slice] = self._encode_sampling_metadata(
+                label, value
+            )
 
         dummy_block_idx = self.kv_block_allocator.dummy_block_idx
         self.request_last_kv_block_id[request_slice] = dummy_block_idx
@@ -2242,7 +2267,7 @@ def add_request(
                     dtype=self.request_metadata[label].dtype,
                 )
 
-            self.request_metadata[label][current_id] = m
+            self.request_metadata[label][current_id] = self._encode_sampling_metadata(label, m)
 
         # Handle length and block assignments.
         self.request_query_lengths[current_id] = effective_prefill_chunk_length

megatron/core/inference/engines/dynamic_engine.py

@@ -389,6 +389,26 @@ def create_cuda_graphs(self, reset_context: bool = True):
             # Forward pass -> logits.
             controller._dynamic_step_forward_logits(input_ids, position_ids)
 
+            # Capture sampling CUDA graphs (plain + filtered) for all dimensions.
+            if controller._enable_cuda_graph:
+                n = cuda_graph_batch_dimension.req_count
+                md = context.active_request_metadata
+                last_token_indices = (
+                    None
+                    if context.config.materialize_only_last_token_logits
+                    else context.active_request_last_token_idxs
+                )
+                for filtered in (False, True):
+                    md["temperature"][:n].fill_(1.0)
+                    md["top_k"][:n].fill_(controller.vocab_size)
+                    md["top_p"][:n].fill_(1.0)
+                    controller._flashinfer_sample(
+                        controller._all_logits_cuda.squeeze(0),
+                        n,
+                        filtered=filtered,
+                        last_token_indices=last_token_indices,
+                    )
+
             context.reset()
 
         # Disable inference dispatcher after graph capture

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -50,7 +50,14 @@
 except ImportError:
     HAVE_TE = False
 
+try:
+    import flashinfer
+except ImportError:
+    flashinfer = None
+
 from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
+from megatron.core.inference.utils import CUDAGraphCache
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
 
 
 # pylint: disable=line-too-long
@@ -128,13 +135,15 @@ def _init_dynamic_sampling_tensors(self):
         device = torch.cuda.current_device()
         logits_dtype = self.inference_wrapped_model.config.params_dtype
 
-        self._sampling_backend = "torch"
-        self._enable_cuda_graph = False
+        self._sampling_backend = context.config.sampling_backend
+        self._enable_cuda_graph = (
+            self._sampling_backend == "flashinfer" and self.model_config.cuda_graph_impl == "local"
+        )
 
         # Initialize bookkeeping tensors.
-        if self._enable_cuda_graph:
-            self._all_logits_cuda = torch.empty(
-                (1, max_logits, vocab_size), dtype=logits_dtype, device=device
+        if self._sampling_backend == "flashinfer":
+            self._all_logits_cuda = torch.zeros(
+                (1, max_logits, self.vocab_size), dtype=logits_dtype, device=device
             )
         else:
             self._all_logits_cuda = None
@@ -146,9 +155,10 @@ def _init_dynamic_sampling_tensors(self):
         # Last accepted sequence indices for serial MTP computation
         self._last_accepted_seq_indices = None
 
-        # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
             self._torch_sampling_buckets: List[Tuple] = []
+        elif self._sampling_backend == "flashinfer":
+            self._sampling_cuda_graphs = CUDAGraphCache()
 
         self._init_mtp_sampling_tensor()
 
@@ -670,7 +680,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
             )
 
         # Copy logits to contiguous buffer.
-        if self._enable_cuda_graph:
+        if self._enable_cuda_graph or self._sampling_backend == "flashinfer":
             self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits)
         else:
             self._all_logits_cuda = logits
@@ -698,6 +708,25 @@ def _dynamic_step_sample_bookkeeping(self):
             self._torch_sampling_buckets = [
                 (indices, *sampling_params) for sampling_params, indices in bucket_map.items()
             ]
+        elif self._sampling_backend == "flashinfer":
+            n = active_request_count
+            padded_n = context.padded_active_request_count
+            md = context.active_request_metadata
+
+            if self._enable_cuda_graph and context.using_cuda_graph_this_step():
+                # Pad slots beyond active with safe defaults so the kernel
+                # doesn't trip on stale values in unused request slots.
+                if padded_n > n:
+                    md["temperature"][n:padded_n].fill_(1.0)
+                    md["top_k"][n:padded_n].fill_(self.vocab_size)
+                    md["top_p"][n:padded_n].fill_(1.0)
+
+            # NOTE: This currently makes a batch-level decision.
+            # If any request in the batch uses topk/topp, the whole batch runs the slower kernel.
+            # Per-request decisions are possible, but not a current priority.
+            self._fi_any_filtered_gpu = (md["top_k"][:n] < self.vocab_size).any() | (
+                md["top_p"][:n] < 1.0
+            ).any()
 
     def _rewind_kv_cache(self):
         """Update the KV cache bookkeeping for speculative decoding.
@@ -820,20 +849,23 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
         Returns:
             Tensor: Sampled tokens of shape [num_requests].
         """
-        spec_token_list = []
-        indices_list = []
-        for request_indices, temp, top_k, top_p in self._torch_sampling_buckets:
-            request_indices_tensor = torch.tensor(
-                request_indices, device=logits_2d.device, dtype=torch.long
-            )
-            spec_token_list.append(
-                self._torch_sampling_func(logits_2d[request_indices_tensor, :], temp, top_k, top_p)
+        n = logits_2d.shape[0]
+        spec_tokens = torch.empty(n, device=logits_2d.device, dtype=torch.int64)
+
+        if self._sampling_backend == "flashinfer":
+            self._flashinfer_sample(
+                logits_2d, n, self._fi_any_filtered_gpu.item(), eager=True
             )
-            indices_list.append(request_indices_tensor)
+            spec_tokens[:n] = self._sampled_tokens_cuda[:n]
+        else:
+            for request_indices, temp, top_k, top_p in self._torch_sampling_buckets:
+                request_indices_tensor = torch.tensor(
+                    request_indices, device=logits_2d.device, dtype=torch.long
+                )
+                spec_tokens[request_indices_tensor] = self._torch_sampling_func(
+                    logits_2d[request_indices_tensor, :], temp, top_k, top_p
+                )
 
-        spec_tokens = torch.empty(logits_2d.shape[0], device=logits_2d.device, dtype=torch.int64)
-        for tokens, indices in zip(spec_token_list, indices_list):
-            spec_tokens[indices] = tokens
         return spec_tokens
 
     def _compute_serial_mtp_and_sample(self):
@@ -982,6 +1014,23 @@ def _sample_speculative_logits(
             repeats,
         )
 
+        if self._sampling_backend == "flashinfer":
+            # Expand per-request sampling params to per-token via gather. Disabled
+            # filters were already remapped to passthrough values by bookkeeping.
+            active_n = len(request_in_prefill_status_tensor)
+            md = self.inference_wrapped_model.inference_context.active_request_metadata
+            temp_per_token = md["temperature"][:active_n][token_to_request_index]
+            top_k_per_token = md["top_k"][:active_n][token_to_request_index]
+            top_p_per_token = md["top_p"][:active_n][token_to_request_index]
+
+            scaled = required_logits.to(torch.float32, copy=True)
+            scaled.div_(temp_per_token.unsqueeze(1))
+            probs = torch.softmax(scaled, dim=-1)
+            output_tokens = flashinfer.sampling.top_k_top_p_sampling_from_probs(
+                probs, top_k_per_token, top_p_per_token, generator=self.sampling_rng
+            )
+            return output_tokens, repeats
+
         output_tokens_jumbled_list = []
         token_order_list = []
 
@@ -1173,15 +1222,74 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
             dim=1
         )
 
+    def _flashinfer_sample(
+        self,
+        logits: torch.Tensor,
+        n: int,
+        filtered: bool,
+        eager: bool = False,
+        last_token_indices: Optional[torch.Tensor] = None,
+    ) -> None:
+        """Sample tokens using FlashInfer.
+
+        Args:
+            logits (torch.Tensor): Logits of shape [>=n, vocab_size].
+            n (int): Number of requests to sample for (padded for graph path).
+            filtered (bool): Whether to use top_k_top_p filtering.
+            eager (bool): If True, run without CUDA graph capture/replay.
+            last_token_indices (Optional[torch.Tensor]): Indices to gather from logits before
+                sampling. Must be a view on a static context tensor for graphability.
+        """
+        pool = CudaGraphManager.global_mempool if not eager else None
+        md = self.inference_wrapped_model.inference_context.active_request_metadata
+
+        for _ in self._sampling_cuda_graphs(n, filtered, pool=pool, eager=eager):
+            if last_token_indices is None:
+                scaled = logits[:n].to(torch.float32, copy=True)
+            else:
+                scaled = logits[last_token_indices[:n], :].to(torch.float32, copy=True)
+            scaled.div_(md["temperature"][:n].unsqueeze(1))
+            probs = torch.softmax(scaled, dim=-1)
+            if filtered:
+                self._sampled_tokens_cuda[:n].copy_(
+                    flashinfer.sampling.top_k_top_p_sampling_from_probs(
+                        probs,
+                        md["top_k"][:n],
+                        md["top_p"][:n],
+                        generator=self.sampling_rng,
+                    )
+                )
+            else:
+                self._sampled_tokens_cuda[:n].copy_(
+                    flashinfer.sampling.sampling_from_probs(probs, generator=self.sampling_rng)
+                )
+
     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.
 
-        # Last token logits.
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
 
+        if self._sampling_backend == "flashinfer":
+            use_graph = self._enable_cuda_graph and context.using_cuda_graph_this_step()
+            n = context.padded_active_request_count if use_graph else active_request_count
+            last_token_indices = (
+                None
+                if context.config.materialize_only_last_token_logits
+                else context.active_request_last_token_idxs
+            )
+            self._flashinfer_sample(
+                self._all_logits_cuda.squeeze(0),
+                n,
+                self._fi_any_filtered_gpu.item(),
+                eager=not use_graph,
+                last_token_indices=last_token_indices,
+            )
+            return
+
+        # Last token logits.
         if context.config.materialize_only_last_token_logits:
             # When materialize_only_last_token_logits is true, last_token_logits is
             # already called in the forward pass of GPT.
@@ -1828,6 +1936,8 @@ async def async_generate_output_tokens_dynamic_batch(
                 else None
             )
 
+            self._dynamic_step_sample_bookkeeping()
+
             # Enable routing recording before forward pass if routing replay is enabled
             config = self.inference_wrapped_model.model.config
             if config.moe_enable_routing_replay:
@@ -1859,8 +1969,6 @@ async def async_generate_output_tokens_dynamic_batch(
         with torch.inference_mode():
             return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
 
-            self._dynamic_step_sample_bookkeeping()
-
             if self.num_speculative_tokens > 0:
                 # Phase 1: Verify speculative tokens using base logits only.
                 self._dynamic_step_sample_logits_and_verify_tokens(input_ids)

megatron/inference/utils.py

@@ -362,6 +362,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args):
         logging_step_interval=args.inference_logging_step_interval,
         num_speculative_tokens=args.num_speculative_tokens,
         use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives,
+        sampling_backend=args.inference_dynamic_batching_sampling_backend,
     )
 
 

megatron/training/arguments.py

@@ -944,6 +944,15 @@ def validate_args(args, defaults={}):
     ):
         raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4")
 
+    if args.inference_dynamic_batching_sampling_backend == 'flashinfer':
+        try:
+            import flashinfer  # noqa: F401
+        except ImportError:
+            raise ValueError(
+                "--inference-dynamic-batching-sampling-backend=flashinfer "
+                "requires flashinfer to be installed."
+            )
+
     if args.use_megatron_fsdp:
         # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions.
         #       Please use `use_megatron_fsdp` instead, as all functionality will be migrated there.
@@ -1908,6 +1917,10 @@ def _add_inference_args(parser):
     group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count',
                        type=int, default=16,
                        help='Number of mixed prefill requests to capture in a cuda graph.')
+    group.add_argument('--inference-dynamic-batching-sampling-backend',
+                       type=str, default='torch',
+                       choices=['torch', 'flashinfer'],
+                       help='Which sampling kernels to use during inference.')
     group.add_argument('--inference-logging-step-interval', type=int, default=0,
                        help='Step interval for logging inference metrics. '
                             'Default to 0 to disable inference logging.')