Publish Mamba state at aligned chunk endpoints

[lmcafee-nvidia]

What does this PR do?

This PR publishes reusable Mamba state whenever a cacheable prefill chunk ends exactly on a KV-block boundary, including when more prompt chunks remain. The previous direct-publication path covered only the final prompt boundary, while intermediate extraction cannot represent a state whose offset is the end of the current chunk. The PR also avoids staging Mamba state when a request has no usable prefix hashes, because such state can never be found by a later request.

Example failure scenario(s)

Scenario 1 (aligned boundary)

• request A has 519 prompt tokens, processed in 256-token chunks.
• its second chunk ends at token 512, with 7 tokens left.
• request B starts with the same first 512 tokens.
• before the fix, it finds saved attention state but no matching Mamba state.
• after the fix, it restores both and skips 512 tokens.

Scenario 2 (request opts out)

• the engine has prefix caching and Mamba storage enabled.
• request A has 512 prompt tokens but opts out of prefix caching.
• request A therefore computes 0 prefix-block hashes.
• before the fix, it still saves Mamba state at tokens 256 and 512.
• after the fix, it saves 0 states because no later request can find them.

How was this bug discovered?

A hybrid chunk test processed a non-final chunk ending exactly on a block boundary, then submitted a follower request with the same prefix. The follower matched the KV blocks but did not restore the saved Mamba state; the expected result was exact restoration at that reusable boundary. The oracle was the allocator's publication/restoration state together with the restored Mamba tensors. A separate cache-disabled check observed durable slot use even though no prefix hashes existed, where the expected slot count was unchanged.

Dependency and scope

Depends on NVIDIA#6418. This fork PR remains based on the corresponding prefix-cache-core-stress branch.

Relative to that base, the PR changes only:

• megatron/core/inference/contexts/dynamic_context.py
• megatron/core/inference/contexts/mamba_slot_allocator.py
• tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py (+85 unit-test lines)

It adds no functional tests, recipes, or unrelated production changes; comprehensive pairwise coverage remains in GPU unit tests.

Validation

• DFW H100 job 15519939 at exact SHA 3acfb17acc: all three focused tests passed on every one of 8 ranks (24 rank-local passes):
  • aligned non-final chunk endpoint publication and exact Mamba-state restoration
  • cache-disabled/no-hash requests do not publish or consume durable state
  • real hybrid/Mamba prefix-cache engine stress row under REF_ZERO pressure
• Black, isort, Ruff, and Python bytecode compilation on all three changed files
• git diff --check fork/prefix-cache-core-stress..HEAD

[lmcafee-nvidia]

Could @santhnm2 and @kvareddy please review this PR?