[TRTLLM-15917][feat] Integrate Sol-Attn sparse attention into VisualGen

[karljang]

Description

Adds Sol-Attn (arXiv:2607.24027) as a third sparse-attention algorithm for VisualGen, alongside skip_softmax and VSA. It folds dynamic block routing, sparse computation, and an approximation-correction term into a single online-softmax pass.

Config surface. SolAttnAttentionConfig (tensorrt_llm/visual_gen/args.py, tensorrt_llm/visual_gen/sparse_attention.py):

Field	Meaning
tau	per-block routing threshold; higher routes more blocks sparse
thresh_type	diag / exact threshold policy
kv_splits	must be 1 on the shipped architectures
disabled_until_timestep	dense-prefix cutoff on the normalized timestep; dense while t >= cutoff
dense_layers	comma/range layer-skip spec, e.g. 0,2-4

disabled_until_timestep deliberately mirrors skip-softmax's field of the same name and sense: the layer runs dense while the normalized timestep is at or above the cutoff, and switches to the sparse kernel below it. The value arrives as a forward kwarg that tensorrt_llm/_torch/visual_gen/modules/attention.py already threads to every backend and every VisualGen pipeline normalizes by num_train_timesteps, so there is no per-pipeline wiring and no process-wide state — tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py is untouched. Dispatch goes through create_attention exactly as skip_softmax and vsa do. Cross-attention (SEPARATE_QKV) falls back to VANILLA; context-parallel (cp_size > 1) and quant_attention_config are both rejected, mirroring VSA's existing guards — Sol-Attn replaces the dense CuTeDSL path, so a quantized-attention request would otherwise be silently ignored.

Kernel scope, and what is deliberately not here

The kernel is vendored from its reference implementation; tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md records the upstream pin, how it was reconstructed, and its currency check against the current upstream tip.

Only the two architectures with hardware evidence behind them are carried: sm100 (B200/GB200) and sm120 (RTX Blackwell). Upstream's sm89 and sm90 kernels and its Triton reference attention path are not included. sm90 in particular would have added ~8,900 lines, since it pulls in 13 further FlashAttention helper files plus a shim standing in for upstream's quack dependency. sm90 covers H100/H200/GH200 and matters — it should return in a follow-up with measurements behind it, rather than ship unvalidated here.

Every vendored file carries an SPDX Apache-2.0 header naming its NVlabs/Sana origin; the two that derive from FlashAttention additionally cite BSD-3-Clause and point at sm100/LICENSE.flash-attention, and the cuDNN Frontend license the SM120 kernel adapts is vendored at sm120/LICENSE.cudnn-frontend at the commit the notices cite.

Upstream also vendors a copy of FlashAttention's CuTe DSL helpers. That copy is not carried. TensorRT-LLM already depends on flash-attn-4 (requirements.txt), which provides the same flash_attn.cute modules, so the kernels import them from that dependency. This was verified on B200: swapping the vendored modules for the installed ones produced bit-identical output across a 12-point (shape, tau) sweep, and afterwards no module resolved out of the vendored tree.

Two things worth knowing that the code makes easy to miss:

• tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/preprocess.py implements the routing/threshold stage in Triton, so Triton is a required runtime dependency on every Sol-Attn path, not only a fallback.
• The vendored tree is linted and formatted to this repository's style rather than kept byte-identical to upstream — no lint exclusion is added, and this PR does not modify pyproject.toml. The consequence is recorded in tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md: a direct diff against upstream now shows formatting noise (upstream wraps at ~80 columns, this repo at 100), so a future currency check should normalize both sides with ruff format before comparing.

Failure behaviour

Inputs the kernel cannot serve — an unsupported architecture, head_dim other than 128, a non-bf16 dtype, or mismatched k/v — fall back to dense SDPA with a warning_once naming the specific reason, and increment dense_fallback_calls alongside kernel_calls. Kernel exceptions take the same path. SOL_ATTN_STRICT=1 raises instead, on both arms.

Because the dense prefix swaps kernels without changing tensor shapes, the two phases must not share a captured CUDA graph: register_cuda_graph_extra_key_fns registers sol_attn_phase from kwargs.get("timestep"), structurally identical to skip_softmax_phase. dense_layers needs no key, being fixed per layer at construction.

The eligibility fallback matters because it is what silently turns Sol-Attn into a no-op for an entire run: on an unsupported GPU or a head_dim-64 model the user would otherwise see only absent speedup, with nothing in the log. The counters are the mechanism by which a reviewer can confirm a benchmark actually exercised the kernel rather than dense SDPA.

Performance (ENGINEERING_CRITERIA.md §3.2)

Denoise time on B200, 50 steps, mean of 2 repetitions after 1 warmup, measured
both with and without CUDA graphs. $S = T_{base}/T_{new}$ is the speedup ratio;
$\Delta = (1 - T_{new}/T_{base}) \times 100$ is the time reduction. Baseline is
dense CuTeDSL (not VANILLA), so the comparison isolates sparsity from the
backend choice.

Model	graphs	baseline	Sol-Attn	$S$	$\Delta$
Wan2.2-TI2V-5B (704×1280, 121f)	off	66.90 s	59.38 s	1.127×	11.2 %
Wan2.2-TI2V-5B	on	67.53 s	56.29 s	1.200×	16.7 %
Wan2.2-T2V-A14B (720×1280, 81f)	off	593.92 s	409.34 s	1.451×	31.1 %
Wan2.2-T2V-A14B	on	593.86 s	422.26 s	1.406×	28.9 %

Run-to-run spread was under 0.06 % on every cell (5B ±0.045 s, A14B ±0.22 s),
so the differences are well outside noise.

Enabling CUDA graphs helps 5B and slightly hurts A14B, and the reason is not
established. On 5B, removing per-launch overhead lifts $S$ from 1.127 to
1.200. On A14B the baseline is unchanged (593.92 → 593.86 s — at ~594 s of
denoise, launch overhead is negligible) while the Sol-Attn arm is 3.2 % slower
(409.34 → 422.26 s). A plausible explanation is that a non-null
disabled_until_timestep forces two captured graphs instead of one, and that
fixed cost is not amortized on the larger model — but that is a hypothesis, not
a measurement, and it should be confirmed before anyone treats graphs-on as the
recommended A14B configuration. The best A14B number remains the graphs-off
1.451×.

Accuracy (lossy work — ENGINEERING_CRITERIA.md §3.3)

Evaluated on two models with an internal matched-protection sweep, gated on worst-prompt LPIPS ≤ 0.25. $S$ is the speedup ratio $T_{base}/T_{new}$; $\Delta$ is time reduction $(1 - T_{new}/T_{base}) \times 100$.

Model	Best certified point	$\Delta$ (denoise)	$S$	worst-prompt LPIPS
Wan2.2-TI2V-5B	tau=2.0, disabled_until_timestep=0.9545	11.24 %	1.13×	0.1120
Wan2.2-T2V-A14B	tau=2.0, disabled_until_timestep=0.9090	29.63 %	1.42×	0.2286 (91.4 % of gate)

Compared head-to-head against skip_softmax at matched dense-prefix protection, Sol-Attn wins on both models — on A14B, $S = 1.42\times$ versus skip_softmax's $1.09\times$.

The operating points were originally tuned as step counts (10 and 12 dense steps). Replicating tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py's scheduler setup — checkpoint scheduler config, set_timesteps(50), normalized by num_train_timesteps — the equivalent cutoffs are 0.9545 and 0.9090. Each sits at the midpoint between consecutive timesteps, so at 50 steps the dense/sparse partition is unchanged and these numbers carry over; unlike a step count, a cutoff keeps its meaning when num_inference_steps changes.

One tuning constraint worth recording: A14B's high-noise/low-noise expert switch is at boundary_ratio = 0.875 (model_index.json), just below the 0.9090 cutoff. The sparse phase therefore begins on the high-noise expert, and lowering the cutoff much further would move the prefix boundary across that switch.

These numbers were measured on the pre-reduction tree. They carry over unchanged because the reduction is bit-identical on sm100 (see below), not because they were re-derived. Full LPIPS-gated results, per-axis sweeps, and a video comparison against the dense reference exist as internal study artifacts and can be shared directly with reviewers on request.

All figures are real GPU evidence (B200, real checkpoints, no synthetic proxies), captured through an evaluation harness with its own integrity checks — backend census, LPIPS gate, denoise-spread. None of this PR's code depends on that harness.

Test Coverage

New: tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py, registered in tests/integration/test_lists/test-db/l0_b200.yml (sm100) and tests/integration/test_lists/test-db/l0_gb202.yml (sm120) so it runs in CI on both shipped architectures. It mirrors tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py in structure and scope: backend-factory dispatch, cross-attention VANILLA fallback, context-parallelism rejection, GQA/MQA rejection, the dense_layers guard, the dense-prefix phase semantics at and either side of the cutoff (including tensor-valued timesteps, which is what pipelines actually pass), fail-open on a missing timestep, both CUDA-graph key cases, kernel-eligibility reasons, SOL_ATTN_STRICT on the eligibility path, dense-fallback numerics and counters, quantized-attention rejection, and arch-list drift between SUPPORTED_ARCHS and _CUTE_BACKENDS. 32 tests.

Hardware validation:

• B200 (sm100) — 32/32 passed in this PR's suite (1 documented skip), and 69 passed overall alongside tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py, which [TRTLLM-15585][feat] Wire SkipSoftmax sparse attention into the CuTeDSL backend #17781 extended — a direct check that this PR's edits to the shared config and CUDA-graph-key paths did not regress that work. Measured on the pushed commit, against a working clone rebuilt on the current base. Kernel output bit-identical to the pre-reduction tree across a 12-point (shape, tau) sweep, and an end-to-end Wan2.2-TI2V-5B generation produced a byte-identical MP4, zero differing pixels. kernel_calls=12, dense_fallback_calls=0, SOL_ATTN_STRICT=1. pre-commit run is clean across all 33 changed files.
• RTX 5090 (sm120) — 14/14 pass (before the graph-key tests were added); resolves to cute_sm120; 9/9 sweep points ran with no dense fallback.

Gap, stated plainly: sm120 has kernel-level evidence only. An end-to-end generation was not possible on the available sm120 hardware — a 5090's 32 GB cannot hold Wan2.2-TI2V-5B (OOM at 27.2 GiB during model load), and no larger sm120 node was available. The untested delta is pipeline integration, which is architecture-independent Python already covered on B200; the architecture-specific part is the kernel, and that is covered. Reviewers who consider that insufficient should say so — dropping sm120 to an sm100-only PR is a small change.

The sm120 figures below predate the rebase and were not re-measured: the local build targets 90-real;100-real, and the only available sm120 hardware (a 32 GB RTX 5090) cannot hold the models used here. The rebase did not touch the vendored kernel tree, so those results should carry over, but that is an inference rather than a measurement.

One test is skipped with a documented reason: GPU kernel-vs-dense numerical equivalence at full routing (the analogue of VSA's test_cute_kernel_matches_dense_at_full_topk). It needs the exact tau/thresh_type combination that guarantees non-sparse routing, which is not simply tau=0 because Sol-Attn's routing is score-derived rather than a plain top-k.

Notes for reviewers

• Docs. docs/source/visual-gen/features/sparse-attention.md gains a sol_attn row and a section covering the YAML surface, the sm100/sm120 + head_dim=128 + bf16 + MHA constraints, the cutoff semantics, and the fallback/SOL_ATTN_STRICT behaviour. Its claim that VSA is the only CUTEDSL algorithm mutually exclusive with quantized attention is corrected, since Sol-Attn now is too.
• API change. This adds one new public sparse_attention_config variant, SolAttnAttentionConfig. It is additive, so api-compatible. Per ENGINEERING_CRITERIA.md §2.1 and §3.3, both the new public API surface and the lossy-work accuracy study need a team-sync review, not only PR approval.
• No new dependencies. flash-attn-4 and triton are already in requirements.txt; this PR consumes them rather than adding anything. The vendored subset ships LICENSE.flash-attention and retains upstream copyright headers; adapted files carry NVIDIA SPDX headers with attribution.
• Scope deliberately excluded. An AttentionConfig.cross_attention_backend override was developed alongside this work and is not included. It applies to all backends, is unrelated to Sol-Attn, and needs its own review — in particular whether it should be allowed to bypass the cross-attention safety fallback. It will follow as a separate PR.
• Follow-ups. Restoring sm90 with H100/H200 evidence, and a per-expert-calibrated checkpoint for dual-expert MoE models (A14B), are future work.