[None][refactor] Refactor Skip Softmax Attention Interface

[bobboli]

Summary

1. Split sparse attention user configs from backend runtime params so AttentionBackend instances store lowered SparseParams, not public config objects.
2. Centralize Skip Softmax Attention backend helpers: ModelOpt checkpoint parsing, calibration formula evaluation, scheduler state, kernel params, and SkipSoftmaxParams.
3. Lower LLM and VisualGen Skip Softmax Attention configs through to_sparse_params() before backend construction, including ModelOpt ignore handling.
4. Keep VisualGen checkpoint calibration per model component, so multi-transformer pipelines resolve each component from its own config.json.
5. Add VisualGen normalized timestep cutoff support through disabled_until_timestep and include the resulting Skip Softmax Attention phase in CUDA graph keys.
6. Refresh sparse attention docs and reorganize LLM/VisualGen tests around Python API, YAML API, checkpoint metadata, layer filtering, and timestep cutoff behavior.

Details

1. Sparse runtime params boundary

   • Adds the sparse backend base carriers SparseParams and SparseMetadataParams: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/sparse/params.py#L25
   • TrtllmAttention accepts and stores lowered sparse runtime params: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/trtllm.py#L1114
   • create_attention() forwards sparse_params into backend constructors instead of exposing user config to backend instances: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/utils.py#L77

2. Shared Skip Softmax Attention backend logic

   • Centralizes ModelOpt checkpoint parsing for sparse_attention_config.config_groups, allowing at most one Skip Softmax Attention group: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py#L60
   • Implements formula parsing and evaluation for calibrated target_sparsity to threshold_scale_factor lowering: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py#L124
   • SkipSoftmaxScheduler resolves concrete kernel params and handles VisualGen timestep-based disablement: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py#L282
   • SkipSoftmaxParams is the runtime object passed to the attention backend: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py#L420

3. LLM lowering path

   • LLM SkipSoftmaxAttentionConfig validates target_sparsity, resolves checkpoint defaults, lowers to SkipSoftmaxParams, and applies checkpoint ignore patterns: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/llmapi/llm_args.py#L668
   • LLM attention modules lower config before backend construction and pass layer identity for checkpoint ignore matching: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/modules/attention.py#L531
   • ModelEngine still owns the global sparse config, but metadata/cache construction receives lowered metadata params rather than backend params: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/pyexecutor/model_engine.py#L1661

4. VisualGen per-layer lowering and checkpoint calibration

   • VisualGen keeps a separate user-facing SkipSoftmaxAttentionConfig for runtime knobs such as threshold_scale_factor, target_sparsity, and disabled_until_timestep: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/visual_gen/sparse_attention.py#L46
   • VisualGen lowering combines user config with each component model's checkpoint metadata and returns either None for ModelOpt-ignored layers or SkipSoftmaxParams for enabled layers: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/visual_gen/sparse_attention.py#L71
   • VisualGen attention modules pass the layer module_name and component pretrained_config into lowering: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/modules/attention.py#L203
   • Component-relative ignore patterns are supported for multi-transformer checkpoints such as transformer and transformer_2: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/visual_gen/sparse_attention.py#L107

5. VisualGen timestep cutoff and CUDA graph keys

   • VisualGen CUDA graph keys are structured as named key parts and can include registered non-shape key functions: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py#L12
   • BaseDiffusionModel registers the Skip Softmax Attention graph phase from the model-forward timestep: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/models/modeling.py#L38
   • VisualGen TRTLLM attention backend consumes timestep as an internal scheduling signal and does not pass it to the raw attention op: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py#L291
   • WAN and Cosmos normalize scheduler timesteps before model forward, while WAN rescales internally for its embedding convention: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py#L554 and https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py#L691

6. Docs and tests

   • LLM sparse-attention docs were refocused around the public LLM API and ModelOpt checkpoint config shape: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/docs/source/features/sparse-attention.md#L1
   • Framework-level sparse-attention implementation notes live in the developer guide: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/docs/source/developer-guide/sparse-attention-development-guide.md#L1
   • VisualGen sparse-attention docs are added as an unindexed draft page for the future VisualGen doc hub: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/docs/source/visual-gen/features/sparse-attention.md#L1
   • VisualGen Skip Softmax Attention tests now cover User API parsing, checkpoint lowering, ModelOpt ignore patterns, timestep cutoff behavior, and per-model checkpoint metadata: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py#L78
   • LLM tests now cover Python/YAML API parsing, target sparsity validation, checkpoint formulas/defaults, ignore patterns, and direct sparse checkpoint config lowering: https://github.com/bobboli/TensorRT-LLM/blob/bli/skip-softmax-decouple-llm-vg/tests/unittest/llmapi/test_llm_args.py#L2359

Follow-ups

• ModelOpt sparse-attention checkpoints use ignore in config.json; this PR treats that as the supported sparse layer-disable path. Existing hf_quant_config.json / exclude_modules handling should remain a quantization compatibility concern, not the primary sparse-attention interface.
• The VisualGen sparse-attention doc is intentionally not indexed yet. It is a draft location for colleagues to continue VisualGen-specific sparse-attention documentation after this PR lands.

Test plan

• Passed: pre-commit run --files tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py tensorrt_llm/_torch/modules/attention.py tensorrt_llm/llmapi/llm_args.py tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py tests/unittest/llmapi/test_llm_args.py
• Passed: commit-time pre-commit hooks for the latest commit.
• Earlier syntax checks passed on the updated test files with python3 -m py_compile and git diff --check.
• Not run locally: targeted pytest suites, because this local shell does not have the full TensorRT-LLM pytest/runtime dependency stack.
• CI on this branch is pending.

[bobboli]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #54722 [ reuse-pipeline ] triggered by Bot. Commit: 6817fb2 Link to invocation

[tensorrt-cicd]

PR_Github #54722 [ reuse-pipeline ] completed with state SUCCESS. Commit: 6817fb2
Reusing PR_Github #54606 for commit 6817fb2

Link to invocation