[4/n] Add joint mask reuse topology selection

[kaix-nv]

What does this PR do?

Type of change: new feature

Adds the third mask-reuse sub-PR: development-only joint discovery of anchor layers, consumer-to-anchor assignments, per-head donor maps, exact fallbacks, target sparsity, and the model-wide BMM1 skip objective. Held-out captures evaluate the frozen topology and cannot retune it.

This PR is stacked on #2236.

Usage

python examples/vllm_serve/calibrate_mask_reuse_topology.py \
  --checkpoint <CKPT> \
  --captures <TOPOLOGY_CAPTURES.jsonl> \
  --capture-manifest <TOPOLOGY_CAPTURES.manifest.json> \
  --vanilla-config <CKPT>/config.json \
  --prompt-plan <PROMPTS.jsonl> \
  --max-anchor-dropped-mass 0.02 \
  --reuse-dropped-mass-report-threshold 0.02 \
  --target-bmm1-skip-ratio 0.50 \
  --output <TOPOLOGY.json>

Testing

• Topology-selection module: 19 passed.
• Complete mask-reuse stack: 110 passed.
• Exact policy parity test verifies anchors [0, 2], nearest assignments, head maps {"1": [0, 0], "3": [0, 0]}, and empty fallback sets.
• The curated implementation, CLI, documentation, and feature-test tree matches kaix/mask-reuse-calib@3b0d268ea byte-for-byte.
• 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 3 of 3. Base PR: #2236. The selected artifact remains topology_candidate_only; promotion is deliberately outside this PR.

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

[codecov]

Codecov Report

❌ Patch coverage is 88.52006% with 83 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.12%. Comparing base (8dcda41) to head (0eba007).

Files with missing lines	Patch %	Lines
...ention_sparsity/calibration/mask_reuse_topology.py	88.50%	83 Missing ⚠️

Additional details and impacted files

@@                         Coverage Diff                         @@
##           kaix/mask-reuse-02-vllm-capture    #2237      +/-   ##
===================================================================
+ Coverage                            79.01%   79.12%   +0.10%     
===================================================================
  Files                                  530      531       +1     
  Lines                                63702    64425     +723     
===================================================================
+ Hits                                 50333    50974     +641     
- Misses                               13369    13451      +82     

Flag	Coverage Δ
unit	57.14% <88.52%> (+0.35%)	⬆️

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.