small updates, mostly documentation

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

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-CuTe DSL Backend for Visual Generation Models
+CuTe DSL (NVIDIA kernels) Backend for Visual Generation Models
 
 CuTeDSLAttention runs the VSA sparse path when sparse_attention_config is set,
 otherwise the dense cubin path (with optional QK16PV8 quantization).
@@ -47,6 +47,17 @@
     cute_dsl = None
     _cute_dsl_import_error = e
 
+_vsa_import_error = None
+try:
+    from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import (
+        block_sparse_attn_from_indices_cute,
+        is_cute_supported,
+    )
+except (ImportError, OSError) as e:
+    block_sparse_attn_from_indices_cute = None
+    is_cute_supported = None
+    _vsa_import_error = e
+
 
 # VSA (Video Sparse Attention) sparse-path helpers
 
@@ -267,8 +278,9 @@ def __init__(
         skip_softmax_threshold_scale: Optional[float] = None,
         **kwargs,
     ):
-        # Dense path requires head_dim=128 (packaged cubins); the VSA sparse
-        # path JIT-compiles per shape, so it has no such restriction.
+        # Dense cubin path is head_dim=128-only (packaged cubins), so enforce it
+        # here. The VSA path needs no check: it is gated at runtime by
+        # is_cute_supported and falls back to dense SDPA when head_dim != 128.
         if sparse_attention_config is None and head_dim != 128:
             raise ValueError(f"CUTEDSL cubins require head_dim=128, got head_dim={head_dim}.")
         self.layer_idx = layer_idx
@@ -479,13 +491,6 @@ def _forward_vsa(
         Returns:
             [B, S, H, D] in the same original token order.
         """
-        # Lazy import: the VSA kernels package is optional and may not be
-        # importable in environments without the cute-dsl runtime.
-        from ..cute_dsl_kernels.blackwell.video_sparse_attention import (
-            block_sparse_attn_from_indices_cute,
-            is_cute_supported,
-        )
-
         if gate_compress is None:
             raise ValueError(
                 "CuTeDSLAttention VSA path requires gate_compress. "
@@ -519,7 +524,9 @@ def _forward_vsa(
         attn_probs_c = scores_c.softmax(dim=-1)
         o_c = torch.einsum("bhnm,bmhd->bnhd", attn_probs_c, v_c)
 
-        use_cute = is_cute_supported(q) and (q.dtype == k.dtype == v.dtype)
+        use_cute = (
+            _vsa_import_error is None and is_cute_supported(q) and (q.dtype == k.dtype == v.dtype)
+        )
         topk_indices = attn_probs_c.topk(cur_topk, dim=-1).indices.to(torch.int32)
 
         o_c_tiled = (

tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py

@@ -63,8 +63,8 @@ def __init__(self, model_config):
         if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa":
             raise ValueError(
                 "Video Sparse Attention (sparse_attention_config.algorithm='vsa') is "
-                "only supported by Wan T2V pipelines. Remove sparse_attention_config "
-                "for FLUX."
+                "only supported by the Wan 2.1 T2V 14B (720P) pipeline. Remove "
+                "sparse_attention_config for FLUX."
             )
 
         super().__init__(model_config)

tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py

@@ -132,8 +132,8 @@ def __init__(self, model_config):
         if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa":
             raise ValueError(
                 "Video Sparse Attention (sparse_attention_config.algorithm='vsa') is "
-                "only supported by Wan T2V pipelines. Remove sparse_attention_config "
-                "for FLUX.2."
+                "only supported by the Wan 2.1 T2V 14B (720P) pipeline. Remove "
+                "sparse_attention_config for FLUX.2."
             )
 
         super().__init__(model_config)

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py

@@ -604,8 +604,8 @@ def __init__(self, model_config):
         if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa":
             raise ValueError(
                 "Video Sparse Attention (sparse_attention_config.algorithm='vsa') is "
-                "only supported by Wan T2V pipelines. Remove sparse_attention_config "
-                "for LTX-2."
+                "only supported by the Wan 2.1 T2V 14B (720P) pipeline. Remove "
+                "sparse_attention_config for LTX-2."
             )
         super().__init__(model_config)
 

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

@@ -105,7 +105,7 @@ def __init__(self, model_config):
         if _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa":
             raise ValueError(
                 "Video Sparse Attention (sparse_attention_config.algorithm='vsa') is "
-                "only supported by Wan T2V pipelines (Wan 2.1 and Wan 2.2). Remove "
+                "only supported by the Wan 2.1 T2V 14B (720P) pipeline. Remove "
                 "sparse_attention_config for Wan I2V."
             )
 

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

@@ -246,6 +246,19 @@ def forward(
         return temb, temb_proj, encoder_hidden_states, encoder_hidden_states_image
 
 
+def _default_vsa_gate(linear: Linear, bias_value: float) -> None:
+    """Default a VSA gate Linear absent from the checkpoint: weight=0, bias=bias_value
+    (G_c=0 / G_f=1, preserving dense behavior at sparsity=0)."""
+    if not linear._weights_created:
+        linear.create_weights()
+    if linear.weight.is_meta:
+        return
+    with torch.no_grad():
+        linear.weight.zero_()
+        if linear.bias is not None:
+            linear.bias.fill_(bias_value)
+
+
 class WanBlock(nn.Module):
     def __init__(
         self,
@@ -417,32 +430,6 @@ def __init__(
             torch.empty(1, 6, hidden_size).normal_(std=hidden_size**-0.5)
         )
 
-    def init_gate_compress_zero(self) -> None:
-        """Zero-initialize to_gate_compress."""
-        if self.to_gate_compress is None:
-            return
-        if not self.to_gate_compress._weights_created:
-            self.to_gate_compress.create_weights()
-        if self.to_gate_compress.weight.is_meta:
-            return
-        with torch.no_grad():
-            self.to_gate_compress.weight.zero_()
-            if self.to_gate_compress.bias is not None:
-                self.to_gate_compress.bias.zero_()
-
-    def init_gate_fine_default(self) -> None:
-        """Initialize to_gate_fine to emit constant 1 (weight=0, bias=1)."""
-        if self.to_gate_fine is None:
-            return
-        if not self.to_gate_fine._weights_created:
-            self.to_gate_fine.create_weights()
-        if self.to_gate_fine.weight.is_meta:
-            return
-        with torch.no_grad():
-            self.to_gate_fine.weight.zero_()
-            if self.to_gate_fine.bias is not None:
-                self.to_gate_fine.bias.fill_(1.0)
-
     def forward(
         self,
         x,
@@ -839,7 +826,6 @@ def load_weights(self, weights: dict) -> None:
         }
         loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map)
 
-        loaded_linears: set = set()
         for name, module in tqdm(self.named_modules(), desc="Loading weights"):
             if len(module._parameters) == 0:
                 continue
@@ -849,7 +835,12 @@ def load_weights(self, weights: dict) -> None:
 
                 if weight_dicts:
                     loader.load_linear_weights(module, name, weight_dicts)
-                    loaded_linears.add(name)
+                # VSA gates absent from the checkpoint default to G_c=0 / G_f=1
+                # (dense behavior at sparsity=0).
+                elif name.endswith(".to_gate_compress"):
+                    _default_vsa_gate(module, 0.0)
+                elif name.endswith(".to_gate_fine"):
+                    _default_vsa_gate(module, 1.0)
                 elif "add_k_proj" in name or "add_v_proj" in name:
                     logger.info(f"[Weight Loading] No weights found for I2V module: {name}")
             elif isinstance(module, RMSNormTPAware):
@@ -863,20 +854,6 @@ def load_weights(self, weights: dict) -> None:
                             module_weights[param_name].to(self.model_config.torch_dtype)
                         )
 
-        # Default any VSA gates not loaded from the checkpoint: G_c=0, G_f=1
-        # (preserves dense behavior at sparsity=0).
-        for name, module in self.named_modules():
-            if not isinstance(module, WanBlock):
-                continue
-            if module.to_gate_compress is not None:
-                gate_path = f"{name}.to_gate_compress" if name else "to_gate_compress"
-                if gate_path not in loaded_linears:
-                    module.init_gate_compress_zero()
-            if module.to_gate_fine is not None:
-                gate_path = f"{name}.to_gate_fine" if name else "to_gate_fine"
-                if gate_path not in loaded_linears:
-                    module.init_gate_fine_default()
-
     def post_load_weights(self) -> None:
         """Call post_load_weights on all Linear modules and convert embedders to target dtype."""
         # Convert condition_embedder components to target dtype

tensorrt_llm/_torch/visual_gen/pipeline_loader.py

@@ -23,6 +23,9 @@
 
 from tensorrt_llm._torch.autotuner import autotune
 from tensorrt_llm._torch.models.modeling_utils import MetaInitMode
+from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import (
+    CUTE_AVAILABLE,
+)
 from tensorrt_llm.llmapi.utils import download_hf_model
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.args import VisualGenArgs
@@ -225,8 +228,6 @@ def load(
             and _sa_cfg is not None
             and getattr(_sa_cfg, "algorithm", None) == "vsa"
         ):
-            from .cute_dsl_kernels.blackwell.video_sparse_attention import CUTE_AVAILABLE
-
             kernel_path = "CuTe DSL block-sparse" if CUTE_AVAILABLE else "dense SDPA fallback"
             logger.info(
                 f"Attention backend: CUTEDSL (algorithm=vsa, "