[2/n] Add sparse softmax to the Triton flash attention kernel (#1078)

### What does this PR do?

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

Type of change: New feature

Add N:M structured sparsity support to the Triton flash attention kernel
(`modelopt/torch/kernels/triton_fa.py`). For every M consecutive key
positions in the attention score tile, keeps the top-N values and sets
the rest to -inf before softmax. This is applied during prefill only.

**Supported patterns:** Any N:M where M=4 (N=1,2,3) or M=8 (N=1..4).

- Sink tokens and dense window blocks for preserving local attention and
attention sinks

 **Performance (TFLOPS at seq_len=16384, RTX 6000):**
  | Pattern | TFLOPS | % of Dense |
  |---------|--------|------------|
  | Dense | 89.3 | 100% |
  | 2:4 (M=4) | 69.5 | 78% |
  | 4:8 (M=8) | 57.3 | 64% |

### Usage

```python
# Add a code snippet demonstrating how to use this
from modelopt.torch.kernels import attention

# 2:4 sparsity (keep top 2 of every 4 K positions)
out = attention(q, k, v, b_start_loc, b_seq_len, max_len,
                sparsity_n=2, sparsity_m=4)

# 4:8 sparsity with sink tokens and dense window
out = attention(q, k, v, b_start_loc, b_seq_len, max_len,
                sparsity_n=4, sparsity_m=8,
                num_sink_tokens=4, dense_window_blocks=2)

# Dense (default, zero overhead)
out = attention(q, k, v, b_start_loc, b_seq_len, max_len)

# Via mtsa.sparsify() on HuggingFace models
import modelopt.torch.sparsity.attention_sparsity as mtsa
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B",
                                              torch_dtype=torch.bfloat16,
                                              device_map="cuda")

# Default config
mtsa.sparsify(model, mtsa.SPARSE_SOFTMAX_DEFAULT)
```

### Testing
<!-- Mention how have you tested your change if applicable. -->

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->

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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* N:M structured sparse softmax for Triton flash-attention prefill with
configurable dense-window and sink-token handling.

* **API**
* attention(...) accepts sparsity_n, sparsity_m, num_sink_tokens,
dense_window_size; HF/Triton prefill path propagates them.

* **Configuration**
* New config fields and exported preset to enable/configure Triton N:M
sparse softmax with validation.

* **Tests**
* Added GPU tests covering N:M behavior, tile structure,
forward/backward correctness.

* **Documentation**
  * CHANGELOG and example docs updated with usage and CLI options.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: kaix-nv

CHANGELOG.rst

@@ -7,6 +7,7 @@ NVIDIA Model Optimizer Changelog
 
 - Support full Transformer Engine spec for Minitron pruning (``mcore_minitron``). Now we no longer need to use custom ModelOpt spec. Note that this does not affect the usage of the pruning workflow but makes pruning slightly faster and may result in slightly different pruned model because of different kernel and numerics.
 - Added iterator interface using CalibrationDataReader in ONNX quantization workflow.
+- Add N:M sparse softmax support to the Triton flash attention kernel (``modelopt.torch.kernels.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
 - Enable PTQ workflow for the Step3.5-Flash MoE model with NVFP4 W4A4 + FP8 KV cache quantization. See `modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml <https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml>`_ for more details.
 
 **Bug Fixes**

examples/llm_sparsity/attention_sparsity/README.md

@@ -1,6 +1,11 @@
 # 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. Two attention backends are supported:
+In this tutorial, we demonstrate how to use NVIDIA Model Optimizer to apply attention sparsity to HuggingFace models. Two sparsity methods are supported:
+
+- **Skip-softmax** (`flash_skip_softmax`): Skips attention tiles whose contribution is negligible, based on a threshold. Based on the [BLASST](https://arxiv.org/pdf/2512.12087) algorithm.
+- **N:M sparse softmax** (`triton_sparse_softmax`): For every M consecutive key positions, keeps the top-N attention scores and sets the rest to -inf before softmax.
+
+Two attention backends are available:
 
 - **pytorch** (default): Patches `F.softmax` to apply skip-softmax sparsity (requires `attn_implementation="eager"`)
 - **triton**: Uses a fused Triton Flash Attention kernel with in-kernel sparsity (uses `attn_implementation="modelopt_triton"`)
@@ -29,9 +34,9 @@ model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
 
 ## Configuration Options
 
-Two pre-defined configurations are available:
+### Skip-Softmax
 
-### 1. Fixed Threshold (SKIP_SOFTMAX_DEFAULT)
+#### 1. Fixed Threshold (SKIP_SOFTMAX_DEFAULT)
 
 Uses a fixed threshold value. Simple but may not be optimal for all sequence lengths.
 
@@ -41,7 +46,7 @@ from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_DEFAU
 model = mtsa.sparsify(model, config=SKIP_SOFTMAX_DEFAULT)
 ```
 
-### 2. Calibrated Threshold (SKIP_SOFTMAX_CALIB)
+#### 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.
 
@@ -51,6 +56,46 @@ from modelopt.torch.sparsity.attention_sparsity.config import SKIP_SOFTMAX_CALIB
 model = mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)
 ```
 
+### N:M Sparse Softmax (SPARSE_SOFTMAX_DEFAULT)
+
+Applies N:M structured sparsity to attention scores using the Triton backend. For every M consecutive key positions, keeps only the top-N scores and sets the rest to -inf. Supports M=4 (N=1,2,3) and M=8 (N=1..7). Attention sinks and a local dense window can be configured to preserve important positions.
+
+```python
+from modelopt.torch.sparsity.attention_sparsity.config import SPARSE_SOFTMAX_DEFAULT
+
+model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Llama-3.1-8B",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+)
+
+model = mtsa.sparsify(model, config=SPARSE_SOFTMAX_DEFAULT)
+```
+
+Custom N:M configuration:
+
+```python
+sparse_cfg = {
+    "sparse_cfg": {
+        "*attn*": {
+            "method": "triton_sparse_softmax",
+            "sparsity_n": 2,            # Keep top-2 of every 4
+            "sparsity_m": 4,            # Group size
+            "num_sink_tokens": 4,       # Keep first 4 tokens dense (attention sinks)
+            "dense_window_size": 128,   # Keep tokens within distance 128 dense
+            "backend": "triton",
+            "enable": True,
+        },
+        "default": {"enable": False},
+    },
+}
+
+model = mtsa.sparsify(model, config=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.
+
 ## Prerequisites
 
 ### Local Installation
@@ -104,8 +149,8 @@ The calibration process:
 | 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) |
+| `--sparse_attn` | `skip_softmax` | Configuration: `skip_softmax`, `skip_softmax_calib`, or `sparse_softmax` |
+| `--backend` | `pytorch` | Backend: `pytorch` (skip-softmax) or `triton` (N:M sparse softmax) |
 | `--seq_len` | `2048` | Maximum sequence length for input prompts |
 | `--export_dir` | `None` | Directory to export the sparsified model |
 
@@ -166,3 +211,4 @@ model = mtsa.sparsify(model, config=custom_config)
 
 - [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)
+- [BLASST: Block-Level Adaptive Structured Sparse Training](https://arxiv.org/pdf/2512.12087) — skip-softmax algorithm

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -31,6 +31,7 @@
 from modelopt.torch.sparsity.attention_sparsity.config import (
     SKIP_SOFTMAX_CALIB,
     SKIP_SOFTMAX_DEFAULT,
+    SPARSE_SOFTMAX_DEFAULT,
 )
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
 
@@ -43,6 +44,7 @@
 SPARSE_ATTN_CFG_CHOICES = {
     "skip_softmax": SKIP_SOFTMAX_DEFAULT,
     "skip_softmax_calib": SKIP_SOFTMAX_CALIB,
+    "sparse_softmax": SPARSE_SOFTMAX_DEFAULT,
 }
 
 
@@ -168,9 +170,10 @@ def main(args):
 
     # Apply CLI overrides to sparse_cfg
     sparse_cfg = sparse_config.get("sparse_cfg", {})
-    for layer_cfg in sparse_cfg.values():
-        if isinstance(layer_cfg, dict) and "method" in layer_cfg:
-            layer_cfg["backend"] = args.backend
+    if args.backend is not None:
+        for layer_cfg in sparse_cfg.values():
+            if isinstance(layer_cfg, dict) and "method" in layer_cfg:
+                layer_cfg["backend"] = args.backend
     if args.target_sparse_ratio is not None:
         calib = sparse_cfg.setdefault("calibration", {})
         assert isinstance(calib, dict)
@@ -240,9 +243,10 @@ def main(args):
     parser.add_argument(
         "--backend",
         type=str,
-        default="pytorch",
+        default=None,
         choices=["pytorch", "triton"],
-        help="Backend for sparse attention (default: pytorch). 'triton' uses the fused Triton kernel.",
+        help="Backend for sparse attention. Overrides the config default if set. "
+        "'triton' uses the fused Triton kernel.",
     )
 
     # Sequence length arguments

modelopt/torch/kernels/hf_triton_attention.py

@@ -105,6 +105,17 @@ def triton_attention_forward(
         kw["b_seq_len_k"] = torch.full((batch,), seq_k, device=device, dtype=torch.int32)
         kw["max_input_len_k"] = seq_k
 
+    # N:M sparse softmax — prefill only (decode should not sparsify KV)
+    if not is_decode and getattr(module, "_apply_sparse_nm", False):
+        # _sparse_method_instance is set by SparseAttentionModule._init_sparse_method()
+        # in modelopt/torch/sparsity/attention_sparsity/sparse_attention.py
+        method = getattr(module, "_sparse_method_instance", None)
+        if method is not None:
+            kw["sparsity_n"] = getattr(method, "sparsity_n", 2)
+            kw["sparsity_m"] = getattr(method, "sparsity_m", 4)
+            kw["num_sink_tokens"] = getattr(method, "num_sink_tokens", 0)
+            kw["dense_window_size"] = getattr(method, "dense_window_size", 64)
+
     o = attention(q, k, v, **kw)
 
     attn_output = o.view(batch, seq_len, num_heads, head_dim)