Add logprobs_mode (raw/processed) to inference config

megatron/core/inference/config.py

@@ -344,6 +344,9 @@ class InferenceConfig:
     sampling_backend: Literal['torch', 'flashinfer'] = 'torch'
     """Which sampling kernels to use during inference."""
 
+    logprobs_mode: Literal['raw_logprobs', 'processed_logprobs'] = 'raw_logprobs'
+    """Whether returned log-probs are modified by the sampling parameters or not."""
+
     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
@@ -387,6 +390,19 @@ def __post_init__(self, verbose: bool):
                 f"got {self.prefix_caching_routing_alpha}"
             )
 
+        if self.logprobs_mode not in ("raw_logprobs", "processed_logprobs"):
+            raise ValueError(
+                f"Unsupported logprobs_mode {self.logprobs_mode!r}. "
+                "Supported modes: raw_logprobs, processed_logprobs."
+            )
+
+        # The speculative log-probs path does not yet apply processed-logprobs.
+        if self.logprobs_mode == "processed_logprobs" and self.num_speculative_tokens > 0:
+            raise ValueError(
+                "logprobs_mode='processed_logprobs' is not yet supported with speculative decoding "
+                "(num_speculative_tokens > 0)."
+            )
+
         if self.sampling_backend == 'flashinfer':
             try:
                 import flashinfer  # noqa: F401

megatron/core/inference/contexts/dynamic_context.py

@@ -22,6 +22,7 @@
     PrefixCachingEvictionPolicy,
 )
 from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling.base import Sampling
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.unified_memory import (
     UnifiedMemoryUnsupportedError,
@@ -3772,8 +3773,38 @@ def update_requests(
             "evict_request_ids": evict_request_ids,
         }
 
+    def _processed_log_probs(
+        self,
+        logits: Tensor,
+        n_active: int,
+        active_query_lengths: Optional[Tensor],
+        sampling: Optional[Sampling],
+    ) -> Tensor:
+        """Sample the logprobs if desired."""
+        if self.config.logprobs_mode == "raw_logprobs":
+            return F.log_softmax(logits, dim=-1)
+
+        assert sampling is not None, "processed_logprobs requires a sampling backend"
+
+        # Map each logits row to its active request.
+        request_idx = torch.arange(n_active, device=logits.device)
+        row_to_request = (
+            request_idx
+            if active_query_lengths is None
+            else request_idx.repeat_interleave(active_query_lengths)
+        )
+        md = self.active_request_metadata
+        temperature = md["temperature"][:n_active].to(logits.device, torch.float32)[row_to_request]
+        top_k = md["top_k"][:n_active].to(logits.device, torch.long)[row_to_request]
+        top_p = md["top_p"][:n_active].to(logits.device, torch.float32)[row_to_request]
+        return sampling.log_probs_kernel(logits, temperature, top_k, top_p)
+
     def calculate_log_probs(
-        self, logits: Tensor, new_tokens: Tensor, only_last_token_logits: Optional[bool] = False
+        self,
+        logits: Tensor,
+        new_tokens: Tensor,
+        only_last_token_logits: Optional[bool] = False,
+        sampling: Optional[Sampling] = None,
     ) -> Tuple[List[List[float]], Tensor]:
         """Calculate log probs for all active requests and return them.
 
@@ -3783,6 +3814,7 @@ def calculate_log_probs(
             logits (Tensor): Raw model output logits with shape [1, sequence_length, vocab_size].
             new_tokens (Tensor): The newly sampled tokens.
             only_last_token_logits (bool): If set, the logits are from only the last token in each request
+            sampling (Optional[Sampling]): Backend used to optionally modify log-probs.
 
         Returns:
             List of lists where each inner list contains log probs for a request in the
@@ -3792,14 +3824,16 @@ def calculate_log_probs(
 
         # Calculate log_probs (sequence_length x vocab_size)
         logits_squeezed = logits.squeeze(0).float()
+        n_active = self.total_request_count - self.paused_request_count
 
         if only_last_token_logits or self.is_decode_only():
             seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device)
-            log_probs = F.log_softmax(logits_squeezed[seq_idx], dim=-1)
+            log_probs = self._processed_log_probs(
+                logits_squeezed[seq_idx], n_active, None, sampling
+            )
             selected_log_probs = log_probs[seq_idx, new_tokens]
             return [[lp] for lp in selected_log_probs.tolist()], log_probs
 
-        log_probs = F.log_softmax(logits_squeezed, dim=-1)
         # Get the selected token ids for all tokens.
         # We shift the active token window left by one to remove the first prompt token for
         # prefill requests and then set the token ids explicitly for the newly generated tokens.
@@ -3823,13 +3857,17 @@ def calculate_log_probs(
         #
         #   active_token_ids[new_token_idx] = new_tokens
         #                       : [ 52 | 12 | 16  3 | 12 72 24 88 86 ]
-        n_active = self.total_request_count - self.paused_request_count
         active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0)
         active_query_lengths = self.gpu_view.request_query_lengths[:n_active]
 
         new_token_idx = active_query_lengths.cumsum(0) - 1
         active_token_ids[new_token_idx] = new_tokens
 
+        # Compute (possibly processed) log-probs over all active-token rows.
+        log_probs = self._processed_log_probs(
+            logits_squeezed, n_active, active_query_lengths, sampling
+        )
+
         # Extract the log probs for only the selected tokens.
         # (sequence_length x vocab_size) -> (sequence_length)
         seq_idx = torch.arange(self.active_token_count, device=log_probs.device)

megatron/core/inference/sampling/base.py

@@ -10,7 +10,8 @@
 class Sampling(ABC):
     """Abstract base for inference sampling backends.
 
-    Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`.
+    Subclasses implement `sample_kernel` and `log_probs_kernel`.
+    CUDA graphs are added via `CudaGraphManager`.
     """
 
     @abstractmethod
@@ -87,3 +88,18 @@ def sample_speculative(
             token_to_request_index=token_to_request_index,
             eager=True,
         )
+
+    @abstractmethod
+    def log_probs_kernel(
+        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+    ) -> Tensor:
+        """Per-row log-probs of the distribution this backend samples from.
+
+        Args:
+            logits: `[num_rows, vocab_size]` raw logits.
+            temperature, top_k, top_p: `[num_rows]` per-row sampling params.
+
+        Returns:
+            `[num_rows, vocab_size]` log-probs; filtered-out tokens are `-inf`.
+        """
+        ...

megatron/core/inference/sampling/flashinfer_sampling.py

@@ -99,3 +99,20 @@ def sample_kernel(
             )
         )
         return output
+
+    def log_probs_kernel(
+        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+    ) -> Tensor:
+        """Per-row log-probs of the FlashInfer top-k / top-p sampling distribution."""
+        temperature = temperature.clamp(min=1e-6)
+        probs = torch.softmax(logits / temperature.unsqueeze(1), 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)
+
+        # Renormalize to the kept set (top-k first, then top-p) to match
+        renormed = flashinfer.sampling.top_k_renorm_probs(probs, top_k_safe)
+        renormed = flashinfer.sampling.top_p_renorm_probs(renormed, top_p_safe)
+        return torch.log(renormed)

megatron/core/inference/sampling/torch_sampling.py

@@ -19,6 +19,56 @@ def __init__(self, rng: torch.Generator, vocab_size: int) -> None:
         self._rng = rng
         self._vocab_size = vocab_size
 
+    @staticmethod
+    def _modify_logits_for_top_k_filtering(logits: Tensor, top_k: int) -> None:
+        """In-place: set logits outside the top-k set to -inf."""
+        filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits.masked_fill_(filter_, float("-Inf"))
+
+    @staticmethod
+    def _modify_logits_for_top_p_filtering(logits: Tensor, top_p: float) -> None:
+        """In-place: set logits outside the top-p (nucleus) set 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"))
+
+    @staticmethod
+    def filter_logits(
+        last_token_logits: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        *,
+        vocab_size: Optional[int] = None,
+    ) -> Tensor:
+        """Temperature-scale then top-k/top-p filter logits; filtered entries become -inf.
+
+        Returns a new tensor (input unmodified). Shared by `sample_from_logits` and
+        `log_probs_kernel` so sampling and processed log-probs apply the same filter.
+        """
+        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]"
+        # Clone needed: .div_() and the filters below modify in-place.
+        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."
+            TorchSampling._modify_logits_for_top_k_filtering(last_token_logits, top_k)
+        elif top_p > 0.0:
+            TorchSampling._modify_logits_for_top_p_filtering(last_token_logits, top_p)
+        return last_token_logits
+
     @staticmethod
     def sample_from_logits(
         last_token_logits: Tensor,
@@ -49,49 +99,45 @@ def sample_from_logits(
         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.
-        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)
+        filtered = TorchSampling.filter_logits(
+            last_token_logits, temperature, top_k, top_p, vocab_size=vocab_size
+        )
+        probabilities = filtered.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 log_probs_kernel(
+        self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor
+    ) -> Tensor:
+        """Per-row log-probs of the temperature, top-k/top-p sampling distribution.
+
+        Buckets rows by identical (temperature, top_k, top_p) and reuses `filter_logits`
+        (the same filter as `sample_from_logits`) so log-probs match how this backend
+        samples. `temperature`/`top_k`/`top_p` are per-row `[num_rows]` tensors.
+        """
+        temps = temperature.tolist()
+        top_ks = top_k.tolist()
+        top_ps = top_p.tolist()
+        buckets: dict = defaultdict(list)
+        for row, key in enumerate(zip(temps, top_ks, top_ps)):
+            buckets[key].append(row)
+
+        log_probs = torch.empty_like(logits)
+        for (t, k, p), rows in buckets.items():
+            idx = torch.tensor(rows, device=logits.device, dtype=torch.long)
+            filtered = TorchSampling.filter_logits(
+                logits[idx], float(t), int(k), float(p), vocab_size=self._vocab_size
+            )
+            log_probs[idx] = torch.log_softmax(filtered, dim=-1)
+        return log_probs
+
     def sample_kernel(
         self,
         logits: Tensor,

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -1157,6 +1157,7 @@ def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]:
             self._all_logits_cuda[:, :logits_seq_len, :],
             self._sampled_tokens_cuda[:active_request_count],
             only_last_token_logits=context.config.materialize_only_last_token_logits,
+            sampling=self._sampling,
         )
 
     def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float]], Tensor]:

megatron/inference/utils.py

@@ -382,6 +382,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args):
         use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives,
         disable_ep_consensus=args.inference_disable_ep_consensus,
         sampling_backend=args.inference_dynamic_batching_sampling_backend,
+        logprobs_mode=args.inference_dynamic_batching_logprobs_mode,
     )
 
 

megatron/training/arguments.py

@@ -2004,6 +2004,12 @@ def _add_inference_args(parser):
                        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-dynamic-batching-logprobs-mode',
+                       type=str, default='raw_logprobs',
+                       choices=['raw_logprobs', 'processed_logprobs'],
+                       help='How returned inference log-probs are computed engine-wide. '
+                            '"raw_logprobs" (default) uses the unmodified model logits; '
+                            '"processed_logprobs" uses temperature and filters by top-k/top-p.')
     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.')