[https://nvbugs/6676511][fix] Reject unsupported speculative outputs

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
 

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