address CodeRabbit comments

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>

Author: o-stoner

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py

@@ -20,6 +20,7 @@
 Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16.
 """
 
+import contextvars
 import math
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -198,22 +199,22 @@ def build(
         )
 
 
-_vsa_forward_context: Optional[VSAMetadata] = None
+_vsa_forward_context_var: contextvars.ContextVar[Optional[VSAMetadata]] = contextvars.ContextVar(
+    "_vsa_forward_context", default=None
+)
 
 
 @contextmanager
 def set_vsa_forward_context(metadata: VSAMetadata):
-    global _vsa_forward_context
-    prev = _vsa_forward_context
-    _vsa_forward_context = metadata
+    token = _vsa_forward_context_var.set(metadata)
     try:
         yield
     finally:
-        _vsa_forward_context = prev
+        _vsa_forward_context_var.reset(token)
 
 
 def get_vsa_forward_context() -> Optional[VSAMetadata]:
-    return _vsa_forward_context
+    return _vsa_forward_context_var.get(None)
 
 
 def _mean_pool_cubes(
@@ -525,7 +526,10 @@ def _forward_vsa(
         o_c = torch.einsum("bhnm,bmhd->bnhd", attn_probs_c, v_c)
 
         use_cute = (
-            _vsa_import_error is None and is_cute_supported(q) and (q.dtype == k.dtype == v.dtype)
+            _vsa_import_error is None
+            and is_cute_supported(q)
+            and (q.dtype == k.dtype == v.dtype)
+            and num_cubes <= VSA_KERNEL_MAX_CUBES
         )
         topk_indices = attn_probs_c.topk(cur_topk, dim=-1).indices.to(torch.int32)
 
@@ -534,11 +538,6 @@ def _forward_vsa(
         )
 
         if use_cute:
-            assert num_cubes <= VSA_KERNEL_MAX_CUBES, (
-                f"VSA CuTe kernel supports at most {VSA_KERNEL_MAX_CUBES} cubes "
-                f"(SMEM-allocated variable_block_sizes buffer); got num_cubes={num_cubes}. "
-                "Lower video resolution/length or fall back to dense SDPA."
-            )
             q_hnd = q_t.transpose(1, 2).contiguous()
             k_hnd = k_t.transpose(1, 2).contiguous()
             v_hnd = v_t.transpose(1, 2).contiguous()

tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/block_sparse_attn_dsl_fwd.py

@@ -62,6 +62,8 @@ class VideoSparseAttentionForwardGroup2QInterleaveKV:
     Q1  0  S1 O1
     """
 
+    MAX_INDICES = 4 * 1024
+
     def __init__(
         self,
         block_m: int,
@@ -413,7 +415,7 @@ def __call__(
             (2, self.mma_tiler_qk[0], self.scale_buffers), (0, 1, 2)
         )
 
-        self.max_indices = 4 * 1024
+        self.max_indices = self.MAX_INDICES
 
         @cute.struct
         class SharedStorage:

tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py

@@ -112,6 +112,15 @@ def block_sparse_attn_from_indices_cute(
             "cutlass-dsl is not importable."
         )
 
+    num_q_blk = variable_block_sizes.shape[0]
+    if num_q_blk > VideoSparseAttentionForward.MAX_INDICES:
+        raise ValueError(
+            f"variable_block_sizes has {num_q_blk} entries but the CuTe kernel "
+            f"supports at most {VideoSparseAttentionForward.MAX_INDICES} "
+            "(SMEM-allocated sVariable_block_sizes buffer). Lower video "
+            "resolution/length or fall back to dense SDPA."
+        )
+
     B, H, T, D = q.shape
     if sm_scale is None:
         sm_scale = 1.0 / math.sqrt(D)

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py

@@ -488,8 +488,10 @@ def forward(
 
         # Apply an explicit user flow_shift override; otherwise keep the checkpoint
         # scheduler default so output matches the reference HuggingFace pipeline.
+        sched_cfg = self.scheduler.config
+        shift_key = None
+        orig_shift = None
         if flow_shift is not None:
-            sched_cfg = self.scheduler.config
             shift_key = (
                 "shift"
                 if "shift" in sched_cfg
@@ -498,10 +500,15 @@ def forward(
                 else None
             )
             if shift_key is not None and sched_cfg[shift_key] != flow_shift:
-                logger.info(f"flow_shift: {sched_cfg[shift_key]} -> {flow_shift} (user)")
+                orig_shift = sched_cfg[shift_key]
+                logger.info(f"flow_shift: {orig_shift} -> {flow_shift} (user)")
                 self.scheduler.register_to_config(**{shift_key: flow_shift})
 
-        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        try:
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        finally:
+            if orig_shift is not None:
+                self.scheduler.register_to_config(**{shift_key: orig_shift})
 
         # Wan2.2 A14B: Calculate boundary timestep for two-stage denoising
         boundary_timestep = None

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py

@@ -540,8 +540,10 @@ def forward(
         )
 
         # Apply an explicit user flow_shift override; otherwise keep the checkpoint scheduler default.
+        sched_cfg = self.scheduler.config
+        shift_key = None
+        orig_shift = None
         if flow_shift is not None:
-            sched_cfg = self.scheduler.config
             shift_key = (
                 "shift"
                 if "shift" in sched_cfg
@@ -550,10 +552,15 @@ def forward(
                 else None
             )
             if shift_key is not None and sched_cfg[shift_key] != flow_shift:
-                logger.info(f"flow_shift: {sched_cfg[shift_key]} -> {flow_shift} (user)")
+                orig_shift = sched_cfg[shift_key]
+                logger.info(f"flow_shift: {orig_shift} -> {flow_shift} (user)")
                 self.scheduler.register_to_config(**{shift_key: flow_shift})
 
-        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        try:
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        finally:
+            if orig_shift is not None:
+                self.scheduler.register_to_config(**{shift_key: orig_shift})
 
         # Wan2.2: Calculate boundary timestep for two-stage denoising
         boundary_timestep = None

tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py

@@ -368,6 +368,7 @@ def __init__(
         )
         if _is_vsa:
             q_dim = num_heads * head_dim
+            gate_tp_mode = TensorParallelMode.COLUMN if tp_size > 1 else None
             self.to_gate_compress = Linear(
                 hidden_size,
                 q_dim,
@@ -377,6 +378,8 @@ def __init__(
                 quant_config=quant_config,
                 skip_create_weights_in_init=skip_create_weights,
                 force_dynamic_quantization=force_dynamic_quant,
+                tensor_parallel_mode=gate_tp_mode,
+                reduce_output=False,
             )
             self.to_gate_fine = Linear(
                 hidden_size,
@@ -387,6 +390,8 @@ def __init__(
                 quant_config=quant_config,
                 skip_create_weights_in_init=skip_create_weights,
                 force_dynamic_quantization=force_dynamic_quant,
+                tensor_parallel_mode=gate_tp_mode,
+                reduce_output=False,
             )
 
         # I2V: Additional K/V projections for image embeddings.

tensorrt_llm/_torch/visual_gen/modules/attention.py

@@ -455,7 +455,7 @@ def _attn_impl(
         seq_len_kv = k.shape[1] if k is not None else seq_len
 
         def _reshape_gate(gate: torch.Tensor) -> torch.Tensor:
-            gate = gate.view(batch_size, -1, self.num_attention_heads, self.head_dim)
+            gate = gate.view(batch_size, -1, self.local_num_attention_heads, self.head_dim)
             if backend_layout == AttentionTensorLayout.HND:
                 gate = gate.transpose(1, 2)
             return gate

tests/unittest/_torch/visual_gen/test_attention_integration.py

@@ -698,6 +698,11 @@ def _build_vsa_setup(sparsity: float, batch_size: int, seed: int):
     integrated = Attention(hidden_size, num_heads, qkv_mode=QKVMode.FUSE_QKV, config=cfg_vsa).to(
         device
     )
+    # Fail loudly if the VSA path silently fell back to dense (which would set
+    # attn_backend to "VANILLA") instead of selecting the CUTEDSL/VSA backend.
+    assert integrated.attn_backend == "CUTEDSL", (
+        f"Expected CUTEDSL (VSA) backend, got {integrated.attn_backend!r}"
+    )
     copy_weights_self_attention(naive, integrated)
     naive.eval()
     integrated.eval()

tests/unittest/_torch/visual_gen/test_attention_perf.py

@@ -515,6 +515,16 @@ def benchmark_single(
             )
             hidden_states, freqs = self.create_test_data(batch_size, seq_len, hidden_size, head_dim)
 
+            # Fail fast if the requested backend silently resolved to something
+            # else (e.g. CUTEDSL/VSA downgraded to VANILLA dense for an
+            # unsupported config).
+            resolved_backend = getattr(model, "attn_backend", backend)
+            if vsa_sparsity is not None:
+                assert resolved_backend == "CUTEDSL", (
+                    f"VSA run requested CUTEDSL but the attention module resolved to "
+                    f"{resolved_backend!r}; refusing to report VSA timings for a dense fallback."
+                )
+
             # VSA needs an active forward context + a zero gate_compress; other
             # backends use neither.
             vsa_metadata = None
@@ -572,6 +582,7 @@ def _forward():
                 "p99_ms": torch.quantile(times_tensor, 0.99).item(),
                 "times_ms": times,
                 "uses_sage": _is_sage_attention_enabled(model),
+                "resolved_backend": resolved_backend,
             }
 
             # Calculate throughput (approximate TOPS)
@@ -587,6 +598,9 @@ def _forward():
 
             return stats
 
+        except AssertionError:
+            # Backend-resolution / fail-fast checks are genuine test failures.
+            raise
         except Exception as e:
             if verbose:
                 print(f"  {backend}: ERROR - {e}")
@@ -1224,6 +1238,9 @@ def test_vsa_module_vs_vanilla_wan22_t2v_14b(self, sparsity: float):
             backend="CUTEDSL", vsa_sparsity=sparsity, latent_shape=self._LATENT_SHAPE, **common
         )
         assert vanilla is not None and vsa is not None, "benchmark returned None (skipped/failed)"
+        assert vsa["resolved_backend"] == "CUTEDSL", (
+            f"VSA timings came from {vsa['resolved_backend']!r}, not the CUTEDSL/VSA path"
+        )
         assert vanilla["avg_ms"] > 0 and vsa["avg_ms"] > 0
 
         speedup = vanilla["avg_ms"] / vsa["avg_ms"]