3 Commits

Include renames

Author	SHA1	Message	Date
3 people	c78e654744	Skip Softmax diffusion export (#1269) ... ### What does this PR do? Type of change: New Feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Adds HuggingFace `config.json` export of skip-softmax sparse-attention calibration for diffusion pipelines (e.g. Wan 2.2), on top of the base skip-softmax work. - **`_export_diffusers_checkpoint`** walks every `nn.Module` component of a diffusers pipeline, calls `export_sparse_attention_config`, and writes the result into that component's `config.json` under the `sparse_attention_config` key. The sparse config lives **only** in `config.json` — there is no standalone `sparse.yaml`. - **`export_sparse_attention_config`** emits a `config_groups` schema where each algorithm's parameters are nested inside its own group; only `config_groups` and `producer` are top-level: - skip-softmax group → `algorithm: "skip_softmax"`, `targets`, `ignore` (layers kept dense — e.g. cross-attention + first/last blocks), `initial_disabled_steps` (opt-in, user-set; emitted only when `> 0`), `threshold_scale_factor` (`a * exp(b * target_sparsity)`), and `target_sparsity`. - N:M group → `algorithm: "sparse_softmax"` with `sparsity_n`/`sparsity_m`, `dense_sink_tokens`, `dense_recent_tokens` flattened into the group. - **Deploy reader** (`modelopt/torch/sparsity/attention_sparsity/plugins/sparse_attn_config.py`) reads these per-group params back, keeping the export↔load round-trip consistent. - **Example wiring**: `examples/diffusers/sparsity/wan22_skip_softmax.py` gains `--export-dir`, `--skip-softmax-threshold`, and `--initial-disabled-steps`. `--export-dir` runs `export_hf_checkpoint(pipe, export_dir=...)` after calibration. - Updated `CHANGELOG.rst`. ### Usage ```bash python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --calibrate --target-sparsity 0.5 --calib-size 4 \ --initial-disabled-steps 5 \ --export-dir ./wan22_skip_softmax_ckpt ``` Resulting layout — a `config.json` per component, **no `sparse.yaml`**: ``` wan22_skip_softmax_ckpt/ ├── transformer/config.json # carries sparse_attention_config ├── transformer_2/config.json # carries sparse_attention_config ├── vae/ … text_encoder/ … tokenizer/ … scheduler/ … └── model_index.json ``` A representative `config.json` entry for a diffusion transformer: ```json "sparse_attention_config": { "config_groups": { "group_0": { "algorithm": "skip_softmax", "targets": ["WanAttention"], "ignore": ["blocks.0.attn1", "blocks.0.attn2", "…"], "initial_disabled_steps": 5, "threshold_scale_factor": { "formula": "a * exp(b * target_sparsity)", "prefill": {"a": 1443.49, "b": 4.30} }, "target_sparsity": {"prefill": 0.5} } }, "producer": {"name": "modelopt", "version": "0.45.0..."} } ``` The N:M variant adds a second group: ```json "group_1": { "algorithm": "sparse_softmax", "targets": ["WanAttention"], "sparsity_n": 2, "sparsity_m": 4, "dense_sink_tokens": 0, "dense_recent_tokens": 64 } ``` ### Testing - `tests/examples/diffusers_sparsity/test_sparsity.py`: baseline / triton-baseline / fixed-threshold runs of the Wan 2.2 example, plus a Python-API calibrate → **export** test asserting the nested `sparse_attention_config` (`threshold_scale_factor`, `target_sparsity`, `ignore`, `initial_disabled_steps`) and the absence of any `sparse.yaml`. - `tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py` and `test_sparse_attn_config.py`: unit coverage of the per-group export schema and the deploy-reader round-trip (writer nests → reader reads from groups → internal mtsa config unchanged). - Validated end-to-end on Wan 2.2 T2V-A14B: full 4-prompt / 40-step / 81-frame calibration; the exported checkpoint carries the nested schema in both `transformer` and `transformer_2` `config.json`, and runtime measurement shows ~47–49% tile sparsity at a 0.5 target. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ❌ The exported `sparse_attention_config` schema was renamed and nested per-group during 0.45.x development, and the loader reads only the new layout — checkpoints exported by earlier 0.45.x builds must be re-exported. No released version is affected. <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-08 16:02:41 -07:00
kaix-nv	c9098b63fb	[4/n] Add vLLM integration for modelopt sparse attention (#1127) ... ### What does this PR do? Type of change: New feature, new example, new tests, documentation. Adds vLLM integration for ModelOpt sparse attention with paged KV cache support. This PR extends the ModelOpt Triton flash attention path so K/V can be read directly from vLLM's paged KV cache through `block_table` lookup. This avoids gather-to-contiguous copies when serving exported sparse-attention checkpoints with vLLM. The vLLM integration swaps vLLM's `FlashAttentionImpl` with `ModelOptSparseAttentionImpl` after model load. The sparse configuration is read from the exported checkpoint's `config.json` `sparse_attention_config` block, written by `examples/llm_sparsity/attention_sparsity/hf_sa.py`. The restored checkpoint metadata supports: - calibrated skip-softmax metadata (`threshold_scale_factor`, `target_sparse_ratio`) - N:M sparse-softmax metadata (`sparsity_n`, `sparsity_m`) - dense token preservation metadata (`dense_sink_tokens`, `dense_recent_tokens`) The vLLM path uses ModelOpt Triton for sparse prefill launches. Decode-only launches, cascade/prefix-cache metadata, and launches without active sparse work delegate back to vLLM FlashAttention. ### Limitations - Sparse attention is enabled for sparse prefill only. - Decode-only launches currently fall back to vLLM FlashAttention. - Attention sinks from vLLM FlashAttention are rejected until the ModelOpt Triton path supports them. - CUDA graph capture is not validated with this sparse attention path yet; use `--enforce-eager`. - Quant-only serving remains covered by `vllm_serve_fakequant.py`. - Combined sparse attention + quantization serving is not handled by this launcher in this PR and is planned as follow-up work. ### Usage Export a checkpoint with calibrated skip-softmax and sparse24 metadata: ```bash python examples/llm_sparsity/attention_sparsity/hf_sa.py \ --pyt_ckpt_path /path/to/hf-model \ --sparse_attn skip_softmax_calib_sparse24 \ --target_sparse_ratio 0.5 \ --calib_samples 64 \ --calib_max_seqlen 16384 \ --calib_chunk_size 4096 \ --seq_len 2048 \ --export_dir /path/to/modelopt-skipsoftmax-sparse24-export ``` Serve the exported checkpoint with the vLLM sparse-attention launcher: ```bash PYTHONPATH=$PWD python examples/vllm_serve/vllm_serve_sparse_attn.py \ /path/to/modelopt-skipsoftmax-sparse24-export \ --tensor-parallel-size 8 \ --host 0.0.0.0 \ --port 8000 \ --trust-remote-code \ --enforce-eager ``` Send a request through the OpenAI-compatible endpoint: ```bash curl http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "/path/to/modelopt-skipsoftmax-sparse24-export", "messages": [{"role": "user", "content": "Explain sparse attention in one paragraph."}], "max_tokens": 128 }' ``` ### Testing GitHub CI on the latest commit is green: - DCO - code-quality - docs build / deploy preview - unit tests, including Linux, Windows, multi-version, partial-install, and launcher jobs - example tests - GPU tests, including required GPU gate - regression tests, including required regression gate - `codecov/project` Focused test coverage added/updated for this PR includes: - `tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py` - `tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_config.py` - `tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py` - `tests/unit/torch/sparsity/attention_sparsity/test_triton_skip_softmax.py` - `tests/gpu/torch/sparsity/attention_sparsity/test_vllm_plugin.py` - `tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py` - `tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py` - `tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py` - `tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py` Manual / NEL eval validation: - Served a ModelOpt exported sparse-attention checkpoint through `examples/vllm_serve/vllm_serve_sparse_attn.py`. - Launched RULER64K NEL evals on DFW with `coreai_nvfm_llm`. - Current partial RULER64K prediction scores, before final `results.yml` is written: - `skipsoftmax-only`: 98.59% over 4500 flushed samples - `skipsoftmax-r0.7`: 99.70% over 1000 flushed samples - `skipsoftmax-r0.9`: 99.70% over 1000 flushed samples ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - no new PIP dependency. - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - documentation, examples, and tests are updated for this vLLM integration path. ### Additional Information Follow-up work: - Validate and enable CUDA graph capture for the sparse vLLM path. - Add combined sparse attention + quantization serving once the combined path is tested. - Investigate whether skip-softmax should also be enabled during decode. --------- Signed-off-by: Kai Xu <kaix@nvidia.com>	2026-05-20 14:11:25 -07:00
jingyu-ml	feec81ad2b	Add the Skip softmax for diffusion (#1166) ... ### What does this PR do? Type of change: new feature, new example <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> <!-- Details about the change. --> ## Summary - Add skip-softmax sparse attention (BLASST) for diffusion models via dedicated Triton kernels — an inference kernel with tile skipping and a calibration kernel with vectorized multi-threshold sparsity measurement - Add `triton_skip_softmax` method with exponential model calibration (`scale_factor = a * exp(b * sparsity)`) and log-space fitting for diffusion models - Add Triton kernel backends for diffusers and LTX attention dispatch - Fix calibration to skip RULER dataset generation when user provides their own `forward_loop` (required for non-LLM models) ## Changes ### Triton kernels (`modelopt/torch/kernels/triton_fa.py`) - **`_attn_fwd`**: Forward kernel with optional tile skipping — tiles whose max attention score is far below the running softmax max are skipped entirely (no V load, no softmax, no accumulation). Runtime sparsity measurement via atomic counters. - **`_attn_fwd_calibrate`**: Calibration kernel that computes full attention while measuring how many tiles would be skipped at each of N thresholds simultaneously. Uses per-program output buffers (zero atomic contention) and vectorized multi-threshold comparison. - **`attention()`** / **`attention_calibrate()`**: Python wrappers for inference and calibration kernels. ### Kernel backends (`modelopt/torch/sparsity/attention_sparsity/kernels/`) - **`diffusers_triton_attention.py`**: Registers `modelopt_triton` backend in diffusers' attention dispatch. Handles [B, S, H, D] → varlen layout conversion, calibration/inference mode switching, thread-local configuration, and counter accumulation. - **`ltx_triton_attention.py`**: Patches `ltx_core.Attention` modules for Triton dispatch with the same calibration/inference modes. ### Method (`modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py`) - `TritonSkipSoftmaxMethod`: Context managers for calibration (→ calibration kernel) and inference (→ forward kernel with tile skipping). Three threshold priority levels: raw threshold > calibrated scale_factor > static threshold. ### Calibration (`modelopt/torch/sparsity/attention_sparsity/calibration/`) - **`calibrator.py`**: `DynamicThresholdCalibrator` with `fit_logspace` option — fits exponential model in log space (minimizes relative error) for diffusion models where scale_factors span many orders of magnitude. Records observed sparsity range for extrapolation warnings. - **`calibrate.py`**: Skips RULER dataset when `forward_loop` is provided; passes `fit_logspace` through from config. ### Config & conversion - **`config.py`**: `CalibrationConfig.fit_logspace` field (default False, recommended True for diffusion models). `skip_softmax_raw_threshold` field for direct threshold mode. - **`conversion.py`**: Auto-registers diffusers/LTX Triton backends on `sparsify()`. Updated summary display. ### Example - **`wan22_skip_softmax.py`**: End-to-end example for WAN 2.2 5B/14B with baseline, raw-threshold, and calibrated modes. Supports runtime sparsity reporting. ## Threshold modes | Mode | How it works | Use case | |------|-------------|----------| | **Raw threshold** (`--raw-threshold -0.7`) | Passed directly to kernel as `skip_threshold_log2` | Quick testing, sweeps | | **Calibrated** (`--calibrate --target-sparsity 0.5`) | `scale_factor = a * exp(b * target)`, then `threshold = scale_factor / seq_k` at runtime | Production use with seqlen adaptation | | **Static** (default `skip_softmax_threshold=0.1`) | `log2(lambda) * sm_scale` | Fallback | ## Usage ```bash # Fixed raw threshold (no calibration) python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --raw-threshold -0.7 \ --prompt "A cat playing piano" --output out.mp4 # With calibration (log-space fit for diffusion models) python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --calibrate --target-sparsity 0.5 \ --prompt "A cat playing piano" --output out.mp4 # Dense baseline for comparison python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --baseline \ --prompt "A cat playing piano" --output baseline.mp4 ``` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added skip-softmax sparse attention support for Diffusers models, enabling efficient video generation * Added support for both eager and Triton attention backends for sparse attention * Added new example script for Wan 2.2 text-to-video generation with sparse attention optimization * **Documentation** * Updated documentation with sparse attention configuration guide and usage examples * **Tests** * Added comprehensive unit tests for kernel backend registration and skip-softmax functionality <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com>	2026-04-18 06:54:19 +00:00