[TRTLLM-16186][feat] Define backend contract for KV cache transfers by Shixiaowei02 · Pull Request #20117 · NVIDIA/TensorRT-LLM

Shixiaowei02 · 2026-10-12T02:34:42Z

Description

This pull request refactors the KV cache transfer backend by separating and clarifying the contract between the backend and the native implementation. It introduces a new extent.py module for native transfer types, moves MemoryKind to its own module, and updates imports throughout the codebase and tests to use the new structure. The changes improve modularity and make the separation between backend contract and native implementation clearer.

Backend contract and native separation:

Created native/extent.py to define native-specific transfer types (Chunk, CacheExtent, Delivered, Failed, Cancelled, Attempt, Outcome) and removed these from the backend contract (base/__init__.py). Imports in native modules and tests are updated to use native.extent for these types. [1] [2]
Updated the backend contract (base/__init__.py) to only export types that are part of the contract, and clarified the module docstring to reflect the new structure.

Memory kind modularization:

Moved MemoryKind to a new base/memory.py module, with its own docstring and enum definition, and updated the backend contract to import it from there. [1] [2]

Import and usage updates across codebase:

Updated all native and test modules to import Chunk, CacheExtent, and related types from native.extent instead of the backend contract, ensuring clear separation of contract and implementation. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14]

Test and contract validation updates:

Updated tests to use the new import paths and revised a test to check that the publish adapter returns an Attempt instance instead of Publishes, aligning with the new contract.

These changes collectively improve code organization, clarify the backend-native interface, and ensure consistency across the codebase.

chienchunhung

Thanks for the PR! Overall LGTM; will take another pass once the related docs are aligned.

chienchunhung

Notes:

View metadata: Keep original group and row metadata stable until result consumption, including after early lease release. Keep this lifetime separate from payload protection.
Hint ownership: State whether the backend copies required hint fields or the caller keeps them stable until use ends.

nv-xtf · 2026-10-14T07:28:50Z

From the Store Backend integration perspective, the current contract supports our example backend. Overall LGTM.

Shixiaowei02 · 2026-10-15T09:51:16Z

/bot run --disable-fail-fast

Shixiaowei02 · 2026-10-15T09:51:45Z

/bot kill

tensorrt-cicd · 2026-10-15T09:59:51Z

PR_Github #78104 [ run ] triggered by Bot. Commit: 660e76e Link to invocation

tensorrt-cicd · 2026-10-15T10:00:11Z

PR_Github #78106 [ kill ] triggered by Bot. Commit: 660e76e Link to invocation

tensorrt-cicd · 2026-10-15T10:00:13Z

PR_Github #78104 [ run ] completed with state ABORTED. Commit: 660e76e

Link to invocation

tensorrt-cicd · 2026-10-19T01:26:39Z

PR_Github #78739 [ run ] triggered by Bot. Commit: 5a67a46 Link to invocation

tensorrt-cicd · 2026-10-19T05:50:11Z

PR_Github #78739 [ run ] completed with state SUCCESS. Commit: 5a67a46
/LLM/main/L0_MergeRequest_PR pipeline #64876 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Define, in tensorrt_llm/_torch/disaggregation/base, the contract between a lease holder, which holds leases from the KV cache lender, and a backend that moves the lent rows to or from a peer or a store. base exports Backend, Capabilities, Attempt, Outcome, View, LayerGroupRows, SlotRegion, MemoryKind, Hint, WakeFn, InvalidSubmission and SubmissionRejected, plus TokenRange and CacheKind, which are not part of the contract. Backend replaces Fetches and Publishes. - MemoryKind, in base/memory.py, is the kind of memory a region lies in: HOST is CPU-addressable and DEVICE is the GPU holding the cache. A backend registers only the regions whose SlotRegion.memory_kind it can access, and rows in a region it has not registered are rows it does not recognize. A row given by an address occupies one slot. - A view's groups and their arrays are an immutable snapshot that stays valid after the lease is returned; the lease bounds only the payload's protection. A hint describes only the other end and is an immutable snapshot; which rows this end reads or writes comes from the view alone. - fetch and publish state what the lease holder guarantees about the memory the view locates until the submission's quiesced() is true, and no two rows of a fetch locate overlapping memory. A raise other than the two rejections may have let access to the memory escape, and the lease holder then keeps the lease until shutdown() returns an empty set. quiesced() covers the whole submission, so the leases lent for a backend that may leave a row not proven ended are kept small. - Conclusion and quiescence do not depend on calls to poll() or quiesced(): a backend advances its own work. A backend calls WakeFn only after every change it reports is visible: an outcome or an access end through poll() or quiesced(), a change an optional protocol names where that protocol says. - An instance serves the submissions of one lender's domain, and a process may hold several. A backend's methods and its attempts' belong to one call sequence. - A backend given rows it does not recognize rejects the submission with InvalidSubmission and never skips them. Delivery masks are bool arrays; anything else raises TypeError. Native's per-request types, Chunk, CacheExtent, Delivered, Failed, Cancelled and its own Attempt and Outcome, move unchanged from base/backend.py to native/extent.py; native and transceiver.py change only their imports, and base/transfer.py imports Chunk from native/extent.py for type checking only. The contract tests import nothing of the lender, and they pin that base imports native only for type checking. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com>

Delivery masks are of dtype bool; anything else raises TypeError. Outcome keeps a read-only bool mask whose type is exactly numpy.ndarray as it is, and its giver writes neither it nor memory it shares afterwards; any other mask is copied. A read-only subclass of numpy.ndarray, such as a masked array, is copied into a plain array, since all() on a masked array skips masked rows. The contract tests add a kept read-only mask beside a copied writable one, a read-only masked array, and read-only int and float masks. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com>

A backend can then mark the masks it is given read-only and hand them to Outcome, which keeps a read-only bool mask without copying it. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com>

Shixiaowei02 · 2026-10-20T07:35:15Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-20T07:43:44Z

PR_Github #79201 [ run ] triggered by Bot. Commit: 8ac96cf Link to invocation

tensorrt-cicd · 2026-10-20T14:09:03Z

PR_Github #79201 [ run ] completed with state SUCCESS. Commit: 8ac96cf
/LLM/main/L0_MergeRequest_PR pipeline #65236 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

… the blob backend The CPU and real-pool tests drive the backend with host slots of their own; none shows that what the staging lender lends is what the backend moves. test_blob_backend_round_trip.py does it end to end on a GPU: one real KV cache manager publishes a request's blocks through its lender's staging slots, another fetches them into a request of its own, and a test lease holder drives the lender as the runtime will (lend, wait for the view, submit it as it stands, wait for quiescence, mark what arrived, release). - Pages come back byte for byte and resume_bounds covers the prompt, over the in-memory store and, with TRTLLM_MOONCAKE_REAL=1, a real Mooncake pool. - A fetch whose later blocks fault marks only the delivered prefix, so resume_bounds stops at the first gap. The in-memory store and driver come from NVIDIA#20217's testing package. The test builds each backend itself, as a factory will: it opens the driver, takes the lender's hold, registers the regions with RegionBinding, starts BlobBackend, and gives the hold back once the backend shut down. The lender's MemoryRegion names no memory kind, which NVIDIA#20117's SlotRegion requires, so the test registers host-memory SlotRegions made from the lender's regions, as a runtime must until the lender's regions name their memory kind. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>

Shixiaowei02 · 2026-10-21T02:05:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-21T02:13:17Z

PR_Github #79469 [ run ] triggered by Bot. Commit: 8ac96cf Link to invocation

nv-xtf

LGTM

tensorrt-cicd · 2026-10-21T08:02:06Z

PR_Github #79469 [ run ] completed with state SUCCESS. Commit: 8ac96cf
/LLM/main/L0_MergeRequest_PR pipeline #65447 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Shixiaowei02 · 2026-10-21T08:02:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-21T08:09:31Z

PR_Github #79589 [ run ] triggered by Bot. Commit: 8ac96cf Link to invocation

tensorrt-cicd · 2026-10-21T11:47:47Z

PR_Github #79589 [ run ] completed with state SUCCESS. Commit: 8ac96cf
/LLM/main/L0_MergeRequest_PR pipeline #65534 completed with status: 'SUCCESS'

CI Report

Link to invocation

… the blob backend The CPU and real-pool tests drive the backend with host slots of their own; none shows that what the staging lender lends is what the backend moves. test_blob_backend_round_trip.py does it end to end on a GPU: one real KV cache manager publishes a request's blocks through its lender's staging slots, another fetches them into a request of its own, and a test lease holder drives the lender as the runtime will (lend, wait for the view, submit it as it stands, wait for quiescence, mark what arrived, release). - Pages come back byte for byte and resume_bounds covers the prompt's three whole blocks, over the in-memory store and, with TRTLLM_MOONCAKE_REAL=1, a real Mooncake pool. - A fetch whose later blocks fault marks only the delivered prefix, so resume_bounds stops at the first gap. The in-memory store and driver come from NVIDIA#20217's testing package. The test builds each backend itself, as a factory will: it opens the driver, takes the lender's hold, registers the regions with RegionBinding, starts BlobBackend, and gives the hold back once the backend shut down. The lender's MemoryRegion names no memory kind, which NVIDIA#20117's SlotRegion requires, so the test registers host-memory SlotRegions made from the lender's regions, as a runtime must until the lender's regions name their memory kind. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>

github-actions Bot assigned Shixiaowei02 Oct 12, 2026

chienchunhung reviewed Oct 12, 2026

View reviewed changes

Shixiaowei02 mentioned this pull request Oct 13, 2026

[None][feat] Model-specific disaggregated KV transfer planners #20207

Draft

1 task done

chienchunhung reviewed Oct 14, 2026

View reviewed changes

This was referenced Oct 14, 2026

[None][feat] Add shared KV runtime assembly #20266

Draft

[None][feat] Add shared KV transfer ownership and progress #20268

Draft

Shixiaowei02 mentioned this pull request Oct 14, 2026

[None][feat] Add the blob store driver boundary and an in-memory test driver #20217

Merged

1 task done

eopXD mentioned this pull request Oct 15, 2026

[None][feat] Add native KV cache retention to KVCM V2 #20086

Merged

6 tasks done

Shixiaowei02 commented Oct 15, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/disaggregation/base/backend.py

Shixiaowei02 commented Oct 15, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/disaggregation/base/backend.py

Comment thread tensorrt_llm/_torch/disaggregation/base/backend.py

Shixiaowei02 marked this pull request as ready for review October 15, 2026 09:56

Shixiaowei02 requested review from a team as code owners October 15, 2026 09:56

Shixiaowei02 requested review from SimengLiu-nv, Tabrizian, YihuiLu512, bo-nv, chienchunhung, chuangz0, lori-ren, nv-xtf, reasonsolo and thorjohnsen October 15, 2026 09:56

Shixiaowei02 added the ci: full pre-merge approved label Oct 19, 2026

Shixiaowei02 force-pushed the dev/xiaoweis/kv-transfer-backend-contract branch from 5a67a46 to 2d05d32 Compare October 20, 2026 07:32

Shixiaowei02 added 3 commits October 20, 2026 15:35

[None][fix] Promise a new mask from every View.row_masks call

8ac96cf

A backend can then mark the masks it is given read-only and hand them to Outcome, which keeps a read-only bool mask without copying it. Signed-off-by: Xiaowei Shi <39303645+Shixiaowei02@users.noreply.github.com>

Shixiaowei02 force-pushed the dev/xiaoweis/kv-transfer-backend-contract branch from 2d05d32 to 8ac96cf Compare October 20, 2026 07:35

chuangz0 reviewed Oct 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/disaggregation/base/backend.py

nv-xtf approved these changes Oct 21, 2026

View reviewed changes

chuangz0 approved these changes Oct 21, 2026

View reviewed changes

Shixiaowei02 enabled auto-merge (squash) October 21, 2026 09:34

Shixiaowei02 merged commit a9000fe into NVIDIA:main Oct 21, 2026
18 checks passed

Shixiaowei02 deleted the dev/xiaoweis/kv-transfer-backend-contract branch October 21, 2026 12:12

chuangz0 mentioned this pull request Oct 22, 2026

[None][chore] Remove the Python KV bounce buffer; kv_cache_bounce_size_mb selects the C++ agent bounce #20049

Open

1 task done

Shixiaowei02 mentioned this pull request Oct 22, 2026

[#18151][feat] Admit trusted KV hints and carry immutable request metadata #20269

Open

1 task done

nv-xtf mentioned this pull request Oct 23, 2026

[None][feat] Add the blob backend with batch registration, bounded queueing and a CPU end-to-end example #20884

Merged

1 task done

Conversation

Shixiaowei02 commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

chienchunhung left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chienchunhung left a comment

Choose a reason for hiding this comment

Uh oh!

nv-xtf commented Oct 14, 2026

Uh oh!

Uh oh!

Shixiaowei02 commented Oct 15, 2026

Uh oh!

Shixiaowei02 commented Oct 15, 2026

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

Shixiaowei02 commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

Shixiaowei02 commented Oct 21, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

Uh oh!

nv-xtf left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

Shixiaowei02 commented Oct 21, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Shixiaowei02 commented Oct 12, 2026 •

edited

Loading

chienchunhung left a comment •

edited

Loading