[None][feat] Optimize CUDA graph memory usage for spec decode cases (NVIDIA#6718)

Signed-off-by: Mike Iovine <[email protected]>

Author: mikeiovine

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -726,8 +726,11 @@ def disable_optimization(backend: Backend):
             # For non-draft model, we also capture the CUDA graph instance for draft length 0,
             # so that when we disable spec decode at runtime, we can still run the captured graph.
             # Note that for one engine mode, we are not able to turn off spec decode at runtime.
-            if not self.is_draft_model and self.max_draft_len > 0 and not self.spec_config.spec_dec_mode.use_one_engine(
-            ):
+            if (not self.is_draft_model and self.max_draft_len > 0
+                    and not self.spec_config.spec_dec_mode.use_one_engine()
+                    # Assume that speculation is always on if the user didn't give us a max_concurrency
+                    # value. This will save on memory.
+                    and self.spec_config.max_concurrency is not None):
                 draft_lengths.append(0)
 
             for bs in cuda_graph_batch_sizes:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -937,11 +937,30 @@ def _executor_loop(self):
                         self.guided_decoder.init_disagg_gen_requests(
                             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)
+                        # When running with an external drafter, only TP rank 0 sends request to drafter
+                        if self.dist.tp_size > 1 and getattr(
+                                self.drafter, 'single_draft_call',
+                                lambda: False)():
+                            if self.dist.rank == 0:
+                                self.drafter.prepare_draft_tokens(
+                                    scheduled_batch, self.resource_manager)
+                                draft_data = {}
+                                for req in scheduled_batch.generation_requests:
+                                    draft_data[
+                                        req.py_request_id] = req.py_draft_tokens
+                                self.dist.tp_broadcast(draft_data, root=0)
+                            else:
+                                draft_data = self.dist.tp_broadcast(None,
+                                                                    root=0)
+                                for req in scheduled_batch.generation_requests:
+                                    req.py_draft_tokens = draft_data[
+                                        req.py_request_id]
+                        else:
+                            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)
 
                     batch_outputs = self._forward_step(scheduled_batch)
                     self._execute_guided_decoder(scheduled_batch,

tensorrt_llm/_torch/speculative/drafter.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import List, Optional, final
 
 from ..pyexecutor.llm_request import LlmRequest
 from ..pyexecutor.resource_manager import ResourceManager
@@ -26,8 +26,13 @@ def prepare_draft_tokens(
         """
         raise NotImplementedError
 
+    @final
     def should_use_spec_decode(self, requests: List[LlmRequest]) -> bool:
-        """Check if spec decode should be used for the current iteration."""
+        """
+        You probably don't want to override this. ModelEngine
+        assumes that speculation is always on if max_concurrency
+        is not specified by the user's spec config.
+        """
         if self.max_concurrency is not None:
             return len(requests) <= self.max_concurrency
         return True

tensorrt_llm/_torch/speculative/external_api.py

@@ -0,0 +1,182 @@
+import asyncio
+import json
+from typing import List
+
+import aiohttp
+
+from tensorrt_llm.logger import logger
+
+from ..pyexecutor.llm_request import *
+from ..pyexecutor.scheduler import ScheduledRequests
+from .drafter import Drafter
+
+
+class APIDrafter(Drafter):
+
+    def __init__(
+        self,
+        spec_config: "ExternalAPIConfig",
+    ):
+        super().__init__()
+        self.max_draft_len = spec_config.max_draft_len
+        self.endpoint = spec_config.endpoint
+        assert self.endpoint is not None, "API endpoint is required for external API speculative decoding."
+        self.template = spec_config.template if spec_config.template is not None else {}
+        self.response_field = spec_config.response_field if spec_config.response_field is not None else "draft_tokens"
+
+    def single_draft_call(self):
+        return True
+
+    def get_nested_field_from_response(self, response: dict) -> List[int]:
+        # Allows for nested fields in the response.
+        # Example: "choices.0.message.content"
+        # Returns the value of the nested field: response["choices"][0]["message"]["content"]
+        keys = self.response_field.split(".")
+        current = response
+
+        for key in keys:
+            try:
+                if key.isdigit():
+                    key = int(key)
+                    if isinstance(current, list) and 0 <= key < len(current):
+                        current = current[key]
+                    else:
+                        logger.warning(
+                            f"Response field {self.response_field} is invalid for response {response}. Index {key} is invalid."
+                        )
+                        return []
+                else:
+                    if isinstance(current, dict) and key in current:
+                        current = current[key]
+                    else:
+                        logger.warning(
+                            f"Response field {self.response_field} is invalid for response {response}. Index {key} is invalid."
+                        )
+                        return []
+
+            except (KeyError, ValueError, IndexError):
+                logger.warning(
+                    f"Response field path is invalid: {self.response_field}")
+                return []
+
+        if not isinstance(current, list):
+            logger.warning(
+                f"API response '{self.response_field}' must be a list. Got type: {type(current)}"
+            )
+            return []
+        return current
+
+    async def get_draft_tokens(
+        self,
+        prefix: list[int],
+        request_id: int,
+        end_id: int,
+        max_sequence_length: int,
+    ) -> List[int]:
+        try:
+            request_data = {
+                "prefix": prefix,
+                "request_id": request_id,
+                "end_id": end_id,
+                "max_sequence_length": max_sequence_length,
+            }
+            if self.template:
+                request_data.update(self.template)
+
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                        url=self.endpoint,
+                        json=request_data,
+                        headers={"Content-Type": "application/json"},
+                        timeout=aiohttp.ClientTimeout(total=10),
+                ) as response:
+
+                    # check for unsuccessful response
+                    if response.status != 200:
+                        logger.error(
+                            f"Failed to get draft tokens. API call failed for request {request_id} with status code {response.status}"
+                        )
+                        return []
+
+                    result = await response.json()
+                    draft_tokens = self.get_nested_field_from_response(result)
+                    if len(draft_tokens) > self.max_draft_len:
+                        draft_tokens = draft_tokens[:self.max_draft_len]
+                    logger.debug(
+                        f"Retrieved draft tokens for request {request_id}: {draft_tokens}"
+                    )
+                    return draft_tokens
+
+        except json.JSONDecodeError as e:
+            logger.error(
+                f"Failed to parse JSON response for request {request_id}: {e}")
+            return []
+
+        except Exception as e:
+            logger.error(
+                f"Failed to get draft tokens. API call failed for request {request_id} with the following error: {e}"
+            )
+            return []
+
+    async def async_prepare_draft_tokens(
+        self,
+        scheduled_requests: ScheduledRequests,
+        resource_manager: None,
+    ) -> None:
+        # Sort by request_id when py_batch_idx is None as a fallback.
+        # This happens in the disagg case: for a set of new requests, we draft
+        # before forward_step, so py_batch_idx is not assigned.
+        sorted_requests = sorted(
+            scheduled_requests.generation_requests,
+            key=lambda r:
+            (r.py_batch_idx is None, r.py_batch_idx or r.request_id),
+        )
+
+        tasks = []
+        for request in sorted_requests:
+            # Add new token to a copy of the generated tokens to find new draft tokens
+            prefix = list(request.get_tokens()[0])  # Get a copy
+            task = self.get_draft_tokens(
+                prefix,
+                request.request_id,
+                request.py_end_id,
+                request.py_orig_prompt_len + request.py_max_new_tokens,
+            )
+            tasks.append(task)
+
+        try:
+            all_draft_tokens = await asyncio.wait_for(asyncio.gather(
+                *tasks, return_exceptions=True),
+                                                      timeout=10.0)
+        except asyncio.TimeoutError:
+            logger.error(
+                f"Timeout occurred while getting draft tokens for batch of requests"
+            )
+            all_draft_tokens = [[] for _ in tasks]
+
+        for request, draft_tokens in zip(sorted_requests, all_draft_tokens):
+            if isinstance(draft_tokens, Exception):
+                logger.error(
+                    f"An exception occurred while getting draft tokens for request {request.request_id}. Set TLLM_LOG_LEVEL for more details."
+                )
+                draft_tokens = []
+            elif len(draft_tokens) == 0:
+                logger.error(
+                    f"Draft tokens could not be generated for request {request.request_id}. Set TLLM_LOG_LEVEL for more details."
+                )
+            else:
+                # Pad length to `self.max_draft_len`
+                if len(draft_tokens) > 0:
+                    pad_length = self.max_draft_len - len(draft_tokens)
+                    draft_tokens.extend([request.py_end_id] * pad_length)
+
+            request.py_draft_tokens = draft_tokens
+
+    def prepare_draft_tokens(
+        self,
+        scheduled_requests: ScheduledRequests,
+        resource_manager: None,
+    ) -> None:
+        asyncio.run(
+            self.async_prepare_draft_tokens(scheduled_requests,
+                                            resource_manager))

tensorrt_llm/_torch/speculative/interface.py

@@ -17,6 +17,7 @@ class SpeculativeDecodingMode(IntEnum):
     NGRAM = auto()
     DRAFT_TARGET = auto()
     USER_PROVIDED = auto()
+    EXTERNAL_API = auto()
     NONE = auto()
     AUTO = auto()
 
@@ -44,6 +45,9 @@ def is_user_provided(self):
     def is_none(self):
         return self == SpeculativeDecodingMode.NONE
 
+    def is_external_api(self):
+        return self == SpeculativeDecodingMode.EXTERNAL_API
+
     def is_draft_target(self):
         return self == SpeculativeDecodingMode.DRAFT_TARGET
 
@@ -79,7 +83,7 @@ def has_spec_decoder(self):
 
     def has_spec_drafter(self):
         return self.is_eagle3() or self.is_draft_target() or self.is_ngram(
-        ) or self.is_user_provided()
+        ) or self.is_user_provided() or self.is_external_api()
 
     def extend_ctx(self, attention_backend: Type[AttentionBackend]):
         """
@@ -91,8 +95,8 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]):
         # Fixme: only trtllm attention backend supports eagle3 generation-phase kernels on blackwell.
         return ((self.is_eagle3() or self.is_draft_target())
                 and not (issubclass(attention_backend, TrtllmAttention)
-                         and get_sm_version() == 100)
-                ) or self.is_ngram() or self.is_user_provided()
+                         and get_sm_version() == 100)) or self.is_ngram(
+                         ) or self.is_user_provided() or self.is_external_api()
 
     def attention_need_spec_dec_mode(self):
         """

tensorrt_llm/_torch/speculative/utils.py

@@ -7,6 +7,7 @@
 from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata,
                      Eagle3OneModelWorker, Eagle3ResourceManager,
                      Eagle3SpecMetadata)
+from .external_api import APIDrafter
 from .model_drafter import ModelDrafter
 from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler,
                   MTPSpecMetadata, MTPWorker)
@@ -50,7 +51,8 @@ def get_spec_metadata(spec_config,
         )
     if  spec_config.spec_dec_mode.is_draft_target() or \
         spec_config.spec_dec_mode.is_ngram() or \
-        spec_config.spec_dec_mode.is_user_provided():
+        spec_config.spec_dec_mode.is_user_provided() or \
+        spec_config.spec_dec_mode.is_external_api():
         return SpecMetadata(
             max_draft_len=spec_config.max_draft_len,
             spec_dec_mode=spec_config.spec_dec_mode,
@@ -99,6 +101,8 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None):
         return NGramPoolManager(spec_config, max_num_requests)
     if spec_dec_mode.is_user_provided():
         return spec_config.resource_manager
+    if spec_dec_mode.is_external_api():
+        return None
     return None
 
 
@@ -142,6 +146,9 @@ def get_spec_drafter(model_engine,
     if spec_config.spec_dec_mode.is_ngram():
         return NGramDrafter(spec_config, spec_resource_manager)
 
+    if spec_config.spec_dec_mode.is_external_api():
+        return APIDrafter(spec_config)
+
     return None
 
 

tensorrt_llm/llmapi/__init__.py

@@ -9,11 +9,11 @@
                        CapacitySchedulerPolicy, ContextChunkingPolicy,
                        CudaGraphConfig, DraftTargetDecodingConfig,
                        DynamicBatchConfig, EagleDecodingConfig,
-                       ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs,
-                       LookaheadDecodingConfig, MedusaDecodingConfig, MoeConfig,
-                       MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig,
-                       TorchCompileConfig, TorchLlmArgs, TrtLlmArgs,
-                       UserProvidedDecodingConfig)
+                       ExtendedRuntimePerfKnobConfig, ExternalAPIConfig,
+                       KvCacheConfig, LlmArgs, LookaheadDecodingConfig,
+                       MedusaDecodingConfig, MoeConfig, MTPDecodingConfig,
+                       NGramDecodingConfig, SchedulerConfig, TorchCompileConfig,
+                       TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig)
 from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo,
                         QuantConfig)
 from .mpi_session import MpiCommSession
@@ -49,6 +49,7 @@
     'CacheTransceiverConfig',
     'NGramDecodingConfig',
     'UserProvidedDecodingConfig',
+    'ExternalAPIConfig',
     'TorchCompileConfig',
     'DraftTargetDecodingConfig',
     'LlmArgs',