diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 8f98f7d4d68a..eb2f3352be9b 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -156,7 +156,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/ * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model. * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config. * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config. -* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a private paged context K/V cache and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. +* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. ```python from tensorrt_llm.llmapi import DFlashDecodingConfig @@ -170,6 +170,37 @@ 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, direct worker launches require +`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` on both the context and generation workers +to enable shared draft-KV transfer. The `disaggregated_mpi_worker` launcher +enables this setting automatically when both servers enable compatible +speculation and use the same `tensor_parallel_size` and +`enable_attention_dp`. + +When sharing is enabled, the context server's draft pages are transferred to +the generation server with the target pages, so the drafter starts generation +with the prompt context instead of an empty one. This requires the draft layers +to have the target's rank-local KV-head count; the draft dtype and head +dimension may differ, except that an NVFP4 target KV cache requires an +identical draft layout. Both servers must use the same layout. + +The drafter keeps a separate draft KV cache manager instead when any of the +following applies: context parallelism larger than 1, a per-layer KV-head count +on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba +models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout +that cannot be transferred with the target's. When the layers are shared, the +worker log reports `target and draft KV layers share one KV cache manager` at +startup. + [DFlash 2](https://inco.ai/blog/dflash2/) is also supported. The same `DFlashDecodingConfig` can be used for DFlash 2; no extra arguments are required. [LiLiCorr](https://arxiv.org/abs/2608.20530) checkpoints use the same diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index c42f48995f46..a3cd0364cbee 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -604,6 +604,12 @@ def _run_mla_decode( params: FmhaParams, kernel_dtype: torch.dtype, ) -> None: + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheManagerV2, + Role, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + attn = params.attn meta = params.meta @@ -680,6 +686,21 @@ def _run_mla_decode( page_table = self._generation_page_table( page_table_layer, meta.num_contexts, batch_size, layer_in_pool ) + if ( + layers_in_pool == 1 + and isinstance(meta.kv_cache_manager, KVCacheManagerV2) + and meta.kv_cache_manager._use_per_layer_page_tables + ): + # V2 per-layer indices include the layer offset, but get_buffers() + # already starts at this layer's KEY buffer. Apply the offset once. + converter = meta.kv_cache_manager.impl.get_page_index_converter( + local_layer_idx, Role.KEY + ) + layer_offset = int(converter.layer_offset) + if layer_offset: + page_table = torch.where( + page_table == BAD_PAGE_INDEX, page_table, page_table - layer_offset + ) # KVCacheManager exposes NHD pages as [num_pages, 1, page_size, 1, head_dim]. # The CuTe DSL kernel consumes a paged [page_size, dim, num_pages] view diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..c3126e9a08da 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -180,6 +180,29 @@ 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; " + "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." + ) + if self._ri.num_draft_layers: + if self._ri.tp_size != peer_ri.tp_size: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs equal tp_size " + f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); " + "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers." + ) + if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs the same attention DP " + f"setting (local={self._ri.attention.enable_attention_dp}, " + f"peer={peer_ri.attention.enable_attention_dp}); unset " + "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers." + ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 87731d303204..358b1b576412 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -49,6 +49,9 @@ class RankInfo: aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None bounded_replay_on_generation: bool = False + # Layers a unified KV cache manager owns for the drafter (layer specs with + # a DRAFT owner). They ride the target's pages, so both peers must agree. + num_draft_layers: int = 0 @property def tp_size_per_dp_group(self) -> int: @@ -114,6 +117,9 @@ def from_kv_cache_manager( aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False), + num_draft_layers=( + len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0 + ), ) @classmethod diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 70c008a4069f..6ea114c61fb1 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -41,6 +41,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 ReusePolicy from ..attention.backends import get_sparse_attn_kv_cache_manager from ..attention.backends.sparse.params import \ @@ -51,9 +52,9 @@ from ..model_config import ModelConfig, _get_kv_cache_layer_specs from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs -from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, - should_use_separate_draft_kv_cache) +from ..speculative import (draft_prompt_lookahead, get_draft_block_width, + get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, should_use_separate_draft_kv_cache) from . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, @@ -66,7 +67,9 @@ from .engine.lora import make_lora_model_config from .engine.multimodal import is_multimodal, use_mrope 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, @@ -824,6 +827,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: ModelConfig, max_seq_len: int, @@ -981,6 +1013,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 @@ -1148,6 +1182,13 @@ def _per_manager_cache_cost(self, self._kv_cache_config) if tokens_per_block is None: tokens_per_block = self._tokens_per_block + 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) + extra_kwargs.update( + shared_draft_block_width=self._unified_drafter_block_width()) if not is_draft: derived_windows = _derive_v2_layer_type_attention_windows( kv_cache_config, manager_cls, model_config, self._max_seq_len) @@ -1656,12 +1697,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 - num_pool_groups = _get_num_pool_groups_for_estimation( - model_cfg, - self._model_engine.max_seq_len, - self._kv_cache_config, - ) + 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 + num_pool_groups = _get_num_pool_groups_for_estimation( + model_cfg, + self._model_engine.max_seq_len, + self._kv_cache_config, + ) num_cache_blocks *= num_pool_groups # Dummy context requests use the configured maximum beam width. Scale @@ -2113,6 +2162,18 @@ 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 + shared_draft_block_width = None + if self._should_unify_dflash_kv_cache(): + unified_layer_specs = self._build_unified_dflash_layer_specs( + kv_cache_config) + shared_draft_block_width = self._unified_drafter_block_width() + logger.info( + f"{self._unified_drafter_name()}: target and draft KV layers " + "share one KV cache manager " + f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})." + ) kv_cache_manager = _create_kv_cache_manager( model_engine=self._model_engine, kv_cache_manager_cls=kv_cache_manager_cls, @@ -2137,6 +2198,8 @@ 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, + shared_draft_block_width=shared_draft_block_width, ) if not self._skip_est: @@ -2166,6 +2229,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 @@ -2191,6 +2256,274 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) + def _unified_drafter_attention_backend(self) -> Optional[str]: + """Resolved block-attention backend of a drafter that reads paged draft KV. + + None for every other drafter. That includes the embedded DSpark draft, + whose context lives in a worker-owned rolling window. + """ + spec_config = self._speculative_config + spec_dec_mode = spec_config.spec_dec_mode + if spec_dec_mode.is_dflash(): + return spec_config.attention_backend + if (not spec_dec_mode.is_dspark() + or spec_config.draft_is_embedded_in_target): + return None + backend = spec_config.attention_backend + if backend == "AUTO": + # The drafter resolves AUTO against its own family and this build. + draft_model = getattr(self._model_engine.model, "draft_model", None) + backend = getattr(draft_model, "dflash_attention_backend", None) + return backend + + def _unified_drafter_block_width(self) -> int: + """Use the loaded checkpoint's slot convention for sizing and allocation.""" + draft_model = getattr(self._model_engine.model, "draft_model", None) + shift_label = getattr(draft_model, "_dspark_shift_label", False) + return get_draft_block_width(self._speculative_config, + shift_label=shift_label) + + def _unified_drafter_name(self) -> str: + return ("DSpark" if self._speculative_config.spec_dec_mode.is_dspark() + else "DFlash") + + 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 + is "1" or "auto". + """ + if self._dflash_unify_blocker() is not None: + return False + return not (self._is_disagg + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") + not in ("1", "auto")) + + def _check_shared_draft_kv_request(self) -> None: + """Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored.""" + 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_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 _has_variable_kv_cache_geometry( + self._model_engine.model.model_config): + return "the target uses different KV cache geometry across layers" + if self._llm_args.kv_cache_compression_config is not None: + return "KV cache compression is enabled" + # Specialized managers add cache roles and lifecycle protocols that an + # external drafter cannot yet describe. Keep their existing fallback. + if self._kv_cache_manager_cls is not KVCacheManagerV2: + return "the target does not use KVCacheManagerV2" + target_layout = self._dflash_kv_layout( + self._model_engine.model.model_config, self._model_engine.dtype) + draft_model_config = self._get_draft_kv_model_config() + draft_layout = self._dflash_kv_layout( + draft_model_config, + draft_model_config.pretrained_config.torch_dtype) + if target_layout is None or draft_layout is None: + logger.info_once( + f"{self._unified_drafter_name()}: variable per-layer KV " + "geometry is not supported in a shared " + "target and draft KV cache; keeping a separate draft KV cache.", + key="dflash_per_layer_kv_heads_fallback") + return "the target or drafter has variable per-layer KV geometry" + # The drafter's shared context block table indexes all its layers + # through one lifecycle; mixed windows need the separate manager. + draft_windows = _resolve_v2_max_attention_window_vec( + self._get_draft_max_attention_window(self._unification_max_seq_len, + self._kv_cache_config), + self._unification_max_seq_len, + list( + self._mapping.pp_layers( + int(draft_model_config.pretrained_config.num_hidden_layers)) + )) + if len(set(draft_windows)) > 1: + return "the draft layers use different attention windows" + if self._is_disagg and not self._dflash_disagg_layouts_compatible( + target_layout, draft_layout): + logger.info_once( + f"{self._unified_drafter_name()}: draft KV layout " + "(kv_heads, head_dim, dtype, k_only)=" + f"{draft_layout} cannot be transferred with the target's " + f"{target_layout}; disaggregated serving keeps a separate " + "draft KV cache.", + key="dflash_disagg_draft_layout_fallback") + return (f"the draft KV layout {draft_layout} cannot be transferred " + f"with the target's {target_layout} " + "(kv_heads, head_dim, dtype, k_only)") + return None + + @staticmethod + def _dflash_disagg_layouts_compatible( + target_layout: tuple[int, int, object, bool], + draft_layout: tuple[int, int, object, bool]) -> bool: + """Whether the transceiver can move draft layers with the target's. + + 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 geometry varies per layer. + """ + config = model_config.pretrained_config + k_only = is_mla(config) + model_layer_specs = _get_kv_cache_layer_specs(model_config) + if k_only: + heads = 1 + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + elif model_layer_specs: + geometries = {(spec.num_kv_heads, spec.head_dim) + for spec in model_layer_specs} + if len(geometries) != 1: + return None + heads, head_dim = next(iter(geometries)) + else: + heads = getattr(config, "num_key_value_heads", + config.num_attention_heads) + if not isinstance(heads, int): + return None + head_dim = 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=ReusePolicy.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=ReusePolicy.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. @@ -2268,13 +2601,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() - or quant_config.quant_mode.has_fp4_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/nvfp4 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 @@ -3327,7 +3659,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, + shared_draft_block_width: Optional[int] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -3369,10 +3703,10 @@ def _create_kv_cache_manager( hidden_size = config.hidden_size - layer_specs = _get_kv_cache_layer_specs(_model_config) - if layer_specs: - head_dim = [spec.head_dim for spec in layer_specs] - num_key_value_heads = [spec.num_kv_heads for spec in layer_specs] + model_layer_specs = _get_kv_cache_layer_specs(_model_config) + if model_layer_specs: + head_dim = [spec.head_dim for spec in model_layer_specs] + num_key_value_heads = [spec.num_kv_heads for spec in model_layer_specs] else: num_attention_heads = config.num_attention_heads num_key_value_heads = num_kv_heads if num_kv_heads is not None else getattr( @@ -3480,6 +3814,9 @@ 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 + manager_extra_kwargs[ + "shared_draft_block_width"] = shared_draft_block_width # Vocab size also enables multimodal event decoding and its per-block # digest scan. Leave it unset for text-only engines. Without an engine # (e.g. separate one-model draft caches), preserve config resolution: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 42b621bde4c1..359f2145c57d 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 @@ -651,10 +651,49 @@ 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, shared_draft_block_width: Optional[int] = None +) -> int: + """Block width of the drafter whose layers share this manager, else 0. + + The creator supplies the checkpoint-resolved width when available. + """ + from tensorrt_llm._torch.speculative import get_draft_block_width + + if layer_specs is not None and any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in layer_specs + ): + if shared_draft_block_width is not None: + return shared_draft_block_width + return get_draft_block_width(spec_config) + return 0 + + +def _get_kv_alloc_extra_tokens( + spec_config, layer_specs, shared_draft_block_width: Optional[int] = None +) -> int: + """Tokens each request is allocated beyond its history, for spec decoding. + + ``get_num_extra_kv_tokens`` is the position-side count. When the drafter's + 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, shared_draft_block_width), + ) + + +def _get_generation_kv_capacity( + spec_config, + *, + is_draft: bool, + layer_specs=None, + shared_draft_block_width: Optional[int] = None, +) -> tuple[int, int]: + """Return draft-token reserve and total capacity headroom over history.""" if spec_config is None: return 0, BASE_GENERATION_TOKEN_COUNT reserve = spec_config.max_total_draft_tokens @@ -667,7 +706,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, shared_draft_block_width) + + spec_config.tokens_per_gen_step + + dynamic_reserve ) return reserve, headroom @@ -1785,6 +1826,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) + # Mixed-precision state as seen by instances built without __init__; # _init_mixed_precision assigns fresh per-instance values. low_precision_dtype: Optional[DataType] = None @@ -1852,6 +1905,7 @@ def __init__( joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> None: self.mapping = mapping @@ -2021,6 +2075,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, shared_draft_block_width + ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 @@ -2047,7 +2104,12 @@ def __init__( ) # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( - _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) + _get_generation_kv_capacity( + spec_config, + is_draft=self.is_draft, + layer_specs=self.layer_specs, + shared_draft_block_width=self._shared_draft_block_width, + ) ) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size @@ -2714,7 +2776,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() @@ -3701,7 +3763,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)] @@ -3722,7 +3784,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, ) ] @@ -4575,7 +4637,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 = ( @@ -5255,7 +5317,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 @@ -5295,7 +5357,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: @@ -5332,7 +5394,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: else (req.total_input_len_cp if self._has_cp_helix else req.prompt_len) ) prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens + target = prompt_len + get_draft_token_length(req) + self._kv_alloc_extra_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -5918,7 +5980,7 @@ def _resize_for_served_prefix(self, req: LlmRequest, kv_cache, position: int, en sliding-window layer group. When the capacity is unchanged the call cannot fail, since no new slot is requested. """ - target = end + self.num_extra_kv_tokens + target = end + self._kv_alloc_extra_tokens if end >= req.prompt_len: target += get_draft_token_length(req) pre_cap = kv_cache.capacity @@ -6064,7 +6126,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( @@ -6942,7 +7004,7 @@ def release_resources( if not success: release_resources(req) return None - 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 @@ -7538,6 +7600,7 @@ def get_cache_size_per_token( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ): if layer_specs is not None: @@ -7573,7 +7636,10 @@ 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, + shared_draft_block_width=shared_draft_block_width, ) ( context_size_per_token, @@ -7656,6 +7722,7 @@ def _get_per_request_layer_tokens( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, ) -> Optional[tuple[List[tuple[int, int, int]], int]]: """Per-local-layer ``(layer_size, retained_tokens, intercept_tokens)`` plus the block-rounded full-sequence token count. @@ -7759,7 +7826,10 @@ def _get_per_request_layer_tokens( else: retained_windows = list(covered_windows) _, generation_capacity_headroom = _get_generation_kv_capacity( - spec_config, is_draft=is_draft + spec_config, + is_draft=is_draft, + layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) full_tokens = ( math.ceil((max_seq_len + generation_capacity_headroom) / tokens_per_block) @@ -7804,6 +7874,7 @@ def get_cache_bytes_per_request( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -7837,6 +7908,7 @@ def get_cache_bytes_per_request( spec_config=spec_config, is_draft=is_draft, layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) if layer_tokens is None: return None @@ -7866,6 +7938,7 @@ def get_cache_quota_per_request( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + shared_draft_block_width: Optional[int] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -7901,6 +7974,7 @@ def get_cache_quota_per_request( spec_config=spec_config, is_draft=is_draft, layer_specs=layer_specs, + shared_draft_block_width=shared_draft_block_width, ) if layer_tokens is None: return None diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 4cf174f17fd0..c9ce9f4b7fa6 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -1018,9 +1018,10 @@ def _try_schedule_disagg_gen_init( Returns ``(action, tokens)``. *tokens* is 0 because disagg requests don't participate in the forward pass token budget. """ - # Cache-transceiver mode disables the separate one-model draft manager, - # so disagg generation init has no paired-reuse path. Supporting one - # would also require draft KV transfer and history_length=prompt_len. + # Only the target manager receives prompt KV. Drafters that need it + # under disagg keep their layers in that manager (MTP/Eagle fold them + # in; DFlash uses explicit draft layer specs). A separate draft manager + # gets no transfer and no paired-reuse path here. if not self.kv_cache_manager.prepare_disagg_gen_init(req): logger.debug("prepare_disagg_gen_init failed for request %s", req.py_request_id) return ScheduleAction.SKIP, 0 diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 382cb09e1571..c3ffe7e28bda 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -18,10 +18,11 @@ SpecSampler) from .spec_tree_manager import SpecTreeManager from .suffix_automaton import SuffixAutomatonManager -from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, get_spec_drafter, - get_spec_metadata, get_spec_resource_manager, - get_spec_worker, update_spec_config_from_draft_model_config, +from .utils import (get_draft_block_width, get_draft_kv_cache_manager, + get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, get_spec_drafter, get_spec_metadata, + get_spec_resource_manager, get_spec_worker, + update_spec_config_from_draft_model_config, update_spec_config_from_loaded_model, update_spec_config_from_model_config) @@ -50,6 +51,7 @@ "SpecWorkerBase", "draft_prompt_lookahead", "get_draft_kv_cache_manager", + "get_draft_block_width", "get_num_extra_kv_tokens", "get_num_spec_layers", "get_spec_decoder", diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fb85eafd3cd8..fbd7c2a1b27b 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -31,6 +31,7 @@ from ..attention.backends import AttentionMetadata from ..modules.logits_processor import LogitsProcessor +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 @@ -615,24 +616,41 @@ def prepare(self): # prefilling, so _store_prefill_context -- the only other assigner -- # never runs and every request would land on the dummy slot. num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + assigned = {} + for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=assigned) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + elif worker._ctx_unified_cache: + # New requests have no in-flight predecessor. Subsequent + # accepted tokens advance the device length during replay. + token_offset = sum(self.seq_lens[:request_idx]) + assigned[slot] = min( + int(self.host_position_ids[token_offset]), worker._max_ctx + ) + worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) - # to dummy slot to avoid corrupting real request's context + # to dummy slot to avoid corrupting real request's context. A + # generation row with the idle id is a dummy even when a context + # warmup stored a prompt under that id: its length is not this row's. mapping = torch.tensor( - [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], + [ + worker._dummy_slot + if i >= num_contexts and rid == ATTENTION_DP_DUMMY_REQUEST_ID + else worker._req_to_slot.get(rid, worker._dummy_slot) + for i, rid in enumerate(self.request_ids) + ], dtype=torch.long, ) worker._batch_to_slot_buffer.update(mapping) @@ -794,6 +812,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 @@ -1083,7 +1104,18 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor # it: _util.py skips that bail for an external drafter, which is # what DSpark is, so this drafter does get a manager under DP. 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) @@ -1116,10 +1148,42 @@ def _init_ctx_block_tables( """ mgr = draft_kv_cache_manager pool_mapping = mgr.kv_cache_pool_mapping - layer_offset = mgr.layer_offsets[0] + first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 + layer_offset = mgr.layer_offsets[first_layer_id] self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - if num_pool_layers != L: + self._ctx_block_offset = 0 + if isinstance(mgr, KVCacheManagerV2): + converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) + layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + unsupported = None + if mgr.enable_swa_scratch_reuse: + unsupported = "DFlash managed context does not support SWA scratch pages" + for layer_id in self._ctx_layer_ids: + local_id = mgr.layer_offsets[layer_id] + if ( + mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id + or mgr.get_layer_page_index_scale(layer_id) != converter.scale + ): + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + break + if unsupported is not None: + # A unified manager was selected for this layout, so the drafter + # cannot move to private storage without breaking allocation, + # reuse and transfer. A separate manager keeps the private arena. + if self._ctx_unified_cache: + raise NotImplementedError(unsupported) + logger.warning(f"{unsupported}; using the private ctx arena.") + return False + if per_layer_tables: + # Attention tables address the physical pool origin, whereas + # get_buffers() starts at this layer's own K buffer. + self._ctx_block_offset = int(converter.layer_offset) + elif num_pool_layers != L: logger.warning( f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " f"the drafter has {L}; using the private ctx arena." @@ -1128,12 +1192,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 " @@ -1157,6 +1232,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. @@ -1166,22 +1258,30 @@ 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. The block count # is taken before the clamp erases the placeholders; it is informative # only on the python V2 backend, the cpp one already mapped # BAD_PAGE_INDEX to 0. + encoded = src[self._ctx_pool_idx, :num_seqs, 0] + if self._ctx_block_offset: + # Use int64 to avoid underflowing negative placeholders, and keep + # the metadata's table unchanged when removing the layer offset. + encoded = encoded.to(torch.int64) - self._ctx_block_offset decode_draft_block_table( - src[self._ctx_pool_idx, :num_seqs, 0], + encoded, self._ctx_block_divisor, self._ctx_block_tables, self._ctx_block_counts, @@ -1192,12 +1292,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." @@ -1212,6 +1318,28 @@ def _lazy_init_ctx_buffers( logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") self._ctx_buf_inited = False + # checkpoint's trained block width + self._resolved_block_size = getattr(draft_model, "block_size", None) or ( + self.max_draft_len + 1 + ) + # what the draft forward actually computes + self._compute_block_size = self._draft_block_width(draft_model) + if self._compute_block_size > self._resolved_block_size: + slack = self._compute_block_size - self.max_draft_len + raise ValueError( + f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." + f"Lower max_draft_len to at most {self._resolved_block_size - slack}." + f"Current max_draft_len={self.max_draft_len}." + ) + if self._ctx_unified_cache: + reserved = getattr(draft_kv_cache_manager, "_shared_draft_block_width", 0) + if self._compute_block_size > reserved: + raise ValueError( + f"Drafter block width {self._compute_block_size} exceeds the " + f"{reserved} positions reserved by the shared KV cache manager. " + "The manager must reserve the drafter's resolved block width." + ) + # Recorded even when the buffers stay private, so that the estimation # pass falling back to the arena still rebinds once the real pool lands. self._ctx_kv_manager = draft_kv_cache_manager @@ -1281,20 +1409,6 @@ def _lazy_init_ctx_buffers( self._req_to_slot = {} self._req_ctx_pos = {} - # checkpoint's trained block width - self._resolved_block_size = getattr(draft_model, "block_size", None) or ( - self.max_draft_len + 1 - ) - # what the draft forward actually computes - self._compute_block_size = self._draft_block_width(draft_model) - if self._compute_block_size > self._resolved_block_size: - slack = self._compute_block_size - self.max_draft_len - raise ValueError( - f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." - f"Lower max_draft_len to at most {self._resolved_block_size - slack}." - f"Current max_draft_len={self.max_draft_len}." - ) - # +block_size slack lets each backend append the per-iteration noise # K/V after the accumulated request context. assert hasattr(draft_model, "_build_fused_kv_buffers"), ( @@ -1306,6 +1420,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 @@ -1361,6 +1476,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 @@ -1399,7 +1516,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( @@ -1819,10 +1936,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. @@ -2067,7 +2192,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 @@ -2342,7 +2472,7 @@ def prepare_1st_drafter_inputs( projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) if self._ctx_block_tables is not None: # Bound by what the manager handed THIS request: the masked @@ -2401,6 +2531,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() @@ -2447,6 +2581,9 @@ def _fused_gen_prepare_applies(self, spec_metadata) -> bool: features, drafting from all ``K + 1`` target tokens.""" return ( self._paged_latent_prepare + # Shared allocations require the generic path's strict capacity + # check; the fused preparation silently caps the context length. + and not self._ctx_unified_cache and spec_metadata.runtime_draft_len + 1 == self._draft_tokens_per_req and spec_metadata.captured_hidden_states is not None ) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 757b84dc3124..fe7a12faaffe 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -1240,6 +1240,15 @@ def bind_draft_lm_head(self, draft_model) -> None: ): self._markov_tp_exchange = tp_pair_exchange(mapping, AllReduceStrategy.NCCL) + @staticmethod + def block_width_for(spec_config, *, shift_label: bool = True) -> int: + """Block slots a DSpark worker computes for ``spec_config``. + + The ``shift_label`` convention reads slots 0..K-1; a checkpoint that + turns it off reads slots 1..K and needs K+1 positions. + """ + return spec_config.max_draft_len + (0 if shift_label else 1) + def _draft_block_width(self, draft_model) -> int: """Block width under the dspark ``shift_label`` convention. @@ -1248,9 +1257,9 @@ def _draft_block_width(self, draft_model) -> int: checkpoint at max_draft_len=7, which is how both published DSpark drafters are meant to run. """ - if getattr(draft_model, "_dspark_shift_label", False): - return self.max_draft_len - return super()._draft_block_width(draft_model) + return self.block_width_for( + self.spec_config, shift_label=getattr(draft_model, "_dspark_shift_label", False) + ) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 97b67a555523..f5c6c1aa82d9 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -179,7 +179,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 @@ -3239,8 +3240,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) @@ -3261,7 +3264,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 f8e1a7340145..94cf65bd7c4d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -786,6 +786,23 @@ def get_num_extra_kv_tokens(spec_config): return 0 +def get_draft_block_width(spec_config, *, shift_label: bool = True) -> int: + """Block width the worker drafting for ``spec_config`` appends per request. + + 0 when that worker does not append a block to a paged draft KV cache. + ``shift_label`` selects the checkpoint's slot convention for standalone DSpark. + """ + if spec_config is None: + return 0 + 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, + shift_label=shift_label) + return 0 + + def get_draft_kv_cache_manager(spec_config, resource_manager): """ Return the one-model draft cache manager, including a shared draft KV view @@ -804,7 +821,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 1a0f80f8c3b5..76144fa44294 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2682,6 +2682,32 @@ def set_cuda_device(): torch.cuda.set_device(device_id) +def _shared_draft_kv_blocker(server_configs) -> Optional[str]: + """What keeps the context and generation servers from sharing draft KV, or None.""" + with_spec = { + cfg.type + for cfg in server_configs if cfg.other_args.get("speculative_config") + } + if with_spec != {"ctx", "gen"}: + return ("speculative_config in both context_servers and " + "generation_servers") + if len({(cfg.other_args.get("tensor_parallel_size", 1), + bool(cfg.other_args.get("enable_attention_dp", False))) + for cfg in server_configs}) != 1: + return ("the same tensor_parallel_size and enable_attention_dp in " + "context_servers and generation_servers") + for cfg in server_configs: + spec = cfg.other_args["speculative_config"] + backend = spec.get("attention_backend") + if not ((spec.get("decoding_type") == "DFlash" and backend == "TRTLLM") + or (spec.get("decoding_type") == "DSpark" + and backend in (None, "AUTO", "TRTLLM"))): + return ("speculative_config to be DFlash with attention_backend " + "TRTLLM, or DSpark with attention_backend AUTO or TRTLLM, " + "in both sections") + return None + + @click.command("disaggregated_mpi_worker") @stability_option( "-c", @@ -2721,6 +2747,14 @@ def disaggregated_mpi_worker(config_file: Optional[str], log_level: str, from tensorrt_llm.llmapi.disagg_utils import split_world_comm disagg_cfg = parse_disagg_config_file(config_file) + blocker = _shared_draft_kv_blocker(disagg_cfg.server_configs) + 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}.") + if os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") in (None, "auto"): + os.environ["TRTLLM_DISAGG_SHARED_DRAFT_KV"] = ("0" + if blocker else "auto") # Run a server with the underlying LLM invokes a RemoteMPISessionClient if os.environ.get(DisaggLauncherEnvs. diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 97cedd3d9a2d..fc3eaad3ac9b 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -137,6 +137,12 @@ ModelingV2DeepseekR10528Nvfp4Sm10xDefault::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 TestModelingV2Glm52Nvfp4Sm103Tep4:: # test_gsm8k_identity_vs_mtp1_and_acceptance. Keyed by target and variant for # the same reason as the ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 4bf85c4041ce..2cdb2f9548ed 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 @@ -1648,6 +1649,96 @@ def test_gen_first(self, enable_block_reuse): run_accuracy_test(llm, self.MODEL_NAME, ["MMLU"]) +@skip_pre_blackwell +@pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) +class TestQwen3_8BDFlash(LlmapiAccuracyTestHarness): + MODEL_NAME = "Qwen3/Qwen3-8B" + MODEL_PATH = f"{llm_models_root()}/Qwen3/Qwen3-8B" + DRAFT_PATH = f"{llm_models_root()}/Qwen3-8B-DFlash-b16" + MAX_DRAFT_LEN = 15 + + @pytest.mark.skip_less_device(2) + def test_gsm8k_1p1d_unified_draft_kv(self): + """DFlash draft KV travels with the target KV from ctx to gen. + + The drafter's layers share the target's KV cache manager, so the + generation worker drafts from the prompt context computed on the + context worker. Verified speculative decoding is lossless, so the + GSM8K score cannot see a drafter that starts from an empty context; + acceptance length can, so it is gated too. + """ + speculative_config = { + "decoding_type": "DFlash", + "max_draft_len": self.MAX_DRAFT_LEN, + "speculative_model": self.DRAFT_PATH, + # The unified draft KV layout requires the TRTLLM drafter backend. + "attention_backend": "TRTLLM", + } + # Both workers run the drafter: the context worker writes the + # drafter's prompt KV into the shared manager, which is then + # transferred with the target KV. + common_server_config = { + "disable_overlap_scheduler": True, + "max_batch_size": 16, + "max_seq_len": 4096, + "max_num_tokens": 4096, + "kv_cache_config": { + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.5, + }, + # KVCacheManagerV2 requires the Python transceiver. + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + }, + "speculative_config": speculative_config, + } + ctx_server_config = {**common_server_config, "cuda_graph_config": None} + gen_server_config = { + **common_server_config, + "cuda_graph_config": { + "max_batch_size": 16, + "enable_padding": True, + }, + # Acceptance length is read back off this worker's /metrics. + "enable_iter_perf_stats": True, + "max_stats_len": -1, + "iter_stats_max_iterations": -1, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + }, + } + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + assert_worker_log_contains= + "target and draft KV layers share one KV cache manager", + extra_env={"TRTLLM_DISAGG_SHARED_DRAFT_KV": "1"}) as llm: + llm.args.speculative_config = DFlashDecodingConfig( + max_draft_len=self.MAX_DRAFT_LEN, + speculative_model=self.DRAFT_PATH, + ) + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + + # Drafting from an empty prompt context costs ~17% acceptance + # length on this model, well below the 95% floor. + acceptance_length = compute_disagg_acceptance_length(llm.serve_url) + print(f"[AL] test_gsm8k_1p1d_unified_draft_kv " + f"acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length( + "TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv", + acceptance_length) + + @skip_pre_blackwell @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) class TestQwen3_30B_A3B(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 5dfdc6fbcdd4..17cff489f10a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -31,6 +31,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/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index a1359591f59f..367337f1a4d9 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX def test_flashinfer_uses_v2_page_index_upper_bound_directly() -> None: @@ -258,6 +259,93 @@ def generation_preprocess(*args: object) -> tuple[object, ...]: assert len(decode_calls) == 1 +@pytest.mark.cpu_only +@pytest.mark.parametrize("local_layer_idx", [0, 1]) +def test_cute_dsl_mla_v2_per_layer_indices_read_correct_pages( + monkeypatch: pytest.MonkeyPatch, local_layer_idx: int +) -> None: + """Per-layer indices and layer-origin buffers must select the same physical page.""" + page_size, latent_dim, rope_dim = 2, 2, 2 + layer_idx = 16 + local_layer_idx + physical_pool = ( + torch.arange(10, dtype=torch.bfloat16) + .view(10, 1, 1, 1, 1) + .expand(-1, 1, page_size, 1, latent_dim + rope_dim) + .contiguous() + ) + # Two MLA target layers share a pool. Mixed target/draft layouts use + # per-layer tables: slot 3 maps to physical pages 6 and 7 for these layers. + block_offsets = torch.tensor( + [[[[6, BAD_PAGE_INDEX]]], [[[7, BAD_PAGE_INDEX]]]], dtype=torch.int32 + ) + original_offsets = block_offsets.clone() + # V2 wraps the shifted raw pointer as storage origin; a slice alone would + # retain the full pool's storage and let an incorrect rebase hide the bug. + layer_buffer = physical_pool[local_layer_idx:].clone() + + def get_buffers(requested_layer_idx: int) -> torch.Tensor: + assert requested_layer_idx == layer_idx + return layer_buffer + + def get_page_index_converter(requested_local_idx: int, role: object) -> SimpleNamespace: + assert (requested_local_idx, role) == (local_layer_idx, Role.KEY) + return SimpleNamespace(layer_offset=local_layer_idx) + + manager = object.__new__(KVCacheManagerV2) + manager._use_per_layer_page_tables = True + manager.layer_offsets = {16: 0, 17: 1} + manager.get_buffers = get_buffers + manager.impl = SimpleNamespace(get_page_index_converter=get_page_index_converter) + metadata = SimpleNamespace( + kv_cache_manager=manager, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_mapping=torch.tensor([[0, 0], [1, 0]], dtype=torch.int32), + num_contexts=0, + max_num_requests=1, + helix_position_offsets=None, + ) + attn = SimpleNamespace( + layer_idx=layer_idx, + get_local_layer_idx=lambda meta: meta.kv_cache_manager.layer_offsets[layer_idx], + num_heads=1, + kv_lora_rank=latent_dim, + qk_rope_head_dim=rope_dim, + qk_nope_head_dim=latent_dim, + q_scaling=1.0, + ) + q = torch.zeros((1, 1, latent_dim + rope_dim), dtype=torch.bfloat16) + output = torch.empty((1, 1, latent_dim), dtype=torch.bfloat16) + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=AttentionForwardArgs( + output=output, attention_input_type=AttentionInputType.generation_only + ), + workspace=torch.empty(0, dtype=torch.uint8), + sequence_lengths=torch.tensor([page_size], dtype=torch.int32), + tokens_per_block=page_size, + num_requests=1, + ) + pages_read: list[torch.Tensor] = [] + + def read_page(q_latent, q_rope, c_latent, c_rope, page_table, *args) -> None: + assert page_table[1, 0].item() == BAD_PAGE_INDEX + page = page_table[0, 0].item() + pages_read.append(torch.cat((c_latent[:, :, page], c_rope[:, :, page]), dim=1)) + + monkeypatch.setattr( + torch.ops.trtllm, "cute_dsl_mla_decode_fp16_blackwell", read_page, raising=False + ) + monkeypatch.setattr(CuteDslMlaFmha, "_required_workspace_size", staticmethod(lambda *args: 0)) + monkeypatch.setattr(CuteDslMlaFmha, "_split_kv_plan", staticmethod(lambda *args: None)) + fmha = object.__new__(CuteDslMlaFmha) + fmha._run_mla_decode(q, output, params, torch.bfloat16) + + assert len(pages_read) == 1 + torch.testing.assert_close(pages_read[0], physical_pool[6 + local_layer_idx, 0, :, 0, :]) + torch.testing.assert_close(block_offsets, original_offsets) + + def _cute_dsl_mla_helix_support( monkeypatch: pytest.MonkeyPatch, *, diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py new file mode 100644 index 000000000000..b0d04068754a --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_dflash_unified_kv_gate.py @@ -0,0 +1,390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Selection of the unified target+DFlash KV cache layout.""" + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import click +import pytest +import torch +import yaml + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.quantization.mode import QuantMode + +pytestmark = pytest.mark.cpu_only + +_SHARED_DRAFT_KV_ENV = "TRTLLM_DISAGG_SHARED_DRAFT_KV" + + +@pytest.fixture(autouse=True) +def _shared_draft_kv_allowed(monkeypatch): + # Tests of the layout rules assume a launcher that allows sharing. + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, "1") + + +def _model_config(*, kv_heads, head_dim, kv_quant_algo=None): + pretrained = SimpleNamespace( + num_hidden_layers=2, + num_attention_heads=32, + num_key_value_heads=kv_heads, + head_dim=head_dim, + hidden_size=32 * head_dim, + torch_dtype=torch.bfloat16, + ) + quant_mode = QuantMode.from_quant_algo(kv_cache_quant_algo=kv_quant_algo) + return SimpleNamespace( + pretrained_config=pretrained, + quant_config=SimpleNamespace(quant_mode=quant_mode), + ) + + +def _creator(*, is_disagg, target, draft, tp_size=1, draft_windows=None): + creator = object.__new__(KvCacheCreator) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dflash=lambda: True, is_external_drafter=lambda: True), + attention_backend="TRTLLM", + _allow_separate_draft_kv_cache=True, + draft_len_schedule=None, + ) + creator._draft_config = draft + creator._mapping = SimpleNamespace( + pp_size=1, + cp_size=1, + tp_size=tp_size, + enable_attention_dp=False, + pp_layers=lambda num_layers: list(range(num_layers)), + ) + creator._max_seq_len = 4096 + creator._kv_cache_config = SimpleNamespace() + creator._get_draft_max_attention_window = lambda max_seq_len, kv_cache_config: draft_windows + creator._llm_args = SimpleNamespace(attn_backend="TRTLLM", kv_cache_compression_config=None) + creator._model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=target), dtype=torch.bfloat16 + ) + creator._kv_cache_manager_cls = KVCacheManagerV2 + creator._kv_cache_config = SimpleNamespace(max_attention_window=None) + creator._unification_max_seq_len = 4096 + creator._is_disagg = is_disagg + creator._get_draft_kv_model_config = lambda: draft + return creator + + +@pytest.mark.parametrize("is_disagg", [False, True]) +def test_matching_draft_layout_unifies(is_disagg): + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + + assert _creator(is_disagg=is_disagg, target=target, draft=draft)._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("shift_label, expected_width", [(True, 7), (False, 8)]) +def test_dspark_checkpoint_width_reaches_shared_cache_cost_and_construction( + monkeypatch, is_disagg, shift_label, expected_width +): + from tensorrt_llm._torch.pyexecutor import _util + + layout = _model_config(kv_heads=8, head_dim=128) + layout.is_generation = True + creator = _creator(is_disagg=is_disagg, target=layout, draft=layout) + creator._speculative_config.spec_dec_mode = SimpleNamespace( + is_dflash=lambda: False, is_dspark=lambda: True, is_external_drafter=lambda: True + ) + creator._speculative_config.max_draft_len = 7 + creator._speculative_config.draft_is_embedded_in_target = False + creator._model_engine.model.draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + creator._kv_cache_config = KvCacheConfig() + creator._tokens_per_block = 32 + creator._max_batch_size = 4 + creator._max_num_tokens = 4096 + creator._max_beam_width = 1 + creator._sparse_attention_config = None + creator._kv_connector_manager = None + creator._execution_stream = None + creator._disable_overlap_scheduler = True + creator._skip_est = True + creator._llm_args.kv_cache_config = creator._kv_cache_config + monkeypatch.setattr(creator, "_get_model_kv_cache_manager_cls", lambda *args: KVCacheManagerV2) + monkeypatch.setattr(creator, "_enable_kv_cache_stats", lambda: False) + + class RecordingManager(KVCacheManagerV2): + get_cache_size_per_token = Mock(return_value=(1, 0)) + get_cache_bytes_per_request = Mock(return_value=1) + get_cache_quota_per_request = Mock(return_value=1) + + assert creator._should_unify_dflash_kv_cache() + assert not creator._should_create_separate_draft_kv_cache() + creator._per_manager_cache_cost(RecordingManager, layout, include_request_cost=True) + create_manager = Mock(return_value=SimpleNamespace(max_seq_len=4096)) + monkeypatch.setattr(_util, "_create_kv_cache_manager", create_manager) + creator._create_kv_cache_manager() + + runtime_specs = create_manager.call_args.kwargs["layer_specs"] + for call in ( + RecordingManager.get_cache_size_per_token, + RecordingManager.get_cache_bytes_per_request, + RecordingManager.get_cache_quota_per_request, + create_manager, + ): + assert call.call_count == 1 + assert call.call_args.kwargs["shared_draft_block_width"] == expected_width + assert call.call_args.kwargs["layer_specs"] == runtime_specs + + +def test_mismatched_draft_kv_heads_unify_only_in_aggregated_serving() -> None: + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=4, head_dim=128) + + 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]) +@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 = attention_dp + + heads = ( + creator._dflash_kv_layout(target, torch.bfloat16)[0], + creator._dflash_kv_layout(draft, torch.bfloat16)[0], + ) + assert heads == ((4, 1) if attention_dp else (1, 1)) + assert creator._should_unify_dflash_kv_cache() == (not (is_disagg and attention_dp)) + + +@pytest.mark.parametrize("is_disagg", [False, True]) +@pytest.mark.parametrize("per_layer_side", ["target", "draft"]) +def test_per_layer_kv_heads_keep_separate_cache(is_disagg, per_layer_side): + uniform = _model_config(kv_heads=8, head_dim=128) + per_layer = _model_config(kv_heads=[8, 4], head_dim=128) + target, draft = (per_layer, uniform) if per_layer_side == "target" else (uniform, per_layer) + + creator = _creator(is_disagg=is_disagg, target=target, draft=draft) + + assert not creator._should_unify_dflash_kv_cache() + assert creator._build_unified_dflash_layer_specs(kv_cache_config=None) == () + + +class _StopLauncher(Exception): + pass + + +def _launch_from_disagg_config(monkeypatch, tmp_path, context_servers, generation_servers): + """Run the real disaggregated_mpi_worker up to the rank-communicator split.""" + from tensorrt_llm.commands import serve + + config = tmp_path / "disagg.yaml" + config.write_text( + yaml.safe_dump( + { + "context_servers": {"num_instances": 1, **context_servers}, + "generation_servers": {"num_instances": 1, **generation_servers}, + } + ) + ) + + def stop(*args, **kwargs): + raise _StopLauncher + + monkeypatch.setattr( + serve._command_telemetry, "apply_disaggregated_telemetry_config", lambda *a, **k: None + ) + monkeypatch.setattr(serve, "set_cuda_device", lambda: None) + monkeypatch.setattr("tensorrt_llm.llmapi.disagg_utils.split_world_comm", stop) + monkeypatch.delenv("TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT", raising=False) + serve.disaggregated_mpi_worker.callback( + config_file=str(config), log_level="info", telemetry=False + ) + + +def test_launcher_rejects_dflash_on_generation_worker_only( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + # The launcher rejects the variable when only one section sets speculative_config. + with pytest.raises(click.UsageError, match="both context_servers and generation_servers"): + _launch_from_disagg_config( + monkeypatch, + tmp_path, + context_servers={}, + generation_servers={"speculative_config": {"decoding_type": "DFlash"}}, + ) + + +def test_disagg_without_shared_draft_kv_env_keeps_separate_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + target = _model_config(kv_heads=8, head_dim=128) + draft = _model_config(kv_heads=8, head_dim=128) + + 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"} + 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}, + ) + + +@pytest.mark.parametrize( + "request_mode,draft_heads,expected_mode", + [("0", 8, "0"), (None, 4, "auto")], + ids=["explicit_opt_out", "auto_fallback"], +) +def test_launcher_preserves_opt_out_and_allows_layout_fallback( + monkeypatch, tmp_path, request_mode, draft_heads, expected_mode +) -> None: + if request_mode is None: + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV, raising=False) + else: + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, request_mode) + config = {"speculative_config": {"decoding_type": "DFlash", "attention_backend": "TRTLLM"}} + + with pytest.raises(_StopLauncher): + _launch_from_disagg_config(monkeypatch, tmp_path, config, config) + + assert os.environ[_SHARED_DRAFT_KV_ENV] == expected_mode + creator = _creator( + is_disagg=True, + target=_model_config(kv_heads=8, head_dim=128), + draft=_model_config(kv_heads=draft_heads, head_dim=128), + ) + creator._check_shared_draft_kv_request() + assert not creator._should_unify_dflash_kv_cache() diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_budget_split.py index e3ed6bf727c7..3fceaa1ada95 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_budget_split.py @@ -336,6 +336,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 drafter outside the engine. mode.use_one_engine.return_value = False + # A truthy DFlash or DSpark check would route costing through the shared + # target and drafter cache path. + mode.is_dflash.return_value = False + mode.is_dspark.return_value = False costed_windows: list[tuple[object, list[int] | None]] = [] class RecordingKVCacheManager(KVCacheManagerV2): diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py index 5c06f0766e20..7abe62ecfe9a 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py @@ -29,6 +29,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 @@ -37,6 +38,8 @@ KVCacheLayerSpec, KVCacheManagerV2, Role, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine @@ -45,6 +48,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( + DSparkDecodingConfig, KvCacheConfig, MTPDecodingConfig, MultimodalConfig, @@ -1310,6 +1314,85 @@ def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: ) == (0, (8 + 23) * 64 * 64) +@pytest.mark.parametrize( + "shared_draft_block_width,full_tokens,window_tokens", + [(None, 96, 64), (7, 96, 64), (8, 128, 96)], + ids=["config_default", "dspark_shifted", "dspark_unshifted"], +) +def test_v2_dspark_resolved_width_sizes_shared_cache_costs( + shared_draft_block_width: int | None, full_tokens: int, window_tokens: int +) -> None: + """One extra draft position can require a full page in every cost estimate.""" + spec_config = DSparkDecodingConfig(max_draft_len=7, speculative_model="draft") + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1024, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + ) + draft = replace( + target, global_layer_id=1, owner=KVCacheLayerOwner.DRAFT, sliding_window_size=19 + ) + kwargs = dict( + tokens_per_block=32, + max_seq_len=81, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + avg_seq_len=81, + max_util_for_resume=1.0, + ), + spec_config=spec_config, + layer_specs=[target, draft], + shared_draft_block_width=shared_draft_block_width, + ) + # A page is 2 MiB, so allocator granularity preserves the boundary difference. + layer_bytes = 64 << 10 + assert KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=32, **kwargs + ) == (layer_bytes, layer_bytes * (2 * window_tokens + 32)) + assert KVCacheManagerV2._get_per_request_layer_tokens(None, Mapping(), **kwargs) == ( + [(layer_bytes, full_tokens, 0), (layer_bytes, window_tokens, window_tokens)], + full_tokens, + ) + assert KVCacheManagerV2.get_cache_bytes_per_request(None, Mapping(), **kwargs) == ( + layer_bytes * full_tokens + ) + assert KVCacheManagerV2.get_cache_quota_per_request(None, Mapping(), **kwargs) == ( + layer_bytes * (full_tokens + window_tokens) + ) + + +@pytest.mark.parametrize("with_target_layer", [False, True]) +def test_v2_resolved_draft_width_does_not_change_target_only_headroom( + with_target_layer: bool, +) -> None: + spec_config = DSparkDecodingConfig(max_draft_len=7, speculative_model="draft") + layer_specs = None + if with_target_layer: + layer_specs = [ + KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + ) + ] + assert _get_kv_alloc_extra_tokens( + spec_config, layer_specs, shared_draft_block_width=8 + ) == _get_kv_alloc_extra_tokens(spec_config, layer_specs) + assert _get_generation_kv_capacity( + spec_config, + is_draft=False, + layer_specs=layer_specs, + shared_draft_block_width=8, + ) == _get_generation_kv_capacity(spec_config, is_draft=False, layer_specs=layer_specs) + + def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) manager.layer_specs = None @@ -1855,6 +1938,10 @@ def _create_manager_and_capture_config( class _RecordingManager(manager_cls): def __init__(self, kv_cache_config: KvCacheConfig, *args: object, **kwargs: object) -> None: captured.append(kv_cache_config) + assert kwargs.get("layer_specs") is factory_overrides.get("layer_specs") + assert kwargs.get("shared_draft_block_width") == factory_overrides.get( + "shared_draft_block_width" + ) model_config = Mock() model_config.pretrained_config = pretrained @@ -2044,6 +2131,58 @@ def test_create_kv_cache_manager_cross_pool_keeps_single_window_default( 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 + + @pytest.mark.parametrize("guard", ["draft", "explicit_window"]) def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) -> None: kv_cache_config = KvCacheConfig( @@ -2061,6 +2200,50 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None) +@pytest.mark.parametrize("publish_layer_specs", [False, True]) +@pytest.mark.parametrize("shared_draft_block_width", [None, 7, 8]) +def test_create_kv_cache_manager_preserves_explicit_storage_specs( + publish_layer_specs: bool, + shared_draft_block_width: int | None, +) -> None: + """Model geometry must not replace an explicit target and draft storage layout.""" + specs = tuple(_fp8_target_bf16_draft_specs(draft_window=128)) + _create_manager_and_capture_config( + _mixed_schedule_pretrained([_FULL] * 3), + KvCacheConfig(), + max_seq_len=2048, + publish_layer_specs=publish_layer_specs, + layer_specs=specs, + shared_draft_block_width=shared_draft_block_width, + ) + + +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_unified_layout_uses_model_provided_geometry(attention_dp: bool) -> None: + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + num_attention_heads=16, num_key_value_heads=8, head_dim=128 + ), + quant_config=None, + kv_cache_layer_specs=(ModelKVCacheLayerSpec(head_dim=256, num_kv_heads=4, num_q_heads=16),) + * 3, + ) + creator = object.__new__(KvCacheCreator) + creator._mapping = SimpleNamespace(tp_size=2, enable_attention_dp=attention_dp) + + assert creator._dflash_kv_layout(model_config, torch.bfloat16) == ( + 4 if attention_dp else 2, + 256, + DataType.BF16, + False, + ) + + model_config.kv_cache_layer_specs += ( + ModelKVCacheLayerSpec(head_dim=512, num_kv_heads=4, num_q_heads=16), + ) + assert creator._dflash_kv_layout(model_config, torch.bfloat16) is None + + def test_profiling_failure_releases_resources_and_restores_response_state(): creator = object.__new__(KvCacheCreator) creator._mapping = SimpleNamespace(cp_config={}) diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 83bd77f0c2f4..e83aab2ca7a3 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -24,6 +24,7 @@ import torch from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.pyexecutor._util import _get_num_storage_pool_groups from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( @@ -33,6 +34,8 @@ KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, + _get_generation_kv_capacity, + _get_kv_alloc_extra_tokens, _KVCacheManagerInitStatus, _swa_branch_snapshot_point, _swa_endpoint_priority, @@ -51,6 +54,7 @@ from tensorrt_llm.llmapi.llm_args import ( BlockReuseConfig, DFlashDecodingConfig, + DSparkDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MTPDecodingConfig, @@ -766,6 +770,85 @@ def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layo manager.shutdown() +def test_dflash_block_tables_with_fp8_target_and_bf16_draft() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + torch.cuda.init() + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=2, + head_dim=8, + buffer_roles=(Role.KEY, Role.VALUE), + ) + # Equal byte sizes coalesce the FP8 target and BF16 draft pages, placing + # the draft's K buffer at a nonzero offset within each physical slot. + draft = replace( + target, + global_layer_id=1, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=1, + ) + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=1, + num_kv_heads=target.num_kv_heads, + head_dim=target.head_dim, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + layer_specs=[target, draft], + ) + try: + assert manager._use_per_layer_page_tables + request_ids = [7, 9] + for request_id, num_blocks in zip(request_ids, (2, 3)): + cache = manager._create_kv_cache(request_id, None, None) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(num_blocks * TOKENS_PER_BLOCK, 0) + + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + pool = worker._managed_ctx_pool( + manager, 1, draft.num_kv_heads, draft.head_dim, torch.bfloat16, 2 + ) + assert pool is not None + assert worker._init_ctx_block_tables( + manager, pool, len(request_ids), 1, draft.num_kv_heads, draft.head_dim, TOKENS_PER_BLOCK + ) + assert worker._ctx_block_offset > 0 + + block_offsets = torch.empty( + (manager.num_attention_op_pools, len(request_ids), 2, manager.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + manager.copy_batch_block_offsets( + block_offsets, request_ids, beam_width=1, num_contexts=0, num_seqs=len(request_ids) + ) + metadata = SimpleNamespace(kv_cache_block_offsets=block_offsets) + assert worker._refresh_ctx_block_tables(metadata, len(request_ids)) + + expected = manager.get_batch_cache_indices(request_ids, layer_idx=draft.global_layer_id) + for row, indices in enumerate(expected): + assert len(indices) == (2, 3)[row] + assert worker._ctx_block_tables[row, : len(indices)].tolist() == indices + finally: + manager.shutdown() + + def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -3913,6 +3996,194 @@ def get_base_page_indices(self, layer_group_id, beam_id=DEFAULT_BEAM_INDEX): 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"]) +@pytest.mark.parametrize( + "spec_config_cls,shared_draft_block_width", + [(DFlashDecodingConfig, None), (DSparkDecodingConfig, 7), (DSparkDecodingConfig, 8)], + ids=["dflash", "dspark_shifted", "dspark_unshifted"], +) +def test_shared_draft_pages_cover_the_drafter_block_after_a_decode_resize( + disagg: bool, spec_config_cls: type, shared_draft_block_width: int | None +) -> None: + """The drafter appends its block after the context, so its pages must reach it. + + After a decode step accepts every draft token the context is ``prompt + K + + 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 = shared_draft_block_width or max_draft_len + 1 + tokens_per_block = 32 + prompt_len = 2 * tokens_per_block - (max_draft_len + 1) - block_width + 1 + 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=spec_config_cls(max_draft_len=max_draft_len, speculative_model="draft"), + shared_draft_block_width=shared_draft_block_width, + vocab_size=4096, + enable_stats=False, + ) + request = _ContextRequest( + 1, + list(range(prompt_len)), + prompt_len, + "conversation", + use_conversation_params=False, + is_disagg_generation_init_state=disagg, + ) + try: + assert manager._shared_draft_block_width == block_width + assert manager._generation_kv_capacity_headroom == max_draft_len + 1 + block_width + if disagg: + request.py_draft_tokens = [] + request.is_disagg_generation_transmission_complete = False + 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 + + def test_generation_token_slice_shares_token_view() -> None: token_view = np.arange(10, dtype=np.int32) request = SimpleNamespace( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py index 7abf67b67d04..7581011fe63b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_arena_sizing.py @@ -76,7 +76,7 @@ def _metadata(*, declares_max_seq_len): def _size_arena(metadata): """Run the sizing path and return the ``_max_ctx`` it settled on.""" worker = _worker() - draft_model = MagicMock() + draft_model = MagicMock(block_size=8) draft_model.config = MagicMock(max_position_embeddings=POSITION_TABLE) spec_metadata = MagicMock(max_num_requests=32) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py new file mode 100644 index 000000000000..62d29a757a0e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py @@ -0,0 +1,280 @@ +# 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 +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.speculative import dflash as dflash_module +from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.utils import get_draft_block_width + +pytestmark = pytest.mark.cpu_only + + +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) + + +@pytest.mark.parametrize("unified", [False, True]) +@pytest.mark.parametrize("block_offset", [0, 12]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_refresh_ctx_block_tables_preserves_layer_offsets_and_source( + monkeypatch: pytest.MonkeyPatch, unified: bool, block_offset: int, dtype: torch.dtype +) -> None: + worker = _worker(unified=unified) + worker._ctx_pool_idx = 1 + worker._ctx_block_offset = block_offset + worker._ctx_block_divisor = 2 + worker._ctx_block_tables = torch.full((3, 4), -7, dtype=torch.int32) + worker._ctx_block_counts = torch.full((3,), -7, dtype=torch.int64) + bad_page_index = -(1 << 31) + source = torch.full((2, 3, 2, 4), 99, dtype=dtype) + source[1, :2, 0] = torch.tensor( + [ + [block_offset, block_offset + 2, -1, bad_page_index], + [block_offset + 4, block_offset + 6, block_offset + 8, -1], + ], + dtype=dtype, + ) + original = source.clone() + other_source = torch.full_like(source, 123) + metadata = SimpleNamespace( + kv_cache_block_offsets=source if unified else other_source, + draft_kv_cache_block_offsets=other_source if unified else source, + ) + + def decode_reference( + encoded: torch.Tensor, divisor: int, table: torch.Tensor, counts: torch.Tensor + ) -> None: + rows = encoded.shape[0] + counts[:rows].copy_((encoded >= 0).sum(dim=1)) + table[:rows].copy_(encoded.clamp(min=0).div(divisor, rounding_mode="floor")) + + decode = Mock(side_effect=decode_reference) + monkeypatch.setattr(dflash_module, "decode_draft_block_table", decode) + + assert worker._refresh_ctx_block_tables(metadata, num_seqs=2) + + decode.assert_called_once() + encoded, divisor, table, counts = decode.call_args.args + assert encoded.dtype == (torch.int64 if block_offset else dtype) + assert torch.equal( + encoded, + torch.tensor( + [ + [0, 2, -1 - block_offset, bad_page_index - block_offset], + [4, 6, 8, -1 - block_offset], + ], + dtype=torch.int64, + ), + ) + assert divisor == 2 + assert table is worker._ctx_block_tables + assert counts is worker._ctx_block_counts + assert torch.equal(table, torch.tensor([[0, 1, 0, 0], [2, 3, 4, 0], [-7, -7, -7, -7]])) + assert torch.equal(counts, torch.tensor([2, 3, -7])) + assert torch.equal(source, original) + assert torch.equal(other_source, torch.full_like(other_source, 123)) + + +@pytest.mark.parametrize("unified", [False, True]) +def test_unified_cache_keeps_strict_allocation_check_in_generation_prepare(unified: bool) -> None: + worker = _worker(unified=unified) + worker._paged_latent_prepare = True + worker._draft_tokens_per_req = 8 + metadata = SimpleNamespace(runtime_draft_len=7, captured_hidden_states=object()) + + assert worker._fused_gen_prepare_applies(metadata) is not unified + + +def test_drafter_block_inside_the_allocation_is_accepted(): + num_ctx = torch.tensor([57, 10]) + allocated = torch.tensor([57, 56]) + + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def test_drafter_block_outside_the_allocation_is_an_error(): + num_ctx = torch.tensor([58, 10]) + allocated = torch.tensor([57, 56]) + + with pytest.raises(RuntimeError, match="block of 7 positions falls outside"): + DFlashWorker._check_block_within_allocation(num_ctx, allocated, block_size=7) + + +def _dspark_worker() -> DSparkWorker: + worker = object.__new__(DSparkWorker) + worker.spec_config = SimpleNamespace( + max_draft_len=7, + tokens_per_gen_step=8, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + draft_is_embedded_in_target=False, + ) + worker._ctx_buf_inited = False + worker._dflash_attention_backend = "TRTLLM" + return worker + + +@pytest.mark.parametrize(("shift_label", "expected"), [(True, 7), (False, 8)]) +def test_dspark_block_width_tracks_checkpoint_slot_convention( + shift_label: bool, expected: int +) -> None: + worker = _dspark_worker() + draft_model = SimpleNamespace(_dspark_shift_label=shift_label) + + assert worker._draft_block_width(draft_model) == expected + assert DSparkWorker.block_width_for(worker.spec_config, shift_label=shift_label) == expected + assert get_draft_block_width(worker.spec_config, shift_label=shift_label) == expected + + +def test_dspark_missing_shift_label_preserves_unshifted_runtime_width() -> None: + worker = _dspark_worker() + + assert worker._draft_block_width(SimpleNamespace()) == 8 + assert DSparkWorker.block_width_for(worker.spec_config) == 7 + assert get_draft_block_width(worker.spec_config) == 7 + + +@pytest.mark.parametrize("shift_label", [True, False]) +def test_draft_block_width_shift_label_only_affects_standalone_dspark(shift_label: bool) -> None: + config = _dspark_worker().spec_config + config.draft_is_embedded_in_target = True + assert get_draft_block_width(config, shift_label=shift_label) == 0 + + config.spec_dec_mode = SpeculativeDecodingMode.DFLASH + assert get_draft_block_width(config, shift_label=shift_label) == 8 + + +def _init_dspark_ctx_buffers( + worker: DSparkWorker, *, shift_label: bool, reserved_width: int, block_size: int = 8 +) -> None: + manager = SimpleNamespace(max_seq_len=64, _shared_draft_block_width=reserved_width) + draft_model = SimpleNamespace( + _dspark_shift_label=shift_label, + block_size=block_size, + dflash_attention_backend="TRTLLM", + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=1), + SimpleNamespace(kv_cache_manager=manager, kv_cache_block_offsets=None), + manager, + ) + + +@pytest.mark.parametrize( + ("shift_label", "reserved_width", "block_size", "error"), + [ + (False, 7, 8, "Drafter block width 8 exceeds the 7 positions"), + (True, 6, 8, "Drafter block width 7 exceeds the 6 positions"), + (False, 8, 7, "checkpoint was trained with block_size=7"), + ], +) +def test_dspark_invalid_width_is_rejected_before_cuda_allocation( + monkeypatch: pytest.MonkeyPatch, + shift_label: bool, + reserved_width: int, + block_size: int, + error: str, +) -> None: + allocate = Mock(side_effect=RuntimeError("buffer allocation reached")) + monkeypatch.setattr(torch, "zeros", allocate) + + with pytest.raises(ValueError, match=error): + _init_dspark_ctx_buffers( + _dspark_worker(), + shift_label=shift_label, + reserved_width=reserved_width, + block_size=block_size, + ) + + allocate.assert_not_called() + + +@pytest.mark.parametrize(("shift_label", "reserved_width"), [(True, 7), (False, 8)]) +def test_dspark_matching_width_reaches_buffer_allocation( + monkeypatch: pytest.MonkeyPatch, shift_label: bool, reserved_width: int +) -> None: + allocate = Mock(side_effect=RuntimeError("buffer allocation reached")) + monkeypatch.setattr(torch, "zeros", allocate) + worker = _dspark_worker() + + with pytest.raises(RuntimeError, match="buffer allocation reached"): + _init_dspark_ctx_buffers(worker, shift_label=shift_label, reserved_width=reserved_width) + + allocate.assert_called_once() + assert worker._ctx_unified_cache + assert worker._compute_block_size == reserved_width diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 15670e3de0c6..fb783480de80 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -1445,3 +1445,29 @@ def test_attention_replicated_validation_rejects_missing_index_key_view(): peer_pt = make_page_table(global_layer_ids=[0]) with pytest.raises(ValueError, match="AttentionPolicy.*differ on overlapping layers"): AttentionPolicy.validate_peer_compatible(self_pt, peer_pt) + + +@pytest.mark.parametrize(("local_draft", "peer_draft"), [(0, 0), (1, 1), (1, 0), (0, 1)]) +def test_peer_registrar_requires_matching_unified_draft_layers( + local_draft: int, peer_draft: int +) -> None: + self_rankinfo = make_rankinfo(instance_name="local") + self_rankinfo.num_draft_layers = local_draft + reg = _make_peer_registrar(self_rankinfo) + peer_ri = make_rankinfo( + instance_name="peer", + instance_rank=1, + layer_num_per_pp=[2], + page_table=make_page_table(), + ) + peer_ri.num_draft_layers = peer_draft + if local_draft != peer_draft: + with pytest.raises( + ValueError, + match=f"draft KV layer count differs.*local={local_draft}, peer={peer_draft}", + ): + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + return + + reg.register(peer_ri.instance_name, peer_ri.instance_rank, peer_ri) + assert reg.get_peer_rank_info("peer", 1) == peer_ri diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index a1d25e01c3f3..a9ef379c93e7 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -15,6 +15,7 @@ from types import SimpleNamespace +import msgpack import numpy as np import pytest @@ -203,3 +204,64 @@ def test_rank_info_represents_cache_element_bytes( restored = RankInfo.from_bytes(rank_info.to_bytes()) assert restored.attention.element_bytes == expected_element_bytes assert isinstance(restored.attention.element_bytes, expected_type) + + +@pytest.mark.parametrize( + ("layer_specs", "draft_layer_ids", "expected"), + [ + (None, (), 0), + # Folded MTP/Eagle draft layers have no layer specs and travel as + # ordinary layers, so they are not advertised. + (None, (36,), 0), + ((object(),), (36, 37), 2), + ], +) +def test_from_kv_cache_manager_reports_unified_draft_layers( + monkeypatch, layer_specs, draft_layer_ids, expected +) -> None: + monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _: None) + manager = SimpleNamespace( + mapping=SimpleNamespace( + rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + dp_size=1, + cp_size=1, + cp_rank=0, + enable_attention_dp=False, + ), + num_kv_heads_per_layer=[8, 8], + pp_layers=[0, 1], + tokens_per_block=32, + head_dim=128, + dtype=bindings.DataType.HALF, + kv_factor=2, + layer_specs=layer_specs, + draft_layer_ids=draft_layer_ids, + ) + + info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) + + assert info.num_draft_layers == expected + assert RankInfo.from_bytes(info.to_bytes()).num_draft_layers == expected + + +def test_rank_info_from_bytes_defaults_missing_draft_layer_count() -> None: + ri = RankInfo( + instance_name="gen_0", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[32], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + ) + payload = msgpack.unpackb(ri.to_bytes(), strict_map_key=False) + del payload["num_draft_layers"] + + assert RankInfo.from_bytes(msgpack.packb(payload)).num_draft_layers == 0