FlashInfer sampling

megatron/core/inference/config.py

@@ -2,7 +2,7 @@
 
 from dataclasses import InitVar, dataclass
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import List, Literal, Optional, Tuple
 
 import torch
 
@@ -297,6 +297,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]]] = None
     """
     A list of the per-request metadata types to track. Each entry is a tuple
@@ -320,3 +323,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.sampling_backend == 'flashinfer':
+            try:
+                import flashinfer  # noqa: F401
+            except ImportError as e:
+                raise ImportError(
+                    "sampling_backend='flashinfer' requires the flashinfer package; "
+                    "install it or set sampling_backend='torch'."
+                ) from e

megatron/core/inference/contexts/dynamic_context.py

@@ -928,16 +928,20 @@ def initialize_all_tensors(self) -> None:
         # transferred to GPU each step via transfer_bookkeeping_to_gpu().
         # Layout matches ContextGPUView._buf so a single cudaMemcpyAsync
         # suffices. Int64 token fields come first (8-byte aligned automatically),
-        # then int32 token fields, then int32 request-staging fields.
-        #   token_to_input_ids                         (int64, max_tokens)
-        #   token_to_pos_ids                           (int64, max_tokens)
-        #   token_to_block_idx                         (int32, max_tokens)
-        #   token_to_local_position_within_kv_block    (int32, max_tokens)
-        #   token_to_request_idx                       (int32, max_tokens)
-        #   token_to_position_in_request               (int32, max_tokens)
-        #   request_in_prefill_status  (staging)       (int32, max_requests)
-        #   request_query_lengths      (staging)       (int32, max_requests)
-        #   request_kv_length_offsets  (staging)       (int32, max_requests)
+        # then int32 token fields, then int32/float32 request-staging fields.
+        #   token_to_input_ids                         (int64,   max_tokens)
+        #   token_to_pos_ids                           (int64,   max_tokens)
+        #   token_to_block_idx                         (int32,   max_tokens)
+        #   token_to_local_position_within_kv_block    (int32,   max_tokens)
+        #   token_to_request_idx                       (int32,   max_tokens)
+        #   token_to_position_in_request               (int32,   max_tokens)
+        #   request_in_prefill_status        (staging) (int32,   max_requests)
+        #   request_query_lengths            (staging) (int32,   max_requests)
+        #   request_kv_length_offsets        (staging) (int32,   max_requests)
+        #   temperature                      (staging) (float32, max_requests)
+        #   top_k                            (staging) (int32,   max_requests)
+        #   top_p                            (staging) (float32, max_requests)
+        #   active_request_last_token_idxs   (alias)   (int32,   max_requests)
         #
         # Token fields are aliased with the source-of-truth attributes
         # (`self.token_to_input_ids`, etc.) because the forward pass reads
@@ -948,7 +952,8 @@ def initialize_all_tensors(self) -> None:
         # slice from the persistent `request_*` tensors above.
         _tok_int64_bytes = self.max_tokens * 8
         _tok_int32_bytes = self.max_tokens * 4
-        _req_int32_bytes = self.max_requests * 4
+        # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields).
+        _req_4byte_bytes = self.max_requests * 4
         # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata. max_bs == max_requests.
         _mha_query_lengths_bytes = self.max_requests * 4
@@ -983,7 +988,7 @@ def initialize_all_tensors(self) -> None:
         _total_bytes = (
             2 * _tok_int64_bytes
             + 4 * _tok_int32_bytes
-            + 3 * _req_int32_bytes
+            + 7 * _req_4byte_bytes
             + _mha_query_lengths_bytes
             + _mha_cu_query_seq_lengths_bytes
             + _mha_kv_seq_lengths_bytes
@@ -1043,17 +1048,40 @@ def initialize_all_tensors(self) -> None:
         # CPU (refreshed from persistent tensors in transfer_bookkeeping_to_gpu);
         # read-only on GPU via matching slots in ContextGPUView._buf.
         self._staging_request_in_prefill_status = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
         self._staging_request_query_lengths = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
         self._staging_request_kv_length_offsets = self._cpu_bookkeeping_buf[
-            _off : _off + _req_int32_bytes
+            _off : _off + _req_4byte_bytes
         ].view(torch.int32)
-        _off += _req_int32_bytes
+        _off += _req_4byte_bytes
+
+        # Sampling-parameter staging slots, refreshed from `active_request_metadata`
+        # in transfer_bookkeeping_to_gpu(). FlashInfer reads these via
+        # `gpu_view.{temperature, top_k, top_p}`.
+        self._staging_temperature = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.float32
+        )
+        _off += _req_4byte_bytes
+        self._staging_top_k = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.int32
+        )
+        _off += _req_4byte_bytes
+        self._staging_top_p = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+            torch.float32
+        )
+        _off += _req_4byte_bytes
+
+        # Per-request last-token row indices. Aliased with the matching gpu_view slot:
+        # build_active_slices/pad_active_slices populate this CPU view.
+        self.active_request_last_token_idxs = self._cpu_bookkeeping_buf[
+            _off : _off + _req_4byte_bytes
+        ].view(torch.int32)
+        _off += _req_4byte_bytes
 
         # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
         max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1)
@@ -1362,6 +1390,13 @@ def build_active_slices(self, batch_size: int):
                 self.request_metadata[label][padded_slice], non_blocking=True
             )
 
+        torch.cumsum(
+            self.request_query_lengths[padded_slice],
+            dim=0,
+            out=self.active_request_last_token_idxs[:batch_size],
+        )
+        self.active_request_last_token_idxs[:batch_size].sub_(1)
+
     def pad_active_slices(self):
         """Pad the active slices of specific tensors."""
         active_request_count = self.total_request_count - self.paused_request_count
@@ -1388,6 +1423,15 @@ def pad_active_slices(self):
 
         self.active_logit_idxs[active_decode_token_count + active_prefill_count :].zero_()
 
+        padding_request_slice = slice(active_request_count, self.padded_active_request_count)
+
+        # Sampling metadata: pad with neutral defaults, so that the kernel early-exits.
+        self.active_request_metadata["temperature"][padding_request_slice].fill_(1.0)
+        self.active_request_metadata["top_k"][padding_request_slice].fill_(0)
+        self.active_request_metadata["top_p"][padding_request_slice].fill_(0.0)
+        # Padded gather indices fan in to row 0 harmlessly when used by FlashInfer.
+        self.active_request_last_token_idxs[padding_request_slice].fill_(0)
+
     def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
         """Append to KV cache.
 
@@ -2226,7 +2270,7 @@ def transfer_bookkeeping_to_gpu(self) -> None:
         All copies use non_blocking=True with pinned CPU memory. CUDA stream
         ordering guarantees the forward pass sees completed transfers.
 
-        The 9 bookkeeping fields are backed by one contiguous pinned CPU buffer
+        The bookkeeping fields are backed by one contiguous pinned CPU buffer
         and one contiguous GPU buffer; a single cudaMemcpyAsync suffices.
         Request-level staging slots are refreshed from the persistent CPU
         tensors immediately before the H2D (GPU reads them at `[:n_active]`
@@ -2237,16 +2281,25 @@ def transfer_bookkeeping_to_gpu(self) -> None:
         padded_active = max(n_active, self.padded_active_request_count)
 
         # Refresh request-level staging slots from the persistent CPU source.
-        # CPU-to-CPU slice assignment on pinned memory (~7.5 KB total for 3
-        # int32 fields at max_requests=624). Negligible vs. the launch overhead
-        # we save by merging 9 H2D memcpys into 1.
+        # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6
+        # 4-byte fields at max_requests=624). Negligible vs. the launch overhead
+        # we save by merging the H2D memcpys into 1.
         self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[
             active_slice
         ]
         self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice]
         self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[
             active_slice
         ]
+        # Sampling-parameter staging slots: read from `active_request_metadata`,
+        # which `build_active_slices` + `pad_active_slices` already populated for
+        # `[:padded_active]` (active values + neutral padding defaults).
+        self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][
+            :padded_active
+        ]
+        self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active]
+        self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active]
+
         # Full-iteration CUDA graphs may have captured GPU consumers with the
         # padded graph request count. Keep those padded staging rows bounded so
         # graph replay never builds indices from stale request lengths.

megatron/core/inference/contexts/gpu_view.py

@@ -17,11 +17,11 @@ class ContextGPUView:
         ``context.foo``      -> CPU (source of truth, used by bookkeeping)
         ``context.gpu_view.foo`` -> GPU (snapshot, used by forward pass)
 
-    Layout note: the 9 bookkeeping fields are backed by a single contiguous
+    Layout note: the bookkeeping fields are backed by a single contiguous
     ``uint8`` buffer (``self._buf``). Each field is a ``view(dtype)`` onto a
     slice of that buffer. This matches the pinned-CPU-buffer layout in
     :class:`DynamicInferenceContext` so that the per-step H2D transfer is a
-    single ``cudaMemcpyAsync`` instead of nine small ones.
+    single ``cudaMemcpyAsync`` instead of one per field.
     """
 
     def __init__(
@@ -39,7 +39,10 @@ def __init__(
         #   max_mamba_chunks == 0).
         tok_int64_bytes = max_tokens * 8  # 2 fields of int64 = 8 bytes/elem
         tok_int32_bytes = max_tokens * 4  # 4 fields of int32 = 4 bytes/elem
-        req_int32_bytes = max_requests * 4  # 3 fields of int32
+        # Request-level fields are all 4 bytes wide. 3 int32 (in_prefill_status,
+        # query_lengths, kv_length_offsets) + 1 int32 (top_k) + 2 float32
+        # (temperature, top_p) + 1 int32 (active_request_last_token_idxs) = 7 fields.
+        req_4byte_bytes = max_requests * 4
 
         # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata
         # (only one is active per step, so sharing storage is fine).
@@ -90,7 +93,7 @@ def __init__(
         total_bytes = (
             2 * tok_int64_bytes
             + 4 * tok_int32_bytes
-            + 3 * req_int32_bytes
+            + 7 * req_4byte_bytes
             + mha_query_lengths_bytes
             + mha_cu_query_seq_lengths_bytes
             + mha_kv_seq_lengths_bytes
@@ -128,12 +131,28 @@ def __init__(
         off += tok_int32_bytes
 
         # Request-level tensors (consumed by sampling, log-probs, speculative verification, MTP).
-        self.request_in_prefill_status = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
-        self.request_query_lengths = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
-        self.request_kv_length_offsets = self._buf[off : off + req_int32_bytes].view(torch.int32)
-        off += req_int32_bytes
+        self.request_in_prefill_status = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.request_query_lengths = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.request_kv_length_offsets = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        # Sampling parameters (consumed by FlashInfer sampling).
+        # Mirror the active slice of `active_request_metadata[{label}]`;
+        # padded slots get neutral defaults from `pad_active_slices` (T=1.0, top_k=0, top_p=0.0).
+        self.temperature = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+        off += req_4byte_bytes
+        self.top_k = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+        off += req_4byte_bytes
+        self.top_p = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+        off += req_4byte_bytes
+        # Per-request last-token row indices (consumed by sampling kernels as `gather_indices`).
+        # The CPU side of this slot IS `context.active_request_last_token_idxs`,
+        # populated by `build_active_slices` and `pad_active_slices`.
+        self.active_request_last_token_idxs = self._buf[off : off + req_4byte_bytes].view(
+            torch.int32
+        )
+        off += req_4byte_bytes
 
         # MHA flash-attention metadata (shared between GraphedMHAMetadata and
         # NonGraphedMHAMetadata — only one is active per step).

megatron/core/inference/engines/dynamic_engine.py

@@ -395,6 +395,12 @@ def create_cuda_graphs(self, reset_context: bool = True):
             with torch.inference_mode():
                 controller._dynamic_step_forward_logits(input_ids, position_ids)
 
+                if controller._sampling_backend == "flashinfer":
+                    if controller.num_speculative_tokens > 0:
+                        controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
+                    else:
+                        controller._dynamic_step_sample_logits()
+
                 # MTP CUDA graph warmup for this batch dimension.
                 if mtp_warmup_enabled:
                     n = cuda_graph_batch_dimension.req_count

megatron/core/inference/sampling/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.inference.sampling.base import Sampling
+from megatron.core.inference.sampling.flashinfer_sampling import FlashInferSampling
+from megatron.core.inference.sampling.torch_sampling import TorchSampling
+
+__all__ = ["Sampling", "TorchSampling", "FlashInferSampling"]

megatron/core/inference/sampling/base.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+
+
+class Sampling(ABC):
+    """Abstract base for inference sampling backends.
+
+    Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`.
+    """
+
+    @abstractmethod
+    def sample_kernel(
+        self,
+        logits: Tensor,
+        n: int,
+        context,
+        *,
+        gather_indices: Optional[Tensor] = None,
+        token_to_request_index: Optional[Tensor] = None,
+        eager: bool = False,
+        cache_key: Any = None,
+    ) -> Tensor:
+        """Sample `n` tokens from `logits` and return them.
+
+        Args:
+            logits: Logits tensor of shape `[>=n, vocab_size]`.
+            n: Number of rows to sample.
+            context: The active DynamicInferenceContext.
+            gather_indices: If provided, only sample from `logits[gather_indices[:n], :]`.
+            token_to_request_index: Per-token request mapping; when set, sampling
+                parameters are gathered per-token instead of per-request.
+            eager, cache_key: Consumed by `CudaGraphManager` when it wraps this kernel.
+
+        Returns:
+            Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
+        """
+        ...
+
+    def sample_speculative(
+        self,
+        required_logits: Tensor,
+        num_decode: int,
+        num_prefill: int,
+        num_speculative_tokens: int,
+        context,
+        *,
+        gather_indices: Optional[Tensor] = None,
+        eager: bool = False,
+        cache_key: Any = None,
+    ) -> Tensor:
+        """Sample tokens for the speculative-verify path.
+
+        Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests contribute 1.
+        Builds the per-token request mapping and dispatches to `sample_kernel`.
+        The `sample_kernel` is forced eager so its own `CudaGraphManager` wrapper does not fire.
+
+        When `gather_indices` is supplied, the kernel selects via `logits[gather_indices[:n], :]`.
+        When `gather_indices` is None, `required_logits` is expected to be already pre-gathered to
+        the layout described above (e.g. when `materialize_only_last_token_logits=True` upstream).
+        """
+        # CudaGraphManager consumes these args, if it exists.
+        del eager, cache_key
+
+        n_spec = num_speculative_tokens
+        num_decode_tokens = num_decode * (1 + n_spec)
+        num_tokens = num_decode_tokens + num_prefill
+        device = required_logits.device
+
+        token_to_request_index = torch.cat(
+            [
+                torch.arange(num_decode, device=device).repeat_interleave(
+                    1 + n_spec, output_size=num_decode_tokens
+                ),
+                torch.arange(num_decode, num_decode + num_prefill, device=device),
+            ]
+        )
+        return self.sample_kernel(
+            required_logits,
+            num_tokens,
+            context,
+            gather_indices=gather_indices,
+            token_to_request_index=token_to_request_index,
+            eager=True,
+        )