From 3bc36a5565adbb54b413aa023049f6b48f055087 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 12 Oct 2026 16:01:23 -0700 Subject: [PATCH 01/35] [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/35] [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/35] [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/35] [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/35] [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/35] 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/35] 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/35] 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/35] 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/35] 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/35] 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/35] [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/35] [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/35] [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/35] [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 b6db612772a44680518cd733ee69569aaa8fe068 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 16 Oct 2026 13:04:52 -0700 Subject: [PATCH 16/35] [TRTLLM-20434][feat] Dflash Unified KVCache Share one KVCacheManagerV2 between compatible DFlash target and draft layers in aggregated and disaggregated serving. Preserve each layer's cache dtype and head dimension, support attention data parallelism, and retain the separate-cache fallback for unsupported layouts. Transfer unified draft layers with target pages, validate peer draft-layer counts, and add unit and B200 integration coverage. Signed-off-by: Allison Lim --- .../_torch/disaggregation/native/peer.py | 8 + .../_torch/disaggregation/native/rank_info.py | 6 + tensorrt_llm/_torch/pyexecutor/_util.py | 204 +++++++++++++++++- .../pyexecutor/scheduler/scheduler_v2.py | 7 +- tensorrt_llm/_torch/speculative/dflash.py | 108 ++++++++-- tensorrt_llm/_torch/speculative/interface.py | 10 +- tensorrt_llm/_torch/speculative/utils.py | 4 +- .../references/acceptance_length.yaml | 6 + .../accuracy/test_disaggregated_serving.py | 91 +++++++- .../test_lists/test-db/l0_dgx_b200.yml | 1 + .../kv_cache/test_dflash_unified_kv_gate.py | 139 ++++++++++++ tests/unittest/disaggregated/test_peer.py | 26 +++ .../unittest/disaggregated/test_rank_info.py | 62 ++++++ 13 files changed, 643 insertions(+), 29 deletions(-) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..6b98a2772ab1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -180,6 +180,14 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + if self._ri.num_draft_layers != peer_ri.num_draft_layers: + # The layer-count tolerance below would transfer only the target + # layers and leave the drafter's prompt KV unset on the receiver. + raise ValueError( + "PeerRegistrar: unified draft KV layer count differs " + f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " + "both peers must use the same speculative decoding KV cache layout." + ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 87731d303204..358b1b576412 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -49,6 +49,9 @@ class RankInfo: aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None bounded_replay_on_generation: bool = False + # Layers a unified KV cache manager owns for the drafter (layer specs with + # a DRAFT owner). They ride the target's pages, so both peers must agree. + num_draft_layers: int = 0 @property def tp_size_per_dp_group(self) -> int: @@ -114,6 +117,9 @@ def from_kv_cache_manager( aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False), + num_draft_layers=( + len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0 + ), ) @classmethod diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index bcfb7529311b..96a160e78670 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -40,6 +40,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy from ..attention.backends import get_sparse_attn_kv_cache_manager from ..attention.backends.sparse.params import \ @@ -65,7 +66,9 @@ from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder -from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache.kv_cache_manager_v2 import (KVCacheLayerOwner, KVCacheLayerSpec, + KVCacheManagerV2, Role, + _resolve_v2_max_attention_window_vec) from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, MambaHybridCacheManagerV2, @@ -1047,6 +1050,11 @@ def _per_manager_cache_cost(self, **extra_kwargs) -> CacheCost: kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) + if not is_draft and self._should_unify_dflash_kv_cache(): + extra_kwargs = dict(extra_kwargs) + extra_kwargs[ + "layer_specs"] = self._build_unified_dflash_layer_specs( + kv_cache_config) if not is_draft: derived_windows = _derive_v2_layer_type_attention_windows( kv_cache_config, manager_cls, model_config, self._max_seq_len) @@ -1972,6 +1980,14 @@ def _create_kv_cache_manager( spec_dec_layer_mask = [True] * num_target_layers estimating_kv_cache = estimating_kv_cache and not self._skip_est + unified_layer_specs = None + if self._should_unify_dflash_kv_cache(): + unified_layer_specs = self._build_unified_dflash_layer_specs( + kv_cache_config) + logger.info( + "DFlash: target and draft KV layers share one KV cache manager " + f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})." + ) kv_cache_manager = _create_kv_cache_manager( model_engine=model_engine, kv_cache_manager_cls=kv_cache_manager_cls, @@ -1996,6 +2012,7 @@ def _create_kv_cache_manager( self._llm_args.kv_cache_config.kv_events_config, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, + layer_specs=unified_layer_specs, ) if not self._skip_est: @@ -2035,6 +2052,8 @@ def _should_create_separate_draft_kv_cache(self) -> bool: if self._speculative_config is None: # No drafter at all, so there is nothing to give a manager to. return False + if self._should_unify_dflash_kv_cache(): + return False # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL # never reach the arena this carve-out exists for. spec_dec_mode = self._speculative_config.spec_dec_mode @@ -2060,6 +2079,180 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) + def _should_unify_dflash_kv_cache(self) -> bool: + """Whether this run can place target and DFlash layers in one V2 manager.""" + spec_config = self._speculative_config + if (spec_config is None or not spec_config.spec_dec_mode.is_dflash() + or spec_config.attention_backend != "TRTLLM" + or not spec_config._allow_separate_draft_kv_cache + or spec_config.draft_len_schedule is not None + or self._draft_model_engine is not None + or self._draft_config is None or self._mapping.pp_size != 1 + or self._mapping.cp_size != 1 + or self._llm_args.attn_backend != "TRTLLM" or is_gemma4_hybrid( + self._model_engine.model.model_config.pretrained_config) + or self._llm_args.kv_cache_compression_config is not None): + return False + # Specialized managers add cache roles and lifecycle protocols that an + # external drafter cannot yet describe. Keep their existing fallback. + if self._kv_cache_manager_cls is not KVCacheManagerV2: + return False + target_layout = self._dflash_kv_layout( + self._model_engine.model.model_config, self._model_engine.dtype) + draft_model_config = self._get_draft_kv_model_config() + draft_layout = self._dflash_kv_layout( + draft_model_config, + draft_model_config.pretrained_config.torch_dtype) + if target_layout is None or draft_layout is None: + logger.info_once( + "DFlash: a per-layer KV-head count is not supported in a shared " + "target and draft KV cache; keeping a separate draft KV cache.", + key="dflash_per_layer_kv_heads_fallback") + return False + if self._is_disagg and not self._dflash_disagg_layouts_compatible( + target_layout, draft_layout): + logger.info_once( + "DFlash: draft KV layout (kv_heads, head_dim, dtype, k_only)=" + f"{draft_layout} cannot be transferred with the target's " + f"{target_layout}; disaggregated serving keeps a separate " + "draft KV cache.", + key="dflash_disagg_draft_layout_fallback") + return False + return True + + @staticmethod + def _dflash_disagg_layouts_compatible( + target_layout: tuple[int, int, object, bool], + draft_layout: tuple[int, int, object, bool]) -> bool: + """Whether the transceiver can move draft layers with the target's. + + Heads are resharded in bytes: a view's bytes per head are its own + per-layer bytes over the rank's KV-head count, so dtype and head_dim + may differ. The KV-head count and K-only flag come from the rank's + single AttentionInfo and must match. NVFP4 adds block-scale views, so + an NVFP4 side is only paired with an identical layout. + """ + target_heads, _, target_dtype, target_k_only = target_layout + draft_heads, _, draft_dtype, draft_k_only = draft_layout + if (target_heads, target_k_only) != (draft_heads, draft_k_only): + return False + nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 + return target_dtype == draft_dtype or nvfp4 not in (target_dtype, + draft_dtype) + + def _dflash_kv_layout( + self, model_config: ModelConfig, + fallback_dtype) -> Optional[tuple[int, int, object, bool]]: + """Rank-local (kv_heads, head_dim, dtype, k_only) of a model's KV layers. + + Returns None when the KV-head count varies per layer. + """ + config = model_config.pretrained_config + k_only = is_mla(config) + heads = 1 if k_only else getattr(config, "num_key_value_heads", + config.num_attention_heads) + if not isinstance(heads, int): + return None + head_dim = (config.kv_lora_rank + config.qk_rope_head_dim + if k_only else getattr(config, "head_dim", None)) + if not isinstance(head_dim, int): + head_dim = config.hidden_size // config.num_attention_heads + dtype = self._layer_spec_dtype(model_config, fallback_dtype) + # Attention DP replicates attention weights, so every rank holds all heads. + tp_size = 1 if self._mapping.enable_attention_dp else self._mapping.tp_size + return ceil_div(heads, tp_size), head_dim, dtype, k_only + + @staticmethod + def _layer_spec_dtype(model_config: ModelConfig, fallback_dtype) -> object: + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( + ): + return tensorrt_llm.bindings.DataType.FP8 + if quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache( + ): + return tensorrt_llm.bindings.DataType.NVFP4 + return str_dtype_to_binding(torch_dtype_to_str(fallback_dtype)) + + @staticmethod + def _layer_spec_roles(dtype, *, k_only: bool) -> tuple: + roles = [Role.KEY] + if not k_only: + roles.append(Role.VALUE) + if dtype == tensorrt_llm.bindings.DataType.NVFP4: + roles.append(Role.KEY_BLOCK_SCALE) + if not k_only: + roles.append(Role.VALUE_BLOCK_SCALE) + return tuple(roles) + + def _build_unified_dflash_layer_specs( + self, + kv_cache_config: KvCacheConfig) -> tuple[KVCacheLayerSpec, ...]: + """Build the rank-local target+DFlash storage layout used by cost and allocation.""" + if not self._should_unify_dflash_kv_cache(): + return () + + target_model_config = self._model_engine.model.model_config + target_config = target_model_config.pretrained_config + target_layers = int(target_config.num_hidden_layers) + target_local_ids = list(self._mapping.pp_layers(target_layers)) + target_heads, target_head_dim, target_dtype, target_k_only = ( + self._dflash_kv_layout(target_model_config, + self._model_engine.dtype)) + target_window_config = kv_cache_config.max_attention_window + if target_window_config is None: + target_window_config = _derive_v2_layer_type_attention_windows( + kv_cache_config, + self._kv_cache_manager_cls, + target_model_config, + self._max_seq_len, + ) + target_windows = _resolve_v2_max_attention_window_vec( + target_window_config, + self._max_seq_len, + target_local_ids, + ) + + draft_model_config = self._get_draft_kv_model_config() + draft_config = draft_model_config.pretrained_config + draft_layers = int(draft_config.num_hidden_layers) + draft_local_ids = list(self._mapping.pp_layers(draft_layers)) + draft_heads, draft_head_dim, draft_dtype, draft_k_only = ( + self._dflash_kv_layout(draft_model_config, + draft_config.torch_dtype)) + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._max_seq_len, + kv_cache_config), + self._max_seq_len, + draft_local_ids, + ) + + specs = [ + KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=target_dtype, + num_kv_heads=target_heads, + head_dim=target_head_dim, + buffer_roles=self._layer_spec_roles(target_dtype, + k_only=target_k_only), + sliding_window_size=window, + reuse_policy=AttentionReusePolicy.REQUIRED, + ) for layer_id, window in zip(target_local_ids, target_windows) + ] + specs.extend( + KVCacheLayerSpec( + global_layer_id=target_layers + layer_id, + owner=KVCacheLayerOwner.DRAFT, + dtype=draft_dtype, + num_kv_heads=draft_heads, + head_dim=draft_head_dim, + buffer_roles=self._layer_spec_roles(draft_dtype, + k_only=draft_k_only), + sliding_window_size=window, + reuse_policy=AttentionReusePolicy.REQUIRED, + ) for layer_id, window in zip(draft_local_ids, draft_windows)) + return tuple(specs) + def _joint_reuse_supported(self) -> bool: """Span half of the pairing decision; ``build_managers`` ANDs it with ``has_separate_one_model_draft``. ``False`` = leave the run unpaired. @@ -2121,12 +2314,12 @@ def _get_draft_kv_model_config(self) -> ModelConfig: if not (spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark()): return effective_draft_config quant_config = getattr(effective_draft_config, "quant_config", None) - if quant_config is None or not quant_config.quant_mode.has_fp8_kv_cache( + if quant_config is None or not quant_config.quant_mode.has_kv_cache_quant( ): return effective_draft_config logger.info( "External drafter KV pool keeps the drafter dtype; dropping " - "the fp8 KV quant algo inherited from the target.") + "the KV quantization inherited from the target.") neutral_quant = copy.copy(quant_config) neutral_quant.kv_cache_quant_algo = None # QuantConfig.quant_mode and .layer_quant_mode are both cached_property @@ -3107,7 +3300,9 @@ def _create_kv_cache_manager( cold_page_codec_provider: Optional[object] = None, kv_events_config: Optional[KVEventsConfig] = None, joint_kv_cache_reuse: bool = False, - max_cuda_graph_batch_size: Optional[int] = None) -> KVCacheManager: + max_cuda_graph_batch_size: Optional[int] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None +) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -3282,6 +3477,7 @@ def _create_kv_cache_manager( manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse manager_extra_kwargs[ "disable_overlap_scheduler"] = disable_overlap_scheduler + manager_extra_kwargs["layer_specs"] = layer_specs # Vocab size also enables multimodal event decoding and its per-block # digest scan. Leave it unset for text-only engines. Without an engine # (e.g. separate one-model draft caches), preserve config resolution: diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 3b7e55c8fe33..0965edfb8dd2 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -837,9 +837,10 @@ def _try_schedule_disagg_gen_init( Returns ``(action, tokens)``. *tokens* is 0 because disagg requests don't participate in the forward pass token budget. """ - # Cache-transceiver mode disables the separate one-model draft manager, - # so disagg generation init has no paired-reuse path. Supporting one - # would also require draft KV transfer and history_length=prompt_len. + # Only the target manager receives prompt KV. Drafters that need it + # under disagg keep their layers in that manager (MTP/Eagle fold them + # in; DFlash uses explicit draft layer specs). A separate draft manager + # gets no transfer and no paired-reuse path here. if not self.kv_cache_manager.prepare_disagg_gen_init(req): logger.debug("prepare_disagg_gen_init failed for request %s", req.py_request_id) return ScheduleAction.SKIP, 0 diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 88da4911cca0..a34b49a35b2b 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -27,6 +27,7 @@ from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata +from ..pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from ..pyexecutor.resource_manager import BaseResourceManager @@ -751,6 +752,9 @@ def __init__( # the table width on the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 + self._ctx_layer_ids = None + self._ctx_unified_cache = False + self._ctx_block_offset = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None self._ctx_kv_last_page_len = None @@ -927,7 +931,18 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor # what DSpark is, so this drafter does get a manager under DP. The # reachable cases are the two-model paths, which never build one. return None - layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)] + draft_layer_ids = tuple(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) + layer_ids = draft_layer_ids if draft_layer_ids else tuple(range(L)) + if len(layer_ids) != L: + logger.warning( + f"DFlash: KV cache manager exposes {len(layer_ids)} draft layers " + f"but the drafter has {L}; using the private ctx arena." + ) + return None + self._ctx_layer_ids = layer_ids + layers = [ + draft_kv_cache_manager.get_buffers(layer_id, kv_layout="HND") for layer_id in layer_ids + ] base = layers[0] # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. expected = (kv_factor, nkv, base.size(-2), hd) @@ -960,10 +975,34 @@ def _init_ctx_block_tables( """ mgr = draft_kv_cache_manager pool_mapping = mgr.kv_cache_pool_mapping - layer_offset = mgr.layer_offsets[0] + first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 + layer_offset = mgr.layer_offsets[first_layer_id] self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - if num_pool_layers != L: + self._ctx_block_offset = 0 + if isinstance(mgr, KVCacheManagerV2): + converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) + layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + if mgr.enable_swa_scratch_reuse: + raise NotImplementedError( + "DFlash managed context does not support SWA scratch pages" + ) + for layer_id in self._ctx_layer_ids: + local_id = mgr.layer_offsets[layer_id] + if ( + mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id + or mgr.get_layer_page_index_scale(layer_id) != converter.scale + ): + raise NotImplementedError( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + if per_layer_tables: + # Attention tables address the physical pool origin, whereas + # get_buffers() starts at this layer's own K buffer. + self._ctx_block_offset = int(converter.layer_offset) + elif num_pool_layers != L: logger.warning( f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " f"the drafter has {L}; using the private ctx arena." @@ -972,12 +1011,23 @@ def _init_ctx_block_tables( # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps # the layer base in a dense view over interleaved pages, so stride(0) is # one layer slot. Guessing wrong is silent: acceptance quietly drops. - layer_block_elems = mgr.kv_factor * nkv * page_size * hd + manager_kv_factor = ( + mgr.get_layer_kv_factor(first_layer_id) + if hasattr(mgr, "get_layer_kv_factor") + else mgr.kv_factor + ) + if manager_kv_factor != getattr(self, "_ctx_kv_factor", manager_kv_factor): + logger.warning( + "DFlash: draft KV layer factor does not match the drafter; " + "using the private ctx arena." + ) + return False + layer_block_elems = manager_kv_factor * nkv * page_size * hd stride0 = pool[0].stride(0) if stride0 == layer_block_elems: - self._ctx_block_divisor = mgr.kv_factor + self._ctx_block_divisor = manager_kv_factor elif stride0 == num_pool_layers * layer_block_elems: - self._ctx_block_divisor = num_pool_layers * mgr.kv_factor + self._ctx_block_divisor = num_pool_layers * manager_kv_factor else: logger.warning( f"DFlash: draft KV pool layer stride {stride0} matches neither V1's " @@ -1010,20 +1060,24 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """ if self._ctx_block_tables is None or num_seqs <= 0: return False - src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + src = getattr(attn_metadata, offsets_name, None) if src is None: # The table is bound but unfillable. Leaving it at zeros would send # every request's every layer to pool block 0 -- silently. raise RuntimeError( "DFlash: bound to the draft KV pool but " f"{type(attn_metadata).__name__} carries no " - "draft_kv_cache_block_offsets this iteration." + f"{offsets_name} this iteration." ) # Field 0 of the K/V pair: both name the same block for a kv_factor=2 # pool, and the drafter addresses blocks, not planes. # clone(): .to() is a no-op for an int64 source, and the in-place ops # below would then edit attn_metadata's own tensor. encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() + encoded.sub_(self._ctx_block_offset) # Before the clamp erases the placeholders. Informative only on the # python V2 backend; the cpp one already mapped BAD_PAGE_INDEX to 0. self._ctx_block_counts[:num_seqs].copy_((encoded >= 0).sum(dim=1)) @@ -1035,12 +1089,18 @@ def _lazy_init_ctx_buffers( self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None ): self._validate_draft_attention_backend(draft_model) - # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets. - # Without it the table stays all-zero and every request would read and - # write pool block 0 -- silently, at acceptance 1.0. - if draft_kv_cache_manager is not None and not hasattr( - attn_metadata, "draft_kv_cache_block_offsets" - ): + self._ctx_unified_cache = ( + draft_kv_cache_manager is not None + and draft_kv_cache_manager is getattr(attn_metadata, "kv_cache_manager", None) + ) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + if draft_kv_cache_manager is not None and not hasattr(attn_metadata, offsets_name): + if self._ctx_unified_cache: + raise NotImplementedError( + f"DFlash unified KV requires {offsets_name} in attention metadata" + ) logger.warning( f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " "using the private ctx arena." @@ -1145,6 +1205,7 @@ def _lazy_init_ctx_buffers( nkv = draft_model._num_kv_heads hd = draft_model._head_dim kv_factor = getattr(draft_model, "_kv_factor", 2) + self._ctx_kv_factor = kv_factor capacity = self._max_ctx + self._compute_block_size # Published before any drafter table is built, so a drafter may read it # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which @@ -1200,6 +1261,8 @@ def _lazy_init_ctx_buffers( ): pool = None if pool is None: + if self._ctx_unified_cache: + raise ValueError("DFlash could not bind the unified manager's draft KV layers") # Per-slot reservation is an arena-only notion; the pool path # takes its pages from the manager one iteration at a time. self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size @@ -1230,7 +1293,7 @@ def _lazy_init_ctx_buffers( # Only a manager that publishes and matches its own blocks can # hand back a reused prefix's drafter K/V; the private arena and # an unpaired draft pool both start every request at 0. - self._ctx_reuse_addressable = bool( + self._ctx_reuse_addressable = self._ctx_unified_cache or bool( getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) ) self._ctx_kv_last_page_len = torch.full( @@ -2100,7 +2163,18 @@ def prepare_1st_drafter_inputs( bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() - ctx_len_gen = self._ctx_len[slots] + if self._ctx_unified_cache: + gen_start = attn_metadata.num_ctx_tokens + # Device positions include overlap-scheduler corrections and + # define the shared request's prefix before this verification. + ctx_len_gen = ( + position_ids[gen_start : gen_start + num_gens * total_tokens_per_req] + .reshape(num_gens, total_tokens_per_req)[:, 0] + .long() + ) + self._ctx_len[slots] = ctx_len_gen + else: + ctx_len_gen = self._ctx_len[slots] j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") @@ -2137,7 +2211,7 @@ def prepare_1st_drafter_inputs( projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) if self._ctx_block_tables is not None: # Bound by what the manager handed THIS request: the masked diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8237c09ee8b7..8ce96346acc1 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -190,7 +190,8 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, Call restore_attn_metadata_after_draft_replay in a finally block. Returns saved state or None if no-op. """ - if draft_kv_cache_manager is None: + if (draft_kv_cache_manager is None + or draft_kv_cache_manager is attn_metadata.kv_cache_manager): return None if not isinstance(attn_metadata, TrtllmAttentionMetadata): return None @@ -3167,8 +3168,10 @@ def get_draft_kv_cache_manager(self, resource_manager): Get the draft KV cache manager if using separate KV cache layouts. """ if self.use_separate_draft_kv_cache and resource_manager is not None: - return resource_manager.get_resource_manager( + draft = resource_manager.get_resource_manager( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if draft is not None: + return draft if resource_manager is not None: target = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER) @@ -3189,7 +3192,8 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): """ # draft_kv_cache_manager is None if using two-engine speculative decoding or not enabling separate draft KV cache. - if draft_kv_cache_manager is None: + if (draft_kv_cache_manager is None + or draft_kv_cache_manager is attn_metadata.kv_cache_manager): yield attn_metadata return diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 813ad824262a..454f2567d000 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -809,7 +809,9 @@ def get_draft_kv_cache_manager(spec_config, resource_manager): target = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER) get_view = getattr(target, "get_draft_kv_cache_view", None) - return get_view() if get_view is not None else None + view = get_view() if get_view is not None else None + # A unified manager's request resources are already handled through target. + return None if view is target else view def update_spec_config_from_model_config(spec_config, diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 8d658362ebc8..f5a1d7fe2e51 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -122,6 +122,12 @@ ModelingV2DeepseekR10528Nvfp4Sm103Dep4::mtp3: TestGLM53FlashFP8::test_mtp: ref_al: 3.384 min_al: 3.215 +# Measured on B200, 1P1D TP1 with target and draft KV in one manager. +# min_al is 95% of the measured value. +TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv: + ref_al: 4.062 + min_al: 3.859 + # Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1:: # test_mtp1_gsm8k_and_acceptance. Keyed by target and variant for the same # reason as the ModelingV2DeepseekR10528Nvfp4Sm103Dep4::mtp3 entry above. diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 01cc40e03e4f..3988bdcc2e1e 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -39,7 +39,8 @@ from tensorrt_llm.executor.result import GenerationResultBase from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams -from tensorrt_llm.llmapi.llm_args import (DSparkDecodingConfig, +from tensorrt_llm.llmapi.llm_args import (DFlashDecodingConfig, + DSparkDecodingConfig, Eagle3DecodingConfig, LlmArgs, MTPDecodingConfig) from tensorrt_llm.llmapi.tokenizer import load_hf_tokenizer @@ -1623,6 +1624,94 @@ def test_gen_first(self, enable_block_reuse): run_accuracy_test(llm, self.MODEL_NAME, ["MMLU"]) +@skip_pre_blackwell +@pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) +class TestQwen3_8BDFlash(LlmapiAccuracyTestHarness): + MODEL_NAME = "Qwen3/Qwen3-8B" + MODEL_PATH = f"{llm_models_root()}/Qwen3/Qwen3-8B" + DRAFT_PATH = f"{llm_models_root()}/Qwen3-8B-DFlash-b16" + MAX_DRAFT_LEN = 15 + + @pytest.mark.skip_less_device(2) + def test_gsm8k_1p1d_unified_draft_kv(self): + """DFlash draft KV travels with the target KV from ctx to gen. + + The drafter's layers share the target's KV cache manager, so the + generation worker drafts from the prompt context computed on the + context worker. Verified speculative decoding is lossless, so the + GSM8K score cannot see a drafter that starts from an empty context; + acceptance length can, so it is gated too. + """ + speculative_config = { + "decoding_type": "DFlash", + "max_draft_len": self.MAX_DRAFT_LEN, + "speculative_model": self.DRAFT_PATH, + # The unified draft KV layout requires the TRTLLM drafter backend. + "attention_backend": "TRTLLM", + } + # Both workers run the drafter: the context worker writes the + # drafter's prompt KV into the shared manager, which is then + # transferred with the target KV. + common_server_config = { + "disable_overlap_scheduler": True, + "max_batch_size": 16, + "max_seq_len": 4096, + "max_num_tokens": 4096, + "kv_cache_config": { + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.5, + }, + # KVCacheManagerV2 requires the Python transceiver. + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "speculative_config": speculative_config, + } + ctx_server_config = {**common_server_config, "cuda_graph_config": None} + gen_server_config = { + **common_server_config, + "cuda_graph_config": { + "max_batch_size": 16, + "enable_padding": True, + }, + # Acceptance length is read back off this worker's /metrics. + "enable_iter_perf_stats": True, + "iter_stats_max_iterations": -1, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + }, + } + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + assert_worker_log_contains= + "target and draft KV layers share one KV cache manager") as llm: + llm.args.speculative_config = DFlashDecodingConfig( + max_draft_len=self.MAX_DRAFT_LEN, + speculative_model=self.DRAFT_PATH, + ) + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + + # Drafting from an empty prompt context costs ~17% acceptance + # length on this model, well below the 95% floor. + acceptance_length = compute_disagg_acceptance_length(llm.serve_url) + print(f"[AL] test_gsm8k_1p1d_unified_draft_kv " + f"acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length( + "TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv", + acceptance_length) + + @skip_pre_blackwell @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) class TestQwen3_30B_A3B(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index a60aa69f4d12..60d82181683c 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -28,6 +28,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam2-reuse] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam4-no-reuse] + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv TIMEOUT (60) # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py new file mode 100644 index 000000000000..fc69a7f37328 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -0,0 +1,139 @@ +# 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. +"""Selection of the unified target+DFlash KV cache layout.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.quantization.mode import QuantMode + +pytestmark = pytest.mark.cpu_only + + +def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): + pretrained = SimpleNamespace( + num_hidden_layers=2, + num_attention_heads=32, + num_key_value_heads=kv_heads, + head_dim=head_dim, + hidden_size=32 * head_dim, + torch_dtype=torch.bfloat16, + ) + quant_mode = QuantMode.from_quant_algo(kv_cache_quant_algo=kv_quant_algo) + return SimpleNamespace( + pretrained_config=pretrained, + quant_config=SimpleNamespace(quant_mode=quant_mode), + ) + + +def _creator(*, is_disagg, target, draft, tp_size=1): + creator = object.__new__(KvCacheCreator) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True), + attention_backend="TRTLLM", + _allow_separate_draft_kv_cache=True, + draft_len_schedule=None, + ) + creator._draft_model_engine = None + creator._draft_config = draft + creator._mapping = SimpleNamespace( + pp_size=1, cp_size=1, tp_size=tp_size, enable_attention_dp=False + ) + creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) + creator._model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 + ) + creator._kv_cache_manager_cls = KVCacheManagerV2 + creator._is_disagg = is_disagg + creator._get_draft_kv_model_config = lambda: draft + return creator + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_matching_draft_layout_unifies(is_disagg): + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +def test_mismatched_draft_kv_heads_unify_only_in_aggregated_serving() -> None: + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + assert _creator(is_disagg=False, target=target, draft=draft)._should_unify_dflash_kv_cache() + assert not _creator(is_disagg=True, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_mismatched_draft_head_dim_unifies(is_disagg: bool) -> None: + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=64) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_fp8_target_with_bf16_draft_unifies(is_disagg: bool) -> None: + target = _model_config(kv_heads=8, head_dim=128, kv_quant_algo=QuantAlgo.FP8) + draft = _model_config(kv_heads=8, head_dim=128) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +def test_layout_compares_rank_local_heads(): + # 8 target heads and 2 draft heads differ per rank at every TP size; at + # TP4 the draft is duplicated to 1 head per rank while the target has 2. + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=2, head_dim=128) + creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=4) + + assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 2 + assert creator._dflash_kv_layout(draft, torch.bfloat16)[0] == 1 + assert not creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_attention_dp_unifies_with_unsharded_heads(is_disagg): + # Under attention DP each rank keeps every KV head, so the TP4 layouts of + # an 8-head target and a 2-head draft are not resharded to 2 and 1. + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=is_disagg, target=target, draft=draft, tp_size=4) + creator._mapping.enable_attention_dp = True + + assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 8 + assert ( + creator._dflash_kv_layout(_model_config(kv_heads=2, head_dim=128), torch.bfloat16)[0] == 2 + ) + assert creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("per_layer_side", ["target", "draft"]) +def test_per_layer_kv_heads_keep_separate_cache(is_disagg, per_layer_side): + uniform = _model_config(kv_heads=8, head_dim=128) + per_layer = _model_config(kv_heads=[8, 4], head_dim=128) + target, draft = (per_layer, uniform) if per_layer_side == "target" else (uniform, per_layer) + + creator = _creator(is_disagg=is_disagg, target=target, draft=draft) + + assert not creator._should_unify_dflash_kv_cache() + assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 7ff5f5e72e7c..8b9d3ca2ffc9 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -1380,3 +1380,29 @@ def test_attention_replicated_validation_rejects_missing_index_key_view(): peer_pt = make_page_table(global_layer_ids=[0]) with pytest.raises(ValueError, match="AttentionPolicy.*differ on overlapping layers"): AttentionPolicy.validate_peer_compatible(self_pt, peer_pt) + + +@pytest.mark.parametrize(("local_draft", "peer_draft"), [(0, 0), (1, 1), (1, 0), (0, 1)]) +def test_peer_registrar_requires_matching_unified_draft_layers( + local_draft: int, peer_draft: int +) -> None: + self_rankinfo = make_rankinfo(instance_name="local") + self_rankinfo.num_draft_layers = local_draft + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=1, + layer_num_per_pp=[2], + page_table=make_page_table(), + ) + peer_ri.num_draft_layers = peer_draft + if local_draft != peer_draft: + with pytest.raises( + ValueError, + match=f"draft KV layer count differs.*local={local_draft}, peer={peer_draft}", + ): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + return + + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + assert reg.get_peer_rank_info("peer", 1) == peer_ri diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index a1d25e01c3f3..a9ef379c93e7 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -15,6 +15,7 @@ from types import SimpleNamespace +import msgpack import numpy as np import pytest @@ -203,3 +204,64 @@ def test_rank_info_represents_cache_element_bytes( restored = RankInfo.from_bytes(rank_info.to_bytes()) assert restored.attention.element_bytes == expected_element_bytes assert isinstance(restored.attention.element_bytes, expected_type) + + +@pytest.mark.parametrize( + ("layer_specs", "draft_layer_ids", "expected"), + [ + (None, (), 0), + # Folded MTP/Eagle draft layers have no layer specs and travel as + # ordinary layers, so they are not advertised. + (None, (36,), 0), + ((object(),), (36, 37), 2), + ], +) +def test_from_kv_cache_manager_reports_unified_draft_layers( + monkeypatch, layer_specs, draft_layer_ids, expected +) -> None: + monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _: None) + manager = SimpleNamespace( + mapping=SimpleNamespace( + rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + dp_size=1, + cp_size=1, + cp_rank=0, + enable_attention_dp=False, + ), + num_kv_heads_per_layer=[8, 8], + pp_layers=[0, 1], + tokens_per_block=32, + head_dim=128, + dtype=bindings.DataType.HALF, + kv_factor=2, + layer_specs=layer_specs, + draft_layer_ids=draft_layer_ids, + ) + + info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) + + assert info.num_draft_layers == expected + assert RankInfo.from_bytes(info.to_bytes()).num_draft_layers == expected + + +def test_rank_info_from_bytes_defaults_missing_draft_layer_count() -> None: + ri = RankInfo( + instance_name="gen_0", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[32], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + ) + payload = msgpack.unpackb(ri.to_bytes(), strict_map_key=False) + del payload["num_draft_layers"] + + assert RankInfo.from_bytes(msgpack.packb(payload)).num_draft_layers == 0 From 39d814dcc64f0c4d274b04d3a78f45c814f2403a Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 17 Oct 2026 21:54:12 -0700 Subject: [PATCH 17/35] [None][fix] share the DFlash draft KV cache in disagg only with speculative_config in both sections In disaggregated serving, workers launched from a disaggregated config file share the draft KV cache only when both context_servers and generation_servers set speculative_config. Other workers keep the separate draft cache. Aggregated serving is unchanged. Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 3 +++ tensorrt_llm/commands/serve.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 96a160e78670..b5a8e71b2aee 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2109,6 +2109,9 @@ def _should_unify_dflash_kv_cache(self) -> bool: "target and draft KV cache; keeping a separate draft KV cache.", key="dflash_per_layer_kv_heads_fallback") return False + if self._is_disagg and os.environ.get( + "TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": + return False if self._is_disagg and not self._dflash_disagg_layouts_compatible( target_layout, draft_layout): logger.info_once( diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ff6d148bc529..692d44440783 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2544,6 +2544,13 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, from tensorrt_llm.llmapi.disagg_utils import split_world_comm disagg_cfg = parse_disagg_config_file(config_file) + with_spec = { + cfg.type + for cfg in disagg_cfg.server_configs + if cfg.other_args.get("speculative_config") + } + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("1" if with_spec + == {"ctx", "gen"} else "0") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. From 22cf25c6f5cc753c50d4c20d072cfc3e92230808 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 17 Oct 2026 21:57:13 -0700 Subject: [PATCH 18/35] [None][fix] gate the disagg shared draft KV cache on matching tensor_parallel_size Workers launched from a disaggregated config file share the draft KV cache only when context_servers and generation_servers use the same tensor_parallel_size, in addition to both setting speculative_config. Signed-off-by: Allison Lim --- tensorrt_llm/commands/serve.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 692d44440783..4bda2c05e02b 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2549,8 +2549,12 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, for cfg in disagg_cfg.server_configs if cfg.other_args.get("speculative_config") } - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("1" if with_spec - == {"ctx", "gen"} else "0") + tp_sizes = { + cfg.other_args.get("tensor_parallel_size", 1) + for cfg in disagg_cfg.server_configs + } + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ( + "1" if with_spec == {"ctx", "gen"} and len(tp_sizes) == 1 else "0") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. From feed06a1b60d6ae397efdd18cdd26bbf8f557762 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 17 Oct 2026 22:10:19 -0700 Subject: [PATCH 19/35] [None][test] cover the disagg shared draft KV gate for generation-only speculative_config and mismatched TP Drive the real disaggregated_mpi_worker against a config and check that the target and draft KV layers keep separate cache managers when only generation_servers sets speculative_config, and when the context and generation TP sizes differ. Signed-off-by: Allison Lim --- .../kv_cache/test_dflash_unified_kv_gate.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index fc69a7f37328..b58fcfdf5bfd 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -18,6 +18,7 @@ import pytest import torch +import yaml from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -26,6 +27,14 @@ pytestmark = pytest.mark.cpu_only +_SHARED_DRAFT_KV_ENV = "TRTLLM_DISAGG_SHARED_DRAFT_KV" + + +@pytest.fixture(autouse=True) +def _shared_draft_kv_allowed(monkeypatch): + # Tests of the layout rules assume a launcher that allows sharing. + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, "1") + def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): pretrained = SimpleNamespace( @@ -137,3 +146,67 @@ def test_per_layer_kv_heads_keep_separate_cache(is_disagg, per_layer_side): assert not creator._should_unify_dflash_kv_cache() assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () + + +class _StopLauncher(Exception): + pass + + +def _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers): + """Run the real disaggregated_mpi_worker up to the rank-communicator split.""" + from tensorrt_llm.commands import serve + + config = tmp_path / "disagg.yaml" + config.write_text( + yaml.safe_dump( + { + "context_servers": {"num_instances": 1, **context_servers}, + "generation_servers": {"num_instances": 1, **generation_servers}, + } + ) + ) + + def stop(*args, **kwargs): + raise _StopLauncher + + monkeypatch.setattr( + serve._command_telemetry, "apply_disaggregated_telemetry_config", lambda *a, **k: None + ) + monkeypatch.setattr(serve, "set_cuda_device", lambda: None) + monkeypatch.setattr("tensorrt_llm.llmapi.disagg_utils.split_world_comm", stop) + monkeypatch.delenv("TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT", raising=False) + with pytest.raises(_StopLauncher): + serve.disaggregated_mpi_worker.callback( + config_file=str(config), log_level="info", telemetry=False + ) + + +def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_path): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={}, + generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, + ) + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + + for _role in ("context", "generation"): + creator = _creator(is_disagg=True, target=target, draft=draft) + assert not creator._should_unify_dflash_kv_cache() + + +def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): + spec = {"decoding_type": "DFlash"} + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, + generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, + ) + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + for tp_size in (4, 8): + creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=tp_size) + assert not creator._should_unify_dflash_kv_cache() From 0319dc361dde45a895bc76f6661f4671d8dfaa59 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 09:11:43 -0700 Subject: [PATCH 20/35] test: use default timeout for unified DFlash KV test Signed-off-by: Allison Lim --- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 60d82181683c..bbd152ad667a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -28,7 +28,7 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam2-reuse] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam4-no-reuse] - - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens From 1f391284f0919d7df83756fdc2719aa8a66d7b61 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 11:22:20 -0700 Subject: [PATCH 21/35] fix: restore DFlash private arena fallback for separate KV caches Signed-off-by: Allison Lim --- tensorrt_llm/_torch/speculative/dflash.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index a34b49a35b2b..fbe296ec50f6 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -994,10 +994,16 @@ def _init_ctx_block_tables( mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id or mgr.get_layer_page_index_scale(layer_id) != converter.scale ): - raise NotImplementedError( - "DFlash's shared context block table requires draft layers " - "with the same lifecycle and page-index scale" + if self._ctx_unified_cache: + raise NotImplementedError( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + logger.warning( + "DFlash: draft KV layers have different lifecycles or page-index " + "scales; using the private ctx arena." ) + return False if per_layer_tables: # Attention tables address the physical pool origin, whereas # get_buffers() starts at this layer's own K buffer. From 97fb214ab1feb74668f9ad310e2e040b238158fa Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 12:04:58 -0700 Subject: [PATCH 22/35] fix: seed unified DFlash context lengths from bounded token counts Signed-off-by: Allison Lim --- tensorrt_llm/_torch/speculative/dflash.py | 35 ++++++++++++++--------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fbe296ec50f6..c7cc1af494be 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -607,19 +607,29 @@ def prepare(self): # prefilling, so _store_prefill_context -- the only other assigner -- # never runs and every request would land on the dummy slot. num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + assigned = {} + for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=assigned) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + elif worker._ctx_unified_cache: + # New requests have no in-flight predecessor. Subsequent + # accepted tokens advance the device length during replay. + token_offset = sum(self.seq_lens[:request_idx]) + assigned[slot] = min( + int(self.host_position_ids[token_offset]), worker._max_ctx + ) + worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context @@ -1705,10 +1715,18 @@ def _forward_impl( # Lazy init buffers and attach worker reference for prepare() draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + previous_ctx_len = self._ctx_len self._lazy_init_ctx_buffers( draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager ) spec_metadata._dflash_worker = self + if ( + self._ctx_unified_cache + and self._ctx_len is not previous_ctx_len + and spec_metadata.request_ids + ): + # New buffers clear request slots; seed them before generation. + spec_metadata.prepare() # Before any store: prefill and decode both address pages through it. # Returning False here means an empty batch -- the missing-offsets case # raises inside, with the metadata type in the message. @@ -2169,18 +2187,7 @@ def prepare_1st_drafter_inputs( bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() - if self._ctx_unified_cache: - gen_start = attn_metadata.num_ctx_tokens - # Device positions include overlap-scheduler corrections and - # define the shared request's prefix before this verification. - ctx_len_gen = ( - position_ids[gen_start : gen_start + num_gens * total_tokens_per_req] - .reshape(num_gens, total_tokens_per_req)[:, 0] - .long() - ) - self._ctx_len[slots] = ctx_len_gen - else: - ctx_len_gen = self._ctx_len[slots] + ctx_len_gen = self._ctx_len[slots] j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") From caf4c45a28d87dc07da5a4cec31d560c61af2ed7 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 12:27:42 -0700 Subject: [PATCH 23/35] test: cover DFlash mixed-precision block tables Signed-off-by: Allison Lim --- .../kv_cache/test_kvcm2_integration.py | 79 +++++++++++++++++++ 1 file changed, 79 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 0ad1910c7894..c38b2e3f8ca2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -748,6 +748,85 @@ def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layo manager.shutdown() +def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + torch.cuda.init() + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=2, + head_dim=8, + buffer_roles=(Role.KEY, Role.VALUE), + ) + # Equal byte sizes coalesce the FP8 target and BF16 draft pages, placing + # the draft's K buffer at a nonzero offset within each physical slot. + draft = replace( + target, + global_layer_id=1, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=1, + ) + 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=1, + num_kv_heads=target.num_kv_heads, + head_dim=target.head_dim, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + layer_specs=[target, draft], + ) + try: + assert manager._use_per_layer_page_tables + request_ids = [7, 9] + for request_id, num_blocks in zip(request_ids, (2, 3)): + cache = manager._create_kv_cache(request_id, None, None) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(num_blocks * TOKENS_PER_BLOCK, 0) + + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + pool = worker._managed_ctx_pool( + manager, 1, draft.num_kv_heads, draft.head_dim, torch.bfloat16, 2 + ) + assert pool is not None + assert worker._init_ctx_block_tables( + manager, pool, len(request_ids), 1, draft.num_kv_heads, draft.head_dim, TOKENS_PER_BLOCK + ) + assert worker._ctx_block_offset > 0 + + block_offsets = torch.empty( + (manager.num_attention_op_pools, len(request_ids), 2, manager.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + manager.copy_batch_block_offsets( + block_offsets, request_ids, beam_width=1, num_contexts=0, num_seqs=len(request_ids) + ) + metadata = SimpleNamespace(kv_cache_block_offsets=block_offsets) + assert worker._refresh_ctx_block_tables(metadata, len(request_ids)) + + expected = manager.get_batch_cache_indices(request_ids, layer_idx=draft.global_layer_id) + for row, indices in enumerate(expected): + assert len(indices) == (2, 3)[row] + assert worker._ctx_block_tables[row, : len(indices)].tolist() == indices + finally: + manager.shutdown() + + def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") From 875eff2ee13361bbee150c3a070edf8347ef5e9c Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 14:24:35 -0700 Subject: [PATCH 24/35] fix: require matching attention DP for shared draft KV Signed-off-by: Allison Lim --- tensorrt_llm/commands/serve.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 4bda2c05e02b..c58c6741dfaf 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2549,12 +2549,14 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, for cfg in disagg_cfg.server_configs if cfg.other_args.get("speculative_config") } - tp_sizes = { - cfg.other_args.get("tensor_parallel_size", 1) - for cfg in disagg_cfg.server_configs - } + attention_parallel_configs = {(cfg.other_args.get("tensor_parallel_size", + 1), + cfg.other_args.get("enable_attention_dp", + False)) + for cfg in disagg_cfg.server_configs} os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ( - "1" if with_spec == {"ctx", "gen"} and len(tp_sizes) == 1 else "0") + "1" if with_spec == {"ctx", "gen"} + and len(attention_parallel_configs) == 1 else "0") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. From 212fdabdf5094687326c973cb0cdcb0c4fe242d3 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 14:27:25 -0700 Subject: [PATCH 25/35] test: enable shared draft KV in DFlash integration test Signed-off-by: Allison Lim --- tests/integration/defs/accuracy/test_disaggregated_serving.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 3988bdcc2e1e..81d893199157 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1694,6 +1694,7 @@ def test_gsm8k_1p1d_unified_draft_kv(self): ctx_server_config, gen_server_config, self.MODEL_PATH, + extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}, assert_worker_log_contains= "target and draft KV layers share one KV cache manager") as llm: llm.args.speculative_config = DFlashDecodingConfig( From 1302d43bf80fa51dc7d542a77b127643d7221c0c Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 15:15:40 -0700 Subject: [PATCH 26/35] fix: keep mixed-window DFlash caches separate Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 13 +++++++++++++ .../kv_cache/test_dflash_unified_kv_gate.py | 4 +++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e9bca37a96f9..eef66a3767e3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -939,6 +939,8 @@ def __init__( self._sparse_attention_config = sparse_attention_config self._tokens_per_block = tokens_per_block self._max_seq_len = max_seq_len + # Estimation can shrink _max_seq_len; cache ownership must stay fixed. + self._unification_max_seq_len = max_seq_len self._max_batch_size = max_batch_size self._net_max_seq_len = net_max_seq_len self._dummy_reqs = None @@ -2195,6 +2197,17 @@ def _should_unify_dflash_kv_cache(self) -> bool: "target and draft KV cache; keeping a separate draft KV cache.", key="dflash_per_layer_kv_heads_fallback") return False + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._unification_max_seq_len, + self._kv_cache_config), + self._unification_max_seq_len, + range(draft_model_config.pretrained_config.num_hidden_layers)) + if len(set(draft_windows)) > 1: + logger.info_once( + "DFlash: draft layers have different effective attention " + "windows; keeping a separate draft KV cache.", + key="dflash_mixed_draft_windows_fallback") + return False if self._is_disagg and os.environ.get( "TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": return False diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index b58fcfdf5bfd..c91507098ac3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -55,7 +55,7 @@ def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): def _creator(*, is_disagg, target, draft, tp_size=1): creator = object.__new__(KvCacheCreator) creator._speculative_config = SimpleNamespace( - spec_dec_mode=SimpleNamespace(is_dflash=lambda: True), + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), attention_backend="TRTLLM", _allow_separate_draft_kv_cache=True, draft_len_schedule=None, @@ -70,6 +70,8 @@ def _creator(*, is_disagg, target, draft, tp_size=1): model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 ) creator._kv_cache_manager_cls = KVCacheManagerV2 + creator._kv_cache_config = SimpleNamespace(max_attention_window=None) + creator._unification_max_seq_len = 4096 creator._is_disagg = is_disagg creator._get_draft_kv_model_config = lambda: draft return creator From 7a85a733a06f4c329e744ddd19ca5698320d4512 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 19 Oct 2026 23:46:10 -0700 Subject: [PATCH 27/35] [None][feat] unify the standalone DSpark target and draft KV cache Keep the DFlash separate-manager fallback and make the disaggregated shared draft KV cache opt-in through TRTLLM_DISAGG_SHARED_DRAFT_KV, which the MPI launcher sets when both sections have speculation and the same (tensor_parallel_size, enable_attention_dp). Derive the V2 estimation pool-group count from the combined target and draft layout, and extend the shared cache to standalone DSpark. Reserve the drafter block (7 positions for DSpark) in the shared allocation, and route generation rows carrying the attention-DP idle request id to the dummy slot so a warmup prefill stored under that id cannot leave a stale length that trips the block-bound check under FP8 target KV. Signed-off-by: Allison Lim --- docs/source/features/speculative-decoding.md | 26 ++- .../_torch/disaggregation/native/peer.py | 17 +- tensorrt_llm/_torch/pyexecutor/_util.py | 170 ++++++++++++---- .../kv_cache/kv_cache_manager_v2.py | 77 ++++++-- tensorrt_llm/_torch/speculative/__init__.py | 10 +- tensorrt_llm/_torch/speculative/dflash.py | 65 +++++-- tensorrt_llm/_torch/speculative/dspark.py | 14 +- tensorrt_llm/_torch/speculative/utils.py | 15 ++ tensorrt_llm/commands/serve.py | 45 +++-- .../accuracy/test_disaggregated_serving.py | 4 +- .../kv_cache/test_dflash_unified_kv_gate.py | 142 +++++++++++--- .../kv_cache/test_kv_cache_budget_split.py | 3 + .../kv_cache/test_kv_cache_estimation.py | 53 +++++ .../kv_cache/test_kvcm2_integration.py | 181 ++++++++++++++++++ .../test_dflash_ctx_block_tables.py | 98 ++++++++++ 15 files changed, 804 insertions(+), 116 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 9e20a8db6083..65b1a52168a8 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -152,7 +152,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/ * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model. * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config. * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config. -* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a private paged context K/V cache and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. +* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. ```python from tensorrt_llm.llmapi import DFlashDecodingConfig @@ -166,6 +166,30 @@ speculative_config = DFlashDecodingConfig( llm = LLM("/path/to/target_model", speculative_config=speculative_config) ``` +#### DFlash draft KV cache + +With `attention_backend="TRTLLM"`, TRTLLM target attention, and the default +`KVCacheManagerV2`, the drafter's layers are allocated in the target's KV cache +manager. Target and draft layers share one page table, so one KV budget covers +both and each request's draft pages are allocated, padded and freed with its +target pages. The draft layers keep the drafter's own dtype, so an FP8 target KV +cache (`kv_cache_config.dtype: fp8`) is paired with a BF16 draft cache. + +In disaggregated serving, the context server's draft pages are transferred to +the generation server with the target pages, so the drafter starts generation +with the prompt context instead of an empty one. This requires the draft layers +to have the target's rank-local KV-head count; the draft dtype and head +dimension may differ, except that an NVFP4 target KV cache requires an +identical draft layout. Both servers must use the same layout. + +The drafter keeps a separate draft KV cache manager instead when any of the +following applies: context parallelism larger than 1, a per-layer KV-head count +on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba +models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout +that cannot be transferred with the target's. When the layers are shared, the +worker log reports `target and draft KV layers share one KV cache manager` at +startup. + [DFlash 2](https://inco.ai/blog/dflash2/) is also supported. The same `DFlashDecodingConfig` can be used for DFlash 2; no extra arguments are required. [LiLiCorr](https://arxiv.org/abs/2608.20530) checkpoints use the same diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 6b98a2772ab1..c3126e9a08da 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -186,8 +186,23 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: raise ValueError( "PeerRegistrar: unified draft KV layer count differs " f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " - "both peers must use the same speculative decoding KV cache layout." + "both peers must use the same speculative decoding KV cache layout; " + "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." ) + if self._ri.num_draft_layers: + if self._ri.tp_size != peer_ri.tp_size: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs equal tp_size " + f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); " + "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers." + ) + if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs the same attention DP " + f"setting (local={self._ri.attention.enable_attention_dp}, " + f"peer={peer_ri.attention.enable_attention_dp}); unset " + "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers." + ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index eef66a3767e3..138cc482847d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -764,6 +764,35 @@ def _derive_v2_layer_type_attention_windows( return derived_windows +def _get_num_storage_pool_groups(layer_specs: Sequence[KVCacheLayerSpec], + tokens_per_block: int) -> int: + """Count the pool groups KVCacheManagerV2 creates for explicit layer specs. + + Mirrors the grouping key of ``createStorageConfig``: a life cycle is a + distinct (window, sink blocks, reuse policy), its slot sizes are the sizes + of its coalesced buffers in descending order, and life cycles with the same + (reuse policy, slot sizes) share one pool group. + """ + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 + life_cycles: dict[tuple, dict[tuple[int, int], int]] = {} + for layer in layer_specs: + life_cycle = (layer.sliding_window_size, + ceil_div(layer.num_sink_tokens or 0, + tokens_per_block), layer.reuse_policy) + coalesced = life_cycles.setdefault(life_cycle, {}) + for role in layer.buffer_roles: + group = 1 + int(role in scale_roles) if layer.dtype == nvfp4 else 0 + key = (group, layer.bytes_per_token(role) * tokens_per_block) + coalesced[key] = coalesced.get(key, 0) + 1 + return len({(life_cycle[2], + tuple( + sorted((size * count + for (_, size), count in coalesced.items()), + reverse=True))) + for life_cycle, coalesced in life_cycles.items()}) + + def _get_num_pool_groups_for_estimation( model_config: object, max_seq_len: int, @@ -1620,12 +1649,20 @@ def _get_token_num_for_estimation(self) -> int: # (distinct layer types without sliding windows). num_pool_groups = 1 if self._is_kv_cache_manager_v2: - model_cfg = self._model_engine.model.model_config.pretrained_config - num_pool_groups = _get_num_pool_groups_for_estimation( - model_cfg, - self._model_engine.max_seq_len, - self._kv_cache_config.max_attention_window, - ) + unified_layer_specs = self._build_unified_dflash_layer_specs( + self._kv_cache_config) + if unified_layer_specs: + # Target and draft layers can differ in page size and life + # cycle, so the group count comes from the combined layout. + num_pool_groups = _get_num_storage_pool_groups( + unified_layer_specs, self._tokens_per_block) + else: + model_cfg = self._model_engine.model.model_config.pretrained_config + num_pool_groups = _get_num_pool_groups_for_estimation( + model_cfg, + self._model_engine.max_seq_len, + self._kv_cache_config.max_attention_window, + ) num_cache_blocks *= num_pool_groups # Dummy context requests use the configured maximum beam width. Scale @@ -2068,12 +2105,14 @@ def _create_kv_cache_manager( spec_dec_layer_mask = [True] * num_target_layers estimating_kv_cache = estimating_kv_cache and not self._skip_est + self._check_shared_draft_kv_request() unified_layer_specs = None if self._should_unify_dflash_kv_cache(): unified_layer_specs = self._build_unified_dflash_layer_specs( kv_cache_config) logger.info( - "DFlash: target and draft KV layers share one KV cache manager " + f"{self._unified_drafter_name()}: target and draft KV layers " + "share one KV cache manager " f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})." ) kv_cache_manager = _create_kv_cache_manager( @@ -2167,24 +2206,85 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) - def _should_unify_dflash_kv_cache(self) -> bool: - """Whether this run can place target and DFlash layers in one V2 manager.""" + def _unified_drafter_attention_backend(self) -> Optional[str]: + """Resolved block-attention backend of a drafter that reads paged draft KV. + + None for every other drafter. That includes the embedded DSpark draft, + whose context lives in a worker-owned rolling window. + """ spec_config = self._speculative_config - if (spec_config is None or not spec_config.spec_dec_mode.is_dflash() - or spec_config.attention_backend != "TRTLLM" - or not spec_config._allow_separate_draft_kv_cache - or spec_config.draft_len_schedule is not None - or self._draft_model_engine is not None - or self._draft_config is None or self._mapping.pp_size != 1 - or self._mapping.cp_size != 1 - or self._llm_args.attn_backend != "TRTLLM" or is_gemma4_hybrid( - self._model_engine.model.model_config.pretrained_config) - or self._llm_args.kv_cache_compression_config is not None): + spec_dec_mode = spec_config.spec_dec_mode + if spec_dec_mode.is_dflash(): + return spec_config.attention_backend + if (not spec_dec_mode.is_dspark() + or spec_config.draft_is_embedded_in_target): + return None + backend = spec_config.attention_backend + if backend == "AUTO": + # The drafter resolves AUTO against its own family and this build. + draft_model = getattr(self._model_engine.model, "draft_model", None) + backend = getattr(draft_model, "dflash_attention_backend", None) + return backend + + def _unified_drafter_name(self) -> str: + return ("DSpark" if self._speculative_config.spec_dec_mode.is_dspark() + else "DFlash") + + def _should_unify_dflash_kv_cache(self) -> bool: + """Whether this run can place target and DFlash or standalone DSpark + layers in one V2 manager. + + Disaggregated serving shares only when TRTLLM_DISAGG_SHARED_DRAFT_KV=1. + """ + if self._dflash_unify_blocker() is not None: return False + return not (self._is_disagg + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1") + + def _check_shared_draft_kv_request(self) -> None: + """Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored.""" + if (not self._is_disagg + or os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1" + or self._speculative_config is None + or self._unified_drafter_attention_backend() is None): + return + blocker = self._dflash_unify_blocker() + if blocker is not None: + raise ValueError( + "TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires the target and draft " + f"KV layers to share one KV cache manager, but {blocker}.") + + def _dflash_unify_blocker(self) -> Optional[str]: + """The reason target and drafter layers cannot share one manager, or None.""" + spec_config = self._speculative_config + if spec_config is None: + return "speculative decoding is not enabled" + backend = self._unified_drafter_attention_backend() + if backend != "TRTLLM": + return f"the drafter attention_backend is {backend}, not TRTLLM" + if not spec_config._allow_separate_draft_kv_cache: + return "a separate draft KV cache is not allowed for this configuration" + if spec_config.draft_len_schedule is not None: + return "draft_len_schedule is set" + if self._draft_model_engine is not None: + return "the drafter runs as a separate draft model engine" + if self._draft_config is None: + return "there is no draft model config" + if self._mapping.pp_size != 1: + return f"pipeline_parallel_size is {self._mapping.pp_size}, not 1" + if self._mapping.cp_size != 1: + return f"context_parallel_size is {self._mapping.cp_size}, not 1" + if self._llm_args.attn_backend != "TRTLLM": + return f"the target attn_backend is {self._llm_args.attn_backend}, not TRTLLM" + if is_gemma4_hybrid( + self._model_engine.model.model_config.pretrained_config): + return "the target uses Gemma 4 hybrid attention" + if self._llm_args.kv_cache_compression_config is not None: + return "KV cache compression is enabled" # Specialized managers add cache roles and lifecycle protocols that an # external drafter cannot yet describe. Keep their existing fallback. if self._kv_cache_manager_cls is not KVCacheManagerV2: - return False + return "the target does not use KVCacheManagerV2" target_layout = self._dflash_kv_layout( self._model_engine.model.model_config, self._model_engine.dtype) draft_model_config = self._get_draft_kv_model_config() @@ -2193,34 +2293,36 @@ def _should_unify_dflash_kv_cache(self) -> bool: draft_model_config.pretrained_config.torch_dtype) if target_layout is None or draft_layout is None: logger.info_once( - "DFlash: a per-layer KV-head count is not supported in a shared " + f"{self._unified_drafter_name()}: a per-layer KV-head count is " + "not supported in a shared " "target and draft KV cache; keeping a separate draft KV cache.", key="dflash_per_layer_kv_heads_fallback") - return False + return "the target or drafter has a per-layer KV-head count" + # The drafter's shared context block table indexes all its layers + # through one lifecycle; mixed windows need the separate manager. draft_windows = _resolve_v2_max_attention_window_vec( self._get_draft_max_attention_window(self._unification_max_seq_len, self._kv_cache_config), self._unification_max_seq_len, - range(draft_model_config.pretrained_config.num_hidden_layers)) + list( + self._mapping.pp_layers( + int(draft_model_config.pretrained_config.num_hidden_layers)) + )) if len(set(draft_windows)) > 1: - logger.info_once( - "DFlash: draft layers have different effective attention " - "windows; keeping a separate draft KV cache.", - key="dflash_mixed_draft_windows_fallback") - return False - if self._is_disagg and os.environ.get( - "TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": - return False + return "the draft layers use different attention windows" if self._is_disagg and not self._dflash_disagg_layouts_compatible( target_layout, draft_layout): logger.info_once( - "DFlash: draft KV layout (kv_heads, head_dim, dtype, k_only)=" + f"{self._unified_drafter_name()}: draft KV layout " + "(kv_heads, head_dim, dtype, k_only)=" f"{draft_layout} cannot be transferred with the target's " f"{target_layout}; disaggregated serving keeps a separate " "draft KV cache.", key="dflash_disagg_draft_layout_fallback") - return False - return True + return (f"the draft KV layout {draft_layout} cannot be transferred " + f"with the target's {target_layout} " + "(kv_heads, head_dim, dtype, k_only)") + return None @staticmethod def _dflash_disagg_layouts_compatible( 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..33604570b342 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 @@ -633,10 +633,36 @@ def _estimate_swa_cache_size( return size_per_token, size_per_request -def _get_generation_kv_capacity(spec_config, *, is_draft: bool) -> tuple[int, int]: - """Return draft-token reserve and total capacity headroom over history.""" +def _get_shared_draft_block_width(spec_config, layer_specs) -> int: + """Block width of the drafter whose layers share this manager, else 0.""" + from tensorrt_llm._torch.speculative import get_draft_block_width + + if layer_specs is not None and any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs + ): + return get_draft_block_width(spec_config) + return 0 + + +def _get_kv_alloc_extra_tokens(spec_config, layer_specs) -> int: + """Tokens each request is allocated beyond its history, for spec decoding. + + ``get_num_extra_kv_tokens`` is the position-side count. When the drafter's + layers share this manager, the drafter appends its whole block after the + request's context, so that block must also fit in the allocation. + """ from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens + return max( + get_num_extra_kv_tokens(spec_config), + _get_shared_draft_block_width(spec_config, layer_specs), + ) + + +def _get_generation_kv_capacity( + spec_config, *, is_draft: bool, layer_specs=None +) -> tuple[int, int]: + """Return draft-token reserve and total capacity headroom over history.""" if spec_config is None: return 0, BASE_GENERATION_TOKEN_COUNT reserve = spec_config.max_total_draft_tokens @@ -649,7 +675,9 @@ def _get_generation_kv_capacity(spec_config, *, is_draft: bool) -> tuple[int, in reserve = max(reserve, draft_loop_tokens) dynamic_reserve = reserve - spec_config.max_total_draft_tokens headroom = ( - get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step + dynamic_reserve + _get_kv_alloc_extra_tokens(spec_config, layer_specs) + + spec_config.tokens_per_gen_step + + dynamic_reserve ) return reserve, headroom @@ -1371,6 +1399,18 @@ class KVCacheManagerV2(BaseResourceManager): # True on hybrid Mamba/GDN managers; __init__ then coerces a positive # recompute tail to a full re-prefill. _has_recurrent_state: bool = False + # Block width of a drafter whose layers share this manager, else 0. + _shared_draft_block_width: int = 0 + + @property + def _kv_alloc_extra_tokens(self) -> int: + """Tokens a request is allocated beyond its history for spec decoding. + + ``num_extra_kv_tokens`` is the position-side count; a shared drafter's + block must also fit in the allocation. + """ + return max(self.num_extra_kv_tokens, self._shared_draft_block_width) + # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1572,6 +1612,9 @@ def __init__( from tensorrt_llm._torch.speculative import draft_prompt_lookahead, get_num_extra_kv_tokens self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + self._shared_draft_block_width = _get_shared_draft_block_width( + spec_config, self.layer_specs + ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 @@ -1598,7 +1641,9 @@ def __init__( ) # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( - _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) + _get_generation_kv_capacity( + spec_config, is_draft=self.is_draft, layer_specs=self.layer_specs + ) ) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size @@ -2220,7 +2265,7 @@ def _reserve_guard_page(self) -> None: # Never committed, so the reuse tree can never hand this page to a # request as a matched prefix. kv_cache.stop_committing() - if not kv_cache.resize(1 + self.num_extra_kv_tokens): + if not kv_cache.resize(1 + self._kv_alloc_extra_tokens): logger.warning("KVCacheManagerV2: could not size the guard page") self.kv_cache_map.pop(_GUARD_PAGE_REQUEST_ID, None) kv_cache.close() @@ -3173,7 +3218,7 @@ def _build_base_config( generation_request_capacity = self._get_generation_request_capacity() context_capacity = ( self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len - ) + self.num_extra_kv_tokens + ) + self._kv_alloc_extra_tokens generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) typical_step = BatchDesc( [KVCacheDesc(capacity=context_capacity, history_length=0)] @@ -3194,7 +3239,7 @@ def _build_base_config( BatchDesc( [ KVCacheDesc( - capacity=self.max_num_tokens + self.num_extra_kv_tokens, + capacity=self.max_num_tokens + self._kv_alloc_extra_tokens, history_length=0, ) ] @@ -3707,7 +3752,7 @@ def get_num_available_tokens( replicates all tokens on every rank, so both bounds constrain the same request-length variable. """ - extra_tokens = self.num_extra_kv_tokens + max_num_draft_tokens + extra_tokens = self._kv_alloc_extra_tokens + max_num_draft_tokens # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. # We need to add extra tokens to the token num upper bound to account for the extra tokens. clamped = ( @@ -4227,7 +4272,7 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: if kv_cache is None: return False - target = req.context_current_position + num_tokens + self.num_extra_kv_tokens + target = req.context_current_position + num_tokens + self._kv_alloc_extra_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -4267,7 +4312,7 @@ def try_allocate_draft_context(self, req: LlmRequest, num_tokens: int) -> bool: if kv_cache is None or not self._resume_and_restore(req.py_request_id, kv_cache): return False - target = req.context_current_position + num_tokens + self.num_extra_kv_tokens + target = req.context_current_position + num_tokens + self._kv_alloc_extra_tokens return kv_cache.resize(max(kv_cache.capacity, target)) def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: @@ -4304,7 +4349,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: else (req.total_input_len_cp if self._has_cp_helix else req.prompt_len) ) prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens + target = prompt_len + get_draft_token_length(req) + self._kv_alloc_extra_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -4875,7 +4920,7 @@ def _resize_for_served_prefix(self, req: LlmRequest, kv_cache, position: int, en sliding-window layer group. When the capacity is unchanged the call cannot fail, since no new slot is requested. """ - target = end + self.num_extra_kv_tokens + target = end + self._kv_alloc_extra_tokens if end >= req.prompt_len: target += get_draft_token_length(req) pre_cap = kv_cache.capacity @@ -5019,7 +5064,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): req.context_current_position + req.context_chunk_size + draft_len - + self.num_extra_kv_tokens + + self._kv_alloc_extra_tokens ) if not kv_cache.resize(capacity): raise RuntimeError( @@ -5865,7 +5910,7 @@ def release_resources( release_resources(req) return None kv_cache.stop_committing() - dummy_capacity = token_num + self.num_extra_kv_tokens + dummy_capacity = token_num + self._kv_alloc_extra_tokens if is_gen and not materialize_history: kv_cache.enable_swa_scratch_reuse = False # Need to hint the committed history to activate stale-block @@ -6438,7 +6483,7 @@ def get_cache_size_per_token( max_seq_len, ) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft + spec_config, is_draft=is_draft, layer_specs=layer_specs ) ( context_size_per_token, @@ -6624,7 +6669,7 @@ def _get_per_request_layer_tokens( else: retained_windows = list(covered_windows) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft + spec_config, is_draft=is_draft, layer_specs=layer_specs ) full_tokens = ( math.ceil((max_seq_len + generation_capacity_headroom) / tokens_per_block) diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 382cb09e1571..c3ffe7e28bda 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -18,10 +18,11 @@ SpecSampler) from .spec_tree_manager import SpecTreeManager from .suffix_automaton import SuffixAutomatonManager -from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, get_spec_drafter, - get_spec_metadata, get_spec_resource_manager, - get_spec_worker, update_spec_config_from_draft_model_config, +from .utils import (get_draft_block_width, get_draft_kv_cache_manager, + get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, get_spec_drafter, get_spec_metadata, + get_spec_resource_manager, get_spec_worker, + update_spec_config_from_draft_model_config, update_spec_config_from_loaded_model, update_spec_config_from_model_config) @@ -50,6 +51,7 @@ "SpecWorkerBase", "draft_prompt_lookahead", "get_draft_kv_cache_manager", + "get_draft_block_width", "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index c7cc1af494be..fbc70a53c8ab 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -632,9 +632,16 @@ def prepare(self): worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) - # to dummy slot to avoid corrupting real request's context + # to dummy slot to avoid corrupting real request's context. A + # generation row with the idle id is a dummy even when a context + # warmup stored a prompt under that id: its length is not this row's. mapping = torch.tensor( - [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], + [ + worker._dummy_slot + if i >= num_contexts and rid == ATTENTION_DP_DUMMY_REQUEST_ID + else worker._req_to_slot.get(rid, worker._dummy_slot) + for i, rid in enumerate(self.request_ids) + ], dtype=torch.long, device="cpu", pin_memory=prefer_pinned(), @@ -994,26 +1001,28 @@ def _init_ctx_block_tables( if isinstance(mgr, KVCacheManagerV2): converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + unsupported = None if mgr.enable_swa_scratch_reuse: - raise NotImplementedError( - "DFlash managed context does not support SWA scratch pages" - ) + unsupported = "DFlash managed context does not support SWA scratch pages" for layer_id in self._ctx_layer_ids: local_id = mgr.layer_offsets[layer_id] if ( mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id or mgr.get_layer_page_index_scale(layer_id) != converter.scale ): - if self._ctx_unified_cache: - raise NotImplementedError( - "DFlash's shared context block table requires draft layers " - "with the same lifecycle and page-index scale" - ) - logger.warning( - "DFlash: draft KV layers have different lifecycles or page-index " - "scales; using the private ctx arena." + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" ) - return False + break + if unsupported is not None: + # A unified manager was selected for this layout, so the drafter + # cannot move to private storage without breaking allocation, + # reuse and transfer. A separate manager keeps the private arena. + if self._ctx_unified_cache: + raise NotImplementedError(unsupported) + logger.warning(f"{unsupported}; using the private ctx arena.") + return False if per_layer_tables: # Attention tables address the physical pool origin, whereas # get_buffers() starts at this layer's own K buffer. @@ -1067,6 +1076,23 @@ def _init_ctx_block_tables( ) return True + @staticmethod + def _check_block_within_allocation( + num_ctx: torch.Tensor, allocated_ctx_limit: torch.Tensor, block_size: int + ) -> None: + """Raise when a request's drafter block would leave its allocated pages. + + The shared manager sizes the request, so a shortfall there would have the + block's scratch write land on the last context position, or on another + request's page. Asynchronous, so it adds no host sync and is safe to + capture in a CUDA graph; on the device it traps instead of raising. + """ + torch._assert_async( + (num_ctx <= allocated_ctx_limit).all(), + f"DFlash: the drafter's block of {block_size} positions falls outside the " + "pages the shared KV cache manager allocated for the request", + ) + def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """Decode this iteration's draft block table into the persistent buffer. @@ -1961,7 +1987,12 @@ def _draft_block_width(self, draft_model) -> int: this alongside :meth:`_draft_slot_ids` — the two must agree, since the width bounds the slots the gather is allowed to name. """ - return self.max_draft_len + 1 + return self.block_width_for(self.spec_config) + + @staticmethod + def block_width_for(spec_config) -> int: + """Block slots a worker of this family computes for ``spec_config``.""" + return spec_config.max_draft_len + 1 def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int @@ -2283,6 +2314,10 @@ def prepare_1st_drafter_inputs( allocated = dflash_allocated_ctx_limit( self._ctx_block_counts, self._ctx_page_size, block_size ) + if self._ctx_unified_cache: + self._check_block_within_allocation( + num_ctx_per_req_t, allocated[gen_rows_out], block_size + ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) noise_embedding = noise_embed_2d query_positions = query_position_ids.long() diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 1839de961f8e..b5e67034b10d 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -1209,6 +1209,16 @@ def set_draft_model(self, draft_model) -> None: f"{type(draft_model).__name__} declares one." ) + @staticmethod + def block_width_for(spec_config) -> int: + """Block slots a DSpark worker computes for ``spec_config``. + + The ``shift_label`` convention, which the DSpark drafter defaults to, + reads slots 0..K-1; a checkpoint that turns it off needs K+1 and is + caught by the block-within-allocation check. + """ + return spec_config.max_draft_len + def _draft_block_width(self, draft_model) -> int: """Block width under the dspark ``shift_label`` convention. @@ -1218,8 +1228,8 @@ def _draft_block_width(self, draft_model) -> int: drafters are meant to run. """ if getattr(draft_model, "_dspark_shift_label", False): - return self.max_draft_len - return super()._draft_block_width(draft_model) + return self.block_width_for(self.spec_config) + return DFlashWorker.block_width_for(self.spec_config) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 454f2567d000..8a5987078b3f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -791,6 +791,21 @@ def get_num_extra_kv_tokens(spec_config): return 0 +def get_draft_block_width(spec_config) -> int: + """Block width the worker drafting for ``spec_config`` appends per request. + + 0 when that worker does not append a block to a paged draft KV cache. + """ + if spec_config is None: + return 0 + if spec_config.spec_dec_mode.is_dflash(): + return DFlashWorker.block_width_for(spec_config) + if (spec_config.spec_dec_mode.is_dspark() + and not spec_config.draft_is_embedded_in_target): + return DSparkWorker.block_width_for(spec_config) + return 0 + + def get_draft_kv_cache_manager(spec_config, resource_manager): """ Return the one-model draft cache manager, including a shared draft KV view diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index c58c6741dfaf..8394109aaeb8 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2505,6 +2505,32 @@ def set_cuda_device(): torch.cuda.set_device(device_id) +def _shared_draft_kv_blocker(server_configs) -> Optional[str]: + """What keeps the context and generation servers from sharing draft KV, or None.""" + with_spec = { + cfg.type + for cfg in server_configs if cfg.other_args.get("speculative_config") + } + if with_spec != {"ctx", "gen"}: + return ("speculative_config in both context_servers and " + "generation_servers") + if len({(cfg.other_args.get("tensor_parallel_size", 1), + bool(cfg.other_args.get("enable_attention_dp", False))) + for cfg in server_configs}) != 1: + return ("the same tensor_parallel_size and enable_attention_dp in " + "context_servers and generation_servers") + for cfg in server_configs: + spec = cfg.other_args["speculative_config"] + backend = spec.get("attention_backend") + if not ((spec.get("decoding_type") == "DFlash" and backend == "TRTLLM") + or (spec.get("decoding_type") == "DSpark" + and backend in (None, "AUTO", "TRTLLM"))): + return ("speculative_config to be DFlash with attention_backend " + "TRTLLM, or DSpark with attention_backend AUTO or TRTLLM, " + "in both sections") + return None + + @click.command("disaggregated_mpi_worker") @stability_option( "-c", @@ -2544,19 +2570,12 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, from tensorrt_llm.llmapi.disagg_utils import split_world_comm disagg_cfg = parse_disagg_config_file(config_file) - with_spec = { - cfg.type - for cfg in disagg_cfg.server_configs - if cfg.other_args.get("speculative_config") - } - attention_parallel_configs = {(cfg.other_args.get("tensor_parallel_size", - 1), - cfg.other_args.get("enable_attention_dp", - False)) - for cfg in disagg_cfg.server_configs} - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ( - "1" if with_spec == {"ctx", "gen"} - and len(attention_parallel_configs) == 1 else "0") + blocker = _shared_draft_kv_blocker(disagg_cfg.server_configs) + if blocker is not None and os.environ.get( + "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1": + raise click.UsageError( + f"TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires {blocker}.") + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" if blocker else "1") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 81d893199157..c89dd523d4b0 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1694,9 +1694,9 @@ def test_gsm8k_1p1d_unified_draft_kv(self): ctx_server_config, gen_server_config, self.MODEL_PATH, - extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}, assert_worker_log_contains= - "target and draft KV layers share one KV cache manager") as llm: + "target and draft KV layers share one KV cache manager", + extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}) as llm: llm.args.speculative_config = DFlashDecodingConfig( max_draft_len=self.MAX_DRAFT_LEN, speculative_model=self.DRAFT_PATH, diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index c91507098ac3..5151d6f5b407 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -16,6 +16,7 @@ from types import SimpleNamespace +import click import pytest import torch import yaml @@ -52,7 +53,7 @@ def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): ) -def _creator(*, is_disagg, target, draft, tp_size=1): +def _creator(*, is_disagg, target, draft, tp_size=1, draft_windows=None): creator = object.__new__(KvCacheCreator) creator._speculative_config = SimpleNamespace( spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), @@ -63,8 +64,15 @@ def _creator(*, is_disagg, target, draft, tp_size=1): creator._draft_model_engine = None creator._draft_config = draft creator._mapping = SimpleNamespace( - pp_size=1, cp_size=1, tp_size=tp_size, enable_attention_dp=False + pp_size=1, + cp_size=1, + tp_size=tp_size, + enable_attention_dp=False, + pp_layers=lambda num_layers: list(range(num_layers)), ) + creator._max_seq_len = 4096 + creator._kv_cache_config = SimpleNamespace() + creator._get_draft_max_attention_window = lambda max_seq_len, kv_cache_config: draft_windows creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) creator._model_engine = SimpleNamespace( model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 @@ -122,19 +130,22 @@ def test_layout_compares_rank_local_heads(): @pytest.mark.parametrize("is_disagg", [False, True]) -def test_attention_dp_unifies_with_unsharded_heads(is_disagg): - # Under attention DP each rank keeps every KV head, so the TP4 layouts of - # an 8-head target and a 2-head draft are not resharded to 2 and 1. - target = _model_config(kv_heads=8, head_dim=128) - draft = _model_config(kv_heads=8, head_dim=128) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_attention_dp_changes_the_disagg_layout_decision(is_disagg, attention_dp): + # At TP4 a 4-head target and a 1-head draft both hold one head per rank, so + # they transfer together. Under attention DP each rank keeps every head (4 + # and 1), so the layouts differ and disaggregated serving keeps a separate cache. + target = _model_config(kv_heads=4, head_dim=128) + draft = _model_config(kv_heads=1, head_dim=128) creator = _creator(is_disagg=is_disagg, target=target, draft=draft, tp_size=4) - creator._mapping.enable_attention_dp = True + creator._mapping.enable_attention_dp = attention_dp - assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 8 - assert ( - creator._dflash_kv_layout(_model_config(kv_heads=2, head_dim=128), torch.bfloat16)[0] == 2 + heads = ( + creator._dflash_kv_layout(target, torch.bfloat16)[0], + creator._dflash_kv_layout(draft, torch.bfloat16)[0], ) - assert creator._should_unify_dflash_kv_cache() + assert heads == ((4, 1) if attention_dp else (1, 1)) + assert creator._should_unify_dflash_kv_cache() == (not (is_disagg and attention_dp)) @pytest.mark.parametrize("is_disagg", [False, True]) @@ -177,19 +188,21 @@ def stop(*args, **kwargs): monkeypatch.setattr(serve, "set_cuda_device", lambda: None) monkeypatch.setattr("tensorrt_llm.llmapi.disagg_utils.split_world_comm", stop) monkeypatch.delenv("TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT", raising=False) - with pytest.raises(_StopLauncher): - serve.disaggregated_mpi_worker.callback( - config_file=str(config), log_level="info", telemetry=False - ) + serve.disaggregated_mpi_worker.callback( + config_file=str(config), log_level="info", telemetry=False + ) def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_path): - _launch_from_disagg_config( - monkeypatch, - tmp_path, - context_servers={}, - generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, - ) + # The launcher rejects the variable when only one section sets speculative_config. + with pytest.raises(click.UsageError, match="both context_servers and generation_servers"): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={}, + generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, + ) + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=8, head_dim=128) @@ -200,15 +213,88 @@ def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_ def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): spec = {"decoding_type": "DFlash"} - _launch_from_disagg_config( - monkeypatch, - tmp_path, - context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, - generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, - ) + with pytest.raises(click.UsageError, match="same tensor_parallel_size and enable_attention_dp"): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, + generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, + ) + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=4, head_dim=128) for tp_size in (4, 8): creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=tp_size) assert not creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_draft_layers_with_different_windows_keep_separate_cache(is_disagg): + layout = _model_config(kv_heads=8, head_dim=128) + + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout, draft_windows=[128, 4096]) + + assert not creator._should_unify_dflash_kv_cache() + assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_draft_layers_with_one_window_still_unify(is_disagg): + layout = _model_config(kv_heads=8, head_dim=128) + + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout, draft_windows=[128, 128]) + + assert creator._should_unify_dflash_kv_cache() + + +def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=[128, 4096]) + + with pytest.raises(ValueError, match="different attention windows"): + creator._check_shared_draft_kv_request() + + +@pytest.mark.parametrize( + "context_servers,generation_servers,match", + [ + ( + {"enable_attention_dp": True}, + {"enable_attention_dp": False}, + "same tensor_parallel_size and enable_attention_dp", + ), + ( + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + "attention_backend TRTLLM", + ), + ( + {"speculative_config": {"decoding_type": "DFlash"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}}, + "attention_backend TRTLLM", + ), + ], + ids=["mismatched_attention_dp", "vanilla_dflash_backend", "default_dflash_backend"], +) +def test_launcher_rejects_shared_cache_for_ineligible_config( + monkeypatch, tmp_path, context_servers, generation_servers, match +): + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + context_servers = {"speculative_config": spec, **context_servers} + generation_servers = {"speculative_config": spec, **generation_servers} + + with pytest.raises(click.UsageError, match=match): + _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers) + + +def test_launcher_accepts_eligible_shared_cache_config(monkeypatch, tmp_path): + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + {"speculative_config": spec, "enable_attention_dp": True}, + {"speculative_config": spec, "enable_attention_dp": True}, + ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 9ba390d40484..0066252e386a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -329,6 +329,9 @@ def get_num_attention_layers(self) -> int: # A Mock's truthy return would send get_num_extra_kv_tokens down the # one-engine arm; this test's config is a two-model drafter. mode.use_one_engine.return_value = False + # Likewise a truthy is_dflash() would route costing through the shared + # target and drafter cache path. + mode.is_dflash.return_value = False costed_windows: list[tuple[object, list[int] | None]] = [] class RecordingKVCacheManager(KVCacheManagerV2): 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 739ca90d6d20..d00eef883c99 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 @@ -27,6 +27,7 @@ _create_kv_cache_manager, _derive_v2_layer_type_attention_windows, _get_num_pool_groups_for_estimation, + _get_num_storage_pool_groups, get_mla_context_workspace_kv_len_cap, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window @@ -1735,3 +1736,55 @@ def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> Non assert manager_config is kv_cache_config assert manager_config.max_attention_window is None + + +def _fp8_target_bf16_draft_specs(*, draft_window): + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + return [ + *(replace(target, global_layer_id=i) for i in range(3)), + replace( + target, + global_layer_id=3, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + sliding_window_size=draft_window, + ), + ] + + +def test_unified_fp8_target_bf16_draft_scales_by_combined_pool_groups(): + """The V2 pool-group count comes from the combined target and draft layout. + + The FP8 target and the windowed BF16 draft differ in life cycle and slot + size, so they land in two pool groups; counting the target's model config + alone gives one. + """ + tpb = 32 + spec_config = SimpleNamespace( + tokens_per_gen_step=1, + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, use_one_engine=lambda: False), + ) + baseline = _make_creator(tpb, [_make_mock_request(128)], False, 1) + baseline._speculative_config = spec_config + baseline._build_unified_dflash_layer_specs = lambda kv_cache_config: () + unified = _make_creator(tpb, [_make_mock_request(128)], False, 1) + unified._speculative_config = spec_config + specs = _fp8_target_bf16_draft_specs(draft_window=128) + unified._build_unified_dflash_layer_specs = lambda kv_cache_config: tuple(specs) + + assert _get_num_storage_pool_groups(specs, tpb) == 2 + assert unified._get_token_num_for_estimation() == 2 * baseline._get_token_num_for_estimation() + + +def test_unified_fp8_target_bf16_draft_in_one_life_cycle_is_one_pool_group(): + """Pages of different sizes in one life cycle share a pool group.""" + specs = _fp8_target_bf16_draft_specs(draft_window=None) + + assert _get_num_storage_pool_groups(specs, 32) == 1 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 78dbec7cb8de..d14f7a3e464a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -24,6 +24,7 @@ import torch from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.pyexecutor._util import _get_num_storage_pool_groups from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory 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 ( @@ -33,6 +34,8 @@ KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _KVCacheManagerInitStatus, _swa_endpoint_priority, _sync_kv_cache_manager_init_status, @@ -3759,3 +3762,181 @@ def get_base_page_indices(self, layer_group_id, beam_id=DEFAULT_BEAM_INDEX): assert len(fake_logger.messages) == 1 assert "unavailable" in fake_logger.messages[0] + + +@pytest.mark.parametrize( + "target_dtype,target_layers,draft_windows", + [ + (DataType.BF16, 2, [None]), + (DataType.FP8, 2, [None]), + (DataType.FP8, 3, [128]), + (DataType.BF16, 2, [128]), + ], + ids=["bf16_same_window", "fp8_same_window", "fp8_windowed_draft", "bf16_windowed_draft"], +) +def test_storage_pool_group_count_matches_manager( + target_dtype, target_layers, draft_windows +) -> None: + """The estimator's pool-group count is what the manager creates.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=target_dtype, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + layer_specs = [replace(target, global_layer_id=i) for i in range(target_layers)] + layer_specs += [ + replace( + target, + global_layer_id=target_layers + i, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + sliding_window_size=window, + ) + for i, window in enumerate(draft_windows) + ] + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_swa_scratch_reuse=False, + ) + manager = KVCacheManagerV2( + kv_cache_config, + CacheType.SELF, + num_layers=len(layer_specs), + num_kv_heads=8, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=2, + max_num_tokens=128, + mapping=Mapping(), + dtype=target_dtype, + layer_specs=layer_specs, + ) + try: + assert _get_num_storage_pool_groups(layer_specs, 32) == len(manager.impl.pool_group_descs) + finally: + manager.shutdown() + + +@pytest.mark.parametrize("disagg", [False, True], ids=["aggregated", "disaggregated"]) +def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg: bool) -> None: + """The drafter appends its block after the context, so its pages must reach it. + + After a decode step accepts every draft token the context is ``prompt + K + + 1`` long and the drafter's block follows it. The prompt is chosen so that + this furthest write is the first position of a third page. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + max_draft_len = 7 + block_width = max_draft_len + 1 + tokens_per_block = 32 + prompt_len = 49 + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft_layer = 2 + layer_specs = [ + target, + replace(target, global_layer_id=1), + replace( + target, global_layer_id=draft_layer, owner=KVCacheLayerOwner.DRAFT, dtype=DataType.BF16 + ), + ] + manager = KVCacheManagerV2( + KvCacheConfig( + enable_block_reuse=False, + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=len(layer_specs), + num_kv_heads=8, + head_dim=128, + tokens_per_block=tokens_per_block, + max_seq_len=1024, + max_batch_size=2, + mapping=Mapping(), + dtype=DataType.FP8, + layer_specs=layer_specs, + spec_config=DFlashDecodingConfig(max_draft_len=max_draft_len, speculative_model="draft"), + vocab_size=4096, + enable_stats=False, + ) + request = _ContextRequest( + 1, + list(range(prompt_len)), + prompt_len, + "conversation", + use_conversation_params=False, + is_disagg_generation_init_state=disagg, + ) + try: + if disagg: + request.py_draft_tokens = [] + request.is_disagg_generation_transmission_complete = False + request.context_phase_params = None + assert manager.prepare_disagg_gen_init(request) + request.is_disagg_generation_init_state = False + else: + _run_context(manager, request) + generation = SimpleNamespace( + py_request_id=request.py_request_id, + py_beam_width=1, + py_draft_tokens=[0] * max_draft_len, + state=LlmRequestState.GENERATION_IN_PROGRESS, + is_dummy_request=False, + is_disagg_generation_transmission_complete=False, + context_phase_params=None, + max_beam_num_tokens=prompt_len + 1, + ) + assert manager.try_allocate_generation(generation) + + furthest_write_end = prompt_len + max_draft_len + 1 + block_width + pages = manager.get_batch_cache_indices([request.py_request_id], layer_idx=draft_layer)[0] + assert len(pages) * tokens_per_block >= furthest_write_end + finally: + manager.free_resources(request) + manager.shutdown() + + +def test_cost_headroom_counts_the_drafter_block_only_for_shared_layers() -> None: + """The budget charges the allocation headroom, not the position-side count.""" + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens + + max_draft_len = 7 + spec_config = DFlashDecodingConfig(max_draft_len=max_draft_len, speculative_model="draft") + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft = replace(target, global_layer_id=1, owner=KVCacheLayerOwner.DRAFT) + + _, own = _get_generation_kv_capacity(spec_config, is_draft=False, layer_specs=[target]) + _, shared = _get_generation_kv_capacity( + spec_config, is_draft=False, layer_specs=[target, draft] + ) + + # The block is K + 1 = 8 wide against K - 1 = 6 extra positions. + assert shared - own == (max_draft_len + 1) - (max_draft_len - 1) + assert _get_kv_alloc_extra_tokens(spec_config, [target, draft]) == max_draft_len + 1 + assert get_num_extra_kv_tokens(spec_config) == max_draft_len - 1 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py new file mode 100644 index 000000000000..10cb0365548e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -0,0 +1,98 @@ +# 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. +"""DFlash shared context block table over a KVCacheManagerV2 draft layout. + +A separate draft manager whose layers cannot share one block table falls back to +the private context arena, as it did before unified caching. A unified manager +keeps its invariant, so the same layout there is an error. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.speculative.dflash import DFlashWorker + +pytestmark = pytest.mark.cpu_only + + +def _draft_manager(*, layer_groups, layer_scales): + mgr = object.__new__(KVCacheManagerV2) + mgr.kv_cache_pool_mapping = torch.zeros(len(layer_groups), 2, dtype=torch.int32) + mgr.layer_offsets = {layer_id: layer_id for layer_id in range(len(layer_groups))} + mgr.layer_to_pool_mapping_dict = dict(enumerate(layer_groups)) + mgr.enable_swa_scratch_reuse = False + mgr.impl = SimpleNamespace( + get_page_index_converter=lambda layer_offset, role: SimpleNamespace( + scale=layer_scales[layer_offset], layer_offset=0 + ), + get_page_index_scale=lambda layer_offset, role: layer_scales[layer_offset], + ) + return mgr + + +def _worker(*, unified): + worker = object.__new__(DFlashWorker) + worker._ctx_layer_ids = (0, 1) + worker._ctx_unified_cache = unified + worker._ctx_kv_factor = 2 + return worker + + +def _init_block_tables(worker, mgr): + pool = [torch.empty(0)] + return worker._init_ctx_block_tables(mgr, pool, num_slots=2, L=2, nkv=1, hd=64, page_size=32) + + +@pytest.mark.parametrize( + "layer_groups,layer_scales", + [([0, 1], [2, 2]), ([0, 0], [2, 4])], + ids=["different_lifecycle", "different_page_index_scale"], +) +def test_separate_v2_draft_manager_with_unsupported_layout_uses_private_arena( + layer_groups, layer_scales +): + mgr = _draft_manager(layer_groups=layer_groups, layer_scales=layer_scales) + + assert _init_block_tables(_worker(unified=False), mgr) is False + + +@pytest.mark.parametrize( + "layer_groups,layer_scales", + [([0, 1], [2, 2]), ([0, 0], [2, 4])], + ids=["different_lifecycle", "different_page_index_scale"], +) +def test_unified_manager_with_unsupported_layout_is_an_error(layer_groups, layer_scales): + mgr = _draft_manager(layer_groups=layer_groups, layer_scales=layer_scales) + + with pytest.raises(NotImplementedError, match="same lifecycle and page-index scale"): + _init_block_tables(_worker(unified=True), mgr) + + +def test_drafter_block_inside_the_allocation_is_accepted(): + num_ctx = torch.tensor([57, 10]) + allocated = torch.tensor([57, 56]) + + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def test_drafter_block_outside_the_allocation_is_an_error(): + num_ctx = torch.tensor([58, 10]) + allocated = torch.tensor([57, 56]) + + with pytest.raises(RuntimeError, match="block of 7 positions falls outside"): + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) From 238ff33ddc63f64eb945e953154a047411b6031e Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 11:47:11 -0700 Subject: [PATCH 28/35] docs: clarify disaggregated draft KV opt-in Signed-off-by: Allison Lim --- docs/source/features/speculative-decoding.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 65b1a52168a8..eea42ab78f40 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -175,7 +175,14 @@ both and each request's draft pages are allocated, padded and freed with its target pages. The draft layers keep the drafter's own dtype, so an FP8 target KV cache (`kv_cache_config.dtype: fp8`) is paired with a BF16 draft cache. -In disaggregated serving, the context server's draft pages are transferred to +In disaggregated serving, direct worker launches require +`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` on both the context and generation workers +to enable shared draft-KV transfer. The `disaggregated_mpi_worker` launcher +enables this setting automatically when both servers enable compatible +speculation and use the same `tensor_parallel_size` and +`enable_attention_dp`. + +When sharing is enabled, the context server's draft pages are transferred to the generation server with the target pages, so the drafter starts generation with the prompt context instead of an empty one. This requires the draft layers to have the target's rank-local KV-head count; the draft dtype and head From 2bcd2b269eee733e0fd1764ad2419898a19779b0 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 11:54:42 -0700 Subject: [PATCH 29/35] test: retain engine stats for DFlash acceptance length Signed-off-by: Allison Lim --- tests/integration/defs/accuracy/test_disaggregated_serving.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index e8d1a8442d25..1a4e60a8b99a 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1695,6 +1695,7 @@ def test_gsm8k_1p1d_unified_draft_kv(self): }, # Acceptance length is read back off this worker's /metrics. "enable_iter_perf_stats": True, + "max_stats_len": -1, "iter_stats_max_iterations": -1, } disaggregated_server_config = { From 76c4cefe737f2d4ab6e9b157e671ab5d55db0eff Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 16:51:15 -0700 Subject: [PATCH 30/35] fix: size shared DSpark KV from checkpoint block width Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 24 +++- .../kv_cache/kv_cache_manager_v2.py | 49 ++++++-- tensorrt_llm/_torch/speculative/dflash.py | 36 +++--- tensorrt_llm/_torch/speculative/dspark.py | 15 ++- tensorrt_llm/_torch/speculative/utils.py | 5 +- .../kv_cache/test_dflash_unified_kv_gate.py | 56 +++++++++ .../kv_cache/test_kv_cache_estimation.py | 88 ++++++++++++++ .../kv_cache/test_kvcm2_integration.py | 19 ++- .../test_dflash_ctx_arena_sizing.py | 2 +- .../test_dflash_ctx_block_tables.py | 109 ++++++++++++++++++ 10 files changed, 359 insertions(+), 44 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0ace73eb17aa..c3d510246ce2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -52,9 +52,9 @@ from ..model_config import ModelConfig, _get_kv_cache_layer_specs from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs -from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, - should_use_separate_draft_kv_cache) +from ..speculative import (draft_prompt_lookahead, get_draft_block_width, + get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, should_use_separate_draft_kv_cache) from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, @@ -1144,6 +1144,8 @@ def _per_manager_cache_cost(self, extra_kwargs[ "layer_specs"] = self._build_unified_dflash_layer_specs( kv_cache_config) + extra_kwargs.update( + shared_draft_block_width=self._unified_drafter_block_width()) if not is_draft: derived_windows = _derive_v2_layer_type_attention_windows( kv_cache_config, manager_cls, model_config, self._max_seq_len) @@ -2111,9 +2113,11 @@ def _create_kv_cache_manager( estimating_kv_cache = estimating_kv_cache and not self._skip_est self._check_shared_draft_kv_request() unified_layer_specs = None + shared_draft_block_width = None if self._should_unify_dflash_kv_cache(): unified_layer_specs = self._build_unified_dflash_layer_specs( kv_cache_config) + shared_draft_block_width = self._unified_drafter_block_width() logger.info( f"{self._unified_drafter_name()}: target and draft KV layers " "share one KV cache manager " @@ -2144,6 +2148,7 @@ def _create_kv_cache_manager( cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, layer_specs=unified_layer_specs, + shared_draft_block_width=shared_draft_block_width, ) if not self._skip_est: @@ -2228,6 +2233,13 @@ def _unified_drafter_attention_backend(self) -> Optional[str]: backend = getattr(draft_model, "dflash_attention_backend", None) return backend + def _unified_drafter_block_width(self) -> int: + """Use the loaded checkpoint's slot convention for sizing and allocation.""" + draft_model = getattr(self._model_engine.model, "draft_model", None) + shift_label = getattr(draft_model, "_dspark_shift_label", False) + return get_draft_block_width(self._speculative_config, + shift_label=shift_label) + def _unified_drafter_name(self) -> str: return ("DSpark" if self._speculative_config.spec_dec_mode.is_dspark() else "DFlash") @@ -3594,8 +3606,8 @@ def _create_kv_cache_manager( kv_events_config: Optional[KVEventsConfig] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, - layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None -) -> KVCacheManager: + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -3749,6 +3761,8 @@ def _create_kv_cache_manager( manager_extra_kwargs[ "disable_overlap_scheduler"] = disable_overlap_scheduler manager_extra_kwargs["layer_specs"] = layer_specs + manager_extra_kwargs[ + "shared_draft_block_width"] = shared_draft_block_width # Vocab size also enables multimodal event decoding and its per-block # digest scan. Leave it unset for text-only engines. Without an engine # (e.g. separate one-model draft caches), preserve config resolution: 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 a2bb78f92253..c17c811abd9f 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 @@ -649,18 +649,27 @@ def _estimate_swa_cache_size( return size_per_token, size_per_request -def _get_shared_draft_block_width(spec_config, layer_specs) -> int: - """Block width of the drafter whose layers share this manager, else 0.""" +def _get_shared_draft_block_width( + spec_config, layer_specs, shared_draft_block_width: Optional[int] = None +) -> int: + """Block width of the drafter whose layers share this manager, else 0. + + The creator supplies the checkpoint-resolved width when available. + """ from tensorrt_llm._torch.speculative import get_draft_block_width if layer_specs is not None and any( layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs ): + if shared_draft_block_width is not None: + return shared_draft_block_width return get_draft_block_width(spec_config) return 0 -def _get_kv_alloc_extra_tokens(spec_config, layer_specs) -> int: +def _get_kv_alloc_extra_tokens( + spec_config, layer_specs, shared_draft_block_width: Optional[int] = None +) -> int: """Tokens each request is allocated beyond its history, for spec decoding. ``get_num_extra_kv_tokens`` is the position-side count. When the drafter's @@ -671,12 +680,16 @@ def _get_kv_alloc_extra_tokens(spec_config, layer_specs) -> int: return max( get_num_extra_kv_tokens(spec_config), - _get_shared_draft_block_width(spec_config, layer_specs), + _get_shared_draft_block_width(spec_config, layer_specs, shared_draft_block_width), ) def _get_generation_kv_capacity( - spec_config, *, is_draft: bool, layer_specs=None + spec_config, + *, + is_draft: bool, + layer_specs=None, + shared_draft_block_width: Optional[int] = None, ) -> tuple[int, int]: """Return draft-token reserve and total capacity headroom over history.""" if spec_config is None: @@ -691,7 +704,7 @@ def _get_generation_kv_capacity( reserve = max(reserve, draft_loop_tokens) dynamic_reserve = reserve - spec_config.max_total_draft_tokens headroom = ( - _get_kv_alloc_extra_tokens(spec_config, layer_specs) + _get_kv_alloc_extra_tokens(spec_config, layer_specs, shared_draft_block_width) + spec_config.tokens_per_gen_step + dynamic_reserve ) @@ -1561,6 +1574,7 @@ def __init__( joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> None: self.mapping = mapping @@ -1726,7 +1740,7 @@ def __init__( self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) self._shared_draft_block_width = _get_shared_draft_block_width( - spec_config, self.layer_specs + spec_config, self.layer_specs, shared_draft_block_width ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. @@ -1755,7 +1769,10 @@ def __init__( # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( _get_generation_kv_capacity( - spec_config, is_draft=self.is_draft, layer_specs=self.layer_specs + spec_config, + is_draft=self.is_draft, + layer_specs=self.layer_specs, + shared_draft_block_width=self._shared_draft_block_width, ) ) @@ -7106,6 +7123,7 @@ def get_cache_size_per_token( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ): if layer_specs is not None: @@ -7141,7 +7159,10 @@ def get_cache_size_per_token( max_seq_len, ) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft, layer_specs=layer_specs + spec_config, + is_draft=is_draft, + layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) ( context_size_per_token, @@ -7224,6 +7245,7 @@ def _get_per_request_layer_tokens( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = 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. @@ -7327,7 +7349,10 @@ def _get_per_request_layer_tokens( else: retained_windows = list(covered_windows) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft, layer_specs=layer_specs + spec_config, + is_draft=is_draft, + layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) full_tokens = ( math.ceil((max_seq_len + generation_capacity_headroom) / tokens_per_block) @@ -7372,6 +7397,7 @@ def get_cache_bytes_per_request( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -7405,6 +7431,7 @@ def get_cache_bytes_per_request( spec_config=spec_config, is_draft=is_draft, layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) if layer_tokens is None: return None @@ -7434,6 +7461,7 @@ def get_cache_quota_per_request( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -7469,6 +7497,7 @@ def get_cache_quota_per_request( spec_config=spec_config, is_draft=is_draft, layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) if layer_tokens is None: return None diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 4a7354d8ac4d..4f8c2113d237 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1152,6 +1152,28 @@ def _lazy_init_ctx_buffers( logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") self._ctx_buf_inited = False + # checkpoint's trained block width + self._resolved_block_size = getattr(draft_model, "block_size", None) or ( + self.max_draft_len + 1 + ) + # what the draft forward actually computes + self._compute_block_size = self._draft_block_width(draft_model) + if self._compute_block_size > self._resolved_block_size: + slack = self._compute_block_size - self.max_draft_len + raise ValueError( + f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." + f"Lower max_draft_len to at most {self._resolved_block_size - slack}." + f"Current max_draft_len={self.max_draft_len}." + ) + if self._ctx_unified_cache: + reserved = getattr(draft_kv_cache_manager, "_shared_draft_block_width", 0) + if self._compute_block_size > reserved: + raise ValueError( + f"Drafter block width {self._compute_block_size} exceeds the " + f"{reserved} positions reserved by the shared KV cache manager. " + "The manager must reserve the drafter's resolved block width." + ) + # Recorded even when the buffers stay private, so that the estimation # pass falling back to the arena still rebinds once the real pool lands. self._ctx_kv_manager = draft_kv_cache_manager @@ -1220,20 +1242,6 @@ def _lazy_init_ctx_buffers( self._req_to_slot = {} self._req_ctx_pos = {} - # checkpoint's trained block width - self._resolved_block_size = getattr(draft_model, "block_size", None) or ( - self.max_draft_len + 1 - ) - # what the draft forward actually computes - self._compute_block_size = self._draft_block_width(draft_model) - if self._compute_block_size > self._resolved_block_size: - slack = self._compute_block_size - self.max_draft_len - raise ValueError( - f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." - f"Lower max_draft_len to at most {self._resolved_block_size - slack}." - f"Current max_draft_len={self.max_draft_len}." - ) - # +block_size slack lets each backend append the per-iteration noise # K/V after the accumulated request context. assert hasattr(draft_model, "_build_fused_kv_buffers"), ( diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index b5e67034b10d..f6f4bde8162f 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -1210,14 +1210,13 @@ def set_draft_model(self, draft_model) -> None: ) @staticmethod - def block_width_for(spec_config) -> int: + def block_width_for(spec_config, *, shift_label: bool = True) -> int: """Block slots a DSpark worker computes for ``spec_config``. - The ``shift_label`` convention, which the DSpark drafter defaults to, - reads slots 0..K-1; a checkpoint that turns it off needs K+1 and is - caught by the block-within-allocation check. + The ``shift_label`` convention reads slots 0..K-1; a checkpoint that + turns it off reads slots 1..K and needs K+1 positions. """ - return spec_config.max_draft_len + return spec_config.max_draft_len + (0 if shift_label else 1) def _draft_block_width(self, draft_model) -> int: """Block width under the dspark ``shift_label`` convention. @@ -1227,9 +1226,9 @@ def _draft_block_width(self, draft_model) -> int: checkpoint at max_draft_len=7, which is how both published DSpark drafters are meant to run. """ - if getattr(draft_model, "_dspark_shift_label", False): - return self.block_width_for(self.spec_config) - return DFlashWorker.block_width_for(self.spec_config) + return self.block_width_for( + self.spec_config, shift_label=getattr(draft_model, "_dspark_shift_label", False) + ) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 9d5efeb597e2..d1ca1e8d3d54 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -786,10 +786,11 @@ def get_num_extra_kv_tokens(spec_config): return 0 -def get_draft_block_width(spec_config) -> int: +def get_draft_block_width(spec_config, *, shift_label: bool = True) -> int: """Block width the worker drafting for ``spec_config`` appends per request. 0 when that worker does not append a block to a paged draft KV cache. + ``shift_label`` selects the checkpoint's slot convention for standalone DSpark. """ if spec_config is None: return 0 @@ -797,7 +798,7 @@ def get_draft_block_width(spec_config) -> int: return DFlashWorker.block_width_for(spec_config) if (spec_config.spec_dec_mode.is_dspark() and not spec_config.draft_is_embedded_in_target): - return DSparkWorker.block_width_for(spec_config) + return DSparkWorker.block_width_for(spec_config, shift_label=shift_label) return 0 diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 3d7e6d996197..98ef0b44f757 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -15,6 +15,7 @@ """Selection of the unified target+DFlash KV cache layout.""" from types import SimpleNamespace +from unittest.mock import Mock import click import pytest @@ -23,6 +24,7 @@ from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.quantization import QuantAlgo from tensorrt_llm.quantization.mode import QuantMode @@ -92,6 +94,60 @@ def test_matching_draft_layout_unifies(is_disagg): assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("shift_label, expected_width", [(True, 7), (False, 8)]) +def test_dspark_checkpoint_width_reaches_shared_cache_cost_and_construction( + monkeypatch, is_disagg, shift_label, expected_width +): + from tensorrt_llm._torch.pyexecutor import _util + + layout = _model_config(kv_heads=8, head_dim=128) + layout.is_generation = True + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout) + creator._speculative_config.spec_dec_mode = SimpleNamespace( + is_dflash=lambda: False, is_dspark=lambda: True, is_external_drafter=lambda: True + ) + creator._speculative_config.max_draft_len = 7 + creator._speculative_config.draft_is_embedded_in_target = False + creator._model_engine.model.draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + creator._kv_cache_config = KvCacheConfig() + creator._tokens_per_block = 32 + creator._max_batch_size = 4 + creator._max_num_tokens = 4096 + creator._max_beam_width = 1 + creator._sparse_attention_config = None + creator._kv_connector_manager = None + creator._execution_stream = None + creator._disable_overlap_scheduler = True + creator._skip_est = True + creator._llm_args.kv_cache_config = creator._kv_cache_config + monkeypatch.setattr(creator, "_get_model_kv_cache_manager_cls", lambda *args: KVCacheManagerV2) + monkeypatch.setattr(creator, "_enable_kv_cache_stats", lambda: False) + + class RecordingManager(KVCacheManagerV2): + get_cache_size_per_token = Mock(return_value=(1, 0)) + get_cache_bytes_per_request = Mock(return_value=1) + get_cache_quota_per_request = Mock(return_value=1) + + assert creator._should_unify_dflash_kv_cache() + assert not creator._should_create_separate_draft_kv_cache() + creator._per_manager_cache_cost(RecordingManager, layout, include_request_cost=True) + create_manager = Mock(return_value=SimpleNamespace(max_seq_len=4096)) + monkeypatch.setattr(_util, "_create_kv_cache_manager", create_manager) + creator._create_kv_cache_manager(creator._model_engine) + + runtime_specs = create_manager.call_args.kwargs["layer_specs"] + for call in ( + RecordingManager.get_cache_size_per_token, + RecordingManager.get_cache_bytes_per_request, + RecordingManager.get_cache_quota_per_request, + create_manager, + ): + assert call.call_count == 1 + assert call.call_args.kwargs["shared_draft_block_width"] == expected_width + assert call.call_args.kwargs["layer_specs"] == runtime_specs + + def test_mismatched_draft_kv_heads_unify_only_in_aggregated_serving() -> None: target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=4, head_dim=128) 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 2bbf534a9a00..760545332075 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 @@ -38,6 +38,8 @@ KVCacheLayerSpec, KVCacheManagerV2, Role, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine @@ -46,6 +48,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( + DSparkDecodingConfig, KvCacheConfig, MTPDecodingConfig, MultimodalConfig, @@ -1311,6 +1314,85 @@ def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: ) == (0, (8 + 23) * 64 * 64) +@pytest.mark.parametrize( + "shared_draft_block_width,full_tokens,window_tokens", + [(None, 96, 64), (7, 96, 64), (8, 128, 96)], + ids=["config_default", "dspark_shifted", "dspark_unshifted"], +) +def test_v2_dspark_resolved_width_sizes_shared_cache_costs( + shared_draft_block_width: int | None, full_tokens: int, window_tokens: int +) -> None: + """One extra draft position can require a full page in every cost estimate.""" + spec_config = DSparkDecodingConfig(max_draft_len=7, speculative_model="draft") + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1024, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft = replace( + target, global_layer_id=1, owner=KVCacheLayerOwner.DRAFT, sliding_window_size=19 + ) + kwargs = dict( + tokens_per_block=32, + max_seq_len=81, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + avg_seq_len=81, + max_util_for_resume=1.0, + ), + spec_config=spec_config, + layer_specs=[target, draft], + shared_draft_block_width=shared_draft_block_width, + ) + # A page is 2 MiB, so allocator granularity preserves the boundary difference. + layer_bytes = 64 << 10 + assert KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=32, **kwargs + ) == (layer_bytes, layer_bytes * (2 * window_tokens + 32)) + assert KVCacheManagerV2._get_per_request_layer_tokens(None, Mapping(), **kwargs) == ( + [(layer_bytes, full_tokens, 0), (layer_bytes, window_tokens, window_tokens)], + full_tokens, + ) + assert KVCacheManagerV2.get_cache_bytes_per_request(None, Mapping(), **kwargs) == ( + layer_bytes * full_tokens + ) + assert KVCacheManagerV2.get_cache_quota_per_request(None, Mapping(), **kwargs) == ( + layer_bytes * (full_tokens + window_tokens) + ) + + +@pytest.mark.parametrize("with_target_layer", [False, True]) +def test_v2_resolved_draft_width_does_not_change_target_only_headroom( + with_target_layer: bool, +) -> None: + spec_config = DSparkDecodingConfig(max_draft_len=7, speculative_model="draft") + layer_specs = None + if with_target_layer: + layer_specs = [ + KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + ) + ] + assert _get_kv_alloc_extra_tokens( + spec_config, layer_specs, shared_draft_block_width=8 + ) == _get_kv_alloc_extra_tokens(spec_config, layer_specs) + assert _get_generation_kv_capacity( + spec_config, + is_draft=False, + layer_specs=layer_specs, + shared_draft_block_width=8, + ) == _get_generation_kv_capacity(spec_config, is_draft=False, layer_specs=layer_specs) + + def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) manager.layer_specs = None @@ -1828,6 +1910,9 @@ class _RecordingManager(manager_cls): def __init__(self, kv_cache_config: KvCacheConfig, *args: object, **kwargs: object) -> None: captured.append(kv_cache_config) assert kwargs.get("layer_specs") is factory_overrides.get("layer_specs") + assert kwargs.get("shared_draft_block_width") == factory_overrides.get( + "shared_draft_block_width" + ) model_config = Mock() model_config.pretrained_config = pretrained @@ -2087,8 +2172,10 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) @pytest.mark.parametrize("publish_layer_specs", [False, True]) +@pytest.mark.parametrize("shared_draft_block_width", [None, 7, 8]) def test_create_kv_cache_manager_preserves_explicit_storage_specs( publish_layer_specs: bool, + shared_draft_block_width: int | None, ) -> None: """Model geometry must not replace an explicit target and draft storage layout.""" specs = tuple(_fp8_target_bf16_draft_specs(draft_window=128)) @@ -2098,6 +2185,7 @@ def test_create_kv_cache_manager_preserves_explicit_storage_specs( max_seq_len=2048, publish_layer_specs=publish_layer_specs, layer_specs=specs, + shared_draft_block_width=shared_draft_block_width, ) 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 2bf853865087..a21605217852 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -54,6 +54,7 @@ from tensorrt_llm.llmapi.llm_args import ( BlockReuseConfig, DFlashDecodingConfig, + DSparkDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MTPDecodingConfig, @@ -4001,7 +4002,14 @@ def test_storage_pool_group_count_matches_manager( @pytest.mark.parametrize("disagg", [False, True], ids=["aggregated", "disaggregated"]) -def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg: bool) -> None: +@pytest.mark.parametrize( + "spec_config_cls,shared_draft_block_width", + [(DFlashDecodingConfig, None), (DSparkDecodingConfig, 7), (DSparkDecodingConfig, 8)], + ids=["dflash", "dspark_shifted", "dspark_unshifted"], +) +def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize( + disagg: bool, spec_config_cls: type, shared_draft_block_width: int | None +) -> None: """The drafter appends its block after the context, so its pages must reach it. After a decode step accepts every draft token the context is ``prompt + K + @@ -4012,9 +4020,9 @@ def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg pytest.skip("requires CUDA") torch.cuda.init() max_draft_len = 7 - block_width = max_draft_len + 1 + block_width = shared_draft_block_width or max_draft_len + 1 tokens_per_block = 32 - prompt_len = 49 + prompt_len = 2 * tokens_per_block - (max_draft_len + 1) - block_width + 1 target = KVCacheLayerSpec( global_layer_id=0, owner=KVCacheLayerOwner.TARGET, @@ -4048,7 +4056,8 @@ def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg mapping=Mapping(), dtype=DataType.FP8, layer_specs=layer_specs, - spec_config=DFlashDecodingConfig(max_draft_len=max_draft_len, speculative_model="draft"), + spec_config=spec_config_cls(max_draft_len=max_draft_len, speculative_model="draft"), + shared_draft_block_width=shared_draft_block_width, vocab_size=4096, enable_stats=False, ) @@ -4061,6 +4070,8 @@ def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg is_disagg_generation_init_state=disagg, ) try: + assert manager._shared_draft_block_width == block_width + assert manager._generation_kv_capacity_headroom == max_draft_len + 1 + block_width if disagg: request.py_draft_tokens = [] request.is_disagg_generation_transmission_complete = False diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py index 7abf67b67d04..7581011fe63b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py @@ -76,7 +76,7 @@ def _metadata(*, declares_max_seq_len): def _size_arena(metadata): """Run the sizing path and return the ``_max_ctx`` it settled on.""" worker = _worker() - draft_model = MagicMock() + draft_model = MagicMock(block_size=8) draft_model.config = MagicMock(max_position_embeddings=POSITION_TABLE) spec_metadata = MagicMock(max_num_requests=32) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py index 10cb0365548e..a9bd05e7f217 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -20,12 +20,16 @@ """ from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.utils import get_draft_block_width pytestmark = pytest.mark.cpu_only @@ -96,3 +100,108 @@ def test_drafter_block_outside_the_allocation_is_an_error(): with pytest.raises(RuntimeError, match="block of 7 positions falls outside"): DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def _dspark_worker() -> DSparkWorker: + worker = object.__new__(DSparkWorker) + worker.spec_config = SimpleNamespace( + max_draft_len=7, + tokens_per_gen_step=8, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + draft_is_embedded_in_target=False, + ) + worker._ctx_buf_inited = False + worker._dflash_attention_backend = "TRTLLM" + return worker + + +@pytest.mark.parametrize(("shift_label", "expected"), [(True, 7), (False, 8)]) +def test_dspark_block_width_tracks_checkpoint_slot_convention( + shift_label: bool, expected: int +) -> None: + worker = _dspark_worker() + draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + + assert worker._draft_block_width(draft_model) == expected + assert DSparkWorker.block_width_for(worker.spec_config, shift_label=shift_label) == expected + assert get_draft_block_width(worker.spec_config, shift_label=shift_label) == expected + + +def test_dspark_missing_shift_label_preserves_unshifted_runtime_width() -> None: + worker = _dspark_worker() + + assert worker._draft_block_width(SimpleNamespace()) == 8 + assert DSparkWorker.block_width_for(worker.spec_config) == 7 + assert get_draft_block_width(worker.spec_config) == 7 + + +@pytest.mark.parametrize("shift_label", [True, False]) +def test_draft_block_width_shift_label_only_affects_standalone_dspark(shift_label: bool) -> None: + config = _dspark_worker().spec_config + config.draft_is_embedded_in_target = True + assert get_draft_block_width(config, shift_label=shift_label) == 0 + + config.spec_dec_mode = SpeculativeDecodingMode.DFLASH + assert get_draft_block_width(config, shift_label=shift_label) == 8 + + +def _init_dspark_ctx_buffers( + worker: DSparkWorker, *, shift_label: bool, reserved_width: int, block_size: int = 8 +) -> None: + manager = SimpleNamespace(max_seq_len=64, _shared_draft_block_width=reserved_width) + draft_model = SimpleNamespace( + _dspark_shift_label=shift_label, + block_size=block_size, + dflash_attention_backend="TRTLLM", + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=1), + SimpleNamespace(kv_cache_manager=manager, kv_cache_block_offsets=None), + manager, + ) + + +@pytest.mark.parametrize( + ("shift_label", "reserved_width", "block_size", "error"), + [ + (False, 7, 8, "Drafter block width 8 exceeds the 7 positions"), + (True, 6, 8, "Drafter block width 7 exceeds the 6 positions"), + (False, 8, 7, "checkpoint was trained with block_size=7"), + ], +) +def test_dspark_invalid_width_is_rejected_before_cuda_allocation( + monkeypatch: pytest.MonkeyPatch, + shift_label: bool, + reserved_width: int, + block_size: int, + error: str, +) -> None: + allocate = Mock(side_effect=RuntimeError("buffer allocation reached")) + monkeypatch.setattr(torch, "zeros", allocate) + + with pytest.raises(ValueError, match=error): + _init_dspark_ctx_buffers( + _dspark_worker(), + shift_label=shift_label, + reserved_width=reserved_width, + block_size=block_size, + ) + + allocate.assert_not_called() + + +@pytest.mark.parametrize(("shift_label", "reserved_width"), [(True, 7), (False, 8)]) +def test_dspark_matching_width_reaches_buffer_allocation( + monkeypatch: pytest.MonkeyPatch, shift_label: bool, reserved_width: int +) -> None: + allocate = Mock(side_effect=RuntimeError("buffer allocation reached")) + monkeypatch.setattr(torch, "zeros", allocate) + worker = _dspark_worker() + + with pytest.raises(RuntimeError, match="buffer allocation reached"): + _init_dspark_ctx_buffers(worker, shift_label=shift_label, reserved_width=reserved_width) + + allocate.assert_called_once() + assert worker._ctx_unified_cache + assert worker._compute_block_size == reserved_width From 6a4fc09b6fbb61c94b6550029fbb515d62730976 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 18:04:11 -0700 Subject: [PATCH 31/35] fix: preserve shared draft KV selection intent Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 6 ++-- tensorrt_llm/commands/serve.py | 4 ++- .../kv_cache/test_dflash_unified_kv_gate.py | 28 +++++++++++++++++++ 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c3d510246ce2..02877dc4be5c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2248,12 +2248,14 @@ def _should_unify_dflash_kv_cache(self) -> bool: """Whether this run can place target and DFlash or standalone DSpark layers in one V2 manager. - Disaggregated serving shares only when TRTLLM_DISAGG_SHARED_DRAFT_KV=1. + Disaggregated serving shares only when TRTLLM_DISAGG_SHARED_DRAFT_KV + is "1" or "auto". """ if self._dflash_unify_blocker() is not None: return False return not (self._is_disagg - and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1") + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") + not in ("1", "auto")) def _check_shared_draft_kv_request(self) -> None: """Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored.""" diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 192ccbd06094..2d1e29b206b1 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2575,7 +2575,9 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1": raise click.UsageError( f"TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires {blocker}.") - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" if blocker else "1") + if os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") in (None, "auto"): + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" + if blocker else "auto") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 98ef0b44f757..f855015a0f06 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -14,6 +14,7 @@ # limitations under the License. """Selection of the unified target+DFlash KV cache layout.""" +import os from types import SimpleNamespace from unittest.mock import Mock @@ -353,3 +354,30 @@ def test_launcher_accepts_eligible_shared_cache_config(monkeypatch, tmp_path): {"speculative_config": spec, "enable_attention_dp": True}, {"speculative_config": spec, "enable_attention_dp": True}, ) + + +@pytest.mark.parametrize( + "request_mode,draft_heads,expected_mode", + [("0", 8, "0"), (None, 4, "auto")], + ids=["explicit_opt_out", "auto_fallback"], +) +def test_launcher_preserves_opt_out_and_allows_layout_fallback( + monkeypatch, tmp_path, request_mode, draft_heads, expected_mode +) -> None: + if request_mode is None: + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV, raising=False) + else: + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, request_mode) + config = {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config(monkeypatch, tmp_path, config, config) + + assert os.environ[_SHARED_DRAFT_KV_ENV] == expected_mode + creator = _creator( + is_disagg=True, + target=_model_config(kv_heads=8, head_dim=128), + draft=_model_config(kv_heads=draft_heads, head_dim=128), + ) + creator._check_shared_draft_kv_request() + assert not creator._should_unify_dflash_kv_cache() From d2b0bb41ee5086733690b023fca3c8b699dc0851 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 18:30:57 -0700 Subject: [PATCH 32/35] style: format DSpark block width helper Signed-off-by: Allison Lim --- tensorrt_llm/_torch/speculative/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index d1ca1e8d3d54..94cf65bd7c4d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -798,7 +798,8 @@ def get_draft_block_width(spec_config, *, shift_label: bool = True) -> int: return DFlashWorker.block_width_for(spec_config) if (spec_config.spec_dec_mode.is_dspark() and not spec_config.draft_is_embedded_in_target): - return DSparkWorker.block_width_for(spec_config, shift_label=shift_label) + return DSparkWorker.block_width_for(spec_config, + shift_label=shift_label) return 0 From 74020518daa08d9ac905f6fbe11363f538602f2d Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 23:15:40 -0700 Subject: [PATCH 33/35] fix: align CuTe MLA per-layer page indices with buffer origin Signed-off-by: Allison Lim --- .../attention/backends/fmha/cute_dsl_mla.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index cf0f56d136ee..a136f95c45bb 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -569,6 +569,12 @@ def _run_mla_decode( params: FmhaParams, kernel_dtype: torch.dtype, ) -> None: + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheManagerV2, + Role, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + attn = params.attn meta = params.meta @@ -649,6 +655,20 @@ def _run_mla_decode( ) if layers_in_pool > 1: page_table = page_table + layer_in_pool + elif ( + isinstance(meta.kv_cache_manager, KVCacheManagerV2) + and meta.kv_cache_manager._use_per_layer_page_tables + ): + # V2 per-layer indices include the layer offset, but get_buffers() + # already starts at this layer's KEY buffer. Apply the offset once. + converter = meta.kv_cache_manager.impl.get_page_index_converter( + local_layer_idx, Role.KEY + ) + layer_offset = int(converter.layer_offset) + if layer_offset: + page_table = torch.where( + page_table == BAD_PAGE_INDEX, page_table, page_table - layer_offset + ) # KVCacheManager exposes NHD pages as [num_pages, 1, page_size, 1, head_dim]. # The CuTe DSL kernel consumes a paged [page_size, dim, num_pages] view From e0171d67015c2f69e13dd31f5638120877fdc580 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 26 Oct 2026 23:41:50 -0700 Subject: [PATCH 34/35] test: cover CuTe MLA reads with nonzero layer offsets Signed-off-by: Allison Lim --- .../_torch/attention/test_fmha_page_index.py | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 31c84af61276..edeeb7a52924 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX def test_flashinfer_uses_v2_page_index_upper_bound_directly() -> None: @@ -258,6 +259,93 @@ def generation_preprocess(*args: object) -> tuple[object, ...]: assert len(decode_calls) == 1 +@pytest.mark.cpu_only +@pytest.mark.parametrize("local_layer_idx", [0, 1]) +def test_cute_dsl_mla_v2_per_layer_indices_read_correct_pages( + monkeypatch: pytest.MonkeyPatch, local_layer_idx: int +) -> None: + """Per-layer indices and layer-origin buffers must select the same physical page.""" + page_size, latent_dim, rope_dim = 2, 2, 2 + layer_idx = 16 + local_layer_idx + physical_pool = ( + torch.arange(10, dtype=torch.bfloat16) + .view(10, 1, 1, 1, 1) + .expand(-1, 1, page_size, 1, latent_dim + rope_dim) + .contiguous() + ) + # Two MLA target layers share a pool. Mixed target/draft layouts use + # per-layer tables: slot 3 maps to physical pages 6 and 7 for these layers. + block_offsets = torch.tensor( + [[[[6, BAD_PAGE_INDEX]]], [[[7, BAD_PAGE_INDEX]]]], dtype=torch.int32 + ) + original_offsets = block_offsets.clone() + # V2 wraps the shifted raw pointer as storage origin; a slice alone would + # retain the full pool's storage and let an incorrect rebase hide the bug. + layer_buffer = physical_pool[local_layer_idx:].clone() + + def get_buffers(requested_layer_idx: int) -> torch.Tensor: + assert requested_layer_idx == layer_idx + return layer_buffer + + def get_page_index_converter(requested_local_idx: int, role: object) -> SimpleNamespace: + assert (requested_local_idx, role) == (local_layer_idx, Role.KEY) + return SimpleNamespace(layer_offset=local_layer_idx) + + manager = object.__new__(KVCacheManagerV2) + manager._use_per_layer_page_tables = True + manager.layer_offsets = {16: 0, 17: 1} + manager.get_buffers = get_buffers + manager.impl = SimpleNamespace(get_page_index_converter=get_page_index_converter) + metadata = SimpleNamespace( + kv_cache_manager=manager, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_mapping=torch.tensor([[0, 0], [1, 0]], dtype=torch.int32), + num_contexts=0, + max_num_requests=1, + helix_position_offsets=None, + ) + attn = SimpleNamespace( + layer_idx=layer_idx, + get_local_layer_idx=lambda meta: meta.kv_cache_manager.layer_offsets[layer_idx], + num_heads=1, + kv_lora_rank=latent_dim, + qk_rope_head_dim=rope_dim, + qk_nope_head_dim=latent_dim, + q_scaling=1.0, + ) + q = torch.zeros((1, 1, latent_dim + rope_dim), dtype=torch.bfloat16) + output = torch.empty((1, 1, latent_dim), dtype=torch.bfloat16) + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=AttentionForwardArgs( + output=output, attention_input_type=AttentionInputType.generation_only + ), + workspace=torch.empty(0, dtype=torch.uint8), + sequence_lengths=torch.tensor([page_size], dtype=torch.int32), + tokens_per_block=page_size, + num_requests=1, + ) + pages_read: list[torch.Tensor] = [] + + def read_page(q_latent, q_rope, c_latent, c_rope, page_table, *args) -> None: + assert page_table[1, 0].item() == BAD_PAGE_INDEX + page = page_table[0, 0].item() + pages_read.append(torch.cat((c_latent[:, :, page], c_rope[:, :, page]), dim=1)) + + monkeypatch.setattr( + torch.ops.trtllm, "cute_dsl_mla_decode_fp16_blackwell", read_page, raising=False + ) + monkeypatch.setattr(CuteDslMlaFmha, "_required_workspace_size", staticmethod(lambda *args: 0)) + monkeypatch.setattr(CuteDslMlaFmha, "_split_kv_plan", staticmethod(lambda *args: None)) + fmha = object.__new__(CuteDslMlaFmha) + fmha._run_mla_decode(q, output, params, torch.bfloat16) + + assert len(pages_read) == 1 + torch.testing.assert_close(pages_read[0], physical_pool[6 + local_layer_idx, 0, :, 0, :]) + torch.testing.assert_close(block_offsets, original_offsets) + + def _cute_dsl_mla_helix_support( monkeypatch: pytest.MonkeyPatch, *, From 4d3904a74e7f4a1797b9d218e17f6dd5f814b5d3 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Fri, 30 Oct 2026 15:21:24 -0700 Subject: [PATCH 35/35] test: split disaggregated draft KV environment fallback case Signed-off-by: Allison Lim --- .../kv_cache/test_dflash_unified_kv_gate.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 0c38cca9114f..b0d04068754a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -15,6 +15,7 @@ """Selection of the unified target+DFlash KV cache layout.""" import os +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock @@ -249,7 +250,9 @@ def stop(*args, **kwargs): ) -def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_path): +def test_launcher_rejects_dflash_on_generation_worker_only( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: # The launcher rejects the variable when only one section sets speculative_config. with pytest.raises(click.UsageError, match="both context_servers and generation_servers"): _launch_from_disagg_config( @@ -258,13 +261,17 @@ def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_ context_servers={}, generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, ) + + +def test_disagg_without_shared_draft_kv_env_keeps_separate_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=8, head_dim=128) - for _role in ("context", "generation"): - creator = _creator(is_disagg=True, target=target, draft=draft) - assert not creator._should_unify_dflash_kv_cache() + creator = _creator(is_disagg=True, target=target, draft=draft) + assert not creator._should_unify_dflash_kv_cache() def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path):