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 bb40369610e3..ddad2dc5de5d 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 bdb429ba9074..0d702f2729bd 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 9363fed79f3f..047f6a8452a0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1597,12 +1597,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/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/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 718c6f397767..401ae37f571d 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. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker and under `modeling_v2`. `"VANILLA"` 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. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0` and under `modeling_v2`. `"VANILLA"` 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,39 @@ 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. Both servers therefore need +the same `speculative_config`: a generation server whose context server keeps a +different number of draft layers, for example one without a drafter, fails +those requests at the first KV transfer. To draft on the generation server +alone, set `attention_backend="VANILLA"` or `TRTLLM_DISAGG_SHARED_DRAFT_KV=0` +there; it then drafts without the prompt context, at a lower acceptance length. +The target and draft layers +are resharded by their own KV-head counts, so the draft's KV-head count, head +dimension and dtype may differ from the target's, and the context and +generation servers may use different tensor-parallel sizes or attention DP +settings. An NVFP4 target KV cache requires a draft with an identical dtype +and KV-head count. The target's and draft's K-only (MLA) or K/V cache kind must +match, and both servers must run the same target and drafter. + +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. ### User-provided drafting diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py index 44bca364de99..49fb9e37a8d4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py @@ -329,6 +329,8 @@ def block_offsets( host is None or getattr(manager, "sparse_metadata_batch", None) is not None or getattr(manager, "_use_per_layer_page_tables", True) + # Shared attention-op tables are not the life cycles' SHARED-path tables. + or getattr(manager, "_attention_op_pool_groups", None) is not None or not dst.is_cuda or dst.dtype != torch.int32 or host.dtype != torch.int32 diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index 5db118320c4e..5f3544a01683 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -1352,20 +1352,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/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 4b68b9cc53aa..2802d2e88f7a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -79,6 +79,11 @@ def get_paged_kv_storage_unsupported_reason( if isinstance(manager, KVCacheManagerV2): if manager.enable_swa_scratch_reuse: return "KVCacheManagerV2 SWA scratch reuse is not supported." + # get_kv_page_offset indexes the life cycles' kv_offset by attention-op pool. + if getattr(manager, "_use_per_layer_page_tables", False) or ( + getattr(manager, "_attention_op_pool_groups", None) is not None + ): + return "KVCacheManagerV2 per-layer or shared attention-op tables are not supported." elif isinstance(manager, KVCacheManager): if manager.num_pools != 1: return "KVCacheManagerV1 with multiple memory pools is not supported." 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/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index c193de0ce996..56806007a4bb 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -235,10 +235,12 @@ def reserve( *, timeout: Optional[float] = None, extra_bytes: int = 0, + writers_split_layers: bool = True, ) -> bool: """Reserve a region and record its address for the senders. False falls back to per-fragment. ``extra_bytes`` sizes the sender's non-paged payload (mamba/KDA recurrent state) that rides - the same coalesced write as the KV blocks.""" + the same coalesced write as the KV blocks. ``writers_split_layers`` says whether the fan-in + writers hold different layers (PP fan-in) rather than a head slice of every layer.""" @abstractmethod def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 1c70019e46bb..f2e6a702de6e 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -235,13 +235,15 @@ def reserve( *, timeout: Optional[float] = _RESERVE_TIMEOUT_S, extra_bytes: int = 0, + writers_split_layers: bool = True, ) -> bool: """Reserve a region and create its state, recording the address for the senders. Returns False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total must divide across the writers. ``extra_bytes`` is the non-paged payload the sender appends to the same coalesced write (mamba/KDA recurrent state, sized by the receiver via ``mamba_receiver_payload_bytes``); the region must cover it or the write would overrun into the - neighboring slot.""" + neighboring slot. ``writers_split_layers`` is False when every fan-in writer sends a head + slice of every layer (TP fan-in), True when the writers hold different layers (PP fan-in).""" total = 0 has_state_group = False # Block tables are positional with -1 holes (see Chunk); only the @@ -311,10 +313,12 @@ def reserve( f"({total % num_writers}B remainder); head-mismatch explosion NOT mitigated", warn_key="kv-bounce-uneven-fanin", ) - if num_writers > 1: + if num_writers > 1 and writers_split_layers: # Fan-in gives each writer an equal share of the region, which only matches where it - # writes when all writers send the same bytes. Equal layer count guarantees that only - # when the per-block sizes match, so require that here, else fall back. + # writes when all writers send the same bytes. A TP fan-in writer sends 1/num_writers + # of every block of every group, so its share is exact for any per-group block sizes. + # PP fan-in writers send different layers, and an equal layer count gives equal bytes + # only when the per-block sizes match, so require that here, else fall back. # Exclude None entries (STATE groups) — they are already guarded above. present_slot_bytes = { self._block_bytes_per_group[g] @@ -475,12 +479,23 @@ def _scatter_loop(self): ok = True try: # Scatter each writer's fragments from its own source, never one global offset, so a - # missing or fallback writer cannot shift where the others are read from. + # missing or fallback writer cannot shift where the others are read from. All writers + # go into ONE batched copy: the copy helper reuses one pinned metadata buffer per + # stream and needs a sync between calls, and one call costs one host round trip + # instead of one per writer. + srcs, dsts, frag_sizes = [], [], [] for src_base, dst_ptrs, sizes in descs: p = Plan(dst_ptrs, dst_ptrs, sizes, int(sizes.sum())) - scatter_contiguous( - src_base, p.dst_ptrs, p.sizes, p.offsets, stream=self._scatter_stream - ) + srcs.append(np.int64(src_base) + p.offsets) + dsts.append(p.dst_ptrs) + frag_sizes.append(p.sizes) + scatter_contiguous( + 0, + np.concatenate(dsts), + np.concatenate(frag_sizes), + np.concatenate(srcs), + stream=self._scatter_stream, + ) CUASSERT(cudart.cudaStreamSynchronize(self._scatter_stream)) except Exception as e: # a scatter failure must not kill the worker nor be reported as success @@ -524,6 +539,7 @@ def reserve( *, timeout: Optional[float] = _RESERVE_TIMEOUT_S, extra_bytes: int = 0, + writers_split_layers: bool = True, ) -> bool: return False diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index f141cce93ff4..86f3be81d92a 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -26,7 +26,12 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import ( + CacheKind, + KVCachePageTable, + MapperKind, + PoolView, +) from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch from tensorrt_llm._utils import nvtx_range @@ -179,9 +184,15 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: Optional[int] = None, + peer_kv_heads: Optional[int] = None, ): self._ri = self_ri self._peer_ri = peer_ri + # Head counts of the matched view pair; rank-wide values are the + # fallback for views from builders that do not record them. + self_kv_heads = self_kv_heads or self_ri.attention.kv_heads_per_rank + peer_kv_heads = peer_kv_heads or peer_ri.attention.kv_heads_per_rank self_tp_per_dp = self_ri.tp_size_per_dp_group peer_tp_per_dp = peer_ri.tp_size_per_dp_group @@ -214,28 +225,21 @@ def __init__( buffers_per_layer=peer_buffers_per_layer, side="peer", ) - bytes_per_head = self._bytes_per_head( - src_buffer_bytes, self._ri.attention.kv_heads_per_rank, side="local" - ) - peer_bytes_per_head = self._bytes_per_head( - dst_buffer_bytes, peer_ri.attention.kv_heads_per_rank, side="peer" - ) + bytes_per_head = self._bytes_per_head(src_buffer_bytes, self_kv_heads, side="local") + peer_bytes_per_head = self._bytes_per_head(dst_buffer_bytes, peer_kv_heads, side="peer") if bytes_per_head != peer_bytes_per_head: raise ValueError( f"HND bytes per head mismatch: local={bytes_per_head}, peer={peer_bytes_per_head}" ) - self._bytes_cont_heads = ( - min(self._ri.attention.kv_heads_per_rank, peer_ri.attention.kv_heads_per_rank) - * bytes_per_head - ) + self._bytes_cont_heads = min(self_kv_heads, peer_kv_heads) * bytes_per_head self._src_head_off, self._dst_head_off = self._compute_head_offsets( self_tp_per_dp, peer_tp_per_dp, self_tp_rank, peer_tp_rank, - self_kv_heads=self._ri.attention.kv_heads_per_rank, - peer_kv_heads=peer_ri.attention.kv_heads_per_rank, + self_kv_heads=self_kv_heads, + peer_kv_heads=peer_kv_heads, bytes_per_head=bytes_per_head, ) @@ -365,6 +369,8 @@ def __init__( peer_bytes_per_layer: int, self_buffers_per_layer: int, peer_buffers_per_layer: int, + self_kv_heads: Optional[int] = None, + peer_kv_heads: Optional[int] = None, ) -> None: # Deliberately do not call HNDHeadMismatchMapper.__init__: its offsets # assume HND-contiguous heads. Initialize the three attributes consumed @@ -377,8 +383,8 @@ def __init__( f"local={self_tpb}, peer={peer_tpb}" ) - self_heads = self_ri.attention.kv_heads_per_rank - peer_heads = peer_ri.attention.kv_heads_per_rank + self_heads = self_kv_heads or self_ri.attention.kv_heads_per_rank + peer_heads = peer_kv_heads or peer_ri.attention.kv_heads_per_rank if self_buffers_per_layer != peer_buffers_per_layer: raise ValueError( "NHD buffer count per layer mismatch: " @@ -484,16 +490,28 @@ def __init__(self, self_rank_info: RankInfo): self._ri = self_rank_info def should_send( - self, peer_overlap, peer_rank_info, *, mapper_kind: Optional[MapperKind] = None + self, + peer_overlap, + peer_rank_info, + *, + mapper_kind: Optional[MapperKind] = None, + duplicate_head_factor: Optional[int] = None, ) -> "bool | None": """Attention uses head-duplication routing. - REPLICATED views (index-key side caches) hold identical bytes on every - TP rank, so they return None to defer to fan-in election instead. + ``duplicate_head_factor`` is the sending-side factor of the matched + view pair; the overlap's rank-wide factor is the default for callers + without a view pair. REPLICATED views (index-key side caches) hold + identical bytes on every TP rank, so they return None to defer to + fan-in election instead. """ if mapper_kind == MapperKind.REPLICATED: return None - dup = peer_overlap.duplicate_head_factor + dup = ( + peer_overlap.duplicate_head_factor + if duplicate_head_factor is None + else duplicate_head_factor + ) if dup <= 1: return True tp_rank = self._ri.tp_rank % self._ri.tp_size_per_dp_group @@ -615,30 +633,55 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool: ) ) - def _head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + @staticmethod + def view_kv_heads(ri: RankInfo, pool_view: Optional[PoolView] = None) -> int: + """Per-rank KV heads of a view; the rank-wide count for legacy views.""" + heads = getattr(pool_view, "num_kv_heads", None) + return heads if heads else ri.attention.kv_heads_per_rank + + def _head_factors( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[int, int]: self_tp = self._tp_per_dp(self._ri) peer_tp = self._tp_per_dp(peer_ri) - a = self._ri.attention - b = peer_ri.attention - return a.kv_heads_per_rank * self_tp, b.kv_heads_per_rank * peer_tp + return ( + self.view_kv_heads(self._ri, self_pv) * self_tp, + self.view_kv_heads(peer_ri, peer_pv) * peer_tp, + ) + + def head_match( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[bool, bool]: + """Whether the matched views hold the same heads per rank. - def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]: - factor_self, factor_peer = self._head_factors(peer_ri) + Without a view pair this compares the rank-wide head counts. + """ + factor_self, factor_peer = self._head_factors(peer_ri, self_pv, peer_pv) is_dup_head = factor_self != factor_peer # Slot layout is compatible if kv_heads_per_rank matches (same slot size). - head_match = ( - self._ri.attention.is_mla - or self._ri.attention.kv_heads_per_rank == peer_ri.attention.kv_heads_per_rank - ) + head_match = self._ri.attention.is_mla or self.view_kv_heads( + self._ri, self_pv + ) == self.view_kv_heads(peer_ri, peer_pv) return head_match, is_dup_head - def duplicate_head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + def duplicate_head_factors( + self, + peer_ri: RankInfo, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, + ) -> tuple[int, int]: # Head duplication only applies to attention KV. In particular, do # not divide 0 / 0 for an attention-free hybrid PP stage; MambaPolicy # computes its TP mapping separately. - if self._ri.attention.kv_heads_per_rank == 0 or peer_ri.attention.kv_heads_per_rank == 0: + if self.view_kv_heads(self._ri, self_pv) == 0 or self.view_kv_heads(peer_ri, peer_pv) == 0: return 1, 1 - factor_self, factor_peer = self._head_factors(peer_ri) + factor_self, factor_peer = self._head_factors(peer_ri, self_pv, peer_pv) dup_head = max(1, factor_self // factor_peer) peer_dup_head = max(1, factor_peer // factor_self) return dup_head, peer_dup_head @@ -683,7 +726,7 @@ def build_mapper( peer_bytes_per_layer, ) - head_match, _ = self.head_match(peer_ri) + head_match, _ = self.head_match(peer_ri, self_pv, peer_pv) if head_match: return IntactMapper( self_layer_offsets, @@ -703,6 +746,8 @@ def build_mapper( peer_bytes_per_layer=peer_bytes_per_layer, self_buffers_per_layer=self_buffers_per_layer, peer_buffers_per_layer=peer_buffers_per_layer, + self_kv_heads=self.view_kv_heads(self._ri, self_pv), + peer_kv_heads=self.view_kv_heads(peer_ri, peer_pv), ) return HNDHeadMismatchMapper( @@ -714,4 +759,6 @@ def build_mapper( peer_bytes_per_layer=peer_bytes_per_layer, self_buffers_per_layer=self_buffers_per_layer, peer_buffers_per_layer=peer_buffers_per_layer, + self_kv_heads=self.view_kv_heads(self._ri, self_pv), + peer_kv_heads=self.view_kv_heads(peer_ri, peer_pv), ) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 72bd85972e0c..d0cc4f6de5bb 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -45,6 +45,16 @@ class PeerOverlap: duplicate_head_factor: int = 1 peer_duplicate_head_factor: int = 1 ranks: List[int] = field(default_factory=list) + # Duplication factors of the matched head-sharded view pairs whose factors + # differ from the rank-wide ones above, keyed by (self layer group, self + # pool view). Views without an entry use the rank-wide factors. + view_duplicate_head_factors: Dict["LGPoolKey", Tuple[int, int]] = field(default_factory=dict) + + def duplicate_head_factors_for(self, pool_key: "LGPoolKey") -> Tuple[int, int]: + """(self, peer) duplication factors of one view pair.""" + return self.view_duplicate_head_factors.get( + pool_key, (self.duplicate_head_factor, self.peer_duplicate_head_factor) + ) class PeerRegistrar: @@ -84,32 +94,37 @@ def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): extractor = KVRegionExtractorV1(peer_ri.page_table) self._peer_ext_cache[key] = extractor - head_match, _ = self._get_policy(CacheKind.PAGED).head_match(peer_ri) - if not head_match: - self_page_table = self._self_ext_cache.page_table - nhd_fragments_per_token = sum( - len( - self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].buffer_entries - ) - for layer_group_id, pool_idx in self.get_pool_mapping(peer_ri) - if self_page_table.layer_groups[layer_group_id].pool_views[pool_idx].mapper_kind - == MapperKind.NHD + policy = self._get_policy(CacheKind.PAGED) + self_page_table = self._self_ext_cache.page_table + nhd_fragments_per_token = 0 + nhd_heads: set[tuple[int, int]] = set() + for (layer_group_id, pool_idx), (peer_lg, peer_pi) in self.get_pool_mapping( + peer_ri + ).items(): + self_pv = self_page_table.layer_groups[layer_group_id].pool_views[pool_idx] + if self_pv.mapper_kind != MapperKind.NHD: + continue + peer_pv = peer_ri.page_table.layer_groups[peer_lg].pool_views[peer_pi] + if policy.head_match(peer_ri, self_pv, peer_pv)[0]: + continue + nhd_fragments_per_token += len(self_pv.buffer_entries) + nhd_heads.add( + (policy.view_kv_heads(self._ri, self_pv), policy.view_kv_heads(peer_ri, peer_pv)) + ) + if nhd_fragments_per_token: + heads_text = ", ".join( + f"local_kv_heads={local}, peer_kv_heads={peer}" for local, peer in sorted(nhd_heads) + ) + heads_key = "_".join(f"{local}-{peer}" for local, peer in sorted(nhd_heads)) + logger.warning_once( + "NHD head-mismatched disaggregated KV transfer has no " + "contiguous staging path and will emit approximately " + f"{nhd_fragments_per_token} NIXL descriptors per transferred " + "token per peer, excluding block-level replicated pools " + f"({heads_text}). " + "Long-context TEP/DEP transfers may have high latency.", + key=f"native-nhd-head-mismatch-{heads_key}-{nhd_fragments_per_token}", ) - if nhd_fragments_per_token: - local_heads = self._ri.attention.kv_heads_per_rank - peer_heads = peer_ri.attention.kv_heads_per_rank - logger.warning_once( - "NHD head-mismatched disaggregated KV transfer has no " - "contiguous staging path and will emit approximately " - f"{nhd_fragments_per_token} NIXL descriptors per transferred " - "token per peer, excluding block-level replicated pools " - f"(local_kv_heads={local_heads}, peer_kv_heads={peer_heads}). " - "Long-context TEP/DEP transfers may have high latency.", - key=( - "native-nhd-head-mismatch-" - f"{local_heads}-{peer_heads}-{nhd_fragments_per_token}" - ), - ) def peer_extractor(self, peer_name: str, peer_rank: int) -> KVRegionExtractorV1: return self._peer_ext_cache[self._unique_key(peer_name, peer_rank)] @@ -170,6 +185,8 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + self._validate_view_head_geometry(peer_ri) + self.validate_draft_layers(self._ri, peer_ri) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -185,6 +202,80 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: return True + @staticmethod + def validate_draft_layers(self_ri: RankInfo, peer_ri: RankInfo) -> None: + """Reject a peer that keeps a different number of draft layers with its target KV. + + The layer-count tolerance of ``_check_peer_compatible`` would transfer + only the target layers and leave the drafter's prompt KV unset on the + generation side, which costs acceptance length without any error. + """ + if self_ri.num_draft_layers != peer_ri.num_draft_layers: + raise ValueError( + "PeerRegistrar: unified draft KV layer count differs " + f"(local={self_ri.num_draft_layers}, peer={peer_ri.num_draft_layers}), so " + "the drafter's prompt KV would not reach the generation worker. Run the " + "context and generation workers with the same speculative_config and the " + "same TRTLLM_DISAGG_SHARED_DRAFT_KV setting. To draft on the generation " + "worker alone, keep its draft KV private with TRTLLM_DISAGG_SHARED_DRAFT_KV=0 " + "or attention_backend VANILLA." + ) + + def _validate_view_head_geometry(self, peer_ri: RankInfo) -> None: + """Check that every matched head-sharded view pair can be resharded. + + Views carry their own KV-head count. A view without one (tables from + older builders) falls back to the rank-wide count, which describes it + only while the rank's views share one head geometry. When either side + publishes counts, the bytes per head of both sides, taken with those + counts, must therefore agree; a legacy side whose views differ in head + geometry fails this check and has to publish per-view counts. + """ + self_page_table = self._self_ext_cache.page_table + peer_page_table = peer_ri.page_table + policy = self._get_policy(CacheKind.PAGED) + for (self_lg, self_pi), (peer_lg, peer_pi) in self.get_pool_mapping(peer_ri).items(): + layer_group = self_page_table.layer_groups[self_lg] + self_pv = layer_group.pool_views[self_pi] + if layer_group.kind != CacheKind.PAGED or self_pv.mapper_kind == MapperKind.REPLICATED: + continue + peer_pv = peer_page_table.layer_groups[peer_lg].pool_views[peer_pi] + if self_pv.num_kv_heads is None and peer_pv.num_kv_heads is None: + continue + bytes_per_head = [] + for side, ri, lg_idx, pi, pool_view in ( + ("local", self._ri, self_lg, self_pi, self_pv), + ("peer", peer_ri, peer_lg, peer_pi, peer_pv), + ): + heads = policy.view_kv_heads(ri, pool_view) + buffers = self._get_buffers_per_layer(pool_view, layer_group_id=lg_idx, pool_idx=pi) + layer_bytes = get_layer_byte_ranges(pool_view)[1] + if heads <= 0 or layer_bytes % (buffers * heads) != 0: + raise ValueError( + f"PeerRegistrar: the {side} view {sorted(pool_view.pool_role)} holds " + f"{layer_bytes} bytes per layer, not a multiple of {buffers} buffers x " + f"{heads} KV heads" + ) + bytes_per_head.append(layer_bytes // (buffers * heads)) + if bytes_per_head[0] != bytes_per_head[1]: + missing = [ + side + for side, pool_view in (("local", self_pv), ("peer", peer_pv)) + if pool_view.num_kv_heads is None + ] + hint = ( + f"; the {' and '.join(missing)} view has no per-view KV-head count and " + "falls back to the rank-wide one, so both peers must publish per-view " + "KV-head counts" + if missing + else "" + ) + raise ValueError( + "PeerRegistrar: bytes per KV head differ for view " + f"{sorted(self_pv.pool_role)} (local={bytes_per_head[0]}, " + f"peer={bytes_per_head[1]}){hint}" + ) + def get_pool_mapping(self, peer_ri: RankInfo) -> Dict[LGPoolKey, LGPoolKey]: """Get mapping from (self_lg_idx, self_pool_idx) -> (peer_lg_idx, peer_pool_idx). @@ -286,7 +377,7 @@ def _peer_l2g(kind: CacheKind) -> Dict[int, int]: # overlapping layer set; other same-role peer pools cover # disjoint layers and fall out via the overlap filter. self_layer_set = set(pv_global_ids) - matched_peer_pi = None + candidates = [] for peer_pi, peer_pv in enumerate(peer_lg.pool_views): if peer_pv.pool_role != self_pv.pool_role: continue @@ -300,11 +391,52 @@ def _peer_l2g(kind: CacheKind) -> Dict[int, int]: f"(local={self_pv.mapper_kind.name}, " f"peer={peer_pv.mapper_kind.name}, peer_pool={peer_pi})" ) - matched_peer_pi = peer_pi - break - - if matched_peer_pi is not None: - mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) + candidates.append((peer_pi, set(peer_global_ids))) + + if not candidates: + unmatched = [ + sorted(peer_pv.pool_role) + for peer_pv in peer_lg.pool_views + if self_layer_set & set(get_pool_view_global_layer_ids(peer_pv, peer_lg)) + ] + if unmatched and self_lg.kind == CacheKind.PAGED: + raise ValueError( + "PeerRegistrar.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}, role={sorted(self_pv.pool_role)}) " + f"has no peer view with the same roles; the peer holds those layers " + f"as {unmatched}" + ) + continue + # The pairing must be exact: one view per side carrying the + # same layers of the role. A self view split over several peer + # views (or the reverse) means the peers partition their + # layers differently, e.g. target and draft layers pooled + # together on one side only; matching the first candidate would + # silently drop the rest of the layers. + if len(candidates) > 1: + raise ValueError( + "PeerRegistrar.get_pool_mapping: pool view " + f"(lg={self_lg_idx}, pool={self_pi}, role={sorted(self_pv.pool_role)}) " + f"overlaps several peer views {[pi for pi, _ in candidates]}; " + "the peers split their layers into views differently" + ) + matched_peer_pi, peer_layer_set = candidates[0] + self_other_layers = { + gid + for other_pi, other_pv in enumerate(self_lg.pool_views) + if other_pi != self_pi and other_pv.pool_role == self_pv.pool_role + for gid in get_pool_view_global_layer_ids(other_pv, self_lg) + } + misplaced = (peer_layer_set & self_other_layers) - self_layer_set + if misplaced: + raise ValueError( + "PeerRegistrar.get_pool_mapping: peer view " + f"(lg={peer_lg_idx}, pool={matched_peer_pi}) also holds layers " + f"{sorted(misplaced)} that this rank keeps in another view of role " + f"{sorted(self_pv.pool_role)}; the peers split their layers into " + "views differently" + ) + mapping[(self_lg_idx, self_pi)] = (peer_lg_idx, matched_peer_pi) self._lg_pool_mapping_cache[key] = mapping return mapping @@ -526,6 +658,7 @@ def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerO ranks.append(pp * peer_ri.tp_size * peer_ri.cp_size + tp * peer_ri.cp_size + cp) dup_head, peer_dup_head = self._get_policy(CacheKind.PAGED).duplicate_head_factors(peer_ri) + view_dup_factors = self._view_duplicate_head_factors(peer_ri, (dup_head, peer_dup_head)) targets = PeerOverlap( overlap_pp_size=overlap_pp_size, @@ -534,10 +667,35 @@ def get_peer_overlap(self, peer_rank_info: RankInfo, peer_dp_rank: int) -> PeerO duplicate_head_factor=dup_head, peer_duplicate_head_factor=peer_dup_head, ranks=ranks, + view_duplicate_head_factors=view_dup_factors, ) self._overlap_cache[key] = targets return targets + def _view_duplicate_head_factors( + self, peer_ri: RankInfo, rank_wide: Tuple[int, int] + ) -> Dict[LGPoolKey, Tuple[int, int]]: + """Head duplication factors of the matched view pairs that differ from ``rank_wide``. + + Each pair uses the head counts of its own two views, so views of one + rank that shard or replicate heads differently (target and draft + layers) get their own factor. + """ + policy = self._get_policy(CacheKind.PAGED) + self_page_table = self._self_ext_cache.page_table + peer_page_table = peer_ri.page_table + factors: Dict[LGPoolKey, Tuple[int, int]] = {} + for (self_lg, self_pi), (peer_lg, peer_pi) in self.get_pool_mapping(peer_ri).items(): + layer_group = self_page_table.layer_groups[self_lg] + self_pv = layer_group.pool_views[self_pi] + if layer_group.kind != CacheKind.PAGED or self_pv.mapper_kind == MapperKind.REPLICATED: + continue + peer_pv = peer_page_table.layer_groups[peer_lg].pool_views[peer_pi] + pair = policy.duplicate_head_factors(peer_ri, self_pv, peer_pv) + if pair != rank_wide: + factors[(self_lg, self_pi)] = pair + return factors + def _owns_tp_fan_in(self, peer_rank_info: RankInfo) -> bool: """Elect one owner when replicated bytes fan in across TP ranks. @@ -579,7 +737,16 @@ def should_send_pool( pool_view = layer_group.pool_views[pool_idx] # Delegate to the policy's should_send. None means fan-in election. policy = self._get_policy(layer_group.kind) - result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) + kwargs = {} + if layer_group.kind == CacheKind.PAGED: + # The sending-side duplication factor of this view pair, not one + # factor for the whole rank. + kwargs["duplicate_head_factor"] = peer_overlap.duplicate_head_factors_for( + (layer_group_id, pool_idx) + )[0] + result = policy.should_send( + peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind, **kwargs + ) return self._owns_tp_fan_in(peer_rank_info) if result is None else result def should_send_aux(self, peer_rank_info: RankInfo) -> bool: diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 8fbbdc6f9583..b3b3fcb4363b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -48,6 +48,9 @@ class RankInfo: attention: Optional[AttentionInfo] = None aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None + # 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: @@ -112,6 +115,9 @@ def from_kv_cache_manager( ), aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), + 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/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 3c7c4a7f45b6..501b0f3dd818 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3138,16 +3138,24 @@ def _fanin_bounce_safe( ) -> bool: """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in - expected_transfers, so the live writers overflow their slots; + * every head duplication factor is 1, on both sides, rank-wide and in each matched view + pair (target and draft views may replicate heads differently) -- else some ranks don't + send KV yet still count in expected_transfers, so the live writers overflow their slots; * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; for PP fan-in, every overlapping stage must hold the same number of layers (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage list isn't available (shorter than the fan-in degree), be conservative and fall back. Otherwise fall back to the per-fragment path (correct, just not coalesced). - Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects - the mismatched case, so this only needs the count to split evenly.""" - if overlap.duplicate_head_factor != 1: + Equal layer count means equal bytes only when the per-block sizes match; for PP fan-in + (writers_split_layers) reserve() rejects the mismatched case, so this only needs the count + to split evenly.""" + # Any duplicated head, rank-wide or in a single view pair, leaves some + # writers without a share (sender side) or doubles one (receiver side). + rank_factors = (overlap.duplicate_head_factor, overlap.peer_duplicate_head_factor) + view_factors = getattr(overlap, "view_duplicate_head_factors", {}).values() + if any(factor != 1 for factor in rank_factors) or any( + factor != 1 for pair in view_factors for factor in pair + ): return False if overlap.overlap_pp_size > 1: lpp = getattr(peer_ri, "layer_num_per_pp", None) @@ -3272,8 +3280,12 @@ def dispatch_task(self, task: KVRecvTask) -> None: receiver_page_table=self._registrar.self_extractor.page_table, dst_slot=mamba_dst_slot, ) + # Fan-in writers split layers only across PP stages; read the overlap only for a fan-in. bounced = allow_bounce and self._bounce.reserve( - receiver_req, task.expected_transfers, extra_bytes=extra_bytes + receiver_req, + task.expected_transfers, + extra_bytes=extra_bytes, + writers_split_layers=task.expected_transfers > 1 and topo_overlap.overlap_pp_size > 1, ) session = self._get_session(task._unique_rid) if session is None: @@ -3378,14 +3390,15 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() - # Recurrent-state (Mamba/KDA) layout gate on the receiver side. - # The sender-side check (PeerRegistrar.register) runs in the - # sender's listener thread, where exceptions are only logged, so - # reject here — before REGISTER_RANK_INFO is even sent, so no - # dealers are connected and no partial registration happens for - # the bad peer. The failure is converted to PeerIncompatibleError - # (handled in dispatch_task) so only requests targeting this peer - # fail, and cached so later requests fail fast. + # Recurrent-state (Mamba/KDA), replicated-role and draft-layer layout + # gates on the receiver side. The sender-side check + # (PeerRegistrar.register) runs in the sender's listener thread, + # where exceptions are only logged, so reject here — before + # REGISTER_RANK_INFO is even sent, so no dealers are connected and + # no partial registration happens for the bad peer. The failure is + # converted to PeerIncompatibleError (handled in dispatch_task) so + # only requests targeting this peer fail, and cached so later + # requests fail fast. try: MambaPolicy.validate_peer_compatible( self._registrar.self_rank_info, @@ -3397,6 +3410,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: self._registrar.self_extractor.page_table, sender_info.page_table, ) + PeerRegistrar.validate_draft_layers(self._registrar.self_rank_info, sender_info) except ValueError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index a5fb8f56c061..2ddf71c268b6 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -14,7 +14,7 @@ # limitations under the License. from collections import defaultdict -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import numpy as np @@ -501,6 +501,7 @@ def _build_pool_views_for_variant( layer_group_id: int, role_layouts: Optional[Dict] = None, ignored_roles: frozenset[DataRole] = frozenset(), + kv_heads_for_layer: Optional[Callable[[int], Optional[int]]] = None, ) -> List[PoolView]: """Bucket one slot-desc variant's coalesced buffers into pool views. @@ -508,11 +509,15 @@ def _build_pool_views_for_variant( kind share a view (KEY+VALUE), roles with different kinds in the same physical pool get separate views (M3 coalesced index-K). ``role_layouts`` supplies the per-role resharding geometry carried onto the view. + ``kv_heads_for_layer`` gives a + layer's per-rank KV-head count, which head-sharded views record; all + layers of a view must agree on it. """ role_layouts = role_layouts or {} bucket_entries: Dict[tuple, list] = defaultdict(list) bucket_roles: Dict[tuple, set] = defaultdict(set) bucket_layouts: Dict[tuple, set] = defaultdict(set) + bucket_layers: Dict[tuple, set] = defaultdict(set) for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): single_buffer_size = int(coalesced_buffer.single_buffer_size) offset = 0 @@ -527,6 +532,7 @@ def _build_pool_views_for_variant( ) bucket_roles[bucket_key].add(str(buffer_id.role)) bucket_layouts[bucket_key].add(role_layouts.get(buffer_id.role)) + bucket_layers[bucket_key].add(int(buffer_id.layer_id)) offset += single_buffer_size # All ordering below is canonicalization — the page table is @@ -570,6 +576,20 @@ def _build_pool_views_for_variant( f"{context} bytes_per_layer={bytes_per_layer} is not a multiple of " f"bytes_per_head={layout.bytes_per_head}" ) + num_kv_heads = None + if kv_heads_for_layer is not None and mapper_kind != MapperKind.REPLICATED: + layer_heads = { + heads + for heads in (kv_heads_for_layer(layer) for layer in bucket_layers[bucket_key]) + if heads + } + if len(layer_heads) > 1: + raise ValueError( + f"{context} mixes layers with different per-rank KV-head counts " + f"{sorted(layer_heads)}; layers with different head geometry must be " + "in separate views" + ) + num_kv_heads = layer_heads.pop() if layer_heads else None pool_views.append( PoolView( pool_idx=pool_idx, @@ -581,6 +601,7 @@ def _build_pool_views_for_variant( list(layout.section_bytes) if layout.section_bytes is not None else None ), bytes_per_head=layout.bytes_per_head, + num_kv_heads=num_kv_heads, ) ) @@ -653,6 +674,10 @@ def _window_size_for_layer(internal_layer_id: int): raise ValueError(f"Cannot resolve layer config for internal layer {internal_layer_id}") + def _kv_heads_for_layer(internal_layer_id: int) -> Optional[int]: + heads = manager.num_kv_heads_per_layer + return int(heads[internal_layer_id]) if internal_layer_id < len(heads) else None + pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) @@ -716,6 +741,7 @@ def _window_size_for_layer(internal_layer_id: int): layer_group_id, role_layouts=role_layouts, ignored_roles=ignored_roles, + kv_heads_for_layer=None if is_recurrent else _kv_heads_for_layer, ) for pool_view in pool_views: if pool_view.mapper_kind not in _MAPPER_KINDS_BY_CACHE_KIND[cache_kind]: diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index cf70b83cbc08..712e51e7b0e2 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -218,7 +218,12 @@ class PoolView: bytes_per_head: INDEXED views whose head count is not derivable from ``RankInfo.attention`` (recurrent SSM state). ``bytes_per_layer // bytes_per_head`` is the per-rank head count. Attention K/V leaves - it ``None`` and uses ``kv_heads_per_rank`` instead. + it ``None``; its head count is ``num_kv_heads``. + num_kv_heads: Attention K/V views only. Per-rank KV-head count shared + by every layer of the view, so views of one rank (e.g. target and + draft layers) may shard or replicate heads differently. ``None`` + in tables from older builders and for views without KV heads; + consumers then fall back to ``AttentionInfo.kv_heads_per_rank``. """ pool_idx: int @@ -228,6 +233,7 @@ class PoolView: bytes_per_layer: Optional[int] = None section_bytes: Optional[List[int]] = None bytes_per_head: Optional[int] = None + num_kv_heads: Optional[int] = None def to_dict(self) -> dict: return { @@ -244,6 +250,7 @@ def to_dict(self) -> dict: "bytes_per_head": ( int(self.bytes_per_head) if self.bytes_per_head is not None else None ), + "num_kv_heads": int(self.num_kv_heads) if self.num_kv_heads is not None else None, } @staticmethod @@ -270,6 +277,9 @@ def from_dict(data: dict) -> "PoolView": bytes_per_head=( int(data["bytes_per_head"]) if data.get("bytes_per_head") is not None else None ), + num_kv_heads=( + int(data["num_kv_heads"]) if data.get("num_kv_heads") is not None else None + ), ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e43785e7e956..68e370b61ed7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -31,7 +31,8 @@ CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig, - SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) + SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy, + disagg_shares_draft_kv) # isort: on from tensorrt_llm._torch.peft.lora.config import ( LoraConfig, get_default_trtllm_modules_to_hf_modules) @@ -64,7 +65,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, @@ -687,6 +690,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), its slot sizes are the sizes of its + coalesced buffers in descending order, and life cycles with the same slot + sizes share one pool group. [agentx port] The agentx bindings have no + per-layer reuse policy (main #19727), so it is not part of either key. + """ + 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)) + 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({ + tuple( + sorted((size * count for (_, size), count in coalesced.items()), + reverse=True)) + for coalesced in life_cycles.values() + }) + + def _get_num_pool_groups_for_estimation( model_config: object, max_seq_len: int, @@ -862,6 +894,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 @@ -1004,6 +1038,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) @@ -1486,12 +1525,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 @@ -1888,6 +1935,16 @@ 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( + 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( model_engine=model_engine, kv_cache_manager_cls=kv_cache_manager_cls, @@ -1912,6 +1969,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: @@ -1951,6 +2009,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 @@ -1976,6 +2036,278 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) + 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 + 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. + + A disaggregated worker shares as ``disagg_shares_draft_kv`` decides. + """ + if self._dflash_unify_blocker() is not None: + return False + return not self._is_disagg or disagg_shares_draft_kv( + self._speculative_config) + + def _check_shared_draft_kv_request(self) -> None: + """Fail at startup when disaggregated draft KV sharing is missing or cannot be honored.""" + if (not self._is_disagg or self._speculative_config is None + or self._unified_drafter_attention_backend() is None): + return + if os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": + # The drafter writes from column 0; a separate pool drops blocks at real positions. + if (self._unified_drafter_attention_backend() == "TRTLLM" + and self._should_create_separate_draft_kv_cache() + and self._get_draft_max_attention_window( + self._unification_max_seq_len, + self._kv_cache_config) is not None): + blocker = self._dflash_unify_blocker() + fix = ( + "use attention_backend VANILLA" if blocker else + "set TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on every worker, or " + "use attention_backend VANILLA") + raise ValueError( + f"{self._unified_drafter_name()} with attention_backend " + "TRTLLM on a disaggregated worker must share its " + "sliding-window draft KV with the target, but " + f"{blocker or 'TRTLLM_DISAGG_SHARED_DRAFT_KV is not 1'}: " + f"{fix}. A separate draft KV cache retires the blocks " + "behind the window at a request's real position while a " + "transferred request's drafter fills it from position 0, " + "so acceptance length collapses once a request passes the " + "window.") + 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 "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() + 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( + 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 "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, + list( + self._mapping.pp_layers( + int(draft_model_config.pretrained_config.num_hidden_layers)) + )) + if len(set(draft_windows)) > 1: + 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( + 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 (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( + 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. + + Every view of a rank's page table carries its own KV-head count and byte + geometry, so the draft's head count (and hence its sharding or + replication across TP), head_dim and dtype may differ from the target's. + The K-only flag decides the buffer count per layer and must match. NVFP4 + adds block-scale views whose heads are not resharded, so an NVFP4 side + is only paired with an identical dtype and head count. + """ + target_heads, _, target_dtype, target_k_only = target_layout + draft_heads, _, draft_dtype, draft_k_only = draft_layout + if target_k_only != draft_k_only: + return False + nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 + if nvfp4 in (target_dtype, draft_dtype): + return target_dtype == draft_dtype and target_heads == draft_heads + return True + + 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, + ) 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, + ) 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. @@ -2037,12 +2369,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 @@ -3117,7 +3449,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 @@ -3292,6 +3626,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/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 37fc9e0e3a6d..68007e9f8717 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 @@ -34,6 +34,8 @@ import numpy as np import torch +import triton +import triton.language as tl from strenum import StrEnum from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout @@ -272,6 +274,117 @@ class BlockReusePolicy(StrEnum): PER_CONVERSATION = "per_conversation" +class KVCacheLayerOwner(StrEnum): + """Model component that reads and writes a cache layer.""" + + TARGET = "target" + DRAFT = "draft" + + +@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. + """ + + 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 + # [agentx port] No reuse_policy field: the agentx KV bindings predate main's per-layer + # AttentionReusePolicy (#19727), so every layer has the default (REQUIRED) reuse. + + 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 {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 scale_roles.intersection(self.buffer_roles) != expected_scale_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}" + ) + + 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}") + + @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 _coalescing_group_kwargs(group: int) -> dict[str, int]: + """[agentx port] ``BufferConfig`` kwargs for a buffer's coalescing group. + + The stack's C++ adds ``BufferConfig.coalescing_group``; the agentx image's + prebuilt bindings do not have it. Group 0 is their implicit behaviour + (coalesce by size alone), so the field is passed only for a nonzero group. + """ + if group == 0: + return {} + if not hasattr(BufferConfig, "coalescing_group"): + raise NotImplementedError( + "NVFP4 KV layer specs need BufferConfig.coalescing_group; rebuild the " + "bindings with this branch's cpp/ changes" + ) + return {"coalescing_group": group} + + def _swa_endpoint_priority( prompt_length: int, tokens_per_block: int, rewind_tokens: int ) -> Callable[[int, object], int]: @@ -507,26 +620,29 @@ 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) if scratch_key in scratch_keys: size_per_request += window_tokens * layer_size else: @@ -535,10 +651,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 @@ -551,7 +693,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 @@ -563,6 +707,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. @@ -572,7 +717,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, @@ -581,6 +731,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, @@ -767,6 +918,30 @@ 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], + *, + 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: + 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], + [ + 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 + ], + ) + + class _MmRunMetadata(NamedTuple): """CPU tensors for exact multimodal run lookup. @@ -1157,6 +1332,75 @@ def _copy_swa_block_offsets_with_scratch_compiled( output.copy_(converted.permute(0, 2, 1, 3)) +@triton.jit(do_not_specialize=["num_seqs", "rows_width"]) +def _convert_batch_block_offsets_kernel( + rows_ptr, + out_ptr, + params_ptr, + num_seqs, + rows_width, + out_pool_stride, + num_blocks, + BLOCK: tl.constexpr, + BAD_PAGE: tl.constexpr, +): + # grid = (num_op_pools, num_seqs, cdiv(num_blocks, BLOCK)): one column tile + # of both lanes of one sequence's row in one attention-op table. + op_pool = tl.program_id(0) + seq = tl.program_id(1).to(tl.int64) + cols = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + num_params = tl.num_programs(0) * 2 + for lane in tl.static_range(2): + param = params_ptr + op_pool * 2 + lane + pool = tl.load(param).to(tl.int64) + scale = tl.load(param + num_params) + offset = tl.load(param + 2 * num_params) + # Columns past the snapshot's width are BAD_PAGE_INDEX, as in the host rows. + src = rows_ptr + (pool * num_seqs + seq) * rows_width + cols + base = tl.load(src, mask=cols < rows_width, other=BAD_PAGE) + value = tl.where(base == BAD_PAGE, BAD_PAGE, base * scale + offset) + dst = out_ptr + op_pool.to(tl.int64) * out_pool_stride + seq * 2 * num_blocks + tl.store(dst + lane * num_blocks + cols, value, mask=cols < num_blocks) + + +def _convert_batch_block_offsets( + rows: torch.Tensor, out: torch.Tensor, params: torch.Tensor, num_seqs: int +) -> None: + """out[t, :num_seqs, lane] = rows[params[0, t, lane]] * params[1, t, lane] + params[2, t, lane]. + + ``rows`` is [num_pools, num_seqs, width] base page indices; BAD_PAGE_INDEX is kept and + written to the columns past ``width``. + """ + num_blocks = out.shape[-1] + assert rows.is_contiguous() and rows.shape[1] == num_seqs and rows.shape[2] <= num_blocks + assert out.shape[2] == 2 and out[0].is_contiguous() + assert params.shape == (3, out.shape[0], 2) and params.dtype == torch.int32 + assert params.is_contiguous() + block = 1024 + grid = (out.shape[0], num_seqs, triton.cdiv(num_blocks, block)) + _convert_batch_block_offsets_kernel[grid]( + rows, + out, + params, + num_seqs, + rows.shape[2], + out.stride(0), + num_blocks, + BLOCK=block, + BAD_PAGE=BAD_PAGE_INDEX, + ) + + +class _AttentionOpPoolGroup(NamedTuple): + """Layers sharing one attention-op page table (see _get_attention_op_pool_groups).""" + + pool_id: int # life cycle whose base page indices fill the table + base_address: int # pool pointer: the first member's page in the physical pool + scale: int # pages per slot of the physical pool + kv_offset: int # V page - K page + pool_mapping_offsets: Dict[int, int] # local layer id -> kv_cache_pool_mapping offset + + def _rewind_context_cursor(req: LlmRequest, reuse: int) -> None: """Pull a context cursor back to a shallower prefix. ``setPrepopulatedPromptLen`` only moves forward, so leave the forward case to it -- it re-floors chunk size too. @@ -1227,6 +1471,20 @@ 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 + # Attention-op tables shared by layer groups instead of one per layer, else None. + _attention_op_pool_groups: Optional[List["_AttentionOpPoolGroup"]] = None + + @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 @@ -1265,10 +1523,13 @@ def __init__( cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> None: self.mapping = mapping self.dtype = dtype + self._declared_target_layer_count = num_layers + 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 @@ -1276,13 +1537,34 @@ 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 + 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 self.layer_specs is not None: + 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") # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the @@ -1293,7 +1575,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) @@ -1379,6 +1666,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 @@ -1396,7 +1686,9 @@ def __init__( self.reuse_match_backoff = 0 # 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 @@ -1419,6 +1711,13 @@ def __init__( self.pp_layers, ) assert len(self.max_attention_window_vec) == self.num_local_layers + if self.layer_specs is not None: + self.max_attention_window_vec = [ + 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 # committed endpoint so commit_min_snapshot still covers the rewind. @@ -1485,7 +1784,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) ] @@ -1515,7 +1819,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, ( @@ -1678,9 +1984,16 @@ 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, + 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, @@ -2048,7 +2361,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() @@ -2192,7 +2505,10 @@ 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) + 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): @@ -2229,12 +2545,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 @@ -2247,10 +2589,20 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] - if self._use_per_layer_page_tables: + 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._attention_op_pool_groups is not None: + kv_cache_pool_mapping_list = [None] * self.num_local_layers + for group_idx, group in enumerate(self._attention_op_pool_groups): + kv_cache_pool_pointers_list.append([group.base_address, 0]) + for layer_id, offset in group.pool_mapping_offsets.items(): + kv_cache_pool_mapping_list[layer_id] = [group_idx, offset] + elif 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( @@ -2259,7 +2611,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( [ @@ -2280,7 +2632,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 @@ -2297,7 +2649,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 = ( @@ -2323,13 +2675,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) @@ -2346,7 +2701,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, ) @@ -2357,7 +2712,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] = [ @@ -2397,7 +2752,17 @@ 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._use_per_layer_page_tables: + 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 + # Layers whose per-layer tables differ only by a page offset share one table. + self._attention_op_pool_groups = ( + self._get_attention_op_pool_groups() if self._use_per_layer_page_tables else None + ) + if self._attention_op_pool_groups is not None: + self._use_per_layer_page_tables = False + self.num_attention_op_pools = len(self._attention_op_pool_groups) + elif self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: self.num_attention_op_pools = self.num_pools @@ -2437,6 +2802,73 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) if self._use_per_layer_page_tables: self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) + elif self._attention_op_pool_groups is not None: + groups = self._attention_op_pool_groups + # Lane l of table t: scale * base index of the group's life cycle + K/V offset. + self._prepare_batch_block_offsets_copy( + index_mapper_capacity * self.max_copy_beam_width, + torch.tensor([[g.pool_id, g.pool_id] for g in groups], dtype=torch.int32), + torch.tensor([[g.scale, g.scale] for g in groups], dtype=torch.int32), + torch.tensor([[0, g.kv_offset] for g in groups], dtype=torch.int32), + ) + + def _get_attention_op_pool_groups(self) -> Optional[List[_AttentionOpPoolGroup]]: + """Group layers that can share one attention-op page table, else None. + + A per-layer table holds ``base * scale + layer_offset`` pages from its + physical pool's origin. Layers of one life cycle and physical pool with + the same page stride, scale, K/V distance and kv factor differ only by + ``layer_offset``, which, as on the SHARED path, can move to the pool + mapping (the attention op adds ``offset * kv_factor`` pages). Unified + target/draft caches then need a table per (life cycle, page size), not + per layer. Not for SWA scratch slots, NVFP4 scale pools, expanded page + indices, or subclasses with their own layouts. + """ + if type(self) is not KVCacheManagerV2 or self.enable_swa_scratch_reuse: + return None + members: Dict[tuple, List[Tuple[int, int]]] = {} + for layer_id in typed_range(LayerId(self.num_local_layers)): + if self._get_layer_dtype_by_local_id(layer_id) == DataType.NVFP4: + return None + pool_id = self.impl.get_layer_group_id(layer_id) + role_a, role_b = self._get_layer_page_table_roles(layer_id, pool_id) + roles = (role_a,) if role_b is None else (role_a, role_b) + key = self.impl.get_page_index_converter(layer_id, role_a) + value = key if role_b is None else self.impl.get_page_index_converter(layer_id, role_b) + strides = {self.impl.get_page_stride(layer_id, r) for r in roles} + origins = { + self.impl.get_mem_pool_base_address(layer_id, r, PageIndexMode.PER_LAYER) + for r in roles + } + # K and V must be addressable from one pool pointer with one stride. + if len(strides) > 1 or len(origins) > 1 or value.scale != key.scale: + return None + if key.expansion != 1 or value.expansion != 1: + return None + kv_factor = self._get_layer_kv_factor_by_local_id(layer_id) + group_key = ( + int(pool_id), + origins.pop(), + strides.pop(), + int(key.scale), + int(value.layer_offset) - int(key.layer_offset), + kv_factor, + int(key.layer_offset) % kv_factor, + ) + members.setdefault(group_key, []).append((int(layer_id), int(key.layer_offset))) + groups = [] + for (pool_id, origin, stride, scale, kv_offset, kv_factor, _), layers in members.items(): + first = min(offset for _, offset in layers) + groups.append( + _AttentionOpPoolGroup( + pool_id, + origin + first * stride, + scale, + kv_offset, + {layer_id: (offset - first) // kv_factor for layer_id, offset in layers}, + ) + ) + return groups def _kv_pool_mapping_offset( self, layer_id: LayerId, layer_group_id: int, key_base_addr: int @@ -2457,11 +2889,16 @@ 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, ) def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: + 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): @@ -2498,6 +2935,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: @@ -2534,6 +2972,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 @@ -2693,7 +3132,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, @@ -2713,6 +3152,12 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: staging_capacity = index_mapper_capacity * self.max_copy_beam_width device = torch.device("cuda", torch.cuda.current_device()) + if not self.enable_swa_scratch_reuse: + # No scratch slots: each step converts only the batch's rows. + self._prepare_batch_block_offsets_copy( + staging_capacity, pool_ids, scales, layer_offsets + ) + return self._device_kv_cache_block_offsets_input = torch.empty_like( self.host_kv_cache_block_offsets, device=device, @@ -2840,6 +3285,61 @@ def _copy_scratch_metadata_to_device( self._device_scratch_slots_staging, ) + def _prepare_batch_block_offsets_copy( + self, + capacity: int, + pool_ids: torch.Tensor, + scales: torch.Tensor, + offsets: torch.Tensor, + ) -> None: + """Buffers for converting only the batch's rows of the attention-op tables. + + Lane ``l`` of table ``t`` is ``scales[t, l] * base + offsets[t, l]``, where + ``base`` is the request's base page index in pool ``pool_ids[t, l]``. + """ + device = torch.device("cuda", torch.cuda.current_device()) + # Device copy of each step's host-row snapshot, viewed at that step's shape. + self._device_batch_block_offsets = torch.empty( + self.num_pools * capacity * self.max_blocks_per_seq, dtype=torch.int32, device=device + ) + self._device_attention_op_params = torch.stack( + [pool_ids.to(torch.int32), scales, offsets] + ).to(device) + + def _copy_batch_block_offsets_by_op_pool( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + copy_idx: torch.Tensor, + num_seqs: int, + ) -> None: + # A captured graph would replay this step's transient snapshot. + assert not torch.cuda.is_current_stream_capturing() + if num_seqs == 0: + return + # A row is BAD_PAGE_INDEX past its cache's num_blocks: snapshot up to the widest. + caches = [self.kv_cache_map.get(request_id) for request_id in request_ids] + width = ( + self.max_blocks_per_seq + if None in caches + else max([1] + [cache.num_blocks for cache in caches]) + ) + # Snapshot the batch's host rows now: an async copy reads its source when it + # runs, and the host table keeps changing for the next step (nvbug 6293536). + # The caching host allocator keeps this buffer until the copy has run. + rows = torch.empty( + self.num_pools, num_seqs, width, dtype=torch.int32, pin_memory=prefer_pinned() + ) + torch.index_select(self.host_kv_cache_block_offsets[:, :, 0, :width], 1, copy_idx, out=rows) + device_rows = self._device_batch_block_offsets[: rows.numel()].view(rows.shape) + device_rows.copy_(rows, non_blocking=True) + _convert_batch_block_offsets( + device_rows, + dst_tensor[: self.num_attention_op_pools], + self._device_attention_op_params, + num_seqs, + ) + def _copy_batch_block_offsets_per_layer( self, dst_tensor: torch.Tensor, @@ -2848,6 +3348,9 @@ def _copy_batch_block_offsets_per_layer( num_contexts: int, num_seqs: int, ) -> None: + if not self.enable_swa_scratch_reuse: + self._copy_batch_block_offsets_by_op_pool(dst_tensor, request_ids, copy_idx, num_seqs) + return device_copy_idx = self._copy_idx_to_device(copy_idx) self._device_kv_cache_block_offsets_input.copy_( self.host_kv_cache_block_offsets, @@ -2882,6 +3385,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) @@ -2947,7 +3452,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)] @@ -2968,50 +3473,78 @@ 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, ) ] ) ) - # 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( + if self.layer_specs is not None: + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + layer_configs = [ AttentionLayerConfig( - layer_id=layer_id, - buffers=buffers, - sliding_window_size=self.max_attention_window_vec[layer_id], - num_sink_tokens=None, + layer_id=LayerId(local_layer_id), + buffers=[ + 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. + # [agentx port] Only NVFP4 groups reach the binding; see + # _coalescing_group_kwargs. + **_coalescing_group_kwargs( + 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 + ), + ) + 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, ) + 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 + # 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. @@ -3167,6 +3700,79 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + def get_disagg_pool_view_partition(self, layer_id: LayerId, role: DataRole): + """Keep target and draft layers in separate transfer views. + + Target and draft layers may share a physical pool, but each model has + its own head geometry, so the transfer layout must keep them apart + whatever the storage coalescing does. Managers without draft-owned + layers need no partition. + """ + layer_specs = getattr(self, "layer_specs", None) + if layer_specs is None or not any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs + ): + return None + return layer_specs[int(layer_id)].owner.value + + 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: + 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, ...]: + 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, ...]: + 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, ...]: + return tuple( + layer_id + for layer_id in self.pp_layers + 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: """ @@ -3182,9 +3788,11 @@ 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) - 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 ) @@ -3196,7 +3804,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 @@ -3204,28 +3812,42 @@ 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, ] - 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, @@ -3361,7 +3983,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 = ( @@ -3858,7 +4480,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 @@ -3898,7 +4520,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: @@ -3929,7 +4551,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: # Helix requests carry the rank-local strided slice in prompt_len; # the global ledger sizes off the full prompt instead. 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 @@ -4531,7 +5153,7 @@ def _resize_for_connector_prefix( 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 @@ -4675,7 +5297,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( @@ -5379,10 +6001,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 @@ -5512,7 +6141,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 @@ -5693,6 +6322,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], @@ -5701,7 +6339,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. @@ -5709,15 +6350,22 @@ def get_batch_cache_indices( if layer_idx is None: pool_id = 0 index_scale = 1 if raw_indices else None + 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) + 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( @@ -5728,16 +6376,17 @@ 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 - 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): @@ -5781,6 +6430,11 @@ def get_batch_cache_indices_flat( if layer_idx is None: pool_id = 0 scale = self._index_scale_ints[pool_id] + 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 @@ -5788,7 +6442,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.kv_factor + 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() @@ -5809,6 +6464,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) @@ -5824,6 +6481,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, @@ -5886,12 +6549,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) @@ -5916,10 +6593,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() @@ -5983,20 +6666,26 @@ def get_cache_size_per_token( max_num_tokens: int = 0, spec_config=None, is_draft: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = 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 layer_specs is not None: + 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, + 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 and kv_cache_config.enable_block_reuse @@ -6013,7 +6702,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, @@ -6027,8 +6716,10 @@ 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, + layer_specs=layer_specs, ) # The affine slope covers all tokens; context additionally retains SWA # pages for the current token batch beyond the generation windows. @@ -6039,16 +6730,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, @@ -6067,6 +6785,7 @@ def _get_per_request_layer_tokens( spec_config=None, is_draft: bool = False, joint_kv_cache_reuse: bool = False, + 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. @@ -6095,7 +6814,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 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 @@ -6104,18 +6823,23 @@ 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 layer_specs is not None: + 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, + 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`` @@ -6136,7 +6860,7 @@ def _get_per_request_layer_tokens( # mirror 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 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 @@ -6167,14 +6891,14 @@ 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) * 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 @@ -6182,15 +6906,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 ) @@ -6210,6 +6936,7 @@ def get_cache_bytes_per_request( spec_config=None, is_draft: bool = False, joint_kv_cache_reuse: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -6243,6 +6970,7 @@ def get_cache_bytes_per_request( spec_config=spec_config, is_draft=is_draft, joint_kv_cache_reuse=joint_kv_cache_reuse, + layer_specs=layer_specs, ) if layer_tokens is None: return None @@ -6272,6 +7000,7 @@ def get_cache_quota_per_request( spec_config=None, is_draft: bool = False, joint_kv_cache_reuse: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -6307,6 +7036,7 @@ def get_cache_quota_per_request( spec_config=spec_config, is_draft=is_draft, joint_kv_cache_reuse=joint_kv_cache_reuse, + layer_specs=layer_specs, ) if layer_tokens is None: return None @@ -6470,6 +7200,9 @@ def copy_batch_block_offsets( dst_tensor, request_ids, copy_idx, num_contexts, num_seqs ) return + if self._attention_op_pool_groups is not None: + self._copy_batch_block_offsets_by_op_pool(dst_tensor, request_ids, copy_idx, num_seqs) + return copy_batch_block_offsets_to_device( self.host_kv_cache_block_offsets, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 01f338aa3725..37c679991f6a 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -1060,9 +1060,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/__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 613071340c06..59333b9278af 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 @@ -519,11 +520,13 @@ def dflash_allocated_ctx_limit( ``_refresh_ctx_block_tables`` clamped to 0 -- another request's block. A silent cross-request write, not a fault. - Per-request only on the ``python`` V2 backend, which propagates - ``BAD_PAGE_INDEX``. The default ``cpp`` one maps it to 0 - (kvCacheManagerV2Utils.cu), every row counts full, and this degenerates to - the block-table-width bound upstream already applied. Do not read a - per-request guarantee into it there. + ``block_counts`` is the extent of each request's allocation (last valid + table entry + 1), not the number of valid entries: a sliding window's + evicted pages precede its live tail. Per-request only on the ``python`` V2 + backend, which propagates ``BAD_PAGE_INDEX``. The default ``cpp`` one maps + it to 0 (kvCacheManagerV2Utils.cu), every row spans its full width, and this + degenerates to the block-table-width bound upstream already applied. Do not + read a per-request guarantee into it there. """ return (block_counts * page_size - block_size).clamp(min=0) @@ -617,24 +620,41 @@ 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 + # 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(), @@ -773,10 +793,15 @@ def __init__( self._ctx_page_table = None self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 self._ctx_block_indptr = None - # Pages the manager handed this request, per batch row. Saturates at - # the table width on the default backend -- dflash_allocated_ctx_limit. + # Pages the manager's allocation spans for this request (last valid + # entry + 1), per batch row. Saturates at the table width on the + # default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None + self._ctx_block_cols = None # 1..max_blocks_per_seq, for the extent 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 @@ -939,7 +964,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) @@ -972,10 +1008,42 @@ 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] + unsupported = None + if mgr.enable_swa_scratch_reuse: + 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 + ): + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + 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. + 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." @@ -984,12 +1052,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 " @@ -1006,6 +1085,7 @@ def _init_ctx_block_tables( 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" ) self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") + self._ctx_block_cols = torch.arange(1, max_blocks + 1, dtype=torch.long, device="cuda") logger.info( f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}, " f"pool_idx={self._ctx_pool_idx}, divisor={self._ctx_block_divisor}, " @@ -1013,6 +1093,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. @@ -1022,23 +1119,32 @@ 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() - # Before the clamp erases the placeholders. Informative only on the + encoded.sub_(self._ctx_block_offset) + # Before the clamp erases the placeholders: the allocation's extent, + # i.e. last valid entry + 1. Not the number of valid entries -- a + # sliding-window layer's evicted pages are placeholders BEFORE its live + # tail (per-layer tables of a unified cache), so a count stops at the + # window and caps every longer context there. 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)) + valid_cols = (encoded >= 0) * self._ctx_block_cols[: encoded.size(1)] + self._ctx_block_counts[:num_seqs].copy_(valid_cols.amax(dim=1)) decoded = encoded.clamp_(min=0).div_(self._ctx_block_divisor, rounding_mode="floor") self._ctx_block_tables[:num_seqs].copy_(decoded.to(torch.int32)) return True @@ -1047,12 +1153,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." @@ -1072,6 +1184,7 @@ def _lazy_init_ctx_buffers( self._ctx_kv_manager = draft_kv_cache_manager self._ctx_block_tables = None self._ctx_block_counts = None + self._ctx_block_cols = None max_batch = spec_metadata.max_num_requests # ctx_len is 1:1 with the target's positions, so max_seq_len bounds it. @@ -1161,6 +1274,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 @@ -1215,6 +1329,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 @@ -1245,7 +1361,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( @@ -1482,6 +1598,44 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) + def _reseed_warmup_generation_slots(self, spec_metadata: "DFlashSpecMetadata") -> None: + """Seed every generation row's slot from this batch's positions; warmup only. + + Only CUDA-graph warmup runs a graph batch eagerly. Its dummy requests reuse + ids across batches, so a slot keeps the previous batch's length -- which the + replay after each capture advances for good -- while this batch may give the + same id far fewer pages. A unified cache's allocation check then traps. + """ + request_ids = spec_metadata.request_ids or [] + num_contexts = max(0, len(request_ids) - spec_metadata.num_generations) + seq_lens = spec_metadata.seq_lens + token_offset = sum(seq_lens[:num_contexts]) + updates = {} + for rid, seq_len in zip(request_ids[num_contexts:], seq_lens[num_contexts:]): + slot = self._req_to_slot.get(rid) + if slot is not None and rid != ATTENTION_DP_DUMMY_REQUEST_ID: + position = int(spec_metadata.host_position_ids[token_offset]) + updates[slot] = min(position, self._max_ctx) + token_offset += seq_len + self._write_ctx_len(updates) + + def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> None: + """Add this step's accepted tokens to the context lengths of ``slots``, on the device. + + Padding, warmup and attention-DP dummy rows share the dummy slot, which no request owns + or frees. Emptying it right after the add keeps it at 0 between steps: those rows attend + no context, and their next context K / V starts at column 0, inside the padding + request's own pages (past them, its page-table row maps page 0, another request's). + The store is to a fixed slot, so a captured step replays it. A unified cache's + allocation check relies on it: a dummy row's pages are the padding request's. + """ + self._ctx_len[slots] += num_accepted + self._ctx_len.clamp_(max=self._max_ctx) + self._ctx_len[self._dummy_slot].zero_() + if self._ctx_base is not None: + # The windowed arena's base moves with the length; empty is (0, 0). + self._ctx_base[self._dummy_slot].zero_() + def _write_ctx_len( self, updates: dict[int, int], base_updates: Optional[dict[int, int]] = None ) -> None: @@ -1779,10 +1933,24 @@ 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() + if ( + self._ctx_unified_cache + and spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + ): + self._reseed_warmup_generation_slots(spec_metadata) # 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. @@ -2029,7 +2197,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 @@ -2306,10 +2479,10 @@ def prepare_1st_drafter_inputs( projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) gen_num_accepted_long = gen_num_accepted.long() if self._ctx_base is None: - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) else: # Windowed arena: physical columns. - col_idx = (self._ctx_len[slots] - self._ctx_base[slots]).unsqueeze( + col_idx = (ctx_len_gen - self._ctx_base[slots]).unsqueeze( 1 ) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) @@ -2361,8 +2534,7 @@ def prepare_1st_drafter_inputs( if v_new is not None: self._ctx_v_buf[slot_long, :, col_long] = v_new - self._ctx_len[slots] += gen_num_accepted_long - self._ctx_len.clamp_(max=self._max_ctx) + self._advance_ctx_len(slots, gen_num_accepted_long) if self._ctx_keep and not self._ctx_paged: self._compact_ctx_window(slots) @@ -2380,6 +2552,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 203edc46a01b..906b1cf1eede 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -943,6 +943,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. @@ -952,8 +962,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/interface.py b/tensorrt_llm/_torch/speculative/interface.py index b6631e5cffce..a4a32cd59229 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -194,7 +194,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 @@ -3184,8 +3185,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) @@ -3206,7 +3209,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 4030e7f7486d..b74889a94a6c 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -786,6 +786,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 @@ -804,7 +819,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/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 5f6ef913bf2c..65e74256d65d 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2447,6 +2447,26 @@ 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") + for cfg in server_configs: + spec = cfg.other_args["speculative_config"] + # AUTO, the default, resolves to the sharing TRTLLM backend for windowed + # drafters. + if (spec.get("decoding_type") not in ("DFlash", "DSpark") or + spec.get("attention_backend") not in (None, "AUTO", "TRTLLM")): + return ("speculative_config to be DFlash or DSpark with " + "attention_backend AUTO or TRTLLM in both sections") + return None + + @click.command("disaggregated_mpi_worker") @stability_option( "-c", @@ -2486,6 +2506,15 @@ 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) + blocker = _shared_draft_kv_blocker(disagg_cfg.server_configs) + shared_draft_kv = os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") + if blocker is not None and shared_draft_kv == "1": + raise click.UsageError( + f"TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires {blocker}.") + # Every worker gets the same value; an explicit 0 keeps the draft KV + # private on all of them. + if shared_draft_kv != "0": + 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/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 96c7b07e41bb..0f25df9d07e0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3316,7 +3316,10 @@ class DFlashDecodingConfig(DecodingBaseConfig): "on the same paged cache and supports SM90 only. AUTO (the default) " "resolves to TRTLLM where it can run and every drafter layer is " "windowed, and to VANILLA otherwise, including on a disaggregated worker " - "and under modeling_v2.") + "with TRTLLM_DISAGG_SHARED_DRAFT_KV=0 and under modeling_v2. In " + "disaggregated serving a TRTLLM drafter keeps its KV with the target's, " + "so the context server's drafter context reaches the generation server; " + "both servers need the same speculative_config.") skip_ctx_buffer_budget_check: bool = Field( default=False, @@ -6003,6 +6006,23 @@ def _check_removed_sampler_type(value: Any) -> None: TORCH_LLMARGS_REMOVED_ARGS = frozenset(_TORCH_LLMARGS_REMOVED_ARG_CHECKS) +def disagg_shares_draft_kv( + speculative_config: Optional[DecodingBaseConfig]) -> bool: + """Whether a disaggregated worker transfers its drafter's KV with the target's. + + TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires the draft KV to share the target's + KV cache manager, and 0 keeps it private. Unset, a DFlash drafter on a + backend that can share it (AUTO or TRTLLM) does, so a context and a + generation server given the same speculative_config agree on their own; + the transceiver rejects a pair that does not. + """ + env = os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") + if env: + return env == "1" + return (isinstance(speculative_config, DFlashDecodingConfig) + and speculative_config.attention_backend in ("AUTO", "TRTLLM")) + + def disagg_dflash_error( speculative_config: Optional[DecodingBaseConfig], cache_transceiver_config: Optional[CacheTransceiverConfig] @@ -7064,14 +7084,34 @@ def validate_speculative_config(self): if isinstance(self.speculative_config, DFlashDecodingConfig): disagg = (self.cache_transceiver_config is not None and self.cache_transceiver_config.backend is not None) - # TRTLLM's separate draft pool collapses on a disaggregated worker; modeling_v2 drafts on VANILLA. - if self.speculative_config.attention_backend == "AUTO" and ( - disagg or self.modeling_v2 != "off"): + shared_draft_kv = disagg_shares_draft_kv( + self.speculative_config) + # An unshared disaggregated worker refuses a windowed TRTLLM draft pool; modeling_v2 drafts on VANILLA. + if (self.speculative_config.attention_backend == "AUTO" + and ((disagg and not shared_draft_kv) + or self.modeling_v2 != "off")): self.speculative_config.attention_backend = "VANILLA" logger.info( "DFlash attention_backend AUTO resolved to VANILLA " - "(disaggregated worker or modeling_v2).") - if disagg: + "(modeling_v2, or a disaggregated worker with " + "TRTLLM_DISAGG_SHARED_DRAFT_KV=0).") + assert self.speculative_config.max_draft_len > 0, "DFlash max_draft_len must be > 0" + # From the disaggregated-serving-benchmarking overlay (TRT-LLM aa7db814): + # the disagg_dflash_error refusal is dropped so gen-only SOL + # ladders can run DFlash on a disaggregated gen. Acceptance is + # forced (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS), so the + # collapsed acceptance the refusal guards against does not + # reach the measurement; the drafter still runs every step. + # Only valid as a throughput measurement, never for quality. + # A Hugging Face repo id is not readable yet: both calls below + # then run without the drafter's config.json (the budget check + # covers only the token budget), and CachedModelLoader repeats + # both after the drafter is downloaded. + self.speculative_config.resolve_from_checkpoint() + # A draft KV shared with the target is transferred with it; + # checked once AUTO is resolved from the checkpoint. + if disagg and not (self.speculative_config.attention_backend in + ("TRTLLM", "AUTO") and shared_draft_kv): # The transceiver moves the target KV cache, but the # drafter's context is built from target hidden states # during prefill and is not transferred with it, so a @@ -7086,20 +7126,8 @@ def validate_speculative_config(self): "prefill and is not transferred, so a generation " "server drafts without the prompt context. Output is " "unaffected; expect a lower acceptance rate than the " - "same configuration run aggregated.") - assert self.speculative_config.max_draft_len > 0, "DFlash max_draft_len must be > 0" - # From the disaggregated-serving-benchmarking overlay (TRT-LLM aa7db814): - # the disagg_dflash_error refusal is dropped so gen-only SOL - # ladders can run DFlash on a disaggregated gen. Acceptance is - # forced (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS), so the - # collapsed acceptance the refusal guards against does not - # reach the measurement; the drafter still runs every step. - # Only valid as a throughput measurement, never for quality. - # A Hugging Face repo id is not readable yet: both calls below - # then run without the drafter's config.json (the budget check - # covers only the token budget), and CachedModelLoader repeats - # both after the drafter is downloaded. - self.speculative_config.resolve_from_checkpoint() + "same configuration run aggregated. To transfer it, " + "run both servers with attention_backend TRTLLM.") self._validate_dflash_ctx_budget() if isinstance(self.speculative_config, DSparkDecodingConfig): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index faaeb3bcfdbe..fc7131bb80d7 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -224,6 +224,7 @@ class BufferConfig: size: int tokens_per_block_override: int | None = None is_sparse: bool = False + coalescing_group: int = 0 @dataclass(slots=True) class AttentionLayerConfig: diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index bc11cdb66db8..e41e0346a7a6 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 b1d3a3ed3a6f..63fd1283c9e4 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 @@ -1663,6 +1664,97 @@ def test_gen_first_kv_cache_v1(self): 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, + # "auto" resolves Qwen3 to V1; the shared draft KV needs V2. + "use_kv_cache_manager_v2": True, + }, + # 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", + 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, + ) + 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_cpu_musespark.yml b/tests/integration/test_lists/test-db/l0_cpu_musespark.yml index aa0432bfffad..fc89e5845f60 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_musespark.yml @@ -18,3 +18,13 @@ l0_cpu_musespark: - unittest/llmapi/apps/test_musespark_serve_extensions.py # Picked up on l0_cpu through the unittest/_torch/modeling directory entry. - unittest/_torch/modeling/test_modeling_musespark_open_preprocess.py + # The cpu_only tests of DFlash on the TRTLLM attention backend, the shared target/draft KV + # cache and the disaggregated KV transfer; l0_gb300_musespark.yml runs the rest. + - unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py + - unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py + - unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py + - unittest/llmapi/test_disagg_dflash_guard.py + - unittest/disaggregated/test_kv_transfer.py + - unittest/disaggregated/test_peer.py + - unittest/disaggregated/test_rank_info.py + - unittest/disaggregated/test_transfer_ownership_regressions.py 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 e075a2ee8153..23efa45be82e 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -20,6 +20,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 # ------------- 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/integration/test_lists/test-db/l0_gb300_musespark.yml b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml index da26c868fdb1..d71a45b69db4 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml @@ -38,3 +38,15 @@ l0_gb300_musespark: - unittest/_torch/modeling/test_musespark_open_moe_prepare_overlap.py - unittest/_torch/modeling/test_musespark_open_fusion_switches.py - unittest/_torch/modeling/test_musespark_open_finalize_local.py + # DFlash on the TRTLLM attention backend, the shared target/draft KV cache, and the + # disaggregated KV transfer that carries it. Pre-merge CI on this branch runs only the + # MuseSpark stages, and these files are listed on no other MuseSpark stage. This stage + # deselects cpu_only tests; those files are on l0_cpu_musespark.yml. + - unittest/_torch/speculative/hw_agnostic/test_dflash_trtllm_window.py + - unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py + - unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py + - unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py + - unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.py + - unittest/disaggregated/test_kv_transfer.py + - unittest/disaggregated/test_bounce.py + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv 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_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py new file mode 100644 index 000000000000..655755ac4cc5 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -0,0 +1,421 @@ +# 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.""" + +import json +import os +from types import SimpleNamespace + +import click +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 +from tensorrt_llm.llmapi import DFlashDecodingConfig +from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.quantization.mode import QuantMode + +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( + 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, draft_windows=None): + creator = object.__new__(KvCacheCreator) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace( + is_dflash=lambda: True, is_dspark=lambda: False, is_external_drafter=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, + 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 + ) + 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 + + +@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() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_mismatched_draft_kv_heads_unify(is_disagg: bool) -> None: + # Each transfer view carries its own KV-head count, so the draft's heads are resharded on their own. + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +def test_nvfp4_with_mismatched_draft_kv_heads_unifies_only_in_aggregated_serving() -> None: + # NVFP4 block-scale views are not resharded, so an NVFP4 side still needs equal heads. + target = _model_config(kv_heads=8, head_dim=128, kv_quant_algo=QuantAlgo.NVFP4) + 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. + # Per-view head counts let the two still transfer together. + 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 creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_attention_dp_keeps_every_head_on_each_rank(is_disagg, attention_dp): + # At TP4 a 4-head target and a 1-head draft both hold one head per rank. + # Under attention DP each rank keeps every head (4 and 1); the layouts differ + # but still transfer together, since each view carries its own head count. + 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 = attention_dp + + heads = ( + creator._dflash_kv_layout(target, torch.bfloat16)[0], + creator._dflash_kv_layout(draft, torch.bfloat16)[0], + ) + assert heads == ((4, 1) if attention_dp else (1, 1)) + 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) == () + + +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) + 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): + generation_only = dict( + 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, **generation_only) + # Unset, it keeps the draft KV private on both servers. + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + with pytest.raises(_StopLauncher): + _launch_from_disagg_config(monkeypatch, tmp_path, **generation_only) + assert os.environ[_SHARED_DRAFT_KV_ENV] == "0" + 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) + creator._speculative_config = DFlashDecodingConfig( + max_draft_len=7, attention_backend="TRTLLM" + ) + assert not creator._should_unify_dflash_kv_cache() + + +def test_launcher_keeps_an_explicit_private_draft_kv(monkeypatch, tmp_path): + spec = {"decoding_type": "DFlash"} + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, "0") + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={"speculative_config": spec}, + generation_servers={"speculative_config": spec}, + ) + assert os.environ[_SHARED_DRAFT_KV_ENV] == "0" + + +def test_dflash_with_different_tp_shares_the_cache(monkeypatch, tmp_path): + # Context and generation may differ in TP: heads are resharded per view. + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + with pytest.raises(_StopLauncher): + _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 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", + [ + ( + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + "attention_backend AUTO or TRTLLM", + ), + ( + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "FA4"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}}, + "attention_backend AUTO or TRTLLM", + ), + ], + ids=["vanilla_dflash_backend", "fa4_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) + + +@pytest.mark.parametrize("attention_backend", [None, "AUTO", "TRTLLM"]) +def test_launcher_shares_draft_kv_for_dflash_in_both_sections( + monkeypatch, tmp_path, attention_backend +): + """AUTO, the default, resolves to TRTLLM for windowed drafters, so it shares like TRTLLM.""" + spec = {"decoding_type": "DFlash"} + if attention_backend is not None: + spec["attention_backend"] = attention_backend + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, tmp_path, {"speculative_config": spec}, {"speculative_config": spec} + ) + assert os.environ[_SHARED_DRAFT_KV_ENV] == "1" + + +@pytest.mark.parametrize( + "context_attention_dp,generation_attention_dp", [(True, True), (True, False)] +) +def test_launcher_accepts_eligible_shared_cache_config( + monkeypatch, tmp_path, context_attention_dp, generation_attention_dp +): + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + {"speculative_config": spec, "enable_attention_dp": context_attention_dp}, + {"speculative_config": spec, "enable_attention_dp": generation_attention_dp}, + ) + + +@pytest.mark.parametrize( + ("shared", "windows", "refused"), + [("0", [2048], True), ("0", None, False), ("1", [2048], False)], +) +def test_disagg_separate_windowed_draft_pool_is_refused(monkeypatch, shared, windows, refused): + """Unshared, a transferred request's drafter would read blocks the pool retired behind the window.""" + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, shared) + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=windows) + creator._speculative_config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + creator._sparse_attention_config = None + if refused: + with pytest.raises(ValueError, match="TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on every worker"): + creator._check_shared_draft_kv_request() + else: + creator._check_shared_draft_kv_request() + + +def test_disagg_dflash_shares_its_draft_kv_without_the_variable(monkeypatch): + """Separately launched servers with the same DFlash config share without a launcher.""" + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=[2048]) + creator._speculative_config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + creator._sparse_attention_config = None + + creator._check_shared_draft_kv_request() + assert creator._should_unify_dflash_kv_cache() + + +def test_disagg_windowed_drafter_that_cannot_share_is_refused(monkeypatch): + """The refusal names what keeps a TRTLLM drafter's window from sharing the target's manager.""" + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=[2048]) + creator._speculative_config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + creator._sparse_attention_config = None + creator._mapping.pp_size = 2 + + with pytest.raises( + ValueError, match="pipeline_parallel_size is 2, not 1: use attention_backend" + ): + creator._check_shared_draft_kv_request() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_rank_pins_auto_to_unify_like_trtllm(tmp_path, monkeypatch, is_disagg): + """An AUTO the frontend could not probe is pinned on the rank, so it unifies like TRTLLM.""" + from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _pin_dflash_auto_attention_backend, + ) + from tensorrt_llm._torch.speculative import dflash + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: None) + windowed = {"layer_types": ["sliding_attention"] * 2, "sliding_window": 16} + (tmp_path / "config.json").write_text(json.dumps(windowed)) + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout, draft_windows=[16]) + creator._speculative_config = DFlashDecodingConfig( + max_draft_len=7, speculative_model=str(tmp_path) + ) + _pin_dflash_auto_attention_backend(creator._speculative_config) + assert creator._speculative_config.attention_backend == "TRTLLM" + assert creator._unified_drafter_attention_backend() == "TRTLLM" + creator._check_shared_draft_kv_request() # TRTLLM_DISAGG_SHARED_DRAFT_KV=1 (autouse) + assert creator._should_unify_dflash_kv_cache() 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 fa7915d27cc8..ecc9f82976ed 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 @@ -319,6 +319,10 @@ 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() or is_dspark() would route costing + # through the shared target and drafter cache path. + mode.is_dflash.return_value = False + mode.is_dspark.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 2d698874e9a8..5cd6a0e3f92d 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 @@ -27,12 +27,19 @@ _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 -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.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, @@ -953,8 +960,76 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): assert cost == CacheCost(slope=20, intercept=21) +@pytest.mark.skip( + reason="[agentx port] needs main's per-layer AttentionReusePolicy (#19727), " + "which the agentx KV bindings lack" +) +def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + + 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 manager._has_cp_helix = False manager.num_local_layers = 3 manager.pp_layers = [4, 5, 6] @@ -1572,3 +1647,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_kv_cache_v2_attention_op_tables.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.py new file mode 100644 index 000000000000..cb2a62943900 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.py @@ -0,0 +1,284 @@ +# 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. +"""Attention-op page tables of a KV cache whose life cycles mix page sizes. + +A unified target/draft cache (FP8 target, BF16 draft sharing a sliding window) +cannot use one table per life cycle. The batch's rows of its tables must match +the old full-capacity per-layer rebuild: bit for bit with per-layer tables, and +page for page (same addresses, same BAD_PAGE_INDEX entries) with shared tables. +They must also be the host rows at enqueue time, however the host table changes +before the asynchronous copy runs. +""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.utils as torch_utils +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, + _copy_swa_block_offsets_with_scratch_compiled, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, PageIndexMode + +TOKENS_PER_BLOCK = 4 +WINDOW = 2 * TOKENS_PER_BLOCK +MAX_SEQ_LEN = 64 +SENTINEL = -7 +# (request id, capacity, history length, dummy): one request far past the +# window, one inside it, and a CUDA-graph padding dummy. +REQUESTS = ((11, 40, 36, False), (12, 8, 0, False), (13, 4, 0, True)) +UNSCHEDULED = (14, 24, 20, False) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _layer_specs(): + # Muse Spark-like: FP8 target layers (windowed and full) plus BF16 draft + # layers with the target's window and twice its page size. + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=2, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=WINDOW, + ) + specs = [ + replace(target, global_layer_id=i, sliding_window_size=None if i % 3 == 2 else WINDOW) + for i in range(6) + ] + draft = replace(target, owner=KVCacheLayerOwner.DRAFT, dtype=DataType.BF16) + return specs + [replace(draft, global_layer_id=6 + i) for i in range(2)] + + +def _make_manager(monkeypatch, grouped: bool) -> KVCacheManagerV2: + torch.cuda.init() + if not grouped: + # raising=False: also runs against a tree without shared tables. + monkeypatch.setattr( + KVCacheManagerV2, "_get_attention_op_pool_groups", lambda self: None, raising=False + ) + 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=6, + num_kv_heads=2, + head_dim=16, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + layer_specs=_layer_specs(), + ) + for request_id, capacity, history, dummy in REQUESTS + (UNSCHEDULED,): + cache = manager._create_kv_cache(request_id, None, None, is_dummy=dummy) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(capacity, history) + return manager + + +def _full_rebuild(manager: KVCacheManagerV2, copy_idx: torch.Tensor) -> torch.Tensor: + """The pre-change per-layer path: convert every IndexMapper row, keep the batch's.""" + params = [] + for layer_id in range(manager.num_local_layers): + pool_id = manager.impl.get_layer_group_id(layer_id) + converters = [ + manager.impl.get_page_index_converter(layer_id, role) for role in (Role.KEY, Role.VALUE) + ] + params.append( + [[pool_id] * 2, [c.scale for c in converters], [c.layer_offset for c in converters]] + ) + params = torch.tensor(params, dtype=torch.int32, device="cuda") # [layers, 3, 2] + host = manager.host_kv_cache_block_offsets + num_pools, capacity, _, num_blocks = host.shape + padded_idx = torch.zeros(capacity, dtype=torch.long, device="cuda") + padded_idx[: copy_idx.numel()] = copy_idx.to(device="cuda", dtype=torch.long) + no_scratch = torch.zeros(num_pools, capacity, dtype=torch.int32, device="cuda") + out = torch.empty(params.shape[0], capacity, 2, num_blocks, dtype=torch.int32, device="cuda") + _copy_swa_block_offsets_with_scratch_compiled( + host.cuda(), + padded_idx, + params[:, 0].long(), + params[:, 1], + params[:, 2], + torch.ones_like(params[:, 1]), + torch.arange(num_blocks, dtype=torch.int32, device="cuda"), + no_scratch, + no_scratch, + no_scratch[:, :, None], + torch.zeros((), dtype=torch.int32, device="cuda"), + out, + ) + return out[:, : copy_idx.numel()] + + +def _copy_batch(manager: KVCacheManagerV2): + request_ids = [request_id for request_id, *_ in REQUESTS] + num_seqs = len(request_ids) + dst = torch.full( + (manager.num_attention_op_pools, num_seqs + 2, 2, manager.max_blocks_per_seq), + SENTINEL, + dtype=torch.int32, + device="cuda", + ) + manager.copy_batch_block_offsets(dst, request_ids, 1, 0, num_seqs) + copy_idx = manager.index_mapper.get_copy_index(request_ids, 0, 1, False) + torch.cuda.synchronize() + # The other rows of the destination are not the batch's and stay untouched. + assert (dst[:, num_seqs:] == SENTINEL).all() + return dst[:, :num_seqs], copy_idx + + +def _check_windowed_out_pages(manager: KVCacheManagerV2, copy_idx: torch.Tensor) -> None: + """The long request's windowed rows have BAD_PAGE_INDEX before their live tail.""" + window_pool = manager.impl.get_layer_group_id(0) + row = manager.host_kv_cache_block_offsets[window_pool, int(copy_idx[0]), 0].tolist() + last_valid = max(i for i, page in enumerate(row) if page != BAD_PAGE_INDEX) + assert BAD_PAGE_INDEX in row[:last_valid] + + +@pytest.fixture(autouse=True) +def eager_full_rebuild(monkeypatch): + # Run the reference rebuild uncompiled (maybe_compile honors this switch). + monkeypatch.setattr(torch_utils, "is_piecewise_running", lambda: True) + + +def test_per_layer_tables_match_full_rebuild(monkeypatch) -> None: + manager = _make_manager(monkeypatch, grouped=False) + try: + assert manager._use_per_layer_page_tables + assert manager.num_attention_op_pools == manager.num_local_layers + # No full-capacity staging without SWA scratch slots. + assert not hasattr(manager, "_device_attention_op_block_offsets_staging") + tables, copy_idx = _copy_batch(manager) + _check_windowed_out_pages(manager, copy_idx) + assert torch.equal(tables, _full_rebuild(manager, copy_idx)) + finally: + manager.shutdown() + + +def test_grouped_tables_address_the_full_rebuild_pages(monkeypatch) -> None: + manager = _make_manager(monkeypatch, grouped=True) + try: + assert not manager._use_per_layer_page_tables + # (window, FP8 pages), (window, BF16 pages), (full, FP8 pages). + assert manager.num_attention_op_pools == 3 + tables, copy_idx = _copy_batch(manager) + _check_windowed_out_pages(manager, copy_idx) + reference = _full_rebuild(manager, copy_idx).cpu() + tables = tables.cpu() + pointers = manager.kv_cache_pool_pointers[:, 0].tolist() + for layer_id in range(manager.num_local_layers): + op_pool, offset = manager.kv_cache_pool_mapping[layer_id].tolist() + stride = manager.impl.get_page_stride(layer_id, Role.KEY) + # The attention op's pool pointer plus the mapping offset is the layer's own base. + layer_base = pointers[op_pool] + offset * 2 * stride + assert layer_base == manager.impl.get_mem_pool_base_address( + layer_id, Role.KEY, PageIndexMode.SHARED + ) + origin = manager.impl.get_mem_pool_base_address( + layer_id, Role.KEY, PageIndexMode.PER_LAYER + ) + expected = reference[layer_id] + actual = tables[op_pool] + bad = expected == BAD_PAGE_INDEX + assert torch.equal(actual == BAD_PAGE_INDEX, bad) + assert torch.equal( + (layer_base + actual.long() * stride)[~bad], + (origin + expected.long() * stride)[~bad], + ) + finally: + manager.shutdown() + + +@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per-layer"]) +def test_dflash_decodes_the_same_draft_tables(monkeypatch, grouped: bool) -> None: + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + manager = _make_manager(monkeypatch, grouped=grouped) + try: + num_draft_layers, nkv, head_dim = 2, 2, 16 + num_seqs = len(REQUESTS) + tables, copy_idx = _copy_batch(manager) + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + pool = worker._managed_ctx_pool(manager, num_draft_layers, nkv, head_dim, torch.bfloat16, 2) + assert pool is not None + assert worker._init_ctx_block_tables( + manager, pool, num_seqs, num_draft_layers, nkv, head_dim, TOKENS_PER_BLOCK + ) + assert worker._refresh_ctx_block_tables( + SimpleNamespace(kv_cache_block_offsets=tables), num_seqs + ) + # Decode the pre-change per-layer tables with the pre-change constants. + draft_layer = manager.layer_offsets[worker._ctx_layer_ids[0]] + reference = object.__new__(DFlashWorker) + reference._ctx_unified_cache = True + reference._ctx_pool_idx = draft_layer + reference._ctx_block_offset = manager.impl.get_page_index_converter( + draft_layer, Role.KEY + ).layer_offset + reference._ctx_block_divisor = worker._ctx_block_divisor + reference._ctx_block_tables = torch.zeros_like(worker._ctx_block_tables) + reference._ctx_block_counts = torch.zeros_like(worker._ctx_block_counts) + reference._ctx_block_cols = worker._ctx_block_cols + assert reference._refresh_ctx_block_tables( + SimpleNamespace(kv_cache_block_offsets=_full_rebuild(manager, copy_idx)), num_seqs + ) + assert torch.equal(worker._ctx_block_tables, reference._ctx_block_tables) + assert torch.equal(worker._ctx_block_counts, reference._ctx_block_counts) + # The long request's extent runs past its windowed-out pages. + assert int(worker._ctx_block_counts[0]) * TOKENS_PER_BLOCK >= 36 + finally: + manager.shutdown() + + +@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "per-layer"]) +def test_copy_is_a_snapshot_of_the_host_table(monkeypatch, grouped: bool) -> None: + """The next step's host-table writes must not reach this step's pending copy.""" + manager = _make_manager(monkeypatch, grouped=grouped) + try: + expected, _ = _copy_batch(manager) + request_ids = [request_id for request_id, *_ in REQUESTS] + dst = torch.full_like(expected, SENTINEL) + host = manager.host_kv_cache_block_offsets + saved = host.clone() + torch.cuda.synchronize() + # Hold the stream so the copy is still pending when the host table changes. + torch.cuda._sleep(1 << 30) + manager.copy_batch_block_offsets(dst, request_ids, 1, 0, len(request_ids)) + host.fill_(7) + torch.cuda.synchronize() + host.copy_(saved) + assert torch.equal(dst, expected) + finally: + manager.shutdown() 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 9ec9df759125..dcc0110b49a3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -24,13 +24,18 @@ 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 import DecoderRunner 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, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _KVCacheManagerInitStatus, _swa_endpoint_priority, _sync_kv_cache_manager_init_status, @@ -41,6 +46,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 ( @@ -52,6 +58,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, @@ -70,6 +77,7 @@ KVCacheManager, KVCacheManagerConfig, LayerId, + PageIndexMode, ReuseScope, SsmLayerConfig, _introspection, @@ -309,6 +317,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 @@ -353,6 +362,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: @@ -416,7 +427,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, @@ -627,6 +639,316 @@ 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 = [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 ( + (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() + + +@pytest.mark.parametrize("grouped", [True, False], ids=["shared-table", "per-layer"]) +def test_dflash_block_tables_with_fp8_target_and_bf16_draft( + monkeypatch: pytest.MonkeyPatch, grouped: bool +) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + if not grouped: + monkeypatch.setattr(KVCacheManagerV2, "_get_attention_op_pool_groups", lambda self: None) + 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: + # Same-size pages share one physical pool: one shared table, or one per layer. + assert manager._use_per_layer_page_tables != grouped + assert manager.num_attention_op_pools == (1 if grouped else 2) + 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 + ) + # Per-layer tables carry the draft K buffer's offset; a shared one maps it. + assert (worker._ctx_block_offset > 0) != grouped + + 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") + 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_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] @@ -1296,6 +1618,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), @@ -2902,6 +3303,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} @@ -3197,6 +3599,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} @@ -3378,3 +3781,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/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 79be6a9e3744..833abc7e8682 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 @@ -1946,6 +1946,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 @@ -1977,6 +1978,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 @@ -2205,6 +2207,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/_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..5c3e9eb7a2c7 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -0,0 +1,125 @@ +# 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) + + +def _refresh_worker(*, num_slots, width, divisor): + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + worker._ctx_pool_idx = 0 + worker._ctx_block_offset = 0 + worker._ctx_block_divisor = divisor + worker._ctx_block_tables = torch.zeros(num_slots, width, dtype=torch.int32) + worker._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long) + worker._ctx_block_cols = torch.arange(1, width + 1, dtype=torch.long) + return worker + + +def test_block_count_is_the_allocation_extent_past_a_sliding_window(): + # Encoded offsets [pool, seq, K/V field, block]. Row 0's first three blocks + # were evicted by a sliding window ahead of its live tail; row 1 is a short + # allocation; row 2 holds nothing. + rows = [[-1, -1, -1, 8, 10, -1], [0, 2, -1, -1, -1, -1], [-1] * 6] + offsets = torch.tensor(rows, dtype=torch.int32)[None, :, None, :].expand(1, 3, 2, 6) + worker = _refresh_worker(num_slots=4, width=6, divisor=2) + + assert worker._refresh_ctx_block_tables(SimpleNamespace(kv_cache_block_offsets=offsets), 3) + + # A count of valid entries would cap row 0 at its two live pages. + assert worker._ctx_block_counts[:3].tolist() == [5, 2, 0] + assert worker._ctx_block_tables[0].tolist() == [0, 0, 0, 4, 5, 0] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 5cca7cb3d1ea..774ed238af78 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -56,17 +56,27 @@ def test_default_dflash_resolves_auto_before_choosing_its_draft_cache( @pytest.mark.parametrize( - ("transceiver", "modeling_v2", "resolved"), - [("NIXL", "off", "VANILLA"), (None, "auto", "VANILLA"), (None, "off", "AUTO")], + ("transceiver", "shared", "modeling_v2", "resolved"), + [ + ("NIXL", "0", "off", "VANILLA"), + ("NIXL", "1", "off", "AUTO"), + ("NIXL", None, "off", "AUTO"), + (None, "0", "auto", "VANILLA"), + (None, "0", "off", "AUTO"), + ], ) -def test_auto_keeps_vanilla_where_trtllm_would_collapse_or_be_bypassed( - tmp_path, monkeypatch, transceiver, modeling_v2, resolved +def test_auto_keeps_vanilla_where_trtllm_is_refused_or_bypassed( + tmp_path, monkeypatch, transceiver, shared, modeling_v2, resolved ) -> None: - """A disaggregated worker gets no draft KV; modeling_v2 DFlash compositions need VANILLA.""" + """Workers keeping a private draft KV refuse TRTLLM; modeling_v2 DFlash compositions need VANILLA.""" from tensorrt_llm.llmapi import CacheTransceiverConfig, TorchLlmArgs, llm_args # No GPU: whatever stays AUTO is left to the ranks' windowed-and-probe rule. monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: False) + if shared is None: + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", shared) (tmp_path / "config.json").write_text(json.dumps(_WINDOWED)) args = TorchLlmArgs( model="dummy", diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py new file mode 100644 index 000000000000..fe29e3395073 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -0,0 +1,104 @@ +# 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's dummy slot (CUDA-graph padding, warmup and attention-DP dummies) is at context length 0 between steps. + +Adapted from main's test of the same name (#19988). A dummy length that kept growing would put the padding rows' +context K / V past the padding request's own pages, and on a unified KV cache trip the drafter's allocation check. +""" + +from functools import partial +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dflash import DFlashWorker + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the slot lengths live on the GPU" +) + + +def _worker(lengths, bases=None): + w = SimpleNamespace( + _dummy_slot=len(lengths) - 1, + _max_ctx=1000, + _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), + _ctx_base=None if bases is None else torch.tensor(bases, dtype=torch.long, device="cuda"), + ) + w._advance_ctx_len = partial(DFlashWorker._advance_ctx_len, w) + return w + + +def _cuda(values): + return torch.tensor(values, dtype=torch.long, device="cuda") + + +def test_padded_steps_keep_the_dummy_slot_empty(): + worker = _worker([40, 300, 7, 0, 0]) + slots = _cuda([0, 1, 2, 4, 4]) # three requests in a graph of five + for step in range(1, 5): + worker._advance_ctx_len(slots, _cuda([8] * 5)) + torch.cuda.synchronize() + assert worker._ctx_len.tolist() == [40 + 8 * step, 300 + 8 * step, 7 + 8 * step, 0, 0], step + + +def test_dummy_slot_reset_replays_in_a_cuda_graph(): + worker = _worker([40, 300, 7, 0, 0]) + slots = _cuda([0, 1, 4]) + accepted = _cuda([2, 3, 5]) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + worker._advance_ctx_len(slots, accepted) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._advance_ctx_len(slots, accepted) + worker._ctx_len.copy_(_cuda([40, 300, 7, 0, 0])) + for step in range(1, 4): + worker._ctx_len[worker._dummy_slot] = 99 # whatever the slot holds before the step + graph.replay() + torch.cuda.synchronize() + assert worker._ctx_len.tolist() == [40 + 2 * step, 300 + 3 * step, 7, 0, 0], step + + +def test_windowed_arena_dummy_slot_is_empty_with_its_base(): + worker = _worker([40, 300, 0], bases=[8, 100, 64]) + worker._advance_ctx_len(_cuda([0, 2, 2]), _cuda([4, 4, 4])) + torch.cuda.synchronize() + assert worker._ctx_len.tolist() == [44, 300, 0] + assert worker._ctx_base.tolist() == [8, 100, 0] + + +def test_warmup_reseeds_generation_slots_from_this_batch(): + # Slot 0 still holds an earlier warmup batch's length for a reused id; the + # padding row (no slot of its own) is left alone. + worker = _worker([500, 7, 0, 0]) + worker._req_to_slot = {11: 0, 12: 1} + worker._ctx_len_host = [500, 7, 0, 0] + worker._write_ctx_len = partial(DFlashWorker._write_ctx_len, worker) + meta = SimpleNamespace( + request_ids=[11, 12, 1 << 62], + num_generations=3, + seq_lens=[8, 8, 8], + host_position_ids=[20] * 8 + [30] * 8 + [0] * 8, + ) + + DFlashWorker._reseed_warmup_generation_slots(worker, meta) + torch.cuda.synchronize() + + assert worker._ctx_len.tolist() == [20, 30, 0, 0] + assert worker._ctx_len_host == [20, 30, 0, 0] diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 70b761b1326f..07f26d9c0355 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -238,8 +238,13 @@ def test_fanin_bounce_safe_gate() -> None: safe = tfr.Receiver._fanin_bounce_safe - def ov(dup: int, pp: int, ranks: tuple[int, ...] = (0,)) -> SimpleNamespace: - return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp, ranks=list(ranks)) + def ov(dup: int, pp: int, ranks: tuple[int, ...] = (0,), peer_dup: int = 1) -> SimpleNamespace: + return SimpleNamespace( + duplicate_head_factor=dup, + peer_duplicate_head_factor=peer_dup, + overlap_pp_size=pp, + ranks=list(ranks), + ) def ri(lpp: list[int], page_table: SimpleNamespace | None = None) -> SimpleNamespace: return SimpleNamespace(layer_num_per_pp=lpp, page_table=page_table) @@ -252,6 +257,10 @@ def pt(mapper_kind: MapperKind) -> SimpleNamespace: assert safe(ov(1, 1), ri([24]), None) is True assert safe(ov(1, 0), ri([24]), None) is True assert safe(ov(2, 1), ri([24]), None) is False # duplicate heads / MLA -> some don't send + # The senders' duplication counts too: with 4 KV heads, ctx TEP8 -> gen TEP2 has 4 writers + # but elects 2 of them; MLA ctx TP8 -> gen TP4 elects 1 of 2. The receiver's own factor is 1. + assert safe(ov(1, 1, ranks=(0, 1, 2, 3), peer_dup=2), ri([24]), None) is False + assert safe(ov(1, 1, ranks=(0, 1), peer_dup=2), ri([24]), None) is False # EVEN PP fan-in (equal layers per overlapping stage) -> allowed assert safe(ov(1, 4), ri([20, 20, 20, 20]), None) is True # UNEVEN PP fan-in -> per-writer sizes differ -> fall back @@ -446,6 +455,17 @@ def test_reserve_heterogeneous_fanin_falls_back(self, monkeypatch): assert t.reserve(req, num_writers=2) is False assert req.bounce_dst_base is None + def test_reserve_heterogeneous_tp_fanin_ok(self, monkeypatch): + # TP fan-in: each writer sends a head slice of every block of every group, so its bytes are + # total // num_writers for any per-group sizes (here a full-attention group of 16 layers and + # a sliding-window group of 46 layers, 64 KiB per layer and block, gathered by 4 writers). + t = _make_transport(monkeypatch, block_bytes_per_group=[16 << 16, 46 << 16]) + req = _recv_req([3, 2]) # total = 3*16 + 2*46 = 140 layer-blocks of 64 KiB + assert t.reserve(req, num_writers=4, writers_split_layers=True) is False + assert t.reserve(req, num_writers=4, writers_split_layers=False) is True + rid_slice = (req.unique_rid, req.slice_id) + assert t.writer_base(rid_slice, 3) == 0x100000 + 3 * ((140 << 16) // 4) + def test_reserve_uniform_multigroup_fanin_ok(self, monkeypatch): # Uniform slot bytes across present groups -> even byte split -> bounce allowed. t = _make_transport(monkeypatch, block_bytes_per_group=[100, 100]) @@ -550,6 +570,45 @@ def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): assert [list(t[1]) for t in descs] == [[10], [20]] # dst_ptrs assert [list(t[2]) for t in descs] == [[8], [8]] # sizes + def test_fanin_scatter_is_one_batched_copy(self, monkeypatch): + # The batched copy reuses one pinned metadata buffer per stream and needs a sync between + # calls; one call per writer before a single sync let a writer's kernel run with the next + # writer's fragment table. All writers must go into one call, each from its own source. + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, num_writers=2) is True + rid_slice = (req.unique_rid, req.slice_id) + for rank, dst, src_base in ((7, 20, 200), (3, 10, 100)): + t.record_result( + rid_slice, + rank, + np.array([dst], dtype=np.int64), + np.array([8], dtype=np.int64), + src_base=src_base, + ) + calls = [] + + def fake_scatter(src_base, dst_ptrs, sizes, offsets, *, stream): + calls.append(("scatter", src_base, list(dst_ptrs), list(sizes), list(offsets))) + + monkeypatch.setattr(btr, "scatter_contiguous", fake_scatter) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace( + cudaSetDevice=lambda dev: (0,), + cudaStreamSynchronize=lambda stream: calls.append(("sync", stream)) or (0,), + ), + ) + monkeypatch.setattr(btr, "CUASSERT", lambda ret: ret) + t._scatter_stream = 0 + t._stop = SimpleNamespace(is_set=lambda: False) + t._scatter_q.put(None) # poison pill behind the queued region + t._scatter_loop() + # one copy: both writers' destinations, each read from its own src_base (+ offset 0) + assert calls == [("scatter", 0, [10, 20], [8, 8], [100, 200]), ("sync", 0)] + assert t._recv_alloc.released # region settled after the scatter + def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch): # Regression: if one fan-in writer falls back to in-place (SUCCESS, empty tail) while a # sibling bounces, the survivor must be scattered from ITS OWN src_base, not packed to 0. diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 92b7a836032f..08c9405f600d 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -2043,6 +2043,58 @@ def raiser(*args, **kwargs): worker.shutdown() +@pytest.mark.timeout(120) +@pytest.mark.parametrize(("ctx_draft_layers", "gen_draft_layers"), [(2, 0), (0, 2)]) +def test_draft_layer_mismatch_fails_requests_at_first_transfer(ctx_draft_layers, gen_draft_layers): + """Peers that disagree on the draft layers riding the target KV fail each request at once. + + Otherwise the sender only logs its rejection and the request waits out + the transfer timeout, or the generation drafter runs without its prompt KV. + """ + setup = create_transfer_worker_setup( + ctx_tp=1, ctx_pp=1, ctx_enable_dp=False, gen_tp=1, gen_pp=1, gen_enable_dp=False + ) + ctx_tw = setup["ctx_transfer_workers"][0] + gen_tw = setup["gen_transfer_workers"][0] + # The context worker's rank-info server publishes this object. + ctx_tw._rank_info.num_draft_layers = ctx_draft_layers + gen_tw._rank_info.num_draft_layers = gen_draft_layers + sc = tensorrt_llm.bindings.SamplingConfig(SamplingParams(temperature=0)._get_sampling_config()) + request = LlmRequest( + request_id=700, + max_new_tokens=1, + input_tokens=list(range(16)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + request.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=700, + ctx_dp_rank=0, + ctx_info_endpoint=setup["ctx_info_endpoint"], + disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF, + ) + page_table = gen_tw._rank_info.page_table + empty_slice = Chunk( + block_ids_per_layer_groups=[np.array([], dtype=np.int64) for _ in page_table.layer_groups], + kind_per_layer_group=[CacheKind.PAGED], + token_range=TokenRange(start=0, end=0), + is_last=True, + ) + try: + rx = gen_tw.create_rx_session(request) + rx.receive(empty_slice) + assert rx.wait_complete(blocking=True) == WaitResult.FAILED + message = str(rx._kv_tasks[0]._exception) + assert f"local={gen_draft_layers}, peer={ctx_draft_layers}" in message + assert "same speculative_config" in message + assert setup["ctx_info_endpoint"] in gen_tw._receiver._incompatible_peers + rx.close() + finally: + for worker in setup["ctx_transfer_workers"] + setup["gen_transfer_workers"]: + worker.shutdown() + + def add_and_verify_pipelined_request( setup, ctx_request_id, 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], diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index dd55120214c0..0634969e8336 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -1364,3 +1364,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 de11fe01f992..73949bdabfcd 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 @@ -194,3 +195,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 diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..9ccb2ace7f59 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -66,8 +66,15 @@ def __init__(self) -> None: def record_failure(self, rid_slice: tuple[int, int], peer_rank: int) -> None: self.failed_writers.append((rid_slice, peer_rank)) - def reserve(self, _receiver_req, _num_writers: int, *, extra_bytes: int = 0) -> bool: - del extra_bytes + def reserve( + self, + _receiver_req, + _num_writers: int, + *, + extra_bytes: int = 0, + writers_split_layers: bool = True, + ) -> bool: + del extra_bytes, writers_split_layers return False def release_idle_reservation(self, _rid_slice: tuple[int, int]) -> None: @@ -107,8 +114,15 @@ def __init__(self) -> None: self.release_count = 0 self.scatter_count = 0 - def reserve(self, receiver_req, num_writers: int, *, extra_bytes: int = 0) -> bool: - del extra_bytes + def reserve( + self, + receiver_req, + num_writers: int, + *, + extra_bytes: int = 0, + writers_split_layers: bool = True, + ) -> bool: + del extra_bytes, writers_split_layers self.context = TransferContext( rid_slice=(receiver_req.unique_rid, receiver_req.slice_id), slot_id=0, @@ -693,6 +707,7 @@ def test_partial_bounced_publication_waits_for_queued_writer_success( overlap = SimpleNamespace( ranks=[0, 1], duplicate_head_factor=1, + peer_duplicate_head_factor=1, overlap_pp_size=1, ) receiver._registrar = SimpleNamespace( diff --git a/tests/unittest/llmapi/test_disagg_dflash_guard.py b/tests/unittest/llmapi/test_disagg_dflash_guard.py index 5e35b5a293bd..4195f22898e6 100644 --- a/tests/unittest/llmapi/test_disagg_dflash_guard.py +++ b/tests/unittest/llmapi/test_disagg_dflash_guard.py @@ -17,10 +17,12 @@ from tensorrt_llm.llmapi import ( CacheTransceiverConfig, DFlashDecodingConfig, + DSparkDecodingConfig, NGramDecodingConfig, TorchLlmArgs, + llm_args, ) -from tensorrt_llm.llmapi.llm_args import disagg_dflash_error +from tensorrt_llm.llmapi.llm_args import disagg_dflash_error, disagg_shares_draft_kv pytestmark = pytest.mark.cpu_only @@ -63,24 +65,45 @@ def test_transceiver_without_backend_is_not_disagg(self): class TestTorchLlmArgsValidation: - """The guard wired into TorchLlmArgs.validate_speculative_config.""" + """What TorchLlmArgs.validate_speculative_config does with DFlash on a disaggregated engine.""" # Keeps the DFlash token-budget rule satisfied so only the disagg guard # is under test: max_batch_size * (1 + max_draft_len) <= max_num_tokens. _SHAPE = dict(max_batch_size=8, max_num_tokens=8192, max_seq_len=8192) - def test_disagg_gen_side_dflash_config_is_refused(self): - # Both disagg sides carry the same cache_transceiver_config (the - # gen-side yaml of a ctx/gen pair looks exactly like this), so this - # is the generation-side shape the guard exists for -- and the - # context-side shape, refused by design (see disagg_dflash_error). - with pytest.raises(ValueError, match="DFlash.*disaggregated"): - TorchLlmArgs( - model="dummy", - speculative_config=_dflash_config(), - cache_transceiver_config=_transceiver_config(), - **self._SHAPE, - ) + @pytest.mark.parametrize( + "attention_backend,shared_draft_kv,warns", + [ + ("VANILLA", None, True), + ("TRTLLM", None, False), + ("TRTLLM", "1", False), + ("TRTLLM", "0", True), + ("AUTO", None, False), + ("AUTO", "1", False), + ("AUTO", "0", True), + ], + ) + def test_disagg_dflash_warns_unless_the_draft_kv_is_shared( + self, monkeypatch, attention_backend, shared_draft_kv, warns + ): + # The drafter's context reaches a generation engine only through a draft + # KV shared with the target: the TRTLLM backend, unless + # TRTLLM_DISAGG_SHARED_DRAFT_KV=0 keeps it private. + warnings = [] + monkeypatch.setattr(llm_args.logger, "warning", lambda msg, *a, **k: warnings.append(msg)) + if shared_draft_kv is None: + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", shared_draft_kv) + TorchLlmArgs( + model="dummy", + speculative_config=DFlashDecodingConfig( + max_draft_len=7, attention_backend=attention_backend + ), + cache_transceiver_config=_transceiver_config(), + **self._SHAPE, + ) + assert any("DFlash acceptance is degraded" in str(w) for w in warnings) == warns def test_aggregated_dflash_config_is_allowed(self): args = TorchLlmArgs(model="dummy", speculative_config=_dflash_config(), **self._SHAPE) @@ -91,3 +114,26 @@ def test_disagg_without_speculation_is_allowed(self): model="dummy", cache_transceiver_config=_transceiver_config(), **self._SHAPE ) assert args.speculative_config is None + + +@pytest.mark.parametrize( + ("decoding_type", "attention_backend", "env", "shares"), + [ + ("DFlash", "AUTO", None, True), + ("DFlash", "TRTLLM", None, True), + ("DFlash", "VANILLA", None, False), + ("DFlash", "AUTO", "0", False), + ("DFlash", "VANILLA", "1", True), + ("DSpark", "AUTO", None, False), + ("DSpark", "AUTO", "1", True), + ], +) +def test_disagg_shares_draft_kv(monkeypatch, decoding_type, attention_backend, env, shares): + """Unset, the variable leaves the choice to a DFlash backend that can share; DSpark needs it.""" + if env is None: + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", env) + config_cls = DFlashDecodingConfig if decoding_type == "DFlash" else DSparkDecodingConfig + config = config_cls(max_draft_len=7, attention_backend=attention_backend) + assert disagg_shares_draft_kv(config) is shares