Skip to content

[https://nvbugs/5534705][fix] Skip unnecessary CUDA graph capture #8050

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 11, 2025
Merged

[https://nvbugs/5534705][fix] Skip unnecessary CUDA graph capture #8050

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f9d0d72
[https://nvbugs/5534705][fix] Skip draft_model forward when draft_len…
ziyixiong-nv Sep 29, 2025
ebad62c
A minor fix to limit CDL to eagle3
ziyixiong-nv Oct 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 20 additions & 13 deletions tensorrt_llm/_torch/pyexecutor/model_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -695,19 +695,25 @@ def release_batch(result: ScheduledRequests | None):
cuda_graph_batch_sizes = sorted(self._cuda_graph_batch_sizes,
reverse=True)
# Create CUDA graphs for different draft lengths
draft_lengths = [self.max_draft_len]
# 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()
# 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)
if self.is_spec_decode and self.is_draft_model and spec_resource_manager is not None and isinstance(
spec_resource_manager, Eagle3ResourceManager):
draft_lengths.append(self.original_max_draft_len)
draft_lengths = []
if self.is_draft_model:
if self.model_is_wrapped and self.is_spec_decode and spec_resource_manager is not None and isinstance(
spec_resource_manager, Eagle3ResourceManager):
# The CDL path uses draft_len > 0 for the number of iterations in the drafting loop.
draft_lengths.append(self.original_max_draft_len)
else:
draft_lengths.append(self.max_draft_len)
else:
# 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 (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)
draft_lengths = [self.max_draft_len]

for bs in cuda_graph_batch_sizes:
if bs > self.batch_size:
Expand All @@ -723,6 +729,7 @@ def release_batch(result: ScheduledRequests | None):
logger.info(
f"Run generation only CUDA graph warmup for batch size={bs}, draft_len={draft_len}"
)

self.enable_spec_decode = draft_len > 0 or self.is_draft_model

def _update_draft_inference_state(is_first_draft: bool,
Expand Down
Loading