Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ layouts or ownership models must not override the implementation.
- `blockRadixTree.*`: the shared prefix-reuse tree and SHA-256 block keys.
- `page.*`, `kvCache.*`, and `kvCacheManager.*`: page lifecycle, per-request
cache state, and the top-level manager.
- `batch.*`: stable request rows, dirty tracking, and raw GPU metadata publication.
- `storage/`, `storageManager.*`, `evictionController.*`, and `copyEngine.*`:
pools, eviction ownership, migration, and data movement.
- `lifeCycleRegistry.*`: layer-group/lifecycle mapping, including attention and
Expand Down Expand Up @@ -79,6 +80,14 @@ slots, schedules pages for eviction, migrates pages between levels, and resizes
pools. `CopyEngine` performs the actual batched transfers; C++ code calls it
directly and must not round-trip through Python bindings.

`Batch` groups requests from one manager across all layer groups. Request changes
invalidate their stable rows; `publish()` uploads final raw indices and eligible
history counts, including post-rollback state. Device addresses remain fixed.
Publish and wait for readiness outside graph capture; after submitting readers
or replaying a graph, call `recordRead()` before mutating requests or membership.
Publication waits for offload and prior readers, and retains each staging buffer
until its upload completes. DLPack views keep the allocation alive, not KV pages.

The dependency direction is broadly:

```text
Expand Down Expand Up @@ -146,6 +155,11 @@ KvCache
|- KvCacheManager (shared; cache keeps manager alive)
`- per-beam/per-block page holders and locks

Batch
|- KvCacheManager (shared)
|- request rows (non-owning; exclusive membership, driven by one owning thread)
`- fixed device metadata and event-protected staging buffers

BlockRadixTree
`- roots -> child Blocks (strong ownership through next maps)
`- lifecycle page entries (raw observer links)
Expand All @@ -159,6 +173,10 @@ Eviction controller
strong ones merely to simplify access.
- `KvCache` keeps its `KvCacheManager` alive. The manager's registry of living
caches must not create the reverse strong-reference cycle.
- Closing a request removes its `Batch` row; closing a batch detaches its live
requests without closing them. Removed rows stay dirty until publication clears
them. Batch destruction waits for uploads and recorded readers before freeing
memory; exported arrays can extend the allocation lifetime beyond `close()`.
- A committed page is referenced by the radix tree without making the tree its
permanent owner. Eviction queues may be the only strong owner of a droppable
page, so never store a raw pointer past the operation that obtained it.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|ARM64")
endif()

set(KV_CACHE_MANAGER_V2_SRCS
kv_cache_manager_v2/batch.cpp
kv_cache_manager_v2/common.cpp
kv_cache_manager_v2/config.cpp
kv_cache_manager_v2/lifeCycleRegistry.cpp
Expand Down
Loading
Loading