[3/n] Add vLLM mask reuse capture

[kaix-nv]

What does this PR do?

Type of change: new feature

Adds the second mask-reuse sub-PR: vLLM V1 dense-reference capture, FA4 source witnessing, provenance validation, rank merging, and the sparse-attention calibration hook needed to collect compact mask-reuse evidence. It does not add topology selection.

This PR is stacked on #2235.

Usage

python examples/vllm_serve/create_fa4_source_witness.py \
  /path/to/flash-attention \
  --expected-commit <COMMIT> \
  --archive-output fa4-runtime-source.tar \
  --manifest-output fa4-runtime-source.manifest.json

python examples/vllm_serve/collect_mask_reuse.py <CKPT> \
  --model-id <MODEL_ID> \
  --checkpoint-manifest-sha256 <SHA256> \
  --plan <PLAN> \
  --fa4-source /path/to/fa4-runtime-source \
  --fa4-source-manifest fa4-runtime-source.manifest.json \
  --fa4-source-manifest-sha256 <SHA256> \
  --fa4-commit <COMMIT> \
  --prompts-jsonl <PROMPTS.jsonl> \
  --vanilla-config <CKPT>/config.json \
  --target-sparsities 0.5 0.6 0.7 \
  --output <CAPTURES.jsonl>

Testing

• python -m pytest -q on the five source-manifest, capture, CLI, sparse-calibration, and vLLM-worker test modules: 55 passed.
• pre-commit run --files ...: passed.
• Commit signature verified.

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?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A — no copied code or new dependency.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ❌ — deferred until the complete stack is ready for review.
• Did you get Claude approval on this PR?: ❌ — pending before ready for review.

Additional Information

Stack 2 of 3. Base PR: #2235. The capture format binds the exact checkpoint, vanilla fit, FA4 source, geometry, prompt identity, and tensor-parallel evidence.

[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

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{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-2236/
Built to branch gh-pages at 2026-08-25 06:38 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 79.44814% with 216 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.01%. Comparing base (2229a8c) to head (8dcda41).

Files with missing lines	Patch %	Lines
...y/attention_sparsity/plugins/mask_reuse_capture.py	79.02%	138 Missing ⚠️
.../attention_sparsity/calibration/source_manifest.py	80.68%	68 Missing ⚠️
...ention_sparsity/plugins/vllm_mask_reuse_capture.py	72.22%	10 Missing ⚠️

Additional details and impacted files

@@                            Coverage Diff                            @@
##           kaix/mask-reuse-01-offline-calibration    #2236     +/-   ##
=========================================================================
  Coverage                                   79.00%   79.01%             
=========================================================================
  Files                                         527      530      +3     
  Lines                                       62652    63702   +1050     
=========================================================================
+ Hits                                        49499    50333    +834     
- Misses                                      13153    13369    +216     

Flag	Coverage Δ
unit	56.78% <79.44%> (+0.37%)	⬆️

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.