tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py (4)

35-37: Type annotation: _out_buffer_flat can be None.

self._out_buffer_flat is initialized to None but annotated as List[torch.Tensor]. This trips type-checkers and can mask bugs.

Apply:

-        self._out_buffer_flat: List[torch.Tensor] = None
+        from typing import Optional
+        self._out_buffer_flat: Optional[List[torch.Tensor]] = None

Optionally assert before use:

-        out_flat = [o_b[:bs].detach().clone() for o_b in self._out_buffer_flat]
+        assert self._out_buffer_flat is not None, "Output buffer not initialized; call capture_graph() first."
+        out_flat = [o_b[:bs].detach().clone() for o_b in self._out_buffer_flat]

---

134-144: Graceful fallback when no captured size ≥ requested bs.

If bs > max(cuda_graph_batch_sizes), round_to_cuda_batch_size returns None and you currently construct shapes containing None. While this works due to the key-miss branch, it's clearer and cheaper to bail out early.

Apply:

-        rounded_shapes = [
-            (self.round_to_cuda_batch_size(input.shape[0]),) + input.shape[1:]
-            for input in args_batched
-        ]
+        rounded_bs = [self.round_to_cuda_batch_size(input.shape[0]) for input in args_batched]
+        if any(rb is None for rb in rounded_bs):
+            return self.model(*args, **kwargs)
+        rounded_shapes = [(rb,) + input.shape[1:] for rb, input in zip(rounded_bs, args_batched)]

---

109-121: Avoid redundant captures when cuda_graph_batch_sizes contain values > max_batch_size.

If a provided batch size exceeds self.max_batch_size, in_buffer[:bs] equals in_buffer[:], leading to duplicate combined_shape keys and redundant capture attempts.

Two options:

• Clamp once in the compiler init (preferred; see next comment).
• Or guard here:

-        for bs in self.cuda_graph_batch_sizes:
+        for bs in self.cuda_graph_batch_sizes:
+            if bs > self.max_batch_size:
+                ad_logger.info(f"Skipping capture for bs={bs} > max_batch_size={self.max_batch_size}")
+                continue

---

190-203: Heuristic helper looks good; add minimal input validation.

Tiny nit: guard invalid inputs and document multiplier behavior for non-divisors of max_bs.

Apply:

     def _get_graph_batch_sizes(
         max_bs: int, extra: Optional[List[int]] = None, multiplier: int = 128
     ) -> List[int]:
-        """Heuristic to set batch sizes for graph capture."""
+        """Heuristic to set batch sizes for graph capture.
+
+        Returns unique sizes in descending order, including 1, max_bs, any `extra`,
+        and all multiples of `multiplier` up to `max_bs`.
+        """
+        assert max_bs >= 1, "max_bs must be >= 1"
+        assert multiplier >= 1, "multiplier must be >= 1"
         # do 1, max_bs, and extra as special batch sizes
         batch_sizes = {1, max_bs, *(extra or [])}

tensorrt_llm/_torch/auto_deploy/compile/backends/torch_compile.py (1)

6-6: Consistency: unify ad_logger import style across backends.

Other backends use relative imports (from ...utils.logger import ad_logger). Consider standardizing on a single absolute style (this one) repo-wide for clarity.

tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py (1)

15-24: Deriving Dynamo recompile_limit from cuda_graph_batch_sizes is sensible; consider headroom.

Setting recompile_limit = max(len(cuda_graph_batch_sizes), existing) avoids premature bailouts, but workloads may still trigger a few extra specializations (e.g., uncommon prompt shapes). Add a small headroom and make it tunable via compiler_kwargs.

Apply:

-        torch._dynamo.config.recompile_limit = max(
-            len(self.cuda_graph_batch_sizes), torch._dynamo.config.recompile_limit
-        )
+        headroom = int(self.compiler_kwargs.get("dynamo_recompile_headroom", 2))
+        target_limit = len(self.cuda_graph_batch_sizes) + max(0, headroom)
+        torch._dynamo.config.recompile_limit = max(target_limit, torch._dynamo.config.recompile_limit)

Keep the informative log (already present).

tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py (1)

590-598: Potential config vs CLI mismatch for max_batch_size.

You set max_batch_size: 128 in YAML, but run_benchmark() is called below without overriding max_batch_size (defaults to 32) and passes --max_batch_size 32 on the CLI. Depending on precedence inside tensorrt_llm.commands.bench, this may cause confusion or unintended overrides.

Consider aligning both sources by passing 128 to run_benchmark here:

-        run_benchmark(model_name, dataset_path, temp_dir)
+        run_benchmark(model_name, dataset_path, temp_dir, max_batch_size=128)

Or drop --max_batch_size for _autodeploy mode inside run_benchmark() and rely solely on YAML to avoid split-brain configuration.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.