[None][feat] Add SkipSoftmax sparse attention support for visual generation

[karljang]

Summary

• Wire SkipSoftmax (BLASST) kernel-level sparse attention into the visual generation pipeline
• Add BaseSparseAttentionConfig with discriminated union pattern (extensible for future algorithms) and SkipSoftmaxConfig
• Auto-detect skip_softmax from ModelOpt checkpoint config.json — no user config needed for calibrated checkpoints
• Accept both ModelOpt diffusion format (log_a) and LLM format (a) coefficients in the YAML / config — auto-normalized to log_a = log(a) at the Pydantic layer
• Self-attention uses TRTLLM with skip_softmax; cross-attention continues using VANILLA

Benchmark

WAN 2.2 T2V-A14B on B200, 720 × 1280, 81 frames, 50 inference steps (production recipe), torch_compile=True, mean across 10 stress prompts. Calibration: ModelOpt sparse-0506.yaml (formula + per-component disabled_layers). Quality measured as LPIPS-AlexNet against the same-precision baseline over all 81 frames; band thresholds are <0.30 same trajectory, 0.30–0.55 "scene preserved", >0.55 "diverged".

Per-precision speedup at three representative targets

In-precision $S = T_\text{baseline (same precision, compile on)} / T_\text{skip}$.

target_sparsity	BF16 $S$	FP8 $S$	NVFP4 $S$	LPIPS (BF16 / FP8 / NVFP4)
0.50	1.08×	1.09×	1.10×	0.45 / 0.46 / 0.45
0.70 (recommended ceiling)	1.15×	1.17×	1.19×	0.54 / 0.55 / 0.53
0.90 (past "scene preserved")	1.28×	1.31×	1.34×	0.71 / 0.71 / 0.67

Dense fill at 1pp resolution from 0.60 → 0.75 (omitted from this table) localises the LPIPS-0.55 cliff between t=0.71 (FP8) and t=0.73 (NVFP4); t=0.70 is the last universally safe ceiling (all three precisions LPIPS mean ≤ 0.547).

Joint speedup vs BF16 no-compile reference (638 s baseline)

Stack on top of BF16 no-compile	$S_{e2e}$
torch_compile=True only	1.18×
Compile + FP8 quant	1.33×
Compile + NVFP4 quant	1.33×
Compile + BF16 + skip-t=0.70	1.36×
Compile + FP8 + skip-t=0.70	1.55×
Compile + NVFP4 + skip-t=0.70	1.59×
Compile + NVFP4 + skip-t=0.90 (diverged quality)	1.78×

Factors compose roughly multiplicatively. The recommended t=0.70 configuration buys ~60 % over the previous-gen no-compile BF16 stack while keeping every prompt's LPIPS in the "scene preserved" band.

Kernel breakdown (nsys, baseline inference, no compile)

Category	% of GPU time
Self-attention (TRTLLM FMHA)	53.88 %
Element-wise	22.34 %
MLP / GEMM (nvjet)	16.71 %
Norm (RMSNorm / LayerNorm)	3.63 %
Cross-attention (VANILLA SDPA)	0.45 %

Self-attention dominates at 720p (quadratic in seq_len ≈ 75 K), which is why skip-softmax has this much headroom on visual-gen workloads. Cross-attention is negligible — not worth targeting.

Use Cases

Case 1: Normal HF checkpoint (no skip_softmax in config.json)

# 1a: Direct threshold — all layers get same value
attention:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: skip_softmax
    threshold_scale_factor: 5000.0

# 1b: Target sparsity without formula → error (need formula from checkpoint or user)

# 1c: Target sparsity with user-provided formula
attention:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: skip_softmax
    target_sparsity: 0.5
    formula:
      a: 0.0003
      b: 7.5

# 1d: With per-layer overrides (fnmatch patterns)
attention:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: skip_softmax
    threshold_scale_factor: 5000.0
    layer_overrides:
      "transformer_blocks.0*": 0            # disable for dense layer 0
      "single_transformer_blocks*": 8000.0  # aggressive for single-stream

Case 2: ModelOpt checkpoint (has calibrated a, b in config.json)

// ModelOpt writes this to checkpoint config.json:
{
  "sparse_attention_config": {
    "threshold_scale_factor": {
      "prefill": {"a": 7.93, "b": 8.61}
    }
  }
}

# 2a: Auto-enable — just load the checkpoint, no user config needed
#     Pipeline auto-detects formula and creates SkipSoftmaxConfig
attention:
  backend: TRTLLM
  # sparse_attention_config not needed — auto-detected from checkpoint

# 2b: User threshold overrides checkpoint formula
attention:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: skip_softmax
    threshold_scale_factor: 3000.0  # user wins over checkpoint

# 2c: User target_sparsity + checkpoint formula
attention:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: skip_softmax
    target_sparsity: 0.5  # uses a, b from checkpoint config.json

Priority Order

1. threshold_scale_factor (if set, used directly — no formula needed)
2. target_sparsity + user formula (user formula takes precedence)
3. target_sparsity + checkpoint formula (auto-detected from config.json)

Changes

File	Change
_torch/visual_gen/config.py	BaseSparseAttentionConfig, SkipSoftmaxConfig, SkipSoftmaxFormula with dual-format (log_a or a) validator, auto_detect_sparse_attention_config(), apply_skip_softmax_overrides(), cached get_or_resolve_threshold() (PrivateAttr), strict sparse.yaml / sparse.*.yaml auto-detect glob
_torch/visual_gen/attention_backend/trtllm.py	Accept sparse_attention_config in __init__; documents the post-construction dynamic-read contract
_torch/visual_gen/modules/attention.py	Resolve config → threshold via cache, thread to backend (no shared-config mutation)
_torch/visual_gen/pipeline_loader.py	Wire apply_skip_softmax_overrides() + auto-detect from checkpoint
tests/unittest/_torch/visual_gen/test_skip_softmax_config.py	43 unit tests covering all use cases + dual-format formula handling

Known Limitations

• Cross-attention: Falls back to VANILLA. TRTLLM C++ requires fused QKV for non-MLA attention (attentionOp.cpp:640). This is a TRT-LLM-wide limitation (also affects LLM models like MLLama). In WAN 2.2 the cost is negligible — cross-attn is 0.45 % of GPU time vs. self-attn at 53.88 % (nsys, 720p).

Test plan

• 43 unit tests: config construction, resolution priority, layer overrides, formula, discriminator, auto-detection, dual-format coefficients, use case scenarios
• WAN 2.2 T2V-A14B E2E on B200, production recipe (50 steps, torch_compile=True): in-precision $S=1.15/1.17/1.19\times$ at target_sparsity=0.70 for BF16/FP8/NVFP4, mean LPIPS ≤ 0.55 (universally safe)
• Joint $S=1.59\times$ vs BF16 no-compile reference for NVFP4 + compile + skip-t=0.70
• Dense fill (1pp resolution from 0.60 → 0.75 across 3 precisions × 10 prompts) localises LPIPS-0.55 cliff between t=0.71 and t=0.73
• nsys kernel breakdown at 720p — self-attn 53.88 %, cross-attn 0.45 %, ratio ~120:1
• SM90 and SM100 both working after upstream [None][bug] fix SM90 full-mask skip-softmax dispatch #13120 merged
• B200 smoke at target_sparsity=0.70, STEPS=10 after PrivateAttr caching refactor: reproduces $S=1.13\times$, enabled=72 / disabled=8 (YAML disabled_layers honored)

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added sparse attention configuration support for visual generation models with skip-softmax optimization
  • Enabled automatic detection and loading of sparse attention settings from YAML files and model checkpoints
  • Added flexible per-layer threshold overrides for granular control over sparsity behavior

• Tests

  • Added comprehensive test coverage for sparse attention configuration and threshold resolution

[chang-l]

++ @lfr-0531 @bobboli for viz

[karljang]

@bobboli , @chang-l , please review~ 👍

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds skip-softmax sparse attention configuration for diffusion visual-generation models. It introduces schema types for threshold formulas and per-layer overrides, auto-detection helpers for YAML and checkpoint metadata, backend parameter support, and pipeline integration to apply resolved thresholds to TRTLLM modules.

Changes

Sparse Attention Configuration and Integration

Layer / File(s)	Summary
Sparse Attention Data Models tensorrt_llm/_torch/visual_gen/config.py	Adds SkipSoftmaxFormula with LLM-to-diffusion format conversion, SkipSoftmaxConfig with cached threshold resolution and fnmatch-based layer-override matching, SparseAttentionConfig discriminated union, and extends AttentionConfig with sparse_attention_config and sparse_config_path fields.
Config Loading and Auto-Detection tensorrt_llm/_torch/visual_gen/config.py	Implements load_sparse_config_from_yaml to parse ModelOpt sparse YAML with disabled-layer handling and formula normalization, auto_detect_sparse_yaml to search checkpoint directories with collision detection, and auto_detect_sparse_attention_config to extract formula coefficients from checkpoint config.json metadata.
Config Application and Model Loading Integration tensorrt_llm/_torch/visual_gen/config.py	Implements apply_skip_softmax_overrides to walk model modules, resolve per-layer thresholds, and update TRTLLM backends. Extends DiffusionModelConfig.from_pretrained to resolve sparse config in precedence order (explicit YAML → auto-detected YAML → checkpoint metadata) with user override merging.
TRTLLM Backend Parameter Support tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py	Adds sparse_attention_config parameter to TrtllmAttention.__init__ and stores it as a runtime-swappable attribute for forward-time kernel use.
Attention Module Sparse Config Resolution tensorrt_llm/_torch/visual_gen/modules/attention.py	Resolves sparse config thresholds from user config and checkpoint metadata, constructs SkipSoftmaxAttentionConfig with prefill/decode thresholds when threshold is positive, and passes it to backend factory.
Pipeline Loader Override Application tensorrt_llm/_torch/visual_gen/pipeline_loader.py	Applies per-layer sparse attention overrides after weight loading by calling apply_skip_softmax_overrides when config contains SkipSoftmaxConfig with layer overrides.
Comprehensive Test Coverage tests/unittest/_torch/visual_gen/test_skip_softmax_config.py	Tests formula validation and normalization, config construction and threshold resolution, YAML/checkpoint loading and auto-detection, layer-override matching with wildcards, and backend module updates via mocked TrtllmAttention.

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 47.46% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main feature: adding SkipSoftmax sparse attention support for visual generation, matching the substantial changes across the codebase.
Description check	✅ Passed	The description is comprehensive and well-structured, covering summary, use cases, changes, test coverage, benchmarks with detailed data, and known limitations. However, the PR Checklist section ("Please check this after reviewing the above items") appears incomplete with unchecked items.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tensorrt_llm/_torch/visual_gen/config.py (1)

408-409: 💤 Low value

Lazy imports create circular dependency risk

The function imports TrtllmAttention and SkipSoftmaxAttentionConfig at call-time (lines 408-409) rather than at module level. This pattern typically indicates a circular import issue between config.py and the attention backend.

While functional, lazy imports add runtime overhead on each call and obscure the dependency graph. If the circular dependency can be resolved (e.g., by moving apply_skip_softmax_overrides to a separate module or using TYPE_CHECKING), that would be cleaner.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/config.py` around lines 408 - 409, The lazy
imports of TrtllmAttention and SkipSoftmaxAttentionConfig inside
apply_skip_softmax_overrides introduce circular-import risk and runtime
overhead; to fix, move those imports to the module top (module-level imports) or
guard them with typing.TYPE_CHECKING, or alternatively refactor
apply_skip_softmax_overrides into a new module (e.g.,
visual_gen/skip_softmax_utils) that can import TrtllmAttention and
SkipSoftmaxAttentionConfig without creating a cycle; update references to
apply_skip_softmax_overrides accordingly and ensure unit tests/imports still
pass.

tests/unittest/_torch/visual_gen/test_skip_softmax_config.py (1)

147-162: 💤 Low value

Optional: Silence Ruff RUF012 for test fixture constant

Ruff flags MODELOPT_CHECKPOINT as a mutable class attribute, but since this is a read-only test fixture, the warning is a false positive.

🔧 Optional fix to silence static analysis

+from typing import Final
+
 class TestUseCaseScenarios:
     ...
-    MODELOPT_CHECKPOINT = {
+    MODELOPT_CHECKPOINT: Final = {
         "sparse_attention_config": {
             ...
         }
     }

Alternatively, move the constant outside the class (module-level) if it's not logically tied to the test class.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/visual_gen/test_skip_softmax_config.py` around lines
147 - 162, MODELOPT_CHECKPOINT is being flagged by Ruff as a mutable class
attribute but it's a read-only test fixture; to silence RUF012 either move
MODELOPT_CHECKPOINT out of the test class to the module level (so it’s a true
constant) or keep it in the class and append a per-line ignore comment for Ruff
(RUF012) to the declaration; update the test file so references to
MODELOPT_CHECKPOINT still resolve and run tests to confirm no behavioral change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py`:
- Line 187: The parameter sparse_attention_config in the function/method defined
in trtllm.py lacks a type annotation; update its signature to include an
explicit type (for example sparse_attention_config: Optional[Dict[str, Any]] =
None) and add the required imports (from typing import Optional, Dict, Any) at
the top of the module; ensure you use the exact parameter name
sparse_attention_config and adjust any related callers or type checks if they
rely on None vs a typed mapping.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/config.py`:
- Around line 408-409: The lazy imports of TrtllmAttention and
SkipSoftmaxAttentionConfig inside apply_skip_softmax_overrides introduce
circular-import risk and runtime overhead; to fix, move those imports to the
module top (module-level imports) or guard them with typing.TYPE_CHECKING, or
alternatively refactor apply_skip_softmax_overrides into a new module (e.g.,
visual_gen/skip_softmax_utils) that can import TrtllmAttention and
SkipSoftmaxAttentionConfig without creating a cycle; update references to
apply_skip_softmax_overrides accordingly and ensure unit tests/imports still
pass.

In `@tests/unittest/_torch/visual_gen/test_skip_softmax_config.py`:
- Around line 147-162: MODELOPT_CHECKPOINT is being flagged by Ruff as a mutable
class attribute but it's a read-only test fixture; to silence RUF012 either move
MODELOPT_CHECKPOINT out of the test class to the module level (so it’s a true
constant) or keep it in the class and append a per-line ignore comment for Ruff
(RUF012) to the declaration; update the test file so references to
MODELOPT_CHECKPOINT still resolve and run tests to confirm no behavioral change.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 871ce345-e692-4912-8ab3-716766dddc17

📥 Commits

Reviewing files that changed from the base of the PR and between f0e9b51 and dade9b5.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py
• tensorrt_llm/_torch/visual_gen/config.py
• tensorrt_llm/_torch/visual_gen/modules/attention.py
• tensorrt_llm/_torch/visual_gen/pipeline_loader.py
• tests/unittest/_torch/visual_gen/test_skip_softmax_config.py

[bobboli]

BTW, could you add a subsection ####Skip Softmax Attention at https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/sparse-attention.md?plain=1#L341, to document all the supported ways of configuring Skip Softmax Attention (via YAML, via LLM API config, via checkpoint)?

[karljang]

BTW, could you add a subsection ####Skip Softmax Attention at https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/sparse-attention.md?plain=1#L341, to document all the supported ways of configuring Skip Softmax Attention (via YAML, via LLM API config, via checkpoint)?

Could you please check the updated sparse-attention.md?

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49248 [ run ] triggered by Bot. Commit: c1f3d6e Link to invocation

[tensorrt-cicd]

PR_Github #49248 [ run ] completed with state FAILURE. Commit: c1f3d6e
/LLM/main/L0_MergeRequest_PR pipeline #38918 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49339 [ run ] triggered by Bot. Commit: c1f3d6e Link to invocation

[tensorrt-cicd]

PR_Github #49339 [ run ] completed with state SUCCESS. Commit: c1f3d6e
/LLM/main/L0_MergeRequest_PR pipeline #38995 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chang-l]

It looks like skip-softmax calibration can come from both ModelOpt sparse.yaml and checkpoint config.json. Why do we need a dedicated sparse_config_path in addition to the normal config path? Is this sparse YAML a separate ModelOpt artifact, or is it supposed to be the same kind of YAML users pass to LLM/server config?

[karljang]

@chang-l ,
I'm waiting for the #14175 PR to be merged. Please review this PR in the meantime

[chang-l]

Thanks for the effort. Can we have a follow-up doc update PR to cover:

• How to enable skipSoftmax from both sources: the user’s server config YAML and the ModelOpt checkpoint config YAML, including how we merge them or which source takes precedence if conflict.
• Some of your skipSoftmax accuracy findings and perf trade-offs for AIGV models (referring to @yibinl-nvidia 's PR that adds the LPIPS metric)

[karljang]

Thanks for the effort. Can we have a follow-up doc update PR to cover:

• How to enable skipSoftmax from both sources: the user’s server config YAML and the ModelOpt checkpoint config YAML, including how we merge them or which source takes precedence if conflict.
• Some of your skipSoftmax accuracy findings and perf trade-offs for AIGV models (referring to @yibinl-nvidiaYL 's PR that adds the LPIPS metric)

Sure, I'll make a follow-up PR to cover those items. Thank you!

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50004 [ run ] triggered by Bot. Commit: ed616dc Link to invocation

[tensorrt-cicd]

PR_Github #50004 [ run ] completed with state FAILURE. Commit: ed616dc
/LLM/main/L0_MergeRequest_PR pipeline #39568 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50062 [ run ] triggered by Bot. Commit: ed616dc Link to invocation

[tensorrt-cicd]

PR_Github #50062 [ run ] completed with state SUCCESS. Commit: ed616dc
/LLM/main/L0_MergeRequest_PR pipeline #39619 completed with status: 'SUCCESS'

CI Report

Link to invocation

[zhenhuaw-me]

Hi @karljang , thanks you for rebasing on the VisualGenArgs PR and getting this PR merged! I apologize that didn't get chance to take a look before, but I have 2 general concerns which I think we should try to address:

1. The API needs rework. I can create a doc to align, but the general idea is that: tensorrt_llm/visual_gen dir is for API, all non-API should be in tensorrt_llm/_torch/visual_gen. Debug purpose knob should be added through env var rather than in VisualGenArgs.
2. No E2E testing. With the new commits we don't know if SkipSoftmax is works while we claim the support in the doc. We should try to add E2E testing (1 test is better than 0) if we are targeting to support as a product feature.

cc @chang-l in case he has other comments. Thanks!

[zhenhuaw-me]

It limits the scope to skip softmax

[karljang]

Sure, I'll address that in a follow up PR~

[zhenhuaw-me]

Why do we need SkipSoftmaxConfig for VisualGen instead of re-exporting SkipSoftmaxAttentionConfig while "there is no new user-facing fields"?

[karljang]

SkipSoftmaxConfig extends the parent with private calibration state: per-layer threshold overrides and formula-based resolution from ModelOpt's diffusion-format coefficients. The parent SkipSoftmaxAttentionConfig only handles uniform thresholds, it has no concept of per-layer disabled layers or component routing.
Still, I agree this class is internal machinery rather than user-facing API. I can move it to tensorrt_llm/_torch/visual_gen/ in a follow-up and keep only the parent SkipSoftmaxAttentionConfig as the user-facing type in the public dir.

[zhenhuaw-me]

What's the purpose of this config? It seems to me that we can auto-detect from the checkpoint and this knob sounds like a debug one?

[karljang]

The sparse_config_path supports the case where users have a separately produced ModelOpt calibration YAML they want to apply, because users don't have to fully calibrate their checkpoints. That said, if we want to minimize the public API surface, we could move this to an env var. Happy to discuss the right placement.

[zhenhuaw-me]

It seems to me that only the users that are as powerful as TRT-LLM/ModelOpt dev team knows how to use this knob? I'd suggest moving to an env var at this stage and bring to be API if we see explicit user requirements.