tensorrt_llm/inputs/__init__.py (1)

2-2: Avoid eager torch import: lazily re‑export EVS functions to cut import‑time overhead.

Importing compute_retained_tokens_count/compute_retention_mask here brings torch into tensorrt_llm.inputs import path unconditionally. Prefer lazy export to keep lightweight imports for non‑EVS users.

Apply this diff to remove the eager import:

-from .evs import compute_retained_tokens_count, compute_retention_mask

Then add a lazy accessor near the bottom of the file (outside the selected range):

# Lazy re-exports to avoid importing torch during package import.
def __getattr__(name):
    if name in ("compute_retained_tokens_count", "compute_retention_mask"):
        from .evs import compute_retained_tokens_count as _cnt, compute_retention_mask as _msk
        globals()["compute_retained_tokens_count"] = _cnt
        globals()["compute_retention_mask"] = _msk
        return globals()[name]
    raise AttributeError(name)

As per coding guidelines

tensorrt_llm/inputs/evs.py (3)

69-74: Cast to float32 for stable cosine similarity.

cosine_similarity on bfloat16 (especially CPU) can be numerically unstable or unsupported on some backends. Cast embeds to float32 for the similarity step.

Apply this diff:

-    # Core EVS
-    similarity = torch.nn.functional.cosine_similarity(video_embeds[1:, ...],
+    # Core EVS
+    _emb = video_embeds.to(torch.float32)
+    similarity = torch.nn.functional.cosine_similarity(_emb[1:, ...],
-                                                       video_embeds[:-1, ...],
+                                                       _emb[:-1, ...],
                                                        dim=-1)

---

81-88: Use topk instead of full sort; clamp k to valid range.

argsort is O(N log N). topk is faster and avoids slicing beyond array bounds if mis‑configured. Also early‑return if k<=0.

Apply this diff:

-    order = torch.argsort(dissimilarity_flat,
-                          dim=-1,
-                          descending=True,
-                          stable=True)
-    retain_num_tokens = compute_retained_tokens_count(video_size,
-                                                      spatial_merge_size, q)
-    topk_indices = order[:retain_num_tokens]
+    retain_num_tokens = compute_retained_tokens_count(video_size,
+                                                      spatial_merge_size, q)
+    k = min(retain_num_tokens, dissimilarity_flat.numel())
+    if k <= 0:
+        mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool)
+        return mask.view(-1) if flatten_output else mask.view(dissimilarity.size())
+    _, topk_indices = torch.topk(dissimilarity_flat, k=k, largest=True, sorted=False)

---

61-68: Shape safety: guard reshape assumptions.

If T * (H//s) * (W//s) != video_embeds.size(0) * (video_embeds.size(1) if video_embeds.dim() > 2 else 1), this will error. Add an assertion to fail fast with a clear message.

Proposed addition:

     video_embeds = video_embeds.reshape(
         T,
         H // spatial_merge_size,
         W // spatial_merge_size,
         video_embeds.size(-1),
     )
+    assert (H // spatial_merge_size) > 0 and (W // spatial_merge_size) > 0, (
+        f"Invalid spatial merge: H={H}, W={W}, spatial_merge_size={spatial_merge_size}"
+    )

tensorrt_llm/_torch/models/modeling_nanov2vlm.py (3)

27-28: Validate and clamp VIDEO_PRUNING_RATIO from env.

Avoid crashes or over‑pruning when env var is malformed or out of [0,1).

Apply this diff:

-VIDEO_PRUNING_RATIO = float(os.getenv("TLLM_VIDEO_PRUNING_RATIO", "0"))
+try:
+    VIDEO_PRUNING_RATIO = float(os.getenv("TLLM_VIDEO_PRUNING_RATIO", "0"))
+except ValueError:
+    VIDEO_PRUNING_RATIO = 0.0
+# Clamp for safety
+VIDEO_PRUNING_RATIO = max(0.0, min(VIDEO_PRUNING_RATIO, 0.99))

---

53-54: Make spatial_merge_size computation robust.

Division and truncation can yield 0 or off‑by‑one for odd ratios. Validate result.

Apply this diff:

-        self.spatial_merge_size = int(self.patch_size / self.downsample_ratio)
+        self.spatial_merge_size = int(round(self.patch_size / self.downsample_ratio))
+        if self.spatial_merge_size <= 0:
+            raise ValueError(f"Invalid spatial_merge_size from patch_size={self.patch_size}, downsample_ratio={self.downsample_ratio}")

Please verify that for your configs: (H // spatial_merge_size) * (W // spatial_merge_size) == self.num_image_token per frame.

---

17-18: Import surface is fine, but consider dependency direction.

Keeping the EVS import local to this model (rather than package‑root re‑export) helps avoid global import side‑effects; if you keep the package re‑export, prefer the lazy import suggested in inputs/init.py.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro