[OMNIML-2850] [3/n] Adds sparse attention calibration (#538)

## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
new feature

**Overview:** ?
- This PR adds the sparse attention calibration algorithm
- Chunked prefill to support long ctx_len
- Separated calibration for prefill and decode

## Usage
<!-- You can potentially add a usage example below. -->

```python
import modelopt.torch.sparsity.attention_sparsity as mtsa

# Apply sparse attention with calibration
model = mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)

# Print summary - now shows actual thresholds
mtsa.print_sparse_attention_summary(model)
# Output:
# Method: flash_skip_softmax, Threshold: Dynamic (λ=437.395926)

# Or llm_eval integration
# HuggingFace sparse attention example
python examples/llm_sparsity/attention_sparsity/hf_sa.py \
    --pyt_ckpt_path Qwen/Qwen3-4B \
    --sparse_attn skip_softmax_calib 
```

# The calibration method
## Calibration Algorithm
- Implemented the Inverse Power model: scale_factor = k / (1 -
sparsity)^p
- Fit model parameters (k, p) per phase using scipy.optimize.curve_fit
- At inference: threshold = k / (1 - target_sparsity)^p / seqlen

## Why Choosing the Inverse Power model?
The inverse power model better fits the relationship between sparsity
ratio and threshold_scale_factor.
<img width="2388" height="1082" alt="sparsity_model_analysis"
src="https://github.com/user-attachments/assets/4dfb45d4-8c16-4f15-a878-c8e08a9b6128"
/>

## Runtime Flexibility
- Target sparsity can be changed at inference time without recalibration
- Users can adjust module._sparse_method_instance.target_sparse_ratio
dynamically
- Threshold automatically adapts to sequence length

## Testing
<!-- Mention how have you tested your change if applicable. -->
The calibration results for `Qwen/Qwen3-30B-A3B-Thinking-2507` are shown
below and are mostly consistent with the ground-truth numbers collected
from the kernel side.
```
Prefill Calibration Results:
  Model: scale_factor = k / (1 - sparsity)^p
  Fitted k: 1003.3990
  Fitted p: 1.2589
  R-squared: 0.827549

Scale factors for different target sparsities:
  Target     Scale Factor
  ---------- ---------------
  50%        2401.35
  70%        4568.26
  80%        7610.98
  90%        18214.70
  95%        43591.65
```

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

---------

Signed-off-by: Kai Xu <kaix@nvidia.com>

Author: kaix-nv

CHANGELOG.rst

@@ -1,6 +1,13 @@
 NVIDIA Model Optimizer Changelog (Linux)
 ========================================
 
+0.43 (2026-03-xx)
+^^^^^^^^^^^^^^^^^
+
+**New Features**
+
+- Add sparse attention optimization for transformer models (``modelopt.torch.sparsity.attention_sparsity``). This reduces computational cost by skipping attention computation. Supports calibration for threshold selection on HuggingFace models. See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
+
 0.42 (2026-02-xx)
 ^^^^^^^^^^^^^^^^^
 

examples/llm_eval/lm_eval_hf.py

@@ -43,9 +43,11 @@
 from lm_eval.api.model import T
 from lm_eval.models.huggingface import HFLM
 from quantization_utils import quantize_model
+from sparse_attention_utils import sparsify_model
 
 import modelopt.torch.opt as mto
 from modelopt.torch.quantization.utils import is_quantized
+from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
 
 
 def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict | None = None) -> T:
@@ -60,6 +62,9 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
     calib_size = arg_dict.pop("calib_size", 512)
     compress = arg_dict.pop("compress", False)
 
+    # Sparse attention arguments
+    sparse_cfg = arg_dict.pop("sparse_cfg", None)
+
     additional_config = {} if additional_config is None else additional_config
     additional_config = {k: v for k, v in additional_config.items() if v is not None}
 
@@ -91,6 +96,15 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
             auto_quantize_checkpoint=auto_quantize_checkpoint,
         )
 
+    if sparse_cfg:
+        if is_attn_sparsified(model_obj.model):
+            warnings.warn("Skipping sparse attention: model already has sparse attention applied.")
+        else:
+            sparsify_model(
+                model=model_obj,
+                sparse_cfg=sparse_cfg,
+            )
+
     return model_obj
 
 
@@ -152,6 +166,11 @@ def setup_parser_with_modelopt_args():
         action="store_true",
         help="Compress the model after quantization",
     )
+    parser.add_argument(
+        "--sparse_cfg",
+        type=str,
+        help="Sparse attention configuration (e.g., SKIP_SOFTMAX_DEFAULT, SKIP_SOFTMAX_CALIB)",
+    )
     return parser
 
 
@@ -177,6 +196,7 @@ def setup_parser_with_modelopt_args():
             "calib_batch_size": args.calib_batch_size,
             "calib_size": args.calib_size,
             "compress": args.compress,
+            "sparse_cfg": args.sparse_cfg,
         }
     )
 

examples/llm_eval/mmlu.py

@@ -48,6 +48,7 @@
 from fire import Fire
 from modeling import EvalModel, select_model
 from quantization_utils import MAX_SEQ_LEN, get_tokenizer, quantize_model
+from sparse_attention_utils import sparsify_model
 from tqdm import tqdm
 
 try:
@@ -56,6 +57,7 @@
     LLM = None  # type: ignore[misc]
 import modelopt.torch.opt as mto
 from modelopt.torch.quantization.utils import is_quantized
+from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -230,6 +232,7 @@ def main(
     auto_quantize_method: str = "gradient",
     auto_quantize_score_size: int = 128,
     auto_quantize_checkpoint: str | None = None,
+    sparse_cfg: str | None = None,
     **kwargs,
 ):
     random.seed(RAND_SEED)
@@ -289,6 +292,20 @@ def main(
                     auto_quantize_checkpoint=auto_quantize_checkpoint,
                 )
 
+        # Apply sparse attention if requested
+        if sparse_cfg:
+            model.load()
+
+            if is_attn_sparsified(model.model):
+                warnings.warn(
+                    "Skipping sparse attention: model already has sparse attention applied."
+                )
+            else:
+                sparsify_model(
+                    model=model,
+                    sparse_cfg=sparse_cfg,
+                )
+
     for subject in tqdm(subjects):
         dev_df = pd.read_csv(os.path.join(data_dir, "dev", subject + "_dev.csv"), header=None)[
             :ntrain

examples/llm_eval/modeling.py

@@ -179,6 +179,7 @@ class SeqToSeqModel(EvalModel):
     lora_path: str = ""
     device: str = "cuda"
     load_8bit: bool = False
+    attn_implementation: str | None = None
 
     def load(self):
         if self.model is None:
@@ -188,6 +189,8 @@ def load(self):
             if self.load_8bit:
                 args.update(device_map="auto", load_in_8bit=True)
             args.update(torch_dtype=getattr(torch, self.dtype) if self.dtype != "auto" else "auto")
+            if self.attn_implementation:
+                args["attn_implementation"] = self.attn_implementation
             self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_path, **args)
             print_gpu_utilization()
             if self.lora_path:
@@ -241,6 +244,8 @@ def load(self):
             if self.load_8bit:
                 args.update(device_map="auto", load_in_8bit=True)
             args.update(torch_dtype=getattr(torch, self.dtype) if self.dtype != "auto" else "auto")
+            if self.attn_implementation:
+                args["attn_implementation"] = self.attn_implementation
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_path, trust_remote_code=True, **args
             )

examples/llm_eval/sparse_attention_utils.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Utilities for sparse attention integration with llm_eval."""
+
+import modelopt.torch.sparsity.attention_sparsity as mtsa
+
+
+def _extract_model(model_obj):
+    """Extract actual model from wrapper (HFLM or EvalModel)."""
+    if hasattr(model_obj, "gpt2"):
+        return model_obj.gpt2
+    elif hasattr(model_obj, "model"):
+        return model_obj.model
+    else:
+        return model_obj
+
+
+def sparsify_model(
+    model,
+    sparse_cfg: str,
+    backend=None,
+):
+    """Apply sparse attention to model with optional RULER calibration.
+
+    Args:
+        model: Model wrapper (HFLM or EvalModel) or raw model
+        sparse_cfg: Sparse attention config name or dict
+        backend: Backend to use (optional, overrides config backend)
+
+    Returns:
+        The model with sparse attention applied
+
+    Note:
+        Calibration is automatically triggered if the config contains a 'calibration' field.
+        The calibration will auto-generate RULER dataset from the model's tokenizer.
+    """
+    # Extract actual model
+    net = _extract_model(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)
+        if mtsa_cfg is None:
+            raise ValueError(f"Unknown sparse_cfg: {sparse_cfg}.")
+    else:
+        mtsa_cfg = sparse_cfg
+
+    # Override backend if specified
+    if backend:
+        if isinstance(mtsa_cfg, dict) and "sparse_cfg" in mtsa_cfg:
+            modified_sparse_cfg = {}
+            for pattern, cfg in mtsa_cfg["sparse_cfg"].items():
+                modified_cfg = cfg.copy() if isinstance(cfg, dict) else cfg
+                if isinstance(modified_cfg, dict):
+                    modified_cfg["backend"] = backend
+                modified_sparse_cfg[pattern] = modified_cfg
+            mtsa_cfg = {"sparse_cfg": modified_sparse_cfg}
+
+    # Apply sparsification
+    print(f"\nApplying sparse attention with config: {sparse_cfg}")
+    mtsa.sparsify(net, mtsa_cfg)
+    print("Sparse attention applied successfully!")
+
+    return model

examples/llm_sparsity/attention_sparsity/.gitignore

@@ -0,0 +1,2 @@
+# Data directory for calibration
+data

examples/llm_sparsity/attention_sparsity/README.md

@@ -0,0 +1,165 @@
+# Attention Sparsity for HuggingFace Models
+
+In this tutorial, we demonstrate how to use NVIDIA Model Optimizer to apply attention sparsity to HuggingFace models. Attention sparsity reduces computational cost by skipping near-zero attention scores during the softmax computation.
+
+## Getting Started
+
+### Quick Example
+
+```python
+import modelopt.torch.sparsity.attention_sparsity as mtsa
+from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_DEFAULT
+
+# Load your model
+model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen3-8B",
+    attn_implementation="eager",  # Required for sparse attention
+    torch_dtype=torch.bfloat16,
+)
+
+# Apply sparse attention
+model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
+```
+
+> [!Note]
+> `attn_implementation="eager"` is required for sparse attention to work properly. Flash Attention 2 or SDPA would bypass the softmax patching needed for stats collection.
+
+## Configuration Options
+
+Two pre-defined configurations are available:
+
+### 1. Fixed Threshold (SKIP_SOFTMAX_DEFAULT)
+
+Uses a fixed threshold value. Simple but may not be optimal for all sequence lengths.
+
+```python
+from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_DEFAULT
+
+model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
+```
+
+### 2. Calibrated Threshold (SKIP_SOFTMAX_CALIB)
+
+Uses RULER-based calibration to determine an optimal dynamic threshold that adapts to sequence length. Recommended for production use.
+
+```python
+from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_CALIB
+
+model = mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)
+```
+
+## Prerequisites
+
+### Local Installation
+
+For Hugging Face models, install Model Optimizer with `hf` dependencies using `pip` from [PyPI](https://pypi.org/project/nvidia-modelopt/) and install the requirements for the example:
+
+```bash
+pip install nvidia-modelopt[hf]
+```
+
+### Download RULER Calibration Data (Required for Calibration)
+
+If using `SKIP_SOFTMAX_CALIB`, you need to download the RULER calibration dataset first:
+
+```bash
+bash ./download_ruler_data.sh
+```
+
+This downloads the Paul Graham essays dataset used for generating calibration samples.
+
+## Run Sparse Attention on HuggingFace Models
+
+### Basic Usage (Without Calibration)
+
+Apply sparse attention with a fixed threshold:
+
+```bash
+python hf_sa.py \
+    --pyt_ckpt_path Qwen/Qwen3-8B \
+    --sparse_attn skip_softmax
+```
+
+### With RULER Calibration
+
+Apply sparse attention with calibrated thresholds for optimal sparsity:
+
+```bash
+python hf_sa.py \
+    --pyt_ckpt_path Qwen/Qwen3-8B \
+    --sparse_attn skip_softmax_calib
+```
+
+The calibration process:
+
+1. Generates RULER calibration samples
+2. Collects attention statistics during forward passes
+3. Determines optimal threshold scale factor for target sparsity ratio
+
+### Command Line Arguments
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--pyt_ckpt_path` | Required | HuggingFace model path or name |
+| `--sparse_attn` | `skip_softmax` | Configuration: `skip_softmax` or `skip_softmax_calib` |
+| `--backend` | `pytorch` | Backend: `pytorch` (only supported backend) |
+| `--seq_len` | `2048` | Maximum sequence length for input prompts |
+| `--export_dir` | `None` | Directory to export the sparsified model |
+
+## Output Comparison
+
+The script automatically compares outputs before and after applying sparse attention:
+
+1. Loads a test sample from the NarrativeQA dataset
+2. Generates text before sparse attention is applied
+3. Applies sparse attention (with optional calibration)
+4. Generates text after sparse attention is applied
+5. Compares and displays both outputs
+
+## Export Model
+
+Export the sparsified model to a HuggingFace checkpoint:
+
+```bash
+python hf_sa.py \
+    --pyt_ckpt_path Qwen/Qwen3-8B \
+    --sparse_attn skip_softmax_calib \
+    --export_dir ./exported_sparse_model
+```
+
+The exported model can be loaded and used with standard HuggingFace APIs.
+
+## Custom Configuration
+
+You can create custom sparse attention configurations:
+
+```python
+custom_config = {
+    "sparse_cfg": {
+        "calibration": {  # Optional: omit for fixed threshold
+            "target_sparse_ratio": {"prefill": 0.5, "decode": 0.5},  # Target 50% sparsity
+            "samples": 128,              # Number of calibration samples
+            "max_seqlen": 8192,          # Maximum sequence length
+            # Optional: customize threshold trials for calibration
+            "threshold_trials": [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1, 2e-1, 3e-1, 5e-1, 7e-1],
+        },
+        "*attn*": {  # Pattern to match attention modules
+            "method": "flash_skip_softmax",
+            "threshold": {"prefill": 1e-3, "decode": 1e-4},  # Phase-specific thresholds (ignored if calibration is used)
+            "br": 128,          # Flash Attention block rows
+            "bc": 128,          # Flash Attention block columns
+            "backend": "pytorch",
+            "collect_stats": True,
+            "enable": True,
+        },
+        "default": {"enable": False},
+    },
+}
+
+model = mtsa.sparsify(model, config=custom_config)
+```
+
+## References
+
+- [Model Optimizer Documentation](https://nvidia.github.io/Model-Optimizer/)
+- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://github.com/NVIDIA/RULER)