From ee3d7c5911e669e19c0fa510755ac124fd896612 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:28:18 -0500
Subject: [PATCH 01/23] Add sampling backend abstraction

---
 megatron/core/inference/sampling/__init__.py  |   6 +
 megatron/core/inference/sampling/base.py      | 112 +++++++++
 .../core/inference/sampling/torch_sampling.py | 179 ++++++++++++++
 .../text_generation_controller.py             | 229 +++---------------
 .../test_mtp_cuda_graph_inference.py          |   8 +-
 .../test_text_generation_controller.py        |  16 +-
 6 files changed, 345 insertions(+), 205 deletions(-)
 create mode 100644 megatron/core/inference/sampling/__init__.py
 create mode 100644 megatron/core/inference/sampling/base.py
 create mode 100644 megatron/core/inference/sampling/torch_sampling.py

diff --git a/megatron/core/inference/sampling/__init__.py b/megatron/core/inference/sampling/__init__.py
new file mode 100644
index 00000000000..97a3a9e327d
--- /dev/null
+++ b/megatron/core/inference/sampling/__init__.py
@@ -0,0 +1,6 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.inference.sampling.base import Sampling
+from megatron.core.inference.sampling.torch_sampling import TorchSampling
+
+__all__ = ["Sampling", "TorchSampling"]
diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
new file mode 100644
index 00000000000..2a06a486aa1
--- /dev/null
+++ b/megatron/core/inference/sampling/base.py
@@ -0,0 +1,112 @@
+# 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 `pre_forward_bookkeeping` (per-step setup) and `sample_kernel`
+    (the GPU kernel). CUDA graph wrapping, when applicable, is added by the subclass via
+    `CudaGraphManager`. The wrapper consumes `eager` and `cache_key` kwargs; concrete
+    subclasses without a wrapper still accept and ignore them.
+    """
+
+    @abstractmethod
+    def pre_forward_bookkeeping(self, context) -> None:
+        """Prepare sampling state before the forward pass."""
+        ...
+
+    @abstractmethod
+    def sample_kernel(
+        self,
+        logits: Tensor,
+        n: int,
+        output: Tensor,
+        context,
+        *,
+        eager: bool = False,
+        cache_key: Any = None,
+        gather_indices: Optional[Tensor] = None,
+        token_to_request_index: Optional[Tensor] = None,
+    ) -> None:
+        """Sample n tokens from logits into `output[:n]`.
+
+        `eager` and `cache_key` are consumed by the `CudaGraphManager` wrapper when one
+        is installed; unwrapped subclasses accept and ignore them.
+
+        Args:
+            logits: Logits tensor of shape `[>=n, vocab_size]`.
+            n: Number of rows to sample.
+            output: Destination buffer for sampled token ids.
+            context: The active DynamicInferenceContext.
+            eager: If True, skip CUDA graph capture/replay (consumed by the wrapper).
+            cache_key: Hashable key for runner lookup (consumed by the wrapper).
+            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.
+        """
+        ...
+
+    def sample(
+        self,
+        logits: Tensor,
+        n: int,
+        output: Tensor,
+        context,
+        *,
+        eager: bool = False,
+        cache_key: Any = None,
+        gather_indices: Optional[Tensor] = None,
+        token_to_request_index: Optional[Tensor] = None,
+    ) -> None:
+        """Sample `n` tokens, optionally with CUDA graph capture/replay."""
+        self.sample_kernel(
+            logits,
+            n,
+            output,
+            context,
+            eager=eager,
+            cache_key=cache_key,
+            gather_indices=gather_indices,
+            token_to_request_index=token_to_request_index,
+        )
+
+    def sample_speculative(
+        self,
+        required_logits: Tensor,
+        request_in_prefill_status: Tensor,
+        num_speculative_tokens: int,
+        context,
+    ) -> Tensor:
+        """Sample tokens for the speculative-verify path.
+
+        Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests
+        contribute one row. Builds the per-token request mapping and dispatches to
+        `sample_kernel(eager=True)` so the call is safe inside an outer captured kernel.
+        """
+        repeats = torch.where(
+            request_in_prefill_status == 0, 1 + num_speculative_tokens, 1
+        )
+        token_to_request_index = torch.repeat_interleave(
+            torch.arange(
+                len(request_in_prefill_status),
+                device=request_in_prefill_status.device,
+            ),
+            repeats,
+        )
+        n = token_to_request_index.shape[0]
+        output_tokens = torch.empty(n, device=required_logits.device, dtype=torch.int64)
+        self.sample_kernel(
+            required_logits,
+            n,
+            output_tokens,
+            context,
+            eager=True,
+            token_to_request_index=token_to_request_index,
+        )
+        return output_tokens
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
new file mode 100644
index 00000000000..18ed84fa1e6
--- /dev/null
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -0,0 +1,179 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from collections import defaultdict
+from typing import Any, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from megatron.core.inference.sampling.base import Sampling
+
+
+class TorchSampling(Sampling):
+    """Sampling via bucketed `torch.multinomial`.
+
+    Groups requests into unique buckets by `(temperature, top_k, top_p)` for separate launches.
+    """
+
+    def __init__(self, rng: torch.Generator, vocab_size: int) -> None:
+        self._rng = rng
+        self._vocab_size = vocab_size
+        self._buckets: List[Tuple] = []
+        self._bucket_index_tensors: List[Tensor] = []
+
+    @staticmethod
+    def sample_from_logits(
+        last_token_logits: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        *,
+        generator: torch.Generator,
+        vocab_size: Optional[int] = None,
+    ) -> Tensor:
+        """Sample tokens from logits with temperature, top-k, and top-p filtering.
+
+        Shared between dynamic batching (`TorchSampling._sampling_func`) and static
+        batching (`TextGenerationController.sample_from_logits`).
+
+        Args:
+            last_token_logits: Logits of shape `[batch_size, vocab_size]`.
+            temperature: Temperature scaling factor.
+            top_k: Top-k filtering value (0 = disabled).
+            top_p: Top-p (nucleus) filtering value (0.0 = disabled).
+            generator: RNG used by `torch.multinomial`.
+            vocab_size: When provided, asserts `top_k < vocab_size` and clamps the
+                sampled ids to `[0, vocab_size - 1]`.
+
+        Returns:
+            Sampled token ids of shape `[batch_size]`.
+        """
+        assert isinstance(top_p, float)
+        assert isinstance(top_k, int)
+        assert not (
+            top_k > 0 and top_p > 0.0
+        ), "Cannot have top-p and top-k both greater than zero"
+        assert top_p <= 1.0, "top-p should be in (0,1]"
+
+        def modify_logits_for_top_k_filtering(logits, top_k):
+            """Set the logits for none top-k values to -inf."""
+            filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
+            logits.masked_fill_(filter_, float("-Inf"))
+
+        def modify_logits_for_top_p_filtering(logits, top_p):
+            """Set the logits for none top-p values to -inf."""
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+            cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+            filter_ = cumulative_probs > top_p
+            # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views;
+            # without clone, each write would corrupt the next read during the shift.
+            filter_[:, 1:] = filter_[:, :-1].clone()
+            filter_[..., 0] = 0
+
+            filter_ = filter_.scatter(1, sorted_indices, filter_)
+            logits.masked_fill_(filter_, float("-Inf"))
+
+        if top_k == 1:
+            return torch.argmax(last_token_logits, dim=-1)
+
+        # Clone needed: .div_() and masked_fill_() below modify in-place,
+        # which would mutate the caller's tensor without this clone.
+        last_token_logits = last_token_logits.clone()
+        if temperature != 1.0:
+            last_token_logits.div_(temperature)
+        if top_k > 1:
+            assert top_k <= last_token_logits.size(1), "top-k is larger than logit size."
+            if vocab_size:
+                assert top_k < vocab_size, "top-k is larger than vocab size."
+            modify_logits_for_top_k_filtering(last_token_logits, top_k)
+        elif top_p > 0.0:
+            modify_logits_for_top_p_filtering(last_token_logits, top_p)
+
+        probabilities = last_token_logits.softmax(dim=-1)
+        sampled = torch.multinomial(probabilities, num_samples=1, generator=generator).view(-1)
+
+        if vocab_size:
+            sampled = torch.clamp(sampled, min=0, max=(vocab_size - 1))
+
+        return sampled
+
+    def pre_forward_bookkeeping(self, context) -> None:
+        """Group active requests into sampling buckets by `(temperature, top_k, top_p)`.
+
+        Pre-computes a GPU index tensor per bucket so subsequent steps avoid H2D copies.
+        """
+        active_request_count = context.total_request_count - context.paused_request_count
+        md = context.active_request_metadata
+        device = torch.cuda.current_device()
+
+        bucket_map: dict = defaultdict(list)
+        temp = md["temperature"][:active_request_count].tolist()
+        top_k = md["top_k"][:active_request_count].tolist()
+        top_p = md["top_p"][:active_request_count].tolist()
+
+        for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
+            bucket_map[(t, k, p)].append(request_index)
+
+        self._buckets = [(indices, *params) for params, indices in bucket_map.items()]
+        self._bucket_index_tensors = [
+            torch.tensor(indices, device=device, dtype=torch.long)
+            for indices, *_ in self._buckets
+        ]
+
+    def sample_kernel(
+        self,
+        logits: Tensor,
+        n: int,
+        output: Tensor,
+        context,
+        *,
+        eager: bool = False,
+        cache_key: Any = None,
+        gather_indices: Optional[Tensor] = None,
+        token_to_request_index: Optional[Tensor] = None,
+    ) -> None:
+        """Sample by iterating over pre-computed sampling buckets.
+
+        Args:
+            logits: Logits tensor of shape `[>=n, vocab_size]`.
+            n: Number of rows to sample.
+            output: Destination buffer; the kernel writes to the rows selected per bucket.
+            context: The active DynamicInferenceContext (unused; kept for ABC parity).
+            eager: Accepted for API symmetry with FlashInfer; ignored (no wrapper here).
+            cache_key: Accepted for API symmetry; ignored.
+            gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
+            token_to_request_index: When set, the loop dispatches per-token rather than
+                per-request (used by the speculative path).
+        """
+        del eager, cache_key
+        if gather_indices is not None:
+            logits = logits[gather_indices[:n], :]
+
+        token_list = []
+        indices_list = []
+        for idx_tensor, (_, temp, top_k, top_p) in zip(
+            self._bucket_index_tensors, self._buckets
+        ):
+            if token_to_request_index is None:
+                row_indices = idx_tensor
+            else:
+                row_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0]
+            token_list.append(self._sampling_func(logits[row_indices, :], temp, top_k, top_p))
+            indices_list.append(row_indices)
+
+        sampled_tokens = torch.cat(token_list, dim=0)
+        sampled_indices = torch.cat(indices_list, dim=0)
+        output[sampled_indices] = sampled_tokens
+
+    def _sampling_func(
+        self, last_token_logits: Tensor, temperature: float, top_k: int, top_p: float
+    ) -> Tensor:
+        return TorchSampling.sample_from_logits(
+            last_token_logits,
+            temperature,
+            top_k,
+            top_p,
+            generator=self._rng,
+            vocab_size=self._vocab_size,
+        )
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index faa9e5babd6..68266fcc732 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -60,6 +60,7 @@
     HAVE_TE = False
 
 from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
+from megatron.core.inference.sampling import Sampling, TorchSampling
 from megatron.core.inference.text_generation_controllers.mtp_utils_pytorch import rewind_kv_cache
 from megatron.core.inference.text_generation_controllers.mtp_utils_triton import (
     mamba_state_selective_copy,
@@ -166,9 +167,8 @@ def _init_dynamic_sampling_tensors(self):
             self._all_logits_cuda = None
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
-        # Used for inefficient torch sampling.
-        if self._sampling_backend == "torch":
-            self._torch_sampling_buckets: List[Tuple] = []
+        # Sampling backend: dispatches per-step bookkeeping and the sampling kernel.
+        self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size)
 
         # Cache values that are constant across inference steps.
         self._unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
@@ -329,86 +329,19 @@ def _torch_sampling_func(
         top_p: float,
         vocab_size: Optional[int] = None,
     ):
-        """Samples the logits to generate outputs
+        """Static-batching sampler shim. Forwards to `TorchSampling.sample_from_logits`.
 
-        Given the logits of the last token, this function samples it
-        according to the parameters defined in sampling_params
-        and returns the samples. If sampling parameters top_n_logprobs > 0
-        at each step it also updates the top_n_logprobs dict.
-
-        Args:
-            last_token_logits (torch.Tensor): The last token logits. A tensor of
-                size [batch_size, vocab_size].
-            temperature (float): The temperature to use for sampling.
-            top_k (int): The top-k value to use for sampling.
-            top_p (float): The top-p value to use for sampling.
-            vocab_size (int): Obtained from the tokenizer. Defaults to None.
-
-        Returns:
-            sampled_logits (torch.Tensor): 1D tensor with [batch_size] elements
+        The dynamic-batching path goes through `self._sampling`; this method is kept
+        for the static `sample_from_logits` flow and for any callers that mock it.
         """
-        assert isinstance(top_p, float)
-        assert isinstance(top_k, int)
-        assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero"
-        assert top_p <= 1.0, "top-p should be in (0,1]"
-
-        def modify_logits_for_top_k_filtering(logits, top_k):
-            """Set the logits for none top-k values to -inf."""
-            filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits.masked_fill_(filter_, float("-Inf"))
-
-        def modify_logits_for_top_p_filtering(logits, top_p):
-            """Set the logits for none top-p values to -inf."""
-            # First sort and calculate cumulative sum of probabilities.
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-            cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
-
-            # Filteration based on the cumulative sum.
-            filter_ = cumulative_probs > top_p
-            # This shift by 1 is weird and I cannot justify it. This existed
-            # in the original implementation:
-            #   https://github.com/ari-holtzman/degen/blob/master/gen.py
-            # and I guess it is needed so keeping it for now.
-            # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views;
-            # without clone, each write would corrupt the next read during the shift.
-            filter_[:, 1:] = filter_[:, :-1].clone()
-            # Make sure we at least have one token to select from.
-            filter_[..., 0] = 0
-
-            # Fill in the filtered part
-            filter_ = filter_.scatter(1, sorted_indices, filter_)
-            logits.masked_fill_(filter_, float("-Inf"))
-
-        # Greedy sampling
-        if top_k == 1:
-            sampled_logits = torch.argmax(last_token_logits, dim=-1)
-        else:
-            # Clone needed: .div_() and masked_fill_() below modify in-place,
-            # which would mutate the caller's tensor without this clone.
-            last_token_logits = last_token_logits.clone()
-            if temperature != 1.0:
-                last_token_logits.div_(temperature)
-            if top_k > 1:
-                assert top_k <= last_token_logits.size(1), "top-k is larger than logit size."
-                if vocab_size:
-                    assert top_k < vocab_size, "top-k is larger than vocab size."
-                modify_logits_for_top_k_filtering(last_token_logits, top_k)
-
-            elif top_p > 0.0:
-                modify_logits_for_top_p_filtering(last_token_logits, top_p)
-
-            # After filtering, we need to recalculate the distribution.
-            probabilities = last_token_logits.softmax(dim=-1)
-
-            sampled_logits = torch.multinomial(
-                probabilities, num_samples=1, generator=self.sampling_rng
-            ).view(-1)
-
-            # If vocab size is provided, make sure the samples are in in the range [0, vocab-size).
-            if vocab_size:
-                sampled_logits = torch.clamp(sampled_logits, min=0, max=(vocab_size - 1))
-
-        return sampled_logits
+        return TorchSampling.sample_from_logits(
+            last_token_logits,
+            temperature,
+            top_k,
+            top_p,
+            generator=self.sampling_rng,
+            vocab_size=vocab_size,
+        )
 
     def sample_from_logits(
         self,
@@ -734,32 +667,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
     def _dynamic_step_sample_bookkeeping(self):
         """Perform bookkeeping necessary to sample logits for dynamic batching."""
         context = self.inference_wrapped_model.inference_context
-        active_request_count = context.total_request_count - context.paused_request_count
-
-        if self._sampling_backend == "torch":
-            # Bucketize the core sampling parameters.
-            # Doing so via list comprehension is orders of magnitude faster than via torch.
-            bucket_map = defaultdict(list)
-
-            # Shorthands for the dictionary comprehension.
-            temp = context.active_request_metadata["temperature"][:active_request_count].tolist()
-            top_k = context.active_request_metadata["top_k"][:active_request_count].tolist()
-            top_p = context.active_request_metadata["top_p"][:active_request_count].tolist()
-
-            for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
-                sampling_params = (t, k, p)
-                bucket_map[sampling_params].append(request_index)
-
-            # Just unpack the key directly!
-            device = torch.cuda.current_device()
-            self._torch_sampling_buckets = [
-                (indices, *sampling_params) for sampling_params, indices in bucket_map.items()
-            ]
-            # Pre-compute index tensors on GPU to avoid per-step H2D copies.
-            self._torch_sampling_bucket_index_tensors = [
-                torch.tensor(indices, device=device, dtype=torch.long)
-                for indices, *_ in self._torch_sampling_buckets
-            ]
+        self._sampling.pre_forward_bookkeeping(context)
 
     def _rewind_kv_cache(self) -> tuple:
         """Update the KV cache bookkeeping for speculative decoding.
@@ -836,18 +744,12 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
         Returns:
             Tensor: Sampled tokens of shape [num_requests].
         """
-        spec_token_list = []
-        for idx_tensor, (_, temp, top_k, top_p) in zip(
-            self._torch_sampling_bucket_index_tensors, self._torch_sampling_buckets
-        ):
-            spec_token_list.append(
-                self._torch_sampling_func(logits_2d[idx_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, self._torch_sampling_bucket_index_tensors):
-            spec_tokens[indices] = tokens
-        return spec_tokens
+        n = logits_2d.shape[0]
+        output = torch.empty(n, device=logits_2d.device, dtype=torch.int64)
+        self._sampling.sample(
+            logits_2d, n, output, self.inference_wrapped_model.inference_context, eager=True
+        )
+        return output
 
     def _compute_serial_mtp_and_sample(self):
         """Compute MTP logits serially after verification and sample speculative tokens.
@@ -995,58 +897,14 @@ def _compute_serial_mtp_and_sample(self):
 
     def _sample_speculative_logits(
         self, required_logits: Tensor, request_in_prefill_status_tensor: Tensor
-    ) -> tuple:
-        """Sample tokens from logits using sampling buckets.
-
-        For torch sampling buckets: [request_indices, temp, top_k, top_p]
-
-        Example with 5 requests:
-            token_to_request_idx :              [ 0    0     0  |  1     1     1     |  2     2     2     |   3    |   4  ]
-            required_logits :                   [ a5l  a6l  a7l |  b3l    b4l  b5l   |  c6l   c7l   c8l   |  d2l   | e4l  ]  # Shape [11, vocab_size]
-
-            Sampling buckets: [[[0,2], temp1, top_k1, top_p1], [[1], temp3, top_k3, top_p3], [[3, 4], temp2, top_k2, top_p2]]
-
-            Final output tokens : [a5s  a6s  a7s  c6s  c7s  c8s  b3s  b4s  b5s  d2s  e4s]  # Shape [11]
-            (Rearranged from sampling bucket order back to input order using token_order)
-
-        Returns:
-            tuple: (output_tokens, repeats) where output_tokens has shape [total_required_tokens]
-        """
-        # request_in_prefill_status_tensor is already on GPU (from gpu_view).
-        repeats = torch.where(
-            request_in_prefill_status_tensor == 0, 1 + self.num_speculative_tokens, 1
+    ) -> Tensor:
+        """Sample speculative-token logits via the active sampling backend."""
+        return self._sampling.sample_speculative(
+            required_logits,
+            request_in_prefill_status_tensor,
+            self.num_speculative_tokens,
+            self.inference_wrapped_model.inference_context,
         )
-        token_to_request_index = torch.repeat_interleave(
-            torch.arange(
-                len(request_in_prefill_status_tensor),
-                device=request_in_prefill_status_tensor.device,
-            ),
-            repeats,
-        )
-
-        output_tokens_jumbled_list = []
-        token_order_list = []
-
-        for idx_tensor, (_, temp, top_k, top_p) in zip(
-            self._torch_sampling_bucket_index_tensors, self._torch_sampling_buckets
-        ):
-            required_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0]
-            output_tokens_jumbled_list.append(
-                self._torch_sampling_func(required_logits[required_indices, :], temp, top_k, top_p)
-            )
-            token_order_list.append(required_indices)
-
-        output_tokens_jumbled = torch.cat(output_tokens_jumbled_list, dim=0)
-        output_tokens = torch.empty(
-            len(output_tokens_jumbled),
-            device=output_tokens_jumbled.device,
-            dtype=output_tokens_jumbled.dtype,
-        )
-        token_order = torch.cat(token_order_list, dim=0)
-        # Rearrange output tokens from sampling_bucket request order back to input ids order
-        output_tokens[token_order] = output_tokens_jumbled
-
-        return output_tokens, repeats
 
     def _verify_speculative_tokens(
         self,
@@ -1096,7 +954,7 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
 
         # Sample tokens from logits
         nvtx_range_push("mtp-spec-decoding/verify/sample")
-        output_tokens, repeats = self._sample_speculative_logits(
+        output_tokens = self._sample_speculative_logits(
             required_logits, request_in_prefill_status_tensor
         )
         nvtx_range_pop("mtp-spec-decoding/verify/sample")
@@ -1168,7 +1026,6 @@ def _dynamic_step_sample_logits(self):
         # 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
 
@@ -1181,27 +1038,13 @@ def _dynamic_step_sample_logits(self):
                 self._all_logits_cuda[:, : context.padded_active_token_count, :]
             )
 
-        if self._sampling_backend == "torch":
-            # Concatenate the outputs once to prevent repeated small writes.
-            token_list = []
-            indices_list = []
-
-            # e.g torch sample buckets will be
-            # i.e (for all unique comibnation of t, topk, topk what are the associated
-            # requests indices (based on the active slices)
-            # [ [req at index 0, req at index 2], t1, topk1, topp1 ]]
-            # [ [req at index 1, req at index 3, req at index 4] , t2, topk2, topp2]
-            for indices, temp, top_k, top_p in self._torch_sampling_buckets:
-                token_list.append(
-                    self._torch_sampling_func(required_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
+        self._sampling.sample(
+            required_token_logits,
+            active_request_count,
+            self._sampled_tokens_cuda,
+            context,
+            eager=True,
+        )
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.
diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
index f4e76f6ab7c..fe1d607e818 100644
--- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
+++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
@@ -420,8 +420,8 @@ def test_cuda_graph_sp_padding_end_to_end(self, mtp_use_repeated_layer):
             ctrl._mtp_resolved_padded_count = padded_count
             context._using_cuda_graph_this_step = True
 
-            ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-            ctrl._torch_sampling_bucket_index_tensors = [
+            ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
+            ctrl._sampling._bucket_index_tensors = [
                 torch.arange(active_request_count, device='cuda', dtype=torch.long)
             ]
 
@@ -521,8 +521,8 @@ def _run_mtp(use_cuda_graph):
                 unwrapped._decoder_hidden_states_cache = local_hidden
 
                 ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda')
-                ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-                ctrl._torch_sampling_bucket_index_tensors = [
+                ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
+                ctrl._sampling._bucket_index_tensors = [
                     torch.arange(active_request_count, device='cuda', dtype=torch.long)
                 ]
 
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 2c1a9902a1f..3e424ef3dac 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -999,11 +999,11 @@ def mock_sampling_func(logits, *args, **kwargs):
                 return torch.zeros((12,), dtype=torch.long, device='cuda')
 
         # Override sampling to return our predictable mock outputs
-        self.text_generation_controller._torch_sampling_buckets = [([0, 1], 1.0, 1, 0.0)]
-        self.text_generation_controller._torch_sampling_bucket_index_tensors = [
+        self.text_generation_controller._sampling._buckets = [([0, 1], 1.0, 1, 0.0)]
+        self.text_generation_controller._sampling._bucket_index_tensors = [
             torch.tensor([0, 1], device='cuda', dtype=torch.long)
         ]
-        self.text_generation_controller._torch_sampling_func = mock.MagicMock(
+        self.text_generation_controller._sampling._sampling_func = mock.MagicMock(
             side_effect=mock_sampling_func
         )
 
@@ -1239,9 +1239,9 @@ def test_speculative_multinomial_sampling(self):
         logits = torch.randn(1, 8, self.vocab_size, device='cuda')
 
         # Set up a bucket that forces multinomial sampling (top_p = 0.9, top_k = 0)
-        # _torch_sampling_buckets format: (indices, temp, top_k, top_p)
-        self.text_generation_controller._torch_sampling_buckets = [([0, 1], 1.0, 0, 0.9)]
-        self.text_generation_controller._torch_sampling_bucket_index_tensors = [
+        # _buckets format: (indices, temp, top_k, top_p)
+        self.text_generation_controller._sampling._buckets = [([0, 1], 1.0, 0, 0.9)]
+        self.text_generation_controller._sampling._bucket_index_tensors = [
             torch.tensor([0, 1], device='cuda', dtype=torch.long)
         ]
 
@@ -1491,8 +1491,8 @@ def test_mtp_sp_padding_real_ranks(self, active_request_count):
         ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda')
 
         # Greedy sampling: top_k=1 selects the argmax token deterministically.
-        ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-        ctrl._torch_sampling_bucket_index_tensors = [
+        ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
+        ctrl._sampling._bucket_index_tensors = [
             torch.arange(active_request_count, device='cuda', dtype=torch.long)
         ]
 

From a225d8cb1aa8ad1de6bf615f7f96e0a294705a84 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:31:55 -0500
Subject: [PATCH 02/23] Add FlashInfer sampling backend

---
 megatron/core/inference/config.py             | 13 +++-
 megatron/core/inference/sampling/__init__.py  |  3 +-
 .../inference/sampling/flashinfer_sampling.py | 76 +++++++++++++++++++
 .../text_generation_controller.py             | 12 ++-
 megatron/inference/utils.py                   |  1 +
 megatron/training/arguments.py                | 13 ++++
 6 files changed, 113 insertions(+), 5 deletions(-)
 create mode 100644 megatron/core/inference/sampling/flashinfer_sampling.py

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index c4b092309c2..86ce980aa07 100644
--- a/megatron/core/inference/config.py
+++ b/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,11 @@ 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."
+                )
diff --git a/megatron/core/inference/sampling/__init__.py b/megatron/core/inference/sampling/__init__.py
index 97a3a9e327d..b2941b33c9e 100644
--- a/megatron/core/inference/sampling/__init__.py
+++ b/megatron/core/inference/sampling/__init__.py
@@ -1,6 +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"]
+__all__ = ["Sampling", "TorchSampling", "FlashInferSampling"]
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
new file mode 100644
index 00000000000..8503b127cd7
--- /dev/null
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+
+try:
+    import flashinfer
+except ImportError:
+    flashinfer = None
+
+from megatron.core.inference.sampling.base import Sampling
+
+
+class FlashInferSampling(Sampling):
+    """Fused FlashInfer sampling.
+
+    Unlike `TorchSampling`, FlashInfer kernels accept per-row parameter tensors
+    (temperature, top_k, top_p) directly, so no bucketing is required.
+    """
+
+    def __init__(self, vocab_size: int, rng: torch.Generator) -> None:
+        self._vocab_size = vocab_size
+        self._rng = rng
+
+    def pre_forward_bookkeeping(self, context) -> None:
+        """No-op; FlashInfer needs no per-step bookkeeping."""
+
+    def sample_kernel(
+        self,
+        logits: Tensor,
+        n: int,
+        output: Tensor,
+        context,
+        *,
+        eager: bool = False,
+        cache_key: Any = None,
+        gather_indices: Optional[Tensor] = None,
+        token_to_request_index: Optional[Tensor] = None,
+    ) -> None:
+        """FlashInfer fused top-k / top-p sampling kernel.
+
+        Reads sampling parameters per-row from `context.active_request_metadata`,
+        applies temperature scaling and top-k/top-p filtering, then samples via
+        `flashinfer.sampling.top_k_top_p_sampling_from_probs`.
+        """
+        del eager, cache_key
+        md = context.active_request_metadata
+        if token_to_request_index is None:
+            temperature = md["temperature"][:n]
+            top_k = md["top_k"][:n]
+            top_p = md["top_p"][:n]
+        else:
+            temperature = md["temperature"][token_to_request_index]
+            top_k = md["top_k"][token_to_request_index]
+            top_p = md["top_p"][token_to_request_index]
+
+        if gather_indices is None:
+            # Slice is a view; clone before in-place div_ to avoid mutating the caller.
+            scaled = logits[:n].clone()
+        else:
+            # Advanced indexing already returns a new tensor.
+            scaled = logits[gather_indices[:n], :]
+        scaled.div_(temperature.unsqueeze(1))
+        probs = torch.softmax(scaled, dim=-1)
+
+        # Sentinel values disable filtering: top_k=vocab_size keeps all
+        # tokens, top_p=1.0 keeps the full probability mass.
+        top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
+        top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
+        output[:n].copy_(
+            flashinfer.sampling.top_k_top_p_sampling_from_probs(
+                probs, top_k_safe, top_p_safe, generator=self._rng
+            )
+        )
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 68266fcc732..b47f76e74a6 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -60,7 +60,7 @@
     HAVE_TE = False
 
 from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
-from megatron.core.inference.sampling import Sampling, TorchSampling
+from megatron.core.inference.sampling import FlashInferSampling, Sampling, TorchSampling
 from megatron.core.inference.text_generation_controllers.mtp_utils_pytorch import rewind_kv_cache
 from megatron.core.inference.text_generation_controllers.mtp_utils_triton import (
     mamba_state_selective_copy,
@@ -155,7 +155,7 @@ 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._sampling_backend = context.config.sampling_backend
         self._enable_cuda_graph = self.model_config.cuda_graph_impl == "local"
 
         # Initialize bookkeeping tensors.
@@ -168,7 +168,10 @@ def _init_dynamic_sampling_tensors(self):
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
         # Sampling backend: dispatches per-step bookkeeping and the sampling kernel.
-        self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size)
+        if self._sampling_backend == "flashinfer":
+            self._sampling: Sampling = FlashInferSampling(self.vocab_size, self.sampling_rng)
+        else:
+            self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size)
 
         # Cache values that are constant across inference steps.
         self._unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
@@ -661,6 +664,9 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
         # Copy logits to contiguous buffer.
         if self._enable_cuda_graph:
             self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits[:, :logits_seq_len, :])
+        elif self._sampling_backend == "flashinfer":
+            # FlashInfer kernels require contiguous inputs.
+            self._all_logits_cuda = logits.contiguous()
         else:
             self._all_logits_cuda = logits
 
diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py
index bb65c754ab1..559ac57f496 100644
--- a/megatron/inference/utils.py
+++ b/megatron/inference/utils.py
@@ -374,6 +374,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,
     )
 
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 17b27187b2e..db976ebbb8b 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1007,6 +1007,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.
@@ -1994,6 +2003,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.')

From c5ebb27ef35568b9d7dd55ed68f24a7d4b998ed3 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:36:41 -0500
Subject: [PATCH 03/23] Cache-key sampling on padded batch dimensions

---
 .../inference/contexts/dynamic_context.py     | 24 +++++
 megatron/core/inference/sampling/base.py      | 38 +++++---
 .../inference/sampling/flashinfer_sampling.py | 24 ++++-
 .../text_generation_controller.py             | 97 ++++++++++++++-----
 4 files changed, 144 insertions(+), 39 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 159e1f90b34..9afdb609f33 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1055,6 +1055,11 @@ def initialize_all_tensors(self) -> None:
         ].view(torch.int32)
         _off += _req_int32_bytes
 
+        # Per-request last-token row indices, used as gather indices by sampling kernels
+        # (e.g. FlashInfer) so the call site can pass a fixed-shape tensor on the padded
+        # batch dimension.
+        self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
+
         # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
         max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1)
         self.active_logit_idxs = torch.zeros(
@@ -1362,6 +1367,15 @@ def build_active_slices(self, batch_size: int):
                 self.request_metadata[label][padded_slice], non_blocking=True
             )
 
+        # Cumsum of query lengths gives per-request last-token row indices.
+        # Padded slots get garbage here; pad_active_slices fills them with 0.
+        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] -= 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 +1402,16 @@ 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: neutral defaults so per-row sampling kernels (FlashInfer)
+        # produce harmless output for padded slots. top_k=0 / top_p=0.0 are sentinels
+        # for "no filter" in the FlashInfer wrapper.
+        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.
 
diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
index 2a06a486aa1..ac82f567649 100644
--- a/megatron/core/inference/sampling/base.py
+++ b/megatron/core/inference/sampling/base.py
@@ -79,34 +79,42 @@ def sample(
     def sample_speculative(
         self,
         required_logits: Tensor,
-        request_in_prefill_status: Tensor,
+        num_decode: int,
+        num_prefill: int,
         num_speculative_tokens: int,
         context,
+        *,
+        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 one row. Builds the per-token request mapping and dispatches to
-        `sample_kernel(eager=True)` so the call is safe inside an outer captured kernel.
+        `sample_kernel`. Callers may use `eager` and `cache_key` to control CUDA graph
+        capture/replay on backends that wrap `sample_kernel`.
         """
-        repeats = torch.where(
-            request_in_prefill_status == 0, 1 + num_speculative_tokens, 1
-        )
-        token_to_request_index = torch.repeat_interleave(
-            torch.arange(
-                len(request_in_prefill_status),
-                device=request_in_prefill_status.device,
-            ),
-            repeats,
+        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),
+            ]
         )
-        n = token_to_request_index.shape[0]
-        output_tokens = torch.empty(n, device=required_logits.device, dtype=torch.int64)
+        output_tokens = torch.empty(num_tokens, device=device, dtype=torch.int64)
         self.sample_kernel(
             required_logits,
-            n,
+            num_tokens,
             output_tokens,
             context,
-            eager=True,
+            eager=eager,
+            cache_key=cache_key,
             token_to_request_index=token_to_request_index,
         )
         return output_tokens
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 8503b127cd7..920e2e38845 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -11,18 +11,34 @@
     flashinfer = None
 
 from megatron.core.inference.sampling.base import Sampling
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
 
 
 class FlashInferSampling(Sampling):
-    """Fused FlashInfer sampling.
+    """Fused FlashInfer sampling, with optional CUDA graph capture/replay.
 
     Unlike `TorchSampling`, FlashInfer kernels accept per-row parameter tensors
     (temperature, top_k, top_p) directly, so no bucketing is required.
     """
 
-    def __init__(self, vocab_size: int, rng: torch.Generator) -> None:
+    def __init__(
+        self,
+        vocab_size: int,
+        rng: torch.Generator,
+        config=None,
+        enable_cuda_graph: bool = False,
+    ) -> None:
         self._vocab_size = vocab_size
         self._rng = rng
+        self._enable_cuda_graph = enable_cuda_graph
+        if enable_cuda_graph and config is not None and config.cuda_graph_impl == "local":
+            CudaGraphManager(
+                config,
+                self,
+                function_name="sample_kernel",
+                need_backward=False,
+                inline_capture=True,
+            )
 
     def pre_forward_bookkeeping(self, context) -> None:
         """No-op; FlashInfer needs no per-step bookkeeping."""
@@ -44,6 +60,10 @@ def sample_kernel(
         Reads sampling parameters per-row from `context.active_request_metadata`,
         applies temperature scaling and top-k/top-p filtering, then samples via
         `flashinfer.sampling.top_k_top_p_sampling_from_probs`.
+
+        When wrapped by `CudaGraphManager`, `eager` and `cache_key` are consumed by
+        the wrapper before this body runs. When unwrapped (no CUDA graphs), they are
+        accepted and ignored so callers can pass them unconditionally.
         """
         del eager, cache_key
         md = context.active_request_metadata
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index b47f76e74a6..6477411ffd4 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -169,7 +169,12 @@ def _init_dynamic_sampling_tensors(self):
 
         # Sampling backend: dispatches per-step bookkeeping and the sampling kernel.
         if self._sampling_backend == "flashinfer":
-            self._sampling: Sampling = FlashInferSampling(self.vocab_size, self.sampling_rng)
+            self._sampling: Sampling = FlashInferSampling(
+                self.vocab_size,
+                self.sampling_rng,
+                config=self.model_config,
+                enable_cuda_graph=self._enable_cuda_graph,
+            )
         else:
             self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size)
 
@@ -902,14 +907,29 @@ def _compute_serial_mtp_and_sample(self):
             del unwrapped_model._decoder_hidden_states_cache
 
     def _sample_speculative_logits(
-        self, required_logits: Tensor, request_in_prefill_status_tensor: Tensor
+        self,
+        required_logits: Tensor,
+        num_decode: int,
+        num_prefill: int,
+        *,
+        eager: bool,
+        cache_key,
     ) -> Tensor:
-        """Sample speculative-token logits via the active sampling backend."""
+        """Sample speculative-token logits via the active sampling backend.
+
+        `num_decode` and `num_prefill` are the request counts the kernel should treat
+        as decode/prefill: padded values when running a captured CUDA graph, actual
+        values otherwise. Padded slots beyond the active count contribute decode-style
+        rows (`pad_active_slices` initialises their query lengths to `1 + n_spec`).
+        """
         return self._sampling.sample_speculative(
             required_logits,
-            request_in_prefill_status_tensor,
+            num_decode,
+            num_prefill,
             self.num_speculative_tokens,
             self.inference_wrapped_model.inference_context,
+            eager=eager,
+            cache_key=cache_key,
         )
 
     def _verify_speculative_tokens(
@@ -936,10 +956,28 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
 
-        # Use gpu_view for data consumed by GPU operations (sampling, verification).
-        request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[
-            :active_request_count
-        ]
+        # Sampling-side request counts: padded when running a captured graph so the
+        # cache key buckets, actual otherwise. Verify uses the actual counts so the
+        # downstream Triton kernels still operate on the real workload.
+        #
+        # Padding-as-decode is only safe when the captured graph is decode-only — in
+        # that case `pad_active_slices` set the padded slots' query lengths to
+        # `1 + n_spec` (decode-style). With a mixed graph, `[actual_decode,
+        # padded_decode)` may contain real prefill rows whose query lengths differ
+        # from `1 + n_spec`, so the kernel's reshape would mix prompt tokens into the
+        # decode chunks. Fall back to actual values + eager sampling in that case.
+        use_graph_for_sampling = (
+            self._sampling_backend == "flashinfer"
+            and self._enable_cuda_graph
+            and context.using_cuda_graph_this_step()
+            and context.padded_batch_dimensions.prefill_req_count == 0
+        )
+        if use_graph_for_sampling:
+            sample_num_decode = context.padded_batch_dimensions.decode_req_count
+            sample_num_prefill = 0
+        else:
+            sample_num_decode = context.num_decode_requests
+            sample_num_prefill = context.num_prefill_requests
 
         # Get the logit indices for tokens that need sampling.
         # These indices are always needed for input_ids slicing and tracking
@@ -961,7 +999,15 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         # Sample tokens from logits
         nvtx_range_push("mtp-spec-decoding/verify/sample")
         output_tokens = self._sample_speculative_logits(
-            required_logits, request_in_prefill_status_tensor
+            required_logits,
+            sample_num_decode,
+            sample_num_prefill,
+            eager=not use_graph_for_sampling,
+            cache_key=(
+                ("sample_speculative", sample_num_decode, sample_num_prefill)
+                if use_graph_for_sampling
+                else None
+            ),
         )
         nvtx_range_pop("mtp-spec-decoding/verify/sample")
 
@@ -1034,22 +1080,29 @@ def _dynamic_step_sample_logits(self):
 
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
-
-        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.
-            required_token_logits = self._all_logits_cuda.squeeze(0)[:active_request_count, :]
-        else:
-            required_token_logits = context.last_token_logits(
-                self._all_logits_cuda[:, : context.padded_active_token_count, :]
-            )
-
+        use_graph = (
+            self._sampling_backend == "flashinfer"
+            and self._enable_cuda_graph
+            and context.using_cuda_graph_this_step()
+        )
+        # Padded count when running a captured graph (cache key buckets); actual otherwise.
+        n = context.padded_active_request_count if use_graph else active_request_count
+        # When `materialize_only_last_token_logits` is true the forward pass already
+        # selected the right rows. Otherwise we point the kernel at the per-request
+        # last-token positions via `gather_indices`; padded slots safely fan in to row 0.
+        gather_indices = (
+            None
+            if context.config.materialize_only_last_token_logits
+            else context.active_request_last_token_idxs
+        )
         self._sampling.sample(
-            required_token_logits,
-            active_request_count,
+            self._all_logits_cuda.squeeze(0),
+            n,
             self._sampled_tokens_cuda,
             context,
-            eager=True,
+            eager=not use_graph,
+            cache_key=("sample", n) if use_graph else None,
+            gather_indices=gather_indices,
         )
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:

From 23b1bf2c531f1e93bec06e18739b6c16555c98f5 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:38:14 -0500
Subject: [PATCH 04/23] Capture sampling graphs during warmup

---
 megatron/core/inference/engines/dynamic_engine.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index ff5454bbaa3..518c4a5c25b 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -395,6 +395,16 @@ def create_cuda_graphs(self, reset_context: bool = True):
             with torch.inference_mode():
                 controller._dynamic_step_forward_logits(input_ids, position_ids)
 
+                # Sampling warm-up. The non-speculative path captures `sample_kernel`
+                # keyed by `("sample", padded_active_request_count)`; the speculative
+                # path additionally captures `("sample_speculative", padded_decode, 0)`
+                # for decode-only graphs. Other backends and mixed graphs run eagerly.
+                controller._dynamic_step_sample_bookkeeping()
+                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

From a2c354f3654651c0b3d1ddd85aaf28d036f6e783 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:39:39 -0500
Subject: [PATCH 05/23] Add tests

---
 .../test_text_generation_controller.py                    | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 3e424ef3dac..a0ca39ef6a0 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -67,6 +67,7 @@ def setup_model(
         expert_model_parallel_size: int = 1,
         num_moe_experts: int = None,
         hybrid_layer_pattern: str = None,
+        sampling_backend: str = 'torch',
         cuda_graph_impl: str = 'none',
     ):
         if use_training_random_init:
@@ -164,6 +165,7 @@ def setup_model(
                     enable_prefix_caching=enable_prefix_caching,
                     max_requests=max_requests,
                     mamba_inference_state_config=mamba_inference_state_config,
+                    sampling_backend=sampling_backend,
                 ),
             )
 
@@ -280,17 +282,20 @@ def detokenize(self, inp, skip_special_tokens=False):
             sampled_logits >= expected_min_value
         ), f"The sampled logits should all be greater than {expected_min_value} but its {sampled_logits}"
 
-    @pytest.mark.parametrize("backend", ["torch"])
+    @pytest.mark.parametrize("backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     def test_sample_from_dynamic_logits(
         self, backend: str, materialize_only_last_token_logits: bool
     ):
+        if backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         batch_size = 12
         self.setup_model(
             torch.float32,
             batch_size=batch_size,
             static=False,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
+            sampling_backend=backend,
         )
         self.mock_tokenizer.eod = self.vocab_size
 
@@ -320,7 +325,6 @@ def test_sample_from_dynamic_logits(
         context.active_request_metadata["temperature"][:batch_size].copy_(temp_values)
         context.active_request_metadata["top_k"][:batch_size].copy_(top_k_values)
         context.active_request_metadata["top_p"][:batch_size].copy_(top_p_values)
-        self.text_generation_controller._sampling_backend = backend
 
         context.padded_active_token_count = batch_size
         context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32, device='cuda')

From 7d315329a592912eabb36d5c7272ac3514ddb398 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 05:40:44 -0500
Subject: [PATCH 06/23] Self-review: misc cleanups

---
 .../text_generation_controller.py             | 61 ++++---------------
 1 file changed, 11 insertions(+), 50 deletions(-)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 6477411ffd4..21bc26de220 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -329,28 +329,6 @@ def detokenize_generations(
 
         return text, prompts_plus_generations_segments
 
-    def _torch_sampling_func(
-        self,
-        last_token_logits: torch.Tensor,
-        temperature: float,
-        top_k: int,
-        top_p: float,
-        vocab_size: Optional[int] = None,
-    ):
-        """Static-batching sampler shim. Forwards to `TorchSampling.sample_from_logits`.
-
-        The dynamic-batching path goes through `self._sampling`; this method is kept
-        for the static `sample_from_logits` flow and for any callers that mock it.
-        """
-        return TorchSampling.sample_from_logits(
-            last_token_logits,
-            temperature,
-            top_k,
-            top_p,
-            generator=self.sampling_rng,
-            vocab_size=vocab_size,
-        )
-
     def sample_from_logits(
         self,
         last_token_logits: torch.Tensor,
@@ -437,7 +415,14 @@ def sample_from_logits(
         top_k = sampling_params.top_k
         temperature = sampling_params.temperature
 
-        return self._torch_sampling_func(last_token_logits, temperature, top_k, top_p, vocab_size)
+        return TorchSampling.sample_from_logits(
+            last_token_logits,
+            temperature,
+            top_k,
+            top_p,
+            generator=self.sampling_rng,
+            vocab_size=vocab_size,
+        )
 
     def update_generation_status(
         self,
@@ -906,32 +891,6 @@ def _compute_serial_mtp_and_sample(self):
         if has_mtp:
             del unwrapped_model._decoder_hidden_states_cache
 
-    def _sample_speculative_logits(
-        self,
-        required_logits: Tensor,
-        num_decode: int,
-        num_prefill: int,
-        *,
-        eager: bool,
-        cache_key,
-    ) -> Tensor:
-        """Sample speculative-token logits via the active sampling backend.
-
-        `num_decode` and `num_prefill` are the request counts the kernel should treat
-        as decode/prefill: padded values when running a captured CUDA graph, actual
-        values otherwise. Padded slots beyond the active count contribute decode-style
-        rows (`pad_active_slices` initialises their query lengths to `1 + n_spec`).
-        """
-        return self._sampling.sample_speculative(
-            required_logits,
-            num_decode,
-            num_prefill,
-            self.num_speculative_tokens,
-            self.inference_wrapped_model.inference_context,
-            eager=eager,
-            cache_key=cache_key,
-        )
-
     def _verify_speculative_tokens(
         self,
         output_tokens: Tensor,
@@ -998,10 +957,12 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
 
         # Sample tokens from logits
         nvtx_range_push("mtp-spec-decoding/verify/sample")
-        output_tokens = self._sample_speculative_logits(
+        output_tokens = self._sampling.sample_speculative(
             required_logits,
             sample_num_decode,
             sample_num_prefill,
+            self.num_speculative_tokens,
+            context,
             eager=not use_graph_for_sampling,
             cache_key=(
                 ("sample_speculative", sample_num_decode, sample_num_prefill)

From 1d6623e793bc32c65e237ae0a95d62b296cc1504 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 06:02:18 -0500
Subject: [PATCH 07/23] Fix: sample_kernel must return its output tensor

---
 megatron/core/inference/sampling/base.py      | 20 ++++++++-----------
 .../inference/sampling/flashinfer_sampling.py | 10 +++++++---
 .../core/inference/sampling/torch_sampling.py |  9 ++++++---
 .../text_generation_controller.py             | 17 +++++++++-------
 4 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
index ac82f567649..ffd3020dd76 100644
--- a/megatron/core/inference/sampling/base.py
+++ b/megatron/core/inference/sampling/base.py
@@ -26,15 +26,14 @@ def sample_kernel(
         self,
         logits: Tensor,
         n: int,
-        output: Tensor,
         context,
         *,
         eager: bool = False,
         cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
-    ) -> None:
-        """Sample n tokens from logits into `output[:n]`.
+    ) -> Tensor:
+        """Sample `n` tokens from `logits` and return them.
 
         `eager` and `cache_key` are consumed by the `CudaGraphManager` wrapper when one
         is installed; unwrapped subclasses accept and ignore them.
@@ -42,13 +41,15 @@ def sample_kernel(
         Args:
             logits: Logits tensor of shape `[>=n, vocab_size]`.
             n: Number of rows to sample.
-            output: Destination buffer for sampled token ids.
             context: The active DynamicInferenceContext.
             eager: If True, skip CUDA graph capture/replay (consumed by the wrapper).
             cache_key: Hashable key for runner lookup (consumed by the wrapper).
             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.
+
+        Returns:
+            Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
         """
         ...
 
@@ -56,19 +57,17 @@ def sample(
         self,
         logits: Tensor,
         n: int,
-        output: Tensor,
         context,
         *,
         eager: bool = False,
         cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
-    ) -> None:
+    ) -> Tensor:
         """Sample `n` tokens, optionally with CUDA graph capture/replay."""
-        self.sample_kernel(
+        return self.sample_kernel(
             logits,
             n,
-            output,
             context,
             eager=eager,
             cache_key=cache_key,
@@ -107,14 +106,11 @@ def sample_speculative(
                 torch.arange(num_decode, num_decode + num_prefill, device=device),
             ]
         )
-        output_tokens = torch.empty(num_tokens, device=device, dtype=torch.int64)
-        self.sample_kernel(
+        return self.sample_kernel(
             required_logits,
             num_tokens,
-            output_tokens,
             context,
             eager=eager,
             cache_key=cache_key,
             token_to_request_index=token_to_request_index,
         )
-        return output_tokens
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 920e2e38845..d999d90394e 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -47,14 +47,13 @@ def sample_kernel(
         self,
         logits: Tensor,
         n: int,
-        output: Tensor,
         context,
         *,
         eager: bool = False,
         cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
-    ) -> None:
+    ) -> Tensor:
         """FlashInfer fused top-k / top-p sampling kernel.
 
         Reads sampling parameters per-row from `context.active_request_metadata`,
@@ -64,6 +63,9 @@ def sample_kernel(
         When wrapped by `CudaGraphManager`, `eager` and `cache_key` are consumed by
         the wrapper before this body runs. When unwrapped (no CUDA graphs), they are
         accepted and ignored so callers can pass them unconditionally.
+
+        Returns:
+            Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
         """
         del eager, cache_key
         md = context.active_request_metadata
@@ -89,8 +91,10 @@ def sample_kernel(
         # tokens, top_p=1.0 keeps the full probability mass.
         top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
         top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
-        output[:n].copy_(
+        output = torch.empty(n, device=logits.device, dtype=torch.int64)
+        output.copy_(
             flashinfer.sampling.top_k_top_p_sampling_from_probs(
                 probs, top_k_safe, top_p_safe, generator=self._rng
             )
         )
+        return output
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
index 18ed84fa1e6..5e2d438219e 100644
--- a/megatron/core/inference/sampling/torch_sampling.py
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -125,31 +125,33 @@ def sample_kernel(
         self,
         logits: Tensor,
         n: int,
-        output: Tensor,
         context,
         *,
         eager: bool = False,
         cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
-    ) -> None:
+    ) -> Tensor:
         """Sample by iterating over pre-computed sampling buckets.
 
         Args:
             logits: Logits tensor of shape `[>=n, vocab_size]`.
             n: Number of rows to sample.
-            output: Destination buffer; the kernel writes to the rows selected per bucket.
             context: The active DynamicInferenceContext (unused; kept for ABC parity).
             eager: Accepted for API symmetry with FlashInfer; ignored (no wrapper here).
             cache_key: Accepted for API symmetry; ignored.
             gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: When set, the loop dispatches per-token rather than
                 per-request (used by the speculative path).
+
+        Returns:
+            Sampled token ids of shape `[n]`.
         """
         del eager, cache_key
         if gather_indices is not None:
             logits = logits[gather_indices[:n], :]
 
+        output = torch.empty(n, device=logits.device, dtype=torch.int64)
         token_list = []
         indices_list = []
         for idx_tensor, (_, temp, top_k, top_p) in zip(
@@ -165,6 +167,7 @@ def sample_kernel(
         sampled_tokens = torch.cat(token_list, dim=0)
         sampled_indices = torch.cat(indices_list, dim=0)
         output[sampled_indices] = sampled_tokens
+        return output
 
     def _sampling_func(
         self, last_token_logits: Tensor, temperature: float, top_k: int, top_p: float
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 21bc26de220..157f0450338 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -740,12 +740,12 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
         Returns:
             Tensor: Sampled tokens of shape [num_requests].
         """
-        n = logits_2d.shape[0]
-        output = torch.empty(n, device=logits_2d.device, dtype=torch.int64)
-        self._sampling.sample(
-            logits_2d, n, output, self.inference_wrapped_model.inference_context, eager=True
+        return self._sampling.sample(
+            logits_2d,
+            logits_2d.shape[0],
+            self.inference_wrapped_model.inference_context,
+            eager=True,
         )
-        return output
 
     def _compute_serial_mtp_and_sample(self):
         """Compute MTP logits serially after verification and sample speculative tokens.
@@ -1056,15 +1056,18 @@ def _dynamic_step_sample_logits(self):
             if context.config.materialize_only_last_token_logits
             else context.active_request_last_token_idxs
         )
-        self._sampling.sample(
+        sampled = self._sampling.sample(
             self._all_logits_cuda.squeeze(0),
             n,
-            self._sampled_tokens_cuda,
             context,
             eager=not use_graph,
             cache_key=("sample", n) if use_graph else None,
             gather_indices=gather_indices,
         )
+        # Copy out of the captured static buffer (under CG) into our persistent
+        # buffer so subsequent steps see the new sampled values, and the static
+        # buffer is free to be reused.
+        self._sampled_tokens_cuda[:n].copy_(sampled)
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.

From 2ec02d7ea69075a9f6970274d056d192227f16e1 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 07:32:57 -0500
Subject: [PATCH 08/23] Remove unnecessary tensor copy

---
 .../text_generation_controller.py                 | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 157f0450338..1c3d366c415 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -165,7 +165,13 @@ def _init_dynamic_sampling_tensors(self):
             )
         else:
             self._all_logits_cuda = None
-        self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
+        # Speculative path:
+        #     - `self._sampled_tokens_cuda` is pre-allocated by `_init_mtp_sampling_tensors`.
+        #     - The tensor cannot be reused between the Triton kernel and the sampling graph.
+        # Non-speculative path:
+        #     - `self._sampled_tokens_cuda` is rebound to the output of `sample()`,
+        #     which uses CudaGraphManager syntactic sugar to keep it as a static tensor.
+        self._sampled_tokens_cuda = None
 
         # Sampling backend: dispatches per-step bookkeeping and the sampling kernel.
         if self._sampling_backend == "flashinfer":
@@ -200,6 +206,7 @@ def _init_mtp_sampling_tensors(self):
         context = self.inference_wrapped_model.inference_context
         max_requests = context.max_requests
         device = torch.cuda.current_device()
+        self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
         self._sampled_mtp_tokens_cuda = torch.empty(
             [self.num_speculative_tokens, max_requests], dtype=torch.int64, device=device
         )
@@ -1056,7 +1063,7 @@ def _dynamic_step_sample_logits(self):
             if context.config.materialize_only_last_token_logits
             else context.active_request_last_token_idxs
         )
-        sampled = self._sampling.sample(
+        self._sampled_tokens_cuda = self._sampling.sample(
             self._all_logits_cuda.squeeze(0),
             n,
             context,
@@ -1064,10 +1071,6 @@ def _dynamic_step_sample_logits(self):
             cache_key=("sample", n) if use_graph else None,
             gather_indices=gather_indices,
         )
-        # Copy out of the captured static buffer (under CG) into our persistent
-        # buffer so subsequent steps see the new sampled values, and the static
-        # buffer is free to be reused.
-        self._sampled_tokens_cuda[:n].copy_(sampled)
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.

From fd31b1287513d565be569c3b5a0eb17a092c291f Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 08:51:15 -0500
Subject: [PATCH 09/23] Fix speculative graphing

---
 megatron/core/inference/sampling/base.py      | 13 ++++---
 .../text_generation_controller.py             | 34 +++++++++----------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
index ffd3020dd76..ebf41fc12c5 100644
--- a/megatron/core/inference/sampling/base.py
+++ b/megatron/core/inference/sampling/base.py
@@ -85,13 +85,17 @@ def sample_speculative(
         *,
         eager: bool = False,
         cache_key: Any = None,
+        gather_indices: Optional[Tensor] = None,
     ) -> Tensor:
         """Sample tokens for the speculative-verify path.
 
-        Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests
-        contribute one row. Builds the per-token request mapping and dispatches to
-        `sample_kernel`. Callers may use `eager` and `cache_key` to control CUDA graph
-        capture/replay on backends that wrap `sample_kernel`.
+        Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests contribute 1.
+        Builds the per-token request mapping and dispatches to `sample_kernel`.
+
+        When `gather_indices` is supplied, `required_logits` is the full per-token logits buffer
+        (constant shape across steps); the kernel selects rows 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).
         """
         n_spec = num_speculative_tokens
         num_decode_tokens = num_decode * (1 + n_spec)
@@ -112,5 +116,6 @@ def sample_speculative(
             context,
             eager=eager,
             cache_key=cache_key,
+            gather_indices=gather_indices,
             token_to_request_index=token_to_request_index,
         )
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 1c3d366c415..fe332995bea 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -925,47 +925,46 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         # Sampling-side request counts: padded when running a captured graph so the
         # cache key buckets, actual otherwise. Verify uses the actual counts so the
         # downstream Triton kernels still operate on the real workload.
-        #
-        # Padding-as-decode is only safe when the captured graph is decode-only — in
-        # that case `pad_active_slices` set the padded slots' query lengths to
-        # `1 + n_spec` (decode-style). With a mixed graph, `[actual_decode,
-        # padded_decode)` may contain real prefill rows whose query lengths differ
-        # from `1 + n_spec`, so the kernel's reshape would mix prompt tokens into the
-        # decode chunks. Fall back to actual values + eager sampling in that case.
         use_graph_for_sampling = (
             self._sampling_backend == "flashinfer"
             and self._enable_cuda_graph
             and context.using_cuda_graph_this_step()
-            and context.padded_batch_dimensions.prefill_req_count == 0
         )
         if use_graph_for_sampling:
             sample_num_decode = context.padded_batch_dimensions.decode_req_count
-            sample_num_prefill = 0
+            sample_num_prefill = context.padded_batch_dimensions.prefill_req_count
         else:
             sample_num_decode = context.num_decode_requests
             sample_num_prefill = context.num_prefill_requests
 
-        # Get the logit indices for tokens that need sampling.
-        # These indices are always needed for input_ids slicing and tracking
-        # accepted sequence positions, even when logits are pre-sliced.
+        # Logit indices for tokens that need sampling.
+        # Padded under graph capture so the captured `gather_indices` input has a stable shape.
+        # Padded slots resolve to row 0; verify and prepare-next read only the actual prefix,
+        # so the padded-row samples produced by the captured kernel are discarded.
         nvtx_range_push("mtp-spec-decoding/verify/logit-indices")
         # Use pre-allocated buffer for CUDA graph compatibility.
         logits = self._all_logits_cuda
+        # `speculative_required_logit_indices()` already returns padded indices when
+        # running a captured graph (`num_last_token_logits` uses the padded counts and
+        # `pad_active_slices` zero-pads the trailing slots), so the call site does not
+        # need to re-pad here.
         required_logit_indices = context.speculative_required_logit_indices()
 
         if context.config.materialize_only_last_token_logits:
             # last_token_logits already selected exactly the required positions.
-            required_logits = logits.squeeze(0)
+            sample_logits = logits.squeeze(0)
+            sample_gather_indices = None
         else:
-            required_logits = logits.squeeze(0)[
-                required_logit_indices, :
-            ]  # Shape [num_required, vocab_size]
+            # Push the gather inside the captured kernel:
+            # pass the full per-token logits buffer (constant shape) plus the padded indices.
+            sample_logits = logits.squeeze(0)
+            sample_gather_indices = required_logit_indices
         nvtx_range_pop("mtp-spec-decoding/verify/logit-indices")
 
         # Sample tokens from logits
         nvtx_range_push("mtp-spec-decoding/verify/sample")
         output_tokens = self._sampling.sample_speculative(
-            required_logits,
+            sample_logits,
             sample_num_decode,
             sample_num_prefill,
             self.num_speculative_tokens,
@@ -976,6 +975,7 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
                 if use_graph_for_sampling
                 else None
             ),
+            gather_indices=sample_gather_indices,
         )
         nvtx_range_pop("mtp-spec-decoding/verify/sample")
 

From 6b274c46a35d828ade4ef03634e31e6ce57c8cbc Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 28 Apr 2026 09:05:06 -0500
Subject: [PATCH 10/23] lint

---
 .../core/inference/sampling/flashinfer_sampling.py    |  6 +-----
 megatron/core/inference/sampling/torch_sampling.py    | 11 +++--------
 2 files changed, 4 insertions(+), 13 deletions(-)

diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index d999d90394e..ca471694f20 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -22,11 +22,7 @@ class FlashInferSampling(Sampling):
     """
 
     def __init__(
-        self,
-        vocab_size: int,
-        rng: torch.Generator,
-        config=None,
-        enable_cuda_graph: bool = False,
+        self, vocab_size: int, rng: torch.Generator, config=None, enable_cuda_graph: bool = False
     ) -> None:
         self._vocab_size = vocab_size
         self._rng = rng
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
index 5e2d438219e..68ff17ee35f 100644
--- a/megatron/core/inference/sampling/torch_sampling.py
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -50,9 +50,7 @@ def sample_from_logits(
         """
         assert isinstance(top_p, float)
         assert isinstance(top_k, int)
-        assert not (
-            top_k > 0 and top_p > 0.0
-        ), "Cannot have top-p and top-k both greater than zero"
+        assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero"
         assert top_p <= 1.0, "top-p should be in (0,1]"
 
         def modify_logits_for_top_k_filtering(logits, top_k):
@@ -117,8 +115,7 @@ def pre_forward_bookkeeping(self, context) -> None:
 
         self._buckets = [(indices, *params) for params, indices in bucket_map.items()]
         self._bucket_index_tensors = [
-            torch.tensor(indices, device=device, dtype=torch.long)
-            for indices, *_ in self._buckets
+            torch.tensor(indices, device=device, dtype=torch.long) for indices, *_ in self._buckets
         ]
 
     def sample_kernel(
@@ -154,9 +151,7 @@ def sample_kernel(
         output = torch.empty(n, device=logits.device, dtype=torch.int64)
         token_list = []
         indices_list = []
-        for idx_tensor, (_, temp, top_k, top_p) in zip(
-            self._bucket_index_tensors, self._buckets
-        ):
+        for idx_tensor, (_, temp, top_k, top_p) in zip(self._bucket_index_tensors, self._buckets):
             if token_to_request_index is None:
                 row_indices = idx_tensor
             else:

From e089240fc9c48e976c8775fbf8db9d0f5d1ebf1a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 29 Apr 2026 12:12:23 -0500
Subject: [PATCH 11/23] Remove unnecessary D2D copy

---
 megatron/core/inference/sampling/flashinfer_sampling.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index ca471694f20..7017c324197 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -75,12 +75,9 @@ def sample_kernel(
             top_p = md["top_p"][token_to_request_index]
 
         if gather_indices is None:
-            # Slice is a view; clone before in-place div_ to avoid mutating the caller.
-            scaled = logits[:n].clone()
+            scaled = logits[:n] / temperature.unsqueeze(1)
         else:
-            # Advanced indexing already returns a new tensor.
-            scaled = logits[gather_indices[:n], :]
-        scaled.div_(temperature.unsqueeze(1))
+            scaled = logits[gather_indices[:n], :] / temperature.unsqueeze(1)
         probs = torch.softmax(scaled, dim=-1)
 
         # Sentinel values disable filtering: top_k=vocab_size keeps all

From bbfa111270c4dfbf957ffb5c265262f16c508777 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 05:21:38 -0500
Subject: [PATCH 12/23] Clean up after merge

---
 .../core/inference/engines/dynamic_engine.py  |  5 --
 megatron/core/inference/sampling/base.py      | 35 ++----------
 .../inference/sampling/flashinfer_sampling.py |  4 +-
 .../core/inference/sampling/torch_sampling.py | 53 +++++++++----------
 .../text_generation_controller.py             | 18 ++-----
 .../test_mtp_cuda_graph_inference.py          | 16 +++---
 .../test_text_generation_controller.py        | 34 +++++-------
 7 files changed, 54 insertions(+), 111 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 518c4a5c25b..f00625d1708 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -395,11 +395,6 @@ def create_cuda_graphs(self, reset_context: bool = True):
             with torch.inference_mode():
                 controller._dynamic_step_forward_logits(input_ids, position_ids)
 
-                # Sampling warm-up. The non-speculative path captures `sample_kernel`
-                # keyed by `("sample", padded_active_request_count)`; the speculative
-                # path additionally captures `("sample_speculative", padded_decode, 0)`
-                # for decode-only graphs. Other backends and mixed graphs run eagerly.
-                controller._dynamic_step_sample_bookkeeping()
                 if controller.num_speculative_tokens > 0:
                     controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
                 else:
diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
index ebf41fc12c5..157d089ec23 100644
--- a/megatron/core/inference/sampling/base.py
+++ b/megatron/core/inference/sampling/base.py
@@ -10,17 +10,12 @@
 class Sampling(ABC):
     """Abstract base for inference sampling backends.
 
-    Subclasses implement `pre_forward_bookkeeping` (per-step setup) and `sample_kernel`
-    (the GPU kernel). CUDA graph wrapping, when applicable, is added by the subclass via
-    `CudaGraphManager`. The wrapper consumes `eager` and `cache_key` kwargs; concrete
-    subclasses without a wrapper still accept and ignore them.
+    Subclasses implement `sample_kernel` (the GPU kernel). CUDA graph wrapping, when
+    applicable, is added by the subclass via `CudaGraphManager`. The wrapper consumes
+    `eager` and `cache_key` kwargs; concrete subclasses without a wrapper still accept
+    and ignore them.
     """
 
-    @abstractmethod
-    def pre_forward_bookkeeping(self, context) -> None:
-        """Prepare sampling state before the forward pass."""
-        ...
-
     @abstractmethod
     def sample_kernel(
         self,
@@ -53,28 +48,6 @@ def sample_kernel(
         """
         ...
 
-    def sample(
-        self,
-        logits: Tensor,
-        n: int,
-        context,
-        *,
-        eager: bool = False,
-        cache_key: Any = None,
-        gather_indices: Optional[Tensor] = None,
-        token_to_request_index: Optional[Tensor] = None,
-    ) -> Tensor:
-        """Sample `n` tokens, optionally with CUDA graph capture/replay."""
-        return self.sample_kernel(
-            logits,
-            n,
-            context,
-            eager=eager,
-            cache_key=cache_key,
-            gather_indices=gather_indices,
-            token_to_request_index=token_to_request_index,
-        )
-
     def sample_speculative(
         self,
         required_logits: Tensor,
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 7017c324197..569942adfd0 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -36,9 +36,6 @@ def __init__(
                 inline_capture=True,
             )
 
-    def pre_forward_bookkeeping(self, context) -> None:
-        """No-op; FlashInfer needs no per-step bookkeeping."""
-
     def sample_kernel(
         self,
         logits: Tensor,
@@ -63,6 +60,7 @@ def sample_kernel(
         Returns:
             Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
         """
+        # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
         md = context.active_request_metadata
         if token_to_request_index is None:
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
index 68ff17ee35f..7d02d61602f 100644
--- a/megatron/core/inference/sampling/torch_sampling.py
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -18,8 +18,6 @@ class TorchSampling(Sampling):
     def __init__(self, rng: torch.Generator, vocab_size: int) -> None:
         self._rng = rng
         self._vocab_size = vocab_size
-        self._buckets: List[Tuple] = []
-        self._bucket_index_tensors: List[Tensor] = []
 
     @staticmethod
     def sample_from_logits(
@@ -33,7 +31,7 @@ def sample_from_logits(
     ) -> Tensor:
         """Sample tokens from logits with temperature, top-k, and top-p filtering.
 
-        Shared between dynamic batching (`TorchSampling._sampling_func`) and static
+        Shared between dynamic batching (`TorchSampling.sample_kernel`) and static
         batching (`TextGenerationController.sample_from_logits`).
 
         Args:
@@ -96,28 +94,6 @@ def modify_logits_for_top_p_filtering(logits, top_p):
 
         return sampled
 
-    def pre_forward_bookkeeping(self, context) -> None:
-        """Group active requests into sampling buckets by `(temperature, top_k, top_p)`.
-
-        Pre-computes a GPU index tensor per bucket so subsequent steps avoid H2D copies.
-        """
-        active_request_count = context.total_request_count - context.paused_request_count
-        md = context.active_request_metadata
-        device = torch.cuda.current_device()
-
-        bucket_map: dict = defaultdict(list)
-        temp = md["temperature"][:active_request_count].tolist()
-        top_k = md["top_k"][:active_request_count].tolist()
-        top_p = md["top_p"][:active_request_count].tolist()
-
-        for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
-            bucket_map[(t, k, p)].append(request_index)
-
-        self._buckets = [(indices, *params) for params, indices in bucket_map.items()]
-        self._bucket_index_tensors = [
-            torch.tensor(indices, device=device, dtype=torch.long) for indices, *_ in self._buckets
-        ]
-
     def sample_kernel(
         self,
         logits: Tensor,
@@ -129,13 +105,13 @@ def sample_kernel(
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
     ) -> Tensor:
-        """Sample by iterating over pre-computed sampling buckets.
+        """Bucket active requests by `(temperature, top_k, top_p)` and sample each bucket.
 
         Args:
             logits: Logits tensor of shape `[>=n, vocab_size]`.
             n: Number of rows to sample.
-            context: The active DynamicInferenceContext (unused; kept for ABC parity).
-            eager: Accepted for API symmetry with FlashInfer; ignored (no wrapper here).
+            context: The active DynamicInferenceContext.
+            eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).
             cache_key: Accepted for API symmetry; ignored.
             gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: When set, the loop dispatches per-token rather than
@@ -144,14 +120,33 @@ def sample_kernel(
         Returns:
             Sampled token ids of shape `[n]`.
         """
+        # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
+
+        # Group active requests into sampling buckets by (temperature, top_k, top_p).
+        active_request_count = context.total_request_count - context.paused_request_count
+        md = context.active_request_metadata
+        device = torch.cuda.current_device()
+
+        bucket_map: dict = defaultdict(list)
+        temp = md["temperature"][:active_request_count].tolist()
+        top_k = md["top_k"][:active_request_count].tolist()
+        top_p = md["top_p"][:active_request_count].tolist()
+        for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
+            bucket_map[(t, k, p)].append(request_index)
+
+        buckets: List[Tuple] = [(indices, *params) for params, indices in bucket_map.items()]
+        bucket_index_tensors: List[Tensor] = [
+            torch.tensor(indices, device=device, dtype=torch.long) for indices, *_ in buckets
+        ]
+
         if gather_indices is not None:
             logits = logits[gather_indices[:n], :]
 
         output = torch.empty(n, device=logits.device, dtype=torch.int64)
         token_list = []
         indices_list = []
-        for idx_tensor, (_, temp, top_k, top_p) in zip(self._bucket_index_tensors, self._buckets):
+        for idx_tensor, (_, temp, top_k, top_p) in zip(bucket_index_tensors, buckets):
             if token_to_request_index is None:
                 row_indices = idx_tensor
             else:
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index fe332995bea..bca82f14c58 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -169,11 +169,11 @@ def _init_dynamic_sampling_tensors(self):
         #     - `self._sampled_tokens_cuda` is pre-allocated by `_init_mtp_sampling_tensors`.
         #     - The tensor cannot be reused between the Triton kernel and the sampling graph.
         # Non-speculative path:
-        #     - `self._sampled_tokens_cuda` is rebound to the output of `sample()`,
+        #     - `self._sampled_tokens_cuda` is rebound to the output of `sample_kernel`,
         #     which uses CudaGraphManager syntactic sugar to keep it as a static tensor.
         self._sampled_tokens_cuda = None
 
-        # Sampling backend: dispatches per-step bookkeeping and the sampling kernel.
+        # Sampling backend: provides the sampling kernel.
         if self._sampling_backend == "flashinfer":
             self._sampling: Sampling = FlashInferSampling(
                 self.vocab_size,
@@ -661,17 +661,9 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
         # Copy logits to contiguous buffer.
         if self._enable_cuda_graph:
             self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits[:, :logits_seq_len, :])
-        elif self._sampling_backend == "flashinfer":
-            # FlashInfer kernels require contiguous inputs.
-            self._all_logits_cuda = logits.contiguous()
         else:
             self._all_logits_cuda = logits
 
-    def _dynamic_step_sample_bookkeeping(self):
-        """Perform bookkeeping necessary to sample logits for dynamic batching."""
-        context = self.inference_wrapped_model.inference_context
-        self._sampling.pre_forward_bookkeeping(context)
-
     def _rewind_kv_cache(self) -> tuple:
         """Update the KV cache bookkeeping for speculative decoding.
 
@@ -747,7 +739,7 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
         Returns:
             Tensor: Sampled tokens of shape [num_requests].
         """
-        return self._sampling.sample(
+        return self._sampling.sample_kernel(
             logits_2d,
             logits_2d.shape[0],
             self.inference_wrapped_model.inference_context,
@@ -1063,7 +1055,7 @@ def _dynamic_step_sample_logits(self):
             if context.config.materialize_only_last_token_logits
             else context.active_request_last_token_idxs
         )
-        self._sampled_tokens_cuda = self._sampling.sample(
+        self._sampled_tokens_cuda = self._sampling.sample_kernel(
             self._all_logits_cuda.squeeze(0),
             n,
             context,
@@ -1770,8 +1762,6 @@ async def async_generate_output_tokens_dynamic_batch(
             range_push("sampling")
             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.
                 nvtx_range_push("mtp-spec-decoding/verify")
diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
index fe1d607e818..f74d1bbed13 100644
--- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
+++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
@@ -420,10 +420,10 @@ def test_cuda_graph_sp_padding_end_to_end(self, mtp_use_repeated_layer):
             ctrl._mtp_resolved_padded_count = padded_count
             context._using_cuda_graph_this_step = True
 
-            ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-            ctrl._sampling._bucket_index_tensors = [
-                torch.arange(active_request_count, device='cuda', dtype=torch.long)
-            ]
+            # Greedy sampling for all active requests.
+            context.active_request_metadata["temperature"][:active_request_count] = 1.0
+            context.active_request_metadata["top_k"][:active_request_count] = 1
+            context.active_request_metadata["top_p"][:active_request_count] = 0.0
 
             ctrl._compute_serial_mtp_and_sample()
 
@@ -521,10 +521,10 @@ def _run_mtp(use_cuda_graph):
                 unwrapped._decoder_hidden_states_cache = local_hidden
 
                 ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda')
-                ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-                ctrl._sampling._bucket_index_tensors = [
-                    torch.arange(active_request_count, device='cuda', dtype=torch.long)
-                ]
+                # Greedy sampling for all active requests.
+                context.active_request_metadata["temperature"][:active_request_count] = 1.0
+                context.active_request_metadata["top_k"][:active_request_count] = 1
+                context.active_request_metadata["top_p"][:active_request_count] = 0.0
 
                 ctrl._compute_serial_mtp_and_sample()
                 # CUDA graph replay is asynchronous, and this test reuses the controller's
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index a0ca39ef6a0..954b9fed915 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -333,9 +333,6 @@ def test_sample_from_dynamic_logits(
         context.num_prefill_requests = 0
         context.pad_active_slices()
 
-        # Bookkeeping.
-        self.text_generation_controller._dynamic_step_sample_bookkeeping()
-
         # Sampling.
         logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda()
         self.text_generation_controller._all_logits_cuda = logits
@@ -1002,11 +999,11 @@ def mock_sampling_func(logits, *args, **kwargs):
                 # The verification logic only uses base tokens, so we can return zeros here.
                 return torch.zeros((12,), dtype=torch.long, device='cuda')
 
-        # Override sampling to return our predictable mock outputs
-        self.text_generation_controller._sampling._buckets = [([0, 1], 1.0, 1, 0.0)]
-        self.text_generation_controller._sampling._bucket_index_tensors = [
-            torch.tensor([0, 1], device='cuda', dtype=torch.long)
-        ]
+        # Drive both requests into a single greedy bucket via metadata, then mock
+        # the per-bucket sampling function to return our predictable outputs.
+        ctx.active_request_metadata["temperature"][:2] = 1.0
+        ctx.active_request_metadata["top_k"][:2] = 1
+        ctx.active_request_metadata["top_p"][:2] = 0.0
         self.text_generation_controller._sampling._sampling_func = mock.MagicMock(
             side_effect=mock_sampling_func
         )
@@ -1242,15 +1239,11 @@ def test_speculative_multinomial_sampling(self):
         # Base logits shape: [1, 8, vocab_size]
         logits = torch.randn(1, 8, self.vocab_size, device='cuda')
 
-        # Set up a bucket that forces multinomial sampling (top_p = 0.9, top_k = 0)
-        # _buckets format: (indices, temp, top_k, top_p)
-        self.text_generation_controller._sampling._buckets = [([0, 1], 1.0, 0, 0.9)]
-        self.text_generation_controller._sampling._bucket_index_tensors = [
-            torch.tensor([0, 1], device='cuda', dtype=torch.long)
-        ]
-
-        # Since we are actually testing the internal math of `_torch_sampling_func` handling the shapes,
-        # we DO NOT mock `_torch_sampling_func` here. We want it to run natively to prove it doesn't crash.
+        # Drive sampling onto the multinomial path (top_p > 0, top_k == 0) via metadata.
+        # We do NOT mock `_sampling_func`: we want it to run natively to prove it doesn't crash.
+        ctx.active_request_metadata["temperature"][:2] = 1.0
+        ctx.active_request_metadata["top_k"][:2] = 0
+        ctx.active_request_metadata["top_p"][:2] = 0.9
 
         self.text_generation_controller._all_logits_cuda = logits
         try:
@@ -1495,10 +1488,9 @@ def test_mtp_sp_padding_real_ranks(self, active_request_count):
         ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda')
 
         # Greedy sampling: top_k=1 selects the argmax token deterministically.
-        ctrl._sampling._buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)]
-        ctrl._sampling._bucket_index_tensors = [
-            torch.arange(active_request_count, device='cuda', dtype=torch.long)
-        ]
+        ctx.active_request_metadata["temperature"][:active_request_count] = 1.0
+        ctx.active_request_metadata["top_k"][:active_request_count] = 1
+        ctx.active_request_metadata["top_p"][:active_request_count] = 0.0
 
         # Run the MTP forward pass
         ctrl._compute_serial_mtp_and_sample()

From ba2566fbf6a2d7d82e4295dc3148dcc758c76aa4 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 05:55:21 -0500
Subject: [PATCH 13/23] Expand test coverage to FlashInfer

---
 .../inference/engines/test_dynamic_engine.py  | 62 ++++++++++++++-----
 .../test_text_generation_controller.py        | 24 ++++---
 2 files changed, 62 insertions(+), 24 deletions(-)

diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index 24efaea9e1d..e5ebaacfaf1 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -2261,7 +2261,8 @@ def set_epoch(epoch):
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
-    def test_speculative_decoding_with_early_termination(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_decoding_with_early_termination(self, materialize_only_last_token_logits):
         """Test that speculative decoding handles premature request termination safely
         (e.g. hitting max_sequence_length mid-speculative-batch)."""
 
@@ -2274,7 +2275,7 @@ def test_speculative_decoding_with_early_termination(self):
             max_sequence_length=7,  # Will force termination after 3 tokens
             model_provider="gpt",
             num_speculative_tokens=3,
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
         )
 
         env = self._build_test_env(test_config)
@@ -2299,6 +2300,8 @@ def mock_mtp_forward(*args, **kwargs):
             unwrapped_model._decoder_hidden_states_cache = torch.zeros(
                 tokens.size(1), 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
             )
+            if test_config.materialize_only_last_token_logits:
+                base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
             return base_logits
 
         def mock_compute_mtp_single_step(
@@ -2335,7 +2338,8 @@ def mock_compute_mtp_single_step(
 
     @pytest.mark.internal
     @torch.inference_mode()
-    def test_speculative_block_boundary_crossing(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_block_boundary_crossing(self, materialize_only_last_token_logits):
         """Test to verify KV cache block boundary crossing logic.
 
         When a request fills exactly one block and speculative decoding generates
@@ -2350,7 +2354,7 @@ def test_speculative_block_boundary_crossing(self):
             context_block_size_tokens=256,  # Exactly matches prompt length
             context_max_requests=16,
             model_provider="gpt",
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             use_fixed_output_lengths=True,
         )
         env = self._build_test_env(test_config)
@@ -2391,7 +2395,8 @@ def test_speculative_block_boundary_crossing(self):
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
-    def test_speculative_stop_word_hit(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_stop_word_hit(self, materialize_only_last_token_logits):
         """Test that if an accepted speculative token completes a stop word,
         the request correctly triggers the stop logic without crashing."""
 
@@ -2401,7 +2406,7 @@ def test_speculative_stop_word_hit(self):
             max_prompt_length=4,
             num_tokens_to_generate=10,
             num_speculative_tokens=2,
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
         )
         env = self._build_test_env(test_config)
@@ -2424,6 +2429,8 @@ def mock_deterministic_forward(*args, **kwargs):
             unwrapped_model._decoder_hidden_states_cache = torch.zeros(
                 s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
             )
+            if test_config.materialize_only_last_token_logits:
+                base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
             return base_logits
 
         def mock_compute_mtp_single_step(
@@ -2477,7 +2484,8 @@ def mock_compute_mtp_single_step(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
-    def test_speculative_long_stop_word_hit(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_long_stop_word_hit(self, materialize_only_last_token_logits):
         """Test that if an accepted speculative token completes a long stop word
         (length > num_speculative_tokens), it is correctly detected."""
 
@@ -2487,7 +2495,7 @@ def test_speculative_long_stop_word_hit(self):
             max_prompt_length=4,
             num_tokens_to_generate=10,
             num_speculative_tokens=2,
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
         )
         env = self._build_test_env(test_config)
@@ -2510,6 +2518,8 @@ def mock_deterministic_forward(*args, **kwargs):
             unwrapped_model._decoder_hidden_states_cache = torch.zeros(
                 s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
             )
+            if test_config.materialize_only_last_token_logits:
+                base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
             return base_logits
 
         def mock_compute_mtp_single_step(
@@ -2559,7 +2569,10 @@ def mock_compute_mtp_single_step(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
-    def test_speculative_stop_word_truncates_trailing_tokens(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_stop_word_truncates_trailing_tokens(
+        self, materialize_only_last_token_logits
+    ):
         """Test that when a stop word lands in the middle of speculative tokens,
         the extra tokens generated after the stop word are removed.
 
@@ -2574,7 +2587,7 @@ def test_speculative_stop_word_truncates_trailing_tokens(self):
             max_prompt_length=4,
             num_tokens_to_generate=10,
             num_speculative_tokens=2,
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
         )
         env = self._build_test_env(test_config)
@@ -2597,6 +2610,8 @@ def mock_deterministic_forward(*args, **kwargs):
             unwrapped_model._decoder_hidden_states_cache = torch.zeros(
                 s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
             )
+            if test_config.materialize_only_last_token_logits:
+                base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
             return base_logits
 
         def mock_compute_mtp_single_step(
@@ -2677,9 +2692,14 @@ def mock_compute_mtp_single_step(
             "non_divisible_boundary",
         ],
     )
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     @torch.inference_mode()
     def test_speculative_tokens_exceed_max_sequence_length(
-        self, prompt_length, num_tokens_to_generate, num_speculative_tokens
+        self,
+        prompt_length,
+        num_tokens_to_generate,
+        num_speculative_tokens,
+        materialize_only_last_token_logits,
     ):
         """Test that speculative decoding correctly trims output when speculative
         tokens would push the sequence beyond max_sequence_length.
@@ -2698,7 +2718,7 @@ def test_speculative_tokens_exceed_max_sequence_length(
             num_tokens_to_generate=num_tokens_to_generate,
             max_sequence_length=max_sequence_length,
             num_speculative_tokens=num_speculative_tokens,
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
             # Disable positional embeddings so speculative position IDs
             # beyond max_sequence_length don't cause out-of-bounds lookups.
@@ -2805,8 +2825,11 @@ def test_detokenize_stop_sequence_flag(self, detokenize_stop_sequence):
     @pytest.mark.parametrize(
         "acceptance_mode", ["all_rejected", "all_accepted"], ids=["all_rejected", "all_accepted"]
     )
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     @torch.inference_mode()
-    def test_speculative_sequence_length_double_counting(self, acceptance_mode):
+    def test_speculative_sequence_length_double_counting(
+        self, acceptance_mode, materialize_only_last_token_logits
+    ):
         """Test to verify active_sequence_lengths is not double-counted.
 
         If active sequence length is double-counted during speculative decoding,
@@ -2827,7 +2850,7 @@ def test_speculative_sequence_length_double_counting(self, acceptance_mode):
             context_max_requests=16,
             num_speculative_tokens=2,
             model_provider="gpt",
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             use_fixed_output_lengths=False,
             context_max_tokens=512,
             position_embedding_type="none",
@@ -2854,6 +2877,8 @@ def mock_mtp_forward(*args, **kwargs):
                 model._decoder_hidden_states_cache = torch.zeros(
                     s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
                 )
+                if test_config.materialize_only_last_token_logits:
+                    base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
                 return base_logits
 
             def mock_compute_mtp(*args_mtp, **kwargs_mtp):
@@ -2925,7 +2950,10 @@ def deterministic_mtp(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
-    def test_speculative_decoding_with_eviction_and_swapping(self):
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_decoding_with_eviction_and_swapping(
+        self, materialize_only_last_token_logits
+    ):
         """Test that speculative decoding works correctly when requests are paused and evicted.
 
         This exercises the `_swap_book_keeping_tensors` logic with the 2D `new_speculative_tokens`
@@ -2942,7 +2970,7 @@ def test_speculative_decoding_with_eviction_and_swapping(self):
             context_buffer_size_gb=0.00064,  # 640 KB
             context_paused_buffer_size_gb=0.0,  # 0 paused buffer forces immediate eviction
             model_provider="gpt",
-            materialize_only_last_token_logits=False,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
             use_fixed_output_lengths=True,
         )
 
@@ -2966,6 +2994,8 @@ def mock_safe_forward(*args, **kwargs):
             unwrapped_model._decoder_hidden_states_cache = torch.zeros(
                 s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16
             )
+            if test_config.materialize_only_last_token_logits:
+                base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0)
             return base_logits
 
         def mock_compute_mtp_single_step(
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 954b9fed915..30cfa72ce89 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -959,10 +959,22 @@ def test_dynamic_top_n_logprobs_calculation(
                     ), f"Request {req_idx}, token {token_idx}: expected {top_n} indices"
 
     @pytest.mark.internal
-    def test_speculative_verify_tokens(self):
+    @pytest.mark.parametrize("backend", ["torch", "flashinfer"])
+    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
+    def test_speculative_verify_tokens(
+        self, backend: str, materialize_only_last_token_logits: bool
+    ):
         """Test consecutive token acceptance logic for speculative decoding."""
+        if backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         self.setup_model(
-            torch.float32, static=False, num_speculative_tokens=2, max_requests=2, mtp_num_layers=2
+            torch.float32,
+            static=False,
+            num_speculative_tokens=2,
+            max_requests=2,
+            mtp_num_layers=2,
+            sampling_backend=backend,
+            materialize_only_last_token_logits=materialize_only_last_token_logits,
         )
 
         # Enable speculative decoding
@@ -999,12 +1011,8 @@ def mock_sampling_func(logits, *args, **kwargs):
                 # The verification logic only uses base tokens, so we can return zeros here.
                 return torch.zeros((12,), dtype=torch.long, device='cuda')
 
-        # Drive both requests into a single greedy bucket via metadata, then mock
-        # the per-bucket sampling function to return our predictable outputs.
-        ctx.active_request_metadata["temperature"][:2] = 1.0
-        ctx.active_request_metadata["top_k"][:2] = 1
-        ctx.active_request_metadata["top_p"][:2] = 0.0
-        self.text_generation_controller._sampling._sampling_func = mock.MagicMock(
+        # Override sampling to return our predictable mock outputs.
+        self.text_generation_controller._sampling.sample_kernel = mock.MagicMock(
             side_effect=mock_sampling_func
         )
 

From ed80ad977baec315c9b464f6f14327888018d4cb Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 08:23:47 -0500
Subject: [PATCH 14/23] Self-review

---
 .../inference/contexts/dynamic_context.py     | 12 ++----
 megatron/core/inference/sampling/base.py      | 27 ++++++-------
 .../inference/sampling/flashinfer_sampling.py | 39 +++++++++++--------
 .../core/inference/sampling/torch_sampling.py | 37 ++++++++----------
 .../text_generation_controller.py             |  9 ++---
 .../test_text_generation_controller.py        |  2 +-
 6 files changed, 58 insertions(+), 68 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 9afdb609f33..d8564a6e0fb 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1055,9 +1055,6 @@ def initialize_all_tensors(self) -> None:
         ].view(torch.int32)
         _off += _req_int32_bytes
 
-        # Per-request last-token row indices, used as gather indices by sampling kernels
-        # (e.g. FlashInfer) so the call site can pass a fixed-shape tensor on the padded
-        # batch dimension.
         self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
 
         # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
@@ -1367,14 +1364,12 @@ def build_active_slices(self, batch_size: int):
                 self.request_metadata[label][padded_slice], non_blocking=True
             )
 
-        # Cumsum of query lengths gives per-request last-token row indices.
-        # Padded slots get garbage here; pad_active_slices fills them with 0.
         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] -= 1
+        self.active_request_last_token_idxs[:batch_size].sub_(1)
 
     def pad_active_slices(self):
         """Pad the active slices of specific tensors."""
@@ -1403,9 +1398,8 @@ 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: neutral defaults so per-row sampling kernels (FlashInfer)
-        # produce harmless output for padded slots. top_k=0 / top_p=0.0 are sentinels
-        # for "no filter" in the FlashInfer wrapper.
+
+        # 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)
diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
index 157d089ec23..8aa4c416c27 100644
--- a/megatron/core/inference/sampling/base.py
+++ b/megatron/core/inference/sampling/base.py
@@ -10,10 +10,7 @@
 class Sampling(ABC):
     """Abstract base for inference sampling backends.
 
-    Subclasses implement `sample_kernel` (the GPU kernel). CUDA graph wrapping, when
-    applicable, is added by the subclass via `CudaGraphManager`. The wrapper consumes
-    `eager` and `cache_key` kwargs; concrete subclasses without a wrapper still accept
-    and ignore them.
+    Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`.
     """
 
     @abstractmethod
@@ -23,25 +20,21 @@ def sample_kernel(
         n: int,
         context,
         *,
-        eager: bool = False,
-        cache_key: Any = None,
         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.
 
-        `eager` and `cache_key` are consumed by the `CudaGraphManager` wrapper when one
-        is installed; unwrapped subclasses accept and ignore them.
-
         Args:
             logits: Logits tensor of shape `[>=n, vocab_size]`.
             n: Number of rows to sample.
             context: The active DynamicInferenceContext.
-            eager: If True, skip CUDA graph capture/replay (consumed by the wrapper).
-            cache_key: Hashable key for runner lookup (consumed by the wrapper).
             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.
@@ -56,20 +49,23 @@ def sample_speculative(
         num_speculative_tokens: int,
         context,
         *,
+        gather_indices: Optional[Tensor] = None,
         eager: bool = False,
         cache_key: Any = None,
-        gather_indices: Optional[Tensor] = 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, `required_logits` is the full per-token logits buffer
-        (constant shape across steps); the kernel selects rows via `logits[gather_indices[:n], :]`.
+        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
@@ -87,8 +83,7 @@ def sample_speculative(
             required_logits,
             num_tokens,
             context,
-            eager=eager,
-            cache_key=cache_key,
             gather_indices=gather_indices,
             token_to_request_index=token_to_request_index,
+            eager=True,
         )
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 569942adfd0..1091f0979d9 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -15,11 +15,7 @@
 
 
 class FlashInferSampling(Sampling):
-    """Fused FlashInfer sampling, with optional CUDA graph capture/replay.
-
-    Unlike `TorchSampling`, FlashInfer kernels accept per-row parameter tensors
-    (temperature, top_k, top_p) directly, so no bucketing is required.
-    """
+    """Fused FlashInfer sampling, with optional CUDA graph capture/replay."""
 
     def __init__(
         self, vocab_size: int, rng: torch.Generator, config=None, enable_cuda_graph: bool = False
@@ -35,6 +31,13 @@ def __init__(
                 need_backward=False,
                 inline_capture=True,
             )
+            CudaGraphManager(
+                config,
+                self,
+                function_name="sample_speculative",
+                need_backward=False,
+                inline_capture=True,
+            )
 
     def sample_kernel(
         self,
@@ -42,26 +45,29 @@ def sample_kernel(
         n: int,
         context,
         *,
-        eager: bool = False,
-        cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
+        eager: bool = False,
+        cache_key: Any = None,
     ) -> Tensor:
         """FlashInfer fused top-k / top-p sampling kernel.
 
-        Reads sampling parameters per-row from `context.active_request_metadata`,
-        applies temperature scaling and top-k/top-p filtering, then samples via
-        `flashinfer.sampling.top_k_top_p_sampling_from_probs`.
-
-        When wrapped by `CudaGraphManager`, `eager` and `cache_key` are consumed by
-        the wrapper before this body runs. When unwrapped (no CUDA graphs), they are
-        accepted and ignored so callers can pass them unconditionally.
+        Args:
+            logits: Logits tensor of shape `[>=n, vocab_size]`.
+            n: Number of rows to sample.
+            context: The active DynamicInferenceContext.
+            gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
+            token_to_request_index: When set, sampling parameters are gathered per-token
+                rather than per-request (used by the speculative path).
+            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.
         """
         # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
+
+        # Read GPU sampling parameters directly from `context.active_request_metadata`.
         md = context.active_request_metadata
         if token_to_request_index is None:
             temperature = md["temperature"][:n]
@@ -78,8 +84,9 @@ def sample_kernel(
             scaled = logits[gather_indices[:n], :] / temperature.unsqueeze(1)
         probs = torch.softmax(scaled, dim=-1)
 
-        # Sentinel values disable filtering: top_k=vocab_size keeps all
-        # tokens, top_p=1.0 keeps the full probability mass.
+        # Sentinel values disable filtering:
+        # top_k=vocab_size keeps all tokens, top_p=1.0 keeps the full probability mass.
+        # TODO: Consider changing the disable flags in the `InferenceRequest`.
         top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
         top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
         output = torch.empty(n, device=logits.device, dtype=torch.int64)
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
index 7d02d61602f..79491add5ab 100644
--- a/megatron/core/inference/sampling/torch_sampling.py
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -31,8 +31,7 @@ def sample_from_logits(
     ) -> Tensor:
         """Sample tokens from logits with temperature, top-k, and top-p filtering.
 
-        Shared between dynamic batching (`TorchSampling.sample_kernel`) and static
-        batching (`TextGenerationController.sample_from_logits`).
+        Shared between dynamic batching and static batching.
 
         Args:
             last_token_logits: Logits of shape `[batch_size, vocab_size]`.
@@ -73,8 +72,7 @@ def modify_logits_for_top_p_filtering(logits, top_p):
         if top_k == 1:
             return torch.argmax(last_token_logits, dim=-1)
 
-        # Clone needed: .div_() and masked_fill_() below modify in-place,
-        # which would mutate the caller's tensor without this clone.
+        # Clone needed: .div_() and masked_fill_() below modify in-place.
         last_token_logits = last_token_logits.clone()
         if temperature != 1.0:
             last_token_logits.div_(temperature)
@@ -100,10 +98,10 @@ def sample_kernel(
         n: int,
         context,
         *,
-        eager: bool = False,
-        cache_key: Any = None,
         gather_indices: Optional[Tensor] = None,
         token_to_request_index: Optional[Tensor] = None,
+        eager: bool = False,
+        cache_key: Any = None,
     ) -> Tensor:
         """Bucket active requests by `(temperature, top_k, top_p)` and sample each bucket.
 
@@ -111,11 +109,11 @@ def sample_kernel(
             logits: Logits tensor of shape `[>=n, vocab_size]`.
             n: Number of rows to sample.
             context: The active DynamicInferenceContext.
-            eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).
-            cache_key: Accepted for API symmetry; ignored.
             gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
             token_to_request_index: When set, the loop dispatches per-token rather than
                 per-request (used by the speculative path).
+            eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).
+            cache_key: Accepted for API symmetry; ignored.
 
         Returns:
             Sampled token ids of shape `[n]`.
@@ -151,22 +149,19 @@ def sample_kernel(
                 row_indices = idx_tensor
             else:
                 row_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0]
-            token_list.append(self._sampling_func(logits[row_indices, :], temp, top_k, top_p))
+            token_list.append(
+                TorchSampling.sample_from_logits(
+                    logits[row_indices, :],
+                    temp,
+                    top_k,
+                    top_p,
+                    generator=self._rng,
+                    vocab_size=self._vocab_size,
+                )
+            )
             indices_list.append(row_indices)
 
         sampled_tokens = torch.cat(token_list, dim=0)
         sampled_indices = torch.cat(indices_list, dim=0)
         output[sampled_indices] = sampled_tokens
         return output
-
-    def _sampling_func(
-        self, last_token_logits: Tensor, temperature: float, top_k: int, top_p: float
-    ) -> Tensor:
-        return TorchSampling.sample_from_logits(
-            last_token_logits,
-            temperature,
-            top_k,
-            top_p,
-            generator=self._rng,
-            vocab_size=self._vocab_size,
-        )
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index bca82f14c58..decc692c4f1 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -914,9 +914,8 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
 
-        # Sampling-side request counts: padded when running a captured graph so the
-        # cache key buckets, actual otherwise. Verify uses the actual counts so the
-        # downstream Triton kernels still operate on the real workload.
+        # Sampling-side request counts: padded when running a captured graph.
+        # Verify uses the actual counts so the Triton kernels operate on the real workload.
         use_graph_for_sampling = (
             self._sampling_backend == "flashinfer"
             and self._enable_cuda_graph
@@ -961,13 +960,13 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
             sample_num_prefill,
             self.num_speculative_tokens,
             context,
+            gather_indices=sample_gather_indices,
             eager=not use_graph_for_sampling,
             cache_key=(
                 ("sample_speculative", sample_num_decode, sample_num_prefill)
                 if use_graph_for_sampling
                 else None
             ),
-            gather_indices=sample_gather_indices,
         )
         nvtx_range_pop("mtp-spec-decoding/verify/sample")
 
@@ -1059,9 +1058,9 @@ def _dynamic_step_sample_logits(self):
             self._all_logits_cuda.squeeze(0),
             n,
             context,
+            gather_indices=gather_indices,
             eager=not use_graph,
             cache_key=("sample", n) if use_graph else None,
-            gather_indices=gather_indices,
         )
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 30cfa72ce89..a4a06eb8897 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -1248,7 +1248,7 @@ def test_speculative_multinomial_sampling(self):
         logits = torch.randn(1, 8, self.vocab_size, device='cuda')
 
         # Drive sampling onto the multinomial path (top_p > 0, top_k == 0) via metadata.
-        # We do NOT mock `_sampling_func`: we want it to run natively to prove it doesn't crash.
+        # We do NOT mock the sampling kernel: we want it to run natively to prove it doesn't crash.
         ctx.active_request_metadata["temperature"][:2] = 1.0
         ctx.active_request_metadata["top_k"][:2] = 0
         ctx.active_request_metadata["top_p"][:2] = 0.9

From 60955a21c865acfa25c70b21885b3e000d397968 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 08:30:32 -0500
Subject: [PATCH 15/23] Default to FlashInfer sampling

---
 megatron/core/inference/config.py             |  9 ++-
 .../inference/sampling/flashinfer_sampling.py |  1 -
 megatron/training/arguments.py                | 14 +++--
 .../inference/engines/test_dynamic_engine.py  | 55 ++++++++++++++++---
 4 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index 86ce980aa07..a13c5686b99 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import warnings
 from dataclasses import InitVar, dataclass
 from enum import Enum
 from typing import List, Literal, Optional, Tuple
@@ -297,7 +298,7 @@ class InferenceConfig:
     Defaults to 0, which means no logging.
     """
 
-    sampling_backend: Literal['torch', 'flashinfer'] = 'torch'
+    sampling_backend: Literal['torch', 'flashinfer'] = 'flashinfer'
     """Which sampling kernels to use during inference."""
 
     request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None
@@ -328,6 +329,8 @@ def __post_init__(self, verbose: bool):
             try:
                 import flashinfer  # noqa: F401
             except ImportError:
-                raise ValueError(
-                    "sampling_backend='flashinfer' requires flashinfer to be installed."
+                warnings.warn(
+                    "sampling_backend='flashinfer' requested but flashinfer is not "
+                    "installed; falling back to 'torch'."
                 )
+                self.sampling_backend = 'torch'
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 1091f0979d9..7b4271274fc 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -22,7 +22,6 @@ def __init__(
     ) -> None:
         self._vocab_size = vocab_size
         self._rng = rng
-        self._enable_cuda_graph = enable_cuda_graph
         if enable_cuda_graph and config is not None and config.cuda_graph_impl == "local":
             CudaGraphManager(
                 config,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index db976ebbb8b..7c144466943 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -9,6 +9,7 @@
 from pathlib import Path
 import re
 import types
+import warnings
 
 import torch
 import torch.nn.functional as F
@@ -1011,10 +1012,11 @@ def validate_args(args, defaults={}):
         try:
             import flashinfer  # noqa: F401
         except ImportError:
-            raise ValueError(
-                "--inference-dynamic-batching-sampling-backend=flashinfer "
-                "requires flashinfer to be installed."
+            warnings.warn(
+                "--inference-dynamic-batching-sampling-backend=flashinfer requested "
+                "but flashinfer is not installed; falling back to 'torch'."
             )
+            args.inference_dynamic_batching_sampling_backend = 'torch'
 
     if args.use_megatron_fsdp:
         # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions.
@@ -2004,9 +2006,11 @@ def _add_inference_args(parser):
                        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',
+                       type=str, default='flashinfer',
                        choices=['torch', 'flashinfer'],
-                       help='Which sampling kernels to use during inference.')
+                       help='Which sampling kernels to use during inference. '
+                            'Falls back to "torch" with a warning if "flashinfer" '
+                            'is requested but the package is not installed.')
     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.')
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index e5ebaacfaf1..38d546fc878 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -146,6 +146,7 @@ class DynamicEngineTestConfig:
     track_generated_token_events: bool = False
     num_speculative_tokens: int = 0
     position_embedding_type: str = "learned_absolute"
+    sampling_backend: str = 'torch'
 
     def __post_init__(self):
 
@@ -274,6 +275,7 @@ def _build_inference_context(
                 unified_memory_level=0,  # unit tests currently broken with UVM
                 track_generated_token_events=test_config.track_generated_token_events,
                 num_speculative_tokens=test_config.num_speculative_tokens,
+                sampling_backend=test_config.sampling_backend,
             ),
         )
 
@@ -2261,10 +2263,15 @@ def set_epoch(epoch):
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_speculative_decoding_with_early_termination(self, materialize_only_last_token_logits):
+    def test_speculative_decoding_with_early_termination(
+        self, materialize_only_last_token_logits, sampling_backend
+    ):
         """Test that speculative decoding handles premature request termination safely
         (e.g. hitting max_sequence_length mid-speculative-batch)."""
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
 
         # Set max_sequence_length tight so it terminates during a speculative step
         test_config = DynamicEngineTestConfig(
@@ -2276,6 +2283,7 @@ def test_speculative_decoding_with_early_termination(self, materialize_only_last
             model_provider="gpt",
             num_speculative_tokens=3,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
+            sampling_backend=sampling_backend,
         )
 
         env = self._build_test_env(test_config)
@@ -2338,13 +2346,18 @@ def mock_compute_mtp_single_step(
 
     @pytest.mark.internal
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_speculative_block_boundary_crossing(self, materialize_only_last_token_logits):
+    def test_speculative_block_boundary_crossing(
+        self, materialize_only_last_token_logits, sampling_backend
+    ):
         """Test to verify KV cache block boundary crossing logic.
 
         When a request fills exactly one block and speculative decoding generates
         multiple tokens, the first new token shouldn't incorrectly overwrite the old block.
         """
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         test_config = DynamicEngineTestConfig(
             num_requests=1,
             min_prompt_length=256,
@@ -2356,6 +2369,7 @@ def test_speculative_block_boundary_crossing(self, materialize_only_last_token_l
             model_provider="gpt",
             materialize_only_last_token_logits=materialize_only_last_token_logits,
             use_fixed_output_lengths=True,
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2395,10 +2409,13 @@ def test_speculative_block_boundary_crossing(self, materialize_only_last_token_l
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_speculative_stop_word_hit(self, materialize_only_last_token_logits):
+    def test_speculative_stop_word_hit(self, materialize_only_last_token_logits, sampling_backend):
         """Test that if an accepted speculative token completes a stop word,
         the request correctly triggers the stop logic without crashing."""
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
 
         test_config = DynamicEngineTestConfig(
             num_requests=0,  # We will manually add our request cleanly
@@ -2408,6 +2425,7 @@ def test_speculative_stop_word_hit(self, materialize_only_last_token_logits):
             num_speculative_tokens=2,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2484,10 +2502,15 @@ def mock_compute_mtp_single_step(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_speculative_long_stop_word_hit(self, materialize_only_last_token_logits):
+    def test_speculative_long_stop_word_hit(
+        self, materialize_only_last_token_logits, sampling_backend
+    ):
         """Test that if an accepted speculative token completes a long stop word
         (length > num_speculative_tokens), it is correctly detected."""
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
 
         test_config = DynamicEngineTestConfig(
             num_requests=0,
@@ -2497,6 +2520,7 @@ def test_speculative_long_stop_word_hit(self, materialize_only_last_token_logits
             num_speculative_tokens=2,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2569,9 +2593,10 @@ def mock_compute_mtp_single_step(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     def test_speculative_stop_word_truncates_trailing_tokens(
-        self, materialize_only_last_token_logits
+        self, materialize_only_last_token_logits, sampling_backend
     ):
         """Test that when a stop word lands in the middle of speculative tokens,
         the extra tokens generated after the stop word are removed.
@@ -2580,6 +2605,8 @@ def test_speculative_stop_word_truncates_trailing_tokens(
         (1 base + 2 speculative). If the stop word is [6] and the engine
         generates [5, 6, 7] in one step, token 7 must be truncated so the
         output ends with the stop word [6]."""
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
 
         test_config = DynamicEngineTestConfig(
             num_requests=0,
@@ -2589,6 +2616,7 @@ def test_speculative_stop_word_truncates_trailing_tokens(
             num_speculative_tokens=2,
             materialize_only_last_token_logits=materialize_only_last_token_logits,
             model_provider="gpt",
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2692,6 +2720,7 @@ def mock_compute_mtp_single_step(
             "non_divisible_boundary",
         ],
     )
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     @torch.inference_mode()
     def test_speculative_tokens_exceed_max_sequence_length(
@@ -2700,6 +2729,7 @@ def test_speculative_tokens_exceed_max_sequence_length(
         num_tokens_to_generate,
         num_speculative_tokens,
         materialize_only_last_token_logits,
+        sampling_backend,
     ):
         """Test that speculative decoding correctly trims output when speculative
         tokens would push the sequence beyond max_sequence_length.
@@ -2709,6 +2739,8 @@ def test_speculative_tokens_exceed_max_sequence_length(
         speculative tokens are accepted and the boundary trimming logic is
         actually exercised.
         """
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         max_sequence_length = prompt_length + num_tokens_to_generate
 
         test_config = DynamicEngineTestConfig(
@@ -2723,6 +2755,7 @@ def test_speculative_tokens_exceed_max_sequence_length(
             # Disable positional embeddings so speculative position IDs
             # beyond max_sequence_length don't cause out-of-bounds lookups.
             position_embedding_type="none",
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2825,10 +2858,11 @@ def test_detokenize_stop_sequence_flag(self, detokenize_stop_sequence):
     @pytest.mark.parametrize(
         "acceptance_mode", ["all_rejected", "all_accepted"], ids=["all_rejected", "all_accepted"]
     )
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     @torch.inference_mode()
     def test_speculative_sequence_length_double_counting(
-        self, acceptance_mode, materialize_only_last_token_logits
+        self, acceptance_mode, materialize_only_last_token_logits, sampling_backend
     ):
         """Test to verify active_sequence_lengths is not double-counted.
 
@@ -2841,6 +2875,8 @@ def test_speculative_sequence_length_double_counting(
         a faulty formula that adds accepted_tokens on top of the KV length will
         over-count by 2 per step, finishing the request after only 4 of 6 tokens.
         """
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         test_config = DynamicEngineTestConfig(
             num_requests=0,
             min_prompt_length=4,
@@ -2854,6 +2890,7 @@ def test_speculative_sequence_length_double_counting(
             use_fixed_output_lengths=False,
             context_max_tokens=512,
             position_embedding_type="none",
+            sampling_backend=sampling_backend,
         )
         env = self._build_test_env(test_config)
 
@@ -2950,15 +2987,18 @@ def deterministic_mtp(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
     )
     @torch.inference_mode()
+    @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"])
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     def test_speculative_decoding_with_eviction_and_swapping(
-        self, materialize_only_last_token_logits
+        self, materialize_only_last_token_logits, sampling_backend
     ):
         """Test that speculative decoding works correctly when requests are paused and evicted.
 
         This exercises the `_swap_book_keeping_tensors` logic with the 2D `new_speculative_tokens`
         tensor, ensuring no dimensional mismatch or index errors occur during tensor swapping.
         """
+        if sampling_backend == "flashinfer":
+            pytest.importorskip("flashinfer")
         # Very constrained memory environment to force pausing and eviction
         test_config = DynamicEngineTestConfig(
             num_requests=3,
@@ -2972,6 +3012,7 @@ def test_speculative_decoding_with_eviction_and_swapping(
             model_provider="gpt",
             materialize_only_last_token_logits=materialize_only_last_token_logits,
             use_fixed_output_lengths=True,
+            sampling_backend=sampling_backend,
         )
 
         env = self._build_test_env(test_config)

From 0c3666c049dfc73e0bfaf777603d44fc78b1638c Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 11:41:56 -0500
Subject: [PATCH 16/23] Fix torch sampling unit tests

---
 megatron/core/inference/engines/dynamic_engine.py        | 9 +++++----
 .../unit_tests/inference/engines/test_dynamic_engine.py  | 1 +
 .../inference/engines/test_hybrid_prefix_caching_e2e.py  | 1 +
 .../inference/test_mtp_cuda_graph_inference.py           | 2 ++
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index f00625d1708..fcee2c1daef 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -395,10 +395,11 @@ def create_cuda_graphs(self, reset_context: bool = True):
             with torch.inference_mode():
                 controller._dynamic_step_forward_logits(input_ids, position_ids)
 
-                if controller.num_speculative_tokens > 0:
-                    controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
-                else:
-                    controller._dynamic_step_sample_logits()
+                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:
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index 38d546fc878..7bcf21882c1 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -4672,6 +4672,7 @@ def _build_engine(self, model, enable_chunked_prefill, num_cuda_graphs, context_
             enable_chunked_prefill=enable_chunked_prefill,
             max_tokens=context_max_tokens,
             max_requests=128,
+            sampling_backend='torch',
         )
         if mamba_config is not None:
             inference_config_kwargs.update(mamba_inference_state_config=mamba_config)
diff --git a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py
index d3f1930561b..92890b22b53 100644
--- a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py
+++ b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py
@@ -200,6 +200,7 @@ def _build_engine(
             enable_prefix_caching=enable_prefix_caching,
             unified_memory_level=0,
             num_cuda_graphs=num_cuda_graphs,
+            sampling_backend='torch',
         )
         if enable_prefix_caching:
             inference_config_kwargs.update(
diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
index f74d1bbed13..8fd1f4a1154 100644
--- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
+++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py
@@ -153,6 +153,7 @@ def _build_engine(
                 block_size_tokens=256,
                 max_requests=max_requests,
                 num_cuda_graphs=-1,
+                sampling_backend='torch',
             ),
         )
         wrapped = GPTInferenceWrapper(model, context)
@@ -900,6 +901,7 @@ def _build_context(
                 num_cuda_graphs=num_cuda_graphs,
                 use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps,
                 max_requests=max_requests,
+                sampling_backend='torch',
             ),
         )
 

From d6affe7fd1a95457d26b52c73f5231d840690852 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 12:21:37 -0500
Subject: [PATCH 17/23] Fix topk & topp tests

---
 .../test_text_generation_controller.py                      | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index a4a06eb8897..f11f74c0504 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -356,13 +356,15 @@ def test_sample_from_dynamic_logits(
         sampled_l.masked_fill_(top_k_mask, 0.0)
         top_p_mask = sampled_l.cumsum(dim=-1) > top_p_values.unsqueeze(1)
 
+        # When `top_p` is enabled, but the cumulative probs don't actually filter anything,
+        # our constraint reduces to top_k alone.
+        start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long()
         first_excluded = torch.where(
             top_p_mask.any(dim=-1),
             top_p_mask.float().argmax(dim=-1),
-            torch.full((batch_size,), self.vocab_size, device=top_p_mask.device),
+            start_idx + 1,
         )
         last_included = torch.clamp(first_excluded - 1, min=0)
-        start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long()
         last_included = torch.max(last_included, start_idx)
         expected_min_values = l.gather(1, last_included.unsqueeze(1)).squeeze(1)
         assert torch.all(

From 583f1d2c8662d7a33482f609777df6aaaaa09f4b Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 12:30:58 -0500
Subject: [PATCH 18/23] lint, lint, lint away

gently down the stream
---
 .../test_text_generation_controller.py                        | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index f11f74c0504..b07041763a5 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -360,9 +360,7 @@ def test_sample_from_dynamic_logits(
         # our constraint reduces to top_k alone.
         start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long()
         first_excluded = torch.where(
-            top_p_mask.any(dim=-1),
-            top_p_mask.float().argmax(dim=-1),
-            start_idx + 1,
+            top_p_mask.any(dim=-1), top_p_mask.float().argmax(dim=-1), start_idx + 1
         )
         last_included = torch.clamp(first_excluded - 1, min=0)
         last_included = torch.max(last_included, start_idx)

From a10e72514117e836e3c4eb3f3171e41e0eeef114 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 30 Apr 2026 14:32:15 -0500
Subject: [PATCH 19/23] fail the user

---
 megatron/core/inference/config.py | 12 +++++-------
 megatron/training/arguments.py    | 13 ++++++-------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index a13c5686b99..2838e43f9ee 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -1,6 +1,5 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-import warnings
 from dataclasses import InitVar, dataclass
 from enum import Enum
 from typing import List, Literal, Optional, Tuple
@@ -328,9 +327,8 @@ def __post_init__(self, verbose: bool):
         if self.sampling_backend == 'flashinfer':
             try:
                 import flashinfer  # noqa: F401
-            except ImportError:
-                warnings.warn(
-                    "sampling_backend='flashinfer' requested but flashinfer is not "
-                    "installed; falling back to 'torch'."
-                )
-                self.sampling_backend = 'torch'
+            except ImportError as e:
+                raise ImportError(
+                    "sampling_backend='flashinfer' requires the flashinfer package; "
+                    "install it or set sampling_backend='torch'."
+                ) from e
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 7c144466943..14251f8df0a 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -9,7 +9,6 @@
 from pathlib import Path
 import re
 import types
-import warnings
 
 import torch
 import torch.nn.functional as F
@@ -1011,12 +1010,12 @@ def validate_args(args, defaults={}):
     if args.inference_dynamic_batching_sampling_backend == 'flashinfer':
         try:
             import flashinfer  # noqa: F401
-        except ImportError:
-            warnings.warn(
-                "--inference-dynamic-batching-sampling-backend=flashinfer requested "
-                "but flashinfer is not installed; falling back to 'torch'."
-            )
-            args.inference_dynamic_batching_sampling_backend = 'torch'
+        except ImportError as e:
+            raise ImportError(
+                "--inference-dynamic-batching-sampling-backend=flashinfer requires "
+                "the flashinfer package; install it or pass "
+                "--inference-dynamic-batching-sampling-backend=torch."
+            ) from e
 
     if args.use_megatron_fsdp:
         # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions.

From 92b50393df6a01fcae0f1e457807f405441473c3 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 1 May 2026 17:16:10 -0500
Subject: [PATCH 20/23] Clamp temperature in FlashInfer sampling

Co-Authored-By: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/inference/sampling/flashinfer_sampling.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index 7b4271274fc..d3720252c4c 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -77,6 +77,8 @@ def sample_kernel(
             top_k = md["top_k"][token_to_request_index]
             top_p = md["top_p"][token_to_request_index]
 
+        # Clamp temperature to avoid division by 0.
+        temperature = temperature.clamp(min=1e-6)
         if gather_indices is None:
             scaled = logits[:n] / temperature.unsqueeze(1)
         else:

From 00829e81b59ae17173ade9ca12c12c4561bb76ec Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 1 May 2026 17:32:31 -0500
Subject: [PATCH 21/23] Cover temp 0 in test_sample_from_dynamic_logits

Co-Authored-By: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../test_text_generation_controller.py                       | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index b07041763a5..5c96836e9ba 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -289,7 +289,7 @@ def test_sample_from_dynamic_logits(
     ):
         if backend == "flashinfer":
             pytest.importorskip("flashinfer")
-        batch_size = 12
+        batch_size = 15
         self.setup_model(
             torch.float32,
             batch_size=batch_size,
@@ -302,11 +302,14 @@ def test_sample_from_dynamic_logits(
         context = self.text_generation_controller.inference_wrapped_model.inference_context
 
         # Prepare sampling params in human-readable format, to aid with test maintenance.
+        # The temperature=0 / top_k=1 bucket exercises the greedy path: torch short-circuits
+        # to argmax, flashinfer relies on its temperature clamp to avoid divide-by-zero.
         sampling_test_cases: List[Tuple[SamplingParams, List[int]]] = [
             (SamplingParams(temperature=0.1, top_p=0.01), [9, 6, 10]),
             (SamplingParams(temperature=5.0, top_k=15), [0, 3, 2]),
             (SamplingParams(top_p=0.8), [4, 1, 7]),
             (SamplingParams(temperature=10.0, top_k=5), [11, 5, 8]),
+            (SamplingParams(temperature=0.0, top_k=1), [12, 13, 14]),
         ]
         # For non-torch backends, test simultaneous top_k and top_p sampling.
         if backend != "torch":

From de5c56de8d8e54db3785c5e8e85cdd94c612d03a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Sun, 3 May 2026 22:44:24 -0500
Subject: [PATCH 22/23] Resolve the rest of the rebase conflict

---
 .../inference/contexts/dynamic_context.py     | 81 +++++++++++++------
 megatron/core/inference/contexts/gpu_view.py  | 39 ++++++---
 .../inference/sampling/flashinfer_sampling.py | 18 +++--
 .../text_generation_controller.py             |  2 +-
 4 files changed, 98 insertions(+), 42 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index d8564a6e0fb..42b5739934c 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/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,19 +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
 
-        self.active_request_last_token_idxs = torch.empty_like(self.request_query_lengths)
+        # 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)
@@ -2244,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]`
@@ -2255,9 +2281,9 @@ 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
         ]
@@ -2265,6 +2291,15 @@ def transfer_bookkeeping_to_gpu(self) -> None:
         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.
diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py
index 17d6e19ea03..65c401163b0 100644
--- a/megatron/core/inference/contexts/gpu_view.py
+++ b/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).
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
index d3720252c4c..c89093daeac 100644
--- a/megatron/core/inference/sampling/flashinfer_sampling.py
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -66,16 +66,18 @@ def sample_kernel(
         # CudaGraphManager consumes these args, if it exists.
         del eager, cache_key
 
-        # Read GPU sampling parameters directly from `context.active_request_metadata`.
-        md = context.active_request_metadata
+        # Read GPU sampling parameters from the per-step gpu_view mirror. The
+        # CPU source-of-truth (`active_request_metadata`) is pinned but resident
+        # on CPU, so reading it here would mix devices with `logits`.
+        gv = context.gpu_view
         if token_to_request_index is None:
-            temperature = md["temperature"][:n]
-            top_k = md["top_k"][:n]
-            top_p = md["top_p"][:n]
+            temperature = gv.temperature[:n]
+            top_k = gv.top_k[:n]
+            top_p = gv.top_p[:n]
         else:
-            temperature = md["temperature"][token_to_request_index]
-            top_k = md["top_k"][token_to_request_index]
-            top_p = md["top_p"][token_to_request_index]
+            temperature = gv.temperature[token_to_request_index]
+            top_k = gv.top_k[token_to_request_index]
+            top_p = gv.top_p[token_to_request_index]
 
         # Clamp temperature to avoid division by 0.
         temperature = temperature.clamp(min=1e-6)
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index decc692c4f1..87edddea566 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1052,7 +1052,7 @@ def _dynamic_step_sample_logits(self):
         gather_indices = (
             None
             if context.config.materialize_only_last_token_logits
-            else context.active_request_last_token_idxs
+            else context.gpu_view.active_request_last_token_idxs
         )
         self._sampled_tokens_cuda = self._sampling.sample_kernel(
             self._all_logits_cuda.squeeze(0),

From e580bde4aef7aef7bdefccce9635857e604dd64f Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 4 May 2026 14:20:42 -0500
Subject: [PATCH 23/23] Revert the default to torch sampling

---
 megatron/core/inference/config.py             |   2 +-
 megatron/training/arguments.py                |   2 +-
 .../golden_values_dev_dgx_h100.json           | 286 ++++++++++++++++++
 .../model_config.yaml                         |  76 +++++
 .../recipes/h100/mamba-dynamic-inference.yaml |   5 +
 5 files changed, 369 insertions(+), 2 deletions(-)
 create mode 100644 tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json
 create mode 100644 tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index 2838e43f9ee..756f4fbe4be 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -297,7 +297,7 @@ class InferenceConfig:
     Defaults to 0, which means no logging.
     """
 
-    sampling_backend: Literal['torch', 'flashinfer'] = 'flashinfer'
+    sampling_backend: Literal['torch', 'flashinfer'] = 'torch'
     """Which sampling kernels to use during inference."""
 
     request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 14251f8df0a..e79334b7c03 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2005,7 +2005,7 @@ def _add_inference_args(parser):
                        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='flashinfer',
+                       type=str, default='torch',
                        choices=['torch', 'flashinfer'],
                        help='Which sampling kernels to use during inference. '
                             'Falls back to "torch" with a warning if "flashinfer" '
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..956754f44c1
--- /dev/null
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json
@@ -0,0 +1,286 @@
+{
+ "0": {
+  "input_prompt": "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies.",
+  "generated_text": " You are not alone. You are not alone. You are not alone. You are not alone. You are not alone. You are not alone.",
+  "generated_tokens": [
+   3213,
+   1584,
+   1605,
+   9412,
+   1046,
+   3213,
+   1584,
+   1605,
+   9412,
+   1046,
+   3213,
+   1584,
+   1605,
+   9412,
+   1046,
+   3213,
+   1584,
+   1605,
+   9412,
+   1046,
+   3213,
+   1584,
+   1605,
+   9412,
+   1046,
+   3213,
+   1584,
+   1605,
+   9412,
+   1046
+  ],
+  "latency": 1.878054141998291,
+  "ttft": 0.07786321640014648,
+  "cuda_graph_request_count_map": null,
+  "step_count": 30,
+  "top_n_logprobs": null,
+  "prompt_top_n_logprobs": null,
+  "prompt_logprobs": [
+   -9.498085021972656,
+   -3.787536859512329,
+   -3.0404648780822754,
+   -1.7445809841156006,
+   -0.29672086238861084,
+   -1.3661342859268188,
+   -2.3458175659179688,
+   -1.83931303024292,
+   -1.4894113540649414,
+   -6.440437316894531,
+   -0.8176816701889038,
+   -1.790361762046814,
+   -3.662419557571411,
+   -3.7036173343658447,
+   -1.6009434461593628,
+   -1.5501081943511963,
+   -2.846059799194336,
+   -6.732302665710449,
+   -0.06605878472328186,
+   -1.334327220916748,
+   -6.083745002746582,
+   -9.440131187438965,
+   -10.473882675170898,
+   -1.5964821577072144,
+   -4.702763557434082,
+   -0.7514524459838867,
+   -2.1461901664733887,
+   -0.012340382672846317,
+   -0.03605639934539795,
+   -3.0907557010650635,
+   -8.744739532470703,
+   -1.5410845279693604,
+   -5.84979772567749,
+   -3.0918972492218018,
+   -3.9814329147338867,
+   -3.78017520904541,
+   -2.5227086544036865,
+   -2.258594036102295,
+   -0.4719255566596985,
+   -1.0329649448394775,
+   -5.3284382820129395,
+   -8.25335693359375,
+   -0.015789249911904335,
+   -2.854100227355957,
+   -1.2236379384994507,
+   -3.905193328857422,
+   -0.9268187284469604,
+   -0.0030202509369701147,
+   -3.224249839782715,
+   -11.11172103881836,
+   -3.8121743202209473,
+   -2.3400487899780273,
+   -4.672845363616943,
+   -0.09729652851819992,
+   -0.06232408434152603,
+   -1.336004614830017,
+   -2.054157257080078,
+   -4.390933036804199,
+   -0.44248226284980774,
+   -3.9417736530303955,
+   -0.5888474583625793,
+   -0.26697415113449097,
+   -2.9271092414855957,
+   -13.515066146850586,
+   -0.10294333100318909,
+   -3.5007452964782715,
+   -0.8535972237586975,
+   -5.173652648925781,
+   -0.330394983291626,
+   -2.304553508758545,
+   -0.5418462753295898,
+   -1.300589919090271,
+   -4.9136152267456055,
+   -15.558022499084473,
+   -4.918652534484863,
+   -0.22206512093544006,
+   -6.589188575744629,
+   -0.9015690684318542,
+   -2.2228457927703857,
+   -1.8689247369766235,
+   -0.2006368339061737,
+   -5.918689727783203,
+   -0.006355076562613249,
+   -7.532094955444336,
+   -3.2708187103271484,
+   -3.743263006210327,
+   -2.011824131011963
+  ],
+  "generated_logprobs": [
+   -3.0331737995147705,
+   -1.9080564975738525,
+   -2.52506947517395,
+   -2.325258493423462,
+   -1.180279016494751,
+   -1.1824196577072144,
+   -0.39788734912872314,
+   -1.110222578048706,
+   -1.5034958124160767,
+   -0.9765141606330872,
+   -0.9300433397293091,
+   -0.15196305513381958,
+   -0.2200203537940979,
+   -0.06051275506615639,
+   -0.6840062737464905,
+   -1.0964292287826538,
+   -0.17654964327812195,
+   -0.18547140061855316,
+   -0.06710249185562134,
+   -0.4758152365684509,
+   -0.6657928228378296,
+   -0.10342729091644287,
+   -0.10059614479541779,
+   -0.046978313475847244,
+   -0.410809725522995,
+   -0.428723007440567,
+   -0.06053968518972397,
+   -0.06518109142780304,
+   -0.030038274824619293,
+   -0.3271780014038086
+  ],
+  "logprobs": [
+   -9.498085021972656,
+   -3.787536859512329,
+   -3.0404648780822754,
+   -1.7445809841156006,
+   -0.29672086238861084,
+   -1.3661342859268188,
+   -2.3458175659179688,
+   -1.83931303024292,
+   -1.4894113540649414,
+   -6.440437316894531,
+   -0.8176816701889038,
+   -1.790361762046814,
+   -3.662419557571411,
+   -3.7036173343658447,
+   -1.6009434461593628,
+   -1.5501081943511963,
+   -2.846059799194336,
+   -6.732302665710449,
+   -0.06605878472328186,
+   -1.334327220916748,
+   -6.083745002746582,
+   -9.440131187438965,
+   -10.473882675170898,
+   -1.5964821577072144,
+   -4.702763557434082,
+   -0.7514524459838867,
+   -2.1461901664733887,
+   -0.012340382672846317,
+   -0.03605639934539795,
+   -3.0907557010650635,
+   -8.744739532470703,
+   -1.5410845279693604,
+   -5.84979772567749,
+   -3.0918972492218018,
+   -3.9814329147338867,
+   -3.78017520904541,
+   -2.5227086544036865,
+   -2.258594036102295,
+   -0.4719255566596985,
+   -1.0329649448394775,
+   -5.3284382820129395,
+   -8.25335693359375,
+   -0.015789249911904335,
+   -2.854100227355957,
+   -1.2236379384994507,
+   -3.905193328857422,
+   -0.9268187284469604,
+   -0.0030202509369701147,
+   -3.224249839782715,
+   -11.11172103881836,
+   -3.8121743202209473,
+   -2.3400487899780273,
+   -4.672845363616943,
+   -0.09729652851819992,
+   -0.06232408434152603,
+   -1.336004614830017,
+   -2.054157257080078,
+   -4.390933036804199,
+   -0.44248226284980774,
+   -3.9417736530303955,
+   -0.5888474583625793,
+   -0.26697415113449097,
+   -2.9271092414855957,
+   -13.515066146850586,
+   -0.10294333100318909,
+   -3.5007452964782715,
+   -0.8535972237586975,
+   -5.173652648925781,
+   -0.330394983291626,
+   -2.304553508758545,
+   -0.5418462753295898,
+   -1.300589919090271,
+   -4.9136152267456055,
+   -15.558022499084473,
+   -4.918652534484863,
+   -0.22206512093544006,
+   -6.589188575744629,
+   -0.9015690684318542,
+   -2.2228457927703857,
+   -1.8689247369766235,
+   -0.2006368339061737,
+   -5.918689727783203,
+   -0.006355076562613249,
+   -7.532094955444336,
+   -3.2708187103271484,
+   -3.743263006210327,
+   -2.011824131011963,
+   -3.0331737995147705,
+   -1.9080564975738525,
+   -2.52506947517395,
+   -2.325258493423462,
+   -1.180279016494751,
+   -1.1824196577072144,
+   -0.39788734912872314,
+   -1.110222578048706,
+   -1.5034958124160767,
+   -0.9765141606330872,
+   -0.9300433397293091,
+   -0.15196305513381958,
+   -0.2200203537940979,
+   -0.06051275506615639,
+   -0.6840062737464905,
+   -1.0964292287826538,
+   -0.17654964327812195,
+   -0.18547140061855316,
+   -0.06710249185562134,
+   -0.4758152365684509,
+   -0.6657928228378296,
+   -0.10342729091644287,
+   -0.10059614479541779,
+   -0.046978313475847244,
+   -0.410809725522995,
+   -0.428723007440567,
+   -0.06053968518972397,
+   -0.06518109142780304,
+   -0.030038274824619293,
+   -0.3271780014038086
+  ]
+ },
+ "mem-max-allocated-bytes": 53350692864,
+ "lifetime_prefill_token_count": 88
+}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml
new file mode 100644
index 00000000000..e989be22f7e
--- /dev/null
+++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml
@@ -0,0 +1,76 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+  TRITON_CACHE_AUTOTUNING: 0
+  MAMBA_DETERMINISTIC: 1
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --log-memory-to-tensorboard: true
+  --timing-log-level: 0
+  --load: ${CHECKPOINT_LOAD_PATH}/model/mamba_hybrid_2b/dcp/mcore-v1_bf16/checkpoint
+  --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/mamba_hybrid_2b/dcp/mcore-v1_bf16/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
+  --tokenizer-type: TikTokenizer
+  --tiktoken-pattern: v2
+  --distributed-backend: nccl
+  --log-interval: 1
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --expert-model-parallel-size: 1
+  --use-mcore-models: true
+  --model-provider: hybrid
+  --init-method-std: 0.0198
+  --untie-embeddings-and-output-weights: true
+  --disable-bias-linear: true
+  --init-method-std: 0.014
+  --position-embedding-type: none
+  --hidden-size: 2048
+  --ffn-hidden-size: 11264
+  --num-attention-heads: 16
+  --kv-channels: 128
+  --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-
+  --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec
+  --normalization: RMSNorm
+  --swiglu: true
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --seq-length: 4096
+  --max-position-embeddings: 4096
+  --micro-batch-size: 1
+  --ckpt-format: torch_dist
+  --ckpt-fully-parallel-save: true
+  --ckpt-fully-parallel-load: true
+  --ckpt-assume-constant-structure: true
+  --dist-ckpt-strictness: log_unexpected
+  --bf16: true
+  --attention-backend: flash
+  --no-create-attention-mask-in-dataloader: true
+  --num-workers: 8
+  --use-checkpoint-args: true
+  --no-use-tokenizer-model-from-checkpoint-args: true
+  --no-load-optim: true
+  --deterministic-mode: true
+  --save-interval: 2000
+  --temperature: 1.0
+  --top_k: 1
+  --inference-dynamic-batching-sampling-backend: flashinfer
+  --return-log-probs: true
+  --num-tokens-to-generate: 30
+  --max-tokens-to-oom: 3600000
+  --inference-max-seq-length: 4096
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
+  --incoming-requests-per-step: 32 
+  --inference-repeat-n: 3
+  --no-record-throughput: true
+  --mamba-inference-conv-states-dtype: fp32
+  --mamba-inference-ssm-states-dtype: fp32
+METRICS:
+  - "generated_tokens"
+  - "logprobs"
diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml
index 495a4d130b0..aa78cdb8316 100644
--- a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml
+++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml
@@ -65,3 +65,8 @@ products:
       - environment: [dev]
         scope: [mr, mr-github]
         platforms: [dgx_h100]
+  - test_case: [hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer]
+    products:
+      - environment: [dev]
+        scope: [mr, mr-github]
+        platforms: [dgx_h100]