Unify weight and attention sparsity under mts.sparsify

CHANGELOG.rst

@@ -34,6 +34,10 @@ Changelog
 - The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``.
 - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for NVIDIA Nemotron 3.5 Super VL, covering chat-data blend preparation, assistant-only distillation, Hugging Face export and vLLM serving. See `examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/>`_ for details.
 
+*Sparsity*
+
+- ``mts.sparsify`` now also applies sparse attention: pass a sparse attention config such as ``mts.SKIP_SOFTMAX_CALIB`` in place of the mode, or combine it with weight sparsity as ``[("sparsegpt", {}), ("sparse_attention", cfg)]``. ``mts.export`` now also handles models with sparse attention, which it keeps.
+
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
@@ -48,6 +52,7 @@ Changelog
 **Deprecations**
 
 - ``modelopt.torch.export.copy_non_safetensor_files_from_ckpt`` is deprecated; use ``copy_non_model_files``.
+- ``modelopt.torch.sparsity.attention_sparsity.sparsify`` (``mtsa.sparsify``) is deprecated; use ``mts.sparsify(model, config, forward_loop=forward_loop)``.
 
 **Bug Fixes**
 

docs/source/guides/6_sparsity.rst

@@ -6,8 +6,8 @@ Introduction
 ============
 
 ModelOpt's Sparsity module (:mod:`modelopt.torch.sparsity <modelopt.torch.sparsity>`) enables
-you to sparsify the weights of your model. This can be useful for reducing the memory footprint of
-your model and can also be used to speed up inference.
+you to sparsify the weights and the attention of your model. This can be useful for reducing the
+memory footprint of your model and can also be used to speed up inference.
 
 
 Follow the steps described below to obtain a model with sparse weights using ModelOpt's Sparsity
@@ -69,6 +69,25 @@ Example usage:
     ``sparsegpt``. `sparse_magnitude` does not require `data_loader` as it uses magnitude-based
     method for thresholding.
 
+Sparse attention
+----------------
+
+Given a sparse attention config instead of a mode, the same API applies sparse attention, e.g.,
+skip-softmax with a fixed threshold (``mts.SKIP_SOFTMAX_DEFAULT``) or with a threshold calibrated
+on generated RULER data (``mts.SKIP_SOFTMAX_CALIB``):
+
+.. code-block:: python
+
+    model = mts.sparsify(model, mts.SKIP_SOFTMAX_CALIB)
+
+    # Weight and attention sparsity can be combined in one call
+    model = mts.sparsify(
+        model, [("sparsegpt", {}), ("sparse_attention", mts.SKIP_SOFTMAX_DEFAULT)], config=sparsity_config
+    )
+
+See the `attention sparsity example <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_
+for all methods and options.
+
 
 Save and restore the sparse model
 ---------------------------------
@@ -107,7 +126,8 @@ To restore the saved sparse model you can use
     :meth:`mts.export() <modelopt.torch.sparsity.sparsification.export>` will export the sparse
     model to a plain pytorch model. The sparse masks will be applied to model weights and all the
     sparse metadata will be removed. After exporting, sparsity will no longer be enforced during
-    subsequent fine-tuning. If you want to continue fine-tuning, do not export the model.
+    subsequent fine-tuning. If you want to continue fine-tuning, do not export the model. Sparse
+    attention needs no export and stays applied.
 
 .. note::
 

examples/diffusers/README.md

@@ -385,7 +385,7 @@ Skip-softmax sparse attention skips KV tiles whose attention scores are negligib
 ### Getting Started
 
 ```python
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 
 # 1. Define config with calibration
 config = {
@@ -410,7 +410,7 @@ def forward_loop(model):
     pipeline(prompt="a cat", num_frames=81, num_inference_steps=40, ...)
 
 # 3. Sparsify + calibrate
-mtsa.sparsify(transformer, config, forward_loop=forward_loop)
+mts.sparsify(transformer, config, forward_loop=forward_loop)
 
 # 4. Generate as usual — sparsity is applied automatically
 output = pipeline(prompt="a dog on the beach", ...)

examples/diffusers/sparsity/wan22_skip_softmax.py

@@ -58,7 +58,7 @@
 from diffusers import AutoencoderKLWan, WanPipeline
 from diffusers.utils import export_to_video
 
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.export import export_hf_checkpoint
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
@@ -426,7 +426,7 @@ def main() -> None:
             num_blocks = _get_num_blocks(transformer)
             print(f"Applying Triton backend to {name} ({num_blocks} blocks)...")
             config = build_sparse_config(args, num_blocks=num_blocks)
-            mtsa.sparsify(transformer, config, forward_loop=None)
+            mts.sparsify(transformer, config, forward_loop=None)
     else:
         # Build calibration forward loop if needed
         forward_loop = None
@@ -461,7 +461,7 @@ def main() -> None:
             num_blocks = _get_num_blocks(transformer)
             print(f"Applying skip-softmax to {name} ({num_blocks} blocks)...")
             config = build_sparse_config(args, num_blocks=num_blocks)
-            mtsa.sparsify(transformer, config, forward_loop=forward_loop)
+            mts.sparsify(transformer, config, forward_loop=forward_loop)
 
     # ---- Free calibration memory before inference ----
     if not args.baseline and not args.triton_baseline and forward_loop is not None:

examples/llm_eval/sparse_attention_utils.py

@@ -15,7 +15,7 @@
 
 """Utilities for sparse attention integration with llm_eval."""
 
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 
 
 def _extract_model(model_obj):
@@ -52,8 +52,8 @@ def sparsify_model(
 
     # Resolve config
     if isinstance(sparse_cfg, str):
-        # Get config from mtsa module (e.g., SKIP_SOFTMAX_CALIB, SKIP_SOFTMAX_DEFAULT)
-        mtsa_cfg = getattr(mtsa, sparse_cfg, None)
+        # Get config from mts module (e.g., SKIP_SOFTMAX_CALIB, SKIP_SOFTMAX_DEFAULT)
+        mtsa_cfg = getattr(mts, sparse_cfg, None)
         if mtsa_cfg is None:
             raise ValueError(f"Unknown sparse_cfg: {sparse_cfg}.")
     else:
@@ -72,7 +72,7 @@ def sparsify_model(
 
     # Apply sparsification
     print(f"\nApplying sparse attention with config: {sparse_cfg}")
-    mtsa.sparsify(net, mtsa_cfg)
+    mts.sparsify(net, mtsa_cfg)
     print("Sparse attention applied successfully!")
 
     return model

examples/llm_sparsity/attention_sparsity/README.md

@@ -15,7 +15,7 @@ Two attention backends are available:
 ### Quick Example
 
 ```python
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_DEFAULT
 
 # Load your model
@@ -26,7 +26,7 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 # Apply sparse attention
-model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
+model = mts.sparsify(model, SKIP_SOFTMAX_DEFAULT)
 ```
 
 > [!Note]
@@ -43,7 +43,7 @@ Uses a fixed threshold value. Simple but may not be optimal for all sequence len
 ```python
 from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_DEFAULT
 
-model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
+model = mts.sparsify(model, SKIP_SOFTMAX_DEFAULT)
 ```
 
 #### 2. Calibrated Threshold (SKIP_SOFTMAX_CALIB)
@@ -53,7 +53,7 @@ Uses RULER-based calibration to determine an optimal dynamic threshold that adap
 ```python
 from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_CALIB
 
-model = mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)
+model = mts.sparsify(model, SKIP_SOFTMAX_CALIB)
 ```
 
 ### N:M Sparse Softmax (SPARSE_SOFTMAX_DEFAULT)
@@ -69,7 +69,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="cuda",
 )
 
-model = mtsa.sparsify(model, config=SPARSE_SOFTMAX_DEFAULT)
+model = mts.sparsify(model, SPARSE_SOFTMAX_DEFAULT)
 ```
 
 Custom N:M configuration:
@@ -90,11 +90,11 @@ sparse_cfg = {
     },
 }
 
-model = mtsa.sparsify(model, config=sparse_cfg)
+model = mts.sparsify(model, sparse_cfg)
 ```
 
 > [!Note]
-> N:M sparse softmax requires the Triton backend (`backend="triton"`). The `attn_implementation` is automatically set to `"modelopt_triton"` by `mtsa.sparsify()`. N:M sparsity is applied during prefill only — decode tokens are not sparsified.
+> N:M sparse softmax requires the Triton backend (`backend="triton"`). The `attn_implementation` is automatically set to `"modelopt_triton"` by `mts.sparsify()`. N:M sparsity is applied during prefill only — decode tokens are not sparsified.
 
 ## Prerequisites
 
@@ -233,7 +233,7 @@ custom_config = {
     },
 }
 
-model = mtsa.sparsify(model, config=custom_config)
+model = mts.sparsify(model, custom_config)
 ```
 
 ## References

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -26,7 +26,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import modelopt.torch.opt as mto
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.export import ensure_local_checkpoint, export_hf_checkpoint
 from modelopt.torch.sparsity.attention_sparsity.config import (
     SKIP_SOFTMAX_CALIB,
@@ -152,7 +152,7 @@ def main(args):
     args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint(args.pyt_ckpt_path)
     print(f"Loading model: {args.local_checkpoint_path}")
 
-    # No need to specify attn_implementation here — mtsa.sparsify() sets it
+    # No need to specify attn_implementation here — mts.sparsify() sets it
     # automatically ("eager" for pytorch backend, "modelopt_triton" for triton).
     model = AutoModelForCausalLM.from_pretrained(
         args.local_checkpoint_path, attn_implementation="eager", dtype="auto", device_map="auto"
@@ -202,7 +202,7 @@ def main(args):
             else str(Path(__file__).parent / "data"),
         )
 
-    model = mtsa.sparsify(model, config=sparse_config)
+    model = mts.sparsify(model, sparse_config)
     print("Sparse attention applied successfully!")
 
     # Generate sample output AFTER sparse attention

modelopt/torch/sparsity/__init__.py

@@ -15,10 +15,13 @@
 
 """API for sparsification algorithms.
 
-This module provides access to both weight sparsity and attention sparsity algorithms.
-For backward compatibility, weight sparsity APIs are re-exported at the module level.
+This module provides access to both weight sparsity and attention sparsity algorithms through
+:meth:`sparsify <modelopt.torch.sparsity.sparsification.sparsify>`.
 """
 
+from . import attention_sparsity
+from .attention_sparsity.config import *
+from .sparsification import *
+
 # Import weight sparsity for backward compatibility
 from .weight_sparsity import mode, module, plugins
-from .weight_sparsity.sparsification import *

modelopt/torch/sparsity/attention_sparsity/conversion.py

@@ -74,7 +74,7 @@ def _set_attn_implementation(model: nn.Module, config: SparseAttentionConfig) ->
       directly in ``SparseAttentionModule.forward()``, so no ``attn_implementation``
       change is needed.
 
-    This is called automatically during ``mtsa.sparsify()`` so users never need
+    This is called automatically during ``mts.sparsify()`` so users never need
     to manually set ``attn_implementation``.
     """
     sparse_cfg = config.sparse_cfg if hasattr(config, "sparse_cfg") else {}
@@ -354,6 +354,21 @@ def update_sparse_attention_metadata(
     )
 
 
+def get_sparse_attention_mode_state(
+    model: nn.Module,
+) -> tuple[SparseAttentionConfig, MetadataDict] | None:
+    """Return the config and metadata of the latest ``sparse_attention`` mode applied to ``model``."""
+    # Wrappers and containers of converted models have no mode state of their own.
+    if not hasattr(model, ModeloptStateManager._state_key):
+        return None
+    mode_states = [
+        (config, metadata)
+        for mode, config, metadata in ModeloptStateManager(model).modes_with_states()
+        if str(mode) == "sparse_attention"
+    ]
+    return mode_states[-1] if mode_states else None
+
+
 def _get_sparse_softmax_export_config(module: SparseAttentionModule) -> dict[str, Any] | None:
     """Return N:M sparse-softmax export metadata for one sparse attention module."""
     method_config = getattr(module, "_method_config", {})
@@ -523,8 +538,9 @@ def disable_sparse_attention(model: nn.Module, wildcard_or_filter_func: str | Ca
             For example: "*lm_head*", "*layer_0*", etc.
 
     Example:
+        >>> import modelopt.torch.sparsity as mts
         >>> import modelopt.torch.sparsity.attention_sparsity as sparse_attn
-        >>> model = sparse_attn.sparsify(model, config)
+        >>> model = mts.sparsify(model, config)
         >>> # Disable sparse attention for lm_head
         >>> sparse_attn.disable_sparse_attention(model, "*lm_head*")
     """
@@ -550,8 +566,9 @@ def enable_sparse_attention(model: nn.Module, wildcard_or_filter_func: str | Cal
             For example: "*attention*", "*attn*", etc.
 
     Example:
+        >>> import modelopt.torch.sparsity as mts
         >>> import modelopt.torch.sparsity.attention_sparsity as sparse_attn
-        >>> model = sparse_attn.sparsify(model, config)
+        >>> model = mts.sparsify(model, config)
         >>> # Re-enable sparse attention for all attention modules
         >>> sparse_attn.enable_sparse_attention(model, "*attention*")
     """

modelopt/torch/sparsity/attention_sparsity/methods/vsa.py

@@ -22,7 +22,7 @@
 Uses the optimized Triton kernel from fastvideo_kernel.
 
 Integration:
-    After ``mtsa.sparsify(model, VSA_DEFAULT)``, each attention layer's
+    After ``mts.sparsify(model, VSA_DEFAULT)``, each attention layer's
     ``F.scaled_dot_product_attention`` call is intercepted and replaced by the VSA
     kernel.  Cross-attention (Q/K have different seq_len) is automatically skipped.
     This works with HF transformers and diffusers.

modelopt/torch/sparsity/attention_sparsity/mode.py

@@ -23,13 +23,15 @@
     UpdateEntrypoint,
     _ModeRegistryCls,
 )
+from modelopt.torch.opt.searcher import BaseSearcher
 
 from .config import SparseAttentionConfig
 from .conversion import (
     convert_to_sparse_attention_model,
     restore_sparse_attention_model,
     update_sparse_attention_metadata,
 )
+from .searcher import SparseAttentionSearcher
 
 # Create registry for sparse attention modes
 SparseAttentionModeRegistry = _ModeRegistryCls("sparse_attention")
@@ -54,15 +56,9 @@ def config_class(self) -> type[ModeloptBaseConfig]:
         return SparseAttentionConfig
 
     @property
-    def next_prohibited_modes(self) -> set[str] | None:
-        """Modes that should not be applied after this mode."""
-        # Can work with quantization but not with weight sparsity
-        return {"sparsity"}
-
-    @property
-    def export_mode(self) -> str | None:
-        """The mode that corresponds to the export mode of this mode."""
-        return "export_sparse_attention"
+    def search_algorithm(self) -> type[BaseSearcher]:
+        """Specifies the search algorithm for the mode."""
+        return SparseAttentionSearcher
 
     @property
     def convert(self) -> ConvertEntrypoint: