fix non-support

Signed-off-by: Enwei Zhu <[email protected]>

Author: syuoni

docs/source/torch/features/feature_combination_matrix.md

@@ -15,4 +15,4 @@
 | KV Cache Reuse             | Yes               | Yes        | Yes                        | Untested              | Yes             | Untested | Yes                       | No                        | Yes           | Yes              | ---            |                        |                       |                 |
 | Slide Window Attention     | Yes               | Yes        | Yes                        | Untested              | No                    | Untested | Untested                  | Untested                  | Yes           | Yes              | WIP            | ---                    |                       |                 |
 | Logits Post Processor      | No                | Yes        | Yes                        | No                    | Untested        | No       | No                        | No                        | Yes           | Yes              | Yes            | Yes                    | ---                   |                 |
-| Guided Decoding            | Yes               | Yes        | Yes                        | No                    | Yes             | No       | No                        | No                        | Yes           | Yes              | Yes            | Yes                    | Yes                   | ---             |
+| Guided Decoding            | Yes               | Yes        | Yes                        | No                    | Yes             | No       | No                        | Yes                       | Yes           | Yes              | Yes            | Yes                    | Yes                   | ---             |

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -933,6 +933,9 @@ def _executor_loop(self):
 
                     self.resource_manager.prepare_resources(scheduled_batch)
                     if self.drafter is not None and self.use_spec_decode:
+                        if self.guided_decoder is not None:
+                            self.guided_decoder.rollback_rejected_tokens(
+                                scheduled_batch)
                         self.drafter.prepare_draft_tokens(
                             scheduled_batch, self.resource_manager)
 

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -331,6 +331,10 @@ def create_py_executor(
             _ExecutorCreationStage.GUIDED_DECODER):
         guided_decoder: Optional[GuidedDecoder] = None
         if executor_config.guided_decoding_config is not None:
+            if spec_config is not None and not has_spec_drafter:
+                raise ValueError(
+                    "Guided decoding is only supported with speculative decoding that has a dedicated drafter (two-model engine)."
+                )
             if mapping.is_last_pp_rank():
                 max_num_draft_tokens = 0
                 if spec_config is not None:

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -340,9 +340,6 @@ def prepare_draft_tokens(
             raise ValueError("Resource manager is required")
 
         try:
-            if self.guided_decoder is not None:
-                self.guided_decoder.rollback_rejected_tokens(scheduled_requests)
-
             draft_batch = self._prepare_draft_batch(scheduled_requests)
 
             if draft_batch.batch_size == 0:

tensorrt_llm/_torch/speculative/ngram.py

@@ -6,7 +6,6 @@
 from tensorrt_llm.llmapi import NGramDecodingConfig
 from tensorrt_llm.logger import logger
 
-from ..pyexecutor.guided_decoder import GuidedDecoder
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState
 from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManager
 from ..pyexecutor.scheduler import ScheduledRequests
@@ -168,22 +167,17 @@ def __init__(
         self,
         spec_config: NGramDecodingConfig,
         ngram_pool_manager: NGramPoolManager = None,
-        guided_decoder: Optional[GuidedDecoder] = None,
     ):
         assert ngram_pool_manager is not None, "NGram needs a resource manager to maintain the pool."
         self.spec_config = spec_config
         self.max_draft_len = spec_config.max_draft_len
         self.spec_resource_manager = ngram_pool_manager
-        self.guided_decoder = guided_decoder
 
     def prepare_draft_tokens(
         self,
         scheduled_requests: ScheduledRequests,
         resource_manager: Optional[ResourceManager] = None,
     ) -> None:
-        if self.guided_decoder is not None:
-            self.guided_decoder.rollback_rejected_tokens(scheduled_requests)
-
         # Disable NGram speculative decoding auto heuristic for batch size > 32.
         if self.spec_config.is_auto_heuristic and len(
                 scheduled_requests.all_requests()) > 32:

tensorrt_llm/_torch/speculative/utils.py

@@ -140,9 +140,7 @@ def get_spec_drafter(model_engine,
                             guided_decoder=guided_decoder)
 
     if spec_config.spec_dec_mode.is_ngram():
-        return NGramDrafter(spec_config,
-                            ngram_pool_manager=spec_resource_manager,
-                            guided_decoder=guided_decoder)
+        return NGramDrafter(spec_config, spec_resource_manager)
 
     return None