[Fix] Vacuous marker writes in the cache tests, and an undebited Mamba admission slot

[alphabetc1]

Two independent pre-existing defects, both found while adding Mamba support to
HiCache buffer-only mode (#36345, stacked on this). They are unrelated to each
other; the two commits are separately reviewable.

---

1. fix: unified radix cache tests write markers into the pool, not a copy

Two helpers in the unified radix cache unit suite write their marker into a
copy of the pool, not the pool:

k_buf[indices].fill_(marker)          # advanced indexing -> new tensor
mamba_cache.temporal[:, mamba_indices].fill_(marker)

buf[indices] with a tensor index is advanced indexing, which returns a fresh
tensor; .fill_() mutates that temporary and the pool keeps its zeros.

Every downstream assertion of the form "the loaded bytes equal the producer's"
is therefore comparing zeros to zeros — test_hicache_load_back_restores_data,
test_hicache_l3_prefetch_roundtrip, test_buffer_only_read_path_roundtrip
and the SWA data checks all pass unconditionally. A load-back that restored the
wrong slot, or restored nothing at all, would not turn any of them red.

Reproduced standalone:

>>> t = torch.zeros(3, 8); t[:, torch.tensor([1, 2])].fill_(5); t.sum()
tensor(0.)

Fixed by assigning through __setitem__, which is an in-place index_put_.

---

2. fix: charge the Mamba slot a prefill admission reserved but never debited

PrefillAdder.add_one_req reads _mamba_gap_budget_for_req(req) twice, and the
two reads straddle a call that changes its answer.

total_tokens += self._mamba_gap_budget_for_req(req)     # gate
...
if req.needs_host_load_back():
    self.tree_cache.init_load_back(...)                 # binds req.mamba_pool_idx
...
self._update_prefill_budget(..., mamba_gap_reserve=self._mamba_gap_budget_for_req(req))

_mamba_gap_budget_for_req returns 0 once req.mamba_pool_idx is not None, and
MambaComponent.prepare_load_back sets exactly that field during
init_load_back. So on any request that takes a host load-back, the gate
reserves a slot's worth of shared-gap bytes and the debit that follows charges 0
— rem_total_tokens, cur_rem_tokens and rem_mamba_slots all stay
un-decremented for a slot the request is now holding.

The next request in the same batch is then admitted against a budget that
believes that slot is still free. On the shared Mamba pool the residual is what
alloc_req_slots' fail-loud RuntimeError backstops.

Fixed by reading it once, before init_load_back, and reusing the value at both
debit sites. No behaviour change for requests that already own a slot, or off
the shared Mamba pool where _mamba_slot_cost is 0.

---

Verification

(1) H200, -k 'buffer or hicache' over the whole unified-radix-cache config
matrix: 418 passed, 403 skipped, 0 failed — identical to main, so the
now-live byte comparisons hold as they stand. The point is that they can fail
from here on.

(2) is inert unless the allocator is the shared
UnifiedMambaTokenToKVPoolAllocator composite, which is only built under
--enable-unified-memory (kv_cache_configurator.py:390). Everywhere else
_mamba_slot_cost is 0, so both reads return 0 and the value passed to
_update_prefill_budget is unchanged. The runs above — and the Inkling classes
in test_unified_radix_cache_kl_hybrid_bitexact.py, which do not set that flag
— therefore show the change is harmless, not that it is exercised. I have
not run a unified-memory hybrid-Mamba configuration; the argument for the change
is the read-after-mutation above, which is visible in the source. Happy to add a
targeted PrefillAdder unit case if reviewers want a guard.

---

CI States

Latest PR Test (Base): ✅ Run #33832677017
Latest PR Test (Extra): ❌ Run #33832676907
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33832677009

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c12432c8b9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[alphabetc1]

/rerun-group radix_cache/unified_radix_tree

[github-actions]

Results for /rerun-group radix_cache/unified_radix_tree:

🚀 4-gpu-h100 (4 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py

🚀 8-gpu-h200 (3 tests): ❌ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4_pp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_nightly.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py

[alphabetc1]

/rerun-group radix_cache/unified_radix_tree

[github-actions]

Results for /rerun-group radix_cache/unified_radix_tree:

🚀 4-gpu-h100 (4 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py

🚀 8-gpu-h200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4_pp.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_glm52.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_swa.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py