[None][perf] Add WAN skip-softmax recipe config

[karljang]

Summary

• Add an 8-GPU BF16 WAN skip-softmax recipe config using TRTLLM attention, target sparsity 0.75, and 16 dense warmup steps.
• Wire first_dense_steps into the VisualGen skip-softmax path so early denoising steps can run dense before sparse attention is enabled.
• Keep the recipe on the GI path for WAN GELU: block-level torch.compile remains enabled and GELU stays in the compiled transformer block path. Self-attention uses the existing fused QK-Norm + RoPE path.
• Keep the recipe config aligned with the current ModelOpt flow: use a checkpoint-root sparse.yaml for calibration/disabled-layer data, with sparse_config_path remaining as an advanced fallback.

Follow-up measurement

• Detailed LPIPS and end-to-end timing results for this recipe will be added after the final validation sweep/report is ready.

Testing

• pre-commit run --files <changed files>
• python3 -m py_compile <changed Python files>
• git diff --check origin/main...HEAD

Focused pytest was not run in this non-Docker shell because mpi4py could not load libmpi.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces step-based skip-softmax scheduling for visual generation by adding a first_dense_steps parameter that controls when sparse attention becomes active during denoising. The changes span configuration schema updates, pipeline integration logic, tensor shape normalization for fused operations, torch.compile stability fixes, and comprehensive test coverage with example configurations.

Changes

Skip-Softmax Scheduling for Visual Generation

Layer / File(s)	Summary
Skip-softmax configuration schema and overrides tensorrt_llm/visual_gen/sparse_attention.py	Adds first_dense_steps field to SkipSoftmaxConfig with validation and threading through factory methods and overrides. Refactors ModelOpt prefill coefficient parsing into _formula_kwargs_from_prefill helper to normalize and validate a/log_a consistency. Introduces set_skip_softmax_enabled function to conditionally apply sparse configs based on enable state and layer threshold resolution.
Pipeline skip-softmax step scheduling tensorrt_llm/_torch/visual_gen/pipeline.py	Adds _set_skip_softmax_for_step method to dynamically toggle skip-softmax based on step index relative to first_dense_steps. Tracks state with _skip_softmax_enabled field. Validates sparse config during CUDA graph setup to reject first_dense_steps > 0 when graph capture is enabled. Enables inductor precision-cast emulation during torch compilation.
RoPE tensor shape helpers for fused operations tensorrt_llm/_torch/visual_gen/modules/attention.py	Adds static helpers _reshape_rope_for_fused_op and _rope_tokens_per_batch to normalize RoPE frequency tensors into flattened 2D form and compute token counts accounting for RoPE row cardinality. Updates apply_packed_qk_norm_rope and apply_split_norm_rope to use these helpers instead of relying on fused ops to handle arbitrary input shapes.
WAN model GELU eager-path control for torch.compile tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py	Adds _wan_gelu_tanh_eager helper decorated with @torch.compiler.disable to ensure GELU tanh approximation stays on eager ATen path. Routes WanImageEmbedding and WanBlock MLP activations through this helper instead of inline F.gelu calls.
Example configuration and documentation examples/visual_gen/configs/wan2.2-t2v-bf16-skip-softmax-8gpu.yaml, docs/source/features/sparse-attention.md	Adds new 8-GPU Wan 2.2 T2V BF16 example configuration with skip-softmax (target_sparsity: 0.75, first_dense_steps: 16), torch compile enabled, and CUDA graphs disabled. Updates sparse-attention documentation to distinguish TRT-LLM runtime knobs from ModelOpt calibration artifacts and explain first_dense_steps role in dense-to-sparse transition.
Unit tests for configuration, scheduling, and example configs tests/unittest/_torch/visual_gen/test_attention_integration.py, tests/unittest/_torch/visual_gen/test_skip_softmax_config.py, tests/unittest/_torch/visual_gen/test_visual_gen_args.py	Validates RoPE reshape and token-per-batch helpers across input layouts; tests first_dense_steps preservation through config construction and overrides; verifies set_skip_softmax_enabled enables/disables sparse configs on model modules; validates ModelOpt coefficient parsing with a/log_a consistency; loads and validates example YAML configuration.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

VisualGen

Suggested reviewers

• arysef
• Shixiaowei02
• o-stoner
• NVShreyas
• chang-l

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% 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	Title is specific and directly related to the main change: adding a WAN skip-softmax recipe configuration for performance optimization.
Description check	✅ Passed	Description covers key aspects (summary, testing, follow-up measurement) but lacks detail in some template sections and uses less formal structure than the template suggests.
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

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/visual_gen/sparse_attention.py (1)

379-386: 💤 Low value

Collapse the duplicated auto_detect branches.

Both branches now build the formula identically via _formula_kwargs_from_prefill(prefill); the only thing that differs is the guard. They can be merged into a single check.

♻️ Proposed simplification

-    if "log_a" in prefill and "b" in prefill:
-        return SkipSoftmaxConfig.with_calibration(
-            formula=SkipSoftmaxFormula(**_formula_kwargs_from_prefill(prefill)),
-        )
-    if "a" in prefill and "b" in prefill:
-        return SkipSoftmaxConfig.with_calibration(
-            formula=SkipSoftmaxFormula(**_formula_kwargs_from_prefill(prefill)),
-        )
-
-    return None
+    if "b" in prefill and ("log_a" in prefill or "a" in prefill):
+        return SkipSoftmaxConfig.with_calibration(
+            formula=SkipSoftmaxFormula(**_formula_kwargs_from_prefill(prefill)),
+        )
+
+    return None

🤖 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/visual_gen/sparse_attention.py` around lines 379 - 386, Two
duplicated branches check prefill guards then call
SkipSoftmaxConfig.with_calibration with the same SkipSoftmaxFormula built from
_formula_kwargs_from_prefill(prefill); collapse them into one branch. Replace
the two separate if blocks with a single condition that tests ("log_a" in
prefill or "a" in prefill) and "b" in prefill, and return
SkipSoftmaxConfig.with_calibration(formula=SkipSoftmaxFormula(**_formula_kwargs_from_prefill(prefill))).
This keeps the same behavior while removing the duplicated construction and
guards.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/visual_gen/sparse_attention.py`:
- Around line 379-386: Two duplicated branches check prefill guards then call
SkipSoftmaxConfig.with_calibration with the same SkipSoftmaxFormula built from
_formula_kwargs_from_prefill(prefill); collapse them into one branch. Replace
the two separate if blocks with a single condition that tests ("log_a" in
prefill or "a" in prefill) and "b" in prefill, and return
SkipSoftmaxConfig.with_calibration(formula=SkipSoftmaxFormula(**_formula_kwargs_from_prefill(prefill))).
This keeps the same behavior while removing the duplicated construction and
guards.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 86690e3e-691b-4bcc-9474-7b55216231c0

📥 Commits

Reviewing files that changed from the base of the PR and between 35cce74 and 93be8cd.

📒 Files selected for processing (9)

• docs/source/features/sparse-attention.md
• examples/visual_gen/configs/wan2.2-t2v-bf16-skip-softmax-8gpu.yaml
• tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py
• tensorrt_llm/_torch/visual_gen/modules/attention.py
• tensorrt_llm/_torch/visual_gen/pipeline.py
• tensorrt_llm/visual_gen/sparse_attention.py
• tests/unittest/_torch/visual_gen/test_attention_integration.py
• tests/unittest/_torch/visual_gen/test_skip_softmax_config.py
• tests/unittest/_torch/visual_gen/test_visual_gen_args.py

[karljang]

Closing in favor of #14687