From 1a1b710a72fe64723445f0c26a3564259fe8de8d Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 20 Oct 2026 23:22:07 -0700 Subject: [PATCH 01/23] [TRTLLM-16722][feat] Dflash Unified KVCache (port of NVIDIA/TensorRT-LLM#20486 incl. #20178 foundation) Squashed net change of PR #20486 at head 714916dc1 relative to the main it last merged (0829689c5), ported onto feat/musespark-open-agentx; unrelated main commits are not included. #20486 carries the unified KV cache manager foundation of PR #20178 through eea6dcbe9 (#20178's later 61248ec36, 77d4eae1a and 636715a1b are not part of the #20709 stack and are not ported). Commits: 3bc36a556 [None][feat] add unified KV cache layout foundation dad9f1f36 [None][refactor] simplify unified KV cache layer specs 6ea5ad816 [None][test] remove unified KV cache foundation tests 090cbc02d [None][docs] remove unified KV cache design document 511975f49 [None][feat] support per-layer storage layouts in KVCacheManagerV2 bfc8bc729 fix: normalize explicit KV cache spec windows 18c71c1a3 fix: use explicit KV cache dtype for cold-page codec ae64dcab3 test: cover explicit KV cache dtype for cold-page codec 701abead2 fix: honor physical slot strides in KV cache views 623816cff test: cover mixed KV cache odd-slot addressing 7cde28eb4 fix: restore KV cache subclass and test compatibility 4a3bb2582 [None][test] fix mixed KV cache slot address arithmetic 9d0107ba4 [None][fix] align NVFP4 data and scale page layouts e273a984d [None][fix] accept mixed KV cache pointer table layouts f6b1d08e2 [None][fix] preserve sink and reuse policy in KV cache costing b6db61277 [TRTLLM-20434][feat] Dflash Unified KVCache 39d814dcc [None][fix] share the DFlash draft KV cache in disagg only with speculative_config in both sections 22cf25c6f [None][fix] gate the disagg shared draft KV cache on matching tensor_parallel_size feed06a1b [None][test] cover the disagg shared draft KV gate for generation-only speculative_config and mismatched TP 0319dc361 test: use default timeout for unified DFlash KV test 1f391284f fix: restore DFlash private arena fallback for separate KV caches 97fb214ab fix: seed unified DFlash context lengths from bounded token counts caf4c45a2 test: cover DFlash mixed-precision block tables 875eff2ee fix: require matching attention DP for shared draft KV 212fdabdf test: enable shared draft KV in DFlash integration test 1302d43bf fix: keep mixed-window DFlash caches separate Port conflict resolutions (base = main 0829689c5, ours = agentx b7979cae8 + !10912): - kv_cache_manager_v2.py: our joint_kv_cache_reuse kwarg and the PR's layer_specs kwarg land at the same spot in _get_per_request_layer_tokens, get_cache_bytes_per_request and get_cache_quota_per_request (5 hunks): keep both. The generic layer_configs loop moves under the PR's `else:`; main's kv_pool_rebalance_config (#19077) is not on agentx and is dropped. - _util.py: agentx's _per_manager_cache_cost has no tokens_per_block parameter (main-only); keep the PR's unified layer_specs injection. - dflash.py: keep agentx's windowed VANILLA arena (_ctx_base) branch; ctx_len_gen (PR) equals self._ctx_len[slots]. - rank_info.py: drop main's bounded_replay_on_generation (not on agentx); keep num_draft_layers. - tests: drop main-only imports (PyTorchModelEngine, KvCacheRetentionConfig). !10912 stays intact: validate_block_attention_windows() is not re-added. Signed-off-by: Allison Lim Signed-off-by: Zheyu Fu --- .../kv_cache_manager_v2/config.h | 4 + .../kv_cache_manager_v2/storage/config.cpp | 14 +- .../batch_manager/kvCacheManagerV2.cpp | 8 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 9 +- .../_torch/attention/backends/flashinfer.py | 11 +- .../sparse/minimax_m3/cache_manager.py | 3 + .../_torch/disaggregation/native/peer.py | 8 + .../_torch/disaggregation/native/rank_info.py | 6 + tensorrt_llm/_torch/pyexecutor/_util.py | 220 +++++- .../kv_cache/kv_cache_manager_v2.py | 685 ++++++++++++++---- .../pyexecutor/scheduler/scheduler_v2.py | 7 +- tensorrt_llm/_torch/speculative/dflash.py | 125 +++- tensorrt_llm/_torch/speculative/interface.py | 10 +- tensorrt_llm/_torch/speculative/utils.py | 4 +- tensorrt_llm/commands/serve.py | 13 + .../runtime/kv_cache_manager_v2/__init__.pyi | 1 + .../references/acceptance_length.yaml | 6 + .../accuracy/test_disaggregated_serving.py | 92 ++- .../test_lists/test-db/l0_dgx_b200.yml | 1 + .../sparse/glm_kpool/test_glm_kpool.py | 1 + .../test_minimax_m3_shared_draft_layers.py | 2 + .../kv_cache/test_dflash_unified_kv_gate.py | 214 ++++++ .../kv_cache/test_kv_cache_estimation.py | 73 +- .../kv_cache/test_kvcm2_integration.py | 394 +++++++++- .../kv_cache/test_mamba_cache_manager.py | 3 + .../test_minimax_m3_kv_transfer.py | 8 + tests/unittest/disaggregated/test_peer.py | 26 + .../unittest/disaggregated/test_rank_info.py | 62 ++ 28 files changed, 1837 insertions(+), 173 deletions(-) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py 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/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/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/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 72bd85972e0c..b2d7451a6dd6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -170,6 +170,14 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + if self._ri.num_draft_layers != peer_ri.num_draft_layers: + # The layer-count tolerance below would transfer only the target + # layers and leave the drafter's prompt KV unset on the receiver. + raise ValueError( + "PeerRegistrar: unified draft KV layer count differs " + f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " + "both peers must use the same speculative decoding KV cache layout." + ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 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/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e43785e7e956..f81e554f802d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -40,6 +40,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy from ..attention.backends import get_sparse_attn_kv_cache_manager from ..disaggregation.kv_cache_transceiver import ( @@ -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, @@ -862,6 +865,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 +1009,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) @@ -1888,6 +1898,14 @@ def _create_kv_cache_manager( spec_dec_layer_mask = [True] * num_target_layers estimating_kv_cache = estimating_kv_cache and not self._skip_est + unified_layer_specs = None + if self._should_unify_dflash_kv_cache(): + unified_layer_specs = self._build_unified_dflash_layer_specs( + kv_cache_config) + logger.info( + "DFlash: target and draft KV layers share one KV cache manager " + f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})." + ) kv_cache_manager = _create_kv_cache_manager( model_engine=model_engine, kv_cache_manager_cls=kv_cache_manager_cls, @@ -1912,6 +1930,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 +1970,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 +1997,194 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) + def _should_unify_dflash_kv_cache(self) -> bool: + """Whether this run can place target and DFlash layers in one V2 manager.""" + spec_config = self._speculative_config + if (spec_config is None or not spec_config.spec_dec_mode.is_dflash() + or spec_config.attention_backend != "TRTLLM" + or not spec_config._allow_separate_draft_kv_cache + or spec_config.draft_len_schedule is not None + or self._draft_model_engine is not None + or self._draft_config is None or self._mapping.pp_size != 1 + or self._mapping.cp_size != 1 + or self._llm_args.attn_backend != "TRTLLM" or is_gemma4_hybrid( + self._model_engine.model.model_config.pretrained_config) + or self._llm_args.kv_cache_compression_config is not None): + return False + # Specialized managers add cache roles and lifecycle protocols that an + # external drafter cannot yet describe. Keep their existing fallback. + if self._kv_cache_manager_cls is not KVCacheManagerV2: + return False + target_layout = self._dflash_kv_layout( + self._model_engine.model.model_config, self._model_engine.dtype) + draft_model_config = self._get_draft_kv_model_config() + draft_layout = self._dflash_kv_layout( + draft_model_config, + draft_model_config.pretrained_config.torch_dtype) + if target_layout is None or draft_layout is None: + logger.info_once( + "DFlash: a per-layer KV-head count is not supported in a shared " + "target and draft KV cache; keeping a separate draft KV cache.", + key="dflash_per_layer_kv_heads_fallback") + return False + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._unification_max_seq_len, + self._kv_cache_config), + self._unification_max_seq_len, + range(draft_model_config.pretrained_config.num_hidden_layers)) + if len(set(draft_windows)) > 1: + logger.info_once( + "DFlash: draft layers have different effective attention " + "windows; keeping a separate draft KV cache.", + key="dflash_mixed_draft_windows_fallback") + return False + if self._is_disagg and os.environ.get( + "TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": + return False + if self._is_disagg and not self._dflash_disagg_layouts_compatible( + target_layout, draft_layout): + logger.info_once( + "DFlash: draft KV layout (kv_heads, head_dim, dtype, k_only)=" + f"{draft_layout} cannot be transferred with the target's " + f"{target_layout}; disaggregated serving keeps a separate " + "draft KV cache.", + key="dflash_disagg_draft_layout_fallback") + return False + return True + + @staticmethod + def _dflash_disagg_layouts_compatible( + target_layout: tuple[int, int, object, bool], + draft_layout: tuple[int, int, object, bool]) -> bool: + """Whether the transceiver can move draft layers with the target's. + + Heads are resharded in bytes: a view's bytes per head are its own + per-layer bytes over the rank's KV-head count, so dtype and head_dim + may differ. The KV-head count and K-only flag come from the rank's + single AttentionInfo and must match. NVFP4 adds block-scale views, so + an NVFP4 side is only paired with an identical layout. + """ + target_heads, _, target_dtype, target_k_only = target_layout + draft_heads, _, draft_dtype, draft_k_only = draft_layout + if (target_heads, target_k_only) != (draft_heads, draft_k_only): + return False + nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 + return target_dtype == draft_dtype or nvfp4 not in (target_dtype, + draft_dtype) + + def _dflash_kv_layout( + self, model_config: ModelConfig, + fallback_dtype) -> Optional[tuple[int, int, object, bool]]: + """Rank-local (kv_heads, head_dim, dtype, k_only) of a model's KV layers. + + Returns None when the KV-head count varies per layer. + """ + config = model_config.pretrained_config + k_only = is_mla(config) + heads = 1 if k_only else getattr(config, "num_key_value_heads", + config.num_attention_heads) + if not isinstance(heads, int): + return None + head_dim = (config.kv_lora_rank + config.qk_rope_head_dim + if k_only else getattr(config, "head_dim", None)) + if not isinstance(head_dim, int): + head_dim = config.hidden_size // config.num_attention_heads + dtype = self._layer_spec_dtype(model_config, fallback_dtype) + # Attention DP replicates attention weights, so every rank holds all heads. + tp_size = 1 if self._mapping.enable_attention_dp else self._mapping.tp_size + return ceil_div(heads, tp_size), head_dim, dtype, k_only + + @staticmethod + def _layer_spec_dtype(model_config: ModelConfig, fallback_dtype) -> object: + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( + ): + return tensorrt_llm.bindings.DataType.FP8 + if quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache( + ): + return tensorrt_llm.bindings.DataType.NVFP4 + return str_dtype_to_binding(torch_dtype_to_str(fallback_dtype)) + + @staticmethod + def _layer_spec_roles(dtype, *, k_only: bool) -> tuple: + roles = [Role.KEY] + if not k_only: + roles.append(Role.VALUE) + if dtype == tensorrt_llm.bindings.DataType.NVFP4: + roles.append(Role.KEY_BLOCK_SCALE) + if not k_only: + roles.append(Role.VALUE_BLOCK_SCALE) + return tuple(roles) + + def _build_unified_dflash_layer_specs( + self, + kv_cache_config: KvCacheConfig) -> tuple[KVCacheLayerSpec, ...]: + """Build the rank-local target+DFlash storage layout used by cost and allocation.""" + if not self._should_unify_dflash_kv_cache(): + return () + + target_model_config = self._model_engine.model.model_config + target_config = target_model_config.pretrained_config + target_layers = int(target_config.num_hidden_layers) + target_local_ids = list(self._mapping.pp_layers(target_layers)) + target_heads, target_head_dim, target_dtype, target_k_only = ( + self._dflash_kv_layout(target_model_config, + self._model_engine.dtype)) + target_window_config = kv_cache_config.max_attention_window + if target_window_config is None: + target_window_config = _derive_v2_layer_type_attention_windows( + kv_cache_config, + self._kv_cache_manager_cls, + target_model_config, + self._max_seq_len, + ) + target_windows = _resolve_v2_max_attention_window_vec( + target_window_config, + self._max_seq_len, + target_local_ids, + ) + + draft_model_config = self._get_draft_kv_model_config() + draft_config = draft_model_config.pretrained_config + draft_layers = int(draft_config.num_hidden_layers) + draft_local_ids = list(self._mapping.pp_layers(draft_layers)) + draft_heads, draft_head_dim, draft_dtype, draft_k_only = ( + self._dflash_kv_layout(draft_model_config, + draft_config.torch_dtype)) + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._max_seq_len, + kv_cache_config), + self._max_seq_len, + draft_local_ids, + ) + + specs = [ + KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=target_dtype, + num_kv_heads=target_heads, + head_dim=target_head_dim, + buffer_roles=self._layer_spec_roles(target_dtype, + k_only=target_k_only), + sliding_window_size=window, + reuse_policy=AttentionReusePolicy.REQUIRED, + ) for layer_id, window in zip(target_local_ids, target_windows) + ] + specs.extend( + KVCacheLayerSpec( + global_layer_id=target_layers + layer_id, + owner=KVCacheLayerOwner.DRAFT, + dtype=draft_dtype, + num_kv_heads=draft_heads, + head_dim=draft_head_dim, + buffer_roles=self._layer_spec_roles(draft_dtype, + k_only=draft_k_only), + sliding_window_size=window, + reuse_policy=AttentionReusePolicy.REQUIRED, + ) for layer_id, window in zip(draft_local_ids, draft_windows)) + return tuple(specs) + def _joint_reuse_supported(self) -> bool: """Span half of the pairing decision; ``build_managers`` ANDs it with ``has_separate_one_model_draft``. ``False`` = leave the run unpaired. @@ -2037,12 +2246,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 +3326,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 +3503,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..634ffa5bb29c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -66,6 +66,7 @@ DEFAULT_BEAM_INDEX, GPU_LEVEL, AttentionLayerConfig, + AttentionReusePolicy, AttnLifeCycle, Batch, BatchDesc, @@ -272,6 +273,99 @@ 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 + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + + def __post_init__(self) -> None: + if self.global_layer_id < 0: + raise ValueError(f"global_layer_id must be non-negative, got {self.global_layer_id}") + if self.num_kv_heads <= 0: + raise ValueError(f"num_kv_heads must be positive, got {self.num_kv_heads}") + if self.head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {self.head_dim}") + if not self.buffer_roles: + raise ValueError("buffer_roles must not be empty") + if len(set(self.buffer_roles)) != len(self.buffer_roles): + raise ValueError( + f"layer {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 _swa_endpoint_priority( prompt_length: int, tokens_per_block: int, rewind_tokens: int ) -> Callable[[int, object], int]: @@ -507,26 +601,34 @@ def _estimate_swa_cache_size( context: bool, scratch: bool, generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> tuple[int, int]: tokens_per_block = int(tokens_per_block) size_per_token = 0 size_per_request = 0 scratch_keys = set() - for layer_size, window_size in zip(layer_sizes, attention_windows): + for layer_idx, (layer_size, window_size) in enumerate(zip(layer_sizes, attention_windows)): if window_size is not None and window_size > 0: + layer = layer_specs[layer_idx] if layer_specs is not None else None + sink_blocks = math.ceil((layer.num_sink_tokens or 0) / tokens_per_block) if layer else 0 # Match AttnLifeCycle.get_stale_range(): the live interval contains # window_size + generation_capacity_headroom - 1 tokens. Across all # page offsets, that interval touches at most the count below. window_blocks = ( math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 ) - window_tokens = window_blocks * tokens_per_block + window_tokens = (sink_blocks + window_blocks) * tokens_per_block if not context: size_per_request += window_tokens * layer_size elif not scratch: size_per_token += layer_size else: - scratch_key = (int(window_size), layer_size) + scratch_key = ( + int(window_size), + layer_size, + sink_blocks, + layer.reuse_policy if layer else None, + ) if scratch_key in scratch_keys: size_per_request += window_tokens * layer_size else: @@ -563,6 +665,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 +675,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 +689,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 +876,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. @@ -1265,10 +1398,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 +1412,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 +1450,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) @@ -1419,6 +1581,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 +1654,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 +1689,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 +1854,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, @@ -2192,7 +2375,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 +2415,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 +2459,14 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] + uses_block_scale_pointers = any( + self._get_layer_dtype_by_local_id(layer_id) == DataType.NVFP4 + for layer_id in range(self.num_local_layers) + ) if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): pool_id = self.impl.get_layer_group_id(layer_id) - role_a, _ = self._get_pool_roles(pool_id) + role_a, _ = self._get_layer_page_table_roles(layer_id, pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( @@ -2259,7 +2475,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 +2496,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 +2513,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 +2539,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 +2565,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 +2576,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,6 +2616,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: > 1 for pool_id, layer_ids in enumerate(self.impl.layer_grouping) ) + if self.layer_specs is not None: + storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} + self._use_per_layer_page_tables |= len(storage_layouts) > 1 if self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: @@ -2457,11 +2679,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 +2725,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 +2762,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 +2922,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, @@ -2882,6 +3111,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) @@ -2975,43 +3206,70 @@ def _build_base_config( ) ) - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional - # per-layer BufferConfig entries — for example a sparse index-K - # buffer — without overriding the K/V/NVFP4 scale wiring above. - # The dict maps local layer id -> list of extra BufferConfig. Each - # extra buffer's role must be unique within the layer (asserted by - # AttentionLayerConfig.__post_init__) and its size must be in bytes - # per block (= bytes_per_token * tokens_per_block). - extra_buffers_per_layer = ( - self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} - ) - - layer_configs: List[AttentionLayerConfig] = [] - for layer_id in typed_range(LayerId(self.num_local_layers)): - buffer_type = self._get_buffer_roles_for_layer(layer_id) - buffers = [ - BufferConfig( - role=role, - size=self.get_layer_bytes_per_token(local_layer_idx=layer_id, data_role=role) - * tokens_per_block, - ) - for role in buffer_type - ] - for extra in extra_buffers_per_layer.get(int(layer_id), ()): - assert extra.role not in buffer_type, ( - f"extra buffer role {extra.role!r} for layer " - f"{int(layer_id)} duplicates a standard K/V/scale role" - ) - buffers.append(extra) - layer_configs.append( + 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. + coalescing_group=( + 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, + reuse_policy=layer.reuse_policy, ) + for local_layer_id, layer in enumerate(self.layer_specs) + ] + else: + # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional + # per-layer BufferConfig entries — for example a sparse index-K + # buffer — without overriding the K/V/NVFP4 scale wiring above. + # The dict maps local layer id -> list of extra BufferConfig. Each + # extra buffer's role must be unique within the layer (asserted by + # AttentionLayerConfig.__post_init__) and its size must be in bytes + # per block (= bytes_per_token * tokens_per_block). + extra_buffers_per_layer = ( + self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} ) + layer_configs: List[AttentionLayerConfig] = [] + for layer_id in typed_range(LayerId(self.num_local_layers)): + buffer_type = self._get_buffer_roles_for_layer(layer_id) + buffers = [ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token( + local_layer_idx=layer_id, data_role=role + ) + * tokens_per_block, + ) + for role in buffer_type + ] + for extra in extra_buffers_per_layer.get(int(layer_id), ()): + assert extra.role not in buffer_type, ( + f"extra buffer role {extra.role!r} for layer " + f"{int(layer_id)} duplicates a standard K/V/scale role" + ) + buffers.append(extra) + layer_configs.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=buffers, + sliding_window_size=self.max_attention_window_vec[layer_id], + num_sink_tokens=None, + ) + ) + return KVCacheManagerConfigPy( # Used by the backend only for token<->block arithmetic and # radix hashing; BufferConfig.size above stays the physical page. @@ -3071,6 +3329,7 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa buffers=buffers, sliding_window_size=layer.sliding_window_size, num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, ) ) else: @@ -3167,6 +3426,64 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + 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 +3499,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 +3515,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 +3523,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, @@ -5379,10 +5712,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 @@ -5693,6 +6033,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 +6050,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 +6061,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 +6087,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 +6141,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 +6153,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 +6175,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 +6192,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 +6260,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 +6304,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 +6377,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 @@ -6027,8 +6427,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 +6441,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 +6496,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 +6525,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 +6534,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 +6571,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 @@ -6174,7 +6609,7 @@ def _get_per_request_layer_tokens( * tokens_per_block ) - def window_tokens(window_size: Optional[int]) -> int: + def window_tokens(window_size: Optional[int], sink_tokens: int) -> int: if window_size is None or window_size <= 0 or window_size >= max_seq_len: return full_tokens # Match _estimate_swa_cache_size: the live interval holds @@ -6182,15 +6617,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 +6647,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 +6681,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 +6711,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 +6747,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 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/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 613071340c06..b3f137e70953 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 @@ -617,19 +618,29 @@ def prepare(self): # prefilling, so _store_prefill_context -- the only other assigner -- # never runs and every request would land on the dummy slot. num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + assigned = {} + for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=assigned) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + elif worker._ctx_unified_cache: + # New requests have no in-flight predecessor. Subsequent + # accepted tokens advance the device length during replay. + token_offset = sum(self.seq_lens[:request_idx]) + assigned[slot] = min( + int(self.host_position_ids[token_offset]), worker._max_ctx + ) + worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context @@ -777,6 +788,9 @@ def __init__( # the table width on the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 + self._ctx_layer_ids = None + self._ctx_unified_cache = False + self._ctx_block_offset = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None self._ctx_kv_last_page_len = None @@ -939,7 +953,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 +997,40 @@ def _init_ctx_block_tables( """ mgr = draft_kv_cache_manager pool_mapping = mgr.kv_cache_pool_mapping - layer_offset = mgr.layer_offsets[0] + first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 + layer_offset = mgr.layer_offsets[first_layer_id] self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - if num_pool_layers != L: + self._ctx_block_offset = 0 + if isinstance(mgr, KVCacheManagerV2): + converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) + layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + if mgr.enable_swa_scratch_reuse: + raise NotImplementedError( + "DFlash managed context does not support SWA scratch pages" + ) + for layer_id in self._ctx_layer_ids: + local_id = mgr.layer_offsets[layer_id] + if ( + mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id + or mgr.get_layer_page_index_scale(layer_id) != converter.scale + ): + if self._ctx_unified_cache: + raise NotImplementedError( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + logger.warning( + "DFlash: draft KV layers have different lifecycles or page-index " + "scales; using the private ctx arena." + ) + return False + if per_layer_tables: + # Attention tables address the physical pool origin, whereas + # get_buffers() starts at this layer's own K buffer. + 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 +1039,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 " @@ -1022,20 +1088,24 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """ if self._ctx_block_tables is None or num_seqs <= 0: return False - src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + src = getattr(attn_metadata, offsets_name, None) if src is None: # The table is bound but unfillable. Leaving it at zeros would send # every request's every layer to pool block 0 -- silently. raise RuntimeError( "DFlash: bound to the draft KV pool but " f"{type(attn_metadata).__name__} carries no " - "draft_kv_cache_block_offsets this iteration." + f"{offsets_name} this iteration." ) # Field 0 of the K/V pair: both name the same block for a kv_factor=2 # pool, and the drafter addresses blocks, not planes. # clone(): .to() is a no-op for an int64 source, and the in-place ops # below would then edit attn_metadata's own tensor. encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() + encoded.sub_(self._ctx_block_offset) # Before the clamp erases the placeholders. Informative only on the # python V2 backend; the cpp one already mapped BAD_PAGE_INDEX to 0. self._ctx_block_counts[:num_seqs].copy_((encoded >= 0).sum(dim=1)) @@ -1047,12 +1117,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." @@ -1161,6 +1237,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 +1292,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 +1324,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( @@ -1779,10 +1858,18 @@ def _forward_impl( # Lazy init buffers and attach worker reference for prepare() draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + previous_ctx_len = self._ctx_len self._lazy_init_ctx_buffers( draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager ) spec_metadata._dflash_worker = self + if ( + self._ctx_unified_cache + and self._ctx_len is not previous_ctx_len + and spec_metadata.request_ids + ): + # New buffers clear request slots; seed them before generation. + spec_metadata.prepare() # Before any store: prefill and decode both address pages through it. # Returning False here means an empty batch -- the missing-offsets case # raises inside, with the metadata type in the message. @@ -2306,10 +2393,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) 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..6e6568afccc6 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -804,7 +804,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..fbdea420106a 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2486,6 +2486,19 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, from tensorrt_llm.llmapi.disagg_utils import split_world_comm disagg_cfg = parse_disagg_config_file(config_file) + with_spec = { + cfg.type + for cfg in disagg_cfg.server_configs + if cfg.other_args.get("speculative_config") + } + attention_parallel_configs = {(cfg.other_args.get("tensor_parallel_size", + 1), + cfg.other_args.get("enable_attention_dp", + False)) + for cfg in disagg_cfg.server_configs} + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ( + "1" if with_spec == {"ctx", "gen"} + and len(attention_parallel_configs) == 1 else "0") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. 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..5dc41d1bf01f 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,95 @@ 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, + }, + # 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, + extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}, + assert_worker_log_contains= + "target and draft KV layers share one KV cache manager") as llm: + llm.args.speculative_config = DFlashDecodingConfig( + max_draft_len=self.MAX_DRAFT_LEN, + speculative_model=self.DRAFT_PATH, + ) + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + + # Drafting from an empty prompt context costs ~17% acceptance + # length on this model, well below the 95% floor. + acceptance_length = compute_disagg_acceptance_length(llm.serve_url) + print(f"[AL] test_gsm8k_1p1d_unified_draft_kv " + f"acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length( + "TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv", + acceptance_length) + + @skip_pre_blackwell @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) class TestQwen3_30B_A3B(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 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/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..c91507098ac3 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Selection of the unified target+DFlash KV cache layout.""" + +from types import SimpleNamespace + +import pytest +import torch +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.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): + creator = object.__new__(KvCacheCreator) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), + attention_backend="TRTLLM", + _allow_separate_draft_kv_cache=True, + draft_len_schedule=None, + ) + creator._draft_model_engine = None + creator._draft_config = draft + creator._mapping = SimpleNamespace( + pp_size=1, cp_size=1, tp_size=tp_size, enable_attention_dp=False + ) + creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) + creator._model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 + ) + creator._kv_cache_manager_cls = KVCacheManagerV2 + creator._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() + + +def test_mismatched_draft_kv_heads_unify_only_in_aggregated_serving() -> None: + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + assert _creator(is_disagg=False, target=target, draft=draft)._should_unify_dflash_kv_cache() + assert not _creator(is_disagg=True, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_mismatched_draft_head_dim_unifies(is_disagg: bool) -> None: + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=64) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_fp8_target_with_bf16_draft_unifies(is_disagg: bool) -> None: + target = _model_config(kv_heads=8, head_dim=128, kv_quant_algo=QuantAlgo.FP8) + draft = _model_config(kv_heads=8, head_dim=128) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +def test_layout_compares_rank_local_heads(): + # 8 target heads and 2 draft heads differ per rank at every TP size; at + # TP4 the draft is duplicated to 1 head per rank while the target has 2. + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=2, head_dim=128) + creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=4) + + assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 2 + assert creator._dflash_kv_layout(draft, torch.bfloat16)[0] == 1 + assert not creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_attention_dp_unifies_with_unsharded_heads(is_disagg): + # Under attention DP each rank keeps every KV head, so the TP4 layouts of + # an 8-head target and a 2-head draft are not resharded to 2 and 1. + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=is_disagg, target=target, draft=draft, tp_size=4) + creator._mapping.enable_attention_dp = True + + assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 8 + assert ( + creator._dflash_kv_layout(_model_config(kv_heads=2, head_dim=128), torch.bfloat16)[0] == 2 + ) + assert creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("per_layer_side", ["target", "draft"]) +def test_per_layer_kv_heads_keep_separate_cache(is_disagg, per_layer_side): + uniform = _model_config(kv_heads=8, head_dim=128) + per_layer = _model_config(kv_heads=[8, 4], head_dim=128) + target, draft = (per_layer, uniform) if per_layer_side == "target" else (uniform, per_layer) + + creator = _creator(is_disagg=is_disagg, target=target, draft=draft) + + assert not creator._should_unify_dflash_kv_cache() + assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () + + +class _StopLauncher(Exception): + pass + + +def _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers): + """Run the real disaggregated_mpi_worker up to the rank-communicator split.""" + from tensorrt_llm.commands import serve + + config = tmp_path / "disagg.yaml" + config.write_text( + yaml.safe_dump( + { + "context_servers": {"num_instances": 1, **context_servers}, + "generation_servers": {"num_instances": 1, **generation_servers}, + } + ) + ) + + def stop(*args, **kwargs): + raise _StopLauncher + + monkeypatch.setattr( + serve._command_telemetry, "apply_disaggregated_telemetry_config", lambda *a, **k: None + ) + monkeypatch.setattr(serve, "set_cuda_device", lambda: None) + monkeypatch.setattr("tensorrt_llm.llmapi.disagg_utils.split_world_comm", stop) + monkeypatch.delenv("TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT", raising=False) + with pytest.raises(_StopLauncher): + serve.disaggregated_mpi_worker.callback( + config_file=str(config), log_level="info", telemetry=False + ) + + +def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_path): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={}, + generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, + ) + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + + for _role in ("context", "generation"): + creator = _creator(is_disagg=True, target=target, draft=draft) + assert not creator._should_unify_dflash_kv_cache() + + +def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): + spec = {"decoding_type": "DFlash"} + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, + generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, + ) + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + for tp_size in (4, 8): + creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=tp_size) + assert not creator._should_unify_dflash_kv_cache() 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..2d0a73e58826 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -10,7 +10,7 @@ share, not all copies. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from unittest.mock import Mock, patch @@ -30,9 +30,15 @@ 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, @@ -41,6 +47,7 @@ TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy pytestmark = pytest.mark.cpu_only @@ -953,8 +960,70 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): assert cost == CacheCost(slope=20, intercept=21) +def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: + windowed = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=128, + num_sink_tokens=65, + ) + layers = [ + windowed, + replace(windowed, global_layer_id=1, reuse_policy=AttentionReusePolicy.PRIVATE), + replace(windowed, global_layer_id=2, num_sink_tokens=None), + replace(windowed, global_layer_id=3, sliding_window_size=None, num_sink_tokens=None), + ] + config = KvCacheConfig(enable_block_reuse=False, enable_swa_scratch_reuse=True) + kwargs = dict( + tokens_per_block=64, + max_seq_len=4096, + kv_cache_config=config, + layer_specs=layers, + ) + slope, fixed = KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=256, **kwargs + ) + # Each layer costs 64 bytes/token. Different sinks or policies cannot share scratch. + # The SWA layers retain 5, 5, and 3 pages, including two pages for 65 sink tokens. + assert slope == 64 + assert fixed == 256 * 192 + 2 * (320 + 320 + 192) * 64 + + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = tuple(layers) + manager._has_cp_helix = False + manager.max_attention_window_vec = [128, 128, 128, None] + manager.tokens_per_block = 64 + manager.max_batch_size = 2 + manager.max_num_tokens = 256 + manager._generation_kv_capacity_headroom = 1 + manager.enable_swa_scratch_reuse = True + quota = manager._get_quota_from_max_tokens(512) + assert quota == slope * 512 + fixed + assert manager._get_max_tokens_from_quota(quota) == 512 + assert KVCacheManagerV2._get_per_request_layer_tokens(None, Mapping(), **kwargs) == ( + [(64, 320, 320), (64, 320, 320), (64, 192, 192), (64, 4160, 0)], + 4160, + ) + + # Draft lifecycles need 7 and 21 usable slots: separate resume rounding + # gives 8 + 23 slots; rounding their average together would give only 30. + draft = replace(windowed, owner=KVCacheLayerOwner.DRAFT, num_sink_tokens=0) + kwargs.update( + layer_specs=[draft, replace(draft, global_layer_id=1, num_sink_tokens=448)], + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_util_for_resume=0.95), + ) + assert KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=64, is_draft=True, **kwargs + ) == (0, (8 + 23) * 64 * 64) + + def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = None manager._has_cp_helix = False manager.num_local_layers = 3 manager.pp_layers = [4, 5, 6] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 9ec9df759125..6dcbcd10a5e3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -28,6 +28,8 @@ from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, @@ -41,6 +43,7 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE +from tensorrt_llm.bindings.internal import thop from tensorrt_llm.bindings.internal.batch_manager import CacheType, LinearCacheType from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( @@ -52,6 +55,7 @@ PARDDecodingConfig, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization import QuantMode from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, @@ -70,6 +74,7 @@ KVCacheManager, KVCacheManagerConfig, LayerId, + PageIndexMode, ReuseScope, SsmLayerConfig, _introspection, @@ -309,6 +314,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 +359,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 +424,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 +636,308 @@ 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() + + +def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + torch.cuda.init() + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=2, + head_dim=8, + buffer_roles=(Role.KEY, Role.VALUE), + ) + # Equal byte sizes coalesce the FP8 target and BF16 draft pages, placing + # the draft's K buffer at a nonzero offset within each physical slot. + draft = replace( + target, + global_layer_id=1, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=1, + ) + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=1, + num_kv_heads=target.num_kv_heads, + head_dim=target.head_dim, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + layer_specs=[target, draft], + ) + try: + assert manager._use_per_layer_page_tables + request_ids = [7, 9] + for request_id, num_blocks in zip(request_ids, (2, 3)): + cache = manager._create_kv_cache(request_id, None, None) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(num_blocks * TOKENS_PER_BLOCK, 0) + + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + pool = worker._managed_ctx_pool( + manager, 1, draft.num_kv_heads, draft.head_dim, torch.bfloat16, 2 + ) + assert pool is not None + assert worker._init_ctx_block_tables( + manager, pool, len(request_ids), 1, draft.num_kv_heads, draft.head_dim, TOKENS_PER_BLOCK + ) + assert worker._ctx_block_offset > 0 + + block_offsets = torch.empty( + (manager.num_attention_op_pools, len(request_ids), 2, manager.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + manager.copy_batch_block_offsets( + block_offsets, request_ids, beam_width=1, num_contexts=0, num_seqs=len(request_ids) + ) + metadata = SimpleNamespace(kv_cache_block_offsets=block_offsets) + assert worker._refresh_ctx_block_tables(metadata, len(request_ids)) + + expected = manager.get_batch_cache_indices(request_ids, layer_idx=draft.global_layer_id) + for row, indices in enumerate(expected): + assert len(indices) == (2, 3)[row] + assert worker._ctx_block_tables[row, : len(indices)].tolist() == indices + finally: + manager.shutdown() + + +def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + 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 +1607,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 +3292,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 +3588,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} 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/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 From 6b40260a15998983171c2ef1ac6187f50b684381 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 20 Oct 2026 23:22:07 -0700 Subject: [PATCH 02/23] [None][feat] Per View KV head transfer (port of draft NVIDIA/TensorRT-LLM#20709) Squashed increment of draft PR #20709 at head 170385120 over #20486 at 714916dc1 (both on main 0829689c5; the #20709 tree contains all of #20486), ported onto feat/musespark-open-agentx. Commits: 7a85a733a [None][feat] unify the standalone DSpark target and draft KV cache d492b5267 [None][feat] reshard unified target/draft KV by per-view attention heads Port conflict resolutions (base = #20486 on main 0829689c5, ours = the #20486 port): - kv_extractor.py: agentx lacks main's view_partition hook (#19727); add only kv_heads_for_layer and the Callable import its annotation needs (Muse target and draft buffers differ in size, so they never share a physical pool/view). - kv_cache_manager_v2.py add_dummy_requests: keep agentx's stop_committing() for the dummy and size it with the PR's _kv_alloc_extra_tokens. - test_kvcm2_integration.py: keep agentx's DecoderRunner import and add the PR's tests only (main's generation-block reuse tests, #19710, are not on agentx). Signed-off-by: Allison Lim Signed-off-by: Zheyu Fu --- docs/source/features/speculative-decoding.md | 29 ++- .../native/mixers/attention/peer.py | 113 ++++++--- .../_torch/disaggregation/native/peer.py | 214 +++++++++++++++--- .../_torch/disaggregation/native/transfer.py | 12 +- .../disaggregation/resource/kv_extractor.py | 28 ++- .../_torch/disaggregation/resource/page.py | 12 +- tensorrt_llm/_torch/pyexecutor/_util.py | 188 +++++++++++---- .../kv_cache/kv_cache_manager_v2.py | 92 ++++++-- tensorrt_llm/_torch/speculative/__init__.py | 10 +- tensorrt_llm/_torch/speculative/dflash.py | 65 ++++-- tensorrt_llm/_torch/speculative/dspark.py | 14 +- tensorrt_llm/_torch/speculative/utils.py | 15 ++ tensorrt_llm/commands/serve.py | 40 ++-- .../accuracy/test_disaggregated_serving.py | 4 +- .../kv_cache/test_dflash_unified_kv_gate.py | 142 +++++++++--- .../kv_cache/test_kv_cache_budget_split.py | 3 + .../kv_cache/test_kv_cache_estimation.py | 53 +++++ .../kv_cache/test_kvcm2_integration.py | 181 +++++++++++++++ .../test_dflash_ctx_block_tables.py | 98 ++++++++ 19 files changed, 1119 insertions(+), 194 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 718c6f397767..a4fb3713758d 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 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,33 @@ 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. 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/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 b2d7451a6dd6..c958d2be9a4b 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,13 +185,15 @@ 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) if self._ri.num_draft_layers != peer_ri.num_draft_layers: # The layer-count tolerance below would transfer only the target # layers and leave the drafter's prompt KV unset on the receiver. raise ValueError( "PeerRegistrar: unified draft KV layer count differs " f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " - "both peers must use the same speculative decoding KV cache layout." + "both peers must use the same speculative decoding KV cache layout; " + "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." ) self_layers = sum(self._ri.layer_num_per_pp) @@ -193,6 +210,61 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: return True + 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). @@ -294,7 +366,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 @@ -308,11 +380,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 @@ -534,6 +647,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, @@ -542,10 +656,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. @@ -587,7 +726,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/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 3c7c4a7f45b6..df3524ab8a65 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3138,8 +3138,9 @@ 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, 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 @@ -3147,7 +3148,12 @@ def _fanin_bounce_safe( 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: + # Any duplicated head, rank-wide or in a single view pair, leaves some + # writers without a share (sender side) or doubles one (receiver side). + view_factors = getattr(overlap, "view_duplicate_head_factors", {}).values() + if overlap.duplicate_head_factor != 1 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) 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 f81e554f802d..b28f9f10db53 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -690,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, reuse policy), its slot sizes are the sizes + of its coalesced buffers in descending order, and life cycles with the same + (reuse policy, slot sizes) share one pool group. + """ + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 + life_cycles: dict[tuple, dict[tuple[int, int], int]] = {} + for layer in layer_specs: + life_cycle = (layer.sliding_window_size, + ceil_div(layer.num_sink_tokens or 0, + tokens_per_block), layer.reuse_policy) + coalesced = life_cycles.setdefault(life_cycle, {}) + for role in layer.buffer_roles: + group = 1 + int(role in scale_roles) if layer.dtype == nvfp4 else 0 + key = (group, layer.bytes_per_token(role) * tokens_per_block) + coalesced[key] = coalesced.get(key, 0) + 1 + return len({(life_cycle[2], + tuple( + sorted((size * count + for (_, size), count in coalesced.items()), + reverse=True))) + for life_cycle, coalesced in life_cycles.items()}) + + def _get_num_pool_groups_for_estimation( model_config: object, max_seq_len: int, @@ -1496,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 @@ -1898,12 +1935,14 @@ def _create_kv_cache_manager( spec_dec_layer_mask = [True] * num_target_layers estimating_kv_cache = estimating_kv_cache and not self._skip_est + self._check_shared_draft_kv_request() unified_layer_specs = None if self._should_unify_dflash_kv_cache(): unified_layer_specs = self._build_unified_dflash_layer_specs( kv_cache_config) logger.info( - "DFlash: target and draft KV layers share one KV cache manager " + f"{self._unified_drafter_name()}: target and draft KV layers " + "share one KV cache manager " f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})." ) kv_cache_manager = _create_kv_cache_manager( @@ -1997,24 +2036,85 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) - def _should_unify_dflash_kv_cache(self) -> bool: - """Whether this run can place target and DFlash layers in one V2 manager.""" + def _unified_drafter_attention_backend(self) -> Optional[str]: + """Resolved block-attention backend of a drafter that reads paged draft KV. + + None for every other drafter. That includes the embedded DSpark draft, + whose context lives in a worker-owned rolling window. + """ spec_config = self._speculative_config - if (spec_config is None or not spec_config.spec_dec_mode.is_dflash() - or spec_config.attention_backend != "TRTLLM" - or not spec_config._allow_separate_draft_kv_cache - or spec_config.draft_len_schedule is not None - or self._draft_model_engine is not None - or self._draft_config is None or self._mapping.pp_size != 1 - or self._mapping.cp_size != 1 - or self._llm_args.attn_backend != "TRTLLM" or is_gemma4_hybrid( - self._model_engine.model.model_config.pretrained_config) - or self._llm_args.kv_cache_compression_config is not None): + spec_dec_mode = spec_config.spec_dec_mode + if spec_dec_mode.is_dflash(): + return spec_config.attention_backend + if (not spec_dec_mode.is_dspark() + or spec_config.draft_is_embedded_in_target): + return None + backend = spec_config.attention_backend + if backend == "AUTO": + # The drafter resolves AUTO against its own family and this build. + draft_model = getattr(self._model_engine.model, "draft_model", None) + backend = getattr(draft_model, "dflash_attention_backend", None) + return backend + + def _unified_drafter_name(self) -> str: + return ("DSpark" if self._speculative_config.spec_dec_mode.is_dspark() + else "DFlash") + + def _should_unify_dflash_kv_cache(self) -> bool: + """Whether this run can place target and DFlash or standalone DSpark + layers in one V2 manager. + + Disaggregated serving shares only when TRTLLM_DISAGG_SHARED_DRAFT_KV=1. + """ + if self._dflash_unify_blocker() is not None: return False + return not (self._is_disagg + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1") + + def _check_shared_draft_kv_request(self) -> None: + """Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored.""" + if (not self._is_disagg + or os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1" + or self._speculative_config is None + or self._unified_drafter_attention_backend() is None): + return + blocker = self._dflash_unify_blocker() + if blocker is not None: + raise ValueError( + "TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires the target and draft " + f"KV layers to share one KV cache manager, but {blocker}.") + + def _dflash_unify_blocker(self) -> Optional[str]: + """The reason target and drafter layers cannot share one manager, or None.""" + spec_config = self._speculative_config + if spec_config is None: + return "speculative decoding is not enabled" + backend = self._unified_drafter_attention_backend() + if backend != "TRTLLM": + return f"the drafter attention_backend is {backend}, not TRTLLM" + if not spec_config._allow_separate_draft_kv_cache: + return "a separate draft KV cache is not allowed for this configuration" + if spec_config.draft_len_schedule is not None: + return "draft_len_schedule is set" + if self._draft_model_engine is not None: + return "the drafter runs as a separate draft model engine" + if self._draft_config is None: + return "there is no draft model config" + if self._mapping.pp_size != 1: + return f"pipeline_parallel_size is {self._mapping.pp_size}, not 1" + if self._mapping.cp_size != 1: + return f"context_parallel_size is {self._mapping.cp_size}, not 1" + if self._llm_args.attn_backend != "TRTLLM": + return f"the target attn_backend is {self._llm_args.attn_backend}, not TRTLLM" + if is_gemma4_hybrid( + self._model_engine.model.model_config.pretrained_config): + return "the target uses Gemma 4 hybrid attention" + if self._llm_args.kv_cache_compression_config is not None: + return "KV cache compression is enabled" # Specialized managers add cache roles and lifecycle protocols that an # external drafter cannot yet describe. Keep their existing fallback. if self._kv_cache_manager_cls is not KVCacheManagerV2: - return False + return "the target does not use KVCacheManagerV2" target_layout = self._dflash_kv_layout( self._model_engine.model.model_config, self._model_engine.dtype) draft_model_config = self._get_draft_kv_model_config() @@ -2023,34 +2123,36 @@ def _should_unify_dflash_kv_cache(self) -> bool: draft_model_config.pretrained_config.torch_dtype) if target_layout is None or draft_layout is None: logger.info_once( - "DFlash: a per-layer KV-head count is not supported in a shared " + f"{self._unified_drafter_name()}: a per-layer KV-head count is " + "not supported in a shared " "target and draft KV cache; keeping a separate draft KV cache.", key="dflash_per_layer_kv_heads_fallback") - return False + return "the target or drafter has a per-layer KV-head count" + # The drafter's shared context block table indexes all its layers + # through one lifecycle; mixed windows need the separate manager. draft_windows = _resolve_v2_max_attention_window_vec( self._get_draft_max_attention_window(self._unification_max_seq_len, self._kv_cache_config), self._unification_max_seq_len, - range(draft_model_config.pretrained_config.num_hidden_layers)) + list( + self._mapping.pp_layers( + int(draft_model_config.pretrained_config.num_hidden_layers)) + )) if len(set(draft_windows)) > 1: - logger.info_once( - "DFlash: draft layers have different effective attention " - "windows; keeping a separate draft KV cache.", - key="dflash_mixed_draft_windows_fallback") - return False - if self._is_disagg and os.environ.get( - "TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": - return False + return "the draft layers use different attention windows" if self._is_disagg and not self._dflash_disagg_layouts_compatible( target_layout, draft_layout): logger.info_once( - "DFlash: draft KV layout (kv_heads, head_dim, dtype, k_only)=" + f"{self._unified_drafter_name()}: draft KV layout " + "(kv_heads, head_dim, dtype, k_only)=" f"{draft_layout} cannot be transferred with the target's " f"{target_layout}; disaggregated serving keeps a separate " "draft KV cache.", key="dflash_disagg_draft_layout_fallback") - return False - return True + return (f"the draft KV layout {draft_layout} cannot be transferred " + f"with the target's {target_layout} " + "(kv_heads, head_dim, dtype, k_only)") + return None @staticmethod def _dflash_disagg_layouts_compatible( @@ -2058,19 +2160,21 @@ def _dflash_disagg_layouts_compatible( draft_layout: tuple[int, int, object, bool]) -> bool: """Whether the transceiver can move draft layers with the target's. - Heads are resharded in bytes: a view's bytes per head are its own - per-layer bytes over the rank's KV-head count, so dtype and head_dim - may differ. The KV-head count and K-only flag come from the rank's - single AttentionInfo and must match. NVFP4 adds block-scale views, so - an NVFP4 side is only paired with an identical layout. + 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_heads, target_k_only) != (draft_heads, draft_k_only): + if target_k_only != draft_k_only: return False nvfp4 = tensorrt_llm.bindings.DataType.NVFP4 - return target_dtype == draft_dtype or nvfp4 not in (target_dtype, - draft_dtype) + 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, 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 634ffa5bb29c..7c9452fe3ff2 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 @@ -637,10 +637,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 @@ -653,7 +679,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 @@ -1360,6 +1388,18 @@ class KVCacheManagerV2(BaseResourceManager): # True on hybrid Mamba/GDN managers; __init__ then coerces a positive # recompute tail to a full re-prefill. _has_recurrent_state: bool = False + # Block width of a drafter whose layers share this manager, else 0. + _shared_draft_block_width: int = 0 + + @property + def _kv_alloc_extra_tokens(self) -> int: + """Tokens a request is allocated beyond its history for spec decoding. + + ``num_extra_kv_tokens`` is the position-side count; a shared drafter's + block must also fit in the allocation. + """ + return max(self.num_extra_kv_tokens, self._shared_draft_block_width) + # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1541,6 +1581,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 @@ -1558,7 +1601,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 @@ -2231,7 +2276,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() @@ -3178,7 +3223,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)] @@ -3199,7 +3244,7 @@ def _build_base_config( BatchDesc( [ KVCacheDesc( - capacity=self.max_num_tokens + self.num_extra_kv_tokens, + capacity=self.max_num_tokens + self._kv_alloc_extra_tokens, history_length=0, ) ] @@ -3426,6 +3471,21 @@ 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: @@ -3694,7 +3754,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 = ( @@ -4191,7 +4251,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 @@ -4231,7 +4291,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: @@ -4262,7 +4322,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 @@ -4864,7 +4924,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 @@ -5008,7 +5068,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( @@ -5852,7 +5912,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 @@ -6413,7 +6473,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, @@ -6602,7 +6662,7 @@ def _get_per_request_layer_tokens( else: retained_windows = list(covered_windows) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft + spec_config, is_draft=is_draft, layer_specs=layer_specs ) full_tokens = ( math.ceil((max_seq_len + generation_capacity_headroom) / tokens_per_block) diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 382cb09e1571..c3ffe7e28bda 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -18,10 +18,11 @@ SpecSampler) from .spec_tree_manager import SpecTreeManager from .suffix_automaton import SuffixAutomatonManager -from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, get_spec_drafter, - get_spec_metadata, get_spec_resource_manager, - get_spec_worker, update_spec_config_from_draft_model_config, +from .utils import (get_draft_block_width, get_draft_kv_cache_manager, + get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, get_spec_drafter, get_spec_metadata, + get_spec_resource_manager, get_spec_worker, + update_spec_config_from_draft_model_config, update_spec_config_from_loaded_model, update_spec_config_from_model_config) @@ -50,6 +51,7 @@ "SpecWorkerBase", "draft_prompt_lookahead", "get_draft_kv_cache_manager", + "get_draft_block_width", "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index b3f137e70953..57fe4516c69f 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -643,9 +643,16 @@ def prepare(self): worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) - # to dummy slot to avoid corrupting real request's context + # to dummy slot to avoid corrupting real request's context. A + # generation row with the idle id is a dummy even when a context + # warmup stored a prompt under that id: its length is not this row's. mapping = torch.tensor( - [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], + [ + worker._dummy_slot + if i >= num_contexts and rid == ATTENTION_DP_DUMMY_REQUEST_ID + else worker._req_to_slot.get(rid, worker._dummy_slot) + for i, rid in enumerate(self.request_ids) + ], dtype=torch.long, device="cpu", pin_memory=prefer_pinned(), @@ -1006,26 +1013,28 @@ def _init_ctx_block_tables( if isinstance(mgr, KVCacheManagerV2): converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + unsupported = None if mgr.enable_swa_scratch_reuse: - raise NotImplementedError( - "DFlash managed context does not support SWA scratch pages" - ) + unsupported = "DFlash managed context does not support SWA scratch pages" for layer_id in self._ctx_layer_ids: local_id = mgr.layer_offsets[layer_id] if ( mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id or mgr.get_layer_page_index_scale(layer_id) != converter.scale ): - if self._ctx_unified_cache: - raise NotImplementedError( - "DFlash's shared context block table requires draft layers " - "with the same lifecycle and page-index scale" - ) - logger.warning( - "DFlash: draft KV layers have different lifecycles or page-index " - "scales; using the private ctx arena." + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" ) - return False + break + if unsupported is not None: + # A unified manager was selected for this layout, so the drafter + # cannot move to private storage without breaking allocation, + # reuse and transfer. A separate manager keeps the private arena. + if self._ctx_unified_cache: + raise NotImplementedError(unsupported) + logger.warning(f"{unsupported}; using the private ctx arena.") + return False if per_layer_tables: # Attention tables address the physical pool origin, whereas # get_buffers() starts at this layer's own K buffer. @@ -1079,6 +1088,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. @@ -2116,7 +2142,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 @@ -2467,6 +2498,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/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 6e6568afccc6..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 diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index fbdea420106a..344012410982 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2447,6 +2447,27 @@ 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"] + backend = spec.get("attention_backend") + if not ((spec.get("decoding_type") == "DFlash" and backend == "TRTLLM") + or (spec.get("decoding_type") == "DSpark" + and backend in (None, "AUTO", "TRTLLM"))): + return ("speculative_config to be DFlash with attention_backend " + "TRTLLM, or DSpark with attention_backend AUTO or TRTLLM, " + "in both sections") + return None + + @click.command("disaggregated_mpi_worker") @stability_option( "-c", @@ -2486,19 +2507,12 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, from tensorrt_llm.llmapi.disagg_utils import split_world_comm disagg_cfg = parse_disagg_config_file(config_file) - with_spec = { - cfg.type - for cfg in disagg_cfg.server_configs - if cfg.other_args.get("speculative_config") - } - attention_parallel_configs = {(cfg.other_args.get("tensor_parallel_size", - 1), - cfg.other_args.get("enable_attention_dp", - False)) - for cfg in disagg_cfg.server_configs} - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ( - "1" if with_spec == {"ctx", "gen"} - and len(attention_parallel_configs) == 1 else "0") + blocker = _shared_draft_kv_blocker(disagg_cfg.server_configs) + if blocker is not None and os.environ.get( + "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1": + raise click.UsageError( + f"TRTLLM_DISAGG_SHARED_DRAFT_KV=1 requires {blocker}.") + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" if blocker else "1") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 5dc41d1bf01f..27641e2f94a0 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1734,9 +1734,9 @@ def test_gsm8k_1p1d_unified_draft_kv(self): ctx_server_config, gen_server_config, self.MODEL_PATH, - extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}, assert_worker_log_contains= - "target and draft KV layers share one KV cache manager") as llm: + "target and draft KV layers share one KV cache manager", + extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}) as llm: llm.args.speculative_config = DFlashDecodingConfig( max_draft_len=self.MAX_DRAFT_LEN, speculative_model=self.DRAFT_PATH, diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index c91507098ac3..5151d6f5b407 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -16,6 +16,7 @@ from types import SimpleNamespace +import click import pytest import torch import yaml @@ -52,7 +53,7 @@ def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): ) -def _creator(*, is_disagg, target, draft, tp_size=1): +def _creator(*, is_disagg, target, draft, tp_size=1, draft_windows=None): creator = object.__new__(KvCacheCreator) creator._speculative_config = SimpleNamespace( spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), @@ -63,8 +64,15 @@ def _creator(*, is_disagg, target, draft, tp_size=1): creator._draft_model_engine = None creator._draft_config = draft creator._mapping = SimpleNamespace( - pp_size=1, cp_size=1, tp_size=tp_size, enable_attention_dp=False + pp_size=1, + cp_size=1, + tp_size=tp_size, + enable_attention_dp=False, + pp_layers=lambda num_layers: list(range(num_layers)), ) + creator._max_seq_len = 4096 + creator._kv_cache_config = SimpleNamespace() + creator._get_draft_max_attention_window = lambda max_seq_len, kv_cache_config: draft_windows creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) creator._model_engine = SimpleNamespace( model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 @@ -122,19 +130,22 @@ def test_layout_compares_rank_local_heads(): @pytest.mark.parametrize("is_disagg", [False, True]) -def test_attention_dp_unifies_with_unsharded_heads(is_disagg): - # Under attention DP each rank keeps every KV head, so the TP4 layouts of - # an 8-head target and a 2-head draft are not resharded to 2 and 1. - target = _model_config(kv_heads=8, head_dim=128) - draft = _model_config(kv_heads=8, head_dim=128) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_attention_dp_changes_the_disagg_layout_decision(is_disagg, attention_dp): + # At TP4 a 4-head target and a 1-head draft both hold one head per rank, so + # they transfer together. Under attention DP each rank keeps every head (4 + # and 1), so the layouts differ and disaggregated serving keeps a separate cache. + target = _model_config(kv_heads=4, head_dim=128) + draft = _model_config(kv_heads=1, head_dim=128) creator = _creator(is_disagg=is_disagg, target=target, draft=draft, tp_size=4) - creator._mapping.enable_attention_dp = True + creator._mapping.enable_attention_dp = attention_dp - assert creator._dflash_kv_layout(target, torch.bfloat16)[0] == 8 - assert ( - creator._dflash_kv_layout(_model_config(kv_heads=2, head_dim=128), torch.bfloat16)[0] == 2 + heads = ( + creator._dflash_kv_layout(target, torch.bfloat16)[0], + creator._dflash_kv_layout(draft, torch.bfloat16)[0], ) - assert creator._should_unify_dflash_kv_cache() + assert heads == ((4, 1) if attention_dp else (1, 1)) + assert creator._should_unify_dflash_kv_cache() == (not (is_disagg and attention_dp)) @pytest.mark.parametrize("is_disagg", [False, True]) @@ -177,19 +188,21 @@ def stop(*args, **kwargs): monkeypatch.setattr(serve, "set_cuda_device", lambda: None) monkeypatch.setattr("tensorrt_llm.llmapi.disagg_utils.split_world_comm", stop) monkeypatch.delenv("TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT", raising=False) - with pytest.raises(_StopLauncher): - serve.disaggregated_mpi_worker.callback( - config_file=str(config), log_level="info", telemetry=False - ) + serve.disaggregated_mpi_worker.callback( + config_file=str(config), log_level="info", telemetry=False + ) def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_path): - _launch_from_disagg_config( - monkeypatch, - tmp_path, - context_servers={}, - generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, - ) + # The launcher rejects the variable when only one section sets speculative_config. + with pytest.raises(click.UsageError, match="both context_servers and generation_servers"): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={}, + generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, + ) + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=8, head_dim=128) @@ -200,15 +213,88 @@ def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_ def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): spec = {"decoding_type": "DFlash"} - _launch_from_disagg_config( - monkeypatch, - tmp_path, - context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, - generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, - ) + with pytest.raises(click.UsageError, match="same tensor_parallel_size and enable_attention_dp"): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={"tensor_parallel_size": 4, "speculative_config": spec}, + generation_servers={"tensor_parallel_size": 8, "speculative_config": spec}, + ) + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=4, head_dim=128) for tp_size in (4, 8): creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=tp_size) assert not creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_draft_layers_with_different_windows_keep_separate_cache(is_disagg): + layout = _model_config(kv_heads=8, head_dim=128) + + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout, draft_windows=[128, 4096]) + + assert not creator._should_unify_dflash_kv_cache() + assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_draft_layers_with_one_window_still_unify(is_disagg): + layout = _model_config(kv_heads=8, head_dim=128) + + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout, draft_windows=[128, 128]) + + assert creator._should_unify_dflash_kv_cache() + + +def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=[128, 4096]) + + with pytest.raises(ValueError, match="different attention windows"): + creator._check_shared_draft_kv_request() + + +@pytest.mark.parametrize( + "context_servers,generation_servers,match", + [ + ( + {"enable_attention_dp": True}, + {"enable_attention_dp": False}, + "same tensor_parallel_size and enable_attention_dp", + ), + ( + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, + "attention_backend TRTLLM", + ), + ( + {"speculative_config": {"decoding_type": "DFlash"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}}, + "attention_backend TRTLLM", + ), + ], + ids=["mismatched_attention_dp", "vanilla_dflash_backend", "default_dflash_backend"], +) +def test_launcher_rejects_shared_cache_for_ineligible_config( + monkeypatch, tmp_path, context_servers, generation_servers, match +): + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + context_servers = {"speculative_config": spec, **context_servers} + generation_servers = {"speculative_config": spec, **generation_servers} + + with pytest.raises(click.UsageError, match=match): + _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers) + + +def test_launcher_accepts_eligible_shared_cache_config(monkeypatch, tmp_path): + spec = {"decoding_type": "DFlash", "attention_backend": "TRTLLM"} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + {"speculative_config": spec, "enable_attention_dp": True}, + {"speculative_config": spec, "enable_attention_dp": True}, + ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index fa7915d27cc8..afa843715431 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,9 @@ def get_num_attention_layers(self) -> int: # A Mock's truthy return would send get_num_extra_kv_tokens down the # one-engine arm; this test's config is a two-model drafter. mode.use_one_engine.return_value = False + # Likewise a truthy is_dflash() would route costing through the shared + # target and drafter cache path. + mode.is_dflash.return_value = False costed_windows: list[tuple[object, list[int] | None]] = [] class RecordingKVCacheManager(KVCacheManagerV2): diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 2d0a73e58826..1a3947737c41 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -27,6 +27,7 @@ _create_kv_cache_manager, _derive_v2_layer_type_attention_windows, _get_num_pool_groups_for_estimation, + _get_num_storage_pool_groups, get_mla_context_workspace_kv_len_cap, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window @@ -1641,3 +1642,55 @@ def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> Non assert manager_config is kv_cache_config assert manager_config.max_attention_window is None + + +def _fp8_target_bf16_draft_specs(*, draft_window): + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + return [ + *(replace(target, global_layer_id=i) for i in range(3)), + replace( + target, + global_layer_id=3, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + sliding_window_size=draft_window, + ), + ] + + +def test_unified_fp8_target_bf16_draft_scales_by_combined_pool_groups(): + """The V2 pool-group count comes from the combined target and draft layout. + + The FP8 target and the windowed BF16 draft differ in life cycle and slot + size, so they land in two pool groups; counting the target's model config + alone gives one. + """ + tpb = 32 + spec_config = SimpleNamespace( + tokens_per_gen_step=1, + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, use_one_engine=lambda: False), + ) + baseline = _make_creator(tpb, [_make_mock_request(128)], False, 1) + baseline._speculative_config = spec_config + baseline._build_unified_dflash_layer_specs = lambda kv_cache_config: () + unified = _make_creator(tpb, [_make_mock_request(128)], False, 1) + unified._speculative_config = spec_config + specs = _fp8_target_bf16_draft_specs(draft_window=128) + unified._build_unified_dflash_layer_specs = lambda kv_cache_config: tuple(specs) + + assert _get_num_storage_pool_groups(specs, tpb) == 2 + assert unified._get_token_num_for_estimation() == 2 * baseline._get_token_num_for_estimation() + + +def test_unified_fp8_target_bf16_draft_in_one_life_cycle_is_one_pool_group(): + """Pages of different sizes in one life cycle share a pool group.""" + specs = _fp8_target_bf16_draft_specs(draft_window=None) + + assert _get_num_storage_pool_groups(specs, 32) == 1 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 6dcbcd10a5e3..84d49baebd64 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -24,6 +24,7 @@ import torch from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.pyexecutor._util import _get_num_storage_pool_groups from tensorrt_llm._torch.pyexecutor.engine.runners.decoder 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 ( @@ -33,6 +34,8 @@ KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _KVCacheManagerInitStatus, _swa_endpoint_priority, _sync_kv_cache_manager_init_status, @@ -3770,3 +3773,181 @@ def get_base_page_indices(self, layer_group_id, beam_id=DEFAULT_BEAM_INDEX): assert len(fake_logger.messages) == 1 assert "unavailable" in fake_logger.messages[0] + + +@pytest.mark.parametrize( + "target_dtype,target_layers,draft_windows", + [ + (DataType.BF16, 2, [None]), + (DataType.FP8, 2, [None]), + (DataType.FP8, 3, [128]), + (DataType.BF16, 2, [128]), + ], + ids=["bf16_same_window", "fp8_same_window", "fp8_windowed_draft", "bf16_windowed_draft"], +) +def test_storage_pool_group_count_matches_manager( + target_dtype, target_layers, draft_windows +) -> None: + """The estimator's pool-group count is what the manager creates.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=target_dtype, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + layer_specs = [replace(target, global_layer_id=i) for i in range(target_layers)] + layer_specs += [ + replace( + target, + global_layer_id=target_layers + i, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + sliding_window_size=window, + ) + for i, window in enumerate(draft_windows) + ] + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_swa_scratch_reuse=False, + ) + manager = KVCacheManagerV2( + kv_cache_config, + CacheType.SELF, + num_layers=len(layer_specs), + num_kv_heads=8, + head_dim=128, + tokens_per_block=32, + max_seq_len=1024, + max_batch_size=2, + max_num_tokens=128, + mapping=Mapping(), + dtype=target_dtype, + layer_specs=layer_specs, + ) + try: + assert _get_num_storage_pool_groups(layer_specs, 32) == len(manager.impl.pool_group_descs) + finally: + manager.shutdown() + + +@pytest.mark.parametrize("disagg", [False, True], ids=["aggregated", "disaggregated"]) +def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize(disagg: bool) -> None: + """The drafter appends its block after the context, so its pages must reach it. + + After a decode step accepts every draft token the context is ``prompt + K + + 1`` long and the drafter's block follows it. The prompt is chosen so that + this furthest write is the first position of a third page. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + max_draft_len = 7 + block_width = max_draft_len + 1 + tokens_per_block = 32 + prompt_len = 49 + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft_layer = 2 + layer_specs = [ + target, + replace(target, global_layer_id=1), + replace( + target, global_layer_id=draft_layer, owner=KVCacheLayerOwner.DRAFT, dtype=DataType.BF16 + ), + ] + manager = KVCacheManagerV2( + KvCacheConfig( + enable_block_reuse=False, + max_gpu_total_bytes=64 << 20, + host_cache_size=0, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=len(layer_specs), + num_kv_heads=8, + head_dim=128, + tokens_per_block=tokens_per_block, + max_seq_len=1024, + max_batch_size=2, + mapping=Mapping(), + dtype=DataType.FP8, + layer_specs=layer_specs, + spec_config=DFlashDecodingConfig(max_draft_len=max_draft_len, speculative_model="draft"), + vocab_size=4096, + enable_stats=False, + ) + request = _ContextRequest( + 1, + list(range(prompt_len)), + prompt_len, + "conversation", + use_conversation_params=False, + is_disagg_generation_init_state=disagg, + ) + try: + if disagg: + request.py_draft_tokens = [] + request.is_disagg_generation_transmission_complete = False + request.context_phase_params = None + assert manager.prepare_disagg_gen_init(request) + request.is_disagg_generation_init_state = False + else: + _run_context(manager, request) + generation = SimpleNamespace( + py_request_id=request.py_request_id, + py_beam_width=1, + py_draft_tokens=[0] * max_draft_len, + state=LlmRequestState.GENERATION_IN_PROGRESS, + is_dummy_request=False, + is_disagg_generation_transmission_complete=False, + context_phase_params=None, + max_beam_num_tokens=prompt_len + 1, + ) + assert manager.try_allocate_generation(generation) + + furthest_write_end = prompt_len + max_draft_len + 1 + block_width + pages = manager.get_batch_cache_indices([request.py_request_id], layer_idx=draft_layer)[0] + assert len(pages) * tokens_per_block >= furthest_write_end + finally: + manager.free_resources(request) + manager.shutdown() + + +def test_cost_headroom_counts_the_drafter_block_only_for_shared_layers() -> None: + """The budget charges the allocation headroom, not the position-side count.""" + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens + + max_draft_len = 7 + spec_config = DFlashDecodingConfig(max_draft_len=max_draft_len, speculative_model="draft") + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=8, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft = replace(target, global_layer_id=1, owner=KVCacheLayerOwner.DRAFT) + + _, own = _get_generation_kv_capacity(spec_config, is_draft=False, layer_specs=[target]) + _, shared = _get_generation_kv_capacity( + spec_config, is_draft=False, layer_specs=[target, draft] + ) + + # The block is K + 1 = 8 wide against K - 1 = 6 extra positions. + assert shared - own == (max_draft_len + 1) - (max_draft_len - 1) + assert _get_kv_alloc_extra_tokens(spec_config, [target, draft]) == max_draft_len + 1 + assert get_num_extra_kv_tokens(spec_config) == max_draft_len - 1 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py new file mode 100644 index 000000000000..10cb0365548e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DFlash shared context block table over a KVCacheManagerV2 draft layout. + +A separate draft manager whose layers cannot share one block table falls back to +the private context arena, as it did before unified caching. A unified manager +keeps its invariant, so the same layout there is an error. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.speculative.dflash import DFlashWorker + +pytestmark = pytest.mark.cpu_only + + +def _draft_manager(*, layer_groups, layer_scales): + mgr = object.__new__(KVCacheManagerV2) + mgr.kv_cache_pool_mapping = torch.zeros(len(layer_groups), 2, dtype=torch.int32) + mgr.layer_offsets = {layer_id: layer_id for layer_id in range(len(layer_groups))} + mgr.layer_to_pool_mapping_dict = dict(enumerate(layer_groups)) + mgr.enable_swa_scratch_reuse = False + mgr.impl = SimpleNamespace( + get_page_index_converter=lambda layer_offset, role: SimpleNamespace( + scale=layer_scales[layer_offset], layer_offset=0 + ), + get_page_index_scale=lambda layer_offset, role: layer_scales[layer_offset], + ) + return mgr + + +def _worker(*, unified): + worker = object.__new__(DFlashWorker) + worker._ctx_layer_ids = (0, 1) + worker._ctx_unified_cache = unified + worker._ctx_kv_factor = 2 + return worker + + +def _init_block_tables(worker, mgr): + pool = [torch.empty(0)] + return worker._init_ctx_block_tables(mgr, pool, num_slots=2, L=2, nkv=1, hd=64, page_size=32) + + +@pytest.mark.parametrize( + "layer_groups,layer_scales", + [([0, 1], [2, 2]), ([0, 0], [2, 4])], + ids=["different_lifecycle", "different_page_index_scale"], +) +def test_separate_v2_draft_manager_with_unsupported_layout_uses_private_arena( + layer_groups, layer_scales +): + mgr = _draft_manager(layer_groups=layer_groups, layer_scales=layer_scales) + + assert _init_block_tables(_worker(unified=False), mgr) is False + + +@pytest.mark.parametrize( + "layer_groups,layer_scales", + [([0, 1], [2, 2]), ([0, 0], [2, 4])], + ids=["different_lifecycle", "different_page_index_scale"], +) +def test_unified_manager_with_unsupported_layout_is_an_error(layer_groups, layer_scales): + mgr = _draft_manager(layer_groups=layer_groups, layer_scales=layer_scales) + + with pytest.raises(NotImplementedError, match="same lifecycle and page-index scale"): + _init_block_tables(_worker(unified=True), mgr) + + +def test_drafter_block_inside_the_allocation_is_accepted(): + num_ctx = torch.tensor([57, 10]) + allocated = torch.tensor([57, 56]) + + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def test_drafter_block_outside_the_allocation_is_an_error(): + num_ctx = torch.tensor([58, 10]) + allocated = torch.tensor([57, 56]) + + with pytest.raises(RuntimeError, match="block of 7 positions falls outside"): + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) From 0e55a6e3d3ebac753c4130233fe20b57ccf3665d Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 20 Oct 2026 23:32:04 -0700 Subject: [PATCH 03/23] [None][chore] Run the unified KV cache stack on the agentx KV cache bindings The ported #20178/#20486/#20709 Python assumes main's KV cache manager bindings. Two pieces are not on agentx: - Per-layer AttentionReusePolicy (main #19727 C++). The stack only ever uses REQUIRED, which is what agentx does for every layer, so the reuse_policy field, its AttentionLayerConfig argument and its place in the scratch and pool-group keys are dropped. The one test that needs PRIVATE is skipped. - BufferConfig.coalescing_group (the stack's own C++, which the agentx image does not carry). Group 0 is what the prebuilt bindings do implicitly, so the argument is passed only for nonzero (NVFP4) groups, which fail loudly unless the bindings are rebuilt from this branch's cpp/. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/pyexecutor/_util.py | 25 ++++++-------- .../kv_cache/kv_cache_manager_v2.py | 34 +++++++++++++------ .../kv_cache/test_kv_cache_estimation.py | 7 +++- 3 files changed, 40 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b28f9f10db53..5620b0cd65ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -40,7 +40,6 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.quantization import QuantAlgo -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy from ..attention.backends import get_sparse_attn_kv_cache_manager from ..disaggregation.kv_cache_transceiver import ( @@ -695,28 +694,28 @@ def _get_num_storage_pool_groups(layer_specs: Sequence[KVCacheLayerSpec], """Count the pool groups KVCacheManagerV2 creates for explicit layer specs. Mirrors the grouping key of ``createStorageConfig``: a life cycle is a - distinct (window, sink blocks, reuse policy), its slot sizes are the sizes - of its coalesced buffers in descending order, and life cycles with the same - (reuse policy, slot sizes) share one pool group. + 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), layer.reuse_policy) + 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({(life_cycle[2], - tuple( - sorted((size * count - for (_, size), count in coalesced.items()), - reverse=True))) - for life_cycle, coalesced in life_cycles.items()}) + 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( @@ -2272,7 +2271,6 @@ def _build_unified_dflash_layer_specs( buffer_roles=self._layer_spec_roles(target_dtype, k_only=target_k_only), sliding_window_size=window, - reuse_policy=AttentionReusePolicy.REQUIRED, ) for layer_id, window in zip(target_local_ids, target_windows) ] specs.extend( @@ -2285,7 +2283,6 @@ def _build_unified_dflash_layer_specs( buffer_roles=self._layer_spec_roles(draft_dtype, k_only=draft_k_only), sliding_window_size=window, - reuse_policy=AttentionReusePolicy.REQUIRED, ) for layer_id, window in zip(draft_local_ids, draft_windows)) return tuple(specs) 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 7c9452fe3ff2..d7b8a9b83095 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -66,7 +66,6 @@ DEFAULT_BEAM_INDEX, GPU_LEVEL, AttentionLayerConfig, - AttentionReusePolicy, AttnLifeCycle, Batch, BatchDesc, @@ -296,7 +295,8 @@ class KVCacheLayerSpec: buffer_roles: tuple[DataRole, ...] sliding_window_size: Optional[int] = None num_sink_tokens: Optional[int] = None - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + # [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: @@ -366,6 +366,23 @@ def kv_factor(self) -> int: 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]: @@ -623,12 +640,7 @@ def _estimate_swa_cache_size( elif not scratch: size_per_token += layer_size else: - scratch_key = ( - int(window_size), - layer_size, - sink_blocks, - layer.reuse_policy if layer else None, - ) + scratch_key = (int(window_size), layer_size, sink_blocks) if scratch_key in scratch_keys: size_per_request += window_tokens * layer_size else: @@ -3262,7 +3274,9 @@ def _build_base_config( size=layer.bytes_per_token(role) * tokens_per_block, # Keep NVFP4 data (1), its scales (2), and other data (0) # separate while preserving matching K/V page order. - coalescing_group=( + # [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 ), ) @@ -3271,7 +3285,6 @@ def _build_base_config( ], 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) ] @@ -3374,7 +3387,6 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa buffers=buffers, sliding_window_size=layer.sliding_window_size, num_sink_tokens=layer.num_sink_tokens, - reuse_policy=layer.reuse_policy, ) ) else: 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 1a3947737c41..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 @@ -48,7 +48,6 @@ TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy pytestmark = pytest.mark.cpu_only @@ -961,7 +960,13 @@ 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, From ff113646c055d14487970b90364e2ddb2a211299 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 20 Oct 2026 23:34:06 -0700 Subject: [PATCH 04/23] [None][fix] Bound DFlash's shared context by the allocation's extent DFlash took the number of valid block-table entries as the pages a request holds. On a unified cache with per-layer page tables (e.g. an FP8 target with a BF16 drafter), a sliding-window draft layer's evicted pages stay BAD_PAGE_INDEX in front of its live tail, so the count stopped at about one window (66 blocks for window 2048, page 32) whatever the context length. The decode-time allocation check then trapped once a context passed ~2.1K tokens, and prefill skipped the drafter's context store for later chunks and for prefix hits past the window. Use the extent of the allocation instead, the last valid entry + 1. Reads below the window are never issued and writes land in the live tail; for a full-attention drafter, and on the cpp backend, the extent equals the count. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/speculative/dflash.py | 29 +++++++++++++------ .../test_dflash_ctx_block_tables.py | 27 +++++++++++++++++ 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 57fe4516c69f..d40c96ee1bfe 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -520,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) @@ -791,9 +793,11 @@ 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 @@ -1081,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}, " @@ -1132,9 +1137,14 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: # below would then edit attn_metadata's own tensor. encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() encoded.sub_(self._ctx_block_offset) - # Before the clamp erases the placeholders. Informative only on the + # 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 @@ -1174,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. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py index 10cb0365548e..5c3e9eb7a2c7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -96,3 +96,30 @@ def test_drafter_block_outside_the_allocation_is_an_error(): with pytest.raises(RuntimeError, match="block of 7 positions falls outside"): DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def _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] From 7be69ca582eb6153a1f6debe06330d57ae9c7740 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 20 Oct 2026 23:49:53 -0700 Subject: [PATCH 05/23] [None][fix] Keep DFlash's dummy slot empty between steps Port of main's DFlashWorker._advance_ctx_len (#19988, c98b9765d), which the unified KV cache stack relies on but agentx lacks. CUDA-graph padding, warmup and attention-DP dummy rows share one dummy slot that no request owns, and each step added the dummies' accepted tokens to its length. On a unified KV cache the drafter's allocation check covers those rows, whose pages are the padding request's, so the growing length tripped the device assert during CUDA-graph capture at server start. Empty the slot right after the add, inside the captured step; with the windowed arena its base too. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/speculative/dflash.py | 20 ++++- .../hw_agnostic/test_dflash_dummy_slot.py | 83 +++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d40c96ee1bfe..ccfb67066c0c 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1598,6 +1598,23 @@ 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 _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: @@ -2490,8 +2507,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) 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..b467fcace365 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -0,0 +1,83 @@ +# 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] From 2b23449e9049baec13905ba01233e6c5368fe111 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 00:19:10 -0700 Subject: [PATCH 06/23] [None][fix] Re-seed DFlash generation slots in CUDA-graph warmup CUDA-graph warmup builds every batch from dummy requests whose ids repeat across batches (0..bs-1), so a reused id keeps its drafter slot and the previous batch's context length, which the replay after each capture advances for good. A later batch can give the same id far fewer pages, and on a unified KV cache the drafter's allocation check then trapped at server start: during the advanced-sampling capture pass the filler rows reached 31..79 positions against a 24-position allocation. Only warmup runs a CUDA-graph batch eagerly, so there re-seed each generation row's slot from this batch's positions, as prepare() already seeds a new request. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/speculative/dflash.py | 25 +++++++++++++++++++ .../hw_agnostic/test_dflash_dummy_slot.py | 21 ++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index ccfb67066c0c..ee112efef111 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1598,6 +1598,25 @@ 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) + updates = {} + for request_idx, rid in enumerate(request_ids[num_contexts:], num_contexts): + slot = self._req_to_slot.get(rid) + if slot is None or rid == ATTENTION_DP_DUMMY_REQUEST_ID: + continue + token_offset = sum(spec_metadata.seq_lens[:request_idx]) + updates[slot] = min(int(spec_metadata.host_position_ids[token_offset]), self._max_ctx) + 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. @@ -1924,6 +1943,12 @@ def _forward_impl( ): # 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. 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 index b467fcace365..fe29e3395073 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -81,3 +81,24 @@ def test_windowed_arena_dummy_slot_is_empty_with_its_base(): 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] From 9862d63170375402c53676fa1f2b9d48cb7d7e45 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 00:43:50 -0700 Subject: [PATCH 07/23] [None][test] Align the unified DFlash KV gate tests with per-view head transfer Per-view head transfer dropped the equal-heads rule from the disaggregated layout check and the launcher's identical tensor_parallel_size and enable_attention_dp rule, but the gate tests still expected those refusals, and their speculative-config mock lacked is_dspark, which the gate now calls to name the drafter. Expect mismatched heads, attention DP and TP to share the cache, keep NVFP4 (whose block-scale views are not resharded) as the remaining head-count refusal, and give the mocks is_dspark (including the derived-window costing test, which reaches the same gate). Signed-off-by: Zheyu Fu --- .../kv_cache/test_dflash_unified_kv_gate.py | 57 ++++++++++++------- .../kv_cache/test_kv_cache_budget_split.py | 5 +- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 5151d6f5b407..98540f020778 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -56,7 +56,9 @@ def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): def _creator(*, is_disagg, target, draft, tp_size=1, draft_windows=None): creator = object.__new__(KvCacheCreator) creator._speculative_config = SimpleNamespace( - spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), + 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, @@ -93,10 +95,20 @@ def test_matching_draft_layout_unifies(is_disagg): assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() -def test_mismatched_draft_kv_heads_unify_only_in_aggregated_serving() -> None: +@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() @@ -120,21 +132,22 @@ def test_fp8_target_with_bf16_draft_unifies(is_disagg: bool) -> None: 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 not creator._should_unify_dflash_kv_cache() + 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_changes_the_disagg_layout_decision(is_disagg, attention_dp): - # At TP4 a 4-head target and a 1-head draft both hold one head per rank, so - # they transfer together. Under attention DP each rank keeps every head (4 - # and 1), so the layouts differ and disaggregated serving keeps a separate cache. +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) @@ -145,7 +158,7 @@ def test_attention_dp_changes_the_disagg_layout_decision(is_disagg, attention_dp 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() == (not (is_disagg and attention_dp)) + assert creator._should_unify_dflash_kv_cache() @pytest.mark.parametrize("is_disagg", [False, True]) @@ -211,22 +224,22 @@ def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_ assert not creator._should_unify_dflash_kv_cache() -def test_dflash_with_different_tp_keeps_separate_cache(monkeypatch, tmp_path): - spec = {"decoding_type": "DFlash"} - with pytest.raises(click.UsageError, match="same tensor_parallel_size and enable_attention_dp"): +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}, ) - monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) target = _model_config(kv_heads=8, head_dim=128) draft = _model_config(kv_heads=4, head_dim=128) for tp_size in (4, 8): creator = _creator(is_disagg=True, target=target, draft=draft, tp_size=tp_size) - assert not creator._should_unify_dflash_kv_cache() + assert creator._should_unify_dflash_kv_cache() @pytest.mark.parametrize("is_disagg", [False, True]) @@ -259,11 +272,6 @@ def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): @pytest.mark.parametrize( "context_servers,generation_servers,match", [ - ( - {"enable_attention_dp": True}, - {"enable_attention_dp": False}, - "same tensor_parallel_size and enable_attention_dp", - ), ( {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, @@ -275,7 +283,7 @@ def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): "attention_backend TRTLLM", ), ], - ids=["mismatched_attention_dp", "vanilla_dflash_backend", "default_dflash_backend"], + ids=["vanilla_dflash_backend", "default_dflash_backend"], ) def test_launcher_rejects_shared_cache_for_ineligible_config( monkeypatch, tmp_path, context_servers, generation_servers, match @@ -288,13 +296,18 @@ def test_launcher_rejects_shared_cache_for_ineligible_config( _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers) -def test_launcher_accepts_eligible_shared_cache_config(monkeypatch, tmp_path): +@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": True}, - {"speculative_config": spec, "enable_attention_dp": True}, + {"speculative_config": spec, "enable_attention_dp": context_attention_dp}, + {"speculative_config": spec, "enable_attention_dp": generation_attention_dp}, ) 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 afa843715431..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,9 +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() would route costing through the shared - # target and drafter cache path. + # 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): From 400ff3028b108ef9d523a4a1434fc8bc55f188fe Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 00:44:05 -0700 Subject: [PATCH 08/23] [None][fix] Refuse an unshared windowed draft pool on disaggregated workers Without TRTLLM_DISAGG_SHARED_DRAFT_KV=1, a DFlash or standalone DSpark drafter on the TRTLLM backend keeps a separate draft KV cache on a disaggregated worker. A transferred request's drafter fills it from column 0, while the cache commits history at the request's real positions and retires the blocks behind a sliding window, which the drafter's table then maps to block 0. Once a request passes the window, acceptance length silently collapses to about 1 (long prompts on the Muse Spark Open DFlash2 drafter: 1.005, VANILLA 3.604). The startup check that refuses an unsatisfiable TRTLLM_DISAGG_SHARED_DRAFT_KV=1 now also refuses this fallback and points at the shared draft KV or attention_backend VANILLA. A worker cannot tell which disaggregated role it serves, so a context worker is refused as well. The llm_args warning that the drafter's context is not transferred no longer fires when the draft KV is shared. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/pyexecutor/_util.py | 24 +++++++++++++++---- tensorrt_llm/llmapi/llm_args.py | 6 ++++- .../kv_cache/test_dflash_unified_kv_gate.py | 19 +++++++++++++++ 3 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5620b0cd65ba..d233f4b76958 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2071,12 +2071,28 @@ def _should_unify_dflash_kv_cache(self) -> bool: and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1") def _check_shared_draft_kv_request(self) -> None: - """Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored.""" - if (not self._is_disagg - or os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1" - or self._speculative_config is None + """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): + 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: set " + "TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on every worker, or use " + "attention_backend VANILLA. 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( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 96c7b07e41bb..a8e31c2b9186 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -7071,7 +7071,11 @@ def validate_speculative_config(self): logger.info( "DFlash attention_backend AUTO resolved to VANILLA " "(disaggregated worker or modeling_v2).") - if disagg: + # A draft KV shared with the target (TRTLLM_DISAGG_SHARED_DRAFT_KV=1) is transferred. + if disagg and not ( + self.speculative_config.attention_backend == "TRTLLM" + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") + == "1"): # 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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 98540f020778..9e81bde0f2e2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.llmapi import DFlashDecodingConfig from tensorrt_llm.quantization import QuantAlgo from tensorrt_llm.quantization.mode import QuantMode @@ -311,3 +312,21 @@ def test_launcher_accepts_eligible_shared_cache_config( {"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() From 93bec9503a730332972d0da1df6cbce6c3a15b45 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 01:35:38 -0700 Subject: [PATCH 09/23] [None][test] Expect the disaggregated DFlash warning only without a shared draft KV feat/musespark-open-agentx turned the refusal of DFlash on a disaggregated engine into a warning (264e60132) but kept the test that expects the refusal, so it failed. A draft KV shared with the target now silences the warning as well. Test the current behaviour: the warning fires for VANILLA, and for TRTLLM without TRTLLM_DISAGG_SHARED_DRAFT_KV=1; it stays silent for TRTLLM with it. Signed-off-by: Zheyu Fu --- .../llmapi/test_disagg_dflash_guard.py | 39 ++++++++++++------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/tests/unittest/llmapi/test_disagg_dflash_guard.py b/tests/unittest/llmapi/test_disagg_dflash_guard.py index 5e35b5a293bd..8534b4d7280e 100644 --- a/tests/unittest/llmapi/test_disagg_dflash_guard.py +++ b/tests/unittest/llmapi/test_disagg_dflash_guard.py @@ -19,6 +19,7 @@ DFlashDecodingConfig, NGramDecodingConfig, TorchLlmArgs, + llm_args, ) from tensorrt_llm.llmapi.llm_args import disagg_dflash_error @@ -63,24 +64,36 @@ 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, True), ("TRTLLM", "1", False)], + ) + 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 (TRTLLM backend, TRTLLM_DISAGG_SHARED_DRAFT_KV=1). + 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) From 70fbd34026abd4ad1e1fd0ec2294bab08a66a245 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 04:09:54 -0700 Subject: [PATCH 10/23] [None][perf] Convert only the batch's rows of per-layer KV page tables Without SWA scratch slots, a per-layer attention-op table is a plain conversion of the request's host page-index row. Every step still copied the whole host table to the device, rebuilt all layers x IndexMapper rows x max_blocks_per_seq into a full-capacity staging buffer, and then copied the batch's rows out. That work scales with max_seq_len and capacity, not with the batch: for a unified Muse Spark cache (67 layers, 129 rows) at max_seq_len 997696 it costs a 2 GiB staging buffer and 3.5-4.2 ms/step. Snapshot only the batch's host rows, up to the widest row's num_blocks (the rest of a row is BAD_PAGE_INDEX), into a fresh pinned buffer, copy it asynchronously and convert it straight into the destination with one Triton kernel. The snapshot also fixes the ordering: the async H2D of the live host table read it when the copy ran, so host-table updates for the next step could leak into this step's tables (nvbug 6293536 on V1). The batch's rows are bit-identical to the full rebuild; other rows are not touched. The copy refuses to run under CUDA-graph capture, which would replay a stale snapshot. The SWA-scratch path (unused by Muse Spark) is left as is. Signed-off-by: Zheyu Fu --- .../kv_cache/kv_cache_manager_v2.py | 115 ++++++++++ .../test_kv_cache_v2_attention_op_tables.py | 199 ++++++++++++++++++ 2 files changed, 314 insertions(+) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.py 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 d7b8a9b83095..de5b31e43b24 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 @@ -1330,6 +1332,55 @@ 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, +): + # 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=-1) + value = tl.where(base < 0, -1, 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 + 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 + ) + + 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. @@ -2999,6 +3050,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, @@ -3126,6 +3183,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, @@ -3134,6 +3246,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, 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..37d44f1e3f87 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.py @@ -0,0 +1,199 @@ +# 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 per-layer tables +must match the old full-capacity rebuild bit for bit, and be the host rows at +enqueue time, however the host table changes before the asynchronous copy runs. +""" + +from dataclasses import replace + +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 + +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() -> KVCacheManagerV2: + torch.cuda.init() + 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() -> None: + manager = _make_manager() + 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_copy_is_a_snapshot_of_the_host_table() -> None: + """The next step's host-table writes must not reach this step's pending copy.""" + manager = _make_manager() + 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() From 87bb9b00c85ed5d17b4cb79a29c7873d1e143424 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 04:09:55 -0700 Subject: [PATCH 11/23] [None][perf] Share attention-op page tables across layers of one KV pool A life cycle that mixes page sizes (a unified FP8 target / BF16 draft cache, heterogeneous head dims) fell back to one attention-op table per layer: 67 tables for Muse Spark, each max_num_sequences x max_blocks_per_seq wide (1 GiB per metadata copy at max_seq_len 997696), rebuilt every step. A per-layer table is base * scale + layer_offset. 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, moves to the kv_cache_pool_mapping offset the attention op already applies. Such layers now share one table: Muse Spark needs 3 (window FP8, window BF16, full FP8), the same layout the separate target and draft managers had. Page addresses and BAD_PAGE_INDEX entries are unchanged, and DFlash's decoded draft tables and extents are identical. SWA scratch reuse, NVFP4 and subclasses keep per-layer tables. Two readers that assume life-cycle tables now decline shared ones: the modeling_v2 Muse Spark block-offset stager, and the opt-in prims_ts FMHA, whose K/V page offset is the life cycle's. Signed-off-by: Zheyu Fu --- .../runtime.py | 2 + .../attention/backends/fmha/prims_ts.py | 5 + .../kv_cache/kv_cache_manager_v2.py | 99 ++++++++++++++++- .../test_kv_cache_v2_attention_op_tables.py | 103 ++++++++++++++++-- .../kv_cache/test_kvcm2_integration.py | 14 ++- 5 files changed, 209 insertions(+), 14 deletions(-) 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/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/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index de5b31e43b24..60fe5da70d95 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 @@ -1381,6 +1381,16 @@ def _convert_batch_block_offsets( ) +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. @@ -1453,6 +1463,8 @@ class KVCacheManagerV2(BaseResourceManager): _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: @@ -2571,7 +2583,13 @@ def _build_pool_mapping_tensors(self): self._get_layer_dtype_by_local_id(layer_id) == DataType.NVFP4 for layer_id in range(self.num_local_layers) ) - if self._use_per_layer_page_tables: + 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_layer_page_table_roles(layer_id, pool_id) @@ -2727,7 +2745,14 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if self.layer_specs is not None: storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} self._use_per_layer_page_tables |= len(storage_layouts) > 1 - if self._use_per_layer_page_tables: + # 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 @@ -2767,6 +2792,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 @@ -7098,6 +7190,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/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 index 37d44f1e3f87..cb2a62943900 100644 --- 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 @@ -15,12 +15,15 @@ """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 per-layer tables -must match the old full-capacity rebuild bit for bit, and be the host rows at -enqueue time, however the host table changes before the asynchronous copy runs. +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 @@ -37,7 +40,7 @@ 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 +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, PageIndexMode TOKENS_PER_BLOCK = 4 WINDOW = 2 * TOKENS_PER_BLOCK @@ -71,8 +74,13 @@ def _layer_specs(): return specs + [replace(draft, global_layer_id=6 + i) for i in range(2)] -def _make_manager() -> KVCacheManagerV2: +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, @@ -164,8 +172,8 @@ def eager_full_rebuild(monkeypatch): monkeypatch.setattr(torch_utils, "is_piecewise_running", lambda: True) -def test_per_layer_tables_match_full_rebuild() -> None: - manager = _make_manager() +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 @@ -178,9 +186,86 @@ def test_per_layer_tables_match_full_rebuild() -> None: manager.shutdown() -def test_copy_is_a_snapshot_of_the_host_table() -> None: +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() + manager = _make_manager(monkeypatch, grouped=grouped) try: expected, _ = _copy_batch(manager) request_ids = [request_id for request_id, *_ in REQUESTS] 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 84d49baebd64..dcc0110b49a3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -746,11 +746,16 @@ def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layo manager.shutdown() -def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: +@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, @@ -788,7 +793,9 @@ def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: layer_specs=[target, draft], ) try: - assert manager._use_per_layer_page_tables + # 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) @@ -804,7 +811,8 @@ def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: assert worker._init_ctx_block_tables( manager, pool, len(request_ids), 1, draft.num_kv_heads, draft.head_dim, TOKENS_PER_BLOCK ) - assert worker._ctx_block_offset > 0 + # 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), From 8abf57835f96ca9caf4bd97d70cdd182eea59d7b Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 16:49:24 -0700 Subject: [PATCH 12/23] [None][feat] Let DFlash AUTO pick TRTLLM on workers that share the draft KV With the unified KV cache, a disaggregated worker that sets TRTLLM_DISAGG_SHARED_DRAFT_KV=1 transfers the draft KV with the target's, so AUTO no longer needs VANILLA there. AUTO now keeps VANILLA on a disaggregated worker only without the variable, where a windowed TRTLLM draft pool is refused at startup; with it, AUTO follows the checkpoint rule as on an aggregated worker, and the warning that the drafter's context is not transferred stays off as it does for TRTLLM. Ranks pin AUTO before the KV cache creator reads the drafter backend, so a pinned AUTO unifies the target and draft KV like TRTLLM; a test covers it. Signed-off-by: Zheyu Fu --- docs/source/features/speculative-decoding.md | 2 +- tensorrt_llm/llmapi/llm_args.py | 20 ++++++++------- .../kv_cache/test_dflash_unified_kv_gate.py | 25 +++++++++++++++++++ .../hw_agnostic/test_dflash_draft_kv.py | 16 ++++++++---- .../llmapi/test_disagg_dflash_guard.py | 8 +++++- 5 files changed, 55 insertions(+), 16 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index a4fb3713758d..673203735bd5 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 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. +* `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 without `TRTLLM_DISAGG_SHARED_DRAFT_KV=1` 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 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a8e31c2b9186..8599f618f0b2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3316,7 +3316,7 @@ 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.") + "without TRTLLM_DISAGG_SHARED_DRAFT_KV=1 and under modeling_v2.") skip_ctx_buffer_budget_check: bool = Field( default=False, @@ -7064,18 +7064,20 @@ 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 = os.environ.get( + "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1" + # 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).") + "(disaggregated worker without " + "TRTLLM_DISAGG_SHARED_DRAFT_KV=1, or modeling_v2).") # A draft KV shared with the target (TRTLLM_DISAGG_SHARED_DRAFT_KV=1) is transferred. - if disagg and not ( - self.speculative_config.attention_backend == "TRTLLM" - and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") - == "1"): + 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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 9e81bde0f2e2..64bccaf646c1 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -14,6 +14,7 @@ # limitations under the License. """Selection of the unified target+DFlash KV cache layout.""" +import json from types import SimpleNamespace import click @@ -330,3 +331,27 @@ def test_disagg_separate_windowed_draft_pool_is_refused(monkeypatch, shared, win creator._check_shared_draft_kv_request() else: 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/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 5cca7cb3d1ea..5a5a05a964f6 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,23 @@ 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"), + (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.""" + """Unshared disaggregated workers 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) + 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/llmapi/test_disagg_dflash_guard.py b/tests/unittest/llmapi/test_disagg_dflash_guard.py index 8534b4d7280e..ce00b1c27ccc 100644 --- a/tests/unittest/llmapi/test_disagg_dflash_guard.py +++ b/tests/unittest/llmapi/test_disagg_dflash_guard.py @@ -72,7 +72,13 @@ class TestTorchLlmArgsValidation: @pytest.mark.parametrize( "attention_backend,shared_draft_kv,warns", - [("VANILLA", None, True), ("TRTLLM", None, True), ("TRTLLM", "1", False)], + [ + ("VANILLA", None, True), + ("TRTLLM", None, True), + ("TRTLLM", "1", False), + ("AUTO", None, True), + ("AUTO", "1", False), + ], ) def test_disagg_dflash_warns_unless_the_draft_kv_is_shared( self, monkeypatch, attention_backend, shared_draft_kv, warns From f01094bb2376bb5a13e8f402394d13f8a2b80711 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Thu, 22 Oct 2026 03:54:35 -0700 Subject: [PATCH 13/23] [None][test] Use KVCacheManagerV2 in the unified draft KV disagg test On this branch kv_cache_config.use_kv_cache_manager_v2 defaults to "auto", which resolves to V1 for Qwen3ForCausalLM. Both workers then refuse TRTLLM_DISAGG_SHARED_DRAFT_KV=1 at startup ("the target does not use KVCacheManagerV2"). Request V2 explicitly; main defaults to it. Signed-off-by: Zheyu Fu --- tests/integration/defs/accuracy/test_disaggregated_serving.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 27641e2f94a0..63fd1283c9e4 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1700,6 +1700,8 @@ def test_gsm8k_1p1d_unified_draft_kv(self): "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": { From 3abddc31571e8ad2d47d9c007d33e5e6c144b9ce Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Fri, 23 Oct 2026 02:20:08 -0700 Subject: [PATCH 14/23] [None][fix] Fail a disaggregated request at once when its peers disagree on draft KV layers A context and a generation worker that keep a different number of draft layers in their unified KV caches cannot move the drafter's prompt KV. Only the sender checked this, in its listener thread, where the error is logged and the request then waits out the KV transfer timeout. The receiver now runs the same check with its other layout gates, so every request to such a context peer fails immediately with the reason, and the peer is cached as incompatible. Signed-off-by: Zheyu Fu --- .../_torch/disaggregation/native/peer.py | 29 +++++++---- .../_torch/disaggregation/native/transfer.py | 18 ++++--- .../disaggregated/test_kv_transfer.py | 52 +++++++++++++++++++ 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index c958d2be9a4b..d0cc4f6de5bb 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -186,15 +186,7 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) self._validate_view_head_geometry(peer_ri) - if self._ri.num_draft_layers != peer_ri.num_draft_layers: - # The layer-count tolerance below would transfer only the target - # layers and leave the drafter's prompt KV unset on the receiver. - raise ValueError( - "PeerRegistrar: unified draft KV layer count differs " - f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " - "both peers must use the same speculative decoding KV cache layout; " - "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." - ) + 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) @@ -210,6 +202,25 @@ 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. diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index df3524ab8a65..07c0eea991fe 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3384,14 +3384,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, @@ -3403,6 +3404,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/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, From d59b14565f017a2f49cc5c0dba121d6c25fdafa5 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Fri, 23 Oct 2026 01:09:27 -0700 Subject: [PATCH 15/23] [None][fix] Share the DFlash draft KV under disaggregated_mpi_worker with the default AUTO backend The launcher turned on draft KV sharing only for a DFlash attention_backend set to TRTLLM. AUTO is the default, so a config with DFlash in both sections and no backend kept the drafter's prompt KV on the context server, and the generation server drafted without it: acceptance dropped with only a warning. AUTO and an unset backend now count like TRTLLM, as they already did for DSpark. An unwindowed drafter, whose AUTO resolves to VANILLA, is refused at startup instead of losing acceptance. Signed-off-by: Zheyu Fu --- tensorrt_llm/commands/serve.py | 13 +++++----- .../kv_cache/test_dflash_unified_kv_gate.py | 26 ++++++++++++++++--- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 344012410982..a339a008d1c8 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2458,13 +2458,12 @@ def _shared_draft_kv_blocker(server_configs) -> Optional[str]: "generation_servers") for cfg in server_configs: spec = cfg.other_args["speculative_config"] - backend = spec.get("attention_backend") - if not ((spec.get("decoding_type") == "DFlash" and backend == "TRTLLM") - or (spec.get("decoding_type") == "DSpark" - and backend in (None, "AUTO", "TRTLLM"))): - return ("speculative_config to be DFlash with attention_backend " - "TRTLLM, or DSpark with attention_backend AUTO or TRTLLM, " - "in both sections") + # 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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 64bccaf646c1..6103dbaa016d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -15,6 +15,7 @@ """Selection of the unified target+DFlash KV cache layout.""" import json +import os from types import SimpleNamespace import click @@ -277,15 +278,15 @@ def test_shared_cache_requested_with_unsupported_layout_fails_at_startup(): ( {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "VANILLA"}}, - "attention_backend TRTLLM", + "attention_backend AUTO or TRTLLM", ), ( - {"speculative_config": {"decoding_type": "DFlash"}}, + {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "FA4"}}, {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}}, - "attention_backend TRTLLM", + "attention_backend AUTO or TRTLLM", ), ], - ids=["vanilla_dflash_backend", "default_dflash_backend"], + ids=["vanilla_dflash_backend", "fa4_dflash_backend"], ) def test_launcher_rejects_shared_cache_for_ineligible_config( monkeypatch, tmp_path, context_servers, generation_servers, match @@ -298,6 +299,23 @@ def test_launcher_rejects_shared_cache_for_ineligible_config( _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)] ) From 98332a7c12db760f38cbda097f4e91ae2d4a59ea Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Fri, 23 Oct 2026 01:18:17 -0700 Subject: [PATCH 16/23] [None][feat] Share the DFlash draft KV on disaggregated workers by default A disaggregated worker kept its drafter's KV with the target's, so that the transceiver moves it, only with TRTLLM_DISAGG_SHARED_DRAFT_KV=1. Workers launched on their own (one trtllm-serve per server) had to set it by hand on every server; without it, AUTO fell back to VANILLA and the generation server drafted without the prompt context, with only a warning. Unset, the variable now leaves the choice to the drafter: DFlash with attention_backend AUTO or TRTLLM shares, so a context and a generation server given the same speculative_config agree on their own. 1 still requires sharing and 0 keeps the draft KV private. A pair that disagrees, such as a drafter on the generation server only, fails at the first KV transfer; to draft there alone, set attention_backend VANILLA or TRTLLM_DISAGG_SHARED_DRAFT_KV=0 on it. DSpark keeps needing the variable. Signed-off-by: Zheyu Fu --- docs/source/features/speculative-decoding.md | 10 ++- tensorrt_llm/_torch/pyexecutor/_util.py | 28 +++++---- tensorrt_llm/llmapi/llm_args.py | 62 +++++++++++++------ .../kv_cache/test_dflash_unified_kv_gate.py | 45 ++++++++++++-- .../hw_agnostic/test_dflash_draft_kv.py | 8 ++- .../llmapi/test_disagg_dflash_guard.py | 35 +++++++++-- 6 files changed, 143 insertions(+), 45 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 673203735bd5..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 without `TRTLLM_DISAGG_SHARED_DRAFT_KV=1` 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. +* `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 @@ -177,7 +177,13 @@ 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. The target and draft layers +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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index d233f4b76958..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) @@ -2063,12 +2064,12 @@ def _should_unify_dflash_kv_cache(self) -> bool: """Whether this run can place target and DFlash or standalone DSpark layers in one V2 manager. - Disaggregated serving shares only when TRTLLM_DISAGG_SHARED_DRAFT_KV=1. + 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 - and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1") + 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.""" @@ -2082,16 +2083,21 @@ def _check_shared_draft_kv_request(self) -> None: 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: set " - "TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on every worker, or use " - "attention_backend VANILLA. 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.") + "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: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 8599f618f0b2..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 " - "without TRTLLM_DISAGG_SHARED_DRAFT_KV=1 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,8 +7084,8 @@ 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) - shared_draft_kv = os.environ.get( - "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1" + 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) @@ -7073,9 +7093,23 @@ def validate_speculative_config(self): self.speculative_config.attention_backend = "VANILLA" logger.info( "DFlash attention_backend AUTO resolved to VANILLA " - "(disaggregated worker without " - "TRTLLM_DISAGG_SHARED_DRAFT_KV=1, or modeling_v2).") - # A draft KV shared with the target (TRTLLM_DISAGG_SHARED_DRAFT_KV=1) is transferred. + "(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 @@ -7092,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/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 6103dbaa016d..930eada27437 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -210,20 +210,26 @@ def stop(*args, **kwargs): 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, - context_servers={}, - generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, - ) + _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() @@ -351,6 +357,33 @@ def test_disagg_separate_windowed_draft_pool_is_refused(monkeypatch, shared, win 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.""" 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 5a5a05a964f6..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 @@ -60,6 +60,7 @@ def test_default_dflash_resolves_auto_before_choosing_its_draft_cache( [ ("NIXL", "0", "off", "VANILLA"), ("NIXL", "1", "off", "AUTO"), + ("NIXL", None, "off", "AUTO"), (None, "0", "auto", "VANILLA"), (None, "0", "off", "AUTO"), ], @@ -67,12 +68,15 @@ def test_default_dflash_resolves_auto_before_choosing_its_draft_cache( def test_auto_keeps_vanilla_where_trtllm_is_refused_or_bypassed( tmp_path, monkeypatch, transceiver, shared, modeling_v2, resolved ) -> None: - """Unshared disaggregated workers refuse TRTLLM; 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) - monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", shared) + 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/llmapi/test_disagg_dflash_guard.py b/tests/unittest/llmapi/test_disagg_dflash_guard.py index ce00b1c27ccc..4195f22898e6 100644 --- a/tests/unittest/llmapi/test_disagg_dflash_guard.py +++ b/tests/unittest/llmapi/test_disagg_dflash_guard.py @@ -17,11 +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 @@ -74,17 +75,20 @@ class TestTorchLlmArgsValidation: "attention_backend,shared_draft_kv,warns", [ ("VANILLA", None, True), - ("TRTLLM", None, True), + ("TRTLLM", None, False), ("TRTLLM", "1", False), - ("AUTO", None, True), + ("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 (TRTLLM backend, TRTLLM_DISAGG_SHARED_DRAFT_KV=1). + # 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: @@ -110,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 From 377a53b58bf548f57d625403926f0ee20a4d64d7 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 20:35:33 -0700 Subject: [PATCH 17/23] [None][fix] Scatter all fan-in writers of a KV-bounce region in one batched copy The receiver's scatter worker issued one batched copy per fan-in writer and synchronized the stream once at the end. The batched copy reuses one pinned metadata buffer per stream and relies on a sync between calls. With several writers the host refilled that buffer for writer k+1 while writer k's async upload was still queued behind writer k-1's kernel, so writer k's kernel ran with another writer's fragment table and writer k's head slice was never scattered: silent KV corruption. Seen on ctx TEP4 -> gen DEP8 (4 writers, ~1 GB each) once those receives were allowed to coalesce: every coalesced request decoded from partly stale KV (acceptance length 7.7-7.9 at K=7, degenerate output). Concatenate every writer's fragment table, each still read from its own source address, and issue one batched copy and one sync per region. A sync after every writer also removes the race, but measured ~117 ms per 4 GiB receive against ~50 ms for the racy single-sync loop (on different nodes), so keep a single copy and a single sync. Signed-off-by: Zheyu Fu --- .../disaggregation/native/bounce/impl.py | 19 +++++++-- tests/unittest/disaggregated/test_bounce.py | 39 +++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 1c70019e46bb..916f2629e219 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -475,12 +475,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 diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 70b761b1326f..20868ce1a4a7 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -550,6 +550,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. From a37c806dd68a211dfdcecb90289bd88788431e5c Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 21 Oct 2026 20:36:19 -0700 Subject: [PATCH 18/23] [None][fix] Allow TP fan-in KV bounce with non-uniform layer-group block sizes reserve() fell back to the per-fragment path for every multi-writer receive whose layer groups differ in per-block bytes. The equal split can only overrun when the writers hold different layers (PP fan-in); a TP fan-in writer sends a head slice of every block of every group, so its share is exactly total // num_writers for any group sizes. Receiver.dispatch_task now passes writers_split_layers (a fan-in whose writers span several PP stages) and reserve() applies the uniform-slot check only then. The default stays True, so other callers keep the old behavior. Depends on the previous commit: before it, coalesced fan-in receives could be scattered with the wrong fragment table. Example: ctx TEP4 -> gen DEP8 (4 writers) on a model with a full-attention and a sliding-window layer group (1 MiB vs 2.9 MiB per block, 3.5 MiB with DFlash draft layers) never coalesced; 4 GiB receives took ~450 ms on the per-fragment path. Signed-off-by: Zheyu Fu --- .../disaggregation/native/bounce/core.py | 4 +++- .../disaggregation/native/bounce/impl.py | 13 +++++++---- .../_torch/disaggregation/native/transfer.py | 11 +++++++--- tests/unittest/disaggregated/test_bounce.py | 11 ++++++++++ .../test_transfer_ownership_regressions.py | 22 +++++++++++++++---- 5 files changed, 49 insertions(+), 12 deletions(-) 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 916f2629e219..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] @@ -535,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/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 07c0eea991fe..5c484f8e56f3 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3146,8 +3146,9 @@ def _fanin_bounce_safe( (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.""" + 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). view_factors = getattr(overlap, "view_duplicate_head_factors", {}).values() @@ -3278,8 +3279,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: diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 20868ce1a4a7..5d66e248397a 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -446,6 +446,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]) diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..42cf8be43fb4 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, From e0d534b5dd620f2c80f2b53197475685b1e883b5 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sun, 25 Oct 2026 19:13:23 -0700 Subject: [PATCH 19/23] [None][test] Run the DFlash and disaggregated KV transfer tests in the MuseSpark pre-merge stage Pre-merge CI on feat/musespark-open-agentx runs only the MuseSpark stages, so the DFlash TRTLLM-backend, shared draft KV and disaggregated KV transfer tests were not run on this branch. List them in l0_gb300_musespark. Signed-off-by: Zheyu Fu --- .../test_lists/test-db/l0_gb300_musespark.yml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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..850db4615086 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,21 @@ 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. + - unittest/_torch/speculative/hw_agnostic/test_dflash_trtllm_window.py + - unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py + - unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.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_dflash_unified_kv_gate.py + - unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.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_bounce.py + - unittest/disaggregated/test_transfer_ownership_regressions.py + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv From 78698002caf6a551860ab8d79b2c5bb5163ec99a Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 26 Oct 2026 18:05:41 -0700 Subject: [PATCH 20/23] [None][test] Run the cpu_only DFlash and KV transfer tests in the MuseSpark CPU stage GPU stages deselect cpu_only tests, so the seven files whose tests are all cpu_only collected nothing on l0_gb300_musespark and failed with pytest's exit code 5, and test_kv_transfer.py's cpu_only tests ran nowhere. List those on l0_cpu_musespark, which runs cpu_only tests, and keep test_kv_transfer.py on both. Signed-off-by: Zheyu Fu --- .../test_lists/test-db/l0_cpu_musespark.yml | 10 ++++++++++ .../test_lists/test-db/l0_gb300_musespark.yml | 10 ++-------- 2 files changed, 12 insertions(+), 8 deletions(-) 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_gb300_musespark.yml b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml index 850db4615086..d71a45b69db4 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml @@ -40,19 +40,13 @@ l0_gb300_musespark: - 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. + # 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_draft_kv.py - - unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.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_dflash_unified_kv_gate.py - unittest/_torch/executor/kv_cache/test_kv_cache_v2_attention_op_tables.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_bounce.py - - unittest/disaggregated/test_transfer_ownership_regressions.py - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv From 31dad771fcfbbaf51149b18c664a4d607fba900d Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 26 Oct 2026 18:06:38 -0700 Subject: [PATCH 21/23] [None][fix] Check the senders' head duplication before a fan-in KV bounce A multi-writer bounce splits the region equally among the writers, which holds only if every writer sends a share. The gate checked the receiver's rank-wide head duplication factor but not the senders', so when the senders replicate KV heads and elect a subset to send, the elected ones overran their neighbours' slots: 4 KV heads with ctx TEP8 -> gen TEP2, or MLA with ctx TP8 -> gen TP4. Check both factors, as the per-view clause already does; those layouts take the per-fragment path. Signed-off-by: Zheyu Fu --- .../_torch/disaggregation/native/transfer.py | 9 +++++---- tests/unittest/disaggregated/test_bounce.py | 13 +++++++++++-- .../test_transfer_ownership_regressions.py | 1 + 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 5c484f8e56f3..501b0f3dd818 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -3138,9 +3138,9 @@ 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: - * every head duplication factor is 1, 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; + * 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 @@ -3151,8 +3151,9 @@ def _fanin_bounce_safe( 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 overlap.duplicate_head_factor != 1 or any( + if any(factor != 1 for factor in rank_factors) or any( factor != 1 for pair in view_factors for factor in pair ): return False diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 5d66e248397a..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 diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index 42cf8be43fb4..9ccb2ace7f59 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -707,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( From bb31e353b0181b50efb4d22eeda87443b5f404ee Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 26 Oct 2026 18:07:16 -0700 Subject: [PATCH 22/23] [None][fix] Keep an explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=0 under disaggregated_mpi_worker The launcher sets one value for every worker it spawns, and wrote 1 whenever both sections could share, overriding an explicit 0. Now that sharing is the default, 0 is how an operator keeps the draft KV private on a pair that would otherwise share; keep it. Signed-off-by: Zheyu Fu --- tensorrt_llm/commands/serve.py | 9 ++++++--- .../kv_cache/test_dflash_unified_kv_gate.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index a339a008d1c8..65e74256d65d 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2507,11 +2507,14 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, disagg_cfg = parse_disagg_config_file(config_file) blocker = _shared_draft_kv_blocker(disagg_cfg.server_configs) - if blocker is not None and os.environ.get( - "TRTLLM_DISAGG_SHARED_DRAFT_KV") == "1": + 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}.") - os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" if blocker else "1") + # 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/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 930eada27437..655755ac4cc5 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -233,6 +233,19 @@ def test_dflash_on_generation_worker_only_keeps_separate_cache(monkeypatch, tmp_ 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"} From 827f52481ef0b55cee05777a16fd877093a09d22 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 26 Oct 2026 18:08:13 -0700 Subject: [PATCH 23/23] [None][chore] Name BAD_PAGE_INDEX in the op-table kernel and walk warmup offsets once The batched op-table conversion kernel now takes BAD_PAGE_INDEX as a constexpr instead of a hardcoded -1, and its wrapper asserts that params is contiguous, as the kernel's pointer arithmetic needs. The warmup slot reseed carries a running token offset instead of re-summing seq_lens for every row. Signed-off-by: Zheyu Fu --- .../pyexecutor/kv_cache/kv_cache_manager_v2.py | 16 +++++++++++++--- tensorrt_llm/_torch/speculative/dflash.py | 12 +++++++----- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 60fe5da70d95..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 @@ -1342,6 +1342,7 @@ def _convert_batch_block_offsets_kernel( 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. @@ -1356,8 +1357,8 @@ def _convert_batch_block_offsets_kernel( 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=-1) - value = tl.where(base < 0, -1, base * scale + offset) + 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) @@ -1374,10 +1375,19 @@ def _convert_batch_block_offsets( 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 + rows, + out, + params, + num_seqs, + rows.shape[2], + out.stride(0), + num_blocks, + BLOCK=block, + BAD_PAGE=BAD_PAGE_INDEX, ) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index ee112efef111..59333b9278af 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1608,13 +1608,15 @@ def _reseed_warmup_generation_slots(self, spec_metadata: "DFlashSpecMetadata") - """ 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 request_idx, rid in enumerate(request_ids[num_contexts:], num_contexts): + for rid, seq_len in zip(request_ids[num_contexts:], seq_lens[num_contexts:]): slot = self._req_to_slot.get(rid) - if slot is None or rid == ATTENTION_DP_DUMMY_REQUEST_ID: - continue - token_offset = sum(spec_metadata.seq_lens[:request_idx]) - updates[slot] = min(int(spec_metadata.host_position_ids[token_offset]), self._max_ctx) + 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: