[None][perf] Autotune TRT-LLM Gen MoE when using CUDA graphs

jinyangyuan-nvidia · jinyangyuan-nvidia · commit 6fd77b08a27c · 2025-08-27T00:40:49.000-07:00
Signed-off-by: Jinyang Yuan &lt;154768711+jinyangyuan-nvidia@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -778,6 +778,21 @@ def release_batch(result: ScheduledRequests | None):
                 continue
 
             for draft_len in draft_lengths:
+                with release_batch(get_cuda_graph_warmup_request(
+                        bs, draft_len)) as batch:
+                    if batch is None:
+                        # No KV cache space!
+                        return
+                    logger.info(
+                        f"Run generation only CUDA graph autotune for batch size={bs}, draft_len={draft_len}"
+                    )
+                    self.enable_spec_decode = draft_len > 0 or self.is_draft_model
+                    with self.no_cuda_graph(), autotune():
+                        self.forward(batch,
+                                     new_tensors_device=None,
+                                     resource_manager=resource_manager)
+                    torch.cuda.synchronize()
+
                 with release_batch(get_cuda_graph_warmup_request(
                         bs, draft_len)) as batch:
                     if batch is None:
