[bo-nv]

@coderabbitai summary

Description

Depends on #20534, which is not merged yet. This branch is built on #20534's head (dbc8015261), so until #20534 merges this PR's diff also contains that commit. Please review only the last commit, 379248caaa. I will rebase onto main once #20534 lands.

This is the first of the RI-10 PRs whose goal is that a cache's layout is declared once, by its manager, and that the lender and every backend read that one description. Supporting a new cache structure then means overriding a few manager hooks.

What changes:

• KVCacheManagerV2.get_page_table() builds the page table (ManagerLayout, new module kv_cache/page_table.py). It reads what the runtime already holds: layer groups, pool groups and their pools, where each buffer lies in its slot, layer configs and the rank's place. It asks hooks only for what the runtime cannot tell, and a subclass overrides only the hooks whose answer differs:
  • _page_buffer_tensors: how a buffer's bytes spread over ranks, as tensors whose axes split over tp/cp. The base declares K/V and their block scales by KV head and index-K replicated. A role that moves and that no hook declares makes get_page_table() fail instead of being guessed.
  • _page_untransferred_roles, _page_rebuilds_untransferred, _virtual_layers, _page_layer_ids, _page_kv_heads, _page_dtype, _page_quant_mode.
• page_layout: BufferDesc.tensors replaces mapper and geometry. The new types are Tensor, Axis and held, with the builders head_major, token_major, sectioned, replicated and owned. LAYOUT_ID_VERSION goes from 3 to 4, so instances of different builds see different layout ids and share no pages.
• The lender's derive_layout returns the manager's table. Its own derivation from the get_disagg_* declarations and from private manager attributes is removed.
• Every manager declares its caches through the hooks: FP4 MLA, CSA2, DeepSeek-V4, MiniMax M3, Mamba V2 (conv state sectioned, SSM state by head), Qwen4Exp (PLE and QSA side planes replicated), KimiK3 and Glm5Next (KDA replay buffers not transferred).

What does not change: the native disaggregated transfer. It still reads the get_disagg_* declarations, which stay for now, and so do MapperKind and RoleLayout.

Follow-ups, in order:

1. Map pages between ranks of different layouts from their page tables, with a differential test against native's current mappers.
2. Move the native transfer onto the page table and that mapping, and remove the get_disagg_* declarations.
3. 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/sharing/test_layout_identity.py (CPU): page table construction, tensor and axis validation, layout ids (version 4 encoding), undeclared roles refused, roles that never move, layer naming, quantization.
• tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py: the base manager's default tensors, and refusal of an undeclared role.
• tests/unittest/_torch/executor/kv_cache/sharing/test_layout_matches_native_page_table.py (GPU): the manager's page table against native's page table across model and parallel layouts.
• tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py, tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py, tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py (GPU): the Mamba/KDA, DeepSeek-V4 and QSA declarations.

Run locally on CPU only:

• The kv_cache and disaggregated CPU suites match [TRTLLM-16186][feat] Lend KV cache pages and recurrent state in place #20534's results, apart from the layout tests this PR replaces.
• The Mamba V2 conv sections and SSM heads match the geometry native declares at TP 1, 2 and 4.

The GPU 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.

[Shixiaowei02]

Thanks for consolidating layout ownership in the manager. To keep this step easy to review, could we preserve the lender’s existing contract and behavior while changing the producer? Consider staging the Tensor/Axis migration, layout-ID version bump, and stricter role admission separately, or explicitly documenting their compatibility effects. Any compatibility projection should derive from one source of truth and cover only layouts it can represent losslessly. Differential checks against the previous implementation would help verify regions, transferred buffers, shards, object keys, and admission decisions for unchanged layouts, alongside existing lifecycle regressions and real-manager GPU checks.

[WeiHaocheng]

Thanks for this PR. I really like the approach of having each specialized manager declare its own page layout through a set of overridable hooks (_page_buffer_tensors, _page_untransferred_roles, _page_layer_ids, ...), with get_page_table() assembling the result. Before this, a manager like Fp4MlaKVCacheManagerV2 had no clear contract for describing what is in its slots. The lender reverse-engineered the layout in derive_layout, and that knowledge ended up split between the manager and the sharing code. With the hooks as the protocol:

• the manager is the single source of truth for its page contents;
• an undeclared transferred role fails loudly in get_page_table();
• about 560 lines of lender-side derivation go away.

That seems clearly right to me.

There are two things I'd like to discuss, both about how cross-rank layout is represented.

1. Should the cross-rank distribution live in the page layout at all?

The per-buffer tensors description (Axis / split / held / grid_place, head_major / token_major / sectioned / owned) describes how a buffer relates to the same tensor on other ranks. That information isn't needed for inference within a rank: attention kernels, block tables and prefix reuse only care about the local bytes in the slot. The lender also works at slot granularity. It lends whole slots keyed by layer group, plus the group's shard count and index, so it needs the shard identity but not per-axis ownership of each buffer.

Coupling the two has some costs:

• Every manager author has to reason about tp/cp placement (_page_head_split, token_major_=True, and so on), even when the manager is never used in disagg.
• The cross-rank description is folded into layout_id, so it becomes part of the reuse identity, even though it doesn't change which bytes sit in the slot.

Could the hooks describe only the single-rank page: layers, roles, offsets, sizes, dtype/quant, and whether each buffer is transferred or rebuilt? The relationship between ranks would then be abstracted later, in the transfer stage that actually consumes it.

2. Should the cross-rank relationship be fixed in the current data form?

The current vocabulary covers the common cases well. It is still a closed format, though:

• ownership must be a progression with a fixed start formula;
• strides must be affine;
• there is one dtype/quant per page and a global rebuilds_untransferred;
• only the tp/cp grid dimensions are supported.

Transfers that need conversion or re-layout don't fit, for example requantization, MLA expansion, swizzled layouts, or layouts that depend on runtime state. Each new case would mean extending the format and bumping LAYOUT_ID_VERSION.

In #20207 I took a different approach: a per-model TransferPlanner code interface (check_peer / targets / pair, plus describe() for the peer handshake). Both sides run the same planner code, so a model with an unusual layout can implement its own pairing without changing a shared schema. A default planner can still cover the common head-split cases.

These seem complementary. The manager hooks from this PR would publish the single-rank page facts. Planners would consume them, e.g. comparing layout_id plus dtype/quant in check_peer and checking transferred before copying whole slots, instead of hard-coding those assumptions. That keeps one source of truth for what's in a slot, while the cross-rank relationship stays flexible.

Happy to discuss. I may be missing a reason the lender or another consumer needs the per-axis ownership at the manager level.