From 3bc36a5565adbb54b413aa023049f6b48f055087 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 16:01:23 -0700 Subject: [PATCH 01/23] [None][feat] add unified KV cache layout foundation Signed-off-by: Allison Lim --- ...ache-v2-target-draft-unification-design.md | 383 ++++++++++++++ .../kv_cache/kv_cache_manager_v2.py | 500 +++++++++++++++--- .../_torch/pyexecutor/kv_cache/layout.py | 266 ++++++++++ .../executor/kv_cache/test_kv_cache_layout.py | 223 ++++++++ 4 files changed, 1308 insertions(+), 64 deletions(-) create mode 100644 docs/source/features/kv-cache-v2-target-draft-unification-design.md create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py diff --git a/docs/source/features/kv-cache-v2-target-draft-unification-design.md b/docs/source/features/kv-cache-v2-target-draft-unification-design.md new file mode 100644 index 000000000000..7b89a4f37b7d --- /dev/null +++ b/docs/source/features/kv-cache-v2-target-draft-unification-design.md @@ -0,0 +1,383 @@ + + + +# Target and draft KV cache unification in KV Cache Manager V2 + +> **Audit basis:** upstream `main` at `dc8451c50864edc07c75dedeafcb91dbb7a5d478` +> (fetched October 12, 2026). + +## Executive verdict + +The central framing is correct: + +> A unified manager owns one committed request timeline. Layer-specific storage +> geometry may differ, but request state, allocation, reuse depth, and admission +> cannot. + +The original report is also correct that reuse policy is already a per-layer +configuration. The audit has four material corrections and one remaining +reuse-policy wiring requirement: + +1. **Per-layer reuse policy already exists in the V2 C++ and Python contracts.** + `AttentionReusePolicy` defines `REQUIRED`, `PRIVATE`, and `OPTIONAL`; + `AttentionLayerConfig.reusePolicy` defaults to `REQUIRED`. The remaining + manager change is to choose and pass the policy explicitly when building + target and draft layer configurations. +2. **`is_draft` should not simply become a per-layer flag.** Some behavior + disappears under unification, some becomes request-wide, and only some + becomes layer metadata. +3. **Per-layer heads, dimensions, and windows are supported by allocation, but + not by all speculative-decoding mutation kernels.** +4. **Mixed dtype requires more than a `get_layer_dtype()` accessor.** Storage is + already byte-oriented, but much of the Python/runtime layer assumes one dtype + and one cache type. +5. **Capacity accounting cannot remain unchanged.** Lifecycle registry behavior + can remain, but estimation, budget splitting, pool sizing, and speculative + headroom must be recomputed for the combined layout. + +The existing limited one-manager speculative path, which appends speculative +layers after the target layers, is a useful starting point. It currently +inherits target-wide assumptions for dtype, layout, and several runtime +operations ([`_util.py`](../../../tensorrt_llm/_torch/pyexecutor/_util.py)). + +## Proposed model + +```mermaid +flowchart TB + subgraph Current["Current separate-manager path"] + R1["Request"] + TM["Target V2 manager
quota + radix tree + request cache"] + DM["Draft V2 manager
quota + radix tree + mirrored request cache"] + TL["Target layers"] + DL["Draft layers"] + + R1 --> TM --> TL + R1 --> DM --> DL + end + + subgraph Unified["Proposed unified path"] + R2["Request"] + UM["One V2 manager
one quota + one radix tree"] + KC["One KvCache
committed history + capacity + beam/state"] + LD["Per-layer descriptors
owner, dtype, geometry, roles,
window, sparse, placement"] + SM["Storage manager"] + P1["Pool/lifecycle group A"] + P2["Pool/lifecycle group B"] + P3["Pool/lifecycle group C"] + + R2 --> UM --> KC + LD --> UM + UM --> SM + SM --> P1 + SM --> P2 + SM --> P3 + end +``` + +A useful refinement is to distinguish: + +- **Committed timeline:** one history length and reuse depth for the request. +- **Speculative scratch tail:** one conservatively sized allocation extent, + within which different layer groups may temporarily use different + token-position mappings. + +That distinction permits one `KvCache` without requiring target and draft +temporary data layouts to be identical. + +## Audit of the original report + +| Claim | Verdict | Audit | +|---|---|---| +| One `KvCache` per request | **Correct** | This is the fundamental unification invariant. Capacity, committed history, beam state, suspension, and request lifetime become shared. | +| One logical block grid and radix tree | **Correct, with nuance** | `tokens_per_block` is manager-wide. Individual buffers may use a `tokensPerBlockOverride`, but it must divide the logical block size ([`config.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h)). | +| One reuse depth | **Correct, with policy semantics** | The request has one claimed reuse endpoint. `REQUIRED` lifecycles constrain that endpoint, `OPTIONAL` lifecycles are checked for complete coverage there without shortening it, and `PRIVATE` lifecycles do not participate ([`kvCache.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp)). | +| One quota, eviction policy, and resume gate | **Correct** | These are manager/request policies and cannot differ between target and draft in a single manager. | +| One sharding scheme | **Correct for the initial implementation** | The current manager has one `Mapping`. Target and draft rank ownership must be representable by that mapping. Helix and other specialized layouts need explicit compatibility checks. | +| One manager specialization | **Mostly correct** | It may be one existing specialization or a composite implementation that supports the union of required hooks. “Same existing class” need not be a permanent restriction. | +| Per-layer heads and head dimension | **Partially correct** | Construction accepts per-layer values, but accepted-token relocation currently requires uniform heads and one local pool ([`kv_cache_manager_v2.py`](../../../tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py)). | +| Per-layer window and sink tokens | **Partially correct** | C++ lifecycle identity supports both window and sink tokens. The generic Python builder currently supplies a per-layer window but always passes `num_sink_tokens=None`. | +| Per-layer buffers and byte sizes | **Correct** | C++ groups buffers by lifecycle and expanded byte size; storage does not require identical layer geometry ([`storage/config.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp)). | +| Sparse and dense layers | **Correct** | `isSparse` is lifecycle identity, and sparse/dense storage pools are separated even when slot sizes match. | +| Per-layer addressing | **Correct, with wording adjustment** | Layers share the request row and committed timeline, but different layers may select different pool groups through `kv_cache_pool_mapping`. They do not necessarily use the same physical block-offset array. | +| Per-layer dtype needs wiring | **Correct, but understated** | C++ storage is byte-oriented, but Python buffer views, cost calculations, attention backends, quantization scale roles, relocation, cold codecs, and disaggregation commonly read a manager-wide dtype. | +| Per-layer reuse policy already exists | **Correct** | `AttentionReusePolicy` and `AttentionLayerConfig.reusePolicy` are present in C++, included in lifecycle identity and pool grouping, and exposed as `reuse_policy` in Python ([`config.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h), [`lifeCycleRegistry.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h), [`__init__.pyi`](../../../tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi)). `BlockReusePolicy` remains a separate manager-wide request-scope policy. | +| Stats, connectors, and events are all skipped for draft | **Partially incorrect** | Request stats and connectors have `is_draft` gates. Event-manager construction itself is not categorically disabled for draft managers. | +| Dynamic-tree headroom becomes request-wide | **Correct** | The unified cache must reserve enough scratch capacity for the larger target/draft requirement. The over-reservation is not necessarily small; it depends on `topK`, tree width, draft length, and layer costs. | +| Compatibility reduces to specialization, sharding, and block size | **Incomplete** | These are necessary but not sufficient. Placement, buffer roles/layout, reuse identity, relocation capability, and transfer/codec support also matter. | +| Registry invariant and capacity accounting stay unchanged | **Half correct** | Lifecycle registry semantics can remain. Capacity and quota accounting must change because two independently sized managers become one composite manager. | + +## Recommended changes + +### 1. Introduce one canonical per-layer cache descriptor + +Create an internal descriptor such as: + +```python +@dataclass(frozen=True) +class LayerCacheSpec: + global_layer_id: int + owner: LayerOwner # TARGET or DRAFT + layer_type: LayerType # ATTENTION or SSM + dtype: DataType + num_kv_heads: int + head_dim: int + buffer_specs: tuple[BufferSpec, ...] + sliding_window_size: int | None + num_sink_tokens: int | None + is_sparse: bool + reuse_policy: AttentionReusePolicy + placement: LayerPlacement +``` + +The exact names are flexible. Manager construction should consume a single +target-plus-draft layout instead of deriving most properties from the target +configuration and appending draft layers afterward. + +`draft_layer_ids` should be derived from `owner`, rather than serving as the +only source of layer identity. + +**Justification:** This prevents dtype, buffer roles, cache type, placement, and +geometry from developing separate parallel lists with inconsistent indexing. + +### 2. Make dtype genuinely per layer + +Add accessors such as: + +```python +get_layer_dtype(global_layer_id) +get_layer_bytes_per_token(global_layer_id, role) +get_layer_buffer_roles(global_layer_id) +``` + +Then update: + +- `BufferConfig` construction. +- Tensor views returned by `get_buffers()`. +- KV byte and capacity calculations. +- NVFP4/FP8 scale-buffer handling. +- Attention backend dtype lookup. +- Cold-page codecs. +- Disaggregated transfer descriptors. +- Accepted-token relocation. + +Do **not** add dtype to `AttnLifeCycle`. Lifecycle identity describes retention +and reuse behavior; storage already separates incompatible byte sizes. + +The existing manager-wide `.dtype` can temporarily remain valid only for +homogeneous layouts. Mixed layouts should fail loudly until every consumer has +moved to a layer-aware accessor. + +**Justification:** Otherwise allocation may be correct while kernels interpret +the bytes using the wrong element type. + +### 3. Replace manager-wide `kv_cache_type` assumptions with per-layer roles + +The wrapper currently derives a manager-wide `kv_factor` from `SELFKONLY` +versus ordinary K/V storage. C++ already understands arbitrary per-layer +`BufferConfig` roles. + +Move toward descriptors such as: + +```text +Target layer 0: K + V +Target MLA layer: K only +Draft layer 48: K + V +FP4 layer: K + V + K_SCALE + V_SCALE +``` + +**Justification:** Dtype unification alone is insufficient if target and draft +use different physical roles or quantization metadata. + +### 4. Split `is_draft` by responsibility + +Do not replace manager-level `is_draft` with a layer-level boolean everywhere. +Use the following disposition: + +| Existing `is_draft` responsibility | Unified disposition | +|---|---| +| Separate request mirror and draft admission | **Eliminate** in the unified path | +| Separate request index and release behavior | **Eliminate** | +| Draft-layer identification | `LayerCacheSpec.owner` | +| Dynamic-tree generation reserve | Request-wide `SpeculativeCapacityPolicy` | +| Accepted-token relocation | Explicit target/draft layer sets or homogeneous layer groups | +| Block publication | One radix tree, with publication governed by each lifecycle's existing reuse policy | +| Reuse backoff | One conservative request-wide backoff | +| Stats | Record once; optionally aggregate by layer owner | +| Events | Emit one unified event stream; owner tags are optional metadata | +| Connectors/disaggregation | Transfer the complete required layer set by default | +| Cold-page codec | Select codec/dtype by layer or buffer | +| SWA scratch reuse | Capability decision per lifecycle; initially disable for mixed target/draft managers | +| Helix and special layouts | Compatibility rejection until their ownership model supports the combined layout | + +**Justification:** Most existing draft branches exist because there are two +managers. Preserving those branches as per-layer behavior would retain +unnecessary complexity. + +### 5. Wire the existing reuse policy and define the draft default + +The current V2 core already implements the required policy semantics: + +- `REQUIRED` lifecycles constrain the common prefix match. +- `OPTIONAL` lifecycles do not shorten the `REQUIRED` match. Their coverage is + checked at that endpoint, and the caller selects complete optional groups on + the first resume. Missing or unselected groups receive fresh pages for + recomputation. +- `PRIVATE` lifecycles remain request-local and are neither published to nor + reused from the radix tree. + +The policy participates in `AttnLifeCycle` identity, and storage keeps distinct +reuse policies in separate pool groups. The request-side implementation is in +[`kvCache.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp), +while the Python contract exposes `reuse_policy` in +[`__init__.pyi`](../../../tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi). + +The manager-first work therefore needs to: + +1. Add `reuse_policy` to the composite layer descriptor. +2. Set `AttentionLayerConfig.reuse_policy` explicitly in the Python config + builder instead of relying on its `REQUIRED` default. +3. Map returned layer-group status back to target/draft ownership. +4. Have the scheduler select available `OPTIONAL` groups on first resume. +5. Ensure missing or unselected draft groups are recomputed before a draft + forward consumes that prefix. + +The appropriate draft policy is an algorithm contract: + +- **`REQUIRED`** is the safest initial default. Draft coverage limits the shared + reuse endpoint, so no catch-up path is needed. This matches the original + report's recommendation when draft KV must be reused with target KV. +- **`OPTIONAL`** allows target KV to determine the reuse endpoint and reuses + draft KV only when complete there. Use it only when the algorithm can + explicitly reconstruct a missing draft prefix from target state. +- **`PRIVATE`** always reconstructs the draft prefix and is useful only when + publishing draft KV is undesirable or unsupported. + +Also ensure the radix reuse identity includes every target and draft setting +that can change generated KV, including model or adapter identity where +applicable. + +### 6. Rework capacity and budget accounting + +For an eligible unified configuration: + +- Do not split `max_gpu_total_bytes` between target and draft. +- Do not split host or disk-tier budgets. +- Build one composite storage configuration containing all layers. +- Run pool-ratio selection over the combined lifecycle and buffer groups. +- Compute per-request fixed and variable cost from the composite descriptor. +- Use the larger speculative scratch requirement across target and draft. +- Apply `max_util_for_resume`, eviction, and quota retuning once. + +The current V2 path explicitly splits draft budgets when a separate manager is +selected ([`_util.py`](../../../tensorrt_llm/_torch/pyexecutor/_util.py)). That +logic must be bypassed for a unified plan. + +**Justification:** Adding the independently calculated target and draft +capacities does not necessarily yield the capacity of one shared request +timeline. + +### 7. Add layer-group-aware relocation primitives + +Current accepted-token relocation assumes: + +- One attention window. +- Uniform KV heads. +- One local KV pool. +- One dtype or element size. +- One broad layer range. + +Those assumptions are enforced directly in +[`_update_kv_cache_draft_token_location()`](../../../tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py). + +Introduce a manager-level operation conceptually like: + +```python +relocate_speculative_tokens( + requests, + layer_ids, + source_positions, + destination_positions, +) +``` + +The manager should partition `layer_ids` into compatible execution groups, for +example by: + +```text +(pool, buffer role, dtype, head geometry, window/layout) +``` + +Each group can invoke an appropriate relocation kernel. Once every required +group has been relocated, the shared `KvCache` can be rewound or advanced once. + +**Justification:** This is the main place where “storage supports heterogeneous +layers” currently fails to imply “speculative decoding supports heterogeneous +layers.” + +### 8. Define an explicit compatibility check + +Return a structured result rather than a boolean: + +```python +UnifiedKvCacheCompatibility( + supported=True, + reasons=(), +) +``` + +A target/draft pair is compatible only when: + +- They can share one committed request timeline and beam width. +- They use the same logical `tokens_per_block`. +- Their layer placement and TP/PP/CP ownership fit one mapping. +- One manager implementation can express the union of buffer roles and hooks. +- All dtype/layout consumers are layer-aware. +- The speculative relocation implementation supports the resulting pool groups. +- Reuse hashing and backoff are valid for both models. +- Connectors, offload codecs, and disaggregation can transfer the required union. +- Specialized features such as Helix, DSA, Mamba state, or model-specific + buffers explicitly declare support. + +In the eventual design, dtype, heads, dimensions, and windows should not +independently reject unification. Until their consumers are migrated, they +remain implementation capability gates. + +## Suggested implementation order + +1. Add `LayerCacheSpec` and composite layout construction without changing + runtime behavior. +2. Centralize all per-layer geometry and byte calculations behind the + descriptor. +3. Implement per-layer dtype and buffer-role accessors. +4. Wire `AttentionLayerConfig.reuse_policy` from the composite descriptor and + add scheduler selection/recomputation handling for `OPTIONAL` draft groups. +5. Add the structured compatibility checker and retain the separate-manager + fallback. +6. Build one V2 C++ storage configuration from target plus draft layers. +7. Replace `is_draft` with request-wide policy, layer ownership, or deletion + according to the table above. +8. Rework composite capacity estimation and stop splitting budgets. +9. Add grouped relocation APIs. +10. Enable unification for the simplest speculative algorithm first. +11. Migrate additional algorithms one by one, expanding the compatibility + checker as each gains support. + +For the initial rollout, use these conservative defaults: + +- Default target and draft lifecycles to `REQUIRED`; enable `OPTIONAL` only for + algorithms with a validated draft-prefix recomputation path. +- Use the maximum required reuse backoff. +- Reserve the maximum speculative scratch width request-wide. +- Disable SWA scratch reuse for mixed managers until validated. +- Reject unsupported CP/Helix/specialized layouts. +- Preserve the existing separate-manager path as the fallback. + +## Revised contract + +> A unified V2 KV-cache manager owns one committed timeline, one request cache, +> one logical block grid, one radix-tree reuse depth, and one memory quota for +> the target and draft layers. Each layer independently describes its buffer +> roles, dtype, geometry, retention lifecycle, sparsity, and physical pool +> mapping. Speculative tails may use layer-specific temporary layouts inside a +> shared request-wide allocation, but committing, rewinding, suspending, +> resuming, and reusing the request occur once and at a common history position. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 8dfd00defdc6..d74901fc27fe 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -131,6 +131,13 @@ request_context, ) from ..scheduler import ScheduledRequests +from .layout import ( + KVCacheLayoutPlan, + LayerBufferSpec, + LayerCacheOwner, + LayerCacheSpec, + LayerCacheTransferPolicy, +) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -270,6 +277,55 @@ class BlockReusePolicy(StrEnum): PER_CONVERSATION = "per_conversation" +class KVCacheLayerView: + """Non-owning layer-scoped view of one :class:`KVCacheManagerV2`. + + The view can be passed to attention metadata while request lifecycle and + transfer remain owned by the underlying manager. It intentionally makes + ``free_resources`` a no-op because legacy speculative call sites may free + a draft manager independently after a forward pass. + """ + + def __init__(self, manager: "KVCacheManagerV2", owner: LayerCacheOwner) -> None: + layout_plan = manager.kv_cache_layout_plan + if layout_plan is None: + raise ValueError("layer views require a KV cache layout plan") + layers = tuple(layer for layer in layout_plan.layers if layer.owner is owner) + if not layers: + raise ValueError(f"KV cache layout contains no {owner.value} layers on this rank") + self._manager = manager + self.owner = owner + self.layer_ids = tuple(layer.global_layer_id for layer in layers) + self.local_layer_ids = tuple(layer.local_layer_id for layer in layers) + + @property + def owner_manager(self) -> "KVCacheManagerV2": + return self._manager + + def get_layer_spec(self, global_layer_id: int) -> LayerCacheSpec: + if global_layer_id not in self.layer_ids: + raise ValueError( + f"layer {global_layer_id} does not belong to the {self.owner.value} cache view" + ) + return self._manager.get_layer_cache_spec(global_layer_id) + + def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + self.get_layer_spec(layer_idx) + return self._manager.get_buffers(layer_idx, kv_layout=kv_layout) + + def free_resources(self, request: LlmRequest) -> None: + """Leave request lifecycle to the owning manager.""" + + def report_batch_to_connector(self, *args, **kwargs) -> None: + raise RuntimeError("KV transfer must be initiated through the owning unified manager") + + def __getattr__(self, name): + manager = self.__dict__.get("_manager") + if manager is None: + raise AttributeError(name) + return getattr(manager, name) + + def _swa_endpoint_priority( prompt_length: int, tokens_per_block: int, rewind_tokens: int ) -> Callable[[int, object], int]: @@ -1290,10 +1346,20 @@ def __init__( cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, + layout_plan: Optional[KVCacheLayoutPlan] = None, **kwargs, ) -> None: self.mapping = mapping self.dtype = dtype + self._declared_target_layer_count = num_layers + self._explicit_layout_plan = layout_plan + self.kv_cache_layout_plan: Optional[KVCacheLayoutPlan] = layout_plan + self._layer_views: dict[LayerCacheOwner, KVCacheLayerView] = {} + if layout_plan is not None and layout_plan.tokens_per_block != tokens_per_block: + raise ValueError( + "layout-plan tokens_per_block does not match the manager: " + f"{layout_plan.tokens_per_block} != {tokens_per_block}" + ) self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager @@ -1308,6 +1374,36 @@ def __init__( layer_mask=layer_mask, ) self.is_draft = is_draft + self._legacy_draft_layer_ids = { + global_layer_id + for global_layer_id in self.pp_layers + if is_draft + or ( + spec_config is not None + and layer_mask is None + and global_layer_id >= self._declared_target_layer_count + ) + } + if layout_plan is not None: + planned_global_layers = [layer.global_layer_id for layer in layout_plan.layers] + if planned_global_layers != self.pp_layers: + raise ValueError( + "layout-plan global layers do not match pipeline placement: " + f"{planned_global_layers} != {self.pp_layers}" + ) + if is_draft and any( + layer.owner is not LayerCacheOwner.DRAFT for layer in layout_plan.layers + ): + raise ValueError("a draft mirror manager cannot own target layout layers") + if layout_plan.has_mixed_dtype: + raise NotImplementedError( + "KVCacheManagerV2 does not yet interpret mixed-dtype layout plans" + ) + if layout_plan.uniform_dtype != dtype: + raise ValueError( + "layout-plan dtype does not match the manager dtype: " + f"{layout_plan.uniform_dtype} != {dtype}" + ) # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the @@ -1449,10 +1545,14 @@ def __init__( self.pp_layers, ) assert len(self.max_attention_window_vec) == self.num_local_layers + if layout_plan is not None: + self.max_attention_window_vec = [ + layer.sliding_window_size for layer in layout_plan.layers + ] # A D-token match backoff needs D more physical SWA history at the # committed endpoint so commit_min_snapshot still covers the rewind. - if kv_cache_config.enable_block_reuse and self.reuse_match_backoff: + if layout_plan is None and kv_cache_config.enable_block_reuse and self.reuse_match_backoff: self.max_attention_window_vec = _extend_swa_windows_for_reuse( self.max_attention_window_vec, self.reuse_match_backoff, @@ -1561,6 +1661,28 @@ def append_to_kv_heads_per_layer( f"unique values={set(self.head_dim_per_layer)}" ) + if layout_plan is not None: + planned_heads = [layer.num_kv_heads for layer in layout_plan.layers] + if planned_heads != self.num_kv_heads_per_layer: + raise ValueError( + "layout-plan KV-head counts do not match the manager: " + f"{planned_heads} != {self.num_kv_heads_per_layer}" + ) + planned_head_dims = [layer.head_dim for layer in layout_plan.layers] + if planned_head_dims != self.head_dim_per_layer: + raise ValueError( + "layout-plan head dimensions do not match the manager: " + f"{planned_head_dims} != {self.head_dim_per_layer}" + ) + planned_kv_factors = { + 2 if Role.VALUE in layer.buffer_roles else 1 for layer in layout_plan.layers + } + if planned_kv_factors != {self.kv_factor}: + raise NotImplementedError( + "KVCacheManagerV2 does not yet interpret mixed or mismatched K/V layouts: " + f"plan factors={sorted(planned_kv_factors)}, manager factor={self.kv_factor}" + ) + self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) max_util_for_resume = kv_cache_config.max_util_for_resume @@ -1701,6 +1823,8 @@ def append_to_kv_heads_per_layer( ) config = self._build_cache_config(config) config = self._remove_zero_size_buffers(config) + self._finalize_layout_plan(config) + self._validate_transfer_layout() has_host_cache_tier = any( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) @@ -1831,6 +1955,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: layer_id: self.impl.get_layer_group_id(layer_id) for layer_id in typed_range(LayerId(num_layers)) } + self.required_transfer_layer_group_ids = self._get_required_transfer_layer_group_ids() self.kv_cache_map: dict[int, _KVCache] = {} self._disagg_receive_ready: dict[int, torch.cuda.Event] = {} @@ -2481,6 +2606,9 @@ def _kv_pool_mapping_offset( ) def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: + explicit_layout_plan = getattr(self, "_explicit_layout_plan", None) + if explicit_layout_plan is not None: + return explicit_layout_plan.cache_size_layer_components() layer_sizes = [] attention_windows = [] for local_layer_idx in range(self.num_local_layers): @@ -2994,43 +3122,49 @@ def _build_base_config( ) ) - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional - # per-layer BufferConfig entries — for example a sparse index-K - # buffer — without overriding the K/V/NVFP4 scale wiring above. - # The dict maps local layer id -> list of extra BufferConfig. Each - # extra buffer's role must be unique within the layer (asserted by - # AttentionLayerConfig.__post_init__) and its size must be in bytes - # per block (= bytes_per_token * tokens_per_block). - extra_buffers_per_layer = ( - self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} - ) - - layer_configs: List[AttentionLayerConfig] = [] - for layer_id in typed_range(LayerId(self.num_local_layers)): - buffer_type = self._get_buffer_roles_for_layer(layer_id) - buffers = [ - BufferConfig( - role=role, - size=self.get_layer_bytes_per_token(local_layer_idx=layer_id, data_role=role) - * tokens_per_block, - ) - for role in buffer_type - ] - for extra in extra_buffers_per_layer.get(int(layer_id), ()): - assert extra.role not in buffer_type, ( - f"extra buffer role {extra.role!r} for layer " - f"{int(layer_id)} duplicates a standard K/V/scale role" - ) - buffers.append(extra) - layer_configs.append( - AttentionLayerConfig( - layer_id=layer_id, - buffers=buffers, - sliding_window_size=self.max_attention_window_vec[layer_id], - num_sink_tokens=None, - ) + explicit_layout_plan = getattr(self, "_explicit_layout_plan", None) + if explicit_layout_plan is not None: + layer_configs = explicit_layout_plan.to_attention_layer_configs() + else: + # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional + # per-layer BufferConfig entries — for example a sparse index-K + # buffer — without overriding the K/V/NVFP4 scale wiring above. + # The dict maps local layer id -> list of extra BufferConfig. Each + # extra buffer's role must be unique within the layer (asserted by + # AttentionLayerConfig.__post_init__) and its size must be in bytes + # per block (= bytes_per_token * tokens_per_block). + extra_buffers_per_layer = ( + self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} ) + layer_configs: List[AttentionLayerConfig] = [] + for layer_id in typed_range(LayerId(self.num_local_layers)): + buffer_type = self._get_buffer_roles_for_layer(layer_id) + buffers = [ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token( + local_layer_idx=layer_id, data_role=role + ) + * tokens_per_block, + ) + for role in buffer_type + ] + for extra in extra_buffers_per_layer.get(int(layer_id), ()): + assert extra.role not in buffer_type, ( + f"extra buffer role {extra.role!r} for layer " + f"{int(layer_id)} duplicates a standard K/V/scale role" + ) + buffers.append(extra) + layer_configs.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=buffers, + sliding_window_size=self.max_attention_window_vec[layer_id], + num_sink_tokens=None, + ) + ) + return KVCacheManagerConfigPy( # Used by the backend only for token<->block arithmetic and # radix hashing; BufferConfig.size above stays the physical page. @@ -3071,6 +3205,157 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config + def _legacy_layer_owner(self, global_layer_id: int) -> LayerCacheOwner: + if global_layer_id in self._legacy_draft_layer_ids: + return LayerCacheOwner.DRAFT + return LayerCacheOwner.TARGET + + def _layout_plan_from_runtime_config( + self, config: KVCacheManagerConfigPy + ) -> Optional[KVCacheLayoutPlan]: + if not all(isinstance(layer, AttentionLayerConfig) for layer in config.layers): + return None + if any(heads <= 0 for heads in self.num_kv_heads_per_layer) or any( + head_dim <= 0 for head_dim in self.head_dim_per_layer + ): + return None + ignored_roles = self.get_disagg_ignored_roles() + layers = [] + for local_layer_id, layer in enumerate(config.layers): + assert isinstance(layer, AttentionLayerConfig) + global_layer_id = self.pp_layers[local_layer_id] + buffers = tuple( + LayerBufferSpec( + role=buffer.role, + size=buffer.size, + tokens_per_block_override=buffer.tokens_per_block_override, + is_sparse=buffer.is_sparse, + transfer_policy=( + LayerCacheTransferPolicy.LOCAL_ONLY + if buffer.role in ignored_roles + else LayerCacheTransferPolicy.REQUIRED + ), + ) + for buffer in layer.buffers + ) + layers.append( + LayerCacheSpec( + local_layer_id=local_layer_id, + global_layer_id=global_layer_id, + owner=self._legacy_layer_owner(global_layer_id), + dtype=self.dtype, + num_kv_heads=self.num_kv_heads_per_layer[local_layer_id], + head_dim=self.head_dim_per_layer[local_layer_id], + buffers=buffers, + sliding_window_size=layer.sliding_window_size, + num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, + ) + ) + return KVCacheLayoutPlan(tokens_per_block=self.tokens_per_block, layers=tuple(layers)) + + @staticmethod + def _buffer_config_signature(buffer: BufferConfig) -> tuple[DataRole, int, int | None, bool]: + return ( + buffer.role, + buffer.size, + buffer.tokens_per_block_override, + buffer.is_sparse, + ) + + def _validate_runtime_config_matches_layout_plan( + self, config: KVCacheManagerConfigPy, layout_plan: KVCacheLayoutPlan + ) -> None: + if len(config.layers) != len(layout_plan.layers): + raise ValueError( + "runtime cache layer count does not match the layout plan: " + f"{len(config.layers)} != {len(layout_plan.layers)}" + ) + for planned, runtime in zip(layout_plan.layers, config.layers): + if not isinstance(runtime, AttentionLayerConfig): + raise ValueError( + "an attention layout plan cannot produce a non-attention runtime layer" + ) + planned_buffers = tuple( + self._buffer_config_signature(buffer.to_buffer_config()) + for buffer in planned.buffers + ) + runtime_buffers = tuple( + self._buffer_config_signature(buffer) for buffer in runtime.buffers + ) + if planned_buffers != runtime_buffers: + raise ValueError( + f"runtime buffers for layer {planned.global_layer_id} do not match " + f"the layout plan: {runtime_buffers} != {planned_buffers}" + ) + runtime_lifecycle = ( + runtime.sliding_window_size, + runtime.num_sink_tokens, + runtime.reuse_policy, + ) + planned_lifecycle = ( + planned.sliding_window_size, + planned.num_sink_tokens, + planned.reuse_policy, + ) + if runtime_lifecycle != planned_lifecycle: + raise ValueError( + f"runtime lifecycle for layer {planned.global_layer_id} does not match " + f"the layout plan: {runtime_lifecycle} != {planned_lifecycle}" + ) + + def _finalize_layout_plan(self, config: KVCacheManagerConfigPy) -> None: + if self._explicit_layout_plan is None: + self.kv_cache_layout_plan = self._layout_plan_from_runtime_config(config) + return + self._validate_runtime_config_matches_layout_plan(config, self._explicit_layout_plan) + self.kv_cache_layout_plan = self._explicit_layout_plan + + def _validate_transfer_layout(self) -> None: + layout_plan = self.kv_cache_layout_plan + if layout_plan is None or (not self.is_disagg and self.kv_connector_manager is None): + return + errors = list(layout_plan.transfer_compatibility_errors()) + ignored_roles = self.get_disagg_ignored_roles() + for layer in layout_plan.layers: + local_only_roles = [ + buffer.role + for buffer in layer.buffers + if buffer.transfer_policy is LayerCacheTransferPolicy.LOCAL_ONLY + and buffer.role not in ignored_roles + ] + if local_only_roles: + errors.append( + f"layer {layer.global_layer_id} marks roles {local_only_roles} local-only, " + "but the manager does not exclude them from disaggregation" + ) + for layer in layout_plan.required_transfer_layers(): + ignored_required_roles = [ + buffer.role + for buffer in layer.buffers + if buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED + and buffer.role in ignored_roles + ] + if ignored_required_roles: + errors.append( + f"layer {layer.global_layer_id} requires transfer of ignored roles " + f"{ignored_required_roles}" + ) + if errors: + raise NotImplementedError( + "the unified KV layout is incompatible with the current transfer metadata: " + + "; ".join(errors) + ) + + def _get_required_transfer_layer_group_ids(self) -> frozenset[int]: + layout_plan = self.kv_cache_layout_plan + if layout_plan is None: + return frozenset() + return frozenset( + int(self.impl.get_layer_group_id(LayerId(layer.local_layer_id))) + for layer in layout_plan.required_transfer_layers() + ) + def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Exclude empty buffers before creating the runtime storage pools.""" if config.layers and all( @@ -3095,6 +3380,7 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa buffers=buffers, sliding_window_size=layer.sliding_window_size, num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, ) ) else: @@ -3191,6 +3477,68 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + def get_layer_cache_spec(self, global_layer_id: int) -> LayerCacheSpec: + """Return the immutable layout entry for one rank-local layer.""" + if self.kv_cache_layout_plan is None: + raise RuntimeError("this specialized manager has no attention-only layout plan") + return self.kv_cache_layout_plan.get_layer(global_layer_id) + + def get_layer_owner(self, global_layer_id: int) -> LayerCacheOwner: + return self.get_layer_cache_spec(global_layer_id).owner + + def get_layer_dtype(self, global_layer_id: int) -> DataType: + return self.get_layer_cache_spec(global_layer_id).dtype + + def get_layer_buffer_roles(self, global_layer_id: int) -> tuple[DataRole, ...]: + return self.get_layer_cache_spec(global_layer_id).buffer_roles + + def get_layer_transfer_policy(self, global_layer_id: int) -> LayerCacheTransferPolicy: + return self.get_layer_cache_spec(global_layer_id).transfer_policy + + @property + def target_layer_ids(self) -> tuple[int, ...]: + if self.kv_cache_layout_plan is None: + return () + return self.kv_cache_layout_plan.target_layer_ids + + @property + def draft_layer_ids(self) -> tuple[int, ...]: + if self.kv_cache_layout_plan is None: + return () + return self.kv_cache_layout_plan.draft_layer_ids + + def get_layer_view(self, owner: LayerCacheOwner) -> Optional[KVCacheLayerView]: + """Return a non-owning attention view for ``owner`` on this rank.""" + layout_plan = self.kv_cache_layout_plan + if layout_plan is None or not any(layer.owner is owner for layer in layout_plan.layers): + return None + view = self._layer_views.get(owner) + if view is None: + view = KVCacheLayerView(self, owner) + self._layer_views[owner] = view + return view + + def get_draft_kv_cache_view(self) -> Optional[KVCacheLayerView]: + """Return the draft view only for an explicitly unified layout. + + Legacy construction may append draft-shaped layers without opting an + algorithm into the unified lifecycle. Requiring an explicit plan keeps + those existing algorithms on their current metadata path. + """ + if self._explicit_layout_plan is None or self.is_draft: + return None + return self.get_layer_view(LayerCacheOwner.DRAFT) + + def get_target_kv_cache_view(self) -> Optional[KVCacheLayerView]: + """Return the target view only for an explicitly unified layout.""" + if self._explicit_layout_plan is None or self.is_draft: + return None + return self.get_layer_view(LayerCacheOwner.TARGET) + + def get_required_transfer_layer_group_ids(self) -> frozenset[int]: + """Layer groups whose state gates request-level transfer completion.""" + return self.required_transfer_layer_group_ids + @property def blocks_in_primary_pool(self) -> int: """ @@ -5979,22 +6327,32 @@ def get_cache_size_per_token( max_num_tokens: int = 0, spec_config=None, is_draft: bool = False, + layout_plan: Optional[KVCacheLayoutPlan] = None, **kwargs, ): - layer_sizes, attention_windows = _get_static_cache_size_layer_components( - model_config, - mapping, - num_layers=num_layers, - max_seq_len=max_seq_len, - kv_cache_config=kv_cache_config, - is_external_draft=( - is_draft - and spec_config is not None - and spec_config.spec_dec_mode.is_external_drafter() - ), - ) + if layout_plan is not None: + if layout_plan.tokens_per_block != tokens_per_block: + raise ValueError( + "layout-plan tokens_per_block does not match cost calculation: " + f"{layout_plan.tokens_per_block} != {tokens_per_block}" + ) + layer_sizes, attention_windows = layout_plan.cache_size_layer_components() + else: + layer_sizes, attention_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + num_layers=num_layers, + max_seq_len=max_seq_len, + kv_cache_config=kv_cache_config, + is_external_draft=( + is_draft + and spec_config is not None + and spec_config.spec_dec_mode.is_external_drafter() + ), + ) reuse_backoff_enabled = ( - kv_cache_config is not None + layout_plan is None + and kv_cache_config is not None and kv_cache_config.enable_block_reuse and cls._supports_reuse_match_backoff ) @@ -6062,6 +6420,7 @@ def _get_per_request_layer_tokens( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layout_plan: Optional[KVCacheLayoutPlan] = None, ) -> Optional[tuple[List[tuple[int, int, int]], int]]: """Per-local-layer ``(layer_size, retained_tokens, intercept_tokens)`` plus the block-rounded full-sequence token count. @@ -6088,7 +6447,7 @@ def _get_per_request_layer_tokens( """ if max_seq_len is None or int(max_seq_len) <= 0: return None - if num_layers is None: + if layout_plan is None and num_layers is None: total_hidden_layers = getattr(model_config.pretrained_config, "num_hidden_layers", None) if ( total_hidden_layers is not None @@ -6097,25 +6456,34 @@ def _get_per_request_layer_tokens( return None max_seq_len = int(max_seq_len) tokens_per_block = int(tokens_per_block) - layer_sizes, covered_windows = _get_static_cache_size_layer_components( - model_config, - mapping, - num_layers=num_layers, - max_seq_len=max_seq_len, - kv_cache_config=kv_cache_config, - is_external_draft=( - is_draft - and spec_config is not None - and spec_config.spec_dec_mode.is_external_drafter() - ), - ) + if layout_plan is not None: + if layout_plan.tokens_per_block != tokens_per_block: + raise ValueError( + "layout-plan tokens_per_block does not match per-request cost calculation: " + f"{layout_plan.tokens_per_block} != {tokens_per_block}" + ) + layer_sizes, covered_windows = layout_plan.cache_size_layer_components() + else: + layer_sizes, covered_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + num_layers=num_layers, + max_seq_len=max_seq_len, + kv_cache_config=kv_cache_config, + is_external_draft=( + is_draft + and spec_config is not None + and spec_config.spec_dec_mode.is_external_drafter() + ), + ) if not layer_sizes: return None # ``covered_windows`` are the windows ``get_cache_size_per_token`` # charges as fixed per-request retention in the affine intercept; keep # them in lockstep with its accounting (reuse backoff included). if ( - kv_cache_config is not None + layout_plan is None + and kv_cache_config is not None and kv_cache_config.enable_block_reuse and cls._supports_reuse_match_backoff ): @@ -6129,7 +6497,7 @@ def _get_per_request_layer_tokens( # against the draft pool), so the drafter's own windows do not # bound per-request retention. retained_windows = [None] * len(layer_sizes) - elif all(window is None for window in covered_windows): + elif layout_plan is None and all(window is None for window in covered_windows): # The static estimator only sees windows the config carries. # Runtime manager creation additionally derives per-layer windows # from ``layer_types`` (see @@ -6202,6 +6570,7 @@ def get_cache_bytes_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layout_plan: Optional[KVCacheLayoutPlan] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -6234,6 +6603,7 @@ def get_cache_bytes_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, + layout_plan=layout_plan, ) if layer_tokens is None: return None @@ -6262,6 +6632,7 @@ def get_cache_quota_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layout_plan: Optional[KVCacheLayoutPlan] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -6296,6 +6667,7 @@ def get_cache_quota_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, + layout_plan=layout_plan, ) if layer_tokens is None: return None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py new file mode 100644 index 000000000000..944e84b916b1 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from enum import Enum +from fractions import Fraction + +from tensorrt_llm.bindings import DataType +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + AttentionReusePolicy, + BufferConfig, + DataRole, + LayerId, +) + + +class LayerCacheOwner(str, Enum): + """Model component that reads and writes a cache layer.""" + + TARGET = "target" + DRAFT = "draft" + + +class LayerCacheTransferPolicy(str, Enum): + """How a cache layer participates in request-level KV transfer.""" + + REQUIRED = "required" + RECOMPUTABLE = "recomputable" + LOCAL_ONLY = "local_only" + + +@dataclass(frozen=True, slots=True) +class LayerBufferSpec: + """Exact physical buffer allocated for one cache layer. + + ``size`` is the byte size of one physical page, matching + :class:`BufferConfig`. Keeping the exact size in the pre-construction plan + lets budgeting and runtime allocation consume the same value. + """ + + role: DataRole + size: int + tokens_per_block_override: int | None = None + is_sparse: bool = False + transfer_policy: LayerCacheTransferPolicy = LayerCacheTransferPolicy.REQUIRED + + def __post_init__(self) -> None: + if self.size <= 0: + raise ValueError(f"layer buffer size must be positive, got {self.size}") + if self.tokens_per_block_override is not None and self.tokens_per_block_override <= 0: + raise ValueError( + "tokens_per_block_override must be positive when set, got " + f"{self.tokens_per_block_override}" + ) + + def bytes_per_token(self, default_tokens_per_block: int) -> Fraction: + """Return the exact average storage bytes per represented token.""" + tokens_per_block = self.tokens_per_block_override or default_tokens_per_block + return Fraction(self.size, tokens_per_block) + + def to_buffer_config(self) -> BufferConfig: + """Convert the semantic plan entry to the bound runtime config.""" + return BufferConfig( + role=self.role, + size=self.size, + tokens_per_block_override=self.tokens_per_block_override, + is_sparse=self.is_sparse, + ) + + +@dataclass(frozen=True, slots=True) +class LayerCacheSpec: + """Semantic and physical cache layout for one rank-local layer.""" + + local_layer_id: int + global_layer_id: int + owner: LayerCacheOwner + dtype: DataType + num_kv_heads: int + head_dim: int + buffers: tuple[LayerBufferSpec, ...] + sliding_window_size: int | None = None + num_sink_tokens: int | None = None + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + transfer_policy: LayerCacheTransferPolicy = LayerCacheTransferPolicy.REQUIRED + + def __post_init__(self) -> None: + if self.local_layer_id < 0: + raise ValueError(f"local_layer_id must be non-negative, got {self.local_layer_id}") + if self.global_layer_id < 0: + raise ValueError(f"global_layer_id must be non-negative, got {self.global_layer_id}") + if self.num_kv_heads <= 0: + raise ValueError(f"num_kv_heads must be positive, got {self.num_kv_heads}") + if self.head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {self.head_dim}") + if not self.buffers: + raise ValueError("a cache layer must contain at least one buffer") + roles = [buffer.role for buffer in self.buffers] + if len(set(roles)) != len(roles): + raise ValueError( + f"layer {self.global_layer_id} contains duplicate buffer roles: {roles}" + ) + sparse_values = {buffer.is_sparse for buffer in self.buffers} + if len(sparse_values) != 1: + raise ValueError(f"layer {self.global_layer_id} mixes sparse and dense buffers") + if self.sliding_window_size is not None and self.sliding_window_size <= 0: + raise ValueError( + f"sliding_window_size must be positive when set, got {self.sliding_window_size}" + ) + if self.num_sink_tokens is not None and self.num_sink_tokens < 0: + raise ValueError( + f"num_sink_tokens must be non-negative when set, got {self.num_sink_tokens}" + ) + + @property + def buffer_roles(self) -> tuple[DataRole, ...]: + return tuple(buffer.role for buffer in self.buffers) + + @property + def is_sparse(self) -> bool: + return self.buffers[0].is_sparse + + def bytes_per_token(self, default_tokens_per_block: int) -> Fraction: + """Return the exact aggregate bytes per represented token.""" + return sum( + (buffer.bytes_per_token(default_tokens_per_block) for buffer in self.buffers), + start=Fraction(), + ) + + def to_attention_layer_config(self) -> AttentionLayerConfig: + """Build the existing C++-facing configuration for this layer.""" + return AttentionLayerConfig( + layer_id=LayerId(self.local_layer_id), + buffers=[buffer.to_buffer_config() for buffer in self.buffers], + sliding_window_size=self.sliding_window_size, + num_sink_tokens=self.num_sink_tokens, + reuse_policy=self.reuse_policy, + ) + + +@dataclass(frozen=True, slots=True) +class KVCacheLayoutPlan: + """Immutable rank-local layout consumed by budgeting and allocation.""" + + tokens_per_block: int + layers: tuple[LayerCacheSpec, ...] + + def __post_init__(self) -> None: + if self.tokens_per_block <= 0: + raise ValueError(f"tokens_per_block must be positive, got {self.tokens_per_block}") + if not self.layers: + raise ValueError("a KV cache layout plan must contain at least one local layer") + local_layer_ids = tuple(layer.local_layer_id for layer in self.layers) + expected_local_ids = tuple(range(len(self.layers))) + if local_layer_ids != expected_local_ids: + raise ValueError( + "layout-plan local layer IDs must be dense and ordered: " + f"expected {expected_local_ids}, got {local_layer_ids}" + ) + global_layer_ids = [layer.global_layer_id for layer in self.layers] + if len(set(global_layer_ids)) != len(global_layer_ids): + raise ValueError(f"layout plan contains duplicate global layer IDs: {global_layer_ids}") + + @property + def target_layer_ids(self) -> tuple[int, ...]: + return tuple( + layer.global_layer_id for layer in self.layers if layer.owner is LayerCacheOwner.TARGET + ) + + @property + def draft_layer_ids(self) -> tuple[int, ...]: + return tuple( + layer.global_layer_id for layer in self.layers if layer.owner is LayerCacheOwner.DRAFT + ) + + @property + def has_mixed_dtype(self) -> bool: + return len({layer.dtype for layer in self.layers}) > 1 + + @property + def uniform_dtype(self) -> DataType | None: + dtypes = {layer.dtype for layer in self.layers} + return next(iter(dtypes)) if len(dtypes) == 1 else None + + def get_layer(self, global_layer_id: int) -> LayerCacheSpec: + for layer in self.layers: + if layer.global_layer_id == global_layer_id: + return layer + raise KeyError(f"global layer {global_layer_id} is not in this rank's KV cache layout") + + def required_transfer_layers(self) -> tuple[LayerCacheSpec, ...]: + return tuple( + layer + for layer in self.layers + if layer.transfer_policy is LayerCacheTransferPolicy.REQUIRED + and any( + buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED + for buffer in layer.buffers + ) + ) + + def cache_size_layer_components(self) -> tuple[list[int], list[int | None]]: + """Return per-layer bytes/token and retention windows for budgeting.""" + layer_sizes = [] + for layer in self.layers: + size = layer.bytes_per_token(self.tokens_per_block) + if size.denominator != 1: + raise ValueError( + f"layer {layer.global_layer_id} has fractional aggregate bytes/token: {size}" + ) + layer_sizes.append(int(size)) + return layer_sizes, [layer.sliding_window_size for layer in self.layers] + + def transfer_compatibility_errors(self) -> tuple[str, ...]: + """Return current native-disaggregation limitations for this plan.""" + layers = self.required_transfer_layers() + if not layers: + return () + errors = [] + if len({layer.dtype for layer in layers}) > 1: + errors.append("required transfer layers use different dtypes") + if len({layer.num_kv_heads for layer in layers}) > 1: + errors.append("required transfer layers use different KV-head counts") + if len({layer.head_dim for layer in layers}) > 1: + errors.append("required transfer layers use different head dimensions") + kv_roles = { + tuple( + buffer.role + for buffer in layer.buffers + if buffer.role in _KV_DATA_ROLES + and buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED + ) + for layer in layers + } + if len(kv_roles) > 1: + errors.append("required transfer layers use different K/V role layouts") + return tuple(errors) + + def to_attention_layer_configs(self) -> list[AttentionLayerConfig]: + """Build the existing per-layer runtime configurations.""" + return [layer.to_attention_layer_config() for layer in self.layers] + + +_KV_DATA_ROLES = frozenset({DataRole("key"), DataRole("value")}) + + +__all__ = [ + "KVCacheLayoutPlan", + "LayerBufferSpec", + "LayerCacheOwner", + "LayerCacheSpec", + "LayerCacheTransferPolicy", +] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py new file mode 100644 index 000000000000..495567afac86 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + BlockReusePolicy, + KVCacheLayerView, + KVCacheManagerV2, + Role, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.layout import ( + KVCacheLayoutPlan, + LayerBufferSpec, + LayerCacheOwner, + LayerCacheSpec, + LayerCacheTransferPolicy, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy, GpuCacheTierConfig + + +def _layer( + local_layer_id: int, + owner: LayerCacheOwner, + *, + dtype: DataType = DataType.HALF, + num_kv_heads: int = 4, + head_dim: int = 64, + window: int | None = None, + sinks: int | None = None, + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED, +) -> LayerCacheSpec: + return LayerCacheSpec( + local_layer_id=local_layer_id, + global_layer_id=local_layer_id, + owner=owner, + dtype=dtype, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + buffers=( + LayerBufferSpec(role=Role.KEY, size=1024), + LayerBufferSpec(role=Role.VALUE, size=1024), + ), + sliding_window_size=window, + num_sink_tokens=sinks, + reuse_policy=reuse_policy, + ) + + +@pytest.mark.cpu_only +def test_layout_plan_builds_exact_runtime_layer_configs() -> None: + plan = KVCacheLayoutPlan( + tokens_per_block=16, + layers=( + _layer(0, LayerCacheOwner.TARGET, window=64, sinks=4), + _layer( + 1, + LayerCacheOwner.DRAFT, + reuse_policy=AttentionReusePolicy.OPTIONAL, + ), + ), + ) + + runtime_layers = plan.to_attention_layer_configs() + + assert plan.target_layer_ids == (0,) + assert plan.draft_layer_ids == (1,) + assert runtime_layers[0].sliding_window_size == 64 + assert runtime_layers[0].num_sink_tokens == 4 + assert runtime_layers[1].reuse_policy == AttentionReusePolicy.OPTIONAL + assert [(buffer.role, buffer.size) for buffer in runtime_layers[1].buffers] == [ + (Role.KEY, 1024), + (Role.VALUE, 1024), + ] + assert plan.cache_size_layer_components() == ([128, 128], [64, None]) + + +@pytest.mark.cpu_only +def test_layout_plan_rejects_non_dense_local_layer_ids() -> None: + with pytest.raises(ValueError, match="dense and ordered"): + KVCacheLayoutPlan( + tokens_per_block=16, + layers=( + _layer(0, LayerCacheOwner.TARGET), + LayerCacheSpec( + local_layer_id=2, + global_layer_id=2, + owner=LayerCacheOwner.DRAFT, + dtype=DataType.HALF, + num_kv_heads=4, + head_dim=64, + buffers=(LayerBufferSpec(role=Role.KEY, size=1024),), + ), + ), + ) + + +@pytest.mark.cpu_only +def test_transfer_compatibility_reports_mixed_required_geometry() -> None: + plan = KVCacheLayoutPlan( + tokens_per_block=16, + layers=( + _layer(0, LayerCacheOwner.TARGET), + _layer( + 1, + LayerCacheOwner.DRAFT, + dtype=DataType.FP8, + num_kv_heads=2, + head_dim=128, + ), + ), + ) + + assert plan.transfer_compatibility_errors() == ( + "required transfer layers use different dtypes", + "required transfer layers use different KV-head counts", + "required transfer layers use different head dimensions", + ) + + +@pytest.mark.cpu_only +def test_base_config_is_derived_from_explicit_layout_plan() -> None: + plan = KVCacheLayoutPlan( + tokens_per_block=16, + layers=( + _layer(0, LayerCacheOwner.TARGET, window=64, sinks=4), + _layer(1, LayerCacheOwner.DRAFT), + ), + ) + manager = object.__new__(KVCacheManagerV2) + manager._explicit_layout_plan = plan + manager.enable_swa_scratch_reuse = False + manager.num_extra_kv_tokens = 0 + manager.max_beam_width = 1 + manager.max_seq_len = 128 + manager.max_batch_size = 2 + manager.max_cuda_graph_batch_size = None + manager.max_num_tokens = 64 + manager.enable_partial_reuse = True + manager.reuse_match_backoff = 0 + manager.enable_stats = False + manager.block_reuse_policy = BlockReusePolicy.ALL_REUSABLE + manager._ledger_tokens_per_block = 16 + + config = manager._build_base_config( + KvCacheConfig(), + tokens_per_block=16, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + ) + + assert [layer.sliding_window_size for layer in config.layers] == [64, None] + assert [layer.num_sink_tokens for layer in config.layers] == [4, None] + assert [ + [(buffer.role, buffer.size) for buffer in layer.buffers] for layer in config.layers + ] == [ + [(Role.KEY, 1024), (Role.VALUE, 1024)], + [(Role.KEY, 1024), (Role.VALUE, 1024)], + ] + + +@pytest.mark.cpu_only +def test_draft_layer_view_does_not_own_lifecycle_or_transfer() -> None: + plan = KVCacheLayoutPlan( + tokens_per_block=16, + layers=( + _layer(0, LayerCacheOwner.TARGET), + _layer(1, LayerCacheOwner.DRAFT), + ), + ) + manager = SimpleNamespace(kv_cache_layout_plan=plan) + manager.get_layer_cache_spec = plan.get_layer + manager.get_buffers = Mock(return_value="draft-buffer") + view = KVCacheLayerView(manager, LayerCacheOwner.DRAFT) + + assert view.layer_ids == (1,) + assert view.get_buffers(1) == "draft-buffer" + view.free_resources(Mock()) + with pytest.raises(RuntimeError, match="owning unified manager"): + view.report_batch_to_connector(Mock()) + with pytest.raises(ValueError, match="does not belong"): + view.get_buffers(0) + + +@pytest.mark.cpu_only +def test_local_only_buffer_does_not_gate_transfer() -> None: + layer = _layer(0, LayerCacheOwner.TARGET) + draft = _layer(1, LayerCacheOwner.DRAFT) + draft = LayerCacheSpec( + local_layer_id=draft.local_layer_id, + global_layer_id=draft.global_layer_id, + owner=draft.owner, + dtype=draft.dtype, + num_kv_heads=draft.num_kv_heads, + head_dim=draft.head_dim, + buffers=draft.buffers + + ( + LayerBufferSpec( + role=Role.INDEX_KEY, + size=256, + transfer_policy=LayerCacheTransferPolicy.LOCAL_ONLY, + ), + ), + ) + plan = KVCacheLayoutPlan(tokens_per_block=16, layers=(layer, draft)) + + assert plan.transfer_compatibility_errors() == () From dad9f1f36a589e437416d85a3194631eeb6f4616 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 16:45:45 -0700 Subject: [PATCH 02/23] [None][refactor] simplify unified KV cache layer specs Signed-off-by: Allison Lim --- ...ache-v2-target-draft-unification-design.md | 20 +- .../kv_cache/kv_cache_manager_v2.py | 491 +++++++----------- .../_torch/pyexecutor/kv_cache/layout.py | 266 ---------- .../executor/kv_cache/test_kv_cache_layout.py | 195 +++---- 4 files changed, 258 insertions(+), 714 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py diff --git a/docs/source/features/kv-cache-v2-target-draft-unification-design.md b/docs/source/features/kv-cache-v2-target-draft-unification-design.md index 7b89a4f37b7d..0575552f8934 100644 --- a/docs/source/features/kv-cache-v2-target-draft-unification-design.md +++ b/docs/source/features/kv-cache-v2-target-draft-unification-design.md @@ -110,28 +110,26 @@ temporary data layouts to be identical. ### 1. Introduce one canonical per-layer cache descriptor -Create an internal descriptor such as: +Define a compact internal descriptor directly in `kv_cache_manager_v2.py`: ```python @dataclass(frozen=True) -class LayerCacheSpec: +class KVCacheLayerSpec: global_layer_id: int owner: LayerOwner # TARGET or DRAFT - layer_type: LayerType # ATTENTION or SSM dtype: DataType num_kv_heads: int head_dim: int - buffer_specs: tuple[BufferSpec, ...] + buffer_roles: tuple[DataRole, ...] sliding_window_size: int | None num_sink_tokens: int | None - is_sparse: bool reuse_policy: AttentionReusePolicy - placement: LayerPlacement ``` -The exact names are flexible. Manager construction should consume a single -target-plus-draft layout instead of deriving most properties from the target -configuration and appending draft layers afterward. +The tuple position is the C++ local layer id. Manager construction and the +pre-construction cost classmethods consume the same target-plus-draft sequence. +There is no separate layout module or transfer-policy abstraction in the first +implementation. `draft_layer_ids` should be derived from `owner`, rather than serving as the only source of layer identity. @@ -197,7 +195,7 @@ Use the following disposition: |---|---| | Separate request mirror and draft admission | **Eliminate** in the unified path | | Separate request index and release behavior | **Eliminate** | -| Draft-layer identification | `LayerCacheSpec.owner` | +| Draft-layer identification | `KVCacheLayerSpec.owner` | | Dynamic-tree generation reserve | Request-wide `SpeculativeCapacityPolicy` | | Accepted-token relocation | Explicit target/draft layer sets or homogeneous layer groups | | Block publication | One radix tree, with publication governed by each lifecycle's existing reuse policy | @@ -344,7 +342,7 @@ remain implementation capability gates. ## Suggested implementation order -1. Add `LayerCacheSpec` and composite layout construction without changing +1. Add `KVCacheLayerSpec` directly to `KVCacheManagerV2` without changing runtime behavior. 2. Centralize all per-layer geometry and byte calculations behind the descriptor. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index d74901fc27fe..9282789ed69b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -66,6 +66,7 @@ DEFAULT_BEAM_INDEX, GPU_LEVEL, AttentionLayerConfig, + AttentionReusePolicy, AttnLifeCycle, Batch, BatchDesc, @@ -131,13 +132,6 @@ request_context, ) from ..scheduler import ScheduledRequests -from .layout import ( - KVCacheLayoutPlan, - LayerBufferSpec, - LayerCacheOwner, - LayerCacheSpec, - LayerCacheTransferPolicy, -) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -277,53 +271,92 @@ class BlockReusePolicy(StrEnum): PER_CONVERSATION = "per_conversation" -class KVCacheLayerView: - """Non-owning layer-scoped view of one :class:`KVCacheManagerV2`. +class KVCacheLayerOwner(StrEnum): + """Model component that reads and writes a cache layer.""" - The view can be passed to attention metadata while request lifecycle and - transfer remain owned by the underlying manager. It intentionally makes - ``free_resources`` a no-op because legacy speculative call sites may free - a draft manager independently after a forward pass. - """ + TARGET = "target" + DRAFT = "draft" - def __init__(self, manager: "KVCacheManagerV2", owner: LayerCacheOwner) -> None: - layout_plan = manager.kv_cache_layout_plan - if layout_plan is None: - raise ValueError("layer views require a KV cache layout plan") - layers = tuple(layer for layer in layout_plan.layers if layer.owner is owner) - if not layers: - raise ValueError(f"KV cache layout contains no {owner.value} layers on this rank") - self._manager = manager - self.owner = owner - self.layer_ids = tuple(layer.global_layer_id for layer in layers) - self.local_layer_ids = tuple(layer.local_layer_id for layer in layers) - @property - def owner_manager(self) -> "KVCacheManagerV2": - return self._manager +@dataclass(frozen=True) +class KVCacheLayerSpec: + """Rank-local attention-layer storage consumed by costing and construction. + + Tuple position is the C++ local layer id. ``global_layer_id`` is the model + layer id used by attention metadata. + """ - def get_layer_spec(self, global_layer_id: int) -> LayerCacheSpec: - if global_layer_id not in self.layer_ids: + global_layer_id: int + owner: KVCacheLayerOwner + dtype: DataType + num_kv_heads: int + head_dim: int + buffer_roles: tuple[DataRole, ...] + sliding_window_size: Optional[int] = None + num_sink_tokens: Optional[int] = None + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + + def __post_init__(self) -> None: + if self.global_layer_id < 0: + raise ValueError(f"global_layer_id must be non-negative, got {self.global_layer_id}") + if self.num_kv_heads <= 0: + raise ValueError(f"num_kv_heads must be positive, got {self.num_kv_heads}") + if self.head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {self.head_dim}") + if not self.buffer_roles: + raise ValueError("buffer_roles must not be empty") + if len(set(self.buffer_roles)) != len(self.buffer_roles): raise ValueError( - f"layer {global_layer_id} does not belong to the {self.owner.value} cache view" + f"layer {self.global_layer_id} contains duplicate buffer roles: {self.buffer_roles}" + ) + supported_roles = { + Role.KEY, + Role.VALUE, + Role.KEY_BLOCK_SCALE, + Role.VALUE_BLOCK_SCALE, + } + unsupported_roles = set(self.buffer_roles) - supported_roles + if unsupported_roles: + raise ValueError( + f"layer {self.global_layer_id} has unsupported standard buffer roles: " + f"{sorted(map(str, unsupported_roles))}" + ) + if Role.KEY not in self.buffer_roles: + raise ValueError(f"layer {self.global_layer_id} must contain a key buffer") + scale_roles = {Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE} + if any(role in scale_roles for role in self.buffer_roles) and self.dtype != DataType.NVFP4: + raise ValueError("block-scale roles require NVFP4 KV-cache dtype") + if self.dtype == DataType.NVFP4: + expected_scale_roles = {Role.KEY_BLOCK_SCALE} + if Role.VALUE in self.buffer_roles: + expected_scale_roles.add(Role.VALUE_BLOCK_SCALE) + if not expected_scale_roles.issubset(self.buffer_roles): + raise ValueError( + f"NVFP4 layer {self.global_layer_id} must contain matching block-scale roles" + ) + if self.sliding_window_size is not None and self.sliding_window_size <= 0: + raise ValueError( + f"sliding_window_size must be positive when set, got {self.sliding_window_size}" + ) + if self.num_sink_tokens is not None and self.num_sink_tokens < 0: + raise ValueError( + f"num_sink_tokens must be non-negative when set, got {self.num_sink_tokens}" ) - return self._manager.get_layer_cache_spec(global_layer_id) - - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: - self.get_layer_spec(layer_idx) - return self._manager.get_buffers(layer_idx, kv_layout=kv_layout) - - def free_resources(self, request: LlmRequest) -> None: - """Leave request lifecycle to the owning manager.""" - - def report_batch_to_connector(self, *args, **kwargs) -> None: - raise RuntimeError("KV transfer must be initiated through the owning unified manager") - def __getattr__(self, name): - manager = self.__dict__.get("_manager") - if manager is None: - raise AttributeError(name) - return getattr(manager, name) + def bytes_per_token(self, role: Optional[DataRole] = None) -> int: + """Return storage bytes per token for one role or the whole layer.""" + if role is None: + return sum(self.bytes_per_token(buffer_role) for buffer_role in self.buffer_roles) + if role not in self.buffer_roles: + raise ValueError(f"role {role} is not present in layer {self.global_layer_id}") + elements = self.num_kv_heads * self.head_dim + if role in (Role.KEY, Role.VALUE): + return get_size_in_bytes(elements, self.dtype) + if role in (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE): + if elements % 16 != 0: + raise ValueError(f"NVFP4 layer {self.global_layer_id} size must be divisible by 16") + return get_size_in_bytes(elements // 16, DataType.FP8) + raise ValueError(f"unsupported buffer role: {role}") def _swa_endpoint_priority( @@ -821,6 +854,21 @@ def normalize_window_size(window_size: Optional[int]) -> Optional[int]: return layer_sizes, attention_windows +def _get_layer_spec_cache_size_components( + layer_specs: Sequence[KVCacheLayerSpec], +) -> tuple[List[int], List[Optional[int]]]: + """Return the shared costing inputs for an explicit layer-spec sequence.""" + if not layer_specs: + raise ValueError("layer_specs must contain at least one rank-local layer") + global_layer_ids = [layer.global_layer_id for layer in layer_specs] + if len(set(global_layer_ids)) != len(global_layer_ids): + raise ValueError(f"layer_specs contains duplicate global layer IDs: {global_layer_ids}") + return ( + [layer.bytes_per_token() for layer in layer_specs], + [layer.sliding_window_size for layer in layer_specs], + ) + + class _MmRunMetadata(NamedTuple): """CPU tensors for exact multimodal run lookup. @@ -1346,20 +1394,13 @@ def __init__( cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, - layout_plan: Optional[KVCacheLayoutPlan] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> None: self.mapping = mapping self.dtype = dtype self._declared_target_layer_count = num_layers - self._explicit_layout_plan = layout_plan - self.kv_cache_layout_plan: Optional[KVCacheLayoutPlan] = layout_plan - self._layer_views: dict[LayerCacheOwner, KVCacheLayerView] = {} - if layout_plan is not None and layout_plan.tokens_per_block != tokens_per_block: - raise ValueError( - "layout-plan tokens_per_block does not match the manager: " - f"{layout_plan.tokens_per_block} != {tokens_per_block}" - ) + self.layer_specs = tuple(layer_specs) if layer_specs is not None else None self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager @@ -1384,25 +1425,26 @@ def __init__( and global_layer_id >= self._declared_target_layer_count ) } - if layout_plan is not None: - planned_global_layers = [layer.global_layer_id for layer in layout_plan.layers] + if self.layer_specs is not None: + planned_global_layers = [layer.global_layer_id for layer in self.layer_specs] if planned_global_layers != self.pp_layers: raise ValueError( - "layout-plan global layers do not match pipeline placement: " + "layer-spec global layers do not match pipeline placement: " f"{planned_global_layers} != {self.pp_layers}" ) if is_draft and any( - layer.owner is not LayerCacheOwner.DRAFT for layer in layout_plan.layers + layer.owner is not KVCacheLayerOwner.DRAFT for layer in self.layer_specs ): - raise ValueError("a draft mirror manager cannot own target layout layers") - if layout_plan.has_mixed_dtype: + raise ValueError("a draft mirror manager cannot own target cache layers") + layer_dtypes = {layer.dtype for layer in self.layer_specs} + if len(layer_dtypes) != 1: raise NotImplementedError( - "KVCacheManagerV2 does not yet interpret mixed-dtype layout plans" + "KVCacheManagerV2 does not yet interpret mixed-dtype layer specs" ) - if layout_plan.uniform_dtype != dtype: + layer_dtype = next(iter(layer_dtypes)) + if layer_dtype != dtype: raise ValueError( - "layout-plan dtype does not match the manager dtype: " - f"{layout_plan.uniform_dtype} != {dtype}" + f"layer-spec dtype does not match the manager dtype: {layer_dtype} != {dtype}" ) # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) @@ -1545,14 +1587,14 @@ def __init__( self.pp_layers, ) assert len(self.max_attention_window_vec) == self.num_local_layers - if layout_plan is not None: + if self.layer_specs is not None: self.max_attention_window_vec = [ - layer.sliding_window_size for layer in layout_plan.layers + layer.sliding_window_size for layer in self.layer_specs ] # A D-token match backoff needs D more physical SWA history at the # committed endpoint so commit_min_snapshot still covers the rewind. - if layout_plan is None and kv_cache_config.enable_block_reuse and self.reuse_match_backoff: + if kv_cache_config.enable_block_reuse and self.reuse_match_backoff: self.max_attention_window_vec = _extend_swa_windows_for_reuse( self.max_attention_window_vec, self.reuse_match_backoff, @@ -1661,21 +1703,21 @@ def append_to_kv_heads_per_layer( f"unique values={set(self.head_dim_per_layer)}" ) - if layout_plan is not None: - planned_heads = [layer.num_kv_heads for layer in layout_plan.layers] + if self.layer_specs is not None: + planned_heads = [layer.num_kv_heads for layer in self.layer_specs] if planned_heads != self.num_kv_heads_per_layer: raise ValueError( - "layout-plan KV-head counts do not match the manager: " + "layer-spec KV-head counts do not match the manager: " f"{planned_heads} != {self.num_kv_heads_per_layer}" ) - planned_head_dims = [layer.head_dim for layer in layout_plan.layers] + planned_head_dims = [layer.head_dim for layer in self.layer_specs] if planned_head_dims != self.head_dim_per_layer: raise ValueError( - "layout-plan head dimensions do not match the manager: " + "layer-spec head dimensions do not match the manager: " f"{planned_head_dims} != {self.head_dim_per_layer}" ) planned_kv_factors = { - 2 if Role.VALUE in layer.buffer_roles else 1 for layer in layout_plan.layers + 2 if Role.VALUE in layer.buffer_roles else 1 for layer in self.layer_specs } if planned_kv_factors != {self.kv_factor}: raise NotImplementedError( @@ -1823,8 +1865,6 @@ def append_to_kv_heads_per_layer( ) config = self._build_cache_config(config) config = self._remove_zero_size_buffers(config) - self._finalize_layout_plan(config) - self._validate_transfer_layout() has_host_cache_tier = any( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) @@ -1955,7 +1995,6 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: layer_id: self.impl.get_layer_group_id(layer_id) for layer_id in typed_range(LayerId(num_layers)) } - self.required_transfer_layer_group_ids = self._get_required_transfer_layer_group_ids() self.kv_cache_map: dict[int, _KVCache] = {} self._disagg_receive_ready: dict[int, torch.cuda.Event] = {} @@ -2606,9 +2645,11 @@ def _kv_pool_mapping_offset( ) def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: - explicit_layout_plan = getattr(self, "_explicit_layout_plan", None) - if explicit_layout_plan is not None: - return explicit_layout_plan.cache_size_layer_components() + if self.layer_specs is not None: + return ( + [layer.bytes_per_token() for layer in self.layer_specs], + list(self.max_attention_window_vec), + ) layer_sizes = [] attention_windows = [] for local_layer_idx in range(self.num_local_layers): @@ -3122,9 +3163,23 @@ def _build_base_config( ) ) - explicit_layout_plan = getattr(self, "_explicit_layout_plan", None) - if explicit_layout_plan is not None: - layer_configs = explicit_layout_plan.to_attention_layer_configs() + if self.layer_specs is not None: + layer_configs = [ + AttentionLayerConfig( + layer_id=LayerId(local_layer_id), + buffers=[ + BufferConfig( + role=role, + size=layer.bytes_per_token(role) * tokens_per_block, + ) + for role in layer.buffer_roles + ], + sliding_window_size=self.max_attention_window_vec[local_layer_id], + num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, + ) + for local_layer_id, layer in enumerate(self.layer_specs) + ] else: # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional # per-layer BufferConfig entries — for example a sparse index-K @@ -3205,157 +3260,6 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config - def _legacy_layer_owner(self, global_layer_id: int) -> LayerCacheOwner: - if global_layer_id in self._legacy_draft_layer_ids: - return LayerCacheOwner.DRAFT - return LayerCacheOwner.TARGET - - def _layout_plan_from_runtime_config( - self, config: KVCacheManagerConfigPy - ) -> Optional[KVCacheLayoutPlan]: - if not all(isinstance(layer, AttentionLayerConfig) for layer in config.layers): - return None - if any(heads <= 0 for heads in self.num_kv_heads_per_layer) or any( - head_dim <= 0 for head_dim in self.head_dim_per_layer - ): - return None - ignored_roles = self.get_disagg_ignored_roles() - layers = [] - for local_layer_id, layer in enumerate(config.layers): - assert isinstance(layer, AttentionLayerConfig) - global_layer_id = self.pp_layers[local_layer_id] - buffers = tuple( - LayerBufferSpec( - role=buffer.role, - size=buffer.size, - tokens_per_block_override=buffer.tokens_per_block_override, - is_sparse=buffer.is_sparse, - transfer_policy=( - LayerCacheTransferPolicy.LOCAL_ONLY - if buffer.role in ignored_roles - else LayerCacheTransferPolicy.REQUIRED - ), - ) - for buffer in layer.buffers - ) - layers.append( - LayerCacheSpec( - local_layer_id=local_layer_id, - global_layer_id=global_layer_id, - owner=self._legacy_layer_owner(global_layer_id), - dtype=self.dtype, - num_kv_heads=self.num_kv_heads_per_layer[local_layer_id], - head_dim=self.head_dim_per_layer[local_layer_id], - buffers=buffers, - sliding_window_size=layer.sliding_window_size, - num_sink_tokens=layer.num_sink_tokens, - reuse_policy=layer.reuse_policy, - ) - ) - return KVCacheLayoutPlan(tokens_per_block=self.tokens_per_block, layers=tuple(layers)) - - @staticmethod - def _buffer_config_signature(buffer: BufferConfig) -> tuple[DataRole, int, int | None, bool]: - return ( - buffer.role, - buffer.size, - buffer.tokens_per_block_override, - buffer.is_sparse, - ) - - def _validate_runtime_config_matches_layout_plan( - self, config: KVCacheManagerConfigPy, layout_plan: KVCacheLayoutPlan - ) -> None: - if len(config.layers) != len(layout_plan.layers): - raise ValueError( - "runtime cache layer count does not match the layout plan: " - f"{len(config.layers)} != {len(layout_plan.layers)}" - ) - for planned, runtime in zip(layout_plan.layers, config.layers): - if not isinstance(runtime, AttentionLayerConfig): - raise ValueError( - "an attention layout plan cannot produce a non-attention runtime layer" - ) - planned_buffers = tuple( - self._buffer_config_signature(buffer.to_buffer_config()) - for buffer in planned.buffers - ) - runtime_buffers = tuple( - self._buffer_config_signature(buffer) for buffer in runtime.buffers - ) - if planned_buffers != runtime_buffers: - raise ValueError( - f"runtime buffers for layer {planned.global_layer_id} do not match " - f"the layout plan: {runtime_buffers} != {planned_buffers}" - ) - runtime_lifecycle = ( - runtime.sliding_window_size, - runtime.num_sink_tokens, - runtime.reuse_policy, - ) - planned_lifecycle = ( - planned.sliding_window_size, - planned.num_sink_tokens, - planned.reuse_policy, - ) - if runtime_lifecycle != planned_lifecycle: - raise ValueError( - f"runtime lifecycle for layer {planned.global_layer_id} does not match " - f"the layout plan: {runtime_lifecycle} != {planned_lifecycle}" - ) - - def _finalize_layout_plan(self, config: KVCacheManagerConfigPy) -> None: - if self._explicit_layout_plan is None: - self.kv_cache_layout_plan = self._layout_plan_from_runtime_config(config) - return - self._validate_runtime_config_matches_layout_plan(config, self._explicit_layout_plan) - self.kv_cache_layout_plan = self._explicit_layout_plan - - def _validate_transfer_layout(self) -> None: - layout_plan = self.kv_cache_layout_plan - if layout_plan is None or (not self.is_disagg and self.kv_connector_manager is None): - return - errors = list(layout_plan.transfer_compatibility_errors()) - ignored_roles = self.get_disagg_ignored_roles() - for layer in layout_plan.layers: - local_only_roles = [ - buffer.role - for buffer in layer.buffers - if buffer.transfer_policy is LayerCacheTransferPolicy.LOCAL_ONLY - and buffer.role not in ignored_roles - ] - if local_only_roles: - errors.append( - f"layer {layer.global_layer_id} marks roles {local_only_roles} local-only, " - "but the manager does not exclude them from disaggregation" - ) - for layer in layout_plan.required_transfer_layers(): - ignored_required_roles = [ - buffer.role - for buffer in layer.buffers - if buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED - and buffer.role in ignored_roles - ] - if ignored_required_roles: - errors.append( - f"layer {layer.global_layer_id} requires transfer of ignored roles " - f"{ignored_required_roles}" - ) - if errors: - raise NotImplementedError( - "the unified KV layout is incompatible with the current transfer metadata: " - + "; ".join(errors) - ) - - def _get_required_transfer_layer_group_ids(self) -> frozenset[int]: - layout_plan = self.kv_cache_layout_plan - if layout_plan is None: - return frozenset() - return frozenset( - int(self.impl.get_layer_group_id(LayerId(layer.local_layer_id))) - for layer in layout_plan.required_transfer_layers() - ) - def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Exclude empty buffers before creating the runtime storage pools.""" if config.layers and all( @@ -3477,14 +3381,24 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() - def get_layer_cache_spec(self, global_layer_id: int) -> LayerCacheSpec: - """Return the immutable layout entry for one rank-local layer.""" - if self.kv_cache_layout_plan is None: - raise RuntimeError("this specialized manager has no attention-only layout plan") - return self.kv_cache_layout_plan.get_layer(global_layer_id) - - def get_layer_owner(self, global_layer_id: int) -> LayerCacheOwner: - return self.get_layer_cache_spec(global_layer_id).owner + def get_layer_cache_spec(self, global_layer_id: int) -> KVCacheLayerSpec: + """Return an explicitly supplied rank-local layer specification.""" + if self.layer_specs is None: + raise RuntimeError("this manager was not constructed with explicit layer specs") + try: + local_layer_id = self.layer_offsets[global_layer_id] + except KeyError as error: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") from error + return self.layer_specs[local_layer_id] + + def get_layer_owner(self, global_layer_id: int) -> KVCacheLayerOwner: + if self.layer_specs is not None: + return self.get_layer_cache_spec(global_layer_id).owner + if global_layer_id in self._legacy_draft_layer_ids: + return KVCacheLayerOwner.DRAFT + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return KVCacheLayerOwner.TARGET def get_layer_dtype(self, global_layer_id: int) -> DataType: return self.get_layer_cache_spec(global_layer_id).dtype @@ -3492,52 +3406,21 @@ def get_layer_dtype(self, global_layer_id: int) -> DataType: def get_layer_buffer_roles(self, global_layer_id: int) -> tuple[DataRole, ...]: return self.get_layer_cache_spec(global_layer_id).buffer_roles - def get_layer_transfer_policy(self, global_layer_id: int) -> LayerCacheTransferPolicy: - return self.get_layer_cache_spec(global_layer_id).transfer_policy - @property def target_layer_ids(self) -> tuple[int, ...]: - if self.kv_cache_layout_plan is None: - return () - return self.kv_cache_layout_plan.target_layer_ids + return tuple( + layer_id + for layer_id in self.pp_layers + if self.get_layer_owner(layer_id) is KVCacheLayerOwner.TARGET + ) @property def draft_layer_ids(self) -> tuple[int, ...]: - if self.kv_cache_layout_plan is None: - return () - return self.kv_cache_layout_plan.draft_layer_ids - - def get_layer_view(self, owner: LayerCacheOwner) -> Optional[KVCacheLayerView]: - """Return a non-owning attention view for ``owner`` on this rank.""" - layout_plan = self.kv_cache_layout_plan - if layout_plan is None or not any(layer.owner is owner for layer in layout_plan.layers): - return None - view = self._layer_views.get(owner) - if view is None: - view = KVCacheLayerView(self, owner) - self._layer_views[owner] = view - return view - - def get_draft_kv_cache_view(self) -> Optional[KVCacheLayerView]: - """Return the draft view only for an explicitly unified layout. - - Legacy construction may append draft-shaped layers without opting an - algorithm into the unified lifecycle. Requiring an explicit plan keeps - those existing algorithms on their current metadata path. - """ - if self._explicit_layout_plan is None or self.is_draft: - return None - return self.get_layer_view(LayerCacheOwner.DRAFT) - - def get_target_kv_cache_view(self) -> Optional[KVCacheLayerView]: - """Return the target view only for an explicitly unified layout.""" - if self._explicit_layout_plan is None or self.is_draft: - return None - return self.get_layer_view(LayerCacheOwner.TARGET) - - def get_required_transfer_layer_group_ids(self) -> frozenset[int]: - """Layer groups whose state gates request-level transfer completion.""" - return self.required_transfer_layer_group_ids + return tuple( + layer_id + for layer_id in self.pp_layers + if self.get_layer_owner(layer_id) is KVCacheLayerOwner.DRAFT + ) @property def blocks_in_primary_pool(self) -> int: @@ -6327,16 +6210,11 @@ def get_cache_size_per_token( max_num_tokens: int = 0, spec_config=None, is_draft: bool = False, - layout_plan: Optional[KVCacheLayoutPlan] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ): - if layout_plan is not None: - if layout_plan.tokens_per_block != tokens_per_block: - raise ValueError( - "layout-plan tokens_per_block does not match cost calculation: " - f"{layout_plan.tokens_per_block} != {tokens_per_block}" - ) - layer_sizes, attention_windows = layout_plan.cache_size_layer_components() + if layer_specs is not None: + layer_sizes, attention_windows = _get_layer_spec_cache_size_components(layer_specs) else: layer_sizes, attention_windows = _get_static_cache_size_layer_components( model_config, @@ -6351,8 +6229,7 @@ def get_cache_size_per_token( ), ) reuse_backoff_enabled = ( - layout_plan is None - and kv_cache_config is not None + kv_cache_config is not None and kv_cache_config.enable_block_reuse and cls._supports_reuse_match_backoff ) @@ -6420,7 +6297,7 @@ def _get_per_request_layer_tokens( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, - layout_plan: Optional[KVCacheLayoutPlan] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> Optional[tuple[List[tuple[int, int, int]], int]]: """Per-local-layer ``(layer_size, retained_tokens, intercept_tokens)`` plus the block-rounded full-sequence token count. @@ -6447,7 +6324,7 @@ def _get_per_request_layer_tokens( """ if max_seq_len is None or int(max_seq_len) <= 0: return None - if layout_plan is None and num_layers is None: + if layer_specs is None and num_layers is None: total_hidden_layers = getattr(model_config.pretrained_config, "num_hidden_layers", None) if ( total_hidden_layers is not None @@ -6456,13 +6333,8 @@ def _get_per_request_layer_tokens( return None max_seq_len = int(max_seq_len) tokens_per_block = int(tokens_per_block) - if layout_plan is not None: - if layout_plan.tokens_per_block != tokens_per_block: - raise ValueError( - "layout-plan tokens_per_block does not match per-request cost calculation: " - f"{layout_plan.tokens_per_block} != {tokens_per_block}" - ) - layer_sizes, covered_windows = layout_plan.cache_size_layer_components() + if layer_specs is not None: + layer_sizes, covered_windows = _get_layer_spec_cache_size_components(layer_specs) else: layer_sizes, covered_windows = _get_static_cache_size_layer_components( model_config, @@ -6482,8 +6354,7 @@ def _get_per_request_layer_tokens( # charges as fixed per-request retention in the affine intercept; keep # them in lockstep with its accounting (reuse backoff included). if ( - layout_plan is None - and kv_cache_config is not None + kv_cache_config is not None and kv_cache_config.enable_block_reuse and cls._supports_reuse_match_backoff ): @@ -6497,7 +6368,7 @@ def _get_per_request_layer_tokens( # against the draft pool), so the drafter's own windows do not # bound per-request retention. retained_windows = [None] * len(layer_sizes) - elif layout_plan is None and all(window is None for window in covered_windows): + elif layer_specs is None and all(window is None for window in covered_windows): # The static estimator only sees windows the config carries. # Runtime manager creation additionally derives per-layer windows # from ``layer_types`` (see @@ -6570,7 +6441,7 @@ def get_cache_bytes_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, - layout_plan: Optional[KVCacheLayoutPlan] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -6603,7 +6474,7 @@ def get_cache_bytes_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, - layout_plan=layout_plan, + layer_specs=layer_specs, ) if layer_tokens is None: return None @@ -6632,7 +6503,7 @@ def get_cache_quota_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, - layout_plan: Optional[KVCacheLayoutPlan] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -6667,7 +6538,7 @@ def get_cache_quota_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, - layout_plan=layout_plan, + layer_specs=layer_specs, ) if layer_tokens is None: return None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py deleted file mode 100644 index 944e84b916b1..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/layout.py +++ /dev/null @@ -1,266 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass -from enum import Enum -from fractions import Fraction - -from tensorrt_llm.bindings import DataType -from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - AttentionLayerConfig, - AttentionReusePolicy, - BufferConfig, - DataRole, - LayerId, -) - - -class LayerCacheOwner(str, Enum): - """Model component that reads and writes a cache layer.""" - - TARGET = "target" - DRAFT = "draft" - - -class LayerCacheTransferPolicy(str, Enum): - """How a cache layer participates in request-level KV transfer.""" - - REQUIRED = "required" - RECOMPUTABLE = "recomputable" - LOCAL_ONLY = "local_only" - - -@dataclass(frozen=True, slots=True) -class LayerBufferSpec: - """Exact physical buffer allocated for one cache layer. - - ``size`` is the byte size of one physical page, matching - :class:`BufferConfig`. Keeping the exact size in the pre-construction plan - lets budgeting and runtime allocation consume the same value. - """ - - role: DataRole - size: int - tokens_per_block_override: int | None = None - is_sparse: bool = False - transfer_policy: LayerCacheTransferPolicy = LayerCacheTransferPolicy.REQUIRED - - def __post_init__(self) -> None: - if self.size <= 0: - raise ValueError(f"layer buffer size must be positive, got {self.size}") - if self.tokens_per_block_override is not None and self.tokens_per_block_override <= 0: - raise ValueError( - "tokens_per_block_override must be positive when set, got " - f"{self.tokens_per_block_override}" - ) - - def bytes_per_token(self, default_tokens_per_block: int) -> Fraction: - """Return the exact average storage bytes per represented token.""" - tokens_per_block = self.tokens_per_block_override or default_tokens_per_block - return Fraction(self.size, tokens_per_block) - - def to_buffer_config(self) -> BufferConfig: - """Convert the semantic plan entry to the bound runtime config.""" - return BufferConfig( - role=self.role, - size=self.size, - tokens_per_block_override=self.tokens_per_block_override, - is_sparse=self.is_sparse, - ) - - -@dataclass(frozen=True, slots=True) -class LayerCacheSpec: - """Semantic and physical cache layout for one rank-local layer.""" - - local_layer_id: int - global_layer_id: int - owner: LayerCacheOwner - dtype: DataType - num_kv_heads: int - head_dim: int - buffers: tuple[LayerBufferSpec, ...] - sliding_window_size: int | None = None - num_sink_tokens: int | None = None - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED - transfer_policy: LayerCacheTransferPolicy = LayerCacheTransferPolicy.REQUIRED - - def __post_init__(self) -> None: - if self.local_layer_id < 0: - raise ValueError(f"local_layer_id must be non-negative, got {self.local_layer_id}") - if self.global_layer_id < 0: - raise ValueError(f"global_layer_id must be non-negative, got {self.global_layer_id}") - if self.num_kv_heads <= 0: - raise ValueError(f"num_kv_heads must be positive, got {self.num_kv_heads}") - if self.head_dim <= 0: - raise ValueError(f"head_dim must be positive, got {self.head_dim}") - if not self.buffers: - raise ValueError("a cache layer must contain at least one buffer") - roles = [buffer.role for buffer in self.buffers] - if len(set(roles)) != len(roles): - raise ValueError( - f"layer {self.global_layer_id} contains duplicate buffer roles: {roles}" - ) - sparse_values = {buffer.is_sparse for buffer in self.buffers} - if len(sparse_values) != 1: - raise ValueError(f"layer {self.global_layer_id} mixes sparse and dense buffers") - if self.sliding_window_size is not None and self.sliding_window_size <= 0: - raise ValueError( - f"sliding_window_size must be positive when set, got {self.sliding_window_size}" - ) - if self.num_sink_tokens is not None and self.num_sink_tokens < 0: - raise ValueError( - f"num_sink_tokens must be non-negative when set, got {self.num_sink_tokens}" - ) - - @property - def buffer_roles(self) -> tuple[DataRole, ...]: - return tuple(buffer.role for buffer in self.buffers) - - @property - def is_sparse(self) -> bool: - return self.buffers[0].is_sparse - - def bytes_per_token(self, default_tokens_per_block: int) -> Fraction: - """Return the exact aggregate bytes per represented token.""" - return sum( - (buffer.bytes_per_token(default_tokens_per_block) for buffer in self.buffers), - start=Fraction(), - ) - - def to_attention_layer_config(self) -> AttentionLayerConfig: - """Build the existing C++-facing configuration for this layer.""" - return AttentionLayerConfig( - layer_id=LayerId(self.local_layer_id), - buffers=[buffer.to_buffer_config() for buffer in self.buffers], - sliding_window_size=self.sliding_window_size, - num_sink_tokens=self.num_sink_tokens, - reuse_policy=self.reuse_policy, - ) - - -@dataclass(frozen=True, slots=True) -class KVCacheLayoutPlan: - """Immutable rank-local layout consumed by budgeting and allocation.""" - - tokens_per_block: int - layers: tuple[LayerCacheSpec, ...] - - def __post_init__(self) -> None: - if self.tokens_per_block <= 0: - raise ValueError(f"tokens_per_block must be positive, got {self.tokens_per_block}") - if not self.layers: - raise ValueError("a KV cache layout plan must contain at least one local layer") - local_layer_ids = tuple(layer.local_layer_id for layer in self.layers) - expected_local_ids = tuple(range(len(self.layers))) - if local_layer_ids != expected_local_ids: - raise ValueError( - "layout-plan local layer IDs must be dense and ordered: " - f"expected {expected_local_ids}, got {local_layer_ids}" - ) - global_layer_ids = [layer.global_layer_id for layer in self.layers] - if len(set(global_layer_ids)) != len(global_layer_ids): - raise ValueError(f"layout plan contains duplicate global layer IDs: {global_layer_ids}") - - @property - def target_layer_ids(self) -> tuple[int, ...]: - return tuple( - layer.global_layer_id for layer in self.layers if layer.owner is LayerCacheOwner.TARGET - ) - - @property - def draft_layer_ids(self) -> tuple[int, ...]: - return tuple( - layer.global_layer_id for layer in self.layers if layer.owner is LayerCacheOwner.DRAFT - ) - - @property - def has_mixed_dtype(self) -> bool: - return len({layer.dtype for layer in self.layers}) > 1 - - @property - def uniform_dtype(self) -> DataType | None: - dtypes = {layer.dtype for layer in self.layers} - return next(iter(dtypes)) if len(dtypes) == 1 else None - - def get_layer(self, global_layer_id: int) -> LayerCacheSpec: - for layer in self.layers: - if layer.global_layer_id == global_layer_id: - return layer - raise KeyError(f"global layer {global_layer_id} is not in this rank's KV cache layout") - - def required_transfer_layers(self) -> tuple[LayerCacheSpec, ...]: - return tuple( - layer - for layer in self.layers - if layer.transfer_policy is LayerCacheTransferPolicy.REQUIRED - and any( - buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED - for buffer in layer.buffers - ) - ) - - def cache_size_layer_components(self) -> tuple[list[int], list[int | None]]: - """Return per-layer bytes/token and retention windows for budgeting.""" - layer_sizes = [] - for layer in self.layers: - size = layer.bytes_per_token(self.tokens_per_block) - if size.denominator != 1: - raise ValueError( - f"layer {layer.global_layer_id} has fractional aggregate bytes/token: {size}" - ) - layer_sizes.append(int(size)) - return layer_sizes, [layer.sliding_window_size for layer in self.layers] - - def transfer_compatibility_errors(self) -> tuple[str, ...]: - """Return current native-disaggregation limitations for this plan.""" - layers = self.required_transfer_layers() - if not layers: - return () - errors = [] - if len({layer.dtype for layer in layers}) > 1: - errors.append("required transfer layers use different dtypes") - if len({layer.num_kv_heads for layer in layers}) > 1: - errors.append("required transfer layers use different KV-head counts") - if len({layer.head_dim for layer in layers}) > 1: - errors.append("required transfer layers use different head dimensions") - kv_roles = { - tuple( - buffer.role - for buffer in layer.buffers - if buffer.role in _KV_DATA_ROLES - and buffer.transfer_policy is LayerCacheTransferPolicy.REQUIRED - ) - for layer in layers - } - if len(kv_roles) > 1: - errors.append("required transfer layers use different K/V role layouts") - return tuple(errors) - - def to_attention_layer_configs(self) -> list[AttentionLayerConfig]: - """Build the existing per-layer runtime configurations.""" - return [layer.to_attention_layer_config() for layer in self.layers] - - -_KV_DATA_ROLES = frozenset({DataRole("key"), DataRole("value")}) - - -__all__ = [ - "KVCacheLayoutPlan", - "LayerBufferSpec", - "LayerCacheOwner", - "LayerCacheSpec", - "LayerCacheTransferPolicy", -] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py index 495567afac86..51c2dd1ac102 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py @@ -13,23 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from types import SimpleNamespace -from unittest.mock import Mock - import pytest from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, - KVCacheLayerView, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, -) -from tensorrt_llm._torch.pyexecutor.kv_cache.layout import ( - KVCacheLayoutPlan, - LayerBufferSpec, - LayerCacheOwner, - LayerCacheSpec, - LayerCacheTransferPolicy, + _get_layer_spec_cache_size_components, ) from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -37,8 +29,8 @@ def _layer( - local_layer_id: int, - owner: LayerCacheOwner, + global_layer_id: int, + owner: KVCacheLayerOwner, *, dtype: DataType = DataType.HALF, num_kv_heads: int = 4, @@ -46,18 +38,17 @@ def _layer( window: int | None = None, sinks: int | None = None, reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED, -) -> LayerCacheSpec: - return LayerCacheSpec( - local_layer_id=local_layer_id, - global_layer_id=local_layer_id, +) -> KVCacheLayerSpec: + roles = (Role.KEY, Role.VALUE) + if dtype == DataType.NVFP4: + roles += (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + return KVCacheLayerSpec( + global_layer_id=global_layer_id, owner=owner, dtype=dtype, num_kv_heads=num_kv_heads, head_dim=head_dim, - buffers=( - LayerBufferSpec(role=Role.KEY, size=1024), - LayerBufferSpec(role=Role.VALUE, size=1024), - ), + buffer_roles=roles, sliding_window_size=window, num_sink_tokens=sinks, reuse_policy=reuse_policy, @@ -65,87 +56,51 @@ def _layer( @pytest.mark.cpu_only -def test_layout_plan_builds_exact_runtime_layer_configs() -> None: - plan = KVCacheLayoutPlan( - tokens_per_block=16, - layers=( - _layer(0, LayerCacheOwner.TARGET, window=64, sinks=4), - _layer( - 1, - LayerCacheOwner.DRAFT, - reuse_policy=AttentionReusePolicy.OPTIONAL, - ), +def test_layer_specs_describe_combined_cache_cost() -> None: + layer_specs = ( + _layer(0, KVCacheLayerOwner.TARGET, window=64, sinks=4), + _layer( + 1, + KVCacheLayerOwner.DRAFT, + reuse_policy=AttentionReusePolicy.OPTIONAL, ), ) - runtime_layers = plan.to_attention_layer_configs() - - assert plan.target_layer_ids == (0,) - assert plan.draft_layer_ids == (1,) - assert runtime_layers[0].sliding_window_size == 64 - assert runtime_layers[0].num_sink_tokens == 4 - assert runtime_layers[1].reuse_policy == AttentionReusePolicy.OPTIONAL - assert [(buffer.role, buffer.size) for buffer in runtime_layers[1].buffers] == [ - (Role.KEY, 1024), - (Role.VALUE, 1024), - ] - assert plan.cache_size_layer_components() == ([128, 128], [64, None]) + assert _get_layer_spec_cache_size_components(layer_specs) == ([1024, 1024], [64, None]) @pytest.mark.cpu_only -def test_layout_plan_rejects_non_dense_local_layer_ids() -> None: - with pytest.raises(ValueError, match="dense and ordered"): - KVCacheLayoutPlan( - tokens_per_block=16, - layers=( - _layer(0, LayerCacheOwner.TARGET), - LayerCacheSpec( - local_layer_id=2, - global_layer_id=2, - owner=LayerCacheOwner.DRAFT, - dtype=DataType.HALF, - num_kv_heads=4, - head_dim=64, - buffers=(LayerBufferSpec(role=Role.KEY, size=1024),), - ), - ), +def test_layer_specs_reject_duplicate_global_layer_ids() -> None: + with pytest.raises(ValueError, match="duplicate global layer IDs"): + _get_layer_spec_cache_size_components( + ( + _layer(0, KVCacheLayerOwner.TARGET), + _layer(0, KVCacheLayerOwner.DRAFT), + ) ) @pytest.mark.cpu_only -def test_transfer_compatibility_reports_mixed_required_geometry() -> None: - plan = KVCacheLayoutPlan( - tokens_per_block=16, - layers=( - _layer(0, LayerCacheOwner.TARGET), - _layer( - 1, - LayerCacheOwner.DRAFT, - dtype=DataType.FP8, - num_kv_heads=2, - head_dim=128, - ), - ), - ) +def test_nvfp4_layer_spec_accounts_for_scale_buffers() -> None: + layer = _layer(0, KVCacheLayerOwner.TARGET, dtype=DataType.NVFP4) - assert plan.transfer_compatibility_errors() == ( - "required transfer layers use different dtypes", - "required transfer layers use different KV-head counts", - "required transfer layers use different head dimensions", - ) + assert layer.bytes_per_token(Role.KEY) == 128 + assert layer.bytes_per_token(Role.KEY_BLOCK_SCALE) == 16 + assert layer.bytes_per_token() == 288 @pytest.mark.cpu_only -def test_base_config_is_derived_from_explicit_layout_plan() -> None: - plan = KVCacheLayoutPlan( - tokens_per_block=16, - layers=( - _layer(0, LayerCacheOwner.TARGET, window=64, sinks=4), - _layer(1, LayerCacheOwner.DRAFT), +def test_base_config_is_derived_from_explicit_layer_specs() -> None: + layer_specs = ( + _layer(0, KVCacheLayerOwner.TARGET, window=64, sinks=4), + _layer( + 1, + KVCacheLayerOwner.DRAFT, + reuse_policy=AttentionReusePolicy.OPTIONAL, ), ) manager = object.__new__(KVCacheManagerV2) - manager._explicit_layout_plan = plan + manager.layer_specs = layer_specs manager.enable_swa_scratch_reuse = False manager.num_extra_kv_tokens = 0 manager.max_beam_width = 1 @@ -167,57 +122,43 @@ def test_base_config_is_derived_from_explicit_layout_plan() -> None: assert [layer.sliding_window_size for layer in config.layers] == [64, None] assert [layer.num_sink_tokens for layer in config.layers] == [4, None] + assert [layer.reuse_policy for layer in config.layers] == [ + AttentionReusePolicy.REQUIRED, + AttentionReusePolicy.OPTIONAL, + ] assert [ [(buffer.role, buffer.size) for buffer in layer.buffers] for layer in config.layers ] == [ - [(Role.KEY, 1024), (Role.VALUE, 1024)], - [(Role.KEY, 1024), (Role.VALUE, 1024)], + [(Role.KEY, 8192), (Role.VALUE, 8192)], + [(Role.KEY, 8192), (Role.VALUE, 8192)], ] @pytest.mark.cpu_only -def test_draft_layer_view_does_not_own_lifecycle_or_transfer() -> None: - plan = KVCacheLayoutPlan( - tokens_per_block=16, - layers=( - _layer(0, LayerCacheOwner.TARGET), - _layer(1, LayerCacheOwner.DRAFT), - ), +def test_manager_reports_explicit_target_and_draft_layers() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = ( + _layer(10, KVCacheLayerOwner.TARGET), + _layer(11, KVCacheLayerOwner.DRAFT), ) - manager = SimpleNamespace(kv_cache_layout_plan=plan) - manager.get_layer_cache_spec = plan.get_layer - manager.get_buffers = Mock(return_value="draft-buffer") - view = KVCacheLayerView(manager, LayerCacheOwner.DRAFT) + manager.pp_layers = [10, 11] + manager.layer_offsets = {10: 0, 11: 1} + manager._legacy_draft_layer_ids = set() - assert view.layer_ids == (1,) - assert view.get_buffers(1) == "draft-buffer" - view.free_resources(Mock()) - with pytest.raises(RuntimeError, match="owning unified manager"): - view.report_batch_to_connector(Mock()) - with pytest.raises(ValueError, match="does not belong"): - view.get_buffers(0) + assert manager.target_layer_ids == (10,) + assert manager.draft_layer_ids == (11,) + assert manager.get_layer_dtype(11) == DataType.HALF + assert manager.get_layer_buffer_roles(11) == (Role.KEY, Role.VALUE) @pytest.mark.cpu_only -def test_local_only_buffer_does_not_gate_transfer() -> None: - layer = _layer(0, LayerCacheOwner.TARGET) - draft = _layer(1, LayerCacheOwner.DRAFT) - draft = LayerCacheSpec( - local_layer_id=draft.local_layer_id, - global_layer_id=draft.global_layer_id, - owner=draft.owner, - dtype=draft.dtype, - num_kv_heads=draft.num_kv_heads, - head_dim=draft.head_dim, - buffers=draft.buffers - + ( - LayerBufferSpec( - role=Role.INDEX_KEY, - size=256, - transfer_policy=LayerCacheTransferPolicy.LOCAL_ONLY, - ), - ), - ) - plan = KVCacheLayoutPlan(tokens_per_block=16, layers=(layer, draft)) - - assert plan.transfer_compatibility_errors() == () +def test_layer_spec_rejects_nonstandard_buffer_roles() -> None: + with pytest.raises(ValueError, match="unsupported standard buffer roles"): + KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.HALF, + num_kv_heads=4, + head_dim=64, + buffer_roles=(Role.KEY, Role.INDEX_KEY), + ) From 6ea5ad8165a444713672f01163ddb33b66d20736 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 16:47:48 -0700 Subject: [PATCH 03/23] [None][test] remove unified KV cache foundation tests Signed-off-by: Allison Lim --- .../executor/kv_cache/test_kv_cache_layout.py | 164 ------------------ 1 file changed, 164 deletions(-) delete mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py deleted file mode 100644 index 51c2dd1ac102..000000000000 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_layout.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import pytest - -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - BlockReusePolicy, - KVCacheLayerOwner, - KVCacheLayerSpec, - KVCacheManagerV2, - Role, - _get_layer_spec_cache_size_components, -) -from tensorrt_llm.bindings import DataType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy, GpuCacheTierConfig - - -def _layer( - global_layer_id: int, - owner: KVCacheLayerOwner, - *, - dtype: DataType = DataType.HALF, - num_kv_heads: int = 4, - head_dim: int = 64, - window: int | None = None, - sinks: int | None = None, - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED, -) -> KVCacheLayerSpec: - roles = (Role.KEY, Role.VALUE) - if dtype == DataType.NVFP4: - roles += (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) - return KVCacheLayerSpec( - global_layer_id=global_layer_id, - owner=owner, - dtype=dtype, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - buffer_roles=roles, - sliding_window_size=window, - num_sink_tokens=sinks, - reuse_policy=reuse_policy, - ) - - -@pytest.mark.cpu_only -def test_layer_specs_describe_combined_cache_cost() -> None: - layer_specs = ( - _layer(0, KVCacheLayerOwner.TARGET, window=64, sinks=4), - _layer( - 1, - KVCacheLayerOwner.DRAFT, - reuse_policy=AttentionReusePolicy.OPTIONAL, - ), - ) - - assert _get_layer_spec_cache_size_components(layer_specs) == ([1024, 1024], [64, None]) - - -@pytest.mark.cpu_only -def test_layer_specs_reject_duplicate_global_layer_ids() -> None: - with pytest.raises(ValueError, match="duplicate global layer IDs"): - _get_layer_spec_cache_size_components( - ( - _layer(0, KVCacheLayerOwner.TARGET), - _layer(0, KVCacheLayerOwner.DRAFT), - ) - ) - - -@pytest.mark.cpu_only -def test_nvfp4_layer_spec_accounts_for_scale_buffers() -> None: - layer = _layer(0, KVCacheLayerOwner.TARGET, dtype=DataType.NVFP4) - - assert layer.bytes_per_token(Role.KEY) == 128 - assert layer.bytes_per_token(Role.KEY_BLOCK_SCALE) == 16 - assert layer.bytes_per_token() == 288 - - -@pytest.mark.cpu_only -def test_base_config_is_derived_from_explicit_layer_specs() -> None: - layer_specs = ( - _layer(0, KVCacheLayerOwner.TARGET, window=64, sinks=4), - _layer( - 1, - KVCacheLayerOwner.DRAFT, - reuse_policy=AttentionReusePolicy.OPTIONAL, - ), - ) - manager = object.__new__(KVCacheManagerV2) - manager.layer_specs = layer_specs - manager.enable_swa_scratch_reuse = False - manager.num_extra_kv_tokens = 0 - manager.max_beam_width = 1 - manager.max_seq_len = 128 - manager.max_batch_size = 2 - manager.max_cuda_graph_batch_size = None - manager.max_num_tokens = 64 - manager.enable_partial_reuse = True - manager.reuse_match_backoff = 0 - manager.enable_stats = False - manager.block_reuse_policy = BlockReusePolicy.ALL_REUSABLE - manager._ledger_tokens_per_block = 16 - - config = manager._build_base_config( - KvCacheConfig(), - tokens_per_block=16, - cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], - ) - - assert [layer.sliding_window_size for layer in config.layers] == [64, None] - assert [layer.num_sink_tokens for layer in config.layers] == [4, None] - assert [layer.reuse_policy for layer in config.layers] == [ - AttentionReusePolicy.REQUIRED, - AttentionReusePolicy.OPTIONAL, - ] - assert [ - [(buffer.role, buffer.size) for buffer in layer.buffers] for layer in config.layers - ] == [ - [(Role.KEY, 8192), (Role.VALUE, 8192)], - [(Role.KEY, 8192), (Role.VALUE, 8192)], - ] - - -@pytest.mark.cpu_only -def test_manager_reports_explicit_target_and_draft_layers() -> None: - manager = object.__new__(KVCacheManagerV2) - manager.layer_specs = ( - _layer(10, KVCacheLayerOwner.TARGET), - _layer(11, KVCacheLayerOwner.DRAFT), - ) - manager.pp_layers = [10, 11] - manager.layer_offsets = {10: 0, 11: 1} - manager._legacy_draft_layer_ids = set() - - assert manager.target_layer_ids == (10,) - assert manager.draft_layer_ids == (11,) - assert manager.get_layer_dtype(11) == DataType.HALF - assert manager.get_layer_buffer_roles(11) == (Role.KEY, Role.VALUE) - - -@pytest.mark.cpu_only -def test_layer_spec_rejects_nonstandard_buffer_roles() -> None: - with pytest.raises(ValueError, match="unsupported standard buffer roles"): - KVCacheLayerSpec( - global_layer_id=0, - owner=KVCacheLayerOwner.TARGET, - dtype=DataType.HALF, - num_kv_heads=4, - head_dim=64, - buffer_roles=(Role.KEY, Role.INDEX_KEY), - ) From 090cbc02dceecee575d28e1839a12d92651076b8 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 16:51:00 -0700 Subject: [PATCH 04/23] [None][docs] remove unified KV cache design document Signed-off-by: Allison Lim --- ...ache-v2-target-draft-unification-design.md | 381 ------------------ 1 file changed, 381 deletions(-) delete mode 100644 docs/source/features/kv-cache-v2-target-draft-unification-design.md diff --git a/docs/source/features/kv-cache-v2-target-draft-unification-design.md b/docs/source/features/kv-cache-v2-target-draft-unification-design.md deleted file mode 100644 index 0575552f8934..000000000000 --- a/docs/source/features/kv-cache-v2-target-draft-unification-design.md +++ /dev/null @@ -1,381 +0,0 @@ - - - -# Target and draft KV cache unification in KV Cache Manager V2 - -> **Audit basis:** upstream `main` at `dc8451c50864edc07c75dedeafcb91dbb7a5d478` -> (fetched October 12, 2026). - -## Executive verdict - -The central framing is correct: - -> A unified manager owns one committed request timeline. Layer-specific storage -> geometry may differ, but request state, allocation, reuse depth, and admission -> cannot. - -The original report is also correct that reuse policy is already a per-layer -configuration. The audit has four material corrections and one remaining -reuse-policy wiring requirement: - -1. **Per-layer reuse policy already exists in the V2 C++ and Python contracts.** - `AttentionReusePolicy` defines `REQUIRED`, `PRIVATE`, and `OPTIONAL`; - `AttentionLayerConfig.reusePolicy` defaults to `REQUIRED`. The remaining - manager change is to choose and pass the policy explicitly when building - target and draft layer configurations. -2. **`is_draft` should not simply become a per-layer flag.** Some behavior - disappears under unification, some becomes request-wide, and only some - becomes layer metadata. -3. **Per-layer heads, dimensions, and windows are supported by allocation, but - not by all speculative-decoding mutation kernels.** -4. **Mixed dtype requires more than a `get_layer_dtype()` accessor.** Storage is - already byte-oriented, but much of the Python/runtime layer assumes one dtype - and one cache type. -5. **Capacity accounting cannot remain unchanged.** Lifecycle registry behavior - can remain, but estimation, budget splitting, pool sizing, and speculative - headroom must be recomputed for the combined layout. - -The existing limited one-manager speculative path, which appends speculative -layers after the target layers, is a useful starting point. It currently -inherits target-wide assumptions for dtype, layout, and several runtime -operations ([`_util.py`](../../../tensorrt_llm/_torch/pyexecutor/_util.py)). - -## Proposed model - -```mermaid -flowchart TB - subgraph Current["Current separate-manager path"] - R1["Request"] - TM["Target V2 manager
quota + radix tree + request cache"] - DM["Draft V2 manager
quota + radix tree + mirrored request cache"] - TL["Target layers"] - DL["Draft layers"] - - R1 --> TM --> TL - R1 --> DM --> DL - end - - subgraph Unified["Proposed unified path"] - R2["Request"] - UM["One V2 manager
one quota + one radix tree"] - KC["One KvCache
committed history + capacity + beam/state"] - LD["Per-layer descriptors
owner, dtype, geometry, roles,
window, sparse, placement"] - SM["Storage manager"] - P1["Pool/lifecycle group A"] - P2["Pool/lifecycle group B"] - P3["Pool/lifecycle group C"] - - R2 --> UM --> KC - LD --> UM - UM --> SM - SM --> P1 - SM --> P2 - SM --> P3 - end -``` - -A useful refinement is to distinguish: - -- **Committed timeline:** one history length and reuse depth for the request. -- **Speculative scratch tail:** one conservatively sized allocation extent, - within which different layer groups may temporarily use different - token-position mappings. - -That distinction permits one `KvCache` without requiring target and draft -temporary data layouts to be identical. - -## Audit of the original report - -| Claim | Verdict | Audit | -|---|---|---| -| One `KvCache` per request | **Correct** | This is the fundamental unification invariant. Capacity, committed history, beam state, suspension, and request lifetime become shared. | -| One logical block grid and radix tree | **Correct, with nuance** | `tokens_per_block` is manager-wide. Individual buffers may use a `tokensPerBlockOverride`, but it must divide the logical block size ([`config.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h)). | -| One reuse depth | **Correct, with policy semantics** | The request has one claimed reuse endpoint. `REQUIRED` lifecycles constrain that endpoint, `OPTIONAL` lifecycles are checked for complete coverage there without shortening it, and `PRIVATE` lifecycles do not participate ([`kvCache.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp)). | -| One quota, eviction policy, and resume gate | **Correct** | These are manager/request policies and cannot differ between target and draft in a single manager. | -| One sharding scheme | **Correct for the initial implementation** | The current manager has one `Mapping`. Target and draft rank ownership must be representable by that mapping. Helix and other specialized layouts need explicit compatibility checks. | -| One manager specialization | **Mostly correct** | It may be one existing specialization or a composite implementation that supports the union of required hooks. “Same existing class” need not be a permanent restriction. | -| Per-layer heads and head dimension | **Partially correct** | Construction accepts per-layer values, but accepted-token relocation currently requires uniform heads and one local pool ([`kv_cache_manager_v2.py`](../../../tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py)). | -| Per-layer window and sink tokens | **Partially correct** | C++ lifecycle identity supports both window and sink tokens. The generic Python builder currently supplies a per-layer window but always passes `num_sink_tokens=None`. | -| Per-layer buffers and byte sizes | **Correct** | C++ groups buffers by lifecycle and expanded byte size; storage does not require identical layer geometry ([`storage/config.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp)). | -| Sparse and dense layers | **Correct** | `isSparse` is lifecycle identity, and sparse/dense storage pools are separated even when slot sizes match. | -| Per-layer addressing | **Correct, with wording adjustment** | Layers share the request row and committed timeline, but different layers may select different pool groups through `kv_cache_pool_mapping`. They do not necessarily use the same physical block-offset array. | -| Per-layer dtype needs wiring | **Correct, but understated** | C++ storage is byte-oriented, but Python buffer views, cost calculations, attention backends, quantization scale roles, relocation, cold codecs, and disaggregation commonly read a manager-wide dtype. | -| Per-layer reuse policy already exists | **Correct** | `AttentionReusePolicy` and `AttentionLayerConfig.reusePolicy` are present in C++, included in lifecycle identity and pool grouping, and exposed as `reuse_policy` in Python ([`config.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h), [`lifeCycleRegistry.h`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h), [`__init__.pyi`](../../../tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi)). `BlockReusePolicy` remains a separate manager-wide request-scope policy. | -| Stats, connectors, and events are all skipped for draft | **Partially incorrect** | Request stats and connectors have `is_draft` gates. Event-manager construction itself is not categorically disabled for draft managers. | -| Dynamic-tree headroom becomes request-wide | **Correct** | The unified cache must reserve enough scratch capacity for the larger target/draft requirement. The over-reservation is not necessarily small; it depends on `topK`, tree width, draft length, and layer costs. | -| Compatibility reduces to specialization, sharding, and block size | **Incomplete** | These are necessary but not sufficient. Placement, buffer roles/layout, reuse identity, relocation capability, and transfer/codec support also matter. | -| Registry invariant and capacity accounting stay unchanged | **Half correct** | Lifecycle registry semantics can remain. Capacity and quota accounting must change because two independently sized managers become one composite manager. | - -## Recommended changes - -### 1. Introduce one canonical per-layer cache descriptor - -Define a compact internal descriptor directly in `kv_cache_manager_v2.py`: - -```python -@dataclass(frozen=True) -class KVCacheLayerSpec: - global_layer_id: int - owner: LayerOwner # TARGET or DRAFT - dtype: DataType - num_kv_heads: int - head_dim: int - buffer_roles: tuple[DataRole, ...] - sliding_window_size: int | None - num_sink_tokens: int | None - reuse_policy: AttentionReusePolicy -``` - -The tuple position is the C++ local layer id. Manager construction and the -pre-construction cost classmethods consume the same target-plus-draft sequence. -There is no separate layout module or transfer-policy abstraction in the first -implementation. - -`draft_layer_ids` should be derived from `owner`, rather than serving as the -only source of layer identity. - -**Justification:** This prevents dtype, buffer roles, cache type, placement, and -geometry from developing separate parallel lists with inconsistent indexing. - -### 2. Make dtype genuinely per layer - -Add accessors such as: - -```python -get_layer_dtype(global_layer_id) -get_layer_bytes_per_token(global_layer_id, role) -get_layer_buffer_roles(global_layer_id) -``` - -Then update: - -- `BufferConfig` construction. -- Tensor views returned by `get_buffers()`. -- KV byte and capacity calculations. -- NVFP4/FP8 scale-buffer handling. -- Attention backend dtype lookup. -- Cold-page codecs. -- Disaggregated transfer descriptors. -- Accepted-token relocation. - -Do **not** add dtype to `AttnLifeCycle`. Lifecycle identity describes retention -and reuse behavior; storage already separates incompatible byte sizes. - -The existing manager-wide `.dtype` can temporarily remain valid only for -homogeneous layouts. Mixed layouts should fail loudly until every consumer has -moved to a layer-aware accessor. - -**Justification:** Otherwise allocation may be correct while kernels interpret -the bytes using the wrong element type. - -### 3. Replace manager-wide `kv_cache_type` assumptions with per-layer roles - -The wrapper currently derives a manager-wide `kv_factor` from `SELFKONLY` -versus ordinary K/V storage. C++ already understands arbitrary per-layer -`BufferConfig` roles. - -Move toward descriptors such as: - -```text -Target layer 0: K + V -Target MLA layer: K only -Draft layer 48: K + V -FP4 layer: K + V + K_SCALE + V_SCALE -``` - -**Justification:** Dtype unification alone is insufficient if target and draft -use different physical roles or quantization metadata. - -### 4. Split `is_draft` by responsibility - -Do not replace manager-level `is_draft` with a layer-level boolean everywhere. -Use the following disposition: - -| Existing `is_draft` responsibility | Unified disposition | -|---|---| -| Separate request mirror and draft admission | **Eliminate** in the unified path | -| Separate request index and release behavior | **Eliminate** | -| Draft-layer identification | `KVCacheLayerSpec.owner` | -| Dynamic-tree generation reserve | Request-wide `SpeculativeCapacityPolicy` | -| Accepted-token relocation | Explicit target/draft layer sets or homogeneous layer groups | -| Block publication | One radix tree, with publication governed by each lifecycle's existing reuse policy | -| Reuse backoff | One conservative request-wide backoff | -| Stats | Record once; optionally aggregate by layer owner | -| Events | Emit one unified event stream; owner tags are optional metadata | -| Connectors/disaggregation | Transfer the complete required layer set by default | -| Cold-page codec | Select codec/dtype by layer or buffer | -| SWA scratch reuse | Capability decision per lifecycle; initially disable for mixed target/draft managers | -| Helix and special layouts | Compatibility rejection until their ownership model supports the combined layout | - -**Justification:** Most existing draft branches exist because there are two -managers. Preserving those branches as per-layer behavior would retain -unnecessary complexity. - -### 5. Wire the existing reuse policy and define the draft default - -The current V2 core already implements the required policy semantics: - -- `REQUIRED` lifecycles constrain the common prefix match. -- `OPTIONAL` lifecycles do not shorten the `REQUIRED` match. Their coverage is - checked at that endpoint, and the caller selects complete optional groups on - the first resume. Missing or unselected groups receive fresh pages for - recomputation. -- `PRIVATE` lifecycles remain request-local and are neither published to nor - reused from the radix tree. - -The policy participates in `AttnLifeCycle` identity, and storage keeps distinct -reuse policies in separate pool groups. The request-side implementation is in -[`kvCache.cpp`](../../../cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp), -while the Python contract exposes `reuse_policy` in -[`__init__.pyi`](../../../tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi). - -The manager-first work therefore needs to: - -1. Add `reuse_policy` to the composite layer descriptor. -2. Set `AttentionLayerConfig.reuse_policy` explicitly in the Python config - builder instead of relying on its `REQUIRED` default. -3. Map returned layer-group status back to target/draft ownership. -4. Have the scheduler select available `OPTIONAL` groups on first resume. -5. Ensure missing or unselected draft groups are recomputed before a draft - forward consumes that prefix. - -The appropriate draft policy is an algorithm contract: - -- **`REQUIRED`** is the safest initial default. Draft coverage limits the shared - reuse endpoint, so no catch-up path is needed. This matches the original - report's recommendation when draft KV must be reused with target KV. -- **`OPTIONAL`** allows target KV to determine the reuse endpoint and reuses - draft KV only when complete there. Use it only when the algorithm can - explicitly reconstruct a missing draft prefix from target state. -- **`PRIVATE`** always reconstructs the draft prefix and is useful only when - publishing draft KV is undesirable or unsupported. - -Also ensure the radix reuse identity includes every target and draft setting -that can change generated KV, including model or adapter identity where -applicable. - -### 6. Rework capacity and budget accounting - -For an eligible unified configuration: - -- Do not split `max_gpu_total_bytes` between target and draft. -- Do not split host or disk-tier budgets. -- Build one composite storage configuration containing all layers. -- Run pool-ratio selection over the combined lifecycle and buffer groups. -- Compute per-request fixed and variable cost from the composite descriptor. -- Use the larger speculative scratch requirement across target and draft. -- Apply `max_util_for_resume`, eviction, and quota retuning once. - -The current V2 path explicitly splits draft budgets when a separate manager is -selected ([`_util.py`](../../../tensorrt_llm/_torch/pyexecutor/_util.py)). That -logic must be bypassed for a unified plan. - -**Justification:** Adding the independently calculated target and draft -capacities does not necessarily yield the capacity of one shared request -timeline. - -### 7. Add layer-group-aware relocation primitives - -Current accepted-token relocation assumes: - -- One attention window. -- Uniform KV heads. -- One local KV pool. -- One dtype or element size. -- One broad layer range. - -Those assumptions are enforced directly in -[`_update_kv_cache_draft_token_location()`](../../../tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py). - -Introduce a manager-level operation conceptually like: - -```python -relocate_speculative_tokens( - requests, - layer_ids, - source_positions, - destination_positions, -) -``` - -The manager should partition `layer_ids` into compatible execution groups, for -example by: - -```text -(pool, buffer role, dtype, head geometry, window/layout) -``` - -Each group can invoke an appropriate relocation kernel. Once every required -group has been relocated, the shared `KvCache` can be rewound or advanced once. - -**Justification:** This is the main place where “storage supports heterogeneous -layers” currently fails to imply “speculative decoding supports heterogeneous -layers.” - -### 8. Define an explicit compatibility check - -Return a structured result rather than a boolean: - -```python -UnifiedKvCacheCompatibility( - supported=True, - reasons=(), -) -``` - -A target/draft pair is compatible only when: - -- They can share one committed request timeline and beam width. -- They use the same logical `tokens_per_block`. -- Their layer placement and TP/PP/CP ownership fit one mapping. -- One manager implementation can express the union of buffer roles and hooks. -- All dtype/layout consumers are layer-aware. -- The speculative relocation implementation supports the resulting pool groups. -- Reuse hashing and backoff are valid for both models. -- Connectors, offload codecs, and disaggregation can transfer the required union. -- Specialized features such as Helix, DSA, Mamba state, or model-specific - buffers explicitly declare support. - -In the eventual design, dtype, heads, dimensions, and windows should not -independently reject unification. Until their consumers are migrated, they -remain implementation capability gates. - -## Suggested implementation order - -1. Add `KVCacheLayerSpec` directly to `KVCacheManagerV2` without changing - runtime behavior. -2. Centralize all per-layer geometry and byte calculations behind the - descriptor. -3. Implement per-layer dtype and buffer-role accessors. -4. Wire `AttentionLayerConfig.reuse_policy` from the composite descriptor and - add scheduler selection/recomputation handling for `OPTIONAL` draft groups. -5. Add the structured compatibility checker and retain the separate-manager - fallback. -6. Build one V2 C++ storage configuration from target plus draft layers. -7. Replace `is_draft` with request-wide policy, layer ownership, or deletion - according to the table above. -8. Rework composite capacity estimation and stop splitting budgets. -9. Add grouped relocation APIs. -10. Enable unification for the simplest speculative algorithm first. -11. Migrate additional algorithms one by one, expanding the compatibility - checker as each gains support. - -For the initial rollout, use these conservative defaults: - -- Default target and draft lifecycles to `REQUIRED`; enable `OPTIONAL` only for - algorithms with a validated draft-prefix recomputation path. -- Use the maximum required reuse backoff. -- Reserve the maximum speculative scratch width request-wide. -- Disable SWA scratch reuse for mixed managers until validated. -- Reject unsupported CP/Helix/specialized layouts. -- Preserve the existing separate-manager path as the fallback. - -## Revised contract - -> A unified V2 KV-cache manager owns one committed timeline, one request cache, -> one logical block grid, one radix-tree reuse depth, and one memory quota for -> the target and draft layers. Each layer independently describes its buffer -> roles, dtype, geometry, retention lifecycle, sparsity, and physical pool -> mapping. Speculative tails may use layer-specific temporary layouts inside a -> shared request-wide allocation, but committing, rewinding, suspending, -> resuming, and reusing the request occur once and at a common history position. From 511975f4949f7d4a2aa8936c94140fe8d643f0d3 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 22:10:16 -0700 Subject: [PATCH 05/23] [None][feat] support per-layer storage layouts in KVCacheManagerV2 Signed-off-by: Allison Lim --- .../kv_cache/kv_cache_manager_v2.py | 200 ++++++++++++------ 1 file changed, 132 insertions(+), 68 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9282789ed69b..637881237797 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -358,6 +358,11 @@ def bytes_per_token(self, role: Optional[DataRole] = None) -> int: return get_size_in_bytes(elements // 16, DataType.FP8) raise ValueError(f"unsupported buffer role: {role}") + @property + def kv_factor(self) -> int: + """Number of data planes exposed by the layer's paged tensor view.""" + return 2 if Role.VALUE in self.buffer_roles else 1 + def _swa_endpoint_priority( prompt_length: int, tokens_per_block: int, rewind_tokens: int @@ -1408,12 +1413,18 @@ def __init__( self._connector_life_cycle_by_group: Optional[List[AttnLifeCycle]] = None self.kv_cache_type = kv_cache_type - self.pp_layers, self.num_layers = get_pp_layers( - num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) + if self.layer_specs is None: + self.pp_layers, self.num_layers = get_pp_layers( + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + else: + self.pp_layers = [layer.global_layer_id for layer in self.layer_specs] + if len(set(self.pp_layers)) != len(self.pp_layers): + raise ValueError(f"layer-spec global layer ids must be unique: {self.pp_layers}") + self.num_layers = max(num_layers, max(self.pp_layers, default=-1) + 1) self.is_draft = is_draft self._legacy_draft_layer_ids = { global_layer_id @@ -1426,26 +1437,10 @@ def __init__( ) } if self.layer_specs is not None: - planned_global_layers = [layer.global_layer_id for layer in self.layer_specs] - if planned_global_layers != self.pp_layers: - raise ValueError( - "layer-spec global layers do not match pipeline placement: " - f"{planned_global_layers} != {self.pp_layers}" - ) if is_draft and any( layer.owner is not KVCacheLayerOwner.DRAFT for layer in self.layer_specs ): raise ValueError("a draft mirror manager cannot own target cache layers") - layer_dtypes = {layer.dtype for layer in self.layer_specs} - if len(layer_dtypes) != 1: - raise NotImplementedError( - "KVCacheManagerV2 does not yet interpret mixed-dtype layer specs" - ) - layer_dtype = next(iter(layer_dtypes)) - if layer_dtype != dtype: - raise ValueError( - f"layer-spec dtype does not match the manager dtype: {layer_dtype} != {dtype}" - ) # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the @@ -1456,7 +1451,12 @@ def __init__( # Set True by a compression manager; generation-step resize then leaves history untouched. self.kv_compression_manages_history: bool = False self.enable_swa_scratch_reuse = ( - kv_cache_config.enable_swa_scratch_reuse and not self.is_draft + kv_cache_config.enable_swa_scratch_reuse + and not self.is_draft + and not ( + self.layer_specs is not None + and any(layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs) + ) ) block_reuse_config = kv_cache_config.block_reuse_config self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) @@ -1657,7 +1657,12 @@ def __init__( mm_token_id_offset=vocab_size, ) - if isinstance(num_kv_heads, int): + if self.layer_specs is not None: + self.num_kv_heads_per_layer = [layer.num_kv_heads for layer in self.layer_specs] + self.total_num_kv_heads_per_layer = [0] * self.num_layers + for layer in self.layer_specs: + self.total_num_kv_heads_per_layer[layer.global_layer_id] = layer.num_kv_heads + elif isinstance(num_kv_heads, int): self.num_kv_heads_per_layer = [ (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_local_layers) ] @@ -1687,7 +1692,9 @@ def append_to_kv_heads_per_layer( append_to_kv_heads_per_layer(self.total_num_kv_heads_per_layer, kv_head) # Build per-layer head_dim (similar to num_kv_heads_per_layer) - if isinstance(head_dim, int): + if self.layer_specs is not None: + self.head_dim_per_layer = [layer.head_dim for layer in self.layer_specs] + elif isinstance(head_dim, int): self.head_dim_per_layer = [head_dim for _ in range(self.num_local_layers)] else: assert len(head_dim) == self.num_layers, ( @@ -1703,28 +1710,6 @@ def append_to_kv_heads_per_layer( f"unique values={set(self.head_dim_per_layer)}" ) - if self.layer_specs is not None: - planned_heads = [layer.num_kv_heads for layer in self.layer_specs] - if planned_heads != self.num_kv_heads_per_layer: - raise ValueError( - "layer-spec KV-head counts do not match the manager: " - f"{planned_heads} != {self.num_kv_heads_per_layer}" - ) - planned_head_dims = [layer.head_dim for layer in self.layer_specs] - if planned_head_dims != self.head_dim_per_layer: - raise ValueError( - "layer-spec head dimensions do not match the manager: " - f"{planned_head_dims} != {self.head_dim_per_layer}" - ) - planned_kv_factors = { - 2 if Role.VALUE in layer.buffer_roles else 1 for layer in self.layer_specs - } - if planned_kv_factors != {self.kv_factor}: - raise NotImplementedError( - "KVCacheManagerV2 does not yet interpret mixed or mismatched K/V layouts: " - f"plan factors={sorted(planned_kv_factors)}, manager factor={self.kv_factor}" - ) - self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) max_util_for_resume = kv_cache_config.max_util_for_resume @@ -1872,6 +1857,13 @@ def append_to_kv_heads_per_layer( def create_cold_page_codec(cache_config: object) -> Optional[object]: if cold_page_codec_provider is None: return None + if ( + self.layer_specs is not None + and len({layer.dtype for layer in self.layer_specs}) > 1 + ): + raise NotImplementedError( + "cold-page compression does not yet support mixed-dtype KV-cache layer specs" + ) return cold_page_codec_provider.create_cold_page_codec( cache_config, runtime_dtype=self.dtype, @@ -2375,7 +2367,9 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: # Same conversion the attention backends index the buffer # with: layers in one pool can carry different scales. scale = self.get_layer_page_index_scale(layer_idx) - pages = [page * scale // self.kv_factor for page in fresh] + layer_offset = self.layer_offsets[layer_idx] + kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) + pages = [page * scale // kv_factor for page in fresh] pages = [page for page in pages if 0 <= page < buffer.shape[0]] if pages: if _fill_kv_pages(buffer, pages, self._fresh_page_fill): @@ -2412,12 +2406,38 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: When present, role B must be addressable from role A using a constant page-index offset. """ + if self.layer_specs is not None: + layer_id = self._pool_layer_ids_by_role[(pool_id, Role.KEY)] + roles = self._get_layer_roles_by_local_id(int(layer_id)) + return Role.KEY, Role.VALUE if Role.VALUE in roles else None role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE return Role.KEY, role_b + def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].kv_factor + return self.kv_factor + + def _get_layer_page_table_roles( + self, local_layer_id: int, pool_id: int + ) -> Tuple[DataRole, Optional[DataRole]]: + if self.layer_specs is None: + return self._get_pool_roles(pool_id) + return Role.KEY, (Role.VALUE if self.layer_specs[local_layer_id].kv_factor == 2 else None) + + def _get_layer_dtype_by_local_id(self, local_layer_id: int) -> DataType: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].dtype + return self.dtype + + def _get_layer_roles_by_local_id(self, local_layer_id: int) -> tuple[DataRole, ...]: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].buffer_roles + return tuple(self._get_buffer_roles_for_layer(local_layer_id)) + def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[DataRole]: """Select block scales for a role and local layer in the pool mapping.""" - if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + if self._get_layer_dtype_by_local_id(layer_id) != DataType.NVFP4 or role_a != Role.KEY: return None return Role.KEY_BLOCK_SCALE @@ -2430,10 +2450,14 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] + uses_block_scale_pointers = any( + self._get_layer_dtype_by_local_id(layer_id) == DataType.NVFP4 + for layer_id in range(self.num_local_layers) + ) if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): pool_id = self.impl.get_layer_group_id(layer_id) - role_a, _ = self._get_pool_roles(pool_id) + role_a, _ = self._get_layer_page_table_roles(layer_id, pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( @@ -2442,7 +2466,7 @@ def _build_pool_mapping_tensors(self): 0, ] ) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: block_scale_role = self._get_block_scale_role(role_a, layer_id) block_scale_pool_pointers_list.append( [ @@ -2463,7 +2487,7 @@ def _build_pool_mapping_tensors(self): layer_id, role_a, PageIndexMode.SHARED ) kv_cache_pool_pointers_list.append([key_base_addr, 0]) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: # The KEY/scale pointers are a (rep-layer, 0-offset) origin # against which each layer's kv_cache_pool_mapping offset is # resolved. The block-scale origin must reproduce the SAME @@ -2480,7 +2504,7 @@ def _build_pool_mapping_tensors(self): rep_offset = self._kv_pool_mapping_offset(layer_id, pool_id, key_base_addr) scale_stride = ( self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block ) scale_base_addr = ( @@ -2506,13 +2530,16 @@ def _build_pool_mapping_tensors(self): ) offset_divisor = self.impl.get_page_stride(layer_id, role_a) if role_b is not None: - offset_divisor *= self.kv_factor + offset_divisor *= self._get_layer_kv_factor_by_local_id(layer_id) offset = exact_div( addr_offset, offset_divisor, ) - if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + if ( + self._get_layer_dtype_by_local_id(layer_id) != DataType.NVFP4 + or role_a != Role.KEY + ): block_scale_offset = None else: block_scale_role = self._get_block_scale_role(role_a, layer_id) @@ -2529,7 +2556,7 @@ def _build_pool_mapping_tensors(self): block_scale_offset = exact_div( block_scale_addr_offset, self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block, ) @@ -2540,7 +2567,7 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_mapping_list.append([layer_group_id, offset]) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: for pool_id, block_scale_pool_pointers in enumerate(block_scale_pool_pointers_list): pool_pointers = kv_cache_pool_pointers_list[pool_id] kv_cache_pool_pointers_list[pool_id] = [ @@ -2580,6 +2607,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: > 1 for pool_id, layer_ids in enumerate(self.impl.layer_grouping) ) + if self.layer_specs is not None: + storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} + self._use_per_layer_page_tables |= len(storage_layouts) > 1 if self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: @@ -2640,7 +2670,7 @@ def _kv_pool_mapping_offset( return exact_div( addr_offset, self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block, ) @@ -2881,7 +2911,7 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) pool_id = self.layer_to_pool_mapping_dict[layer_id] - role_a, role_b = self._get_pool_roles(pool_id) + role_a, role_b = self._get_layer_page_table_roles(layer_id, pool_id) roles = [ role_a, role_a if role_b is None else role_b, @@ -3070,6 +3100,8 @@ def _copy_batch_block_offsets_per_layer( def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: """Return the primary cache roles allocated for one local layer.""" + if self.layer_specs is not None: + return list(self.layer_specs[local_layer_idx].buffer_roles) roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: roles.append(Role.VALUE) @@ -3401,10 +3433,21 @@ def get_layer_owner(self, global_layer_id: int) -> KVCacheLayerOwner: return KVCacheLayerOwner.TARGET def get_layer_dtype(self, global_layer_id: int) -> DataType: + if self.layer_specs is None: + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self.dtype return self.get_layer_cache_spec(global_layer_id).dtype def get_layer_buffer_roles(self, global_layer_id: int) -> tuple[DataRole, ...]: - return self.get_layer_cache_spec(global_layer_id).buffer_roles + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self._get_layer_roles_by_local_id(self.layer_offsets[global_layer_id]) + + def get_layer_kv_factor(self, global_layer_id: int) -> int: + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self._get_layer_kv_factor_by_local_id(self.layer_offsets[global_layer_id]) @property def target_layer_ids(self) -> tuple[int, ...]: @@ -3422,6 +3465,12 @@ def draft_layer_ids(self) -> tuple[int, ...]: if self.get_layer_owner(layer_id) is KVCacheLayerOwner.DRAFT ) + def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: + """Return this manager when explicit draft layers share its request timeline.""" + if self.layer_specs is None or not self.draft_layer_ids: + return None + return self + @property def blocks_in_primary_pool(self) -> int: """ @@ -3439,7 +3488,8 @@ def is_attention_layer(self, layer_idx: int) -> bool: def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: layer_offset = self.layer_offsets[layer_idx] addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + layer_kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) + if layer_kv_factor == 2: addr_value = self.impl.get_mem_pool_base_address( layer_offset, Role.VALUE, PageIndexMode.SHARED ) @@ -3451,7 +3501,7 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" element_per_container = 1 - dtype = self.dtype + dtype = self._get_layer_dtype_by_local_id(layer_offset) if dtype == DataType.NVFP4: element_per_container = 2 dtype = torch.int8 @@ -3459,16 +3509,16 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. layer_head_dim = self.head_dim_per_layer[layer_offset] if kv_layout == "NHD": shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, - self.kv_factor, + self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // layer_kv_factor, + layer_kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], layer_head_dim // element_per_container, ] else: shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, - self.kv_factor, + self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // layer_kv_factor, + layer_kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, layer_head_dim // element_per_container, @@ -5936,15 +5986,18 @@ def get_batch_cache_indices( if layer_idx is None: pool_id = 0 index_scale = 1 if raw_indices else None + kv_factor = self.kv_factor else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) + kv_factor = self.get_layer_kv_factor(layer_idx) return self._get_batch_cache_indices_by_pool_id( request_ids, pool_id=pool_id, is_kv_aggregate=not raw_indices, num_blocks_per_seq=num_blocks_per_seq, index_scale=index_scale, + kv_factor=kv_factor, ) def _get_batch_cache_indices_by_pool_id( @@ -5955,11 +6008,12 @@ def _get_batch_cache_indices_by_pool_id( is_kv_aggregate: bool = True, num_blocks_per_seq: Optional[Sequence[int]] = None, index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, ) -> List[List[int]]: if is_kv_aggregate: # Div by kv_factor to index kv cache with size # [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] - div_factor = self.kv_factor + div_factor = self.kv_factor if kv_factor is None else kv_factor else: div_factor = 1 @@ -6008,6 +6062,7 @@ def get_batch_cache_indices_flat( if layer_idx is None: pool_id = 0 scale = self._index_scale_ints[pool_id] + div_factor = self.kv_factor else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] # Layers sharing a pool can still require different page-index @@ -6015,7 +6070,7 @@ def get_batch_cache_indices_flat( # scale so this flat block table matches get_batch_cache_indices() # and never feeds out-of-range page ids to FlashInfer. scale = self.get_layer_page_index_scale(layer_idx) - div_factor = self.kv_factor + div_factor = self.get_layer_kv_factor(layer_idx) out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) out = out_tensor.numpy() @@ -6036,6 +6091,8 @@ def get_batch_cache_indices_flat( return out_tensor def get_cache_bytes_per_token(self) -> int: + if self.layer_specs is not None: + return sum(layer.bytes_per_token() for layer in self.layer_specs) data_roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: data_roles.append(Role.VALUE) @@ -6051,6 +6108,12 @@ def get_cache_bytes_per_token(self) -> int: ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): + if self.layer_specs is not None: + layer = self.layer_specs[local_layer_idx] + if data_role == Role.ALL: + return layer.bytes_per_token() + return layer.bytes_per_token(data_role) + if self.dtype not in ( DataType.FP8, DataType.HALF, @@ -6258,6 +6321,7 @@ def get_cache_size_per_token( kv_cache_config is not None and kv_cache_config.enable_swa_scratch_reuse and not is_draft + and not any(spec.owner == KVCacheLayerOwner.DRAFT for spec in (layer_specs or ())) ), generation_capacity_headroom=generation_capacity_headroom, ) From bfc8bc7292ea2378b1446e7b2138168b0a3d89b8 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 14:13:59 -0700 Subject: [PATCH 06/23] fix: normalize explicit KV cache spec windows Signed-off-by: Allison Lim --- .../kv_cache/kv_cache_manager_v2.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 637881237797..3aa4e5b16560 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -861,6 +861,8 @@ def normalize_window_size(window_size: Optional[int]) -> Optional[int]: def _get_layer_spec_cache_size_components( layer_specs: Sequence[KVCacheLayerSpec], + *, + max_seq_len: Optional[int] = None, ) -> tuple[List[int], List[Optional[int]]]: """Return the shared costing inputs for an explicit layer-spec sequence.""" if not layer_specs: @@ -870,7 +872,14 @@ def _get_layer_spec_cache_size_components( raise ValueError(f"layer_specs contains duplicate global layer IDs: {global_layer_ids}") return ( [layer.bytes_per_token() for layer in layer_specs], - [layer.sliding_window_size for layer in layer_specs], + [ + None + if max_seq_len is not None + and layer.sliding_window_size is not None + and layer.sliding_window_size >= max_seq_len + else layer.sliding_window_size + for layer in layer_specs + ], ) @@ -1589,7 +1598,10 @@ def __init__( assert len(self.max_attention_window_vec) == self.num_local_layers if self.layer_specs is not None: self.max_attention_window_vec = [ - layer.sliding_window_size for layer in self.layer_specs + None + if layer.sliding_window_size is None or layer.sliding_window_size >= max_seq_len + else layer.sliding_window_size + for layer in self.layer_specs ] # A D-token match backoff needs D more physical SWA history at the @@ -6277,7 +6289,9 @@ def get_cache_size_per_token( **kwargs, ): if layer_specs is not None: - layer_sizes, attention_windows = _get_layer_spec_cache_size_components(layer_specs) + layer_sizes, attention_windows = _get_layer_spec_cache_size_components( + layer_specs, max_seq_len=max_seq_len + ) else: layer_sizes, attention_windows = _get_static_cache_size_layer_components( model_config, @@ -6398,7 +6412,9 @@ def _get_per_request_layer_tokens( max_seq_len = int(max_seq_len) tokens_per_block = int(tokens_per_block) if layer_specs is not None: - layer_sizes, covered_windows = _get_layer_spec_cache_size_components(layer_specs) + layer_sizes, covered_windows = _get_layer_spec_cache_size_components( + layer_specs, max_seq_len=max_seq_len + ) else: layer_sizes, covered_windows = _get_static_cache_size_layer_components( model_config, From 18c71c1a359fde0dc8d2cec7cede23bb60326308 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 14:43:09 -0700 Subject: [PATCH 07/23] fix: use explicit KV cache dtype for cold-page codec Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 3aa4e5b16560..e4706931b30e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1878,7 +1878,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) return cold_page_codec_provider.create_cold_page_codec( cache_config, - runtime_dtype=self.dtype, + runtime_dtype=self.layer_specs[0].dtype if self.layer_specs else self.dtype, pp_layers=self.pp_layers, num_kv_heads_per_layer=self.num_kv_heads_per_layer, head_dim_per_layer=self.head_dim_per_layer, From ae64dcab3edc349eccafaabed9596eaed842de72 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 14:57:12 -0700 Subject: [PATCH 08/23] test: cover explicit KV cache dtype for cold-page codec Signed-off-by: Allison Lim --- .../kv_cache/test_kvcm2_integration.py | 86 ++++++++++++++++++- 1 file changed, 85 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 3b1fdde85f9a..79c82df60e2a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -28,6 +28,8 @@ from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, @@ -358,6 +360,8 @@ def _make_manager_for_cache_tier_test( joint_reuse: bool = False, mapping: Mapping | None = None, vocab_size: int | None = 16, + dtype: DataType = DataType.HALF, + layer_specs: list[KVCacheLayerSpec] | None = None, ) -> tuple[KVCacheManagerV2, Mock]: impl_constructor = Mock(side_effect=impl_side_effect) if mapping is None: @@ -421,7 +425,8 @@ def build_cache_config( max_seq_len=MAX_SEQ_LEN, max_batch_size=1, mapping=mapping, - dtype=DataType.HALF, + dtype=dtype, + layer_specs=layer_specs, spec_config=spec_config, is_draft=is_draft, is_disagg=is_disagg, @@ -1301,6 +1306,85 @@ def test_cold_codec_provider_receives_draft_role() -> None: assert codec_provider.create_cold_page_codec.call_args.kwargs["is_draft"] is True +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "manager_dtype,layer_dtype", + [ + (DataType.HALF, DataType.BF16), + (DataType.BF16, DataType.HALF), + (DataType.HALF, DataType.FP8), + ], +) +def test_cold_codec_provider_receives_layer_spec_dtype( + manager_dtype: DataType, layer_dtype: DataType +) -> None: + codec_provider = Mock() + layer = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=layer_dtype, + num_kv_heads=1, + head_dim=1, + buffer_roles=(Role.KEY,), + ) + manager, _ = _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + dtype=manager_dtype, + layer_specs=[layer], + ) + + assert manager.dtype == manager_dtype + codec_provider.create_cold_page_codec.assert_called_once() + assert codec_provider.create_cold_page_codec.call_args.kwargs["runtime_dtype"] == layer_dtype + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("manager_dtype", [DataType.HALF, DataType.BF16]) +def test_cold_codec_provider_receives_legacy_manager_dtype(manager_dtype: DataType) -> None: + codec_provider = Mock() + _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + dtype=manager_dtype, + ) + + codec_provider.create_cold_page_codec.assert_called_once() + assert codec_provider.create_cold_page_codec.call_args.kwargs["runtime_dtype"] == manager_dtype + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "layer_dtypes", [(DataType.BF16, DataType.HALF), (DataType.BF16, DataType.FP8)] +) +def test_cold_codec_provider_rejects_mixed_layer_spec_dtypes( + layer_dtypes: tuple[DataType, DataType], +) -> None: + codec_provider = Mock() + layer_specs = [ + KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=layer_dtype, + num_kv_heads=1, + head_dim=1, + buffer_roles=(Role.KEY,), + ) + for layer_id, layer_dtype in enumerate(layer_dtypes) + ] + with pytest.raises(NotImplementedError, match="mixed-dtype KV-cache layer specs"): + _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + layer_specs=layer_specs, + ) + + codec_provider.create_cold_page_codec.assert_not_called() + + def test_extra_tokens_are_in_context_capacity() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=264), From 701abead2a49e0d36e1ef528918dd5d31ac4f5de Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 15:58:55 -0700 Subject: [PATCH 09/23] fix: honor physical slot strides in KV cache views Signed-off-by: Allison Lim --- .../_torch/attention/backends/flashinfer.py | 11 +- .../kv_cache/kv_cache_manager_v2.py | 114 ++++++++++++++---- .../kv_cache/test_kvcm2_integration.py | 2 + 3 files changed, 98 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index e3373666692c..7d6996b2d386 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -1205,20 +1205,23 @@ def _post_init_with_buffers(self, buffers) -> None: self._max_num_blocks_per_seq = self.kv_cache_manager.max_blocks_per_seq # Layers may share one page-index list only when they are in the - # same pool AND have the same page-index scale: VSWA splits pools, + # same pool, page-index scale, and KV factor: VSWA splits pools, # and per-layer geometry (e.g. Gemma4 sliding/global head_dim) # splits scales even when all windows collapse to max_seq_len and - # is_vswa is False. Guarded on V2-specific attributes (V1 managers - # lack the per-pool infrastructure). + # is_vswa is False. Mixed K-only/KV layers can use different view + # indices even with the same physical scale. Guarded on V2-specific + # attributes (V1 managers lack the per-pool infrastructure). mgr = self.kv_cache_manager get_scale = getattr(mgr, 'get_layer_page_index_scale', None) + get_kv_factor = getattr(mgr, 'get_layer_kv_factor', None) layer_space: Dict[int, int] = {} if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale: space_ids = {} for layer_idx in attention_layer_indices: layer_offset = mgr.layer_offsets[layer_idx] key = (mgr.layer_to_pool_mapping_dict[layer_offset], - get_scale(layer_idx)) + get_scale(layer_idx), + get_kv_factor(layer_idx) if get_kv_factor else None) space_ids.setdefault(key, len(space_ids)) layer_space[layer_idx] = space_ids[key] if layer_space: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index e4706931b30e..ef85247b18e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2381,7 +2381,8 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: scale = self.get_layer_page_index_scale(layer_idx) layer_offset = self.layer_offsets[layer_idx] kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) - pages = [page * scale // kv_factor for page in fresh] + view_stride = self._get_cache_view_page_stride(scale, kv_factor) + pages = [page * scale // view_stride for page in fresh] pages = [page for page in pages if 0 <= page < buffer.shape[0]] if pages: if _fill_kv_pages(buffer, pages, self._fresh_page_fill): @@ -3498,6 +3499,7 @@ def is_attention_layer(self, layer_idx: int) -> bool: ) def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + """Return a zero-copy KV view; its slot dimension may be strided.""" layer_offset = self.layer_offsets[layer_idx] addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) layer_kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) @@ -3536,13 +3538,27 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. layer_head_dim // element_per_container, ] - return convert_to_torch_tensor( - TensorWrapper( - addr_key, - dtype, - shape, + wrapper = TensorWrapper(addr_key, dtype, shape) + scale = self.get_layer_page_index_scale(layer_idx) + view_stride = self._get_cache_view_page_stride(scale, layer_kv_factor) + if view_stride != layer_kv_factor: + # A slot containing, for example, three buffers cannot be indexed + # as contiguous K/V pairs. Expose only this layer's buffers with + # the physical slot stride, indexed by raw slot IDs. + converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) + shape[0] = exact_div(page_upper + int(converter.layer_offset), scale) + element_size = torch.empty(0, dtype=wrapper.dtype).element_size() + page_elements = exact_div( + self.impl.get_page_stride(layer_offset, Role.KEY), element_size ) - ) + wrapper.shape = shape + wrapper.strides = [ + view_stride * page_elements, + page_elements, + *[math.prod(shape[i + 1 :]) for i in range(2, len(shape))], + ] + return convert_to_torch_tensor(wrapper) def get_index_k_buffer( self, @@ -5677,10 +5693,17 @@ def get_iteration_stats(self): ) def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor: - block_ids_per_seq = self.get_batch_cache_indices(request_ids) + if self.layer_specs is not None: + block_ids_per_seq = self.get_batch_cache_indices(request_ids, raw_indices=True) + div_factor = 1 + else: + block_ids_per_seq = self.get_batch_cache_indices(request_ids) + div_factor = self.num_local_layers + if self.kv_factor > 1 and self._index_scale_ints[0] % self.kv_factor != 0: + div_factor = 1 block_ids_per_seq_tensors = [ torch.tensor( - [i // self.num_local_layers if i != BAD_PAGE_INDEX else 0 for i in sublist], + [i // div_factor if i != BAD_PAGE_INDEX else 0 for i in sublist], dtype=torch.int, ) for sublist in block_ids_per_seq @@ -5982,6 +6005,15 @@ def get_layer_page_index_scale(self, layer_idx: int) -> int: must not use the pool-level scale.""" return int(self.impl.get_page_index_scale(self.layer_offsets[layer_idx], Role.KEY)) + @staticmethod + def _get_cache_view_page_stride(index_scale: int, kv_factor: int) -> int: + """Physical buffer pages between rows of the KV tensor view. + + Contiguous K/V rows are usable only when every slot starts on a row + boundary. Otherwise, rows span whole slots and use unscaled slot IDs. + """ + return kv_factor if index_scale % kv_factor == 0 else index_scale + def get_batch_cache_indices( self, request_ids: List[int], @@ -5990,7 +6022,10 @@ def get_batch_cache_indices( *, raw_indices: bool = False, ) -> List[List[int]]: - """Return cache indices for a layer, or pool 0 when no layer is given. + """Return indices into a layer's ``get_buffers()`` view. + + With explicit specs, omitting the layer selects pool 0's KEY + representative. Callers accessing other layers must pass their IDs. Set raw_indices for slot-major views that need base slot IDs without page-index scaling or KV aggregation. @@ -5998,7 +6033,11 @@ def get_batch_cache_indices( if layer_idx is None: pool_id = 0 index_scale = 1 if raw_indices else None - kv_factor = self.kv_factor + kv_factor = ( + self._get_layer_kv_factor_by_local_id(self._pool_layer_ids_by_role[(0, Role.KEY)]) + if self.layer_specs is not None + else self.kv_factor + ) else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) @@ -6022,15 +6061,15 @@ def _get_batch_cache_indices_by_pool_id( index_scale: Optional[int] = None, kv_factor: Optional[int] = None, ) -> List[List[int]]: - if is_kv_aggregate: - # Div by kv_factor to index kv cache with size - # [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] - div_factor = self.kv_factor if kv_factor is None else kv_factor - else: - div_factor = 1 - if index_scale is None: index_scale = int(self.index_scales[pool_id]) + div_factor = ( + self._get_cache_view_page_stride( + index_scale, self.kv_factor if kv_factor is None else kv_factor + ) + if is_kv_aggregate + else 1 + ) res = [] for req_idx, req_id in enumerate(request_ids): @@ -6074,7 +6113,11 @@ def get_batch_cache_indices_flat( if layer_idx is None: pool_id = 0 scale = self._index_scale_ints[pool_id] - div_factor = self.kv_factor + kv_factor = ( + self._get_layer_kv_factor_by_local_id(self._pool_layer_ids_by_role[(0, Role.KEY)]) + if self.layer_specs is not None + else self.kv_factor + ) else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] # Layers sharing a pool can still require different page-index @@ -6082,7 +6125,8 @@ def get_batch_cache_indices_flat( # scale so this flat block table matches get_batch_cache_indices() # and never feeds out-of-range page ids to FlashInfer. scale = self.get_layer_page_index_scale(layer_idx) - div_factor = self.get_layer_kv_factor(layer_idx) + kv_factor = self.get_layer_kv_factor(layer_idx) + div_factor = self._get_cache_view_page_stride(scale, kv_factor) out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) out = out_tensor.numpy() @@ -6188,12 +6232,26 @@ def calculate_scaling_factor_size_bytes( def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Tensor]]: pool_handled = set() for layer_id, layer_offset in self.layer_offsets.items(): - pool_id = self.layer_to_pool_mapping_dict[layer_offset] - if pool_id in pool_handled: - continue buffer = self.get_buffers(layer_id) if buffer is None: continue + if self.layer_specs is not None: + # Inspect each explicit layer's own rows, without reading + # another layer's bytes through the wrong dtype. + scale = self.get_layer_page_index_scale(layer_id) + view_stride = self._get_cache_view_page_stride( + scale, self.get_layer_kv_factor(layer_id) + ) + yield layer_id, buffer[:: exact_div(scale, view_stride)] + continue + if buffer.stride(0) != math.prod(buffer.shape[1:]): + # Slot-strided views cover only one layer, including legacy + # layouts that coalesce extra buffers alongside K/V. + yield layer_id, buffer + continue + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled: + continue yield layer_id, buffer pool_handled.add(pool_id) @@ -6218,10 +6276,16 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool # before inference, but that must not wipe the guard-page # sentinel -- doing so would leave the startup warning claiming # a guard that no longer exists. Rewrite the guard page after - # the zero. Each pool is zeroed on one representative layer, so - # only that layer's guard page needs restoring here. + # the zero. Explicit specs expose one view per layer here; + # contiguous legacy views use a single representative. guard_page = self._guard_page_by_layer.get(layer_id) if guard_page is not None and self._guard_page_value is not None: + if self.layer_specs is not None: + scale = self.get_layer_page_index_scale(layer_id) + view_stride = self._get_cache_view_page_stride( + scale, self.get_layer_kv_factor(layer_id) + ) + guard_page = exact_div(guard_page, exact_div(scale, view_stride)) _fill_kv_pages(buffer, [guard_page], self._guard_page_value) torch.cuda.synchronize() diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 79c82df60e2a..c43712e296b7 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -2992,6 +2992,7 @@ def _make_fresh_fill_manager( monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) manager = object.__new__(KVCacheManagerV2) manager._fresh_page_fill = fill_value + manager.layer_specs = None manager._fresh_pages_filled = {} manager._fresh_fill_announced = True manager.kv_cache_map = {1: kv_cache} @@ -3287,6 +3288,7 @@ def test_warmup_zeroing_preserves_guard_page(monkeypatch: pytest.MonkeyPatch) -> manager = object.__new__(KVCacheManagerV2) manager.layer_offsets = {0: 0} + manager.layer_specs = None manager.layer_to_pool_mapping_dict = {0: 0} manager.get_buffers = lambda layer_idx, kv_layout="NHD": buffer manager._guard_page_by_layer = {0: guard_page} From 623816cff862ffaa263b3b3c6e3293f3f47471d9 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 17:19:01 -0700 Subject: [PATCH 10/23] test: cover mixed KV cache odd-slot addressing Signed-off-by: Allison Lim --- .../kv_cache/test_kvcm2_integration.py | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index c43712e296b7..8a7974c8d20a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -72,6 +72,7 @@ KVCacheManager, KVCacheManagerConfig, LayerId, + PageIndexMode, ReuseScope, SsmLayerConfig, _introspection, @@ -637,6 +638,113 @@ def test_zero_size_layers_are_removed_before_pool_allocation( manager.shutdown() +@pytest.mark.parametrize("kv_first", [False, True], ids=["key-first", "kv-first"]) +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) +def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layout: str) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + key_layer_id, kv_layer_id = 0, 1 + key_spec = KVCacheLayerSpec( + global_layer_id=key_layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.HALF, + num_kv_heads=2, + head_dim=8, + buffer_roles=(Role.KEY,), + ) + kv_spec = replace(key_spec, global_layer_id=kv_layer_id, buffer_roles=(Role.KEY, Role.VALUE)) + layer_specs = [kv_spec, key_spec] if kv_first else [key_spec, kv_spec] + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=2, + head_dim=8, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + layer_specs=layer_specs, + ) + try: + key_local_id = manager.layer_offsets[key_layer_id] + kv_local_id = manager.layer_offsets[kv_layer_id] + pool_id = manager.layer_to_pool_mapping_dict[key_local_id] + assert manager.layer_to_pool_mapping_dict[kv_local_id] == pool_id + assert manager.num_pools == 1 + assert manager.get_layer_page_index_scale(key_layer_id) == 3 + assert manager.get_layer_page_index_scale(kv_layer_id) == 3 + + key_buffer = manager.get_buffers(key_layer_id, kv_layout=kv_layout) + kv_buffer = manager.get_buffers(kv_layer_id, kv_layout=kv_layout) + page_stride = manager.impl.get_page_stride(kv_local_id, Role.KEY) + key_base = manager.impl.get_mem_pool_base_address( + kv_local_id, Role.KEY, PageIndexMode.SHARED + ) + value_base = manager.impl.get_mem_pool_base_address( + kv_local_id, Role.VALUE, PageIndexMode.SHARED + ) + page_shape = (TOKENS_PER_BLOCK, 2, 8) if kv_layout == "NHD" else (2, TOKENS_PER_BLOCK, 8) + assert kv_buffer.shape[1:] == (2, *page_shape) + assert kv_buffer.stride(0) * kv_buffer.element_size() == 3 * page_stride + # Slot 1 begins three physical pages later, even though this layer + # exposes only two roles. Reversing the specs also exercises a base + # pointer offset within the coalesced slot. + assert kv_buffer[1, 0].data_ptr() == key_base + 3 * page_stride + assert kv_buffer[1, 1].data_ptr() == value_base + 3 * page_stride + assert key_buffer.stride(0) * key_buffer.element_size() == page_stride + + request_id = 7 + cache = manager._create_kv_cache(request_id, None, None) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(3 * TOKENS_PER_BLOCK, 0) + raw_slots = list(cache.get_base_page_indices(pool_id)[: cache.num_blocks]) + assert len(raw_slots) == 3 + assert any(slot % 2 == 1 for slot in raw_slots) + for layer_id, buffer, expected in ( + (key_layer_id, key_buffer, [3 * slot for slot in raw_slots]), + (kv_layer_id, kv_buffer, raw_slots), + ): + batch_indices = manager.get_batch_cache_indices([request_id], layer_idx=layer_id) + flat_indices = manager.get_batch_cache_indices_flat( + [request_id], [len(raw_slots)], layer_idx=layer_id + ).tolist() + assert batch_indices == [expected] + assert flat_indices == expected + local_id = manager.layer_offsets[layer_id] + for indices in (batch_indices[0], flat_indices): + for slot, index in zip(raw_slots, indices): + for role_index, role in enumerate(manager.get_layer_buffer_roles(layer_id)): + role_base = manager.impl.get_mem_pool_base_address( + local_id, role, PageIndexMode.SHARED + ) + assert ( + buffer[index, role_index].data_ptr() + == role_base + 3 * slot * page_stride + ) + + # Give each role and slot a distinct value. Writing an odd K/V slot + # must preserve the K-only page and both neighboring K/V slots. + for ordinal, slot in enumerate(raw_slots): + key_buffer[3 * slot, 0].fill_(10 + ordinal) + for ordinal, slot in enumerate(raw_slots): + kv_buffer[slot, 0].fill_(20 + ordinal) + kv_buffer[slot, 1].fill_(30 + ordinal) + for ordinal, slot in enumerate(raw_slots): + assert torch.all(key_buffer[3 * slot, 0] == 10 + ordinal) + assert torch.all(kv_buffer[slot, 0] == 20 + ordinal) + assert torch.all(kv_buffer[slot, 1] == 30 + ordinal) + finally: + manager.shutdown() + + def test_zero_size_filter_preserves_ssm_and_nonempty_extra_buffers() -> None: manager = object.__new__(KVCacheManagerV2) manager.pp_layers = [3, 4, 5] From 7cde28eb4d803af19f026af422679e6d7055ed6a Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 13 Oct 2026 17:32:28 -0700 Subject: [PATCH 11/23] fix: restore KV cache subclass and test compatibility Signed-off-by: Allison Lim --- .../attention/backends/sparse/minimax_m3/cache_manager.py | 3 +++ .../_torch/attention/sparse/glm_kpool/test_glm_kpool.py | 1 + .../sparse/msa/test_minimax_m3_shared_draft_layers.py | 2 ++ .../_torch/executor/kv_cache/test_kv_cache_estimation.py | 1 + .../_torch/executor/kv_cache/test_kvcm2_integration.py | 1 + .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 3 +++ .../unittest/disaggregated/test_minimax_m3_kv_transfer.py | 8 ++++++++ 7 files changed, 19 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 841bb0366237..f8f67585e1ee 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -877,6 +877,7 @@ def _get_batch_cache_indices_by_pool_id( is_kv_aggregate: bool = True, num_blocks_per_seq: Optional[Sequence[int]] = None, index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, ): """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1). @@ -905,6 +906,8 @@ def _get_batch_cache_indices_by_pool_id( index_scale: Kept for compatibility with the base virtual method; M3 bypasses the V1 block-id conversion entirely, so any caller-supplied scale is ignored alongside ``index_scales``. + kv_factor: Kept for compatibility with the base virtual method; + M3 returns raw slot IDs without KV aggregation. """ res = [] for req_idx, req_id in enumerate(request_ids): diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index e989db333636..e522fd422acf 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -146,6 +146,7 @@ def test_indexer_runtime_cache_cost_matches_registered_buffers(fp8, snapshot_int from tensorrt_llm.mapping import Mapping manager = object.__new__(Glm5NextCacheManager) + manager.layer_specs = None manager.mapping = Mapping(world_size=1, tp_size=1) manager.pp_layers = [0, 1] manager.layer_offsets = {0: 0, 1: 1} diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 3364da7089a5..b446432876b4 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -179,6 +179,7 @@ def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables( @pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) def test_block_scale_role_uses_global_layer_ids(dtype: DataType) -> None: manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None manager.dtype = dtype manager.pp_layers = [60, 7, 2] manager.sparse_layer_ids = {7} @@ -187,6 +188,7 @@ def test_block_scale_role_uses_global_layer_ids(dtype: DataType) -> None: assert manager._get_block_scale_role(Role.KEY, local_layer) == expected assert manager._get_block_scale_role(Role.VALUE, local_layer) is None generic = KVCacheManagerV2.__new__(KVCacheManagerV2) + generic.layer_specs = None generic.dtype = dtype expected = Role.KEY_BLOCK_SCALE if dtype == DataType.NVFP4 else None assert generic._get_block_scale_role(Role.KEY, 0) == expected diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index d706fb826687..2f9e3e00c9f4 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -997,6 +997,7 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = None manager._has_cp_helix = False manager.num_local_layers = 3 manager.pp_layers = [4, 5, 6] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 8a7974c8d20a..a795572d2f05 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -312,6 +312,7 @@ def _make_cache_config_for_test( assert len(max_attention_window_vec) == len(pp_layers) cache_manager = object.__new__(KVCacheManagerV2) + cache_manager.layer_specs = None cache_manager.max_beam_width = 1 cache_manager.kv_cache_type = kv_cache_type cache_manager.dtype = dtype diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 2bdea0b58426..284f5c552229 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -1955,6 +1955,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 @@ -1986,6 +1987,7 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 @@ -2214,6 +2216,7 @@ def test_expect_snapshot_points_binding_round_trip(): def test_v2_hybrid_pool_ratio_controls_allocated_memory(): def allocated_memory(pool_ratio): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index 844f4d0e1f18..ecf31101246b 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -70,6 +70,8 @@ def get_base_page_indices(_pool_id): def test_minimax_cache_indices_support_block_count_limit() -> None: manager = object.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None + manager.kv_factor = 2 manager.kv_cache_map = {7: _FakeKVCache(), 8: _ShortFakeKVCache()} assert manager._get_batch_cache_indices_by_pool_id([7]) == [[4, 5, 6, -1]] @@ -79,6 +81,12 @@ def test_minimax_cache_indices_support_block_count_limit() -> None: [4, 5, 6, -1], [7, 8, -1, -1], ] + assert manager._get_batch_cache_indices_by_pool_id([7], index_scale=3, kv_factor=2) == [ + [4, 5, 6, -1] + ] + assert manager.get_batch_cache_indices([7, 8]) == [[4, 5, 6, -1], [7, 8, -1, -1]] + assert manager.get_batch_cache_indices([7], num_blocks_per_seq=[2]) == [[4, 5]] + assert manager.get_batch_cache_indices([7], raw_indices=True) == [[4, 5, 6, -1]] assert manager.get_block_ids_per_seq([7]).tolist() == [[4, 5, 6, 0]] assert manager.get_block_ids_per_seq([7, 8]).tolist() == [ [4, 5, 6, 0], From 4a3bb25826dcb6fa1dc78cc9767b9cfdf754d970 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 14 Oct 2026 00:15:48 -0700 Subject: [PATCH 12/23] [None][test] fix mixed KV cache slot address arithmetic Signed-off-by: Allison Lim --- .../unittest/_torch/executor/kv_cache/test_kvcm2_integration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index a795572d2f05..376b501fcc1f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -706,7 +706,7 @@ def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layo cache = manager._create_kv_cache(request_id, None, None) assert cache.resume(manager._stream.cuda_stream) assert cache.resize(3 * TOKENS_PER_BLOCK, 0) - raw_slots = list(cache.get_base_page_indices(pool_id)[: cache.num_blocks]) + raw_slots = [int(slot) for slot in cache.get_base_page_indices(pool_id)[: cache.num_blocks]] assert len(raw_slots) == 3 assert any(slot % 2 == 1 for slot in raw_slots) for layer_id, buffer, expected in ( From 9d0107ba42b7379930de128b84d6cb45d9437a57 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Thu, 15 Oct 2026 16:18:33 -0700 Subject: [PATCH 13/23] [None][fix] align NVFP4 data and scale page layouts Separate NVFP4 data and block scales from other equal-size buffers while keeping shared lifecycle slot allocation. Preserve matching role order and require matching scale roles. Add one regression test for mixed NVFP4 and BF16 storage, checking scale addresses at slot 1 and the final slot. Signed-off-by: Allison Lim --- .../kv_cache_manager_v2/config.h | 4 ++ .../kv_cache_manager_v2/storage/config.cpp | 14 ++-- .../batch_manager/kvCacheManagerV2.cpp | 8 ++- .../kv_cache/kv_cache_manager_v2.py | 11 ++- .../runtime/kv_cache_manager_v2/__init__.pyi | 1 + .../kv_cache/test_kvcm2_integration.py | 67 +++++++++++++++++++ 6 files changed, 93 insertions(+), 12 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index a6cf8530f21b..a834025483b2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -107,6 +107,10 @@ struct BufferConfig std::optional tokensPerBlockOverride; bool isSparse = false; //!< Whether history uses the sparse attention lifecycle. + + //! Buffers may coalesce only when their group and expanded size match. + //! This does not change their lifecycle or shared slot allocation. + int coalescingGroup = 0; }; // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index c9e3718b14b9..a5d17a7817e1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -154,13 +154,13 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) LifeCycleRegistry registry{config}; int tokensPerBlock = config.tokensPerBlock; - // Map: lifeCycleId → (bufferSize → list of BufferId). - // Outer map key is LifeCycleId; inner map key is expanded buffer size. + // Map: lifeCycleId → ((coalescingGroup, expanded buffer size) → list of BufferId). // NOTE: Python uses insertion-ordered dict/defaultdict here. The std::map // grouping below intentionally remains sorted, which can make observable // pool-group indices/layout order differ from Python. Later lookups use // StorageConfig mappings, so this is not expected to affect correctness. - std::map>> bufferGroups; + using CoalescingKey = std::pair; + std::map>> bufferGroups; std::unordered_map expansionMap; for (auto const& layer : config.layers) @@ -179,7 +179,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) BufferId bid{cfg.layerId, buf.role}; expansionMap[bid] = exp; - bufferGroups[lcId][expandedSize].push_back(bid); + bufferGroups[lcId][{buf.coalescingGroup, expandedSize}].push_back(bid); } }, layer); @@ -189,14 +189,14 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) // Each variant: coalesced buffers sorted by size descending. std::vector slotGroups; slotGroups.reserve(bufferGroups.size()); - for (auto const& [lcId, sizeMap] : bufferGroups) + for (auto const& [lcId, coalescingMap] : bufferGroups) { SlotDescVariant var; var.lifeCycleId = lcId; - for (auto const& [sz, bufIds] : sizeMap) + for (auto const& [key, bufIds] : coalescingMap) { CoalescedBuffer cb; - cb.singleBufferSize = sz; + cb.singleBufferSize = key.second; cb.bufferIds = bufIds; var.coalescedBuffers.push_back(std::move(cb)); } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 7c4ddaae40b5..389085e7eeb0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1605,12 +1605,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("assert_valid", &kv::DiskCacheTierConfig::assertValid) DEF_COPY(kv::DiskCacheTierConfig); nb::class_(m, "BufferConfig") - .def(nb::init, bool>(), nb::arg("role"), nb::arg("size"), - nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false) + .def(nb::init, bool, int>(), nb::arg("role"), nb::arg("size"), + nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, + nb::arg("coalescing_group") = 0) .def_rw("role", &kv::BufferConfig::role) .def_rw("size", &kv::BufferConfig::size) .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) - .def_rw("is_sparse", &kv::BufferConfig::isSparse) DEF_COPY(kv::BufferConfig); + .def_rw("is_sparse", &kv::BufferConfig::isSparse) + .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) DEF_COPY(kv::BufferConfig); nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index ef85247b18e9..50a65ce7cc2e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -330,7 +330,7 @@ def __post_init__(self) -> None: expected_scale_roles = {Role.KEY_BLOCK_SCALE} if Role.VALUE in self.buffer_roles: expected_scale_roles.add(Role.VALUE_BLOCK_SCALE) - if not expected_scale_roles.issubset(self.buffer_roles): + if scale_roles.intersection(self.buffer_roles) != expected_scale_roles: raise ValueError( f"NVFP4 layer {self.global_layer_id} must contain matching block-scale roles" ) @@ -3209,6 +3209,7 @@ def _build_base_config( ) if self.layer_specs is not None: + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) layer_configs = [ AttentionLayerConfig( layer_id=LayerId(local_layer_id), @@ -3216,8 +3217,14 @@ def _build_base_config( BufferConfig( role=role, size=layer.bytes_per_token(role) * tokens_per_block, + # Keep NVFP4 data (1), its scales (2), and other data (0) + # separate while preserving matching K/V page order. + coalescing_group=( + 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 + ), ) - for role in layer.buffer_roles + for role in (Role.KEY, Role.VALUE, *scale_roles) + if role in layer.buffer_roles ], sliding_window_size=self.max_attention_window_vec[local_layer_id], num_sink_tokens=layer.num_sink_tokens, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 90be81694868..aafc31f26635 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -234,6 +234,7 @@ class BufferConfig: size: int tokens_per_block_override: int | None = None is_sparse: bool = False + coalescing_group: int = 0 class AttentionReusePolicy(enum.IntEnum): REQUIRED = 0 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 376b501fcc1f..c91bd7a3beec 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -746,6 +746,73 @@ def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layo manager.shutdown() +def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + nvfp4 = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.NVFP4, + num_kv_heads=4, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE), + ) + bf16 = replace( + nvfp4, + global_layer_id=1, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=1, + buffer_roles=(Role.KEY, Role.VALUE), + ) + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=4, + head_dim=128, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.NVFP4, + layer_specs=[nvfp4, bf16], + ) + try: + assert manager.num_pools == 1 + (group,) = manager.impl.pool_group_descs + assert len(group.pools) == 3 + assert group.num_slots > 1 + slots = [1, int(group.num_slots) - 1] + for offset, (data_role, scale_role) in enumerate( + ((Role.KEY, Role.KEY_BLOCK_SCALE), (Role.VALUE, Role.VALUE_BLOCK_SCALE)) + ): + data_converter = manager.impl.get_page_index_converter(0, data_role) + scale_converter = manager.impl.get_page_index_converter(0, scale_role) + pages = data_converter(slots, PageIndexMode.PER_LAYER) + assert pages == [2 * slot + offset for slot in slots] + assert scale_converter(slots, PageIndexMode.PER_LAYER) == pages + + scale_base = manager.impl.get_mem_pool_base_address( + 0, scale_role, PageIndexMode.PER_LAYER + ) + scale_pool = next(pool for pool in group.pools if pool.base_address == scale_base) + page_bytes = manager.impl.get_page_stride(0, scale_role) + # Attention reuses data indices for scales, including at the final slot. + for slot, page in zip(slots, pages): + address = scale_base + page * page_bytes + assert address == scale_base + slot * scale_pool.slot_bytes + offset * page_bytes + assert address + page_bytes <= scale_base + group.num_slots * scale_pool.slot_bytes + finally: + manager.shutdown() + + def test_zero_size_filter_preserves_ssm_and_nonempty_extra_buffers() -> None: manager = object.__new__(KVCacheManagerV2) manager.pp_layers = [3, 4, 5] From e273a984d9d0d0eaab42e087c8b4d17788f01a73 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Thu, 15 Oct 2026 19:57:39 -0700 Subject: [PATCH 14/23] [None][fix] accept mixed KV cache pointer table layouts Read the data-only view of three-dimensional pointer tables for non-FP4 layers while preserving NVFP4 scale handling and legacy two-dimensional tables. Add one native-binding regression covering mixed BF16/NVFP4 pointer tables and legacy BF16 tables. Signed-off-by: Allison Lim --- cpp/tensorrt_llm/thop/attentionOp.cpp | 9 ++-- .../kv_cache/test_kvcm2_integration.py | 51 +++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 4d1dd59cf305..f351dc779da6 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1659,14 +1659,17 @@ KvCachePoolPointers buildKvCachePoolPointers(at::Tensor const& hostKvCachePoolPo } else { - TORCH_CHECK(hostKvCachePoolPointers.dim() == 2); + TORCH_CHECK(hostKvCachePoolPointers.dim() == 2 || hostKvCachePoolPointers.dim() == 3); + // Non-FP4 layers use only the data pointers in a mixed-format pool table. + auto const dataPointers + = hostKvCachePoolPointers.dim() == 3 ? hostKvCachePoolPointers.select(2, 0) : hostKvCachePoolPointers; pointers.primaryPoolPtr = reinterpret_cast( reinterpret_cast( - readHostTensor2D(hostKvCachePoolPointers, poolIndex, 0, "host_kv_cache_pool_pointers")) + readHostTensor2D(dataPointers, poolIndex, 0, "host_kv_cache_pool_pointers")) + intraPoolOffset); pointers.secondaryPoolPtr = reinterpret_cast( reinterpret_cast( - readHostTensor2D(hostKvCachePoolPointers, poolIndex, 1, "host_kv_cache_pool_pointers")) + readHostTensor2D(dataPointers, poolIndex, 1, "host_kv_cache_pool_pointers")) + intraPoolOffset); } return pointers; diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index c91bd7a3beec..6ae92297182c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -43,6 +43,7 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE +from tensorrt_llm.bindings.internal import thop from tensorrt_llm.bindings.internal.batch_manager import CacheType, LinearCacheType from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( @@ -54,6 +55,7 @@ PARDDecodingConfig, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization import QuantMode from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, @@ -813,6 +815,55 @@ def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: manager.shutdown() +def test_mixed_nvfp4_bf16_pool_pointer_views() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + bf16 = torch.empty((2, 1, 32, 128), dtype=torch.bfloat16, device="cuda") + nvfp4 = torch.empty((2, 1, 32, 64), dtype=torch.uint8, device="cuda") + scales = torch.empty((2, 1, 32, 8), dtype=torch.float8_e4m3fn, device="cuda") + pointers = torch.tensor( + [ + [[nvfp4.data_ptr(), scales.data_ptr()], [0, 0]], + # A nonzero scale pointer must be ignored for the BF16 layer. + [[bf16.data_ptr(), scales.data_ptr()], [0, 0]], + ], + dtype=torch.int64, + device="cpu", + ) + mapping = torch.tensor([[0, 0], [1, 0]], dtype=torch.int32, device="cpu") + block_offsets = torch.zeros((2, 1, 2, 1), dtype=torch.int32, device="cuda") + for table, layer_id, quant_mode, expected_data, expected_scales in ( + (pointers, 1, 0, bf16, None), + (pointers, 0, int(QuantMode.NVFP4_KV_CACHE), nvfp4, scales), + (pointers[:, :, 0].contiguous(), 1, 0, bf16, None), + ): + data, blocks, block_scales = thop.build_trtllm_gen_kv_cache_metadata( + host_kv_cache_pool_pointers=table, + host_kv_cache_pool_mapping=mapping, + kv_cache_block_offsets=block_offsets, + layer_idx=layer_id, + num_kv_heads=1, + tokens_per_block=32, + head_dim=128, + kv_factor=2, + total_num_blocks=2, + kv_cache_quant_mode=quant_mode, + batch_start=0, + batch_size=1, + dtype=torch.bfloat16, + ) + assert data.data_ptr() == expected_data.data_ptr() + assert data.shape == expected_data.shape + assert data.dtype == expected_data.dtype + assert blocks.data_ptr() == block_offsets[layer_id].data_ptr() + if expected_scales is None: + assert block_scales is None + else: + assert block_scales is not None + assert block_scales.data_ptr() == expected_scales.data_ptr() + assert block_scales.shape == expected_scales.shape + + def test_zero_size_filter_preserves_ssm_and_nonempty_extra_buffers() -> None: manager = object.__new__(KVCacheManagerV2) manager.pp_layers = [3, 4, 5] From f6b1d08e2f27e62c5a45eae2952049bb829f102c Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Thu, 15 Oct 2026 22:12:28 -0700 Subject: [PATCH 15/23] [None][fix] preserve sink and reuse policy in KV cache costing Signed-off-by: Allison Lim --- .../kv_cache/kv_cache_manager_v2.py | 73 +++++++++++++++---- .../kv_cache/test_kv_cache_estimation.py | 72 +++++++++++++++++- 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 50a65ce7cc2e..ab04c0562d73 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -599,26 +599,34 @@ def _estimate_swa_cache_size( context: bool, scratch: bool, generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> tuple[int, int]: tokens_per_block = int(tokens_per_block) size_per_token = 0 size_per_request = 0 scratch_keys = set() - for layer_size, window_size in zip(layer_sizes, attention_windows): + for layer_idx, (layer_size, window_size) in enumerate(zip(layer_sizes, attention_windows)): if window_size is not None and window_size > 0: + layer = layer_specs[layer_idx] if layer_specs is not None else None + sink_blocks = math.ceil((layer.num_sink_tokens or 0) / tokens_per_block) if layer else 0 # Match AttnLifeCycle.get_stale_range(): the live interval contains # window_size + generation_capacity_headroom - 1 tokens. Across all # page offsets, that interval touches at most the count below. window_blocks = ( math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 ) - window_tokens = window_blocks * tokens_per_block + window_tokens = (sink_blocks + window_blocks) * tokens_per_block if not context: size_per_request += window_tokens * layer_size elif not scratch: size_per_token += layer_size else: - scratch_key = (int(window_size), layer_size) + scratch_key = ( + int(window_size), + layer_size, + sink_blocks, + layer.reuse_policy if layer else None, + ) if scratch_key in scratch_keys: size_per_request += window_tokens * layer_size else: @@ -655,6 +663,7 @@ def _estimate_cache_size_components( *, scratch: bool, generation_capacity_headroom: int, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> tuple[int, int, int]: """Return context/generation bytes per token and generation bytes per request. @@ -664,7 +673,12 @@ def _estimate_cache_size_components( """ full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows) context_swa_size, _ = _estimate_swa_cache_size( - layer_sizes, attention_windows, tokens_per_block, context=True, scratch=scratch + layer_sizes, + attention_windows, + tokens_per_block, + context=True, + scratch=scratch, + layer_specs=layer_specs, ) generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size( layer_sizes, @@ -673,6 +687,7 @@ def _estimate_cache_size_components( context=False, scratch=False, generation_capacity_headroom=generation_capacity_headroom, + layer_specs=layer_specs, ) return ( full_attn_size + context_swa_size, @@ -2729,6 +2744,7 @@ def _get_max_tokens_from_quota_impl(self, quota: int) -> float: self.tokens_per_block, scratch=self.enable_swa_scratch_reuse, generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, ) size_per_batch = self._get_generation_request_capacity() * generation_swa_size_per_request if quota < size_per_batch: @@ -2765,6 +2781,7 @@ def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: self.tokens_per_block, scratch=self.enable_swa_scratch_reuse, generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, ) context_tokens = min(max_tokens, self.max_num_tokens) generation_tokens = max_tokens - context_tokens @@ -6409,6 +6426,7 @@ def get_cache_size_per_token( and not any(spec.owner == KVCacheLayerOwner.DRAFT for spec in (layer_specs or ())) ), generation_capacity_headroom=generation_capacity_headroom, + layer_specs=layer_specs, ) # The affine slope covers all tokens; context additionally retains SWA # pages for the current token batch beyond the generation windows. @@ -6419,16 +6437,43 @@ def get_cache_size_per_token( bytes_per_slot = _get_single_swa_pool_slot_bytes( layer_sizes, attention_windows, tokens_per_block ) - if is_draft and fixed_cost > 0 and bytes_per_slot is not None: + if is_draft and fixed_cost > 0: + costs = [(fixed_cost, bytes_per_slot)] + if layer_specs is not None: + # Explicit lifecycles can have different slot counts. Normalize + # each layer separately so pool rounding cannot hide a deficit. + costs = [] + for layer, layer_size, window in zip(layer_specs, layer_sizes, attention_windows): + if window is None or window <= 0: + continue + _, retained_bytes = _estimate_swa_cache_size( + [layer_size], + [window], + tokens_per_block, + context=False, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + layer_specs=[layer], + ) + costs.append( + ( + layer_size * max_num_tokens + retained_bytes * max_batch_size, + layer_size * tokens_per_block, + ) + ) # Without a config, all windows are full attention, so no SWA slot exists. assert kv_cache_config is not None # The affine intercept is the configured quota needed to preserve # the fixed usable capacity at V2's resume watermark. Keep the # allocator's page rounding inside the manager estimator rather # than extending the generic CacheCost model with pool geometry. - required_slots = math.ceil(fixed_cost / bytes_per_slot) resume_util = float(np.float32(kv_cache_config.max_util_for_resume)) - fixed_cost = math.ceil(required_slots / resume_util) * bytes_per_slot + fixed_cost = sum( + cost + if slot_bytes is None + else math.ceil(math.ceil(cost / slot_bytes) / resume_util) * slot_bytes + for cost, slot_bytes in costs + ) return ( cache_size_per_token, fixed_cost, @@ -6557,7 +6602,7 @@ def _get_per_request_layer_tokens( * tokens_per_block ) - def window_tokens(window_size: Optional[int]) -> int: + def window_tokens(window_size: Optional[int], sink_tokens: int) -> int: if window_size is None or window_size <= 0 or window_size >= max_seq_len: return full_tokens # Match _estimate_swa_cache_size: the live interval holds @@ -6565,15 +6610,17 @@ def window_tokens(window_size: Optional[int]) -> int: window_blocks = ( math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 ) - return min(window_blocks * tokens_per_block, full_tokens) + sink_blocks = math.ceil(sink_tokens / tokens_per_block) + return min((sink_blocks + window_blocks) * tokens_per_block, full_tokens) entries = [] - for layer_size, retained_window, covered_window in zip( - layer_sizes, retained_windows, covered_windows + for layer_idx, (layer_size, retained_window, covered_window) in enumerate( + zip(layer_sizes, retained_windows, covered_windows) ): - retained = window_tokens(retained_window) + sink_tokens = (layer_specs[layer_idx].num_sink_tokens or 0) if layer_specs else 0 + retained = window_tokens(retained_window, sink_tokens) covered = ( - window_tokens(covered_window) + window_tokens(covered_window, sink_tokens) if covered_window is not None and 0 < covered_window < max_seq_len else 0 ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 2f9e3e00c9f4..739ca90d6d20 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -10,7 +10,7 @@ share, not all copies. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from unittest.mock import Mock, patch @@ -30,10 +30,16 @@ get_mla_context_workspace_kv_len_cap, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( KvCacheConfig, @@ -42,6 +48,7 @@ TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy pytestmark = pytest.mark.cpu_only @@ -995,6 +1002,67 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): assert cost == CacheCost(slope=20, intercept=21) +def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: + windowed = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=128, + num_sink_tokens=65, + ) + layers = [ + windowed, + replace(windowed, global_layer_id=1, reuse_policy=AttentionReusePolicy.PRIVATE), + replace(windowed, global_layer_id=2, num_sink_tokens=None), + replace(windowed, global_layer_id=3, sliding_window_size=None, num_sink_tokens=None), + ] + config = KvCacheConfig(enable_block_reuse=False, enable_swa_scratch_reuse=True) + kwargs = dict( + tokens_per_block=64, + max_seq_len=4096, + kv_cache_config=config, + layer_specs=layers, + ) + slope, fixed = KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=256, **kwargs + ) + # Each layer costs 64 bytes/token. Different sinks or policies cannot share scratch. + # The SWA layers retain 5, 5, and 3 pages, including two pages for 65 sink tokens. + assert slope == 64 + assert fixed == 256 * 192 + 2 * (320 + 320 + 192) * 64 + + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = tuple(layers) + manager._has_cp_helix = False + manager.max_attention_window_vec = [128, 128, 128, None] + manager.tokens_per_block = 64 + manager.max_batch_size = 2 + manager.max_num_tokens = 256 + manager._generation_kv_capacity_headroom = 1 + manager.enable_swa_scratch_reuse = True + quota = manager._get_quota_from_max_tokens(512) + assert quota == slope * 512 + fixed + assert manager._get_max_tokens_from_quota(quota) == 512 + assert KVCacheManagerV2._get_per_request_layer_tokens(None, Mapping(), **kwargs) == ( + [(64, 320, 320), (64, 320, 320), (64, 192, 192), (64, 4160, 0)], + 4160, + ) + + # Draft lifecycles need 7 and 21 usable slots: separate resume rounding + # gives 8 + 23 slots; rounding their average together would give only 30. + draft = replace(windowed, owner=KVCacheLayerOwner.DRAFT, num_sink_tokens=0) + kwargs.update( + layer_specs=[draft, replace(draft, global_layer_id=1, num_sink_tokens=448)], + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_util_for_resume=0.95), + ) + assert KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=64, is_draft=True, **kwargs + ) == (0, (8 + 23) * 64 * 64) + + def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) manager.layer_specs = None From 636715a1b7612c83c4cf56806da1575e847f3499 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 16 Oct 2026 09:19:42 -0700 Subject: [PATCH 16/23] fix: initialize layer specs in FP4 MLA sizing fixture Signed-off-by: Allison Lim --- tests/unittest/_torch/attention/test_fp4_mla.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 4d48ceaee9cc..c94c6e2e69fc 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -621,6 +621,7 @@ def test_fp4_mla_v2_cache_size_accounts_for_bf16_fallback( def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) + manager.layer_specs = None manager.max_batch_size = 3 manager.mapping = SimpleNamespace(pp_size=2) manager.max_num_tokens = 100 From 77d4eae1a3e15c8e794f2e7a6ea89feff6c86f36 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 17 Oct 2026 21:27:19 -0700 Subject: [PATCH 17/23] [None][fix] Release DSA test cache managers with metadata Pass the cache manager explicitly to MockMetadata so the local class does not retain it through the constructor closure. This avoids deferring CUDA resource destruction to cyclic garbage collection. Signed-off-by: Allison Lim --- .../_torch/attention/sparse/dsa/test_dsa_indexer.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 18880ca96dcc..e6d5a3f38121 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -1774,7 +1774,7 @@ def __init__(self): class MockMetadata(DSAtrtllmAttentionMetadata): """Mock DSA attention metadata for unit testing indexer operations.""" - def __init__(self): + def __init__(self, kv_cache_manager: DSACacheManager | DSACacheManagerV2) -> None: """Initialize mock metadata with buffers, block tables, and indexer state.""" self.num_sms = deep_gemm.get_num_sms() self.max_num_tokens = max_num_tokens @@ -1796,7 +1796,9 @@ def __init__(self): self.seq_lens = seq_lens.cpu() if seq_lens.is_cuda else seq_lens self.kv_lens = kv_lens self.kv_cache_params = MockKVCacheParams() - self.kv_cache_manager = cache_manager + # Keep the manager out of the local class's reference cycle so its + # CUDA resources are not left for garbage collection on another thread. + self.kv_cache_manager = kv_cache_manager # Effective tokens-per-block for the indexer k-cache slot mapping, # mirroring DSAtrtllmAttentionMetadata.__post_init__ (dsa.py). tpb = self.kv_cache_manager.tokens_per_block @@ -1806,7 +1808,7 @@ def __init__(self): self._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() self.indexer_k_cache_block_offsets = torch.zeros( - [batch_size, cache_manager.max_blocks_per_seq], + [batch_size, self.kv_cache_manager.max_blocks_per_seq], device="cuda", dtype=torch.int32, ) @@ -1815,7 +1817,7 @@ def __init__(self): device="cpu", pin_memory=True, ) - block_ids = cache_manager.get_batch_cache_indices(request_ids) + block_ids = self.kv_cache_manager.get_batch_cache_indices(request_ids) for i in range(len(block_ids)): self.host_indexer_k_cache_block_offsets[i, : len(block_ids[i])] = torch.tensor( block_ids[i], dtype=torch.int32 @@ -2014,7 +2016,7 @@ def num_seqs(self) -> int: """ return self._num_seqs - return MockMetadata() + return MockMetadata(cache_manager) @pytest.mark.parametrize("disabled_ratio", [0, 1]) From 61248ec36110e8077a0e605a2be292ad1acfb633 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 09:33:55 -0700 Subject: [PATCH 18/23] [None][fix] Reject odd NVFP4 layer-spec head dimensions Preserve the legacy NVFP4 packing constraint for explicit layer specs. Add a CPU regression for sixteen heads with dimension one, which passes total-element divisibility but cannot form a per-head packed buffer view. Signed-off-by: Allison Lim --- .../pyexecutor/kv_cache/kv_cache_manager_v2.py | 2 ++ .../executor/kv_cache/test_kvcm2_integration.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 5e0756a7d80d..235ea84e6c4d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -326,6 +326,8 @@ def __post_init__(self) -> None: if any(role in scale_roles for role in self.buffer_roles) and self.dtype != DataType.NVFP4: raise ValueError("block-scale roles require NVFP4 KV-cache dtype") if self.dtype == DataType.NVFP4: + if self.head_dim % 2 != 0: + raise ValueError(f"NVFP4 head_dim must be even, got {self.head_dim}") expected_scale_roles = {Role.KEY_BLOCK_SCALE} if Role.VALUE in self.buffer_roles: expected_scale_roles.add(Role.VALUE_BLOCK_SCALE) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index c8bc40f2f163..12419c1ef216 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -84,6 +84,19 @@ MAX_SEQ_LEN = 16 +@pytest.mark.cpu_only +def test_nvfp4_layer_spec_rejects_odd_head_dim() -> None: + with pytest.raises(ValueError, match="NVFP4 head_dim must be even"): + KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.NVFP4, + num_kv_heads=16, + head_dim=1, + buffer_roles=(Role.KEY, Role.KEY_BLOCK_SCALE), + ) + + @pytest.mark.cpu_only @pytest.mark.parametrize( "prompt,rewind,ordinal,window,sinks,expected", From d1ab5712030ce6f91cffd814817363442c57693a Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 23 Oct 2026 14:18:16 -0700 Subject: [PATCH 19/23] fix: reject draft-owned layer specs with Helix context parallelism Signed-off-by: Allison Lim --- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 1d49a4d85ce4..dd10838ff33d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1574,10 +1574,12 @@ def __init__( "all CP ranks, so publishing it requires a group-wide " "protocol. Set kv_cache_config.enable_block_reuse=False." ) - if is_draft: + if is_draft or any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs or () + ): raise ValueError( "KVCacheManagerV2 does not support a draft cache " - "manager with helix context parallelism." + "manager or draft-owned layers with helix context parallelism." ) if mapping.enable_attention_dp: raise ValueError( From 62061183d87f910daf39c0cd0a25d0867d7613ed Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 23 Oct 2026 14:19:19 -0700 Subject: [PATCH 20/23] fix: reject layer specs for cache managers with extra buffers Signed-off-by: Allison Lim --- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index dd10838ff33d..f3d0fad68975 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1455,6 +1455,14 @@ def __init__( self.dtype = dtype self._declared_target_layer_count = num_layers self.layer_specs = tuple(layer_specs) if layer_specs is not None else None + if ( + self.layer_specs is not None + and type(self)._extra_buffers_per_layer is not KVCacheManagerV2._extra_buffers_per_layer + ): + raise ValueError( + f"{type(self).__name__} does not support layer_specs with " + "an overridden _extra_buffers_per_layer." + ) self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager From b069cdf5785146861c980f0c8f975de3711b02aa Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 23 Oct 2026 14:47:47 -0700 Subject: [PATCH 21/23] fix: initialize layer specs in KDA draft record test fixture Signed-off-by: Allison Lim --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 07c3bfd778b8..f3cfcb9c16d6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3905,6 +3905,7 @@ def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention layer.""" mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF From eadaf057d2a48e84e4e2ae664e01c35c715e9b3a Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 23 Oct 2026 16:07:44 -0700 Subject: [PATCH 22/23] fix: count explicit KV cache capacity from allocator slots Derive warmup capacity from available GPU slots and per-pool lifecycle demand for explicit layer specs. Preserve the legacy page-index calculation. Signed-off-by: Allison Lim --- .../kv_cache/kv_cache_manager_v2.py | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 8dc019829676..bc870400a56e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -4123,18 +4123,33 @@ def get_num_available_tokens( return clamped def get_num_free_blocks(self) -> int: - # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks. - # However, since we only use this function when the kv cache manager is empty, so it is safe to do so. - # - # "Empty" means no request holds a cache. The guard page, when it is - # enabled, is held by a permanent reservation made at construction, not - # by a request, so it is excluded from the check and subtracted from - # the count: a caller sizing itself off this number must not plan to - # use the page the guard owns. + """Estimate warmup block capacity when only permanent reservations remain. + + Explicit specs count independent one-block sequences across all required + GPU pools. Longer sequences and resume limits are checked during allocation. + """ reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS assert not set(self.kv_cache_map) - reserved, ( "get_num_free_blocks is only used when the kv cache manager is empty" ) + if self.layer_specs is not None: + # Each lifecycle needs one slot per block. Coalesced layers share a + # lifecycle, while separate lifecycles can compete for the same pool. + slots_per_block: dict[int, int] = defaultdict(int) + for pool_id in self.impl.get_life_cycle_pool_group_indices(GPU_LEVEL): + slots_per_block[int(pool_id)] += 1 + + # Available slots already exclude the guard's permanent reservation. + stats = self._get_storage_statistics(GPU_LEVEL) + return min( + ( + int(stats[pool_id].available) // slot_count + for pool_id, slot_count in slots_per_block.items() + ), + default=0, + ) + + # Legacy managers estimate primary-pool capacity from page indices. max_num_pages = max( [ self.impl.get_page_index_upper_bound(layer_id, Role.KEY) From ded3a12ed6c0fce16aa7128d3d4d253447a5c20c Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 23 Oct 2026 16:17:45 -0700 Subject: [PATCH 23/23] fix: use explicit KV geometry for accepted-token relocation Derive relocation head bytes from explicit layer dimensions and dtype. Reject layouts that do not match the native uniform contiguous K/V representation, while preserving the legacy calculation. Signed-off-by: Allison Lim --- .../kv_cache/kv_cache_manager_v2.py | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index bc870400a56e..54b9deae7fc3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1260,6 +1260,31 @@ def _update_kv_cache_draft_token_location( "update_kv_cache_draft_token_location requires all local layers " f"in one KV pool, got pools {sorted(local_pool_ids)}" ) + layer_specs = getattr(cache_manager, "layer_specs", None) + if layer_specs is not None: + if len({(layer.head_dim, layer.dtype) for layer in layer_specs}) != 1: + raise ValueError( + "update_kv_cache_draft_token_location requires uniform head_dim and dtype " + "for explicit layer specs" + ) + if any(set(layer.buffer_roles) != {Role.KEY, Role.VALUE} for layer in layer_specs): + raise ValueError( + "update_kv_cache_draft_token_location requires explicit layer specs " + "with only KEY and VALUE buffers" + ) + # The kernel advances each layer base by one contiguous K/V pair. + if any( + int(cache_manager.kv_cache_pool_mapping[layer_idx][1]) != layer_idx + for layer_idx in range(cache_manager.num_local_layers) + ): + raise ValueError( + "update_kv_cache_draft_token_location requires explicit layers " + "in contiguous local-layer order within the KV pool" + ) + layer = layer_specs[0] + head_size_in_bytes = get_size_in_bytes(layer.head_dim, layer.dtype) + else: + head_size_in_bytes = int(cache_manager.head_dim * kv_cache_dtype_byte_size) pool_idx = local_pool_ids.pop() torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( accepted_draft_token_offsets, @@ -1270,7 +1295,7 @@ def _update_kv_cache_draft_token_location( # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded # total so the C++ kernel computes correct strides and grid dims. cache_manager.num_kv_heads_per_layer[0], - int(cache_manager.head_dim * kv_cache_dtype_byte_size), + head_size_in_bytes, cache_manager.max_total_draft_tokens, max_attention_window, rewind_draft_token_separate_adjustments,