[OMNIML-2852] [2/n] Add Core Sparse Attention Infrastructure

[kaix-nv]

What does this PR do?

Type of change: ?
New feature

Overview: ?
This PR provides a sparse attention support in ModelOpt for applying attention sparsity through skip softmax method, enabling inference speedups for LLMs.

Key Features:

• Skip softmax support
• Sparse attention config
• Extensible method registry for future sparse attention algorithms
• HuggingFace Transformers integration
• Phase-aware thresholds (separate prefill/decode)

Design doc

Usage

import torch
import modelopt.torch.sparsity.attention_sparsity as mts
from transformers import AutoModelForCausalLM

# Load model (must use eager attention for softmax patching)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    attn_implementation="eager",  # Required!
    torch_dtype=torch.bfloat16,
)

# Use pre-defined configuration
from modelopt.torch.sparsity.attention_sparsity import SKIP_SOFTMAX_DEFAULT
model = mts.sparsify(model, SKIP_SOFTMAX_DEFAULT)

Testing

Unit Test

pytest tests/unit/torch/sparsity/attention_sparsity -v
pytest tests/gpu/torch/sparsity/attention_sparsity -v
pytest tests/examples/llm_sparsity/attention_sparsity -v

ALL PASSED.

Accuracy

Benchmark: MMLU
Model: Qwen/Qwen3-4B
Cmd: python mmlu.py --model_name causal --model_path Qwen/Qwen3-4B --sparse_cfg SKIP_SOFTMAX_DEFAULT

	MMLU
BF16	69.96
SKIP_SOFTMAX_DEFAULT	69.86

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

[codecov]

Codecov Report

❌ Patch coverage is 89.67254% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.95%. Comparing base (fa84955) to head (cd6fce2).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...ch/sparsity/attention_sparsity/sparse_attention.py	71.42%	16 Missing ⚠️
...pt/torch/sparsity/attention_sparsity/conversion.py	93.02%	9 Missing ⚠️
...y/attention_sparsity/methods/flash_skip_softmax.py	90.90%	8 Missing ⚠️
...ch/sparsity/attention_sparsity/methods/registry.py	88.88%	3 Missing ⚠️
modelopt/torch/sparsity/attention_sparsity/mode.py	90.32%	3 Missing ⚠️
...delopt/torch/sparsity/attention_sparsity/config.py	95.91%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #527      +/-   ##
==========================================
+ Coverage   74.64%   74.95%   +0.31%     
==========================================
  Files         183      192       +9     
  Lines       18542    18939     +397     
==========================================
+ Hits        13840    14196     +356     
- Misses       4702     4743      +41     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Hi @kaix-nv could you further split this code change? This PR has 3000+ lines of code change and many file moves

[kevalmorabia97]

Can we use a single registry for all Sparsity algorithms and modes and then use top-level mts.sparsify(model, mode=...) so all algorithms (e.g. weight or attention sparsify) are invoked by single shared API instead of separate API per algorithm?

[kevalmorabia97]

No need for 1-gpu marker. All tests are run on 1 or more gpus only

[kevalmorabia97]

tests inside tests/gpu dont need a gpu check. Its assumed it is run only on gpu enabled machines.

Same applies to all test files

[kevalmorabia97]

why do we need such large hidden and intermediate size? Can we use 32/64 instead?

[kevalmorabia97]

these tests seem unnecessary