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 7722f76ea664..db0f24b29a6a 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 @@ -28,12 +28,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple +import dataclasses +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union import torch from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, KVCacheManagerV2, Role, _BasePageTableMaterializer, @@ -63,6 +65,7 @@ if TYPE_CHECKING: from transformers import PretrainedConfig + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig @@ -336,6 +339,10 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): # Extra page sizes trtllm-gen may use with this manager (see # FlashInferTrtllmGenFmha); set with the virtual pools. trtllm_gen_extra_tokens_per_block: frozenset = frozenset() + # A standalone DFlash/DSpark drafter's layers may join as explicit draft layer specs: their own + # lifecycle keeps them out of the index-K pools, and size-only coalescing keeps the target pools. + hosts_unified_draft_layers = True + _separate_nvfp4_coalescing = False def __init__( self, @@ -377,10 +384,26 @@ def __init__( kwargs.get("num_kv_heads"), shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")), ) + layer_specs = kwargs.get("layer_specs") + # Explicit specs declare the draft layers, which are never sparse or native target layers. + draft_ids = { + spec.global_layer_id + for spec in (layer_specs or ()) + if spec.owner is KVCacheLayerOwner.DRAFT + } + if layer_specs is not None: + # Layer specs may be a PP rank's subset: the checkpoint's layer count is the global + # target count, and without it draft layers still follow every target layer. + num_hidden_layers = getattr(kwargs.get("pretrained_config"), "num_hidden_layers", None) + if num_hidden_layers is not None: + num_target_layers = int(num_hidden_layers) + elif draft_ids: + num_target_layers = min(draft_ids) if sparse_layer_ids is None: sparse_layer_ids = default_sparse_layer_ids( kwargs.get("pretrained_config"), num_target_layers ) + sparse_layer_ids = [i for i in sparse_layer_ids if int(i) not in draft_ids] if disable_index_value_layer_ids is None: disable_index_value_layer_ids = list(sparse_layer_ids) @@ -390,7 +413,9 @@ def __init__( self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) if implementation == "triton" and native_dense_cache and num_target_layers is not None: self._dense_native_layer_ids = tuple( - i for i in range(num_target_layers) if i not in self.sparse_layer_ids + i + for i in range(num_target_layers) + if i not in self.sparse_layer_ids and i not in draft_ids ) self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) self.sparse_index_dim = sparse_index_dim @@ -403,12 +428,61 @@ def __init__( # Filled at the end of __init__; the base __init__ may already call get_buffers. self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {} self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {} + if layer_specs is not None and any( + spec.owner is KVCacheLayerOwner.TARGET + and spec.dtype != kwargs.get("dtype", DataType.HALF) + for spec in layer_specs + ): + # Target views and block scales follow the manager dtype. + raise ValueError(f"{type(self).__name__} needs target layer specs in the cache dtype.") + if layer_specs is not None and any( + spec.owner is KVCacheLayerOwner.DRAFT and spec.dtype == DataType.NVFP4 + for spec in layer_specs + ): + # Block scales here belong to the MSA sparse target layers only. + raise ValueError(f"{type(self).__name__} does not host NVFP4 draft layers.") + hosts_draft_layers = layer_specs is not None and any( + spec.owner is KVCacheLayerOwner.DRAFT for spec in layer_specs + ) + if ( + hosts_draft_layers + and kwargs.get("is_disagg") + and self._main_kv_mapper_kind != MapperKind.HND + ): + # Drafters write head-major K/V, but transfer reads every K/V buffer with the main mapper. + raise ValueError( + f"{type(self).__name__} transfers draft layers only with the HND K/V layout." + ) + if layer_specs is not None and kwargs.get("dtype") == DataType.NVFP4: + # The hybrid layout of _build_cache_config as specs: FP8 on the dense target layers. + kwargs["layer_specs"] = tuple( + dataclasses.replace(spec, dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE)) + if spec.owner is KVCacheLayerOwner.TARGET + and spec.dtype == DataType.NVFP4 + and spec.global_layer_id not in self.sparse_layer_ids + else spec + for spec in layer_specs + ) super().__init__(*args, **kwargs) + if hosts_draft_layers: + groups = self.layer_to_pool_mapping_dict + target_groups = {groups[self.layer_offsets[layer]] for layer in self.target_layer_ids} + if any( + groups[self.layer_offsets[layer]] in target_groups for layer in self.draft_layer_ids + ): + # Coalesced with target buffers, draft layers would get target-relative pool offsets. + self.shutdown() + raise ValueError( + f"{type(self).__name__} hosts draft layers only in their own lifecycle." + ) + if self.dtype == DataType.NVFP4: dense_target_layers = [ layer - for layer in self.layer_offsets + for layer in ( + self.target_layer_ids if self.layer_specs is not None else self.layer_offsets + ) if layer not in self.sparse_layer_ids and layer not in self._shared_draft_layer_ids ] logger.info( @@ -513,8 +587,16 @@ def _build_cache_config(self, config): ] return super()._build_cache_config(config) - def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): + def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role) -> int: """Report the hybrid sparse-NVFP4 / dense-FP8 storage footprint.""" + if self.layer_specs is not None: + # Explicit layer specs carry K/V and scales; index-K is this manager's extra buffer. + if data_role == Role.INDEX_KEY: + if int(self.pp_layers[int(local_layer_idx)]) not in self.sparse_layer_ids: + return 0 + index_dtype = self._torch_dtype_for_index_cache() + return self.sparse_index_dim * torch.tensor([], dtype=index_dtype).element_size() + return super().get_layer_bytes_per_token(local_layer_idx, data_role) if self.dtype != DataType.NVFP4: return super().get_layer_bytes_per_token(local_layer_idx, data_role) global_layer_idx = int(self.pp_layers[int(local_layer_idx)]) @@ -556,17 +638,23 @@ def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[Dat return None return super()._get_block_scale_role(role_a, layer_id) - def get_draft_kv_cache_view(self) -> Optional["MiniMaxM3DraftKVCacheView"]: + def get_draft_kv_cache_view( + self, + ) -> Optional[Union["MiniMaxM3DraftKVCacheView", "MiniMaxM3KVCacheManagerV2"]]: """FP8/P128 view of the shared draft layer in a hybrid NVFP4 cache, or None. FP8/BF16 caches address the draft layer through this manager's own - attention-op pools instead. + attention-op pools instead. With explicit draft layer specs, this + manager itself is the draft view. Only meaningful on a target manager carrying appended one-model draft layers; built lazily so the manager's page tables exist. A method rather than a property so ``getattr`` fetches it without executing it (see ``get_draft_kv_cache_manager``). """ + if self.layer_specs is not None: + # Explicit draft layers are addressed through this manager itself. + return super().get_draft_kv_cache_view() if self.dtype != DataType.NVFP4 or self.is_draft or not self._shared_draft_layer_ids: return None if self.enable_swa_scratch_reuse: @@ -589,15 +677,19 @@ def get_draft_kv_cache_view(self) -> Optional["MiniMaxM3DraftKVCacheView"]: ) return self._draft_kv_cache_view - def add_dummy_requests(self, *args, **kwargs): + def add_dummy_requests(self, *args: Any, **kwargs: Any) -> Optional[List[LlmRequest]]: """Drop the draft KV cache view before delegating. The base method mirrors dummy KV caches into a *separate* draft manager. With shared draft layers a dummy request's blocks already span the drafter's pool (pools allocate in lockstep per logical - block), and the view owns no block lifecycle. + block), and the view owns no block lifecycle. Explicit draft layers + are this manager's own, so their view is this manager itself. """ - if isinstance(kwargs.get("draft_kv_cache_manager"), MiniMaxM3DraftKVCacheView): + draft_kv_cache_manager = kwargs.get("draft_kv_cache_manager") + if draft_kv_cache_manager is self or isinstance( + draft_kv_cache_manager, MiniMaxM3DraftKVCacheView + ): kwargs["draft_kv_cache_manager"] = None return super().add_dummy_requests(*args, **kwargs) @@ -608,7 +700,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: pool's slot ids; :meth:`copy_batch_block_offsets` scales them. """ super()._prepare_page_table_tensor(index_mapper_capacity) - if self.dtype == DataType.NVFP4: + # Explicit draft layers keep the pools: the unified drafter reads their block offsets. + if self.dtype == DataType.NVFP4 and self.layer_specs is None: # Nothing reads a hybrid cache's own block offsets: MSA layers address pages # by slot id and the shared draft layer reads MiniMaxM3DraftKVCacheView's # table. Without pools, a new reader fails instead of reading unfilled offsets. @@ -683,8 +776,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: - # Staging for the per-layer block-offset copy, which a hybrid NVFP4 cache never runs. - if self.dtype != DataType.NVFP4: + # Staging for the per-layer block-offset copy, which a pool-less hybrid NVFP4 cache never runs. + if self.dtype != DataType.NVFP4 or self.layer_specs is not None: super()._prepare_swa_scratch_copy_tensors(index_mapper_capacity) def copy_batch_block_offsets( @@ -896,7 +989,15 @@ def get_buffers( at K's base, then slices ``[:, :2]`` to extract K+V. The slice preserves the dim-0 stride (``scale * page_stride``), so ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. + + Explicit draft layers keep the base view: their own lifecycle's pool holds only their K/V. """ + if self.layer_specs is not None and self.get_layer_owner(layer_idx) is ( + KVCacheLayerOwner.DRAFT + ): + return super().get_buffers( + layer_idx, kv_layout if kv_layout is not None else self._main_kv_layout_name() + ) if kv_layout is None: kv_layout = self._main_kv_view_layout if kv_layout not in ("HND", "NHD"): @@ -1040,6 +1141,37 @@ def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> in ) return layers_by_addr.index(int(layer_id)) + def get_batch_cache_indices( + self, + request_ids: List[int], + layer_idx: Optional[int] = None, + num_blocks_per_seq: Optional[Sequence[int]] = None, + *, + raw_indices: bool = False, + ) -> List[List[int]]: + """Return indices into a layer's :meth:`get_buffers` view. + + Explicit draft layers keep the base view, so their indices take the base + conversion; every other layer gets raw slot ids. + """ + if ( + self.layer_specs is None + or layer_idx is None + or self.get_layer_owner(layer_idx) is not KVCacheLayerOwner.DRAFT + ): + return super().get_batch_cache_indices( + request_ids, layer_idx, num_blocks_per_seq, raw_indices=raw_indices + ) + return KVCacheManagerV2._get_batch_cache_indices_by_pool_id( + self, + request_ids, + pool_id=self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]], + is_kv_aggregate=not raw_indices, + num_blocks_per_seq=num_blocks_per_seq, + index_scale=1 if raw_indices else self.get_layer_page_index_scale(layer_idx), + kv_factor=self.get_layer_kv_factor(layer_idx), + ) + def _get_batch_cache_indices_by_pool_id( self, request_ids, 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 42b621bde4c1..bc1107bcda8d 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 @@ -1816,6 +1816,12 @@ class KVCacheManagerV2(BaseResourceManager): # need snapshotted (request_prefix_snapshot); kept apart from the request's # own expect_snapshot_points, which other code rewrites. Created on first use. _waiter_snapshot_points: Optional[Dict[int, set[int]]] = None + # Declared on the class so managers built without __init__ have no layer specs. + layer_specs: Optional[Tuple["KVCacheLayerSpec", ...]] = None + # A subclass whose extra buffers may join explicit layer specs (draft layers then get their own lifecycle). + hosts_unified_draft_layers: bool = False + # Layer specs coalesce NVFP4 data and scales apart; a subclass addressing its pools by size opts out. + _separate_nvfp4_coalescing: bool = True def __init__( self, @@ -1861,6 +1867,7 @@ def __init__( if ( self.layer_specs is not None and type(self)._extra_buffers_per_layer is not KVCacheManagerV2._extra_buffers_per_layer + and not self.hosts_unified_draft_layers ): raise ValueError( f"{type(self).__name__} does not support layer_specs with " @@ -3729,8 +3736,13 @@ def _build_base_config( ) ) + initial_pool_ratio = kv_cache_config.pool_ratio if self.layer_specs is not None: scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + # Subclass buffers beyond K/V/scales (e.g. MiniMax-M3's index-K) follow the spec-built ones. + extra_buffers_per_layer = ( + self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} + ) layer_configs = [ AttentionLayerConfig( layer_id=LayerId(local_layer_id), @@ -3741,18 +3753,41 @@ def _build_base_config( # Keep NVFP4 data (1), its scales (2), and other data (0) # separate while preserving matching K/V page order. coalescing_group=( - 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 + 1 + int(role in scale_roles) + if layer.dtype == DataType.NVFP4 and self._separate_nvfp4_coalescing + else 0 ), ) for role in (Role.KEY, Role.VALUE, *scale_roles) if role in layer.buffer_roles - ], + ] + + list(extra_buffers_per_layer.get(local_layer_id, ())), sliding_window_size=self.max_attention_window_vec[local_layer_id], num_sink_tokens=layer.num_sink_tokens, reuse_policy=layer.reuse_policy, ) for local_layer_id, layer in enumerate(self.layer_specs) ] + if initial_pool_ratio is None and typical_step is None: + initial_pool_ratio = self._windowed_draft_pool_ratio( + layer_configs, + cache_tiers[0].quota, + tokens_per_block, + kv_cache_config.max_util_for_resume, + ) + if initial_pool_ratio is not None: + # Allocation rounds the ratio's byte target; this floor keeps one fresh + # context chunk admissible in every lifecycle. + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) + ) else: # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional # per-layer BufferConfig entries — for example a sparse index-K @@ -3827,7 +3862,7 @@ def _build_base_config( kv_cache_config.enable_block_reuse and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE ), - initial_pool_ratio=kv_cache_config.pool_ratio, + initial_pool_ratio=initial_pool_ratio, pool_rebalance=PoolRebalanceConfig( min_sampled_kv_caches=rebalance_cfg.min_sampled_kv_caches, cooldown_secs=rebalance_cfg.cooldown_secs, @@ -3891,6 +3926,85 @@ def _allocate_cache_pools( f"external cache storage ({external_bytes} bytes)." ) + def _windowed_draft_pool_ratio( + self, + layer_configs: List[AttentionLayerConfig], + gpu_quota: int, + tokens_per_block: int, + max_util_for_resume: float = 1.0, + ) -> Optional[List[float]]: + """Hot-tier ``[target, draft]`` byte ratios when the draft layers form one windowed lifecycle. + + Without a pool ratio or typical step the backend splits the hot tier for one 2048-token + request, which says nothing about how many sequences hold a draft window at once. The + draft layers need every admissible sequence's window plus one context chunk, which is + what the KV cache cost already charges for them, over the resume watermark and capped + at an equal-token share. + """ + if self.max_num_tokens is None or gpu_quota <= 0: + # Without a context-chunk bound or a quota the backend's split is the safe one. + return None + owners = [spec.owner for spec in self.layer_specs] + # Key layers as the backend groups them into lifecycles (sinks count in whole blocks). + life_cycles = [ + AttnLifeCycle.make( + c.sliding_window_size, c.num_sink_tokens, tokens_per_block, c.reuse_policy + ) + for c in layer_configs + ] + keys = [(lc.window_size, lc.num_sink_blocks, lc.reuse_policy) for lc in life_cycles] + target_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.TARGET} + draft_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.DRAFT} + # Lifecycle ids follow first appearance, so the target's must come first. + if ( + len(target_keys) != 1 + or len(draft_keys) != 1 + or target_keys == draft_keys + or not next(iter(draft_keys))[0] + or owners[0] is not KVCacheLayerOwner.TARGET + ): + return None + draft_specs = [spec for spec in self.layer_specs if spec.owner is KVCacheLayerOwner.DRAFT] + draft_sizes = [spec.bytes_per_token() for spec in draft_specs] + _, window_bytes_per_sequence = _estimate_swa_cache_size( + draft_sizes, + [next(iter(draft_keys))[0]] * len(draft_specs), + tokens_per_block, + context=False, + scratch=False, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=draft_specs, + ) + # Sequences that can hold KV at once (see max_admissible_sequences), counting the overlap + # scheduler's retiring request: the pool is sized before the scheduler mode is known. + num_sequences = ( + self.max_batch_size + * self.mapping.pp_size + * (2 if self.is_disagg or not self.mapping.has_pp() else 1) + ) + draft_bytes = ( + window_bytes_per_sequence * num_sequences + sum(draft_sizes) * self.max_num_tokens + ) / max_util_for_resume + bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0} + for config, owner in zip(layer_configs, owners): + bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers) + # The context-chunk floor declared with this ratio: every lifecycle holds one fresh chunk + # below the resume watermark, each pool rounded to the allocation granularity (bounded + # here per buffer). Without room for it the backend keeps its own split. + chunk_blocks = -(-(self.max_num_tokens + self.num_extra_kv_tokens) // tokens_per_block) + floor_slots = math.ceil(chunk_blocks / float(np.float32(max_util_for_resume))) + granularity = _gpu_alloc_granularity(gpu_quota) + floor_bytes = sum( + -(-floor_slots * buffer.size // granularity) * granularity + for config in layer_configs + for buffer in config.buffers + ) + if floor_bytes > gpu_quota: + return None + token_share = bytes_per_block[KVCacheLayerOwner.DRAFT] / sum(bytes_per_block.values()) + draft_ratio = min(draft_bytes / gpu_quota, token_share) + return [1.0 - draft_ratio, draft_ratio] + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Customize the general cache config for a specialized cache manager.""" return config diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py new file mode 100644 index 000000000000..fd26e42e5447 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py @@ -0,0 +1,430 @@ +# 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. +"""A standalone DSpark drafter's layers as explicit draft layer specs of the MiniMax-M3 manager.""" + +from dataclasses import replace +from types import SimpleNamespace +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( + MiniMaxM3KVCacheManagerV2, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MiniMaxM3SparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +NUM_TARGET = 4 # M3 convention without a checkpoint: layers 0..2 dense, 3 sparse +NUM_DRAFT = 2 +DRAFT_WINDOW = 256 +HEADS = 2 +HEAD_DIM = 128 + + +def _specs(draft_window: Optional[int] = DRAFT_WINDOW) -> tuple[KVCacheLayerSpec, ...]: + """Target layers in the NVFP4 cache dtype, then a BF16 drafter (windowed by default).""" + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + target = [ + KVCacheLayerSpec( + global_layer_id=layer, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.NVFP4, + num_kv_heads=HEADS, + head_dim=HEAD_DIM, + buffer_roles=(Role.KEY, Role.VALUE, *scale_roles), + ) + for layer in range(NUM_TARGET) + ] + draft = [ + KVCacheLayerSpec( + global_layer_id=NUM_TARGET + layer, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=HEADS, + head_dim=HEAD_DIM, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=draft_window, + ) + for layer in range(NUM_DRAFT) + ] + return tuple(target + draft) + + +def _manager( + layer_specs: Optional[tuple[KVCacheLayerSpec, ...]], + dtype: DataType = DataType.NVFP4, + implementation: str = "msa", + is_disagg: bool = False, +) -> MiniMaxM3KVCacheManagerV2: + """A MiniMax-M3 manager (hybrid NVFP4 MSA by default), with or without draft layer specs.""" + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=NUM_TARGET, + num_kv_heads=HEADS, + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=1024, + max_batch_size=2, + max_num_tokens=256, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=dtype, + vocab_size=16, + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation=implementation, + indexer_kv_dtype="fp8" if implementation == "msa" else "bf16", + ), + layer_specs=layer_specs, + is_disagg=is_disagg, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None: + """Draft layers join the hybrid manager without changing the target's pages.""" + torch.cuda.init() + plain = _manager(None) + unified = _manager(_specs()) + try: + assert unified.target_layer_ids == tuple(range(NUM_TARGET)) + assert unified.draft_layer_ids == tuple(range(NUM_TARGET, NUM_TARGET + NUM_DRAFT)) + # Dense target layers keep the FP8 half of the hybrid layout; draft layers their own dtype. + dtypes = [unified.get_layer_dtype(layer) for layer in unified.pp_layers] + assert dtypes == [DataType.FP8] * 3 + [DataType.NVFP4] + [DataType.BF16] * NUM_DRAFT + for layer in range(NUM_TARGET): + for role in (Role.KEY, Role.VALUE): + assert unified.impl.get_page_stride( + unified.layer_offsets[layer], role + ) == plain.impl.get_page_stride(plain.layer_offsets[layer], role) + assert unified.get_index_k_buffer(3).shape[1:] == plain.get_index_k_buffer(3).shape[1:] + # Index-K is the manager's own buffer, not a layer-spec role. + index_k_bytes = [ + unified.get_layer_bytes_per_token(unified.layer_offsets[layer], Role.INDEX_KEY) + for layer in unified.pp_layers + ] + assert index_k_bytes == [0, 0, 0, 128] + [0] * NUM_DRAFT + + # Draft layers: their own lifecycle, the base view in their dtype. + groups = unified.layer_to_pool_mapping_dict + draft_groups = {groups[unified.layer_offsets[layer]] for layer in unified.draft_layer_ids} + target_groups = {groups[unified.layer_offsets[layer]] for layer in unified.target_layer_ids} + assert len(draft_groups) == 1 and not draft_groups & target_groups + view = unified.get_buffers(NUM_TARGET, kv_layout="HND") + assert view.dtype == torch.bfloat16 + assert tuple(view.shape[1:]) == (2, HEADS, 128, HEAD_DIM) + assert unified.get_draft_kv_cache_view() is unified + assert plain.get_draft_kv_cache_view() is None + + # The unified drafter reads its layers' block offsets, so they keep attention-op pools. + assert plain.num_attention_op_pools == 0 + assert unified.num_attention_op_pools > 0 + # CUDA-graph padding passes the draft view (this manager) back in. + (request,) = unified.add_dummy_requests( + [7], token_nums=[200], draft_kv_cache_manager=unified.get_draft_kv_cache_view() + ) + offsets = torch.zeros( + unified.num_attention_op_pools, + 1, + 2, + unified.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + unified.copy_batch_block_offsets(offsets, [7], 1, 1, 1) + torch.cuda.synchronize() + pool = int(unified.kv_cache_pool_mapping[unified.layer_offsets[NUM_TARGET], 0]) + slots = [ + s + for s in unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET, raw_indices=True)[0] + if s >= 0 + ] + assert len(slots) == 2 + # The row the drafter reads holds the draft layer's K page of each slot, V right after it. + scale = unified.get_layer_page_index_scale(NUM_TARGET) + pages = [[slot * scale + plane for slot in slots] for plane in (0, 1)] + assert offsets[pool, 0, :, : len(slots)].tolist() == pages + # Draft-layer indices address the base view's contiguous K/V rows, like the flat table. + rows = [page // 2 for page in pages[0]] + assert unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET)[0][: len(slots)] == rows + flat = unified.get_batch_cache_indices_flat([7], [len(slots)], layer_idx=NUM_TARGET) + assert flat.tolist() == rows + # Target layers keep MiniMax's raw slot ids. + sparse_slots = unified.get_batch_cache_indices([7], layer_idx=3)[0] + assert ( + sparse_slots == unified.get_batch_cache_indices([7], layer_idx=3, raw_indices=True)[0] + ) + unified.free_resources(request) + + # The draft lifecycle is sized by its window, not by its share of the bytes per token. + target_ratio, draft_ratio = unified.kv_cache_manager_py_config.initial_pool_ratio + assert target_ratio + draft_ratio == pytest.approx(1.0) + assert ( + 0 + < draft_ratio + < 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT / (unified.get_cache_bytes_per_token()) + ) + finally: + unified.shutdown() + plain.shutdown() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_page_indices_follow_the_layer_owner() -> None: + """Target layers get raw slot ids and draft layers base view rows.""" + torch.cuda.init() + manager = _manager(_specs()) + try: + (request,) = manager.add_dummy_requests([7], token_nums=[200]) + for layer in manager.pp_layers: + raw = manager.get_batch_cache_indices([7], layer_idx=layer, raw_indices=True)[0] + indices = manager.get_batch_cache_indices([7], layer_idx=layer)[0] + if layer in manager.draft_layer_ids: + scale = manager.get_layer_page_index_scale(layer) + expected = [s * scale // 2 if s >= 0 else s for s in raw] + else: + expected = raw + assert indices[: len(raw)] == expected, layer + manager.free_resources(request) + finally: + manager.shutdown() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_nvfp4_draft_layers_are_rejected() -> None: + """Block scales belong to the sparse target layers, so an NVFP4 drafter cannot join.""" + specs = _specs() + nvfp4_draft = replace( + specs[NUM_TARGET], + dtype=DataType.NVFP4, + buffer_roles=(Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE), + ) + with pytest.raises(ValueError, match="does not host NVFP4 draft layers"): + _manager((*specs[:NUM_TARGET], nvfp4_draft, *specs[NUM_TARGET + 1 :])) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +@pytest.mark.parametrize( + "sparse_attention_freq,num_specs,expected", + [ + (None, NUM_TARGET + NUM_DRAFT, [3]), + ([0, 1, 0, 1], NUM_TARGET + NUM_DRAFT, [1, 3]), + ([0, 1, 0, 1], NUM_TARGET, [1, 3]), + ], + ids=["default", "config", "config-target-only"], +) +def test_specs_decide_the_sparse_target_layers( + sparse_attention_freq: Optional[list[int]], num_specs: int, expected: list[int] +) -> None: + """A per-layer head list spanning the drafter keeps the target's sparse split.""" + torch.cuda.init() + pretrained_config = None + if sparse_attention_freq is not None: + pretrained_config = SimpleNamespace( + num_hidden_layers=NUM_TARGET, + sparse_attention_config={"sparse_attention_freq": sparse_attention_freq}, + ) + manager = MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=NUM_TARGET, + num_kv_heads=[HEADS] * (NUM_TARGET + NUM_DRAFT), + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=1024, + max_batch_size=2, + max_num_tokens=256, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.NVFP4, + vocab_size=16, + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa", indexer_kv_dtype="fp8" + ), + layer_specs=_specs()[:num_specs], + pretrained_config=pretrained_config, + ) + try: + assert manager.sparse_layer_ids == expected + finally: + manager.shutdown() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_target_specs_need_the_cache_dtype() -> None: + """Target views and block scales follow the manager dtype, so target specs must use it.""" + specs = _specs() + fp8_target = replace(specs[3], dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE)) + with pytest.raises(ValueError, match="target layer specs in the cache dtype"): + _manager((*specs[:3], fp8_target, *specs[NUM_TARGET:])) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_draft_layers_need_their_own_lifecycle() -> None: + """Unwindowed draft layers would share the target's lifecycle and pool offsets.""" + torch.cuda.init() + with pytest.raises(ValueError, match="only in their own lifecycle"): + _manager(_specs(draft_window=None)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_disagg_draft_layers_need_the_hnd_layout() -> None: + """A Triton FP8 cache transfers token-major K/V, but drafters write head-major K/V.""" + specs = tuple( + replace(spec, dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE)) + if spec.owner is KVCacheLayerOwner.TARGET + else spec + for spec in _specs() + ) + with pytest.raises(ValueError, match="only with the HND K/V layout"): + _manager(specs, dtype=DataType.FP8, implementation="triton", is_disagg=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") +def test_window_sized_draft_pool_admits_a_full_context_chunk() -> None: + """Allocation rounds the window-sized draft share, but one full chunk still fits.""" + torch.cuda.init() + max_num_tokens = 8320 + specs = tuple( + KVCacheLayerSpec( + global_layer_id=layer, + owner=owner, + dtype=DataType.BF16, + num_kv_heads=heads, + head_dim=HEAD_DIM, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=window, + ) + for layer, owner, heads, window in ( + (0, KVCacheLayerOwner.TARGET, 2 * HEADS, None), + (1, KVCacheLayerOwner.DRAFT, HEADS, DRAFT_WINDOW), + ) + ) + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=1 << 30, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=1, + num_kv_heads=[2 * HEADS, HEADS], + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=16384, + max_batch_size=1, + max_num_tokens=max_num_tokens, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + vocab_size=16, + layer_specs=specs, + ) + try: + assert manager.kv_cache_manager_py_config.initial_pool_ratio is not None + # Different page sizes keep the draft lifecycle in its own physical pool. + assert int(manager.kv_cache_pool_mapping[0, 0]) != int(manager.kv_cache_pool_mapping[1, 0]) + requests = manager.add_dummy_requests([7], token_nums=[max_num_tokens]) + assert requests is not None + manager.free_resources(requests[0]) + finally: + manager.shutdown() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("pp_size", [1, 2]) +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_window_sized_draft_ratio(is_disagg: bool, pp_size: int) -> None: + """A windowed draft lifecycle gets its windows plus one chunk, not a bytes-per-token share.""" + + def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] = None) -> list: + return [ + SimpleNamespace( + sliding_window_size=spec.sliding_window_size, + num_sink_tokens=target_sinks if spec.owner is KVCacheLayerOwner.TARGET else None, + reuse_policy=spec.reuse_policy, + buffers=[SimpleNamespace(size=spec.bytes_per_token() * 128)], + ) + for spec in specs + ] + + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = _specs()[NUM_TARGET - 1 :] + manager.max_batch_size = 2 + manager.max_num_tokens = 256 + manager.mapping = Mapping(world_size=pp_size, rank=0, pp_size=pp_size) + manager.is_disagg = is_disagg + manager._generation_kv_capacity_headroom = 1 + manager.num_extra_kv_tokens = 0 + configs = _configs(manager.layer_specs) + draft_bytes_per_token = 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT + # Admissible sequences (batch x PP stages, doubled for the overlap scheduler's retiring + # request unless an aggregated PP pipeline) hold three pages of window each, plus one chunk. + num_sequences = 2 * pp_size * (2 if is_disagg or pp_size == 1 else 1) + draft_bytes = (num_sequences * 3 * 128 + 256) * draft_bytes_per_token + assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) == pytest.approx( + [1 - draft_bytes / (1 << 30), draft_bytes / (1 << 30)] + ) + # Resume refuses a pool above its watermark, so the share covers the windows below it. + assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128, 0.5)[1] == pytest.approx( + 2 * draft_bytes / (1 << 30) + ) + # A quota too small for every window keeps the bytes-per-token split. + token_share = draft_bytes_per_token / sum( + spec.bytes_per_token() for spec in manager.layer_specs + ) + manager.max_batch_size = 64 + assert manager._windowed_draft_pool_ratio(configs, 8 << 20, 128)[1] == pytest.approx( + token_share + ) + manager.max_batch_size = 2 + # The quota must hold the context-chunk floor: two blocks per lifecycle over the resume + # watermark, each buffer rounded to the 2 MiB allocation granularity; otherwise the backend + # keeps its own split. + floor_bytes = len(configs) * (2 << 20) + assert manager._windowed_draft_pool_ratio(configs, floor_bytes - 1, 128, 0.95) is None + assert manager._windowed_draft_pool_ratio(configs, floor_bytes, 128, 0.95) is not None + # Without a context-chunk bound the backend keeps its own split. + manager.max_num_tokens = None + assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) is None + manager.max_num_tokens = 256 + # Without a windowed draft lifecycle the backend keeps its own split. + manager.layer_specs = _specs(draft_window=None)[NUM_TARGET - 1 :] + assert manager._windowed_draft_pool_ratio(_configs(manager.layer_specs), 1 << 30, 128) is None + # A target sharing the drafter's window shares its lifecycle too: sink tokens that round to + # no sink block do not make a second one. + manager.layer_specs = tuple( + replace(spec, sliding_window_size=DRAFT_WINDOW) for spec in _specs()[NUM_TARGET - 1 :] + ) + assert ( + manager._windowed_draft_pool_ratio(_configs(manager.layer_specs, 0), 1 << 30, 128) is None + )