tensorrt_llm/_torch/autotuner.py (4)

89-91: Fix inputs_pre_hook docs to match implementation and wrap lines (ruff E501).

The docstring currently says the hook takes (profile, inputs) and is called “before the input tensors are prepared,” but the implementation calls it after _prepare_input_tensors and passes only inputs. Please align the docs with the actual behavior and keep lines ≤120 chars.

Apply:

-        inputs_pre_hook (Callable): A function that takes a profile and a list of input tensors, and returns a list of input tensors.
-            The function is called before the input tensors are prepared for the tuning process to match the real data distribution.
+        inputs_pre_hook (Callable | None): Optional callable invoked after synthetic
+            input tensors are prepared for tuning, but before profiling. It should accept
+            either (inputs) or (inputs, profile) and return a list of input tensors with
+            the same shapes, dtypes, and devices.

---

95-95: Type inputs_pre_hook precisely and make it Optional.

Using a bare Callable loses intent and None is not captured. Recommend Optional with a variadic signature to support both hook forms.

-    inputs_pre_hook: Callable = None
+    inputs_pre_hook: Optional[Callable[..., List[torch.Tensor]]] = None

Additionally, add the missing import at the top of the file:

from typing import Optional

---

453-456: Support both hook signatures and validate outputs to avoid hard-to-debug failures.

Right now we always call hook(inputs) and don’t validate the result. Add support for both (inputs) and (inputs, profile) forms and verify that shapes/dtypes/devices remain unchanged.

-        # If the inputs_pre_hook is provided, it will be called before profiling.
-        if tuning_config.inputs_pre_hook is not None:
-            input_tensors = tuning_config.inputs_pre_hook(input_tensors)
+        # If the inputs_pre_hook is provided, it will be called before profiling.
+        # Support both signatures: hook(inputs) and hook(inputs, profile).
+        if tuning_config.inputs_pre_hook is not None:
+            hook = tuning_config.inputs_pre_hook
+            try:
+                sig = inspect.signature(hook)
+                if len(sig.parameters) >= 2:
+                    new_inputs = hook(input_tensors, profile)
+                else:
+                    new_inputs = hook(input_tensors)
+            except Exception as e:
+                logger.debug(f"[Autotuner] inputs_pre_hook invocation failed: {e}")
+                raise
+            # Validate hook output
+            if not isinstance(new_inputs, list):
+                raise TypeError("inputs_pre_hook must return a List[torch.Tensor]")
+            if len(new_inputs) != len(input_tensors):
+                raise ValueError("inputs_pre_hook must return the same number of tensors")
+            for old, new in zip(input_tensors, new_inputs):
+                if not isinstance(new, torch.Tensor):
+                    raise TypeError("inputs_pre_hook must return only torch.Tensor elements")
+                if old.size() != new.size() or old.dtype != new.dtype or old.device != new.device:
+                    raise ValueError(
+                        "inputs_pre_hook must preserve shapes, dtype and device for all inputs"
+                    )
+            input_tensors = new_inputs

---

1-1: Missing NVIDIA copyright header.

Per coding guidelines, prepend the current-year NVIDIA copyright header.

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

tests/unittest/_torch/test_autotuner.py (3)

341-349: Avoid mutable default for tactic dict.

Using {} as a default is a Python footgun. Use None and normalize inside the function; keep -1 behavior intact.

-        tactic: dict = {},
+        tactic: Dict[str, int] = None,
     ) -> torch.Tensor:
         # Notice that in fallback case tactic is -1
-        if tactic == -1:
+        if tactic in (-1, None):
             # assign default configs for fallback case
             block_size, tactic_id = 128, -1
         else:
             block_size, tactic_id = tactic["block_size"], tactic["tactic_id"]

---

352-359: Clarify pre-hook and avoid implicit dtype cast.

Minor cleanup: add return type, and set the sentinel value with the correct dtype/device explicitly to avoid implicit int→float cast.

     @staticmethod
-    def inputs_pre_hook(inputs: List[torch.Tensor]):
-        # always set the first element to bo iota in x
+    def inputs_pre_hook(inputs: List[torch.Tensor]) -> List[torch.Tensor]:
+        # Set the last row/first column of x to the current token count
         x, w = inputs
         x_hooked = torch.zeros_like(x)
-        x_hooked[-1, 0] = x.shape[0]
+        x_hooked[-1, 0] = torch.tensor(x.shape[0], dtype=x.dtype, device=x.device)
         return [x_hooked, w]

---

1-1: Missing NVIDIA copyright header.

Per coding guidelines, prepend the current-year NVIDIA copyright header.

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

Configuration used: .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.

tensorrt_llm/_torch/autotuner.py (1)

89-90: Docstring/API mismatch and >120-char lines (Ruff E501) — clarify signature and wrap lines

Docstring states the hook takes (profile, inputs), but implementation currently passes only inputs. Also both lines exceed 120 chars.

Apply this diff to (a) document the intended signature clearly and (b) fix line length:

-        inputs_pre_hook (Callable): A function that takes a profile and a list of input tensors, and returns a list of input tensors.
-            The function is called before the input tensors are prepared for the tuning process to match the real data distribution.
+        inputs_pre_hook (Callable): A function that takes (profile, inputs) and returns a new list
+            of input tensors. It is invoked before preparing the tuning tensors to better match
+            real data distributions.

tests/unittest/_torch/test_autotuner.py (1)

323-329: This assertion will fail until _optimization_profiles is gated by tune_max_num_tokens (duplicate of prior review)

_optimization_profiles currently unconditionally adds the mapped current input size to opt_shapes, which introduces 64 when tune_max_num_tokens=32. That will make inputs[0].shape[0] == 64 for some profiles, tripping this assertion.

Proposed fix in tensorrt_llm/_torch/autotuner.py (context ~Lines 605–611), gate the addition:

# Add the current input value as one of the opt values, but don't exceed the cap.
opt_shapes = set(opt_shapes)
mapped_val = spec.map_to_tuning_buckets(
    base_profile.shapes[spec.input_idx][spec.dim_idx].val
)
if tuning_config.tune_max_num_tokens is None or mapped_val <= tuning_config.tune_max_num_tokens:
    opt_shapes.add(mapped_val)
opt_shapes = sorted(list(opt_shapes))

Once applied, this test assertion should pass consistently.

tensorrt_llm/_torch/autotuner.py (1)

95-95: Type the new field as optional and narrow the callable’s contract

Typing this as Optional and clarifying the callable’s shape improves readability and tooling. If you keep supporting only (inputs)->inputs for now, please reflect that in types.

Suggested change (requires adding Optional to typing imports at Line 8):

-    inputs_pre_hook: Callable = None
+    # Accepts either (profile, inputs) -> inputs or (inputs) -> inputs
+    inputs_pre_hook: Optional[Callable] = None

Optionally, narrow further to Callable[..., List[torch.Tensor]] to reflect return type.

tests/unittest/_torch/test_autotuner.py (3)

341-349: Avoid mutable default argument for tactic; accept None and handle fallback cleanly

tactic: dict = {} uses a mutable default and can lead to subtle bugs. Use None and treat None like fallback (-1). Optionally widen the type to allow either -1 or a dict.

Apply this diff (add Optional to imports at the top):

-    def forward(
+    def forward(
         self,
         /,
         inputs: List[torch.Tensor],
         *,
-        tactic: dict = {},
+        tactic: dict = None,
     ) -> torch.Tensor:
-        # Notice that in fallback case tactic is -1
-        if tactic == -1:
+        # Notice that in fallback case tactic is -1 or tactic is None
+        if tactic in (-1, None):
             # assign default configs for fallback case
             block_size, tactic_id = 128, -1
         else:
             block_size, tactic_id = tactic["block_size"], tactic["tactic_id"]

If you prefer stricter typing:

• from typing import Optional, Union
• tactic: Union[int, Dict[str, int]] = -1
• if tactic == -1: ...

---

352-359: Align test pre-hook signature with documented API (profile, inputs)

To match the TuningConfig doc and the suggested change in autotuner, accept profile as the first arg. This keeps tests representative of intended usage.

Apply this diff:

     @staticmethod
-    def inputs_pre_hook(inputs: List[torch.Tensor]):
-        # always set the first element to bo iota in x
-        x, w = inputs
+    def inputs_pre_hook(profile: OptimizationProfile, inputs: List[torch.Tensor]):
+        # always set the first element to an identifying value in x
+        x, w = inputs
         x_hooked = torch.zeros_like(x)
         x_hooked[-1, 0] = x.shape[0]
         return [x_hooked, w]

---

354-355: Clarify comment wording (“bo iota”)

“bo iota” is unclear and looks like a typo. Consider clarifying the intent (e.g., “set a sentinel value based on sequence length”).

Suggested:

• “Set a sentinel value in x based on the sequence length to verify the pre-hook ran.”

Configuration used: .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.

tensorrt_llm/_torch/autotuner.py (1)

453-456: Pass profile into inputs_pre_hook (per API/intent) and validate return; optionally support both signatures

Currently the hook is called with only inputs, diverging from the documented API and limiting usefulness. Pass the profile too and validate the return value. For backward compatibility you can fall back to the 1-arg call.

Apply this diff:

-        # If the inputs_pre_hook is provided, it will be called before profiling.
-        if tuning_config.inputs_pre_hook is not None:
-            input_tensors = tuning_config.inputs_pre_hook(input_tensors)
+        # If provided, preprocess inputs before profiling to better match real data distributions.
+        if tuning_config.inputs_pre_hook is not None:
+            hook = tuning_config.inputs_pre_hook
+            # Prefer the documented signature (profile, inputs); fall back to (inputs) for compat.
+            try:
+                input_tensors = hook(profile, input_tensors)
+            except TypeError:
+                input_tensors = hook(input_tensors)
+            # Basic validation: ensure we got back a list with the same arity.
+            if not isinstance(input_tensors, list) or len(input_tensors) != len(profile.shapes):
+                raise ValueError(
+                    "inputs_pre_hook must return a list of tensors with the same arity as inputs"
+                )

tests/unittest/_torch/test_autotuner.py (1)

322-329: Fix: don't add base (current-input) opt value when it exceeds tuning_config.tune_max_num_tokens

The autotuner unconditionally adds the mapped current input size into opt_shapes (tensorrt_llm/_torch/autotuner.py ~lines 605–610). When tuning_config.tune_max_num_tokens is set (e.g. 32) that can introduce opt values > tune_max_num_tokens (e.g. 64) and cause the test assertion in tests/unittest/_torch/test_autotuner.py (lines ~322–329) to fail.

The assertion should pass after applying the suggested fix in autotuner.py to gate the addition of the base shape by tune_max_num_tokens.

#!/bin/bash
# Description: Verify that _optimization_profiles properly respects tune_max_num_tokens
# Expected: The mapped current input size should only be added if it's <= tune_max_num_tokens

# Search for the problematic code in _optimization_profiles
rg -n -A 5 -B 5 'opt_shapes\.add' tensorrt_llm/_torch/autotuner.py

tensorrt_llm/_torch/autotuner.py (1)

89-90: Fix line length violations.

Lines 89-90 exceed the 120-character limit. Consider breaking these lines appropriately.

Apply this diff to fix the line length issues:

-        inputs_pre_hook (Callable): A function that takes a list of input tensors, returns a list of modified input tensors.
-            It is called before the input tensors are prepared for the tuning process to match the real data distribution.
+        inputs_pre_hook (Callable): A function that takes a list of input tensors, 
+            returns a list of modified input tensors. It is called before the input
+            tensors are prepared for the tuning process to match the real data distribution.

Configuration used: .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.