[OMNIML-2850] [3/n] Adds sparse attention calibration

[kaix-nv]

What does this PR do?

Type of change: ?
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

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.

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

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"

• 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 68.67196% with 276 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.54%. Comparing base (3801923) to head (179e8dd).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...arsity/attention_sparsity/calibration/calibrate.py	26.24%	104 Missing ⚠️
...rsity/attention_sparsity/calibration/calibrator.py	44.00%	70 Missing ⚠️
...ty/attention_sparsity/calibration/ruler_dataset.py	84.90%	53 Missing ⚠️
...pt/torch/sparsity/attention_sparsity/conversion.py	59.45%	30 Missing ⚠️
...ch/sparsity/attention_sparsity/sparse_attention.py	63.15%	7 Missing ⚠️
...delopt/torch/sparsity/attention_sparsity/config.py	90.76%	6 Missing ⚠️
...y/attention_sparsity/methods/flash_skip_softmax.py	90.00%	5 Missing ⚠️
...ch/sparsity/attention_sparsity/methods/registry.py	85.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #538      +/-   ##
==========================================
- Coverage   73.74%   73.54%   -0.21%     
==========================================
  Files         199      205       +6     
  Lines       21183    22000     +817     
==========================================
+ Hits        15621    16179     +558     
- Misses       5562     5821     +259     

☔ 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.

[kevalmorabia97]

@kaix-nv github is showing 7000+ lines of code as part of this PR. Is that accurate?
It shouldn’t be that much. Less than half of the code should remain after rebasing on the preceding PR.

[kevalmorabia97]

How long does running pytest tests/gpu/torch/sparsity take?

[kevalmorabia97]

+1

[kaix-nv]

@kevalmorabia97 All feedback has been addressed. Please take another look. Thanks.

[kevalmorabia97]

Some minor comments. Otherwise LGTM. Thanks for addressing my comments

[kevalmorabia97]

Can we merge this file with ruler_utils.py and name ruler_dataset.py? Both are specific for ruler dataset only

[kaix-nv]

Updated

[Edwardf0t1]

Add codeowner approval.

@kaix-nv Will the VSA support be your next PR? =)

[kaix-nv]

Add codeowner approval.

@kaix-nv Will the VSA support be your next PR? =)

Yes, the VSA PR will be submitted soon.