From 39adea61b71999e938d7373ec3d491a7ad7d3ba2 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Fri, 28 Aug 2026 17:57:33 -0700
Subject: [PATCH] [https://nvbugs/6676511][fix] Reject unsupported speculative
 outputs

* Why?

Speculative decoding does not support context or generation logits and
their derived log probabilities. Allowing these requests to reach the
executor can crash forked postprocessing workers and leave serving
requests permanently hung.

* What?

Reject unsupported logits and log-probability output options at the LLM
request boundary whenever speculative decoding is configured. Preserve
these options for non-speculative requests and treat top-0 log-probability
requests as explicitly requested outputs.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 .../_torch/speculative/spec_sampler_base.py   | 39 +++++++++----------
 tests/unittest/llmapi/test_sampling_params.py | 22 ++++++++---
 2 files changed, 34 insertions(+), 27 deletions(-)

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