From 80df9cb1dd6f8e4e54a0af37d248b79b1de656e5 Mon Sep 17 00:00:00 2001
From: Rohan Joshi <rohjoshi@nvidia.com>
Date: Mon, 26 Oct 2026 23:57:04 +0000
Subject: [PATCH 1/2] Calibrate sparse attention through a mode searcher and
 allow it after weight sparsity

First step toward a single mts.sparsify() for weight and attention sparsity; no public API changes.

The sparse_attention mode now has a search algorithm, SparseAttentionSearcher, that runs the existing threshold calibration, so the mode can be driven like the weight-sparsity modes. mtsa.sparsify uses it. On Qwen2.5-0.5B (6 layers, CPU, prefill target 0.5) the fitted skip-softmax coefficients are identical to the previous path.

sparse_attention may now follow sparse_magnitude and sparsegpt. The mode no longer declares the unregistered export_sparse_attention export mode, which made mts.export fail on models with attention sparsity applied after weight sparsity. It also drops the registry name "sparsity" from its prohibited next modes, which never matched a mode.

calibrate() now writes its results into the sparse_attention mode metadata. Saving refreshes only the last mode, so results were lost when another mode had been applied after sparse_attention.

modelopt.torch.sparsity imports attention_sparsity directly, so the sparse_attention mode is registered without relying on the indirect import through modelopt.torch.export.

Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 modelopt/torch/sparsity/__init__.py           |  2 +
 .../sparsity/attention_sparsity/conversion.py | 17 ++++++
 .../torch/sparsity/attention_sparsity/mode.py | 14 ++---
 .../attention_sparsity/model_sparsify.py      | 15 +++--
 .../sparsity/attention_sparsity/searcher.py   | 43 ++++++++++++++
 .../torch/sparsity/weight_sparsity/mode.py    |  2 +-
 .../test_sparse_attention_conversion.py       | 57 +++++++++++++++++++
 7 files changed, 136 insertions(+), 14 deletions(-)
 create mode 100644 modelopt/torch/sparsity/attention_sparsity/searcher.py

diff --git a/modelopt/torch/sparsity/__init__.py b/modelopt/torch/sparsity/__init__.py
index 2013fded1ae..8e0836431a6 100644
--- a/modelopt/torch/sparsity/__init__.py
+++ b/modelopt/torch/sparsity/__init__.py
@@ -19,6 +19,8 @@
 For backward compatibility, weight sparsity APIs are re-exported at the module level.
 """
 
+from . import attention_sparsity
+
 # Import weight sparsity for backward compatibility
 from .weight_sparsity import mode, module, plugins
 from .weight_sparsity.sparsification import *
diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py
index 38944517740..8d146256a8c 100644
--- a/modelopt/torch/sparsity/attention_sparsity/conversion.py
+++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py
@@ -354,6 +354,23 @@ 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``."""
+    # A wrapper whose submodule holds the modelopt state has no mode state of its own.
+    if not ModeloptStateManager.is_converted(model) or any(
+        ModeloptStateManager.is_converted(child) for child in model.children()
+    ):
+        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", {})
diff --git a/modelopt/torch/sparsity/attention_sparsity/mode.py b/modelopt/torch/sparsity/attention_sparsity/mode.py
index f389509a592..06210d78cea 100644
--- a/modelopt/torch/sparsity/attention_sparsity/mode.py
+++ b/modelopt/torch/sparsity/attention_sparsity/mode.py
@@ -23,6 +23,7 @@
     UpdateEntrypoint,
     _ModeRegistryCls,
 )
+from modelopt.torch.opt.searcher import BaseSearcher
 
 from .config import SparseAttentionConfig
 from .conversion import (
@@ -30,6 +31,7 @@
     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:
diff --git a/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py b/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
index a33938b0576..a038b30ce4a 100644
--- a/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
+++ b/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
@@ -24,7 +24,9 @@
 
 from .calibration import calibrate_sparse_attention
 from .config import SparseAttentionConfig
+from .conversion import get_sparse_attention_mode_state, update_sparse_attention_metadata
 from .mode import SparseAttentionModeRegistry
+from .searcher import SparseAttentionSearcher
 
 __all__ = [
     "calibrate",
@@ -101,7 +103,9 @@ def sparsify(
 
             This requires the model to be loaded with ``attn_implementation="eager"``.
 
-        forward_loop: Reserved for future use.
+        forward_loop: Callable that forwards calibration data through the model, used when
+            ``sparse_cfg`` contains a ``"calibration"`` entry. If None, a RULER calibration
+            dataset is generated from the model's tokenizer.
 
             Here are a few examples for correct ``forward_loop`` definitions:
 
@@ -151,9 +155,8 @@ def forward_loop(model) -> float:
     model = apply_mode(
         model, mode=[("sparse_attention", config)], registry=SparseAttentionModeRegistry
     )
-
-    # Calibrate the sparsity ratio of the attention modules
-    return calibrate(model, config, forward_loop=forward_loop)
+    SparseAttentionSearcher().search(model, {}, (), {"forward_loop": forward_loop})
+    return model
 
 
 def calibrate(
@@ -173,5 +176,9 @@ def calibrate(
     Returns:
         The calibrated model with optimized sparse attention thresholds.
     """
+    mode_state = get_sparse_attention_mode_state(model)
     calibrate_sparse_attention(model, config, forward_loop=forward_loop)
+    # Saving refreshes only the last mode's metadata, which may not be sparse_attention.
+    if mode_state is not None:
+        update_sparse_attention_metadata(model, *mode_state)
     return model
diff --git a/modelopt/torch/sparsity/attention_sparsity/searcher.py b/modelopt/torch/sparsity/attention_sparsity/searcher.py
new file mode 100644
index 00000000000..ac623f2c238
--- /dev/null
+++ b/modelopt/torch/sparsity/attention_sparsity/searcher.py
@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Searcher that calibrates sparse attention thresholds."""
+
+from modelopt.torch.opt.searcher import BaseSearcher, SearchStateDict
+
+from .calibration import calibrate_sparse_attention
+from .conversion import get_sparse_attention_mode_state
+
+__all__ = ["SparseAttentionSearcher"]
+
+
+class SparseAttentionSearcher(BaseSearcher):
+    """Calibrate sparse attention thresholds with the forward loop of the search config.
+
+    Calibration is configured by ``sparse_cfg["calibration"]`` of the model's ``sparse_attention``
+    mode config and is skipped without it. Without a ``forward_loop`` or ``data_loader``, a RULER
+    calibration dataset is generated from the model's tokenizer.
+    """
+
+    @property
+    def default_state_dict(self) -> SearchStateDict:
+        """Return default state dict."""
+        return {}
+
+    def run_search(self) -> None:
+        """Calibrate the sparse attention modules of the model."""
+        mode_state = get_sparse_attention_mode_state(self.model)
+        assert mode_state is not None, "Model has no sparse_attention mode to calibrate!"
+        calibrate_sparse_attention(self.model, mode_state[0], forward_loop=self.forward_loop)
diff --git a/modelopt/torch/sparsity/weight_sparsity/mode.py b/modelopt/torch/sparsity/weight_sparsity/mode.py
index db7e1f33253..a963007d56c 100644
--- a/modelopt/torch/sparsity/weight_sparsity/mode.py
+++ b/modelopt/torch/sparsity/weight_sparsity/mode.py
@@ -117,7 +117,7 @@ def config_class(self) -> type[ModeloptBaseConfig]:
     @property
     def next_modes(self) -> set[str] | None:
         """Specifies the next modes for the mode."""
-        return {"export_sparse", "kd_loss", "quantize"}
+        return {"export_sparse", "kd_loss", "quantize", "sparse_attention"}
 
     @property
     def export_mode(self) -> str | None:
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
index e70f2b7de65..096d291321b 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
@@ -23,6 +23,7 @@
 
 from unittest.mock import MagicMock, patch
 
+import torch
 import torch.nn as nn
 from _test_utils.torch.sparsity.sparse_attention_common import (
     FLASH_SKIP_SOFTMAX_DEFAULT_CFG,
@@ -31,6 +32,7 @@
 )
 
 import modelopt.torch.opt as mto
+import modelopt.torch.sparsity as mts
 import modelopt.torch.sparsity.attention_sparsity as sparse_attn
 from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_CALIB_SPARSE24
 from modelopt.torch.sparsity.attention_sparsity.conversion import (
@@ -41,6 +43,7 @@
     print_sparse_attention_summary,
 )
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
+from modelopt.torch.sparsity.weight_sparsity.module import SparseModule
 
 
 class TestSparseAttentionReplacement:
@@ -212,6 +215,60 @@ def test_restore_keeps_calibration_params(self, tmp_path):
             assert method.target_sparse_ratio == target_sparse_ratio
             assert module.get_threshold_info()["type"] == "dynamic_calibrated"
 
+    def test_calibrate_after_later_mode_survives_restore(self):
+        """calibrate() persists its results even when sparse_attention is not the last mode."""
+        calibration_params = {"prefill": {"a": 3.14, "b": 7.5}}
+
+        def fake_calibration(model, config, forward_loop=None):
+            for module in model.modules():
+                if isinstance(module, SparseAttentionModule):
+                    module._sparse_method_instance.calibration_params = calibration_params
+
+        model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, "sparse_magnitude")
+        with patch(
+            "modelopt.torch.sparsity.attention_sparsity.model_sparsify.calibrate_sparse_attention",
+            side_effect=fake_calibration,
+        ):
+            sparse_attn.calibrate(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+
+        model_restored = mto.restore_from_modelopt_state(
+            SimpleAttentionModel(), mto.modelopt_state(model)
+        )
+        sparse_modules = [
+            m for m in model_restored.modules() if isinstance(m, SparseAttentionModule)
+        ]
+        assert sparse_modules
+        for module in sparse_modules:
+            assert module._sparse_method_instance.calibration_params == calibration_params
+
+    @pytest.mark.parametrize("weight_first", [True, False])
+    def test_chain_with_weight_sparsity(self, weight_first):
+        """Weight and attention sparsity chain in either order; export folds only weight masks."""
+        steps = [
+            lambda m: mts.sparsify(m, "sparse_magnitude"),
+            lambda m: sparse_attn.sparsify(m, FLASH_SKIP_SOFTMAX_DEFAULT_CFG),
+        ]
+        model = SimpleAttentionModel().eval()
+        for step in steps if weight_first else steps[::-1]:
+            model = step(model)
+        x = SimpleAttentionModel.get_input()
+        with torch.no_grad():
+            output = model(x)
+
+        model_restored = mto.restore_from_modelopt_state(
+            SimpleAttentionModel(), mto.modelopt_state(model)
+        ).eval()
+        model_restored.load_state_dict(model.state_dict())
+        with torch.no_grad():
+            assert torch.allclose(model_restored(x), output)
+
+        mts.export(model)
+        assert not any(isinstance(m, SparseModule) for m in model.modules())
+        assert any(isinstance(m, SparseAttentionModule) for m in model.modules())
+        with torch.no_grad():
+            assert torch.allclose(model(x), output)
+
     def test_restore_keeps_disabled_modules(self):
         """Modules disabled after sparsify stay disabled after restore."""
         model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)

From cd58ab3f535bdbbff6b47ffb49e51ce588781d05 Mon Sep 17 00:00:00 2001
From: Rohan Joshi <rohjoshi@nvidia.com>
Date: Tue, 27 Oct 2026 22:54:20 +0000
Subject: [PATCH 2/2] Unify weight and attention sparsity under mts.sparsify

mts.sparsify now applies sparse attention as well as weight sparsity. A sparse attention config such as mts.SKIP_SOFTMAX_CALIB can be passed in place of the mode, and modes can be combined, e.g. [("sparsegpt", {}), ("sparse_attention", cfg)]. Existing weight-sparsity calls are unchanged.

The new modelopt.torch.sparsity.sparsification module applies each mode and runs its searcher before applying the next one, because applying a mode freezes the metadata of the previous one, which holds the attention calibration results. Search options are routed to the searchers that accept them, and a new forward_loop keyword is a shorthand for config["forward_loop"]. All modes, mode configs and option keys are validated before the model is modified.

mts.export folds weight-sparsity masks as before and keeps sparse attention, which has nothing to export and whose modules export_hf_checkpoint reads. A model left with only sparse attention is returned unchanged.

mts re-exports the sparse attention presets and config classes, and calibrate for re-calibrating sparse attention. mtsa.sparsify is deprecated in favor of mts.sparsify; examples, docs and tests now use mts.sparsify.

Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CHANGELOG.rst                                 |   5 +
 docs/source/guides/6_sparsity.rst             |  26 +++-
 examples/diffusers/README.md                  |   4 +-
 .../diffusers/sparsity/wan22_skip_softmax.py  |   6 +-
 examples/llm_eval/sparse_attention_utils.py   |   8 +-
 .../llm_sparsity/attention_sparsity/README.md |  16 +--
 .../llm_sparsity/attention_sparsity/hf_sa.py  |   6 +-
 modelopt/torch/sparsity/__init__.py           |   7 +-
 .../sparsity/attention_sparsity/conversion.py |  14 +--
 .../attention_sparsity/methods/vsa.py         |   2 +-
 .../attention_sparsity/model_sparsify.py      |  10 ++
 modelopt/torch/sparsity/sparsification.py     | 112 ++++++++++++++++++
 .../torch/sparsity/sparse_attention_common.py |   4 +-
 .../diffusers/sparsity/test_sparsity.py       |   4 +-
 .../attention/test_triton_fa_skip_softmax.py  |  10 +-
 .../attention/test_triton_fa_sparse_nm.py     |   8 +-
 .../test_attention_sparsity_gpu.py            |   8 +-
 .../test_calibration_gpu.py                   |   2 +-
 .../test_integration_gpu.py                   |  10 +-
 .../test_triton_calibration_gpu.py            |   4 +-
 .../test_wan22_skip_softmax.py                |   8 +-
 .../test_calibrator_fitting.py                |   2 +-
 .../test_sparse_attention_calibration.py      |   2 +-
 .../test_sparse_attention_conversion.py       |  81 +++++++++----
 .../attention_sparsity/test_threshold_info.py |   2 +-
 .../sparsity/attention_sparsity/test_vsa.py   |  16 +--
 26 files changed, 278 insertions(+), 99 deletions(-)
 create mode 100644 modelopt/torch/sparsity/sparsification.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 66f4015c760..e71d601b16d 100755
--- a/CHANGELOG.rst
+++ b/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**
 
diff --git a/docs/source/guides/6_sparsity.rst b/docs/source/guides/6_sparsity.rst
index 12a17b3c66c..ca759a70079 100644
--- a/docs/source/guides/6_sparsity.rst
+++ b/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::
 
diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md
index b67d9630b8d..37fdff50689 100644
--- a/examples/diffusers/README.md
+++ b/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", ...)
diff --git a/examples/diffusers/sparsity/wan22_skip_softmax.py b/examples/diffusers/sparsity/wan22_skip_softmax.py
index 890cec547ff..38d01c35a6e 100644
--- a/examples/diffusers/sparsity/wan22_skip_softmax.py
+++ b/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:
diff --git a/examples/llm_eval/sparse_attention_utils.py b/examples/llm_eval/sparse_attention_utils.py
index 8dc560851f1..66530293055 100644
--- a/examples/llm_eval/sparse_attention_utils.py
+++ b/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
diff --git a/examples/llm_sparsity/attention_sparsity/README.md b/examples/llm_sparsity/attention_sparsity/README.md
index 00cf6e0ac21..e430a11f802 100644
--- a/examples/llm_sparsity/attention_sparsity/README.md
+++ b/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
diff --git a/examples/llm_sparsity/attention_sparsity/hf_sa.py b/examples/llm_sparsity/attention_sparsity/hf_sa.py
index b94238b8cb1..4fea29180af 100644
--- a/examples/llm_sparsity/attention_sparsity/hf_sa.py
+++ b/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
diff --git a/modelopt/torch/sparsity/__init__.py b/modelopt/torch/sparsity/__init__.py
index 8e0836431a6..f0789748d21 100644
--- a/modelopt/torch/sparsity/__init__.py
+++ b/modelopt/torch/sparsity/__init__.py
@@ -15,12 +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 *
diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py
index 8d146256a8c..a51d77dc263 100644
--- a/modelopt/torch/sparsity/attention_sparsity/conversion.py
+++ b/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 {}
@@ -358,10 +358,8 @@ 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``."""
-    # A wrapper whose submodule holds the modelopt state has no mode state of its own.
-    if not ModeloptStateManager.is_converted(model) or any(
-        ModeloptStateManager.is_converted(child) for child in model.children()
-    ):
+    # 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)
@@ -540,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*")
     """
@@ -567,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*")
     """
diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/vsa.py b/modelopt/torch/sparsity/attention_sparsity/methods/vsa.py
index 66acfb510ce..d050da52125 100644
--- a/modelopt/torch/sparsity/attention_sparsity/methods/vsa.py
+++ b/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.
diff --git a/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py b/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
index a038b30ce4a..079750e67af 100644
--- a/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
+++ b/modelopt/torch/sparsity/attention_sparsity/model_sparsify.py
@@ -15,6 +15,7 @@
 
 """Main API functions for sparse attention optimization."""
 
+import warnings
 from typing import Any
 
 import torch
@@ -41,6 +42,10 @@ def sparsify(
 ) -> torch.nn.Module:
     """Applies sparse attention optimization to the model in-place.
 
+    .. deprecated:: 0.49.0
+        Use :meth:`mts.sparsify(model, config, forward_loop=forward_loop)
+        <modelopt.torch.sparsity.sparsification.sparsify>` instead.
+
     This method performs replacement of attention modules with their sparse counterparts.
 
     Args:
@@ -152,6 +157,11 @@ def forward_loop(model) -> float:
     Returns:
         A pytorch model which has sparse attention applied and optionally calibrated.
     """
+    warnings.warn(
+        "mtsa.sparsify() is deprecated; use modelopt.torch.sparsity.sparsify() instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     model = apply_mode(
         model, mode=[("sparse_attention", config)], registry=SparseAttentionModeRegistry
     )
diff --git a/modelopt/torch/sparsity/sparsification.py b/modelopt/torch/sparsity/sparsification.py
new file mode 100644
index 00000000000..e6fd914fe90
--- /dev/null
+++ b/modelopt/torch/sparsity/sparsification.py
@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""High-level API to sparsify the weights and the attention of a model."""
+
+from typing import Any
+
+from torch import nn
+
+from modelopt.torch.opt.conversion import apply_mode
+from modelopt.torch.opt.mode import ModeLike, get_mode_config
+from modelopt.torch.opt.searcher import ForwardLoop, SearchConfig
+
+from .attention_sparsity.config import SparseAttentionConfig
+from .attention_sparsity.conversion import is_attn_sparsified
+from .attention_sparsity.mode import SparseAttentionModeRegistry
+from .attention_sparsity.model_sparsify import calibrate
+from .weight_sparsity import sparsification as weight_sparsification
+from .weight_sparsity.mode import SparsityModeRegistry
+from .weight_sparsity.module import SparseModule
+
+__all__ = ["calibrate", "export", "sparsify"]
+
+_REGISTRIES = (SparsityModeRegistry, SparseAttentionModeRegistry)
+
+
+def sparsify(
+    model: nn.Module,
+    mode: ModeLike | dict[str, Any] | SparseAttentionConfig,
+    config: SearchConfig | None = None,
+    *,
+    forward_loop: ForwardLoop | None = None,
+) -> nn.Module:
+    """Sparsify the weights and/or the attention of a model in-place.
+
+    Example::
+
+        model = mts.sparsify(model, "sparsegpt", config={"data_loader": data_loader})
+        model = mts.sparsify(model, mts.SKIP_SOFTMAX_CALIB)  # calibrates on RULER data
+
+    Args:
+        model: The model to sparsify.
+        mode: A mode, a list of modes, or a list of ``(mode, mode_config)`` tuples. Modes are
+            applied and then searched or calibrated one at a time, in order:
+
+            * ``"sparse_magnitude"``: 2:4 weight sparsity by weight magnitude.
+            * ``"sparsegpt"``: 2:4 weight sparsity with SparseGPT weight updates.
+            * ``"sparse_attention"``: sparse attention configured by a
+              :class:`SparseAttentionConfig <modelopt.torch.sparsity.attention_sparsity.config.SparseAttentionConfig>`
+              such as ``mts.SKIP_SOFTMAX_CALIB``, which may also be passed directly as ``mode``.
+        config: Search options such as ``forward_loop``, ``data_loader``, ``collect_func`` and
+            ``verbose``, plus algorithm-specific ones such as SparseGPT's ``device``. Each mode
+            receives the options its algorithm accepts.
+        forward_loop: Shorthand for ``config["forward_loop"]``. Sparse attention calibration
+            accepts a ``forward_loop`` or ``data_loader`` only with a decode target sparsity of 0;
+            without either, it calibrates on generated RULER data.
+
+    Returns:
+        The sparsified model.
+    """
+    if isinstance(mode, (dict, SparseAttentionConfig)):
+        mode = [("sparse_attention", dict(mode))]
+    config = dict(config or {})
+    if forward_loop is not None:
+        if config.get("forward_loop") is not None:
+            raise ValueError("Pass `forward_loop` either as an argument or in `config`, not both.")
+        config["forward_loop"] = forward_loop
+    if config.get("forward_loop") is not None and config.get("data_loader") is not None:
+        raise ValueError("Pass either `forward_loop` or `data_loader`, not both.")
+
+    # Validate all modes and config keys before modifying the model.
+    modes = get_mode_config(mode)
+    searchers = []
+    for m, mode_config in modes:
+        registry = next((r for r in _REGISTRIES if m in r), None)
+        if registry is None:
+            raise ValueError(f"{m} is not a sparsity mode.")
+        registry[m].config_class(**mode_config)
+        searchers.append(registry[m].search_algorithm())
+    supported_keys = set().union(*(s.default_search_config for s in searchers))
+    if unsupported_keys := config.keys() - supported_keys:
+        raise ValueError(f"Unsupported config keys: {sorted(unsupported_keys)}")
+
+    # The next mode freezes the metadata of the previous one, so each mode is searched right away.
+    for (m, mode_config), searcher in zip(modes, searchers):
+        model = apply_mode(model, [(m, mode_config)])
+        search_keys = searcher.default_search_config.keys()
+        searcher.search(model, {}, (), {k: v for k, v in config.items() if k in search_keys})
+    return model
+
+
+def export(model: nn.Module) -> nn.Module:
+    """Export a sparse model to a regular model by folding the weight sparsity masks into the weights.
+
+    Export after fine-tuning: the masks are no longer enforced afterwards. Sparse attention needs no
+    export and stays applied.
+    """
+    if is_attn_sparsified(model) and not any(isinstance(m, SparseModule) for m in model.modules()):
+        return model
+    return weight_sparsification.export(model)
diff --git a/tests/_test_utils/torch/sparsity/sparse_attention_common.py b/tests/_test_utils/torch/sparsity/sparse_attention_common.py
index 58d711affa2..9efa7a6d8fc 100644
--- a/tests/_test_utils/torch/sparsity/sparse_attention_common.py
+++ b/tests/_test_utils/torch/sparsity/sparse_attention_common.py
@@ -19,7 +19,7 @@
 import torch.nn as nn
 
 import modelopt.torch.opt as mto
-import modelopt.torch.sparsity.attention_sparsity as sparse_attn
+import modelopt.torch.sparsity as mts
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
 
@@ -141,7 +141,7 @@ def forward_loop(model):
             model(batch)
 
     # Apply sparse attention
-    model = sparse_attn.sparsify(model, config, forward_loop=forward_loop)
+    model = mts.sparsify(model, config, forward_loop=forward_loop)
 
     # Verify sparse attention modules were inserted
     assert any(isinstance(m, SparseAttentionModule) for m in model.modules()), (
diff --git a/tests/examples/diffusers/sparsity/test_sparsity.py b/tests/examples/diffusers/sparsity/test_sparsity.py
index 1727ee39b7f..f296437d55f 100644
--- a/tests/examples/diffusers/sparsity/test_sparsity.py
+++ b/tests/examples/diffusers/sparsity/test_sparsity.py
@@ -30,7 +30,7 @@
 from _test_utils.examples.run_command import run_example_command
 from diffusers import AutoencoderKLWan, WanPipeline
 
-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
 
@@ -183,7 +183,7 @@ def test_wan22_calibrated_export(tiny_wan22_path, tmp_path):
     }
     config = {"sparse_cfg": sparse_cfg}
     for transformer in [pipe.transformer, pipe.transformer_2]:
-        mtsa.sparsify(transformer, config)
+        mts.sparsify(transformer, config)
 
     # Inject calibration params (simulating a successful log-space calibration).
     test_log_a = math.log(1.5)
diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py
index 033952d34eb..133444ee24a 100644
--- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py
+++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py
@@ -20,7 +20,7 @@
 from conftest import make_varlen_meta
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE
 from modelopt.torch.sparsity.attention_sparsity.methods.flash_skip_softmax import FlashSkipSoftmax
 
@@ -246,10 +246,10 @@ def test_triton_matches_pytorch_reference(self):
 
 @pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton")
 class TestSkipSoftmaxHFIntegration:
-    """HF integration for skip-softmax via mtsa.sparsify()."""
+    """HF integration for skip-softmax via mts.sparsify()."""
 
     def test_skip_softmax_via_sparsify(self, tiny_llama_dir):
-        """mtsa.sparsify() with triton_skip_softmax produces finite logits."""
+        """mts.sparsify() with triton_skip_softmax produces finite logits."""
         tok = AutoTokenizer.from_pretrained(tiny_llama_dir)
         if tok.pad_token_id is None:
             tok.pad_token_id = tok.eos_token_id
@@ -267,13 +267,13 @@ def test_skip_softmax_via_sparsify(self, tiny_llama_dir):
             logits_dense = model_dense(input_ids=ids).logits
         del model_dense
 
-        # Skip-softmax via mtsa.sparsify()
+        # Skip-softmax via mts.sparsify()
         model_skip = AutoModelForCausalLM.from_pretrained(
             tiny_llama_dir,
             torch_dtype=torch.bfloat16,
             device_map="cuda",
         )
-        mtsa.sparsify(model_skip, mtsa.SKIP_SOFTMAX_TRITON_DEFAULT)
+        mts.sparsify(model_skip, mts.SKIP_SOFTMAX_TRITON_DEFAULT)
         model_skip.eval()
         with torch.no_grad():
             logits_skip = model_skip(input_ids=ids).logits
diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py
index 651aeab4116..538e6597def 100644
--- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py
+++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py
@@ -22,7 +22,7 @@
 from conftest import make_qkv, make_varlen_meta
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE
 
 if TRITON_KERNEL_AVAILABLE:
@@ -417,7 +417,7 @@ def test_sparse_disabled_matches_dense(self):
         assert torch.equal(out_dense, out_n0)
 
     def test_sparse_nm_via_sparsify(self, tiny_llama_dir):
-        """mtsa.sparsify() with N:M sparse softmax produces finite logits that differ from dense."""
+        """mts.sparsify() with N:M sparse softmax produces finite logits that differ from dense."""
         tok = AutoTokenizer.from_pretrained(tiny_llama_dir)
         if tok.pad_token_id is None:
             tok.pad_token_id = tok.eos_token_id
@@ -435,7 +435,7 @@ def test_sparse_nm_via_sparsify(self, tiny_llama_dir):
             logits_dense = model_dense(input_ids=ids).logits
         del model_dense
 
-        # Sparse via mtsa.sparsify() with dense_recent_tokens=0 to force sparsity on all tiles
+        # Sparse via mts.sparsify() with dense_recent_tokens=0 to force sparsity on all tiles
         sparse_cfg = {
             "sparse_cfg": {
                 "*attn*": {
@@ -455,7 +455,7 @@ def test_sparse_nm_via_sparsify(self, tiny_llama_dir):
             torch_dtype=torch.bfloat16,
             device_map="cuda",
         )
-        mtsa.sparsify(model_sparse, sparse_cfg)
+        mts.sparsify(model_sparse, sparse_cfg)
         assert model_sparse.config._attn_implementation == "modelopt_triton"
         model_sparse.eval()
         with torch.no_grad():
diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_attention_sparsity_gpu.py b/tests/gpu/torch/sparsity/attention_sparsity/test_attention_sparsity_gpu.py
index ce762267a23..ab26a810d7e 100644
--- a/tests/gpu/torch/sparsity/attention_sparsity/test_attention_sparsity_gpu.py
+++ b/tests/gpu/torch/sparsity/attention_sparsity/test_attention_sparsity_gpu.py
@@ -27,7 +27,7 @@
     sparsify_model_and_forward,
 )
 
-import modelopt.torch.sparsity.attention_sparsity as sparse_attn
+import modelopt.torch.sparsity as mts
 
 
 class TestAttentionSparsityGPU:
@@ -86,7 +86,7 @@ def test_different_dtypes(self, dtype):
     def test_backward_pass(self):
         """Test that gradients flow correctly through sparse attention."""
         model = SimpleAttentionModel(hidden_size=128, num_heads=4).to(self.device)
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         # Enable training mode
         model.train()
@@ -114,7 +114,7 @@ def test_backward_pass(self):
     def test_various_sequence_lengths(self, seq_len):
         """Test sparse attention with various sequence lengths."""
         model = SimpleAttentionModel(hidden_size=128, num_heads=4).to(self.device)
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         x = model.get_input(hidden_size=128, seq_len=seq_len, batch_size=1).to(self.device)
 
@@ -129,7 +129,7 @@ def test_various_sequence_lengths(self, seq_len):
     def test_various_batch_sizes(self, batch_size):
         """Test sparse attention with various batch sizes."""
         model = SimpleTransformerEncoderLayer(d_model=128, nhead=4).to(self.device)
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         x = model.get_input(d_model=128, seq_len=64, batch_size=batch_size).to(self.device)
 
diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_calibration_gpu.py b/tests/gpu/torch/sparsity/attention_sparsity/test_calibration_gpu.py
index 39f9837667a..b8edfe44c77 100644
--- a/tests/gpu/torch/sparsity/attention_sparsity/test_calibration_gpu.py
+++ b/tests/gpu/torch/sparsity/attention_sparsity/test_calibration_gpu.py
@@ -20,7 +20,7 @@
 from _test_utils.torch.sparsity.sparse_attention_common import SimpleTransformerEncoderLayer
 
 import modelopt.torch.opt as mto
-from modelopt.torch.sparsity.attention_sparsity import sparsify
+from modelopt.torch.sparsity import sparsify
 from modelopt.torch.sparsity.attention_sparsity.calibration import RulerDatasetBuilder
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_integration_gpu.py b/tests/gpu/torch/sparsity/attention_sparsity/test_integration_gpu.py
index 80408024a2e..55e03a1eb53 100644
--- a/tests/gpu/torch/sparsity/attention_sparsity/test_integration_gpu.py
+++ b/tests/gpu/torch/sparsity/attention_sparsity/test_integration_gpu.py
@@ -20,7 +20,7 @@
 from _test_utils.torch.transformers_models import create_tiny_llama_dir
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-import modelopt.torch.sparsity.attention_sparsity as sparse_attn
+import modelopt.torch.sparsity as mts
 from modelopt.torch.sparsity.attention_sparsity import SparseAttentionConfig
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
@@ -75,7 +75,7 @@ def test_load_and_sparsify(self, tinyllama_model):
             },
         )
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Verify sparse attention modules were added
         sparse_count = sum(
@@ -101,7 +101,7 @@ def test_forward_prefill(self, tinyllama_model, tinyllama_tokenizer):
             },
         )
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Create prefill input (seq_len > 1)
         test_text = "Once upon a time in a land far away"
@@ -134,7 +134,7 @@ def test_forward_decode(self, tinyllama_model):
             },
         )
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Create decode input (seq_len = 1)
         vocab_size = model.config.vocab_size
@@ -174,7 +174,7 @@ def test_gqa_attention(self, tinyllama_model):
             },
         )
 
-        sparse_model = sparse_attn.sparsify(model, sparse_config)
+        sparse_model = mts.sparsify(model, sparse_config)
 
         # Test forward pass with GQA
         input_ids = torch.randint(0, config.vocab_size, (1, 32), device="cuda")
diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_calibration_gpu.py b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_calibration_gpu.py
index 949e67b2cd8..c83a3aba707 100644
--- a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_calibration_gpu.py
+++ b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_calibration_gpu.py
@@ -29,7 +29,7 @@
 from _test_utils.torch.transformers_models import create_tiny_llama_dir
 from transformers import AutoModelForCausalLM
 
-import modelopt.torch.sparsity.attention_sparsity as mtsa
+import modelopt.torch.sparsity as mts
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE
 from modelopt.torch.kernels.common.attention.hf_triton_attention import triton_attention_forward
 from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_TRITON_CALIB
@@ -111,7 +111,7 @@ def test_calibrated_model_inference(self, tiny_llama_dir):
         config["sparse_cfg"]["calibration"]["target_sparse_ratio"] = {"prefill": 0.5}
 
         forward_loop = _make_forward_loop(model.config.vocab_size)
-        sparse_model = mtsa.sparsify(model, config, forward_loop=forward_loop)
+        sparse_model = mts.sparsify(model, config, forward_loop=forward_loop)
         assert sparse_model.config._attn_implementation == "modelopt_triton"
 
         sparse_model.eval()
diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py
index 814a03ba564..841e12162e9 100644
--- a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py
+++ b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py
@@ -17,7 +17,7 @@
 
 Uses ``create_tiny_wan22_pipeline_dir`` (matches diffusers' own tiny Wan 2.2
 test config): dual 2-layer transformer, tiny VAE, tiny T5 text encoder.
-The full ``WanPipeline`` is loaded, sparsified with ``mtsa.sparsify``, and
+The full ``WanPipeline`` is loaded, sparsified with ``mts.sparsify``, and
 run end-to-end with a 2-step denoising loop — asserting no NaN/Inf in the
 pipeline output.
 """
@@ -39,7 +39,7 @@
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE
 
 if TRITON_KERNEL_AVAILABLE:
-    import modelopt.torch.sparsity.attention_sparsity as mtsa
+    import modelopt.torch.sparsity as mts
     from modelopt.torch.sparsity.attention_sparsity.methods.triton_skip_softmax import (
         TritonSkipSoftmaxMethod,
     )
@@ -92,8 +92,8 @@ def _skip_softmax_cfg(threshold=0.03125):
 
 def _sparsify_both_transformers(pipe, cfg):
     """Apply sparsify to both transformer and transformer_2 (14B-style dual)."""
-    mtsa.sparsify(pipe.transformer, cfg)
-    mtsa.sparsify(pipe.transformer_2, cfg)
+    mts.sparsify(pipe.transformer, cfg)
+    mts.sparsify(pipe.transformer_2, cfg)
 
 
 def _run_pipe(pipe, seed=0):
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py b/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py
index c7c6f56928b..f818f249082 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py
@@ -27,7 +27,7 @@
 
 from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel
 
-from modelopt.torch.sparsity.attention_sparsity import sparsify
+from modelopt.torch.sparsity import sparsify
 from modelopt.torch.sparsity.attention_sparsity.calibration.calibrator import (
     DynamicThresholdCalibrator,
 )
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py
index e14d498a415..cc9228b900f 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py
@@ -26,7 +26,7 @@
 from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel
 from pydantic import ValidationError
 
-from modelopt.torch.sparsity.attention_sparsity import sparsify
+from modelopt.torch.sparsity import sparsify
 from modelopt.torch.sparsity.attention_sparsity.calibration import (
     DynamicThresholdCalibrator,
     RulerDatasetBuilder,
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
index 096d291321b..c26482472ac 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
@@ -60,7 +60,7 @@ def test_basic_replacement(self):
         assert original_attention_count > 0
 
         # Apply sparse attention
-        sparse_model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        sparse_model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         # Count sparse attention modules
         sparse_attention_count = sum(
@@ -88,7 +88,7 @@ def test_pattern_based_replacement(self):
             },
         }
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Verify sparse modules exist
         has_sparse = any(isinstance(m, SparseAttentionModule) for m in sparse_model.modules())
@@ -116,7 +116,7 @@ def filter_func(name):
             },
         }
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
         has_sparse = any(isinstance(m, SparseAttentionModule) for m in sparse_model.modules())
         assert has_sparse
 
@@ -135,13 +135,13 @@ def test_no_matching_modules(self):
         }
 
         # Should not error, even with no matches
-        sparse_attn.sparsify(model, config)
+        mts.sparsify(model, config)
 
     def test_disable_enable_functions(self):
         """Test disable/enable utility functions."""
 
         model = SimpleAttentionModel()
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         # Disable all
         disable_sparse_attention(model, "*")
@@ -158,7 +158,7 @@ def test_disable_enable_functions(self):
     def test_print_sparse_attention_summary(self, capsys):
         """Test print_sparse_attention_summary function."""
         model = SimpleAttentionModel()
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         # Print summary
         print_sparse_attention_summary(model)
@@ -172,7 +172,7 @@ def test_restore_sparse_attention_model(self):
         """Test save/restore via modelopt_state."""
         # Create and sparsify original model
         model_orig = SimpleAttentionModel()
-        model_orig = sparse_attn.sparsify(model_orig, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model_orig = mts.sparsify(model_orig, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         # Save state
         state_dict = mto.modelopt_state(model_orig)
@@ -196,7 +196,7 @@ def test_restore_keeps_calibration_params(self, tmp_path):
         calibration_params = {"prefill": {"a": 3.14, "b": 7.5}, "decode": {"a": 0.5, "b": 9.0}}
         target_sparse_ratio = {"prefill": 0.4, "decode": 0.6}
 
-        model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
         for module in model.modules():
             if isinstance(module, SparseAttentionModule):
                 module._sparse_method_instance.calibration_params = calibration_params
@@ -224,7 +224,7 @@ def fake_calibration(model, config, forward_loop=None):
                 if isinstance(module, SparseAttentionModule):
                     module._sparse_method_instance.calibration_params = calibration_params
 
-        model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
         model = mts.sparsify(model, "sparse_magnitude")
         with patch(
             "modelopt.torch.sparsity.attention_sparsity.model_sparsify.calibrate_sparse_attention",
@@ -242,17 +242,25 @@ def fake_calibration(model, config, forward_loop=None):
         for module in sparse_modules:
             assert module._sparse_method_instance.calibration_params == calibration_params
 
+    def test_calibrate_container_of_converted_models(self):
+        """calibrate() works on a module holding several separately converted models."""
+        models = nn.ModuleList(
+            mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) for _ in range(2)
+        )
+        sparse_attn.calibrate(models, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+
     @pytest.mark.parametrize("weight_first", [True, False])
     def test_chain_with_weight_sparsity(self, weight_first):
-        """Weight and attention sparsity chain in either order; export folds only weight masks."""
-        steps = [
-            lambda m: mts.sparsify(m, "sparse_magnitude"),
-            lambda m: sparse_attn.sparsify(m, FLASH_SKIP_SOFTMAX_DEFAULT_CFG),
-        ]
-        model = SimpleAttentionModel().eval()
-        for step in steps if weight_first else steps[::-1]:
-            model = step(model)
+        """Weight and attention sparsity in one call, either order; export folds only weight masks."""
+        modes = [("sparsegpt", {}), ("sparse_attention", FLASH_SKIP_SOFTMAX_DEFAULT_CFG)]
         x = SimpleAttentionModel.get_input()
+        model = mts.sparsify(
+            SimpleAttentionModel().eval(),
+            modes if weight_first else modes[::-1],
+            config={"device": "cpu"},  # SparseGPT-only option
+            forward_loop=lambda m: m(x),
+        )
+        assert (model.fc.weight == 0).float().mean().item() == 0.5
         with torch.no_grad():
             output = model(x)
 
@@ -266,12 +274,35 @@ def test_chain_with_weight_sparsity(self, weight_first):
         mts.export(model)
         assert not any(isinstance(m, SparseModule) for m in model.modules())
         assert any(isinstance(m, SparseAttentionModule) for m in model.modules())
+        assert mts.export(model) is model  # only sparse attention is left, which needs no export
         with torch.no_grad():
             assert torch.allclose(model(x), output)
 
+    @pytest.mark.parametrize(
+        ("mode", "config", "forward_loop"),
+        [
+            (FLASH_SKIP_SOFTMAX_DEFAULT_CFG, {"devcie": "cpu"}, None),
+            (FLASH_SKIP_SOFTMAX_DEFAULT_CFG, {"forward_loop": print}, print),
+            (FLASH_SKIP_SOFTMAX_DEFAULT_CFG, {"data_loader": [None]}, print),
+            ("quantize", None, None),
+        ],
+    )
+    def test_sparsify_rejects_invalid_input_before_converting(self, mode, config, forward_loop):
+        model = SimpleAttentionModel()
+        with pytest.raises(ValueError):
+            mts.sparsify(model, mode, config, forward_loop=forward_loop)
+        assert not mto.ModeloptStateManager.is_converted(model)
+
+    def test_deprecated_sparsify(self):
+        """The deprecated mtsa.sparsify warns and matches mts.sparsify."""
+        with pytest.deprecated_call():
+            model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model_ref = mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        assert mto.modelopt_state(model) == mto.modelopt_state(model_ref)
+
     def test_restore_keeps_disabled_modules(self):
         """Modules disabled after sparsify stay disabled after restore."""
-        model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
         disable_sparse_attention(model, "*attention*")
 
         model_restored = mto.restore_from_modelopt_state(
@@ -286,7 +317,7 @@ def test_restore_keeps_disabled_modules(self):
 
     def test_restore_legacy_state_without_calibration_keys(self):
         """Metadata saved before calibration state was persisted still restores."""
-        model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
         modelopt_state = mto.modelopt_state(model)
         for _, mode_state in modelopt_state["modelopt_state_dict"]:
             for module_state in mode_state["metadata"]["sparse_attention_state"].values():
@@ -319,7 +350,7 @@ def test_get_stats_with_stats_manager(self):
             },
         }
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Find sparse module
         sparse_module = None
@@ -355,7 +386,7 @@ def test_get_stats_without_stats_manager(self):
             },
         }
 
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # Find sparse module
         for module in sparse_model.modules():
@@ -468,14 +499,14 @@ class TestExportSparseAttentionConfig:
     def test_returns_none_without_calibration(self):
         """When no module has calibration params, returns None."""
         model = SimpleAttentionModel()
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
         out = export_sparse_attention_config(model)
         assert out is None
 
     def test_exports_when_calibration_present(self):
         """Calibration params are reflected in the exported config."""
         model = SimpleAttentionModel()
-        model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
+        model = mts.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG)
 
         for module in model.modules():
             if isinstance(module, SparseAttentionModule):
@@ -529,7 +560,7 @@ def test_exports_sparse_softmax_metadata(self):
             "modelopt.torch.kernels.sparsity.attention.register_triton_attention",
             MagicMock(return_value=True),
         ):
-            model = sparse_attn.sparsify(model, config)
+            model = mts.sparsify(model, config)
 
         out = export_sparse_attention_config(model)
 
@@ -548,7 +579,7 @@ def test_exports_calibrated_skip_softmax_with_sparse_softmax_overlay(self):
         config = copy.deepcopy(SKIP_SOFTMAX_CALIB_SPARSE24)
         config["sparse_cfg"].pop("calibration")
         config["sparse_cfg"]["*attention*"] = config["sparse_cfg"].pop("*attn*")
-        model = sparse_attn.sparsify(model, config)
+        model = mts.sparsify(model, config)
 
         for module in model.modules():
             if isinstance(module, SparseAttentionModule):
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_threshold_info.py b/tests/unit/torch/sparsity/attention_sparsity/test_threshold_info.py
index 00622d365d4..78849e1652f 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_threshold_info.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_threshold_info.py
@@ -21,7 +21,7 @@
 
 from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel
 
-from modelopt.torch.sparsity.attention_sparsity import sparsify
+from modelopt.torch.sparsity import sparsify
 from modelopt.torch.sparsity.attention_sparsity.methods.flash_skip_softmax import FlashSkipSoftmax
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_vsa.py b/tests/unit/torch/sparsity/attention_sparsity/test_vsa.py
index 09621656e8c..af53ca5bc77 100644
--- a/tests/unit/torch/sparsity/attention_sparsity/test_vsa.py
+++ b/tests/unit/torch/sparsity/attention_sparsity/test_vsa.py
@@ -322,7 +322,7 @@ def test_vsa_config_defaults(self):
 from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel
 
 import modelopt.torch.opt as mto
-import modelopt.torch.sparsity.attention_sparsity as sparse_attn
+import modelopt.torch.sparsity as mts
 from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule
 
 VSA_TEST_CFG = {
@@ -344,7 +344,7 @@ class TestVSASparsifyIntegration:
     def test_sparsify_creates_sparse_modules(self):
         """sparsify() with VSA config replaces attention modules."""
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         sparse_modules = [m for m in sparse_model.modules() if isinstance(m, SparseAttentionModule)]
         assert len(sparse_modules) > 0
@@ -352,7 +352,7 @@ def test_sparsify_creates_sparse_modules(self):
     def test_sparse_module_has_vsa_method(self):
         """Replaced modules are configured with VSA method."""
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         for module in sparse_model.modules():
             if isinstance(module, SparseAttentionModule):
@@ -364,7 +364,7 @@ def test_sparse_module_has_vsa_method(self):
     def test_enable_disable(self):
         """Enable/disable works on VSA sparse modules."""
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         for module in sparse_model.modules():
             if isinstance(module, SparseAttentionModule):
@@ -377,7 +377,7 @@ def test_enable_disable(self):
     def test_threshold_info(self):
         """VSA sparse modules report correct threshold info."""
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         for module in sparse_model.modules():
             if isinstance(module, SparseAttentionModule):
@@ -388,7 +388,7 @@ def test_threshold_info(self):
     def test_save_restore(self):
         """VSA modelopt_state can be saved and restored."""
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         state = mto.modelopt_state(sparse_model)
 
@@ -416,7 +416,7 @@ def test_pattern_matching(self):
                 "default": {"enable": False},
             },
         }
-        sparse_model = sparse_attn.sparsify(model, config)
+        sparse_model = mts.sparsify(model, config)
 
         # No modules should have VSA enabled
         for module in sparse_model.modules():
@@ -431,7 +431,7 @@ def test_forward_patches_sdpa(self):
         ImportError — proving the interception works.
         """
         model = SimpleAttentionModel()
-        sparse_model = sparse_attn.sparsify(model, VSA_TEST_CFG)
+        sparse_model = mts.sparsify(model, VSA_TEST_CFG)
 
         # Set video_shape so metadata can be computed.
         # seq_len=64, video_shape (4,4,4) -> T*H*W=64