[TRTLLM-12891][feat] Include KV connector prefixes in V2 scheduler budgeting

docs/source/features/kv-cache-connector.md

@@ -33,9 +33,18 @@ These methods run on the leader process and drive the connector's behavior.
   * **Returns**: An arbitrary metadata object (picklable) that describes the tasks for the workers. This object is broadcasted to all workers.
 
 * **`get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> tuple[int, bool]`**
-  * **Description**: Called when a new request arrives. It checks to see if any KV cache can be loaded from an external KV store.
+  * **Description**: Queries external KV after the compute batch is selected on the legacy path. Connectors that implement the complete reservation protocol use `reserve_prefix` during admission when prefix-aware scheduling and KV cache manager V2 are enabled.
   * **Returns**: A tuple `(num_tokens, is_async)`. `num_tokens` is the number of tokens found in the external cache. `is_async` indicates if the loading will happen asynchronously (background) or requires blocking.
 
+* **`reserve_prefix(self, request: LlmRequest, num_computed_tokens: int, reservation_id: int) -> tuple[int, bool]`**
+  * **Description**: Reserves an additional contiguous prefix beginning at `num_computed_tokens` during batch construction. A positive answer protects that source range against mutation and eviction. The query must start no KV transmission or destination writes.
+  * **Returns**: `(additional_tokens, is_async)`. The runtime may accept a shorter, block-aligned range and keeps at least the final prompt token for local computation. `is_async=True` parks an accepted request until all workers finish its load.
+  * **Identity**: The runtime assigns a fresh `reservation_id` to each attempt. Keep reservation state by this identity, including when the same request is retried.
+
+* **`release_prefix_reservation(self, request: LlmRequest, reservation_id: int, start: int, end: int) -> None`**
+  * **Description**: Releases the named reservation's protection for the absolute half-open token interval `[start, end)`. The runtime releases rejected or clipped portions before transmission and the accepted portion after all workers report completion. Overlapping reservations must keep their own protection.
+  * **Lifetime**: This callback never asks the connector to interrupt an active transfer. Client cancellation after dispatch drains the load before releasing its source and destination resources.
+
 * **`request_finished(self, request: LlmRequest, cache_block_ids: list[int]) -> bool`**
   * **Description**: Called when a request completes generation.
   * **Returns**: A boolean indicating if an asynchronous save operation is underway. If `True`, the system waits for the operation to complete before releasing the KV cache blocks.
@@ -175,13 +184,13 @@ KV pool; the scheduler's exhaustion error says so directly when a connector is a
 
 The connector holds page indices across iterations, and `RequestData` reports only the pages appended
 since the last call. Anything that hands a slot the connector already knows about to a different
-request therefore goes unreported and corrupts the next transfer against it. Three configurations do
-that, and each is rejected at bring-up.
+request therefore requires resetting the connector state before replay. The runtime applies the
+following configuration limits at bring-up.
 
 | Configuration | Mechanism |
 |---|---|
 | Speculative decoding | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. |
-| A capacity scheduler policy other than `GUARANTEED_NO_EVICT` | A destroyed-and-replayed request comes back on different pages, and the connector's per-request block delta is then measured against pages that were freed with it. |
+| A capacity scheduler policy other than `GUARANTEED_NO_EVICT` with KV cache manager V1 | V1 does not reset the connector block delta when a request is destroyed and replayed. V2 resets that state and permits replay; an active connector load retains its allocation until completion. |
 | A host or disk cache tier | Tier eviction reassigns the GPU slot. See [KV cache tiers are GPU-only under a connector](#kv-cache-tiers-are-gpu-only-under-a-connector). |
 
 The exact set the runtime refuses depends on your cache configuration; the bring-up error is
@@ -238,16 +247,60 @@ Variable sliding-window attention is the case where that stops working, because
 
 A model whose layers all share one sliding window stays a single layer group, so the flat callbacks still apply and such a connector is not refused. What differs is that the callbacks cover the **live window only**. Blocks the window has passed report `-1` (`BAD_PAGE_INDEX`) in place — the list stays aligned to block ordinals, so entry `i` still describes tokens `[i * tokens_per_block, (i+1) * tokens_per_block)`, but the up-front blocks carry no page and are not available to load into or save from. Filter with `valid_page_slots`, described in [Blocks with no page](#blocks-with-no-page); without it a `-1` resolves to the last page slot of the pool. A warning naming the window size is logged at start-up when a flat-only connector is attached to such a model.
 
-`get_num_new_matched_tokens` is asked once the batch for the upcoming forward pass is final. A request that is asked is therefore a request that runs, and the connector can take ownership of remote blocks in the query and release it in `request_finished`.
+##### Reserving a prefix during admission
+
+For connector implementers, scheduler participation requires three optional methods together:
+`reserve_prefix` and `release_prefix_reservation` on the scheduler, and
+`get_finished_prefix_loads` on the worker. Bring-up rejects a partial implementation. KV cache
+manager V2 enables this protocol when `SchedulerConfig.enable_prefix_aware_scheduling=True`.
+Connectors using the existing methods retain the final-batch query path.
+
+The scheduler charges compute tokens after the local and reserved external prefix. The complete
+prefix still needs destination KV pages, so a token-budget hit can remain limited by KV capacity.
+After the final admission checks, `SchedulerOutput.prefix_loads` authorizes the exact loads:
+
+| `PrefixLoad` field | Meaning |
+|---|---|
+| `reservation_id` | The identity supplied to `reserve_prefix`, also returned on completion. |
+| `request_id` | The request owning the allocation. |
+| `start`, `end` | Absolute half-open token interval to load. |
+| `is_async` | Whether the request waits outside the compute batch. |
+| `block_ids_by_layer_group` | Complete destination page lists, indexed by group and block ordinal. Filter entries through `valid_page_slots`. |
+| `tokens`, `cache_salt` | Request tokens and cache-key isolation salt. |
+
+Build transfers from this collection. A confirmed async load appears here even when it is absent
+from `new_requests` and `cached_requests`, and even when the compute batch is empty. Those two
+lists continue to describe computation and its token/block deltas. The reservation itself supplies
+no permission to write destination memory.
+
+| Event | Connector and runtime behavior |
+|---|---|
+| Candidate queried | Connector protects the promised source; no transmission starts. |
+| Candidate rejected or offer shortened | Runtime releases the unused source interval. A later attempt gets a new reservation identity. |
+| Accepted load dispatched | Connector starts only the confirmed interval; runtime retains its destination allocation. |
+| Client cancels during transmission | Runtime keeps the allocation while the connector completes the load. |
+| Every worker completes | Runtime releases the source reservation, resumes a live request, or finalizes a cancelled request and frees its allocation. |
+
+Each worker reports completion after its destination writes and source reads finish. Reports are
+sent to the leader without a collective. Once every worker has reported, an ordered control item
+retires the load on all ranks before scheduling. This releases the source reservation and allows
+an async request to rejoin the compute batch, or a cancelled request to release its allocation.
+The runtime continues polling when no forward pass is scheduled. A transfer that never completes
+retains its resources; elapsed time alone cannot make pages safe to reuse.
 
-Two things are worth knowing when tuning a deployment.
+For a load admitted with computation, `wait_for_layer_load` establishes the dependency on that
+worker's transfer stream. Inference does not wait for the retirement control item.
 
-* **The runtime may honour less than you offer.** With chunked prefill the cache is allocated per context chunk, which is what bounds its memory, so an offer reaching past the current chunk requires the runtime to grow the allocation and that can fail under pressure. The runtime then serves the part it can cover and computes the rest locally. The amount actually served is what `RequestData.computed_position` reflects; the unserved remainder needs no action from the connector beyond its usual `request_finished` cleanup.
-* **The query is not part of the scheduler's budget.** The scheduler sizes a request's chunk as if the connector will serve nothing, so a served prefix reduces the work in the forward pass but does not free budget for another request in the same iteration.
+##### Legacy final-batch queries
 
-Specify `enable_block_reuse=True` alongside the connector for any of this to run; see [Block reuse alongside the connector](#block-reuse-alongside-the-connector).
+`get_num_new_matched_tokens` is called at most once per KV allocation after batch selection.
+Its offer can reduce computation for that request, but does not free token budget for another
+request in the same iteration. With chunked prefill, allocation growth can limit the served prefix;
+`RequestData.computed_position` reflects the load interval that the runtime honors.
 
-`get_num_new_matched_tokens` is called **at most once per KV allocation**. This is the precise form of the "once per request" rule: if a request's KV cache is destroyed and the request is replayed -- which `MAX_UTILIZATION` does under memory pressure -- the replay asks again, because the pages the first answer described are gone.
+Destroying an allocation clears its connector state, so replay queries again. Specify
+`enable_block_reuse=True` alongside the connector; see
+[Block reuse alongside the connector](#block-reuse-alongside-the-connector).
 
 **Deployment note.** Under a connector, a workload that was token-bound becomes KV-bound: the connector removes forward-pass tokens but its prefix still occupies GPU pages. Lowering `max_num_tokens` to hand memory back to the KV pool is usually the right adjustment, the opposite of the guidance for a connector-free deployment.
 
@@ -280,32 +333,38 @@ These methods run on all workers (GPU processes) and interact with the actual GP
   * **Description**: Polled by the runtime to check the status of asynchronous operations.
   * **Returns**: Two lists of request IDs: those that have finished saving, and those that have finished loading.
 
+* **`get_finished_prefix_loads(self) -> list[int]`**
+  * **Description**: Reports locally completed reservation identities for loads dispatched through `SchedulerOutput.prefix_loads`, including synchronous loads. Completion means this worker has finished all reads and writes and established the required CUDA stream visibility. This method must not wait for other workers. The runtime collects reports asynchronously and distributes retirement decisions before scheduling; parked async requests resume and shared resources are released only after every worker has reported.
+  * **Compatibility**: Legacy request-ID load and save completions continue through `get_finished`.
+
 ## Example Implementation
 
 The file `examples/llm-api/llm_kv_cache_connector.py` provides a reference implementation of a **Persistent KV Cache**.
 
 ### Overview
 
 This example implements a file-system based KV cache.
-1. **Save**: When a request finishes or needs to be swapped out, its KV blocks are saved to disk as `.pt` files.
+1. **Save**: After prefill, complete computed KV blocks are saved to disk as immutable `.pt` files.
 2. **Load**: When a new request arrives with the same prompt prefix, the connector identifies the cached files and loads them back into GPU memory, skipping re-computation.
 
 ### Implementation Details
 
-* **Metadata**: The example defines a `PersistentKvCacheConnectorMetadata` dataclass containing lists of `(file_path, block_id)` tuples for both loading and saving. This simple structure allows the Scheduler to tell the Worker exactly which file corresponds to which GPU block index.
+* **Metadata**: `PersistentKvCacheConnectorMetadata` carries `(file_path, block_id)` load/save targets and accepted reservation identities. The worker returns those identities after the copies finish.
+
+* **Source protection**: Saves publish complete files atomically and never overwrite an existing inode. `reserve_prefix` creates private hard links to these immutable files without reading KV data. Releasing a range removes only its reservation links, so another reservation keeps its source available. External cache management must preserve this immutability: remove or replace a path rather than writing into a published file.
 
-* **Hashing Strategy**: The `PersistentKvCacheConnectorLeader` hashes the token sequence of a block to generate a unique filename (e.g., `hash_value.pt`). This acts as the lookup key.
+* **Hashing Strategy**: `PersistentKvCacheConnectorLeader` hashes the entire token prefix through each block, together with `cache_salt`, using SHA-256. This distinguishes identical blocks reached through different preceding tokens and remains stable across Python processes. Use separate cache directories for different models and KV layouts.
 
 * **Worker Logic**:
-  * `start_load_kv`: Iterates through the load list provided in the metadata, loads the `.pt` file to CPU, and copies it to the specific `block_id` in the GPU tensor.
+  * `start_load_kv`: Reads only accepted load targets, copies their `.pt` data to GPU with blocking copies, and records their reservation identities for `get_finished_prefix_loads`.
   * `wait_for_save`: Performs the reverse. It copies data from the GPU `block_id` to CPU and saves it to disk using `torch.save`.
 
 ### Limitations & Patterns
 
 This example illustrates the API mechanics but has several limitations that make it unsuitable for high-performance production use without modification:
 
 1. **Blocking I/O**: The example uses `torch.load` and `torch.save` synchronously. In a real implementation, these should be offloaded to a background thread or asynchronous I/O handler to avoid stalling the GPU.
-2. **Simplified Block Matching**: The `get_num_new_matched_tokens` implementation in the example only matches full blocks. It does not handle partial cache hits.
+2. **Simplified Block Matching**: The example matches and loads complete blocks, requires one layer group, and saves only the first context chunk. It does not demonstrate chunked-prefill persistence.
 3. **FileSystem Latency**: Storing one file per block can create high filesystem overhead.
 
 ### Usage