[isytwu]

Motivation

#38778 stores the MLA/DSA KV of rank-replicated pool groups once in the UMBP direct linker. The FULL + Mamba group that #44977 added for Kimi K3 cannot be rank-replicated, since its Mamba state is TP-sharded, so every attention TP rank still writes its own copy of the MLA KV (#44977's first known limitation): at TP8 the store holds eight copies of each page and starts evicting shared prefixes once it fills. On Kimi K3 at TP8 with DSpark (the InferenceMINI agentic TP8 c14 recipe) and 480 GiB of UMBP DRAM, the same as HiCache's host pools, this PR against its base on the same node raised the prefix hit rate from 90.7% to 94.2% and throughput by 14%, and cut the mean TTFT from 3.2 s to 1.9 s (P90 from 7.9 s to 3.1 s).

Modifications

• A DevicePoolEntry can declare the shard it holds: ranks that declare the same shard hold identical objects. The FULL + Mamba group's MLA KV declares shard 0; the Mamba state and the DSpark draft stay per rank.
• The UMBP direct linker keys a pool that declares a shard with the shard in place of the TP rank, so the store keeps one copy. Every rank still writes it, and mori acknowledges a put of a key that is already stored or being stored (fix(umbp): stop false put-failure noise (sampled allocation WARN; in-flight key counts as stored) ROCm/mori#777).

Tests

• test_linker_pool_assembler.py: the shards of the FULL + Mamba group.
• test_umbp_direct_linker_rank_keys.py: the linker's object keys put the MLA KV's shard in place of the TP rank and keep the Mamba and draft keys per rank.

Accuracy Tests

Kimi-Linear-48B-A3B (MLA + KDA) at TP4 on MI355X with a local UMBP standalone server, so the four ranks restore one shared copy of the MLA KV. The KL test from dc291b5, modeled on test_unified_cache_linker_kl_inkling.py, offloads turn 1, drops the device tree with /flush_cache, restores turn 2 from UMBP alone (asserts cached_tokens_details.host == cached_tokens) and replays it against recomputed logprobs right after the restore boundary, since further out the MLA layers mask a wrong KDA state. It is not kept in CI: ROCm is not bit-exact, so its threshold (1e-2) only catches a grossly wrong restore. 4 samples:

Run	avg KL
This PR, 8 runs	1.4e-3, 1.8e-3, 9.5e-4, 8.7e-4, 1.2e-3, 1.4e-3, 7.1e-4, 1.5e-3
Negative control: KDA state not restored (slot zeroed)	4.8e-2, 6.2e-2 (fails)

Kimi K3 at TP8 on 8x MI355X:

• Teacher-forced restore probe (logprobs over a prefix restored from UMBP after /flush_cache, against the device tree's run-to-run noise): passes.
• GSM8K with the few-shot harness of test_kimi_k3_eval_mi35x.py, run twice with /flush_cache in between so that the second pass resumes from UMBP: 0.956, then 0.954.

Related: #38778, #44977, #44929, #41343. Supersedes #39511.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests.
• Provide accuracy evaluation results.

---

CI States

Latest PR Test (Base): ❌ Run #39567129814
Latest PR Test (Extra): ❌ Run #39567129825
Latest PR Test (AMD ROCm 10): ❌ Run #39567129836