[bo-nv]

@coderabbitai summary

Description

Depends on #21184, which depends on #20534; neither is merged yet. This branch is built on #21184's head, so until both merge this PR's diff also contains their commits. Please review only the last commit, 6c25942f3a. I will rebase as they land.

This is the second of the RI-10 PRs. #21184 lets every cache manager describe its pages once, as a page table built from a few hooks. This PR adds the algorithm that reads two ranks' page tables and says which bytes one sends the other, and checks it against what native transfers today. Nothing calls it in production yet; the next PR moves the native transfer onto it.

What it adds:

• kv_cache/layout_mapping.py:
  • map_buffers(src, dst, key=...): for a source and a destination rank's page layouts, the byte runs the source sends of every buffer both hold (same global layer and role). These are the indices both hold that the source is elected to send, as runs contiguous on both sides.
  • fan(src_size, dst_size, dst_rank): the source ranks a destination pairs with along one parallel dimension. A destination only ever hears from these.
  • Election: bytes several paired ranks hold alike go from candidates[key % n], key being the destination's attention data-parallel rank. This is the same owner native's fan-in election picks.
  • A destination holding an index no paired source holds is refused instead of left without it.
  • Also adds check_compatible, merge_runs, layer_group_of and group_kinds.
  • Nothing in the module knows a cache kind or a model: K/V heads, MLA latents, sectioned conv state, head-split SSM state and replicated side caches are all tensors whose axes split over tp/cp.
• tests/unittest/utils/page_layout_oracle.py: a byte oracle. Each rank's page holds, at every byte, the id of the global element there. After every source sends what map_buffers gives it, every byte each destination holds must have arrived exactly once, with the right element, and only from ranks it pairs with.

Tests:

• tests/unittest/_torch/executor/kv_cache/test_layout_mapping.py (CPU) covers fans, elections, cyclic and uneven splits, refusals and run merging, all against the oracle.
• tests/unittest/disaggregated/test_layout_mapping_matches_native.py (CPU) is a differential test. For the same managers it compares, byte for byte, what native's Sender write plan writes from every context rank into every generation rank with what map_buffers writes from the managers' page tables.
  • The managers are the real classes over a stand-in for the runtime storage, so both sides read the declarations those classes make. Covered: base K/V (MHA, GQA with duplicated heads, MLA), VSWA with separate and shared pool groups, an uneven PP split, MiniMax M3 (index-K separate and coalesced, HND and NHD), Mamba V2, Qwen4Exp with PLE state, FP4 MLA, CSA2 and DeepSeek-V4.
  • Topologies: context TP 1/2/4, PP 2, TP×PP and attention DP, against the same on the generation side plus Helix CP 2 receivers.
• The declarations are also checked against the oracle: test_layout_identity.py across TP sizes, and kv_transfer_harness.py before every GPU transfer test runs.

A gap in native that the differential test documents: native pairs a Helix receiver's recurrent state with the sender holding the same share of the tensor×context grid, but elects the sender of replicated recurrent state (Qwen4Exp's PLE) among tensor ranks alone. A receiver paired with an unelected sender therefore gets that state from no rank. map_buffers sends it. The test accepts exactly this difference, and only where native writes everything else identically. The native switch in the next PR removes the gap.

Follow-ups, in order:

1. Move the native transfer onto the page table and map_buffers, and remove the get_disagg_* declarations and native's per-kind mappers.
2. Move the per-request transfer rules (transfer length, buffers left to the receiver, first tokens, state slot, rebuild after import) into manager hooks.

Test Coverage

• tests/unittest/_torch/executor/kv_cache/test_layout_mapping.py (CPU, cpu_only)
• tests/unittest/disaggregated/test_layout_mapping_matches_native.py (CPU)
• tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py::test_declared_layouts_reshard_byte_for_byte (CPU)
• tests/unittest/disaggregated/kv_transfer_harness.py: the oracle check runs inside the existing GPU transfer tests (DeepSeek-V4, MiniMax M3)

Run locally on CPU only, with the GPU-free stub environment used for #21184:

• The differential test passes for every model and topology: 112 cases, each covering all 9 generation topologies. Mutating the election in map_buffers makes 27 of them fail.
• The kv_cache and disaggregated CPU suites match [TRTLLM-17449][feat] Let KVCacheManagerV2 build its page table from layout hooks #21184's results, apart from the tests this PR adds.

The GPU transfer tests have not been run yet.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.