tensorrt_llm/_torch/auto_deploy/llm_args.py (2)

197-203: Align field description with enforced behavior.

Clarify defaulting vs. requirement for triton/torch backends.

Apply:

-        description="Page size for attention (tokens_per_block). For triton and torch "
-        "backends, this should equal max_seq_len. Temporary field until tokens_per_block gets "
-        "properly passed through.",
+        description=(
+            "Page size for attention (tokens_per_block). For triton and torch backends, "
+            "this must equal max_seq_len. If unset, it defaults to max_seq_len. Temporary "
+            "field until tokens_per_block is plumbed end-to-end."
+        ),

---

206-213: Gate override of attn_page_size and error on mismatch

Only force-set attn_page_size to max_seq_len for "triton"/"torch" when not explicitly provided; otherwise raise if it doesn’t match:

 tensorrt_llm/_torch/auto_deploy/llm_args.py
 @model_validator(mode="after")
 # TODO: discuss what to do with this once we fully transition to the new inference optimizer
 def update_attn_page_size(self):
-    # NOTE force attn_page_size to equal max_seq_len for triton backend
-    # TODO: maybe don't do this and rely on slot_idx instead??
-    if self.attn_backend == "triton" or self.attn_backend == "torch":
-        self.attn_page_size = self.max_seq_len
+    # Keep attn_page_size == max_seq_len for backends that require fixed tokens_per_block,
+    # but do not silently override explicit user input.
+    if self.attn_backend in ("triton", "torch"):
+        if "attn_page_size" not in self.model_fields_set:
+            self.attn_page_size = self.max_seq_len
+        elif self.attn_page_size != self.max_seq_len:
+            raise ValueError(
+                "attn_page_size must equal max_seq_len for attn_backend in {'triton','torch'}."
+            )
     return self

All existing tests (including default flashinfer/triton logic) remain green under this change.

tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py (3)

52-53: free_mem_ratio: clarify intent and use 1e-4 literal.

Tiny ratio is fine for CI; add a brief comment and prefer scientific literal for readability.

             compile_backend="torch-opt",
-            free_mem_ratio=0.0001,
+            # Keep memory footprint minimal for CI to avoid OOM on single-GPU runners.
+            free_mem_ratio=1e-4,

---

98-102: Guard Nemotron-Nano (flashinfer) with a skip-if when flashinfer is unavailable.

Prevents spurious CI failures on environments without flashinfer.

         get_small_model_config_pytest_param(
             "nvidia/NVIDIA-Nemotron-Nano-12B-v2",
+            pytest_param_kwargs={
+                "marks": pytest.mark.skipif(
+                    not HAS_FLASHINFER, reason="flashinfer not available"
+                )
+            },
             attn_backend="flashinfer",
             compile_backend="torch-simple",
         ),

Add the following helper near the imports (once per file):

# at top-level, after imports
try:
    import flashinfer  # noqa: F401
    HAS_FLASHINFER = True
except Exception:
    HAS_FLASHINFER = False

Consider applying the same skip-if to other flashinfer-backed params in this file for consistency.
Based on learnings

---

106-112: Clarify operator precedence and skip semantics (transformers-only vs always).

Current expression relies on and > or and yields:

• Always skip DeepSeek-V3 and Phi-3 (both modes).
• Skip Nemotron-Nano only in transformers mode.
  If that’s intended, make it explicit and fix the reason string; else, apply mode == "transformers" to all three.

Recommended (mode-gated for all three):

-    if (
-        "DeepSeek-V3" in experiment_config["args"]["model"]
-        or "Phi-3-mini-4k-instruct" in experiment_config["args"]["model"]
-        or "NVIDIA-Nemotron-Nano-12B-v2" in experiment_config["args"]["model"]
-        and mode == "transformers"
-    ):
+    model_id = experiment_config["args"]["model"]
+    if mode == "transformers" and any(
+        name in model_id
+        for name in ("DeepSeek-V3", "Phi-3-mini-4k-instruct", "NVIDIA-Nemotron-Nano-12B-v2")
+    ):
         pytest.skip(f"{experiment_config['args']['model']} is not supported in transformers mode")

Alternative (keep current semantics but avoid precedence pitfalls and fix reason when mode != "transformers"):

-    if (
-        "DeepSeek-V3" in experiment_config["args"]["model"]
-        or "Phi-3-mini-4k-instruct" in experiment_config["args"]["model"]
-        or "NVIDIA-Nemotron-Nano-12B-v2" in experiment_config["args"]["model"]
-        and mode == "transformers"
-    ):
-        pytest.skip(f"{experiment_config['args']['model']} is not supported in transformers mode")
+    model_id = experiment_config["args"]["model"]
+    if (
+        "DeepSeek-V3" in model_id
+        or "Phi-3-mini-4k-instruct" in model_id
+        or ("NVIDIA-Nemotron-Nano-12B-v2" in model_id and mode == "transformers")
+    ):
+        reason = "transformers mode" if mode == "transformers" else "this test configuration"
+        pytest.skip(f"{model_id} is not supported in {reason}")

Please confirm the intended behavior.

tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py (2)

286-294: Run forward in inference_mode and sync to stabilize memory readings.

Prevents grad allocations and ensures mem_get_info sees settled allocations.

-        self._run_forward(gm, cm)
+        with torch.inference_mode():
+            self._run_forward(gm, cm)
+        torch.cuda.synchronize()

---

10-12: Good switch to instance-scoped logging; consider logging final size only on rank 0.

Prevents N-way duplicated logs in multi-GPU runs.

-from ...distributed.common import all_gather_object, get_world_size
+from ...distributed.common import all_gather_object, get_world_size, get_rank
@@
-        self._log_info(
-            f"Final KV cache size after resize: {final_cache_size_gb:.2f} GiB ({new_num_pages} pages)"
-        )
+        if get_rank() == 0:
+            self._log_info(
+                f"Final KV cache size after resize: {final_cache_size_gb:.2f} GiB ({new_num_pages} pages)"
+            )

Also applies to: 312-314

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (4)

161-161: Avoid fixed global seed; make per-rank and seed CUDA too.

A fixed global seed can produce identical sampling across ranks and ignores CUDA RNG.

-        torch.manual_seed(42)
+        # Per-rank seed; also seed CUDA RNGs
+        _seed = 42 + mpi_rank()
+        torch.manual_seed(_seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(_seed)

Would you prefer making the base seed configurable via LlmArgs/AutoDeployConfig?

---

185-187: slot_idx collection looks good; add a quick invariant check before nesting.

Keeps inputs aligned and catches scheduler mismatches early.

         extra_args: Dict[str, List[torch.Tensor]] = defaultdict(list)
@@
         # update the sequence info object now
+        # Sanity: ensure alignment of per-sequence lists
+        assert len(slot_idx) == len(input_ids) == len(input_pos) == len(page_assignments), (
+            f"len mismatch: slot_idx={len(slot_idx)}, input_ids={len(input_ids)}, "
+            f"input_pos={len(input_pos)}, pages={len(page_assignments)}"
+        )

---

204-206: Guard seq_slot presence for context requests.

Make the contract explicit; avoids None/invalid types propagating into nesting.

-            slot_idx.append(request.seq_slot)
+            assert request.seq_slot is not None, "seq_slot must be set for context requests"
+            slot_idx.append(int(request.seq_slot))

---

226-229: Clarify overlap-scheduler semantics for gen path; assert seq_slot.

Given the TODO, add a minimal guard now; verify that slot_idx should reflect the current (post-schedule) slot.

-            # TODO: double-check if this is correct for the overlap scheduler
-            slot_idx.append(request.seq_slot)
+            # TODO: confirm correctness for the overlap scheduler; slot_idx should map to current slot
+            assert request.seq_slot is not None, "seq_slot must be set for generation requests"
+            slot_idx.append(int(request.seq_slot))

Would you like a focused unit test covering overlap scheduling with rescatter?

tensorrt_llm/_torch/auto_deploy/shim/demollm.py (1)

145-149: Pass slot_idx on subsequent nest_sequences calls during decode.

Maintain the same slot mapping across steps to match cache/state expectations.

Apply:

-            sequence_info.nest_sequences(
+            sequence_info.nest_sequences(
                 token_ids,
                 input_pos=input_pos_next,
-                page_assignments=self._assign_pages(total_lens_next),
+                page_assignments=self._assign_pages(total_lens_next),
+                slot_idx=slot_idx,
             )

tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (6)

7-10: Avoid hard dependency on einops (optional)

_rms_norm_ref only needs simple reshapes; consider replacing einops with view/reshape to drop an extra dependency or lazy-import it inside the fallback path.

---

21-45: RMSNorm: add shape checks and prefer rsqrt; optional: drop einops

• Guard group_size divisibility and optional z shape.
• Use torch.rsqrt for better numerics.

 def _rms_norm_ref(
     x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True
 ):
-    dtype = x.dtype
-    # N = x.shape[-1]
+    dtype = x.dtype
+    hidden = x.shape[-1]
     weight = weight.float()
     bias = bias.float() if bias is not None else None
     if upcast:
         x = x.float()
         z = z.float() if z is not None else z
     if z is not None and not norm_before_gate:
         x = x * F.silu(z)
     if group_size is None:
-        rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps)
+        rstd = torch.rsqrt((x.square()).mean(dim=-1, keepdim=True) + eps)
         out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight)
     else:
-        x_group = rearrange(x, "... (g d) -> ... g d", d=group_size)
-        rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps)
-        out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight
+        if hidden % group_size != 0:
+            raise ValueError(f"group_size ({group_size}) must divide last dim ({hidden})")
+        # If einops stays: keep rearrange; otherwise use reshape:
+        x_group = rearrange(x, "... (g d) -> ... g d", d=group_size)
+        rstd = torch.rsqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps)
+        out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight
         if bias is not None:
             out = out + bias
     if z is not None and norm_before_gate:
         out *= F.silu(z)
     return out.to(dtype)

---

49-56: Silence unused-arg warnings on stubbed mask updaters

These are intentionally no-ops; add noqa or consume args to appease linters.

-def _nemotron_h_model_update_mamba_mask(self, attention_mask, cache_position):
+def _nemotron_h_model_update_mamba_mask(self, attention_mask, cache_position):  # noqa: ARG001
     return None
@@
-def _nemotron_h_model_update_causal_mask(self, attention_mask, input_tensor, cache_position):
+def _nemotron_h_model_update_causal_mask(self, attention_mask, input_tensor, cache_position):  # noqa: ARG001
     # Force attention to use causal mode without explicit masks
     return None

---

103-111: Minor: avoid .keys() and micro‑cleanup

Use membership directly; also consider guarding repeat patching.

-    for _, module in model.named_modules():
-        if (module_name := type(module).__name__) in CUSTOM_MODULE_PATCHES.keys():
+    for _, module in model.named_modules():
+        if (module_name := type(module).__name__) in CUSTOM_MODULE_PATCHES:
             patches = CUSTOM_MODULE_PATCHES[module_name]
             for method_name, method_patch in patches:
-                setattr(module, method_name, types.MethodType(method_patch, module))
+                setattr(module, method_name, types.MethodType(method_patch, module))

---

115-117: Global monkey‑patch risk: confine scope

Overriding transformers.AutoModelForCausalLM.from_config globally can affect unrelated code in‑process and tests. Gate behind an env flag or provide a scoped context manager.

-AutoModelForCausalLM.from_config = get_model_from_config_patched
+import os
+if os.getenv("TLLM_ENABLE_NEMOTRON_PATCH", "1") == "1":
+    AutoModelForCausalLM.from_config = get_model_from_config_patched

---

80-85: Optional: use a dedicated exception type

TRY003 flagged the ValueError message; if this is a hot path, consider a small custom exception with a fixed message to satisfy the rule. Otherwise, ignore.

tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py (3)

55-59: Seed change is fine; centralize seeding for full reproducibility.

Use a small helper to seed PyTorch (and CUDA) consistently; call it in both tests to avoid drift.

@@
-    seed = 42  # Set random seed for model param init
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
+    seed = 42  # Set random seed for model param init
+    set_seed(seed)

Add once near imports:

@@
+def set_seed(seed: int) -> None:
+    import random
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)

---

14-25: Nit: class name casing.

Consider TransformerLikeModelWithFakeCachePool for readability.

---

88-95: Minor: unify seeding and keep sampling deterministic checks minimal.

Use set_seed(seed) for this test too; current assertions are shape/range only, so no need to force global determinism beyond that.

@@
-    seed = 0
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
+    set_seed(0)

Also applies to: 116-137

tensorrt_llm/_torch/auto_deploy/models/patches/bamba.py (4)

22-28: Silence intentional unused-arg warnings without changing signatures.

These functions must keep their signatures for compatibility, but several parameters are intentionally unused.

Apply inline noqa markers on the def lines:

-def _bamba_mixer_torch_forward(
+def _bamba_mixer_torch_forward(  # noqa: ARG001
     self,
     input_states,
     cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
     cache_position: Optional[torch.LongTensor] = None,
     attention_mask: Optional[torch.Tensor] = None,
 ):
@@
-def _bamba_model_update_mamba_mask(self, attention_mask, cache_position):
+def _bamba_model_update_mamba_mask(self, attention_mask, cache_position):  # noqa: ARG001
     return None
@@
-def _bamba_model_update_causal_mask(
+def _bamba_model_update_causal_mask(  # noqa: ARG001
     self,
     attention_mask,
     input_tensor,
     cache_position,
     past_key_values,
     output_attentions,
 ):
@@
-def _cache_bool(self) -> bool:
+def _cache_bool(self) -> bool:  # noqa: ARG001
     return True

Based on coding guidelines.

Also applies to: 149-151, 153-163, 179-187

---

46-51: Ensure metadata dtypes/shapes match custom op contracts; cast explicitly.

seq_len/seq_start are created with torch.int (int32) while slot_idx is int64. Please confirm the op signatures; mis‑typed metadata will silently upcast on CPU but can fail AOT/export or CUDA kernels.

Proposed explicit casts (adjust if your op expects different types):

-        seq_len_t = torch.full((batch_size,), seq_len, device=input_states.device, dtype=torch.int)
-        seq_start_t = torch.arange(
-            0, batch_size * seq_len, seq_len, device=input_states.device, dtype=torch.int
-        )
-        slot_idx_t = torch.arange(batch_size, device=input_states.device, dtype=torch.long)
+        seq_len_t = torch.full(
+            (batch_size,), seq_len, device=input_states.device, dtype=torch.int32
+        )
+        seq_start_t = torch.arange(
+            0, batch_size * seq_len, seq_len, device=input_states.device, dtype=torch.int32
+        )
+        slot_idx_t = torch.arange(
+            batch_size, device=input_states.device, dtype=torch.int64
+        )

Also, if time_step_limit can be None, guard its conversion to list:

-            time_step_limit=list(self.time_step_limit),
+            time_step_limit=(list(self.time_step_limit) if self.time_step_limit is not None else []),

Please verify against the C++/Triton op signatures and adjust accordingly.

Also applies to: 52-72, 101-121

---

135-136: Prefer reshape over view after custom ops.

y may be non‑contiguous; view can error. reshape is safer.

-    y = y.view(batch_size, seq_len, -1)
+    y = y.reshape(batch_size, seq_len, -1)

---

7-15: Optional: avoid importing concrete classes directly from deep paths.

To reduce fragility against upstream refactors, consider importing the module and referencing names off it.

Example:

-from transformers.models.bamba.modeling_bamba import (
-    BambaMixer,
-    BambaModel,
-    BambaPreTrainedModel,
-    HybridMambaAttentionDynamicCache,
-    apply_mask_to_padding_states,
-)
+from transformers.models.bamba import modeling_bamba as bamba_mod
+from transformers.models.bamba.modeling_bamba import apply_mask_to_padding_states
+# Usage: bamba_mod.BambaMixer, bamba_mod.BambaModel, ...

As per coding guidelines.

tensorrt_llm/_torch/auto_deploy/custom_ops/mla.py (2)

1-1: Add NVIDIA Apache-2.0 header (2025) at top of file.

Repository guideline requires header on all .py files.

Apply:

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

As per coding guidelines.

---

184-186: Silence unused slot_idx parameter (Ruff ARG001).

Parameter is intentionally unused; rename to underscore to avoid lint noise.

-    slot_idx: torch.Tensor,
+    _slot_idx: torch.Tensor,

tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py (2)

1-1: Add NVIDIA Apache-2.0 header.

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License"); see LICENSE file for details.

As per coding guidelines.

---

8-10: Rename unused loop variable to underscore.

Small lint cleanup; functional no-op.

-for _, module_name, is_pkg in pkgutil.iter_modules(__path__):
+for _, module_name, _is_pkg in pkgutil.iter_modules(__path__):

tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py (2)

1-1: Add NVIDIA Apache-2.0 header.

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License"); see LICENSE file for details.

As per coding guidelines.

---

293-295: Silence unused metadata params (Ruff ARG001).

slot_idx/page_size/position_ids/pages_per_seq are intentionally unused here.

-def prepare_fused_mha_metadata(
+def prepare_fused_mha_metadata(
@@
-    slot_idx: torch.Tensor,
-    page_size: int,
+    _slot_idx: torch.Tensor,
+    _page_size: int,
@@
-def prepare_fused_mha_metadata_fake(
-    input_ids, position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, page_size
+def prepare_fused_mha_metadata_fake(
+    input_ids, _position_ids, seq_len, input_pos, cache_loc, _pages_per_seq, _slot_idx, _page_size
 ):

Also applies to: 312-313

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.py (3)

1-1: Add NVIDIA Apache-2.0 header.

+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License"); see LICENSE file for details.

As per coding guidelines.

---

12-12: Remove unused noqa; keep side-effect import comment.

-import tensorrt_llm._torch.auto_deploy  # noqa: F401
+import tensorrt_llm._torch.auto_deploy  # side-effect import to register custom ops

---

88-88: Drop unused E501 noquas.

Ruff flagged non-enabled noqa directives.

-        tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_decode(  # type: ignore  # noqa: E501
+        tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_decode(  # type: ignore
@@
-            tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_prefill(  # type: ignore  # noqa: E501
+            tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_prefill(  # type: ignore

Also applies to: 151-151

tensorrt_llm/commands/serve.py (1)

37-71: Signal handler performs non–async-signal-safe work (logging, waits) and exits the parent

Logging and multi-step process control inside a signal handler can be unsafe; also the unconditional sys.exit may tear down the parent in contexts where cleanup should happen in finally.

Minimize the handler to set a flag and send a signal to child; move waits/logging to normal control flow (e.g., check flag around blocking sections). At minimum, replace logger calls with os.write to stderr or defer logs.

tensorrt_llm/_torch/auto_deploy/llm.py (1)

46-48: Over‑strict assert on “multi_modal_data” with “messages”

Asserting absence can break valid callers; prefer ignoring or merging multimodal payload.

Replace assert with a guard/warning or merge behavior. Example:

-            assert "multi_modal_data" not in inputs, f"unexpected multi_modal_data key in {inputs=}"
+            # If present, multi_modal_data will be handled by processor; ignore here.

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py (2)

365-372: slot_idx accepted but unused — silence lints and document

Ruff flags unused args. Keep signature but explicitly discard to avoid warnings.

 def torch_backend_prepare_metadata(
@@
-    slot_idx: torch.Tensor,
+    slot_idx: torch.Tensor,
     page_size: int,
 ) -> List[torch.Tensor]:
     """Prepare metadata for torch backend attention (similar to triton backend)."""
+    # Note: slot_idx and page_size are reserved for future use in torch backend.
+    # Keep signature stable; explicitly mark as unused.
+    _ = (slot_idx, page_size)

---

380-390: Fake variant: discard unused parameters to appease linters

 def torch_backend_prepare_metadata_fake(
-    input_ids, position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, page_size
+    input_ids, position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, page_size
 ):
+    _ = (position_ids, pages_per_seq, slot_idx, page_size)
     num_seq = SequenceInfo._get_sanitized_num_sequences(input_ids, seq_len)

tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py (4)

156-166: slot_idx unused in prepare_flashinfer_metadata; keep signature but silence linters.

The backend doesn't need slot_idx (uniform signature across ops), but Ruff flags it.

Use a throwaway name to document intent:

-    pages_per_seq: torch.Tensor,
-    slot_idx: torch.Tensor,
+    pages_per_seq: torch.Tensor,
+    _slot_idx: torch.Tensor,

---

215-228: Fake variant: multiple unused args; rename to underscore to avoid ARG001.

No functional change; keeps signatures aligned.

-@prepare_flashinfer_metadata.register_fake
-def prepare_flashinfer_metadata_fake(
-    input_ids, position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, page_size
-):
+@prepare_flashinfer_metadata.register_fake
+def prepare_flashinfer_metadata_fake(
+    input_ids, _position_ids, seq_len, _input_pos, cache_loc, _pages_per_seq, _slot_idx, _page_size
+):

---

418-421: Scale type guard is too strict; accept ints/real numbers.

Some callers may pass an int or NumPy scalar. Treat any real number as valid.

-from typing import ...
+import numbers
...
-        if not (isinstance(scale, float) or scale is None):
+        if not (isinstance(scale, numbers.Real) or scale is None):
             ad_logger.warning(f"Provided {scale=}, is not a float. Using default scale instead.")
             scale = None

---

278-283: Shadowing function args k_scale/v_scale; use locals to avoid confusion.

You immediately overwrite k_scale/v_scale, then pass them to wrapper.run. Use local names to make intent explicit and prevent accidental reliance on call-site constants.

-    # Assuming k_scale = v_scale = 1.0, we just have to cast k and v to fp8 before appending to kv cache
-    k_scale, v_scale = 1.0, 1.0
+    # Assuming fixed scales of 1.0 for FP8 cache append
+    k_scale_local, v_scale_local = 1.0, 1.0
     if k_cache.dtype == torch.float8_e4m3fn:
-        k = (k / k_scale).to(torch.float8_e4m3fn)
-        v = (v / v_scale).to(torch.float8_e4m3fn)
+        k = (k / k_scale_local).to(torch.float8_e4m3fn)
+        v = (v / v_scale_local).to(torch.float8_e4m3fn)
...
-    y = wrapper.run(q, (k_cache, v_cache), k_scale=k_scale, v_scale=v_scale)
+    y = wrapper.run(q, (k_cache, v_cache), k_scale=k_scale_local, v_scale=v_scale_local)

If dynamic scaling becomes needed, consider plumbing per‑tensor scales and removing these constants from get_constants.

Also applies to: 303-305

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.py (2)

12-12: Remove unused noqa directive.

Ruff flags RUF100 since F401 isn’t enabled. Keep the import (side-effect registration), but drop the noqa.

-import tensorrt_llm._torch.auto_deploy  # noqa: F401
+import tensorrt_llm._torch.auto_deploy

Optionally add a clarifying comment: “import for side-effects: operator registration.”

---

153-156: Remove unused noqa E501.

Ruff indicates the directive is unused; line fits current config.

-            tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_prefill(  # type: ignore  # noqa: E501
+            tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_causal_conv._torch_causal_conv1d_prefill(  # type: ignore

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.py (2)

4-4: Remove unused noqa directive.

Keep the import for side-effects; drop # noqa: F401 per RUF100.

-import tensorrt_llm._torch.auto_deploy  # noqa: F401
+import tensorrt_llm._torch.auto_deploy

---

166-170: Rename unused loop variable to underscore.

Silences Ruff B007 and clarifies intent.

-    for i, ln in enumerate(lens):
+    for i, _ln in enumerate(lens):

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_causal_conv.py (2)

20-21: Use explicit error instead of assert for padding_mode.

Asserts can be compiled out; prefer a runtime check.

-    assert padding_mode == "zeros", "padding_mode must be zeros"
+    if padding_mode != "zeros":
+        raise ValueError("padding_mode must be 'zeros'")

---

22-22: Silence unused variable to satisfy linters.

-    batch_size, seq_len, _ = input.shape
+    _, seq_len, _ = input.shape

tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_hybrid_patches.py (2)

1-1: Remove unnecessary # noqa directives.

They hide useful linting and are not needed here.

-import torch  # noqa
+import torch
-from torch.export import Dim  # noqa
+from torch.export import Dim
-from tensorrt_llm._torch.auto_deploy.transformations._graph import move_to_device  # noqa
+from tensorrt_llm._torch.auto_deploy.transformations._graph import move_to_device

Also applies to: 3-3, 8-8

---

66-67: Hard-coded CUDA usage can fail on CPU envs; use detected device.

Tests live in singlegpu, but device plumbing is trivial to make robust.

-    inputs = tokenizer(message, return_tensors="pt", return_token_type_ids=False).to("cuda")
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    inputs = tokenizer(message, return_tensors="pt", return_token_type_ids=False).to(dev)

tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.py (2)

9-13: Skip module if CUDA unavailable.

Avoid failures on CPU CI.

 import pytest
 import torch
 
 import tensorrt_llm._torch.auto_deploy  # noqa: F401
+
+# Skip all tests here if CUDA is unavailable
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")

---

12-12: Remove redundant # noqa markers.

They refer to non-enabled rules and aren’t needed.

-import tensorrt_llm._torch.auto_deploy  # noqa: F401
+import tensorrt_llm._torch.auto_deploy
-        tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_mamba._torch_cached_ssm_transform_decode(  # type: ignore  # noqa: E501
+        tensorrt_llm._torch.auto_deploy.custom_ops.torch_backend_mamba._torch_cached_ssm_transform_decode(  # type: ignore
-            tensorrt_llm._torch.auto_deploy.custom_ops.torch_mamba._torch_ssm_transform_prefill(  # type: ignore  # noqa: E501
+            tensorrt_llm._torch.auto_deploy.custom_ops.torch_mamba._torch_ssm_transform_prefill(  # type: ignore

Also applies to: 98-98, 171-171

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_mamba.py (1)

165-183: Op wiring and meta path look consistent with prefill returning fp32.

No functional concerns; consider brief Google-style docstrings on public custom ops for clarity (args/dtypes/shapes).

Also applies to: 185-197

tensorrt_llm/_torch/auto_deploy/custom_ops/triton_backend_mamba.py (2)

125-156: Clamp/softplus path is sound; minor perf nit.

dt_bias_zero is constructed but unused by kernel (dt_softplus=False). Safe, but you can avoid allocating it by passing None if selective_state_update permits optional dt_bias when dt_softplus=False.

---

224-249: Cache initializer: simplify ssm_state_size detection.

The else branch max(1, B_fake.shape[-1]) is redundant; shape[-1] already ≥1. Also guard for dtype/device consistency from cache_config.

-        if B_fake.ndim >= 4:
-            ssm_state_size = B_fake.shape[-1]
-        else:
-            ssm_state_size = max(1, B_fake.shape[-1])
+        ssm_state_size = B_fake.shape[-1]

tensorrt_llm/_torch/auto_deploy/transform/library/attention.py (5)

337-348: Pattern 6 passes enable_gqa=True in search; may miss graphs without this kwarg.

If upstream sdpa.default omits enable_gqa (defaults), the pattern may not match. Consider registering a variant without enable_gqa or ignore this kwarg during matching.

-            register_ad_pattern(
+            # Variant that ignores the 'enable_gqa' kwarg to improve match robustness
+            register_ad_pattern(
                 search_fn=_grouped_attn_pattern_6,
                 replace_fn=_grouped_attn_replacement_6,
                 patterns=patterns,
                 dummy_args=dummy_args_2,
-                scalar_workaround={"scale": scale, "dropout_p": dropout},
+                scalar_workaround={"scale": scale, "dropout_p": dropout},
+                op_ignore_types={torch.ops.auto_deploy.torch_attention_sdpa.default: (bool,)},
             )

---

356-367: Same as above for Pattern 7 (causal).

Apply the same op_ignore_types mapping as in Pattern 6 to avoid brittle matches.

---

375-392: Patterns 8/9 omit attn_mask and add enable_gqa; add robust variant or ignore bools.

Mirror the op_ignore_types addition here too to match callsites lacking the kwarg.

---

413-435: Replacement 10 ignores n_rep (benign) — suppress lint or document.

n_rep is unused in replacement_10 by design (grouped op handles repetition). Add a comment or underscore the param to avoid ARG001 noise.

-def _grouped_attn_replacement_10(q, k, v, n_rep, dropout_p):
+def _grouped_attn_replacement_10(q, k, v, _n_rep, dropout_p):
     return torch.ops.auto_deploy.torch_attention_grouped_sdpa.default(

---

668-672: MatchAttentionLayout only targets grouped_sdpa.

If graphs still contain torch_attention_sdpa at this stage, they won't be transposed for bsnd. Consider including sdpa as well or document the pass ordering guarantee.

tensorrt_llm/_torch/auto_deploy/custom_ops/cuda_backend_causal_conv.py (3)

139-145: Weight shape handling: assert depthwise contract or document supported cases.

You assume weight.ndim==3 with shape[-2]==1 or 2D [C,K]. If other layouts sneak in, assert early with a clear error.

-    if weight.ndim == 3:
+    if weight.ndim == 3:
         assert weight.shape[-2] == 1
         w2d = weight.squeeze(-2)
-    else:
+    elif weight.ndim == 2:
         w2d = weight
+    else:
+        raise ValueError(f"Unsupported weight shape {tuple(weight.shape)}; expected [C,K] or [C,1,K].")

---

288-293: Constants are ignored by CUDA path; validate defaults to avoid semantic drift.

If source op uses non-default stride/padding/dilation/groups/padding_mode, cached CUDA path will not honor them. Either plumb them or guard with a check that they are default values.

-        stride, padding, dilation, groups, padding_mode = extract_op_args(
+        stride, padding, dilation, groups, padding_mode = extract_op_args(
             source_attn_node, "stride", "padding", "dilation", "groups", "padding_mode"
         )
+        # Optional: enforce supported defaults
+        if (stride != 1) or (padding != 0) or (dilation != 1) or (groups != 1) or (padding_mode not in (None, "zeros")):
+            ad_logger.warning("cuda_cached_causal_conv1d does not honor non-default conv params; results may differ.")
         return [stride, padding, dilation, groups, padding_mode]

---

36-50: Unused helper _build_conv_state_from_sequence.

If not needed, drop it to reduce surface area.

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_causal_conv.py (5)

98-106: Decode path constraints not enforced by constants; document/guard in metadata.

Decode asserts stride==1 and dilation==1. If source op is configured differently, cached path will raise at runtime.

• Consider validating these constants once in prepare/transform pass and fail early.
• Alternatively, specialize kernel/constants so prefill uses provided stride/dilation but decode enforces 1.

If you want, I can add a compile-time check in the pass that inserts cached ops.

Also applies to: 101-103

---

65-76: Unused locals (batch_size, seq_len).

Minor lint from Ruff (RUF059). Replace with underscores to avoid warnings.

-    batch_size, seq_len, _ = input.shape
+    _batch_size, _seq_len, _ = input.shape

---

104-116: Unused local (batch_size).

Same nit in decode.

-    batch_size, seq_len, _ = input.shape
+    _batch_size, _seq_len, _ = input.shape

---

112-121: Cache update via roll is O(K). Consider deque-like assignment.

Rolling shifts the entire [C_in, K] window every token. For large K this is avoidable.

Use a circular index tracked per slot or maintain a write pointer buffer; if that’s out of scope now, keep as-is and revisit for perf.

---

300-310: get_num_qkv_args=3 with optional bias; ensure callsites pass None safely.

If bias is None, confirm the replacement pass supplies an Optional[Tensor] acceptable by custom_op schema; otherwise wrap a registered None-like tensor.

I can add a small adapter in the transform to coerce None to a zero bias tensor of correct dtype/device.

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (2)

663-666: Zero out stale tail of slot_idx for fewer-than-max batches.

_seq_len is reset, but slot_idx isn’t. While sanitized consumers slice to num_seq, zeroing avoids accidental misuse elsewhere.

-        if slot_idx is not None:
-            self._store_arg("slot_idx", slot_idx)
+        if slot_idx is not None:
+            self._store_arg("slot_idx", slot_idx, reset=True)

---

427-466: Return types in sanitization helpers: ensure Python ints for slicing.

_get_sanitized_num_sequences returns a Tensor when s > 1; later code uses it in slicing. Convert to int to avoid brittle behavior.

-        if s > 1:
-            num_seq = torch.sum(seq_len > 0)
+        if s > 1:
+            num_seq = int(torch.sum(seq_len > 0).item())

Do the same wherever num_seq is produced and used as a slice bound. Based on learnings.

tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_mamba.py (2)

31-43: Decode helper’s chunk_size is unused.

Lint noise and possible confusion.

-    chunk_size: int,
+    _chunk_size: int,

Alternatively, remove the parameter if not needed.

---

105-107: Unused helper _update_ssm_state_cache.

Not referenced; remove to reduce surface area.

-def _update_ssm_state_cache(ssm_cache: torch.Tensor, ssm_state: torch.Tensor) -> None:
-    ssm_cache.copy_(ssm_state)
+# Removed unused helper _update_ssm_state_cache