From 8341fb1c2abbc49c2cf89f1b6ea659375ebb1820 Mon Sep 17 00:00:00 2001 From: BenjaminBraunDev Date: Wed, 22 Jul 2026 09:22:28 -0700 Subject: [PATCH 1/5] [#12595][feat] Emit initial KV cache stats at startup for external metric scrapers (#12596) Signed-off-by: BenjaminBraunDev --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 20 +++++++++++++ tensorrt_llm/serve/openai_server.py | 4 +++ .../llmapi/apps/_test_openai_metrics.py | 29 +++++++++++++++++++ .../llmapi/apps/_test_openai_prometheus.py | 25 ++++++++++++++++ tests/unittest/llmapi/test_llm_pytorch.py | 29 ++++++++++++++----- 5 files changed, 99 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4f4fc081a1d1..9253b2f41fbd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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. diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 71664fffb7f0..153b35b5f6cb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -401,6 +401,10 @@ async def lifespan(app: FastAPI): self._iteration_stats_buffer = deque(maxlen=max_buf) self._iteration_stats_collector_task = asyncio.create_task( self._iteration_stats_collector_loop()) + # Wake up the collector immediately so it processes the + # initial stats emitted by the executor at startup (e.g. + # cache_config_info). + self._iteration_stats_wakeup_event.set() logger.info( "Started background iteration stats collector task") diff --git a/tests/unittest/llmapi/apps/_test_openai_metrics.py b/tests/unittest/llmapi/apps/_test_openai_metrics.py index 12ce84f63646..957297f5a5ca 100644 --- a/tests/unittest/llmapi/apps/_test_openai_metrics.py +++ b/tests/unittest/llmapi/apps/_test_openai_metrics.py @@ -1,5 +1,6 @@ """Test the metrics endpoint when using OpenAI API to send requests""" +import time from unittest.mock import patch import pytest @@ -51,6 +52,34 @@ def test_version(client): assert response.status_code == 200 +def test_metrics_available_before_first_request(client): + """Verify that KV cache config stats are available at startup, + before any inference request is sent. This is critical for external + metric scrapers (e.g. the Kubernetes Inference Gateway EPP) that need + cache_config_info immediately to make routing decisions.""" + # Poll until the background stats collector processes the initial stats + deadline = time.time() + 5.0 + stats = [] + while time.time() < deadline: + response = client.get("/metrics") + assert response.status_code == 200 + stats = response.json() + if stats: + break + time.sleep(0.1) + assert stats, "Expected initial stats before first request" + response_dict = stats[0] + assert "kvCacheStats" in response_dict, \ + "kvCacheStats should be present before first request" + kv_stats = response_dict["kvCacheStats"] + assert "maxNumBlocks" in kv_stats + assert "tokensPerBlock" in kv_stats + assert kv_stats["maxNumBlocks"] > 0, \ + "maxNumBlocks should be positive at startup" + assert kv_stats["tokensPerBlock"] > 0, \ + "tokensPerBlock should be positive at startup" + + def test_metrics(client): response = client.post("/v1/completions", json={ diff --git a/tests/unittest/llmapi/apps/_test_openai_prometheus.py b/tests/unittest/llmapi/apps/_test_openai_prometheus.py index 48abe35263aa..a113274ad0c2 100644 --- a/tests/unittest/llmapi/apps/_test_openai_prometheus.py +++ b/tests/unittest/llmapi/apps/_test_openai_prometheus.py @@ -127,6 +127,31 @@ def _parse_all_kv_metrics(data: str, prefix: str) -> Dict[str, float | None]: return {name: _parse_prometheus_sample(data, name) for name in names} +def test_kv_cache_metrics_available_before_first_request( + server: RemoteOpenAIServer): + """Verify that KV cache metrics are available at startup, before any + inference request. External scrapers (e.g. the Kubernetes Inference + Gateway EPP) rely on these metrics for routing decisions.""" + metric_prefix = "trtllm_" + max_wait_time = 10.0 + poll_interval = 0.5 + start_time = time.time() + metrics_found = False + + while time.time() - start_time < max_wait_time: + response = urlopen(f'{server.url_root}/prometheus/metrics') + assert response.status == 200 + data = response.read().decode("utf-8") + if metric_prefix + "kv_cache_utilization" in data: + metrics_found = True + break + time.sleep(poll_interval) + + assert metrics_found, \ + (f"{metric_prefix}kv_cache_utilization not found in /prometheus/metrics " + f"after {max_wait_time}s — it should be available before any request") + + def test_metrics_endpoint(server: RemoteOpenAIServer): """Verify that Prometheus metrics are correctly exposed after serving requests. diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index f1fb240e1395..eaadc09df8af 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1546,18 +1546,22 @@ def test_llm_context_only_timed_out(transceiver_runtime): disaggregated_params=disaggregated_params): print(output) + # Wait until the context-only request has allocated KV cache blocks max_retries = 10 + all_results = [] for _ in range(max_retries): results = llm.get_stats(2) - if len(results) == 1: + all_results.extend(results) + if all_results and all_results[-1]["kvCacheStats"]["usedNumBlocks"] > 0: break time.sleep(1) else: pytest.fail( - f"Failed to get stats with len==1 after {max_retries} retries") + f"Context-only KV cache blocks not allocated after {max_retries} retries" + ) + results = all_results - assert len(results) == 1 - context_only_used_num_blocks = results[0]["kvCacheStats"]["usedNumBlocks"] + context_only_used_num_blocks = results[-1]["kvCacheStats"]["usedNumBlocks"] print(f"Context only used num blocks: {context_only_used_num_blocks}") # Sleep 5 seconds to allow context only request to time out @@ -1567,11 +1571,20 @@ def test_llm_context_only_timed_out(transceiver_runtime): for output in llm.generate(prompts0, sampling_params=sampling_params): print(output) - # Get number of allocated blocks - results = llm.get_stats(2) - assert len(results) == 1 - final_used_num_blocks = results[0]["kvCacheStats"]["usedNumBlocks"] + # Wait until KV cache blocks are released (usedNumBlocks == 0) + max_retries = 10 + all_results = [] + for _ in range(max_retries): + results = llm.get_stats(2) + all_results.extend(results) + if all_results and all_results[-1]["kvCacheStats"]["usedNumBlocks"] == 0: + break + time.sleep(1) + else: + pytest.fail(f"KV cache blocks not released after {max_retries} retries") + results = all_results + final_used_num_blocks = results[-1]["kvCacheStats"]["usedNumBlocks"] assert final_used_num_blocks == 0 From 9e96f8b9e72baef11145105d20e350436fe16aa8 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:36:11 -0700 Subject: [PATCH 2/5] [TRTLLM-14255][fix] migrate MiniMax M3 to loader v2 for TP8 support (#16468) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../checkpoints/hf/minimaxm3_weight_mapper.py | 57 +++++ .../_torch/models/modeling_minimaxm3.py | 55 +++-- .../hf/test_minimaxm3_weight_mapper.py | 198 ++++++++++++++++++ .../unittest/_torch/models/test_minimax_m3.py | 27 ++- 4 files changed, 304 insertions(+), 33 deletions(-) create mode 100644 tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py create mode 100644 tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/minimaxm3_weight_mapper.py new file mode 100644 index 000000000000..bb6d58957fd3 --- /dev/null +++ b/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() + } diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 2416c13665df..f10696c97bb9 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/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, diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py new file mode 100644 index 000000000000..896539beaa71 --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -0,0 +1,198 @@ +# 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. + +import inspect +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.auto_mapper import AutoCheckpointMapper +from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( + MiniMaxM3HfWeightMapper, +) +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + MiniMaxM3VLForConditionalGeneration, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +_NUM_KV_HEADS = 4 +_ROWS_PER_HEAD = 2 + + +def _make_mapper(tp_size: int = 8, num_kv_heads: int = _NUM_KV_HEADS) -> MiniMaxM3HfWeightMapper: + config = SimpleNamespace(num_key_value_heads=num_kv_heads, num_attention_heads=8) + model_config = ModelConfig( + pretrained_config=config, + mapping=Mapping(world_size=tp_size, rank=0, tp_size=tp_size), + ) + model = SimpleNamespace(model_config=model_config, config=config) + mapper = MiniMaxM3HfWeightMapper() + mapper.init_model_and_config(model, model_config) + return mapper + + +def _duplicate_heads(tensor: torch.Tensor, repetitions: int) -> torch.Tensor: + return ( + tensor.reshape(_NUM_KV_HEADS, _ROWS_PER_HEAD, -1) + .repeat_interleave(repetitions, dim=0) + .reshape(_NUM_KV_HEADS * repetitions * _ROWS_PER_HEAD, -1) + ) + + +@pytest.mark.parametrize( + "architecture", + [ + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + ], +) +def test_mapper_registration_and_mx_fallback(architecture: str) -> None: + assert isinstance(AutoCheckpointMapper.get("HF", architecture), MiniMaxM3HfWeightMapper) + assert isinstance(AutoCheckpointMapper.get("MX", architecture), MiniMaxM3HfWeightMapper) + + +@pytest.mark.parametrize( + "model_class", + [MiniMaxM3ForCausalLM, MiniMaxM3VLForConditionalGeneration], +) +def test_load_weights_exposes_weight_mapper(model_class: type) -> None: + assert "weight_mapper" in inspect.signature(model_class.load_weights).parameters + + +@pytest.mark.parametrize("scale_name", ["weight_scale_inv", "weight_scale"]) +def test_tp8_mxfp8_duplicates_kv_weight_and_scale_via_callbacks(scale_name: str) -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.MXFP8)) + prefix = "model.layers.0.self_attn" + + weights = {} + sources = {} + for offset, projection in enumerate(("k_proj", "v_proj")): + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + offset * 100 + scale = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + offset * 32 + weights[f"{prefix}.{projection}.weight"] = weight + weights[f"{prefix}.{projection}.{scale_name}"] = scale + sources[projection] = {"weight": weight, scale_name: scale} + + mapped = mapper.apply_callbacks(module, "qkv_proj", prefix.split("."), weights) + + assert mapped[0] == {} + for projection, projected in zip(("k_proj", "v_proj"), mapped[1:]): + for name, source in sources[projection].items(): + torch.testing.assert_close(projected[name], _duplicate_heads(source, repetitions=2)) + + +def test_nvfp4_scale_behavior_is_preserved() -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4)) + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + weight_scale = torch.arange(16, dtype=torch.float32).reshape(8, 2) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "k_proj", + { + "weight": weight, + "weight_scale": weight_scale, + "weight_scale_inv": scale_inv, + }, + ) + + torch.testing.assert_close(mapped["weight"], _duplicate_heads(weight, repetitions=2)) + torch.testing.assert_close( + mapped["weight_scale"], _duplicate_heads(weight_scale, repetitions=2) + ) + assert mapped["weight_scale_inv"] is scale_inv + + +def test_quant_config_none_is_guarded() -> None: + mapper = _make_mapper() + module = SimpleNamespace(quant_config=None) + weight = torch.arange(24, dtype=torch.float32).reshape(8, 3) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "k_proj", + { + "weight": weight, + "weight_scale_inv": scale_inv, + }, + ) + + torch.testing.assert_close(mapped["weight"], _duplicate_heads(weight, repetitions=2)) + assert mapped["weight_scale_inv"] is scale_inv + + +def test_kv_scale_is_not_expanded_when_kv_heads_cover_tp() -> None: + mapper = _make_mapper(tp_size=2) + module = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.MXFP8)) + scale_inv = torch.arange(16, dtype=torch.uint8).reshape(8, 2) + + mapped = mapper._duplicate_kv_weights( + module, + "v_proj", + {"weight_scale_inv": scale_inv}, + ) + + torch.testing.assert_close(mapped["weight_scale_inv"], scale_inv) + + +def test_gate_bias_params_map() -> None: + mapper = MiniMaxM3HfWeightMapper() + source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias" + target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias" + bias = torch.arange(4, dtype=torch.float32) + gate_weight = torch.ones(4, 4) + + renamed = mapper.rename_by_params_map( + mapper.params_map, + { + source_name: bias, + "model.layers.3.block_sparse_moe.gate.weight": gate_weight, + }, + ) + + assert source_name not in renamed + assert renamed[target_name] is bias + assert renamed["model.layers.3.block_sparse_moe.gate.weight"] is gate_weight + + +def test_load_weights_accepts_base_mapper_without_params_map() -> None: + config = SimpleNamespace(num_key_value_heads=_NUM_KV_HEADS, num_attention_heads=8) + model_config = ModelConfig(pretrained_config=config, mapping=Mapping()) + model = object.__new__(MiniMaxM3ForCausalLM) + torch.nn.Module.__init__(model) + model.model_config = model_config + mapper = HfWeightMapper() + source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias" + target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias" + bias = torch.arange(4, dtype=torch.float32) + + with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights: + model.load_weights({source_name: bias}, weight_mapper=mapper) + + call_kwargs = base_load_weights.call_args.kwargs + assert call_kwargs["weight_mapper"] is mapper + renamed = mapper.rename_by_params_map(call_kwargs["params_map"], {source_name: bias}) + assert renamed[target_name] is bias diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index ca2ed046e57c..1383b5577ccc 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -32,6 +32,9 @@ from utils.llm_data import llm_models_root from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( + MiniMaxM3HfWeightMapper, +) from tensorrt_llm._torch.models.modeling_minimaxm3 import ( MiniMaxM3Attention, _build_swiglu_oai_dense_mlp, @@ -43,7 +46,7 @@ get_text_config, is_minimax_m3_vl_config, ) -from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl +from tensorrt_llm._torch.models.modeling_utils import _load_weights_impl_v2 from tensorrt_llm._torch.modules.fused_moe.routing import ( MiniMaxM2MoeRoutingMethod, MiniMaxM3MoeRoutingMethod, @@ -763,8 +766,8 @@ def test_minimax_m3_routing_method_default_scale_is_identity(): @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 needs CUDA") -def test_text_norm_weights_real_loader_smoke(): - """real ``_load_weights_impl`` populates norm parameters. +def test_text_norm_weights_real_loader_smoke(monkeypatch: pytest.MonkeyPatch): + """real ``_load_weights_impl_v2`` populates norm parameters. Constructs a memory-safe stub containing the top-level ``model.norm`` and the first decoder layer's ``input_layernorm`` and @@ -772,7 +775,7 @@ def test_text_norm_weights_real_loader_smoke(): :class:`RMSNorm`, ~12 KB on CUDA), reads the corresponding tensors from the real checkpoint via ``safetensors``, strips the ``language_model.`` prefix exactly as the M3 VL wrapper does, and - invokes :func:`_load_weights_impl` end-to-end. The test fails if any + invokes :func:`_load_weights_impl_v2` end-to-end. The test fails if any target parameter remains at its zero-initialisation, proving the canonical loader walks the module tree and copies the correct source keys for these BF16 parameters. @@ -871,17 +874,21 @@ def __init__(self) -> None: "model.layers.0.post_attention_layernorm.weight", } - # Invoke the canonical loader. `_load_weights_impl` walks the stub's + # Invoke the canonical loader. `_load_weights_impl_v2` walks the stub's # module tree and uses the generic per-parameter copy fallback because # RMSNorm does not define ``load_weights``. Disable the parallel # executor so a failure surfaces immediately rather than as a thread # traceback (the parallel path is exercised in production; for this # tiny 3-module slice the serial walk is what the test should observe). - os.environ["TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL"] = "True" - try: - _load_weights_impl(stub, text_weights, allow_partial_loading=True) - finally: - os.environ.pop("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", None) + monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") + weight_mapper = MiniMaxM3HfWeightMapper() + weight_mapper.init_model_and_config(stub, stub.model_config) + _load_weights_impl_v2( + stub, + text_weights, + weight_mapper, + allow_partial_loading=True, + ) # The three norms should now hold the source tensors' values. torch.testing.assert_close( From 5e877639fcd5536eaf934aaf157b651b04ad954b Mon Sep 17 00:00:00 2001 From: zhaoyangwang-nvidia Date: Thu, 23 Jul 2026 00:41:57 +0800 Subject: [PATCH 3/5] [TRTLLM-13231][feat] Support top_p_decay in the PyTorch TorchSampler (#16184) Signed-off-by: ZhaoyangWang --- docs/source/features/sampling.md | 16 +- .../_torch/pyexecutor/sampler/ops/vanilla.py | 102 +++++ .../_torch/pyexecutor/sampler/sampler.py | 373 ++++++++++++++++++ .../pyexecutor/sampler/sampling_utils.py | 87 +++- tensorrt_llm/sampling_params.py | 78 +++- .../_torch/sampler/test_torch_sampler.py | 137 +++++++ 6 files changed, 775 insertions(+), 18 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index d5a9eaef0685..72ad4a62ad41 100644 --- a/docs/source/features/sampling.md +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index 5307876ae61d..8cfcca1e13b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2f5529bcb2c8..041cb31f97be 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -97,11 +97,13 @@ GenericStrategyKeyType, Strategy, StrategyMetadata, + TopPDecayMetadata, UtilsSamplingParams, get_rejected_indices, resolve_sampling_strategy, sample, sample_rejected, + top_p_decay_active, ) if sys.version_info[:2] >= (3, 12): @@ -475,9 +477,17 @@ def _get_max_beam_width(request: LlmRequest) -> int: def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: sampling_config = request.sampling_config + # These sampling fields live on the C++ SamplingConfig as optional> + # (a shape designed for the batched TRT-LLM sampler); the torch sampler consumes + # them per request, so we unwrap the singleton lists into scalars here. When the + # TRT-LLM sampler is removed, this SamplingConfig-based plumbing should be removed + # too in favor of reading the values directly from the per-request params. temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) + top_p_decay = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_decay)) + top_p_min = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_min)) + top_p_reset_ids = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_p_reset_ids)) beam_width_out = _get_beam_width_out(request) beam_width_in = _get_beam_width_in(request) use_beam_search = _get_max_beam_width(request) > 1 @@ -489,6 +499,9 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: beam_width_in=beam_width_in, beam_width_out=beam_width_out, use_beam_search=use_beam_search, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, ) @@ -2247,6 +2260,79 @@ class LogProbsStore: """Shape: batch_size, max_tokens, max_topk_logprobs Usage: Stores the values of the topk logprobs""" + @dataclass(kw_only=True) + class TopPDecayStore: + """Per-slot runtime state for Top-P Decay -- the single source of truth + for the feature's semantics on the torch path. + + Semantics (matching the legacy C++ ``computeToppDecay`` kernel): after + every sampled token of a decay-active request:: + + runtime_top_p = initial_top_p if token == reset_id + = max(runtime_top_p * top_p_decay, top_p_min) otherwise + + A negative ``reset_ids`` sentinel (-1, "reset disabled") never matches, + since sampled token ids are non-negative. Decay is active iff + ``top_p_decay`` is set and < 1.0 (``SamplingParams. + params_imply_top_p_decay_active``); an active decay forces a + top-p-capable strategy even for an otherwise implicitly-greedy request + (initial top-p defaults to 1.0), while explicit greedy controls win. + Beam search and speculative draft tokens are rejected at admission + (``_validate_top_p_decay_request``); parameter ranges are enforced by + ``SamplingParams._validate`` / the executor::SamplingConfig constructor. + + Lifecycle per slot (each tensor has shape ``(max_num_sequences,)``): + + 1. Admission (``_setup_top_p_decay_for_new_requests``): membership is + cleared then re-set for the newly-admitted slots -- both the host-side + ``TorchSampler._top_p_decay_slots`` set (an O(1) hot-path early-out) + and its device mirror ``is_top_p_decay_slot_cuda`` (the gate the + fused ops use, so the hot path needs no host-side filtering) -- + and the per-slot buffers are initialized. This clear-then-set also + covers slot reuse: stale entries from a prior occupant are never + consumed. + 2. Pre-sample (``_build_top_p_decay_metadata`` -> ``TopPDecayMetadata`` + -> ``TopPDecayMixin``): the per-row top-p fed to top_p / + top_k_top_p sampling is overridden with the decayed runtime value + for decay-active rows (fused gather, ``top_p_decay_gather``). + 3. Post-sample (``_update_top_p_decay_after_sample``): the recurrence + above is applied in place for the sampled decay-active slots (fused + update, ``top_p_decay_update``). + 4. Finish (``_retire_top_p_decay_slot``): the slot leaves the + membership set so the early-outs re-arm; the device buffers need no + cleanup (a freed slot is never sampled, reuse re-initializes it). + """ + + runtime_top_p_decay_cuda: torch.Tensor + """The current (decaying) top-p per slot; mutated post-sample each step.""" + initial_top_p_decay_cuda: torch.Tensor + """The initial top-p per slot; used to reset on a reset-id match.""" + top_p_decay_cuda: torch.Tensor + """Per-slot multiplicative decay factor.""" + top_p_decay_min_cuda: torch.Tensor + """Per-slot lower bound for the decayed top-p.""" + top_p_decay_reset_ids_cuda: torch.Tensor + """Per-slot reset token id (< 0 never matches a sampled token).""" + is_top_p_decay_slot_cuda: torch.Tensor + """Per-slot bool gate (device mirror of ``_top_p_decay_slots``). Lets the + fused post-sample update op filter decay-active slots on the GPU, + avoiding a host-side ``.tolist()`` / set intersection each step.""" + + @classmethod + def create(cls, max_num_sequences: int) -> "TorchSampler.TopPDecayStore": + n = (max_num_sequences,) + return cls( + runtime_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + initial_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_min_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_reset_ids_cuda=torch.empty(n, dtype=torch.int, device="cuda"), + # The gate buffer IS the gate, so (unlike the others) it must start + # False: a slot that was never admitted as decay-active must read + # False. + is_top_p_decay_slot_cuda=torch.zeros(n, dtype=torch.bool, device="cuda"), + ) + @dataclass(kw_only=True) class Store: new_tokens: torch.Tensor @@ -2258,6 +2344,8 @@ class Store: """Holds data related to beam search.""" log_probs_store: "TorchSampler.LogProbsStore" """Holds data related to log-probs handling.""" + top_p_decay_store: "TorchSampler.TopPDecayStore" + """Holds per-slot runtime state for Top-P Decay.""" def _create_store(self) -> Store: # Tensors necessary for all sampling methods @@ -2308,10 +2396,15 @@ def _create_store(self) -> Store: seq_offsets=seq_offsets, beam_idx_arange=beam_idx_arange, ) + # Per-slot Top-P Decay runtime state (FlashInfer path). Allocated for all + # sampler instances; only slots in self._top_p_decay_slots are ever read. + top_p_decay_store = self.TopPDecayStore.create(self.max_num_sequences) + return self.Store( new_tokens=new_tokens, log_probs_store=log_probs_store, beam_search_store=beam_search_store, + top_p_decay_store=top_p_decay_store, ) @dataclass(frozen=True, kw_only=True) @@ -2364,6 +2457,10 @@ def __init__(self, args: Args): "in requirements.txt." ) self._grouped_sampler_cls = FlashInferGroupedStrategySampler + # Slots with an active top-p-decay request. Sole gate for reading/updating + # the per-slot TopPDecayStore buffers; discarded on slot reuse so stale + # buffer entries are never consumed. + self._top_p_decay_slots: set[int] = set() # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. @@ -2464,6 +2561,48 @@ def get_spec_tree_manager( def _use_beam_search(self) -> bool: return self.max_beam_width > 1 + def _validate_top_p_decay_request(self, request: LlmRequest) -> None: + """Reject unsupported combinations for a top-p-decay-active request. + + Top-p decay is supported only for single-token decode steps without beam + search. Called from validate_request (request admission), so a violating + request is failed individually instead of aborting the whole batch. + """ + params = _request_get_sampling_params(request) + # NB: value ranges need no re-check here. Every request enters through + # the executor::SamplingConfig constructor, which hard-validates + # top_p_decay in (0, 1], top_p_min in (0, 1] and top_p_reset_ids >= 0 + # (samplingConfig.cpp check* helpers) for all frontends. A reset id + # >= vocab_size is not checked anywhere but is semantically inert: it + # can never match a sampled token, i.e. it behaves as "reset disabled". + if not top_p_decay_active(params): + return + if params.use_beam_search: + raise ValueError("top_p_decay is not supported with beam search.") + # A non-zero draft length means the request carries speculative draft + # tokens and produces multiple tokens per step (req_num_steps = + # 1 + draft_token_length). One-model speculation (vanilla MTP, one-model + # Eagle3 / MTP-Eagle, SA, draft-target-one-model) uses its own + # SpecSamplerBase-derived sampler and never reaches TorchSampler; the + # drafter-based modes that DO flow draft tokens through TorchSampler + # (two-model draft-target, NGram, user-provided, two-model Eagle3 / + # MTP-Eagle) are what can make this length non-zero. top-p decay does not + # support these multi-token steps. + # NB: at admission time the draft tokens of drafter-based modes are + # usually not attached yet, so this check is best-effort. Two-model + # speculation (the only source of such requests in TorchSampler) is + # slated for removal, so in practice no speculative request reaches the + # decay path; a debug assert in _build_top_p_decay_metadata guards the + # invariant at sample time. + if get_draft_token_length(request) > 0: + raise ValueError( + "top_p_decay is not supported for requests carrying speculative " + "draft tokens (req_num_steps > 1). This covers the drafter-based " + "modes routed through TorchSampler (two-model draft-target, NGram, " + "user-provided, two-model Eagle3 / MTP-Eagle); one-model " + "speculation uses its own sampler and is unaffected." + ) + def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool: """ Check if we can use the fast argmax path for greedy sampling. @@ -2924,6 +3063,10 @@ def _collect_new_requests_for_setup( @override def validate_request(self, request: LlmRequest) -> None: + # Reject unsupported top-p-decay combinations at admission, so only the + # offending request fails (raising later, inside setup_sampler_step or + # sampling, would abort the whole executor step). + self._validate_top_p_decay_request(request) if self._use_beam_search: if request.py_return_log_probs: if request.py_num_logprobs > 1: @@ -2998,6 +3141,10 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: all_sampling_requests=new_requests + scheduled_requests.generation_requests, ) + self._setup_top_p_decay_for_new_requests( + new_requests, new_seq_slots_cuda_long=seq_slots_tensor_cuda_long + ) + if self._use_beam_search: beam_search_store = self.store.beam_search_store assert beam_search_store is not None @@ -3008,6 +3155,106 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_prompt_len=max_prompt_len, ) + def _setup_top_p_decay_for_new_requests( + self, + new_requests: list[LlmRequest], + *, + new_seq_slots_cuda_long: torch.Tensor, + ) -> None: + """Refresh top-p-decay membership and per-slot buffers for admitted requests + (lifecycle step 1, see ``TopPDecayStore``). + + Drops stale membership from prior occupants of the newly-admitted slots + (host set and device gate), then re-admits the decay-active requests and + initializes their per-slot store entries. Unsupported decay combinations + were already rejected per-request in validate_request at admission. + """ + # Clear the device decay gate for every newly-admitted slot (covers slot + # reuse: a slot previously decay-active but reused by a non-decay request + # must read False). Decay-active slots are then set True below. + decay_gate = self.store.top_p_decay_store.is_top_p_decay_slot_cuda + decay_gate.index_fill_(0, new_seq_slots_cuda_long, False) + + decay_seq_slots: list[int] = [] + initial_top_p: list[float] = [] + top_p_decay: list[float] = [] + top_p_min: list[float] = [] + top_p_reset_ids: list[int] = [] + for request in new_requests: + slot = request.py_seq_slot + assert slot is not None + self._top_p_decay_slots.discard(slot) + sampling_params = _request_get_sampling_params(request) + if not top_p_decay_active(sampling_params): + continue + self._top_p_decay_slots.add(slot) + decay_seq_slots.append(slot) + # Initial runtime top-p defaults to 1.0 when top_p is unset. + initial_top_p.append( + sampling_params.top_p if sampling_params.top_p is not None else 1.0 + ) + # decay is guaranteed non-None and < 1.0 here (top_p_decay_active); + # min/reset fall back to the C++ runtime defaults when unset. + assert sampling_params.top_p_decay is not None + top_p_decay.append(sampling_params.top_p_decay) + top_p_min.append( + sampling_params.top_p_min if sampling_params.top_p_min is not None else 1e-6 + ) + top_p_reset_ids.append( + sampling_params.top_p_reset_ids + if sampling_params.top_p_reset_ids is not None + else -1 + ) + + if decay_seq_slots: + self._update_top_p_decay_store_for_new_requests( + decay_seq_slots=decay_seq_slots, + initial_top_p=initial_top_p, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, + ) + + def _update_top_p_decay_store_for_new_requests( + self, + *, + decay_seq_slots: list[int], + initial_top_p: list[float], + top_p_decay: list[float], + top_p_min: list[float], + top_p_reset_ids: list[int], + ) -> None: + """Initialize per-slot Top-P Decay buffers for newly admitted decay requests. + + runtime_top_p and initial_top_p both start at the effective initial top-p; + the runtime value is decayed post-sample each step. + """ + store = self.store.top_p_decay_store + device = store.runtime_top_p_decay_cuda.device + slots_cuda = torch.tensor( + decay_seq_slots, device="cpu", dtype=torch.int64, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + floats_host = torch.tensor( + [initial_top_p, top_p_decay, top_p_min], + device="cpu", + dtype=torch.float32, + pin_memory=prefer_pinned(), + ) + floats_cuda = floats_host.to(device, non_blocking=True) + initial_cuda = floats_cuda[0] + reset_ids_cuda = torch.tensor( + top_p_reset_ids, device="cpu", dtype=torch.int32, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + + store.runtime_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.initial_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.top_p_decay_cuda.index_copy_(0, slots_cuda, floats_cuda[1]) + store.top_p_decay_min_cuda.index_copy_(0, slots_cuda, floats_cuda[2]) + store.top_p_decay_reset_ids_cuda.index_copy_(0, slots_cuda, reset_ids_cuda) + # Enable the device gate for these decay-active slots (cleared for all new + # slots in setup_sampler_step just before this call). + store.is_top_p_decay_slot_cuda.index_fill_(0, slots_cuda, True) + @staticmethod def _prepare_beam_search( beam_search_store: BeamSearchStore, @@ -3493,6 +3740,7 @@ def _add_metadata_to_grouped_requests( *, seq_slots_cuda: torch.Tensor, seq_lens_cuda: torch.Tensor, + req_num_steps: torch.Tensor, ) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata]: grouped_requests_with_metadata: dict[ RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata @@ -3502,6 +3750,7 @@ def _add_metadata_to_grouped_requests( num_requests = len(requests) for key, value in grouped_requests.items(): metadata_type = get_metadata_type_for_group_fn(key.strategy_key) + metadata: StrategyMetadata | None if metadata_type is BeamSearchMetadata: assert beam_search_store is not None assert seq_lens is not None, "seq_lens is required for beam search" @@ -3530,6 +3779,13 @@ def _add_metadata_to_grouped_requests( seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, ) + elif metadata_type is TopPDecayMetadata: + metadata = self._build_top_p_decay_metadata( + group_req_indices=value.indices, + req_num_steps=req_num_steps, + seq_slots=seq_slots, + seq_slots_cuda=seq_slots_cuda, + ) elif metadata_type is None: metadata = None else: @@ -3727,6 +3983,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: + self._retire_top_p_decay_slot(req) continue if req.py_beam_width > 1: @@ -3793,6 +4050,20 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: self._finish_reasons_handler.store.num_accepted_draft_tokens_host[ req.py_seq_slot ] = req.py_num_accepted_draft_tokens + if req.state == LlmRequestState.GENERATION_COMPLETE: + self._retire_top_p_decay_slot(req) + + def _retire_top_p_decay_slot(self, req: LlmRequest) -> None: + """Retire a finished request's slot from the top-p-decay membership set + (lifecycle step 4, see ``TopPDecayStore``), so the O(1) hot-path + early-outs re-arm once decay traffic drains. + + Callers ensure ``req`` has finished. Requests that finish outside the + sampler (e.g. cancellation) are covered by the slot-reuse cleanup at + admission instead. + """ + if self._top_p_decay_slots and req.py_seq_slot is not None: + self._top_p_decay_slots.discard(req.py_seq_slot) def _return_log_probs(self, requests: list[LlmRequest]) -> bool: return any(req.py_return_log_probs for req in requests) @@ -4060,6 +4331,65 @@ def provision_bias_index() -> int: # sharing). logits[logits_bias_mask_cuda] += biases_tensor_cuda + def _build_top_p_decay_metadata( + self, + *, + group_req_indices: torch.Tensor, + req_num_steps: torch.Tensor, + seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, + ) -> Optional[TopPDecayMetadata]: + """Build the Top-P Decay metadata for a top_p / top_k_top_p group. + + Lifecycle step 2, see ``TopPDecayStore``. Returns None when no request + currently uses decay. The metadata's ``slots`` tensor is aligned to the + group's per-STEP row order (matching group_strategies_per_step); + non-decay rows (possibly multi-step draft rows) are gated out on-device + by ``is_decay_slot``, so decay presence in the group is not checked + host-side: a group without decay rows samples every row with its static + top-p -- same result as returning None. + """ + if not self._top_p_decay_slots: + return None + store = self.store.top_p_decay_store + # Fast path (steady-state decoding): if every row in the group is + # single-token, the per-STEP row order equals the per-request order, and + # (group_req_indices being sorted ascending) a contiguous group's slots + # are just a slice of seq_slots_cuda -- no host layout build and no H2D + # copy. + first_req = int(group_req_indices[0].item()) + last_req = int(group_req_indices[-1].item()) + group_steps = req_num_steps[group_req_indices] + if last_req - first_req + 1 == group_req_indices.size(0) and ( + group_steps.max().item() == 1 + ): + per_step_slots_cuda = seq_slots_cuda[first_req : last_req + 1] + else: + # Build the per-STEP slot layout (each request contributes + # req_num_steps rows). + group_seq_slots = seq_slots[group_req_indices] + if __debug__: + # Internal invariant (stripped under python -O): a decay-active + # row is always single-token -- the only source of multi-step + # rows in TorchSampler is two-model speculation (slated for + # removal), and decay + draft tokens is rejected per-request at + # admission in validate_request. + decay_row_steps = group_steps[ + torch.isin(group_seq_slots, torch.tensor(list(self._top_p_decay_slots))) + ] + assert decay_row_steps.numel() == 0 or decay_row_steps.max().item() == 1, ( + "top_p_decay row with req_num_steps != 1; decay + draft tokens " + "should have been rejected at admission" + ) + per_step_slots_cuda = torch.repeat_interleave(group_seq_slots.long(), group_steps).to( + seq_slots_cuda.device, non_blocking=True + ) + return TopPDecayMetadata( + slots=per_step_slots_cuda, + runtime_top_p=store.runtime_top_p_decay_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + ) + @nvtx_range("sample_batched_by_strategy") @torch.inference_mode() def _sample_batched_by_strategy( @@ -4096,6 +4426,7 @@ def _sample_batched_by_strategy( get_metadata_type_for_group_fn=self._grouped_sampler_cls.get_metadata_type_for_group, seq_slots_cuda=seq_slots_cuda, seq_lens_cuda=seq_lens_cuda, + req_num_steps=req_num_steps, ) generator_cuda = self.get_generator(cuda_device) @@ -4329,6 +4660,7 @@ def _unbatch_sampling_results( new_tokens_cuda: torch.Tensor, req_num_generated_tokens: torch.Tensor, seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, ) -> torch.Tensor: batch_req_indices = batched_sampling_result.batch_req_indices batch_next_tokens_cuda_int = batched_sampling_result.batch_next_tokens_cuda_int @@ -4362,8 +4694,48 @@ def _dims_canonically_ordered(t: torch.Tensor) -> bool: new_tokens_cuda.view(-1, *new_tokens_cuda.shape[2:]).scatter_( 0, batch_dest_indices_1d_cuda, batch_next_tokens_cuda_int ) + # Post-sample: decay the runtime top-p for any decay-active slots that were + # sampled this iteration (must run after tokens land in new_tokens_cuda). + # batch_req_indices is a permutation of all sampled requests, so the set of + # sampled slots is exactly seq_slots (the kernel updates each slot + # independently; order is irrelevant) -- pass the resident device copy + # instead of gathering seq_slots[batch_req_indices] on host and copying it. + self._update_top_p_decay_after_sample( + new_tokens_cuda=new_tokens_cuda, + sampled_slots_cuda=seq_slots_cuda, + ) return self._copy_to_host(new_tokens_cuda) + def _update_top_p_decay_after_sample( + self, + *, + new_tokens_cuda: torch.Tensor, + sampled_slots_cuda: torch.Tensor, + ) -> None: + """Apply the post-sample decay recurrence for sampled decay-active slots. + + See ``TopPDecayStore`` for the feature-level semantics (lifecycle + step 3). Restricting to the sampled slots avoids reading stale + new_tokens_cuda for slots that were not scheduled this iteration. + """ + # Host-side O(1) early-out: skip the kernel launch when no request uses + # decay; otherwise a single fused (torch.compile) op gates on + # is_decay_slot on-device and gathers the sampled token in place (decay is + # single-token-only, so the token is at local step 0, beam 0). + if not self._top_p_decay_slots: + return + store = self.store.top_p_decay_store + Fusions.top_p_decay_update( + runtime_top_p=store.runtime_top_p_decay_cuda, + initial_top_p=store.initial_top_p_decay_cuda, + top_p_decay=store.top_p_decay_cuda, + top_p_min=store.top_p_decay_min_cuda, + reset_ids=store.top_p_decay_reset_ids_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + step_tokens=new_tokens_cuda[DEFAULT_STEP_IDX, :, DEFAULT_BEAM_IDX], + sampled_slots=sampled_slots_cuda, + ) + @staticmethod @torch.inference_mode() def _apply_min_length_penalty( @@ -4998,6 +5370,7 @@ def _process_requests( new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=sampling_requests_metadata.req_num_generated_tokens, seq_slots=seq_slots_host, + seq_slots_cuda=seq_slots_cuda, ) # NB: update_requests syncs w/ device computation and async D2H copies diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index f99946ffd5a1..7ca58c7edc08 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -106,6 +106,10 @@ class UtilsSamplingParams: use_beam_search: Whether to use beam search. beam_width_in: The beam_width of a request before the sampling step. beam_width_out: The beam_width of a request after the sampling step. + top_p_decay: Per-step multiplicative decay applied to the runtime top-p. + top_p_min: Lower bound for the decayed runtime top-p. + top_p_reset_ids: Token id which, when sampled, resets the runtime top-p to + its initial value. A value < 0 never matches a token. """ temperature: Optional[float] @@ -114,6 +118,38 @@ class UtilsSamplingParams: use_beam_search: Optional[bool] beam_width_in: Optional[int] = None beam_width_out: Optional[int] = None + top_p_decay: Optional[float] = None + top_p_min: Optional[float] = None + top_p_reset_ids: Optional[int] = None + + +@dataclass(kw_only=True) +class TopPDecayMetadata(StrategyMetadata): + """Per-group runtime top-p override for Top-P Decay (attached to top_p / + top_k_top_p groups via the ``StrategyMetadata`` mechanism). + + ``slots`` maps each per-step group row to its sequence slot; the decayed + per-row top-p is gathered on-device from the per-slot ``runtime_top_p`` + store, gated by ``is_decay_slot`` (non-decay rows keep their static top-p). + Consumed by the TopP*/TopKTopP* strategy impls in ``sample()``. See + ``TorchSampler.TopPDecayStore`` for the feature-level semantics. + """ + + slots: torch.Tensor + """Per-step group rows' sequence slots (int64, device).""" + runtime_top_p: torch.Tensor + """Per-slot runtime (decayed) top-p store (float32, device).""" + is_decay_slot: torch.Tensor + """Per-slot decay-active gate (bool, device).""" + + +def top_p_decay_active(params: UtilsSamplingParams) -> bool: + """Whether dynamic top-p decay is active for a request. + + Delegates to the single-source predicate on SamplingParams; note that + ``top_p_min`` / ``top_p_reset_ids`` alone do not activate dynamic behavior. + """ + return SamplingParams.params_imply_top_p_decay_active(params.top_p_decay) def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: @@ -124,11 +160,15 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - top_p = params.top_p top_k = params.top_k + # The greedy verdict (including the top-p-decay override of the implicit + # all-unset greedy default, and explicit greedy controls winning over decay) + # is single-sourced in SamplingParams.params_imply_greedy_decoding. if SamplingParams.params_imply_greedy_decoding( temperature=temperature, top_p=top_p, top_k=top_k, use_beam_search=use_beam_search, + top_p_decay=params.top_p_decay, ): return GREEDY @@ -152,7 +192,10 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - assert top_k > 1, "non-greedy sampling requires valid top_k" need_top_k = top_k < vocab_size assert top_p > 0, "non-greedy sampling requires valid top_p" - need_top_p = top_p < 1 + # A decay-active request must go through a top-p-capable path even when its + # initial top_p is 1.0, so the runtime top-p (sourced per-row at sample time) + # can shrink the nucleus on later steps. + need_top_p = top_p < 1 or top_p_decay_active(params) if need_top_p: if need_top_k: @@ -346,6 +389,34 @@ def _sample_with_probs( new_tokens = cls._sample_from_probs(probs, generator=generator) return new_tokens, probs + class TopPDecayMixin: + """Mixed into the TopP*/TopKTopP* impls (the owners of a per-row + ``_top_p`` tensor) to consume ``TopPDecayMetadata``.""" + + _top_p: torch.Tensor + + def _maybe_apply_top_p_decay(self, group_metadata: Optional[StrategyMetadata]) -> None: + """Override the per-row static top-p with the decayed runtime top-p. + + Only decay-active rows (per the on-device ``is_decay_slot`` gate) are + overridden, so a group mixing top-p-decay and plain top-p requests + keeps each row's correct value. The overridden ``self._top_p`` tensor + then feeds both sampling and ``top_p_renorm_probs_op`` (so processed + logprobs match). Fused via torch.compile (gather + gate + select). + """ + if not isinstance(group_metadata, TopPDecayMetadata): + return + assert self._top_p.shape == group_metadata.slots.shape, ( + self._top_p.shape, + group_metadata.slots.shape, + ) + self._top_p = Fusions.top_p_decay_gather( + runtime_top_p=group_metadata.runtime_top_p, + is_decay_slot=group_metadata.is_decay_slot, + static_top_p=self._top_p, + slots=group_metadata.slots, + ) + class StrategyImplWithProbs(StrategyImpl): @override @classmethod @@ -374,7 +445,7 @@ def sample( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: return self._sample_greedy_with_probs(logits, group_logit_indices=group_logit_indices) - class TopKTopPWithProbs(StrategyImplWithProbs): + class TopKTopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -400,6 +471,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -442,7 +514,7 @@ def sample( generator=generator, ) - class TopPWithProbs(StrategyImplWithProbs): + class TopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -466,6 +538,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -534,7 +607,7 @@ def sample( logits = logits[group_logit_indices] return torch.argmax(logits, dim=-1), None - class TopKTopPSampleOnly(StrategyImplSampleOnly): + class TopKTopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -560,6 +633,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) logits = self._prepare_logits_with_temperature( logits, group_logit_indices, self._temperature ) @@ -605,7 +679,7 @@ def sample( check_nan=self._flashinfer_check_nans(probs), ), None - class TopPSampleOnly(StrategyImplSampleOnly): + class TopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -629,6 +703,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) probs = self._prepare_probs_with_temperature( logits, group_logit_indices, self._temperature ) @@ -755,6 +830,8 @@ def get_metadata_type_for_group( match strategy_key: case ("beam_search", _, _): return BeamSearchMetadata + case "top_p" | "top_k_top_p": + return TopPDecayMetadata case _: return None diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index bded9477cd3d..e7fed5e2ced0 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -210,9 +210,9 @@ class SamplingParams: If neither temperature, top_p, nor top_k are specified, sampling is greedy. If temperature > 0 and/or top_k > 1 are specified, sampling will proceed accordingly and top_p will default to top_p = 1. Setting top_p = 0 should result in greedy sampling, but is currently disallowed in the backend. - top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. None means using C++ runtime default 1.e-6. Defaults to None. - top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. Indicates where to reset the decay. None means using C++ runtime default 1. Defaults to None. - top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None. + top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.e-6. Defaults to None. + top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. The token id which, when sampled, resets the decayed top-P to its initial value. Must be >= 0; invalid values are rejected. None means using C++ runtime default -1 (which never matches a token). Defaults to None. + top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.f. Defaults to None. seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None. temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None. The value None is treated as "not specified" in the following. @@ -379,6 +379,20 @@ def _validate(self): if self.temperature is not None and self.temperature < 0: raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + # Top-p decay param ranges mirror the hard checks in the + # executor::SamplingConfig constructor (samplingConfig.cpp check* + # helpers); rejecting here gives a clear, early error instead of a + # RuntimeError from the C++ boundary. Note top_p_min > top_p is + # intentionally allowed (the runtime top-p may rise toward top_p_min). + if self.top_p_decay is not None and not 0.0 < self.top_p_decay <= 1.0: + raise ValueError(f"require 0 < top_p_decay <= 1, got top_p_decay={self.top_p_decay}") + if self.top_p_min is not None and not 0.0 < self.top_p_min <= 1.0: + raise ValueError(f"require 0 < top_p_min <= 1, got top_p_min={self.top_p_min}") + if self.top_p_reset_ids is not None and self.top_p_reset_ids < 0: + raise ValueError( + f"require top_p_reset_ids >= 0, got top_p_reset_ids={self.top_p_reset_ids}" + ) + if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") @@ -427,8 +441,37 @@ def _validate(self): if self.logprobs_simple_format and self.use_beam_search: raise ValueError("logprobs_simple_format is not supported with beam search") - # NB: Static, because downstream code only holds instances of - # bindings.SamplingConfig (not SamplingParams). + # NB: The predicates below are static because downstream code (e.g. + # sampling_utils.resolve_sampling_strategy) only holds instances of + # bindings.SamplingConfig (not SamplingParams). They are the single + # source of truth for the greedy / top-p-decay resolution shared by + # _greedy_decoding and the torch sampler. + + @staticmethod + def params_imply_top_p_decay_active(top_p_decay: Optional[float]) -> bool: + """Whether dynamic top-p decay is active. + + Active iff ``top_p_decay`` is explicitly set and ``< 1.0``; a decay of + ``1.0`` (the C++ default) is a no-op. Values outside ``(0, 1]`` are + rejected up front (_validate and the executor::SamplingConfig + constructor), so they never reach this predicate. + """ + return top_p_decay is not None and top_p_decay < 1.0 + + @staticmethod + def params_imply_explicit_greedy( + *, + temperature: Optional[float], + top_p: Optional[float], + top_k: Optional[int], + ) -> bool: + """Whether the request carries an explicit greedy control. + + Explicit means top_k == 1, top_p == 0.0, or temperature == 0, as opposed + to the implicit "all params unset" greedy default. + """ + return top_k == 1 or top_p == 0.0 or temperature == 0 + @staticmethod def params_imply_greedy_decoding( *, @@ -436,13 +479,23 @@ def params_imply_greedy_decoding( top_p: Optional[float], top_k: Optional[int], use_beam_search: bool | None, - ): - return (not use_beam_search) and ( - (temperature is None and top_p is None and top_k is None) - or top_k == 1 - or top_p == 0.0 - or temperature == 0 - ) + top_p_decay: Optional[float] = None, + ) -> bool: + """Whether the parameters resolve to greedy decoding. + + An explicit greedy control always wins. The implicit "all params unset" + greedy default is overridden by an active top-p decay (which implies + top-p sampling so the decayed runtime top-p can take effect); callers + that do not support decay may omit ``top_p_decay``. + """ + if use_beam_search: + return False + if SamplingParams.params_imply_explicit_greedy( + temperature=temperature, top_p=top_p, top_k=top_k + ): + return True + implicitly_greedy = temperature is None and top_p is None and top_k is None + return implicitly_greedy and not SamplingParams.params_imply_top_p_decay_active(top_p_decay) @property def _greedy_decoding(self) -> bool: @@ -451,6 +504,7 @@ def _greedy_decoding(self) -> bool: top_p=self.top_p, top_k=self.top_k, use_beam_search=self.use_beam_search, + top_p_decay=self.top_p_decay, ) @property diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e2e0d1bc9402..892b93357fb0 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -59,6 +59,7 @@ TopKTopP, TopP, UtilsSamplingParams, + resolve_sampling_strategy, ) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import SamplingConfig @@ -2577,6 +2578,11 @@ class UutResultWrapper: result: Optional[UutResult] = None res = UutResultWrapper() + # Precomputed outside the no-sync region (mirrors the production + # resident device copy of seq_slots). + seq_slots_tensor_cuda = ( + seq_slots_tensor.to(torch.int64).pin_memory().to("cuda", non_blocking=True) + ) def _uut(res=res): new_tokens_host = sampler._unbatch_sampling_results( @@ -2584,6 +2590,7 @@ def _uut(res=res): new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=req_num_steps, seq_slots=seq_slots_tensor, + seq_slots_cuda=seq_slots_tensor_cuda, ) res.result = UutResult(new_tokens_host=new_tokens_host) @@ -2756,3 +2763,133 @@ def test_stale_conditional_and_unconditional_same_cell(self): logits = self._run_stale([req], [True], {0: 9}) assert self._banned_cols(logits[0]) == {2} assert not torch.isnan(logits).any() + + +class TestTopPDecay: + """Minimal functional guards for Top-P Decay in TorchSampler. + + Covers strategy routing, the post-sample runtime update (parity with the + C++ computeToppDecay recurrence; cases ported from + topPSamplingLayerTest.cpp), and per-request rejection of unsupported + combinations. + """ + + VOCAB_SIZE = 1000 + + @staticmethod + def _params(**kw) -> UtilsSamplingParams: + base = dict(temperature=None, top_p=None, top_k=None, use_beam_search=False) + base.update(kw) + return UtilsSamplingParams(**base) + + @staticmethod + def _make_sampler(*, max_draft_len=0): + return TorchSampler( + TorchSampler.Args( + max_seq_len=128, + max_draft_len=max_draft_len, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=max_draft_len, + disable_overlap_scheduler=True, + ) + ) + + def test_strategy_routing(self): + # Active decay (set and < 1.0) forces a top-p-capable strategy even for + # an otherwise-greedy request (initial top-p defaults to 1.0), so the + # decayed runtime value can take effect on later steps. + s = resolve_sampling_strategy(self._params(top_p_decay=0.5), vocab_size=self.VOCAB_SIZE) + assert s[0] == "top_p" and s[1] == pytest.approx(1.0) + s = resolve_sampling_strategy( + self._params(top_k=50, top_p=0.9, top_p_decay=0.8), vocab_size=self.VOCAB_SIZE + ) + assert s[0] == "top_k_top_p" + # decay == 1.0 (the C++ default) is a no-op and does not activate... + s = resolve_sampling_strategy(self._params(top_p_decay=1.0), vocab_size=self.VOCAB_SIZE) + assert s is GREEDY + # ...and an explicit greedy control wins over an active decay. + s = resolve_sampling_strategy( + self._params(top_p_decay=0.5, top_k=1), vocab_size=self.VOCAB_SIZE + ) + assert s is GREEDY + + def test_runtime_update_parity(self): + # Post-sample update parity with the C++ computeToppDecay recurrence + # (a negative reset_id never matches, since token ids are non-negative): + # runtime = initial if token == reset_id + # = max(runtime * decay, min) otherwise + sampler = self._make_sampler() + store = sampler.store.top_p_decay_store + configs = [ + dict(initial=0.8, decay=0.3, top_p_min=0.5, reset_id=2), # decay, then reset + dict(initial=0.2, decay=0.9, top_p_min=0.1, reset_id=-1), # plain decay, floored + dict(initial=0.3, decay=0.5, top_p_min=0.6, reset_id=-1), # min > initial: rises + ] + token_steps = [[1, 2, 3], [9, 9, 9], [9, 9, 9]] + slots = list(range(len(configs))) + for slot, cfg in zip(slots, configs): + sampler._top_p_decay_slots.add(slot) + store.runtime_top_p_decay_cuda[slot] = cfg["initial"] + store.initial_top_p_decay_cuda[slot] = cfg["initial"] + store.top_p_decay_cuda[slot] = cfg["decay"] + store.top_p_decay_min_cuda[slot] = cfg["top_p_min"] + store.top_p_decay_reset_ids_cuda[slot] = cfg["reset_id"] + store.is_top_p_decay_slot_cuda[slot] = True + + runtime = [cfg["initial"] for cfg in configs] + slots_cuda = torch.tensor(slots, dtype=torch.int64, device="cuda") + for step in range(3): + for slot in slots: + sampler.store.new_tokens[0, slot, 0] = token_steps[slot][step] + sampler._update_top_p_decay_after_sample( + new_tokens_cuda=sampler.store.new_tokens, sampled_slots_cuda=slots_cuda + ) + got = store.runtime_top_p_decay_cuda.cpu() + for slot, cfg in zip(slots, configs): + tok = token_steps[slot][step] + if tok == cfg["reset_id"]: + runtime[slot] = cfg["initial"] + else: + runtime[slot] = max(runtime[slot] * cfg["decay"], cfg["top_p_min"]) + assert got[slot].item() == pytest.approx(runtime[slot], abs=1e-6), (step, slot) + + @staticmethod + def _mock_request(params: SamplingParams, *, draft_tokens=None): + params._validate() + req = SimpleNamespace( + sampling_config=SamplingConfig(params._get_sampling_config()), + is_context_init_state=False, + py_sampling_strategy=None, + py_draft_tokens=draft_tokens, + ) + req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 + return cast(LlmRequest, req) + + @pytest.mark.parametrize( + "bad_kwargs", + [ + {"top_p_decay": 1.5}, + {"top_p_decay": -0.5}, + {"top_p_decay": 0.0}, + {"top_p_min": 0.0}, + {"top_p_min": 1.5}, + {"top_p_reset_ids": -1}, + ], + ) + def test_out_of_range_decay_params_rejected(self, bad_kwargs): + # Out-of-range decay params raise (mirroring the executor::SamplingConfig + # constructor's hard checks) instead of the former warn-and-default. + with pytest.raises(ValueError): + SamplingParams(**bad_kwargs) + + def test_reject_speculative_draft_tokens(self): + # Decay + draft tokens through TorchSampler is rejected per-request at + # admission (validate_request), so only the offending request fails. + sampler = self._make_sampler(max_draft_len=4) + with pytest.raises(ValueError, match="speculative"): + sampler.validate_request( + self._mock_request(SamplingParams(top_p=0.9, top_p_decay=0.5), draft_tokens=[1, 2]) + ) + # Same request without decay is accepted. + sampler.validate_request(self._mock_request(SamplingParams(top_p=0.9), draft_tokens=[1, 2])) From 1fbd240363e73c6d2835e084f327fcb66fb04662 Mon Sep 17 00:00:00 2001 From: sunnyqgg <159101675+sunnyqgg@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:43:14 +0800 Subject: [PATCH 4/5] [https://nvbugs/6442074][fix] Rename MTPEagleDynamicTreeWorker.forward to _forward_impl (#16724) Signed-off-by: qgai --- tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 5dca84abae8e..053d23362f98 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -636,7 +636,7 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): # Top-level forward # # ------------------------------------------------------------------ # @nvtx_range("mtp_dyn.forward") - def forward( + def _forward_impl( self, input_ids, position_ids, From b294868a77bfb688e3399027c63726cae2da18f1 Mon Sep 17 00:00:00 2001 From: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:26:12 -0700 Subject: [PATCH 5/5] [None][feat] Add BaseMultimodalDummyInputsBuilder to minimaxm3_vl (#16696) Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index fce64e1a561c..55ce331d176c 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/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"