Correct prefix-caching ref-count accounting

[tdene]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your 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.

[kvareddy]

good find.

[kvareddy]

@lmcafee-nvidia, we missed this in our testing.

[sidsingh-nvidia]

Great find, thanks!

[sidsingh-nvidia]

/claude strict-review

[claude]

Strict review passed — no significant issues found. LGTM

Verified the fix addresses a real ref-count leak: self.block_ref_counts[blocks] -= 1 only decrements once per unique index under advanced indexing, so a prefix block shared by multiple requests finishing in the same batched release was under-decremented and leaked. scatter_add_ correctly accumulates per occurrence (the deterministic-accumulate op).

Also confirmed:

• The added torch.unique on both the REF_ZERO deregister path and the LRU pool-return path fixes a companion double-free (a duplicated block ID reaching ref 0 would otherwise be returned to the pool / deregistered twice).
• dtype/device are consistent: int32 self + int32 source + int64 index on CPU, matching scatter_add_ requirements; no new device constraint over the prior indexing.
• -1 sentinels are filtered by the caller before reaching the scatter, so no out-of-range index.

The new unit test is thorough — staged multi-owner release, in-batch duplication with a private block mixed in, and both eviction policies, including a set-uniqueness assertion that guards against double-returned IDs.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30383126962