[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.