[None][feat] add VisualGen VSA and SOL sparse attention on PrimTS block-sparse FMHA

docs/source/features/visualgen-sparse-attention.md

@@ -6,56 +6,23 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti
 
 - [Overview](#overview)
   - [Algorithms](#algorithms)
-  - [Sol-Attn](#sol-attn)
 - [Skip Softmax Attention](#skip-softmax-attention)
+- [SOL Attention](#sol-attention)
 - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa)
 
 ## Overview
 
 Visual generation models naturally operate on long image or video token sequences. Each denoising step is closer to a full-context prefill pass than to autoregressive decoding, and attention can dominate runtime for high-resolution image generation or long video generation.
 
-Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config. Checkpoint calibration metadata remains internal and is lowered into per-attention-backend `SparseParams` when each attention module is constructed.
+Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config, while `attention_config.backend` selects the kernel family. Algorithms produce their block-sparse routes through the core `block_sparse_attn_predict` hook: a backend either predicts inside that hook from the flattened Q/K/V, or predicts before the core forward and hands the complete `BlockSparseForwardInputs` through `AttentionForwardArgs.sparse_backend_args`, which the default hook passes through. The VisualGen TRTLLM wrapper owns the timestep schedule those backends consult: it prepares the denoising timestep once per eager call and answers whether a layer runs sparse for it (dense layers and dense timestep phases do not). `SparseRuntimeParams` is the single lowered runtime carrier passed as `AttentionForwardArgs.sparse_runtime_params`; its optional `block_sparse_inputs` field nests the algorithm-neutral routes for the general block-sparse FMHA. `None` means prediction has not run, while an empty `SparseRuntimeParams()` records that prediction ran without a sparse payload.
 
 ### Algorithms
 
 | `algorithm` | Config class | Status |
 |---|---|---|
 | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported |
-| `vsa` | `VideoSparseAttentionConfig` | Supported (CUTEDSL) |
-| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) |
-
-### Sol-Attn
-
-Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block
-routing, sparse computation, and an approximation-correction term into one
-online-softmax pass. It runs on the **CUTEDSL** backend only, on datacenter
-Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- and requires
-`head_dim=128`, bfloat16, and MHA
-(`num_kv_heads == num_heads`).
-
-```yaml
-attention_config:
-  backend: CUTEDSL
-  sparse_attention_config:
-    algorithm: sol_attn
-    tau: 2.0                        # routing threshold; higher routes more blocks sparse
-    thresh_type: diag               # or "exact"
-    disabled_until_timestep: 0.9090 # dense while normalized timestep >= cutoff
-    dense_layers: [0]               # optional: layer indices forced dense
-```
-
-`disabled_until_timestep` has the same meaning as it does for Skip Softmax:
-attention runs dense while the normalized denoising timestep is at or above the
-cutoff, protecting the high-noise prefix, and switches to the sparse kernel
-below it. Use `None` rather than `0.0` to disable the prefix.
-
-On an input the kernel is known not to serve — an unsupported architecture, a
-`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn runs dense
-attention instead (the configured backend's dense kernel where available,
-torch SDPA otherwise), logs the specific reason once, and counts the fallback.
-Set `TRTLLM_SOL_ATTN_STRICT=1` to raise instead, which is useful when
-benchmarking to confirm the kernel actually ran. Errors raised by the kernel
-itself are not caught.
+| `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) |
+| `sol_attn` | `SolAttentionConfig` | Supported (`TRTLLM`, `CUTEDSL`) |
 
 ## Skip Softmax Attention
 
@@ -127,7 +94,7 @@ User configuration is supplied through Python or YAML and controls how the check
 
 `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers.
 
-Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn are currently the two supported sparse-attention algorithms, both available only through the CuTeDSL attention backend; quantized attention is not yet supported or enabled with either mode.
+Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and SOL cannot be combined with quantized attention on any backend.
 
 #### Mapping `disabled_until_timestep` to Actual Denoising Steps
 
@@ -251,6 +218,110 @@ attention_config:
 
 Graphs are captured lazily. The first denoising step seen for a given tensor shape and sparse-attention phase captures a graph; later steps with the same shape and phase replay that graph. When denoising crosses the cutoff, the phase key changes, so VisualGen captures a second graph for the enabled phase instead of replaying the graph from the disabled phase.
 
+## SOL Attention
+
+SOL ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) routes attention
+blocks dynamically: blocks whose scores stand out are computed exactly, the
+rest are approximated from compact K/V proxies and folded back into the online
+softmax. VisualGen serves it on two backends behind one `SolAttentionConfig`:
+
+- **TRTLLM** runs SOL in two stages. A TRT-LLM-owned predictor first produces
+  an exact block bitmask and K/V proxy summaries from Q/K/V; the shared
+  `PrimsTSBlockSparseFmha` library then executes that route from the shared
+  `SparseRuntimeParams`. `SOLTrtllmAttention` is only the VisualGen bridge: it
+  overrides the core `block_sparse_attn_predict` hook, so prediction runs
+  inside the core forward from the flattened Q/K/V, the batch layout in the
+  attention metadata and the `timestep` in the forward arguments; dense layers
+  and dense timestep phases return no routes. The VisualGen wrapper compacts
+  fused projection split views once and shares those tensors between
+  prediction and the block-sparse FMHA. Cross-attention, context parallelism,
+  attention quantization and unsupported tensor envelopes raise an error
+  instead of silently falling back to dense attention.
+- **CUTEDSL** runs the fused kernel vendored from the reference implementation
+  (`SOLCuTeDSLAttention`), which folds routing, sparse computation and the
+  approximation correction into one online-softmax pass. It runs on datacenter
+  Blackwell, sm100 (B200/GB200) and sm103 (B300/GB300), and requires
+  `head_dim=128`, bfloat16 and MHA (`num_kv_heads == num_heads`). On an input
+  the kernel is known not to serve, such as an unsupported architecture, a
+  `head_dim` other than 128 or a non-bfloat16 dtype, it runs dense attention
+  instead (the CuTe DSL FMHA where available, torch SDPA otherwise), logs the
+  reason once and counts the fallback; set `TRTLLM_SOL_ATTN_STRICT=1` to raise
+  instead, which is useful when benchmarking to confirm the kernel actually
+  ran. Errors raised by the kernel itself are not caught. Cross-attention and
+  masked calls are delegated to dense attention per call.
+
+Configure SOL with `SolAttentionConfig` and either backend:
+
+```python
+from tensorrt_llm.visual_gen import AttentionConfig, SolAttentionConfig
+
+attention_config = AttentionConfig(
+    backend="TRTLLM",  # or "CUTEDSL"
+    sparse_attention_config=SolAttentionConfig(
+        tau=1.0,
+        disabled_until_timestep=0.6,
+        dense_layers=[0, 2, 3, 4],
+    ),
+)
+```
+
+The equivalent YAML is:
+
+```yaml
+attention_config:
+  backend: TRTLLM                 # or CUTEDSL
+  sparse_attention_config:
+    algorithm: sol_attn
+    tau: 1.0                      # routing threshold; higher tau routes more blocks sparse
+    disabled_until_timestep: 0.6  # dense while the normalized timestep >= cutoff
+    dense_layers: [0, 2, 3, 4]    # optional: layer indices forced dense
+    thresh_type: diag             # block threshold policy: "diag" or "exact"
+```
+
+- `tau` is the routing threshold in standard deviations above the mean block
+  score; higher values route more blocks sparse.
+- `disabled_until_timestep` has the same meaning as it does for Skip Softmax
+  Attention: attention runs dense while the normalized denoising timestep is at
+  or above the cutoff, protecting the high-noise prefix, and switches to SOL
+  below it. Use `None` rather than `0.0` to disable the prefix.
+- `dense_layers` lists zero-based layer indices that always use dense
+  attention.
+- `thresh_type` selects how the routing threshold models the key blocks:
+  `diag` treats every key channel independently, `exact` uses the full key
+  covariance. Both backends implement both policies from the same per-block
+  statistics.
+
+The TRTLLM envelope is full-mask BF16 self-attention on SM100 or SM103 with 4-D
+BSHD Q/K/V tensors, equal Q/K/V shapes and head dimension 128. The predictor is
+one graph-visible operator that allocates its route and summary tensors per
+call, so it composes with CUDA Graph capture and with torch.compile, including
+`fullgraph`, without keeping any state between calls.
+
+When a cutoff is configured, VisualGen includes the dense-or-sparse phase in
+the CUDA Graph key. The TRTLLM attention metadata reduces the timestep to a host
+value during graph warmup, keeps it in the component attention state and reuses
+it during capture for every timestep-scheduled algorithm (Skip Softmax Attention
+and SOL), while the SOL predictor's outputs are allocated inside the captured
+graph and replayed with it; the CuTeDSL backends read the
+phase the CUDA Graph runner resolved for the graph key instead of the device
+tensor. Per-token timesteps, such as Wan I2V where the conditioning frame stays
+at timestep zero, reduce to their largest live value, so the schedule stays
+dense until every token is below the cutoff. A dense capture therefore cannot
+be reused for the sparse phase.
+
 ## Video Sparse Attention (VSA)
 
-TODO
+VSA combines a coarse mean-pooled branch with a top-K block-sparse fine branch. Select either `CUTEDSL` for the CuTe DSL kernel or `TRTLLM` for PrimTS block-sparse attention. If the selected sparse kernel is unavailable or the known VSA tensor envelope is not met, the fine branch uses the compact Q/K/V tensors with that backend's dense path. VSA cannot be combined with `quant_attention_config`.
+
+VSA retains shape-dependent metadata and route tensors so CUDA Graph replay can reuse stable addresses. A pipeline instance keeps one set of these tensors per distinct shape profile for as long as the CUDA Graphs that reference them, so their footprint grows with the number of served resolution/frame profiles exactly like the graphs do.
+
+Both VSA backends use the same VisualGen-owned predictor implementation, one
+instance per attention layer, and identical post-processing. The `TRTLLM` path runs the coarse stage before the core
+forward, hands the predicted `BlockSparseForwardInputs` (including the
+tile-padding validity bits only the VSA predictor knows) through
+`sparse_backend_args`, lets the default core prediction hook pass them to the
+general block-sparse FMHA, and then blends the fine and coarse outputs. Its
+compact dense fallback passes no sparse inputs, so the core runs dense attention
+and VSA post-processing still runs. `CUTEDSL` retains only its backend-specific
+fine-attention execution. The core FMHA registry owns the reusable block-sparse
+implementation rather than a VSA-specific lifecycle.

docs/source/models/visual-generation.md

@@ -360,12 +360,15 @@ args = VisualGenArgs(
 
 ### Video Sparse Attention (VSA)
 
-VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs a block-sparse CuTe kernel over only those blocks. The two outputs are blended with learned gates.
+VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs the selected backend's block-sparse kernel over only those blocks. The two outputs are blended with learned gates.
+
+VisualGen owns VSA route prediction and coarse/fine post-processing. With the `TRTLLM` backend, it nests the predicted routes in `SparseRuntimeParams.block_sparse_inputs` and passes those precomputed runtime parameters through the normal core attention forward. The core `PrimsTSBlockSparseFmha` owns the general block-sparse execution contract; it does not own VSA-specific prediction or blending.
 
 **Requirements:**
 - VSA-fine-tuned checkpoint: [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not have the learned VSA gates.
-- Blackwell GPU (sm_100+) for the CuTe JIT kernel. Falls back to dense SDPA on older hardware with no accuracy loss.
-- `CUTEDSL` attention backend.
+- `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected blocks through the generic PrimTS block-sparse FMHA contract.
+- A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its supported envelope, the fine branch uses the selected backend's compact dense path (`SDPA` for `CUTEDSL`, TRTLLM attention for `TRTLLM`).
+- VSA cannot be combined with `quant_attention_config`.
 - Not compatible with Ring attention or Attention2D (VSA does not produce per-split LSE). Ulysses is supported.
 
 **`vsa_sparsity`** controls the fraction of K/V blocks skipped in the fine branch (0.0 = dense, 0.9 = 90% blocks skipped). Higher sparsity gives more speedup at the cost of some quality.
@@ -379,7 +382,7 @@ from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionCo
 args = VisualGenArgs(
     model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers",
     attention_config=AttentionConfig(
-        backend="CUTEDSL",
+        backend="TRTLLM",  # Use "CUTEDSL" for the CuTe DSL fine-stage kernel.
         sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9),
     ),
 )
@@ -389,7 +392,7 @@ YAML (for use with `--visual_gen_args` or `trtllm-serve`):
 
 ```yaml
 attention_config:
-  backend: CUTEDSL
+  backend: TRTLLM  # CUTEDSL is also supported.
   sparse_attention_config:
     algorithm: vsa
     vsa_sparsity: 0.90

tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md

@@ -224,11 +224,13 @@ routes, and optional token-validity bits mask ragged KV tails. Plans contain
 only static format, proxy, geometry, and capacity choices; every run receives
 the live routes, summaries, validity bits, page tables, and sequence lengths.
 
-`PrimsTSBlockSparseFmha` owns its wrapper-plan cache by default. Integrations
-whose attention layers execute serially may explicitly bind a model-scoped
-cache to reuse graph-stable route workspaces across compatible layers. The
-cache must not be shared by concurrent forwards; each independent model
-component must own separate state.
+`PrimsTSBlockSparseFmha` keeps its wrapper-plan cache in the attention's
+`fmha_state` dict (`TrtllmAttention(fmha_state=...)`), which defaults to a
+private dict per layer. Integrations whose attention layers execute serially
+hand one dict to every layer of a component to plan each static profile once
+and reuse its graph-stable route workspace across compatible layers; the
+sharing granularity is the lifetime of that dict. A shared dict must not serve
+concurrent forwards, so each independent model component owns separate state.
 
 `TrtllmAttention.block_sparse_attn_predict(q, k, v, metadata, forward_args)`
 is the backend hook that produces this payload; `prepare_sparse_runtime_params`

tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py

@@ -190,29 +190,24 @@ def _uniform_seq_len_q(
 class PrimsTSBlockSparseFmha(PrimsTSFmha):
     """Contiguous context and fixed-Q paged generation block-sparse FMHA."""
 
+    PLAN_CACHE_KEY = "prims_ts_block_sparse"
+
     supports_block_sparse_inputs = True
 
     def __init__(self, attn: "TrtllmAttention") -> None:
         super().__init__(attn)
-        self._contiguous_wrappers: dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"] = {}
-        self._paged_wrappers: dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"] = {}
-
-    def bind_plan_cache(self, cache_state: dict[str, object]) -> None:
-        """Share planned wrappers with every adapter bound to ``cache_state``.
-
-        Attention layers that execute serially, such as the blocks of one
-        diffusion transformer, see identical static profiles. Binding them to
-        one model-scoped container plans each profile once and allocates its
-        route workspace once. Call before the first forward.
-        """
-
+        # Planned wrappers live in the attention's FMHA state. Layers constructed
+        # with one shared state, such as the blocks of one diffusion transformer,
+        # plan each static profile once and allocate its route workspace once;
+        # the default per-instance state keeps the caches private to this layer.
+        caches = attn.fmha_state.setdefault(self.PLAN_CACHE_KEY, {})
         self._contiguous_wrappers = cast(
             dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"],
-            cache_state.setdefault("contiguous_wrappers", {}),
+            caches.setdefault("contiguous_wrappers", {}),
         )
         self._paged_wrappers = cast(
             dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"],
-            cache_state.setdefault("paged_wrappers", {}),
+            caches.setdefault("paged_wrappers", {}),
         )
 
     def _is_supported(