Fix for CUDA graph padding

Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>

Author: ziyixiong-nv

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -16,6 +16,7 @@
 from ..modules.multi_stream_utils import with_multi_stream
 from ..speculative.eagle3 import Eagle3ResourceManager
 from ..speculative.mtp import SampleStateTensorsMTP
+from ..speculative.utils import get_draft_kv_cache_manager
 from ..utils import make_weak_ref, piecewise_cuda_graph
 from .llm_request import get_draft_token_length
 from .mamba_cache_manager import MambaCacheManager
@@ -439,12 +440,19 @@ def _get_padded_batch(self, batch: ScheduledRequests,
             if available_blocks < 1:
                 return 0
 
+            # Get draft KV cache manager only for one-model speculative decoding.
+            # In two-model mode, each model has its own KV cache manager, so
+            # draft_kv_cache_manager should be None.
+            draft_kv_cache_manager = get_draft_kv_cache_manager(
+                self.spec_config, resource_manager)
+
             self.padding_dummy_request = kv_cache_manager.add_dummy_requests(
                 [CUDA_GRAPH_DUMMY_REQUEST_ID],
                 is_gen=True,
                 max_num_draft_tokens=runtime_draft_len,
                 use_mrope=self.config.use_mrope,
-                max_beam_width=self.config.max_beam_width)[0]
+                max_beam_width=self.config.max_beam_width,
+                draft_kv_cache_manager=draft_kv_cache_manager)[0]
             self.padding_dummy_request.is_cuda_graph_dummy = True
             spec_res_mgr = resource_manager.get_resource_manager(
                 ResourceManagerType.SPEC_RESOURCE_MANAGER)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -44,8 +44,8 @@
 from ..models.modeling_utils import DecoderModelForCausalLM
 from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer,
                                                    MoeLoadBalancerIterContext)
-from ..speculative import (SpecMetadata, get_num_extra_kv_tokens,
-                           get_spec_metadata,
+from ..speculative import (SpecMetadata, get_draft_kv_cache_manager,
+                           get_num_extra_kv_tokens, get_spec_metadata,
                            update_spec_config_from_model_config)
 from ..speculative.drafting_loops import BaseDraftingLoopWrapper
 from ..speculative.eagle3 import (Eagle3OneModelSpecMetadata,
@@ -550,12 +550,7 @@ def _get_draft_kv_cache_manager(
         Returns the draft KV cache manager only in one-model speculative decoding
         mode where the target model manages a separate draft KV cache.
         """
-        if self.spec_config is None:
-            return None
-        if not self.spec_config.spec_dec_mode.use_one_engine():
-            return None
-        return resource_manager.get_resource_manager(
-            ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+        return get_draft_kv_cache_manager(self.spec_config, resource_manager)
 
     @contextmanager
     def set_warmup_flag(self):

tensorrt_llm/_torch/speculative/utils.py

@@ -247,6 +247,21 @@ def get_num_extra_kv_tokens(spec_config):
     return 0
 
 
+def get_draft_kv_cache_manager(spec_config, resource_manager):
+    """
+    Returns the draft KV cache manager only in one-model speculative decoding
+    mode where the target model manages a separate draft KV cache.
+    """
+    from ..pyexecutor.resource_manager import ResourceManagerType
+
+    if spec_config is None:
+        return None
+    if not spec_config.spec_dec_mode.use_one_engine():
+        return None
+    return resource_manager.get_resource_manager(
+        ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+
+
 def update_spec_config_from_model_config(spec_config, model_config):
     if spec_config.spec_dec_mode.is_mtp_one_model():
         # Use `max_draft_len` for several low-level APIs. TODO: Remove this after distinguishing them.