[pull] main from NVIDIA:main

docs/source/features/sampling.md

@@ -75,7 +75,8 @@ llm.generate(["Hello, my name is",
 
 * The sampling is controlled via `SamplingParams`.
 
-* By default (`temperature = top_p = top_k = None`), greedy sampling is used.
+* By default (`temperature = top_p = top_k = None`), greedy sampling is used
+  (unless top-p decay is active, see below).
 
 * If either `temperature = 0`, `top_p = 0`, and/or `top_k = 1`, is specified, sampling is greedy,
   irrespective of the values of the remaining parameters.
@@ -101,6 +102,19 @@ llm.generate(["Hello, my name is",
 
   * The implementation does not guarantee any particular treatment of tied probabilities.
 
+* Top-P decay is supported: if `top_p_decay < 1` is specified, the effective `top_p` is
+  multiplied by `top_p_decay` after every sampled token, bounded from below by `top_p_min`
+  (default `1e-6`), and reset to the initial `top_p` whenever the token `top_p_reset_ids`
+  is sampled (default `-1`, which never matches a token). Out-of-range values
+  (`top_p_decay` or `top_p_min` outside `(0, 1]`, negative `top_p_reset_ids`) are rejected.
+
+  * An active top-p decay implies top-p sampling even if `top_p` is unspecified or `top_p = 1`
+    (the initial `top_p` then defaults to 1). However, explicitly requested greedy sampling
+    (`temperature = 0`, `top_p = 0`, and/or `top_k = 1`) takes precedence over top-p decay.
+
+  * Top-P decay is not supported in combination with beam search or with speculative decoding
+    modes that route draft tokens through the Torch Sampler; such requests are rejected.
+
 ### Performance
 
 The Torch Sampler leverages the optimized sampling kernels provided by

tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from torch import Tensor, nn
+
+from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
+from tensorrt_llm._torch.models.modeling_utils import register_mapper
+
+MINIMAX_M3_PARAMS_MAP = {
+    r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": r"\1.gate.e_score_correction_bias",
+}
+
+
+@register_mapper("HF", "MiniMaxM3SparseForCausalLM")
+@register_mapper("HF", "MiniMaxM3SparseForConditionalGeneration")
+class MiniMaxM3HfWeightMapper(HfWeightMapper):
+    """Handle M3 gate naming and MXFP8 GQA duplication for loader v2."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.params_map = MINIMAX_M3_PARAMS_MAP
+
+    def _duplicate_kv_weights(
+        self, module: nn.Module, new_name: str, weights: dict[str, Tensor]
+    ) -> dict[str, Tensor]:
+        if new_name not in ["k_proj", "v_proj"]:
+            return weights
+
+        duplicated_keys = ["weight", "bias"]
+        quant_config = getattr(module, "quant_config", None)
+        if quant_config is not None:
+            quant_mode = quant_config.quant_mode
+            if quant_mode.has_nvfp4():
+                duplicated_keys.append("weight_scale")
+            if quant_mode.has_mxfp8():
+                duplicated_keys.extend(["weight_scale", "weight_scale_inv"])
+
+        return {
+            key: self._duplicate_kv(
+                weight=value[:], num_kv_heads=self._num_kv_heads, tensor_parallel_size=self._tp_size
+            )
+            if key in duplicated_keys
+            else value
+            for key, value in weights.items()
+        }

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -62,6 +62,8 @@
     get_model_extra_attrs,
     is_torch_compiling,
 )
+from .checkpoints.base_weight_mapper import BaseWeightMapper
+from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper
 from .modeling_utils import DecoderModel, DecoderModelForCausalLM, ModelConfig, register_auto_model
 
 # Dense layers use SDPA with non-contiguous Q/K/V and a bool attn_mask.
@@ -1471,20 +1473,6 @@ def forward(
         return hidden_states
 
 
-# HF MiniMax-M3 stores the routed score-correction bias one level above
-# the router weight (``block_sparse_moe.e_score_correction_bias``,
-# sibling of ``block_sparse_moe.gate.weight``). The TRT-LLM module tree
-# binds it to :class:`MiniMaxM3Gate`, so the generic loader expects to
-# see it at ``block_sparse_moe.gate.e_score_correction_bias``. The
-# regex below moves the key into the gate's prefix before the loader
-# dispatches; this lets ``mark_consumed("...gate")`` cleanly remove
-# both the weight and the bias together without disturbing the sibling
-# ``block_sparse_moe.experts.*`` backend subtree.
-_M3_GATE_BIAS_RENAME_MAP = {
-    r"^(.*\.block_sparse_moe)\.e_score_correction_bias$": (r"\1.gate.e_score_correction_bias"),
-}
-
-
 @register_auto_model("MiniMaxM3SparseForCausalLM")
 class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]):
     """Text-only M3 model."""
@@ -1500,13 +1488,23 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"):
             vocab_size=model_config.pretrained_config.vocab_size,
         )
 
-    def load_weights(self, weights, *args, **kwargs):
-        # Merge the M3-specific gate-bias rename into any caller-
-        # supplied ``params_map`` so the VL wrapper and any downstream
-        # tooling that already passes one keep working.
-        params_map = kwargs.pop("params_map", None) or {}
-        merged = {**_M3_GATE_BIAS_RENAME_MAP, **params_map}
-        return super().load_weights(weights, *args, params_map=merged, **kwargs)
+    def load_weights(
+        self,
+        weights: Dict,
+        weight_mapper: Optional[BaseWeightMapper] = None,
+        params_map: Optional[Dict[str, str]] = None,
+        allow_partial_loading: bool = False,
+    ) -> None:
+        if weight_mapper is None:
+            weight_mapper = MiniMaxM3HfWeightMapper()
+        weight_mapper.init_model_and_config(self, self.model_config)
+        merged_params_map = {**MINIMAX_M3_PARAMS_MAP, **(params_map or {})}
+        super().load_weights(
+            weights=weights,
+            weight_mapper=weight_mapper,
+            params_map=merged_params_map,
+            allow_partial_loading=allow_partial_loading,
+        )
 
 
 def _strip_language_model_prefix(
@@ -1642,7 +1640,13 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"):
         self.last_loaded_vision_keys = []
         self.last_missing_vision_keys = []
 
-    def load_weights(self, weights, *args, **kwargs):
+    def load_weights(
+        self,
+        weights: Dict,
+        weight_mapper: Optional[BaseWeightMapper] = None,
+        params_map: Optional[Dict[str, str]] = None,
+        allow_partial_loading: bool = False,
+    ) -> None:
         text_cfg = self.config
         if is_minimax_m3_vl_config(text_cfg):
             text_cfg = get_text_config(text_cfg)
@@ -1665,7 +1669,12 @@ def load_weights(self, weights, *args, **kwargs):
             self.last_loaded_vision_keys = loaded
             self.last_missing_vision_keys = missing
 
-        return super().load_weights(text_weights, *args, **kwargs)
+        super().load_weights(
+            weights=text_weights,
+            weight_mapper=weight_mapper,
+            params_map=params_map,
+            allow_partial_loading=allow_partial_loading,
+        )
 
     def forward(
         self,

tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py

@@ -2071,9 +2071,14 @@ def get_minimax_m3_vl_input_processor_cls():
     path. Re-type by dynamic subclassing here so the registered class
     inherits from the base.
     """
-    from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor
+    from tensorrt_llm.inputs.registry import (
+        BaseMultimodalDummyInputsBuilder,
+        BaseMultimodalInputProcessor,
+    )
 
-    class _Registered(MiniMaxM3VLInputProcessor, BaseMultimodalInputProcessor):
+    class _Registered(
+        MiniMaxM3VLInputProcessor, BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder
+    ):
         pass
 
     _Registered.__name__ = "MiniMaxM3VLInputProcessor"

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -914,6 +914,7 @@ def on_detected():
         # under steady state — see ping-pong comment in _profiler).
         self._latest_host_step_time_ms: Optional[float] = None
         self._latest_prev_device_step_time_ms: Optional[float] = None
+        self._emit_initial_stats()
         self.gather_all_responses = False
 
         self.kv_cache_transceiver = kv_cache_transceiver
@@ -3941,6 +3942,25 @@ def _handle_disagg_cache_errors_synced(self):
             charge_budget=False,
         )
 
+    def _emit_initial_stats(self) -> None:
+        """Emit a startup stats snapshot so that cache_config_info is
+        immediately available to external metric scrapers (e.g. the
+        Kubernetes Inference Gateway EPP) before any inference request."""
+        if not self.enable_iter_perf_stats:
+            return
+        stats = self._get_init_iter_stats(0, 0)
+        kv_cache_manager = self.resource_manager.resource_managers.get(
+            ResourceManagerType.KV_CACHE_MANAGER)
+        if kv_cache_manager is not None:
+            kv_stats = kv_cache_manager.get_kv_cache_stats()
+            kv_stats_to_save = KvCacheStats()
+            kv_stats_to_save.max_num_blocks = kv_stats.max_num_blocks
+            kv_stats_to_save.tokens_per_block = kv_stats.tokens_per_block
+            kv_stats_to_save.free_num_blocks = kv_stats.free_num_blocks
+            kv_stats_to_save.used_num_blocks = kv_stats.used_num_blocks
+            stats.kv_cache_stats = kv_stats_to_save
+        self._append_iter_stats(stats)
+
     def _executor_loop(self):
         torch.cuda.set_device(self.device_id)
         # ensure the context is created, otherwise, some MPI calls will fail.

tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py

@@ -373,3 +373,105 @@ def gather_log_softmax(inputs_cuda: torch.Tensor, indices_cuda: torch.Tensor) ->
         torch._dynamo.mark_dynamic(inputs_cuda, 0)
         torch._dynamo.mark_dynamic(indices_cuda, 0)
         return Fusions._gather_log_softmax_impl(inputs_cuda, indices_cuda)
+
+    # --- Top-P Decay ops ---------------------------------------------------
+    # Host-launch-bound per-step ops (a few dozen elements per row), fused with
+    # Inductor to keep the launch count low. mode="max-autotune-no-cudagraphs":
+    # cudagraphs is unsafe here (the update mutates persistent per-slot state
+    # in place and the gather's output is consumed outside the compiled region;
+    # cudagraph static output buffers get overwritten by subsequent replays).
+    # mark_dynamic on the batch-varying dims avoids recompilation as the batch
+    # composition changes. Compilation is lazy: the first decay-active request
+    # pays it (roughly a second); non-decay workloads never trigger it. See
+    # TorchSampler.TopPDecayStore for the feature-level semantics.
+
+    @staticmethod
+    @torch.compile(mode="max-autotune-no-cudagraphs")
+    def _top_p_decay_update_impl(
+        runtime_top_p: torch.Tensor,
+        initial_top_p: torch.Tensor,
+        top_p_decay: torch.Tensor,
+        top_p_min: torch.Tensor,
+        reset_ids: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        step_tokens: torch.Tensor,
+        sampled_slots: torch.Tensor,
+    ) -> None:
+        active = is_decay_slot[sampled_slots]
+        current = runtime_top_p[sampled_slots]
+        updated = torch.where(
+            step_tokens[sampled_slots] == reset_ids[sampled_slots],
+            initial_top_p[sampled_slots],
+            torch.maximum(current * top_p_decay[sampled_slots], top_p_min[sampled_slots]),
+        )
+        runtime_top_p[sampled_slots] = torch.where(active, updated, current)
+
+    @staticmethod
+    def top_p_decay_update(
+        *,
+        runtime_top_p: torch.Tensor,
+        initial_top_p: torch.Tensor,
+        top_p_decay: torch.Tensor,
+        top_p_min: torch.Tensor,
+        reset_ids: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        step_tokens: torch.Tensor,
+        sampled_slots: torch.Tensor,
+    ) -> None:
+        """Fused in-place update of ``runtime_top_p`` for the sampled decay slots.
+
+        Applies the Top-P Decay recurrence (see ``TorchSampler.TopPDecayStore``
+        for the feature-level semantics) to every sampled row whose slot is
+        decay-active per ``is_decay_slot``.
+
+        All per-slot tensors are 1-D of length ``max_num_sequences``;
+        ``step_tokens`` is a slot-indexed 1-D (possibly strided) int32 view of
+        the new-tokens buffer for a fixed step/beam
+        (``new_tokens[step, :, beam]``); ``sampled_slots`` is 1-D of length
+        ``num_sampled`` (this iteration's rows). ``runtime_top_p`` is mutated
+        in place; nothing is returned.
+        """
+        torch._dynamo.mark_dynamic(sampled_slots, 0)
+        Fusions._top_p_decay_update_impl(
+            runtime_top_p,
+            initial_top_p,
+            top_p_decay,
+            top_p_min,
+            reset_ids,
+            is_decay_slot,
+            step_tokens,
+            sampled_slots,
+        )
+
+    @staticmethod
+    @torch.compile(mode="max-autotune-no-cudagraphs")
+    def _top_p_decay_gather_impl(
+        runtime_top_p: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        static_top_p: torch.Tensor,
+        slots: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.where(is_decay_slot[slots], runtime_top_p[slots], static_top_p)
+
+    @staticmethod
+    def top_p_decay_gather(
+        *,
+        runtime_top_p: torch.Tensor,
+        is_decay_slot: torch.Tensor,
+        static_top_p: torch.Tensor,
+        slots: torch.Tensor,
+    ) -> torch.Tensor:
+        """Fused pre-sample per-row top-p gather for decay-active rows.
+
+        Returns a new per-row tensor::
+
+            row_top_p[i] = runtime_top_p[slots[i]]  if is_decay_slot[slots[i]]
+                         = static_top_p[i]          otherwise
+
+        ``runtime_top_p`` / ``is_decay_slot`` are per-slot arrays;
+        ``static_top_p`` and ``slots`` are per-row (length = the group's
+        per-step row count).
+        """
+        torch._dynamo.mark_dynamic(slots, 0)
+        torch._dynamo.mark_dynamic(static_top_p, 0)
+        return Fusions._top_p_decay_gather_impl(runtime_top_p, is_decay_slot, static_top_p, slots)