Add fp16 mixed precision softmax

[kaix-nv]

What does this PR do?

Type of change: ?
New feature.
Type of change: New feature

Adds opt-in mixed-FP16 online-softmax support to the compact NVFP4 attention path.

• Adds MODELOPT_ATTN_SOFTMAX_MODE=fp32|mixed_fp16; default behavior remains fp32.

• Uses native FP16 exp2 for tile probabilities and online-max corrections while retaining FP32 denominator state,
  accumulators, and split-K reconciliation.

• Supports Triton prefill and fixed 32-split decode.

• Propagates the mode through FlashAttention and FlashInfer adapters, including mixed decode/prefill batches.

• Composes with NVFP4 Q/K/P/V fake quantization and optional prefill N:M sparsity.

• Keeps mixed mode inference-only and rejects autograd before kernel launch.

• Propagates the selector to Ray workers across supported vLLM executor layouts.

• Documents hardware and serving limitations. The compact NVFP4 worker requires SM89+.

Usage

cd examples/vllm_serve

  MODELOPT_ATTN_SOFTMAX_MODE=mixed_fp16 \
  python vllm_serve_sparse_attn.py <MODEL_PATH> -tp 8 \
    --enforce-eager \
    --no-enable-prefix-caching \
    --worker-cls sparse_attn_worker.QuantSparseAttnWorker

  Use MODELOPT_ATTN_SOFTMAX_MODE=fp32 or leave it unset to preserve the existing FP32 softmax behavior.

Testing

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ / ❌ / N/A
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ / ❌ / N/A
• Did you write any new necessary tests?: ✅ / ❌ / N/A
• Did you update Changelog?: ✅ / ❌ / N/A
• Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8204d5e5-7ca1-4ed8-9749-229c3c9fdc20

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch kaix/sparse_attn_quant_mixed_fp16

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1950/
Built to branch gh-pages at 2026-07-09 16:51 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 54.81799% with 211 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.95%. Comparing base (75b5803) to head (a3f54ac).
⚠️ Report is 20 commits behind head on main.

Files with missing lines	Patch %	Lines
...torch/kernels/common/attention/decode_attention.py	13.11%	106 Missing ⚠️
...t/torch/kernels/quantization/attention/bmm2_qdq.py	0.00%	47 Missing ⚠️
...delopt/torch/kernels/common/attention/triton_fa.py	37.70%	38 Missing ⚠️
.../torch/sparsity/attention_sparsity/plugins/vllm.py	93.13%	14 Missing ⚠️
...t/torch/kernels/quantization/common/nvfp4_quant.py	0.00%	4 Missing ⚠️
...ernels/quantization/attention/softmax_fakequant.py	90.00%	1 Missing ⚠️
modelopt/torch/quantization/plugins/vllm.py	94.73%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1950      +/-   ##
==========================================
+ Coverage   61.21%   67.95%   +6.73%     
==========================================
  Files         515      523       +8     
  Lines       57245    62047    +4802     
==========================================
+ Hits        35043    42164    +7121     
+ Misses      22202    19883    -2319     

Flag	Coverage Δ
unit	55.01% <6.00%> (+0.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.