diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 1c2d8015be6b..683bca121dd4 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,8 +27,6 @@ import torch -from tensorrt_llm.logger import logger - from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.sampler import ( @@ -98,9 +96,25 @@ def validate_request(self, request: LlmRequest) -> None: buffer there, so it would be silently dropped and the request would decode from a different distribution than the user asked for. Threading it through costs measurable throughput on the rejection path, so reject - instead. Raised from validate_request (request admission), so only the - offending request fails rather than the whole executor step. + instead. This sampler also does not return context logits, generation + logits, or log probabilities. Raised from validate_request (request + admission), so only the offending request fails rather than the whole + executor step. """ + requested_outputs = ( + ("return_context_logits / prompt_logprobs", request.py_return_context_logits), + ("return_generation_logits", request.py_return_generation_logits), + ("logprobs", request.py_return_log_probs), + ) + unsupported_outputs = [name for name, requested in requested_outputs if requested] + if unsupported_outputs: + raise ValueError( + "The following output options are not supported with " + "one-model speculative decoding: " + f"{', '.join(unsupported_outputs)}. Drop these options from " + "the request, or disable speculative decoding." + ) + sampling_config = request.sampling_config if sampling_config is None: return @@ -264,23 +278,6 @@ def _request_common_handling( runtime_draft_len: Optional[int], ) -> None: """Common handling for both context and generation requests.""" - if request.py_return_context_logits: - logger.warning( - "return_context_logits not supported with speculative decoding, " - "skipping for request %s", - request.py_request_id, - ) - if request.py_return_generation_logits: - logger.warning( - "return_generation_logits not supported with speculative decoding, " - "skipping for request %s", - request.py_request_id, - ) - if request.py_return_log_probs: - logger.warning( - "return_log_probs not supported with speculative decoding, skipping for request %s", - request.py_request_id, - ) request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py index 889a8abcee9a..3c42fc097cf5 100644 --- a/tests/unittest/llmapi/test_sampling_params.py +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -14,12 +14,13 @@ # limitations under the License. import asyncio import json -from types import SimpleNamespace +from typing import Any import pytest import torch from tensorrt_llm.llmapi.llm import BaseLLM +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.thinking_budget import ( ThinkingBudgetLogitsProcessor, add_thinking_budget_logits_processor, @@ -37,16 +38,25 @@ pytestmark = pytest.mark.cpu_only +# BaseLLM.__init__ builds a model and executor. This lightweight subclass keeps +# these method-level tests CPU-only while using the production argument types. +class _TestLLM(BaseLLM): + def __init__(self, **args_overrides: Any) -> None: + self.args = TorchLlmArgs( + model="dummy", + skip_tokenizer_init=True, + **args_overrides, + ) + + def _apply_generation_config_sampling_defaults( mode: str, sampling_params: SamplingParams, generation_config_explicit_values: dict, ) -> SamplingParams: - llm = SimpleNamespace( - args=SimpleNamespace(backend="pytorch", generation_config=mode), - _generation_config_explicit_values=generation_config_explicit_values, - ) - BaseLLM._apply_generation_config_sampling_defaults(llm, sampling_params) + llm = _TestLLM(generation_config=mode) + llm._generation_config_explicit_values = generation_config_explicit_values + llm._apply_generation_config_sampling_defaults(sampling_params) return sampling_params