diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index b5f585450493..c80823654312 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -152,7 +152,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/ * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model. * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config. * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config. -* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0`. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. +* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"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 the SM100 family only. Because the drafter's context lives in the KV cache, a prefix-cache hit reuses it and a disaggregated context server transfers it to the generation server. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache outside the KV cache manager and runs anywhere, but a prefix-cache hit then drafts without the reused prefix and a disaggregated generation server drafts without the prompt, both at a lower acceptance length; it is deprecated wherever TRTLLM can run the drafter. `"FA4"` uses the flash-attn CuTe DSL kernels on a private paged cache and supports SM90 only. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run the drafter, and to `"VANILLA"` where it cannot (another GPU architecture, no FlashInfer, or a head geometry the TRTLLM kernels do not take), with KV cache compression (the drafter then cannot share the target's KV cache manager, see [Draft KV cache](#dflash-draft-kv-cache)), and for an unwindowed drafter on an aggregated server, on which TRTLLM's context kernel is slower. A disaggregated server that shares the draft KV (unless `TRTLLM_DISAGG_SHARED_DRAFT_KV=0`) resolves `"AUTO"` to `"TRTLLM"` for windowed and unwindowed drafters alike. ```python from tensorrt_llm.llmapi import DFlashDecodingConfig @@ -194,8 +194,10 @@ match, and both servers must run the same target and drafter. The drafter keeps a separate draft KV cache manager instead when any of the following applies: context parallelism larger than 1, a per-layer KV-head count on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba -models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout -that cannot be transferred with the target's. When the layers are shared, the +models), Gemma 4 hybrid attention, KV cache compression, or, in disaggregated +serving, a draft layout that cannot be transferred with the target's. With +cold-page KV cache compression the target's cold pages are compressed and the +separate draft manager's pages are kept at their own dtype. When the layers are shared, the worker log reports `target and draft KV layers share one KV cache manager` at startup. diff --git a/examples/kv_cache_compression/nvfp4_cold_page.md b/examples/kv_cache_compression/nvfp4_cold_page.md index b2f8d234c112..93830839a685 100644 --- a/examples/kv_cache_compression/nvfp4_cold_page.md +++ b/examples/kv_cache_compression/nvfp4_cold_page.md @@ -68,7 +68,11 @@ The current implementation requires the PyTorch backend, native C++ KVCacheManagerV2, and an SM100, SM103, or SM107 GPU. Hot Attention KV can use FP16, BF16, or FP8. Host and Disk cold tiers share the same compact representation, and KV-cache block reuse remains supported because Page and token identity are -unchanged. One-model MTP-EAGLE and EAGLE3 are supported; HELIX context +unchanged. One-model MTP-EAGLE and EAGLE3 are supported. DFlash is supported +on an aggregated server: only the target's cold Pages are compressed. Its +default `attention_backend` (`AUTO`) is `VANILLA` here, whose drafter context +stays outside the KV cache manager; an explicit `TRTLLM` drafter keeps a +separate draft KV cache, which moves its Pages losslessly. HELIX context parallelism is not currently supported. Set `kv_cache_config.use_kv_cache_manager_v2: true` explicitly. A nonzero Host diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 68e370b61ed7..1bf6ab010b6c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2077,12 +2077,24 @@ def _check_shared_draft_kv_request(self) -> None: or self._unified_drafter_attention_backend() is None): return if os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "1": + separate_trtllm_pool = ( + self._unified_drafter_attention_backend() == "TRTLLM" + and self._should_create_separate_draft_kv_cache()) + windowed = separate_trtllm_pool and self._get_draft_max_attention_window( + self._unification_max_seq_len, + self._kv_cache_config) is not None + if (separate_trtllm_pool and not windowed + and os.environ.get("TRTLLM_DISAGG_SHARED_DRAFT_KV") != "0"): + reason = (self._dflash_unify_blocker() + or "TRTLLM_DISAGG_SHARED_DRAFT_KV is not 1") + logger.warning( + f"{self._unified_drafter_name()} with attention_backend " + "TRTLLM on a disaggregated worker keeps a separate draft " + f"KV cache because {reason}. It is not transferred with " + "the target's, so the generation worker drafts without " + "the prompt, at a lower acceptance length.") # The drafter writes from column 0; a separate pool drops blocks at real positions. - if (self._unified_drafter_attention_backend() == "TRTLLM" - and self._should_create_separate_draft_kv_cache() - and self._get_draft_max_attention_window( - self._unification_max_seq_len, - self._kv_cache_config) is not None): + if windowed: blocker = self._dflash_unify_blocker() fix = ( "use attention_backend VANILLA" if blocker else @@ -3284,7 +3296,8 @@ def build_managers(self, original_max_seq_len, estimating_kv_cache, kv_cache_config_override=draft_build_kv_cache_config, - cold_page_codec_provider=cold_page_codec_provider) + cold_page_codec_provider=draft_kv_cold_page_codec_provider( + cold_page_codec_provider, self._speculative_config)) # Encoder-decoder cross-attention pool cross_kv_cache_manager = None @@ -4106,6 +4119,12 @@ def validate_kv_cache_compression_compatibility( "support speculative decoding with its current configuration" f"{guidance}") mode = spec_config.spec_dec_mode + if config.algorithm == "quantization_for_cold_page" and mode.is_dflash(): + # The codec only ever sees target pages: VANILLA and FA4 keep the + # drafter's context K/V in worker-owned buffers, and a TRTLLM drafter's + # separate draft KVCM is built without a codec + # (draft_kv_cold_page_codec_provider). + return if config.algorithm == "quantization_for_cold_page": supported = mode.is_mtp_eagle_one_model() or mode.is_eagle3_one_model() guidance = "one-model MTP-EAGLE or EAGLE3" @@ -4118,6 +4137,29 @@ def validate_kv_cache_compression_compatibility( f"mode {mode.name}; use {guidance}") +def draft_kv_cold_page_codec_provider( + cold_page_codec_provider: Optional[object], + spec_config: Optional[SpeculativeConfig]) -> Optional[object]: + """The cold-page codec provider for a one-model drafter's separate KV cache manager. + + A DFlash drafter's manager gets none, so its pages move to and from the + Host and Disk tiers losslessly, like the private context arena of the + other DFlash backends, which never leaves the GPU. The codec holds only + the target's config and would quantize the drafter's K/V as whole + vectors, which is not validated for acceptance length. + + Args: + cold_page_codec_provider: The provider the target manager was given. + spec_config: The speculative decoding config. + + Returns: + The provider for the draft manager, or None for the lossless default. + """ + if spec_config is not None and spec_config.spec_dec_mode.is_dflash(): + return None + return cold_page_codec_provider + + def create_kv_cache_compression_manager( config: Optional[KvCacheCompressionConfig], *, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 095424b7f37a..0d06ed04e739 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -24,7 +24,8 @@ ContextChunkingPolicy, ExecutorMemoryType, GuidedDecodingConfig, KvCacheConfig, - SpeculativeConfig, TorchLlmArgs) + SpeculativeConfig, TorchLlmArgs, + dflash_auto_shares_draft_kv) from tensorrt_llm.llmapi.tokenizer import (TokenizerBase, _llguidance_tokenizer_info, _xgrammar_tokenizer_info) @@ -224,11 +225,20 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType): )) -def _pin_dflash_auto_attention_backend(spec_config) -> None: - """Resolve a DFlash AUTO a GPU-less frontend left open, before the model and KV caches read it.""" +def _pin_dflash_auto_attention_backend(spec_config, llm_args=None) -> None: + """Resolve a DFlash AUTO a GPU-less frontend left open, before the model and KV caches read it. + + Resolves it as the frontend would have (``dflash_auto_shares_draft_kv`` + and the KV cache compression setting). + """ if (spec_config is not None and spec_config.spec_dec_mode.is_dflash() and spec_config.attention_backend == "AUTO"): - spec_config.resolve_from_checkpoint() + spec_config.resolve_from_checkpoint( + disagg_shared_draft_kv=llm_args is not None + and dflash_auto_shares_draft_kv(llm_args, spec_config), + kv_cache_compression=getattr(llm_args, + "kv_cache_compression_config", + None) is not None) def _flashinfer_one_engine_spec_supported(attn_backend: str, @@ -620,7 +630,7 @@ def _create_py_executor( if spec_config is not None and spec_config.decoding_type == "AUTO": from tensorrt_llm._torch.speculative import suggest_spec_config spec_config = suggest_spec_config(max_batch_size) - _pin_dflash_auto_attention_backend(spec_config) + _pin_dflash_auto_attention_backend(spec_config, llm_args) if not llm_args.disable_overlap_scheduler and spec_config is not None: if not spec_config.spec_dec_mode.support_overlap_scheduler(): diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index ee0e8fb35024..1ec6c7c6590d 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -81,10 +81,71 @@ def drafter_config_window(draft_config: Optional[dict]) -> Optional[int]: return None -def resolve_dflash_auto_attention_backend(draft_config: Optional[dict]) -> str: - """What AUTO means for a DFlash drafter's config.json; DSpark drafters keep their own AUTO.""" - # TRTLLM runs an unwindowed non-causal layer on its context kernel, 2-9x slower than VANILLA. - if drafter_config_window(draft_config) is None: +def _dflash_trtllm_gen_shape_unsupported_reason(draft_config: Optional[dict]) -> Optional[str]: + """Why the drafter's head geometry rules out the TRTLLM kernels, or None. + + Mirrors the geometry half of ``validate_dflash_trtllm_gen_runtime`` from + the drafter's config.json, so AUTO does not pick a backend the drafter + then refuses at its first forward. A missing key counts as supported. + """ + if not draft_config: + return None + from ..attention.backends.fmha.flashinfer_trtllm_gen import FlashInferTrtllmGenFmha + + num_heads = draft_config.get("num_attention_heads") + num_kv_heads = draft_config.get("num_key_value_heads") or num_heads + if num_heads and num_kv_heads: + if num_heads % num_kv_heads != 0: + return f"num_attention_heads {num_heads} is not a multiple of {num_kv_heads}" + if num_heads // num_kv_heads > FlashInferTrtllmGenFmha.MAX_HEADS_RATIO_GENERATION: + return f"a Q/KV head ratio of {num_heads // num_kv_heads}" + head_dim = draft_config.get("head_dim") + if head_dim is None and num_heads and draft_config.get("hidden_size"): + head_dim = draft_config["hidden_size"] // num_heads + # The non-causal draft block runs on the context kernel, which every + # DFlash drafter but a causal one uses. + if head_dim in FlashInferTrtllmGenFmha.UNSUPPORTED_HEAD_SIZES_CONTEXT: + return f"head_dim {head_dim} on the context kernel" + return None + + +def resolve_dflash_auto_attention_backend( + draft_config: Optional[dict], + *, + disagg_shared_draft_kv: bool = False, + kv_cache_compression: bool = False, +) -> str: + """What AUTO means for a DFlash drafter's config.json; DSpark drafters keep their own AUTO. + + TRTLLM keeps the drafter's context K/V in the KV cache manager, so a + prefix-cache hit reuses it and a disaggregated context worker transfers + it; VANILLA keeps it in a private arena that neither can reach. AUTO + therefore resolves to TRTLLM wherever its kernels run the drafter, with + two exceptions: an unwindowed drafter on an aggregated worker, where TRTLLM + runs the non-causal layer on its context kernel, 2-9x slower than VANILLA, + and KV cache compression, under which the drafter cannot share the + target's KV cache manager, so a TRTLLM drafter would keep a separate draft + KV cache that prefix reuse and the transfer do not reach either. + + Args: + draft_config: The drafter checkpoint's config.json contents. + disagg_shared_draft_kv: The worker is disaggregated and transfers the + draft KV with the target's (``disagg_shares_draft_kv``). A + VANILLA generation worker would draft without the prompt, so any + drafter takes TRTLLM there. + kv_cache_compression: ``kv_cache_compression_config`` is set. + + Returns: + ``"TRTLLM"`` or ``"VANILLA"``. + """ + if kv_cache_compression: + logger.info_once( + "DFlash attention_backend AUTO keeps VANILLA: with KV cache compression " + "the drafter cannot share the target's KV cache manager.", + key="dflash_auto_kv_cache_compression", + ) + return "VANILLA" + if drafter_config_window(draft_config) is None and not disagg_shared_draft_kv: return "VANILLA" if (reason := dflash_trtllm_gen_unavailability_reason()) is not None: logger.info_once( @@ -92,6 +153,12 @@ def resolve_dflash_auto_attention_backend(draft_config: Optional[dict]) -> str: key="dflash_auto_trtllm_unavailable", ) return "VANILLA" + if (reason := _dflash_trtllm_gen_shape_unsupported_reason(draft_config)) is not None: + logger.info_once( + f"DFlash attention_backend AUTO keeps VANILLA: TRTLLM does not support {reason}.", + key="dflash_auto_trtllm_shape", + ) + return "VANILLA" return "TRTLLM" diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 65e74256d65d..92bea3adf0a3 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -2458,8 +2458,8 @@ def _shared_draft_kv_blocker(server_configs) -> Optional[str]: "generation_servers") for cfg in server_configs: spec = cfg.other_args["speculative_config"] - # AUTO, the default, resolves to the sharing TRTLLM backend for windowed - # drafters. + # AUTO, the default, resolves to the sharing TRTLLM backend wherever + # TRTLLM can run the drafter. if (spec.get("decoding_type") not in ("DFlash", "DSpark") or spec.get("attention_backend") not in (None, "AUTO", "TRTLLM")): return ("speculative_config to be DFlash or DSpark with " diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5a6c913bdfe1..811d6fd03d13 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3309,17 +3309,23 @@ class DFlashDecodingConfig(DecodingBaseConfig): default="AUTO", description= "Attention backend for DFlash pooled-context cross-attention, independent " - "of the drafter's standard attention modules. VANILLA uses FlashAttention " - "with a contiguous context K/V cache and runs anywhere. TRTLLM uses " - "TRTLLM-Gen FMHA (via FlashInfer) over a private paged context K/V cache " - "and supports SM100/SM103 only. FA4 uses the flash-attn CuTe DSL kernels " - "on the same paged cache and supports SM90 only. AUTO (the default) " - "resolves to TRTLLM where it can run and every drafter layer is " - "windowed, and to VANILLA otherwise, including on a disaggregated worker " - "with TRTLLM_DISAGG_SHARED_DRAFT_KV=0. In " - "disaggregated serving a TRTLLM drafter keeps its KV with the target's, " - "so the context server's drafter context reaches the generation server; " - "both servers need the same speculative_config.") + "of the drafter's standard attention modules. TRTLLM uses TRTLLM-Gen FMHA " + "(via FlashInfer) over a paged context K/V cache held by the KV cache " + "manager, so a prefix-cache hit reuses the drafter's context and a " + "disaggregated context server transfers it; it supports the SM100 family " + "only. VANILLA uses FlashAttention with a contiguous context K/V cache " + "outside the KV cache manager and runs anywhere, but a prefix-cache hit " + "drafts without the reused prefix and a disaggregated generation server " + "drafts without the prompt, both of which lower acceptance length. FA4 " + "uses the flash-attn CuTe DSL kernels on a private paged cache and " + "supports SM90 only. AUTO (the default) resolves to TRTLLM where it can " + "run the drafter, and to VANILLA where it cannot, with KV cache " + "compression (the drafter then cannot share the target's KV cache " + "manager), and for an unwindowed drafter on an aggregated server, on " + "which TRTLLM's context kernel is slower. A disaggregated server that shares the draft KV (the default; " + "TRTLLM_DISAGG_SHARED_DRAFT_KV=0 opts out) resolves AUTO to TRTLLM for " + "any drafter; the context and generation servers need the same " + "speculative_config.") skip_ctx_buffer_budget_check: bool = Field( default=False, @@ -3334,13 +3340,24 @@ def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len return self - def resolve_from_checkpoint(self) -> None: + def resolve_from_checkpoint(self, + *, + disagg_shared_draft_kv: bool = False, + kv_cache_compression: bool = False) -> None: """Fill unset fields from the draft checkpoint's ``config.json``. Must run once ``speculative_model`` names a local directory: ``target_layer_ids`` has no usable default, and without it no target hidden states are captured. Called from validation for a local path, and again from model loading once a HF repo id has been downloaded. + + Args: + disagg_shared_draft_kv: The server is disaggregated and transfers + the draft KV with the target's (see + ``dflash_auto_shares_draft_kv``); AUTO then resolves to TRTLLM + for any drafter TRTLLM can run. + kv_cache_compression: The server sets + ``kv_cache_compression_config``; AUTO then resolves to VANILLA. """ if self.speculative_model is None: return @@ -3357,7 +3374,9 @@ def resolve_from_checkpoint(self) -> None: from tensorrt_llm._torch.speculative.dflash import \ resolve_dflash_auto_attention_backend self.attention_backend = resolve_dflash_auto_attention_backend( - draft_cfg) + draft_cfg, + disagg_shared_draft_kv=disagg_shared_draft_kv, + kv_cache_compression=kv_cache_compression) logger.info( f"DFlash attention_backend AUTO resolved to {self.attention_backend}." ) @@ -6023,6 +6042,39 @@ def disagg_shares_draft_kv( and speculative_config.attention_backend in ("AUTO", "TRTLLM")) +def dflash_auto_shares_draft_kv( + llm_args: "BaseLlmArgs", + speculative_config: Optional[DecodingBaseConfig] = None) -> bool: + """Whether a DFlash AUTO backend resolves for a server that transfers its draft KV. + + True on a disaggregated server (a cache transceiver backend is set) whose + drafter shares its KV with the target's per ``disagg_shares_draft_kv``, + unless a reason ``KvCacheCreator._dflash_unify_blocker`` refuses one + shared KV cache manager is already known from the arguments: a target + attention backend other than TRTLLM, a KV cache manager not explicitly + V2, pipeline or context parallelism, KV cache compression, or a + ``draft_len_schedule``. Every place that pins AUTO passes this to + ``resolve_from_checkpoint``, so a frontend and its ranks resolve alike. + + Args: + llm_args: The server's arguments. + speculative_config: Overrides ``llm_args.speculative_config``. + """ + spec_config = (speculative_config if speculative_config is not None else + llm_args.speculative_config) + transceiver = getattr(llm_args, "cache_transceiver_config", None) + if transceiver is None or transceiver.backend is None: + return False + if not disagg_shares_draft_kv(spec_config): + return False + return (getattr(llm_args, "attn_backend", None) == "TRTLLM" + and llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + and llm_args.pipeline_parallel_size == 1 + and llm_args.context_parallel_size == 1 + and getattr(llm_args, "kv_cache_compression_config", None) is None + and spec_config.draft_len_schedule is None) + + def disagg_dflash_error( speculative_config: Optional[DecodingBaseConfig], cache_transceiver_config: Optional[CacheTransceiverConfig] @@ -6826,6 +6878,46 @@ def _kv_cache_estimation_runs(self) -> bool: return False return True + def _warn_if_dflash_vanilla_overrides_trtllm(self, disagg: bool) -> None: + """Warn when an explicit DFlash VANILLA replaces what AUTO resolves to TRTLLM. + + VANILLA keeps the drafter's context K/V outside the KV cache manager, + so a prefix-cache hit drafts without the reused prefix and a + disaggregated generation server drafts without the prompt. Silent + unless AUTO would run this drafter on TRTLLM here, which needs a GPU + and the drafter's config.json at validation time. + """ + spec_cfg = self.speculative_config + if spec_cfg.speculative_model is None or not torch.cuda.is_available(): + return + draft_config_path = os.path.join(str(spec_cfg.speculative_model), + "config.json") + if not os.path.exists(draft_config_path): + return + with open(draft_config_path) as f: + draft_config = json.load(f) + from tensorrt_llm._torch.speculative.dflash import \ + resolve_dflash_auto_attention_backend + auto_shares = dflash_auto_shares_draft_kv( + self, spec_cfg.model_copy(update={"attention_backend": "AUTO"})) + if disagg and not auto_shares: + # A TRTLLM drafter could not share its KV here either. + return + if resolve_dflash_auto_attention_backend( + draft_config, + disagg_shared_draft_kv=auto_shares, + kv_cache_compression=self.kv_cache_compression_config + is not None) != "TRTLLM": + return + lost = ("the generation server drafts without the prompt" if disagg else + "a prefix-cache hit drafts without the reused prefix") + logger.warning( + "DFlash attention_backend VANILLA is set where the default (AUTO) " + "runs this drafter on TRTLLM. VANILLA keeps the drafter's context " + f"K/V outside the KV cache manager, so {lost}, which lowers " + "acceptance length. VANILLA is deprecated wherever TRTLLM can run " + "the drafter; leave attention_backend unset to use TRTLLM.") + def _validate_dflash_ctx_budget(self, memory_budget_bytes: Optional[int] = None ) -> None: @@ -7095,6 +7187,8 @@ def validate_speculative_config(self): "(a disaggregated worker with " "TRTLLM_DISAGG_SHARED_DRAFT_KV=0).") assert self.speculative_config.max_draft_len > 0, "DFlash max_draft_len must be > 0" + pinned_vanilla = ( + self.speculative_config.attention_backend == "VANILLA") # From the disaggregated-serving-benchmarking overlay (TRT-LLM aa7db814): # the disagg_dflash_error refusal is dropped so gen-only SOL # ladders can run DFlash on a disaggregated gen. Acceptance is @@ -7106,7 +7200,12 @@ def validate_speculative_config(self): # then run without the drafter's config.json (the budget check # covers only the token budget), and CachedModelLoader repeats # both after the drafter is downloaded. - self.speculative_config.resolve_from_checkpoint() + self.speculative_config.resolve_from_checkpoint( + disagg_shared_draft_kv=dflash_auto_shares_draft_kv(self), + kv_cache_compression=self.kv_cache_compression_config + is not None) + if pinned_vanilla: + self._warn_if_dflash_vanilla_overrides_trtllm(disagg) # A draft KV shared with the target is transferred with it; # checked once AUTO is resolved from the checkpoint. if disagg and not (self.speculative_config.attention_backend in diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index b402f433e360..2c3fbc6db39d 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -28,7 +28,8 @@ LlmArgs, MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, UserProvidedDecodingConfig, _ModelWrapper, - _ParallelConfig, update_llm_args_with_extra_dict, + _ParallelConfig, dflash_auto_shares_draft_kv, + update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable from .mpi_session import MpiSession @@ -412,11 +413,19 @@ def __call__(self) -> Optional[Path]: # repo id had no readable config.json at validation time; a local # path was already resolved back then. if was_hub_model: - resolve_from_checkpoint = getattr( - self.llm_args.speculative_config, 'resolve_from_checkpoint', - None) - if resolve_from_checkpoint is not None: - resolve_from_checkpoint() + spec_config = self.llm_args.speculative_config + if isinstance(spec_config, DFlashDecodingConfig): + spec_config.resolve_from_checkpoint( + disagg_shared_draft_kv=dflash_auto_shares_draft_kv( + self.llm_args), + kv_cache_compression=getattr( + self.llm_args, "kv_cache_compression_config", + None) is not None) + else: + resolve_from_checkpoint = getattr( + spec_config, 'resolve_from_checkpoint', None) + if resolve_from_checkpoint is not None: + resolve_from_checkpoint() # The DFlash buffer-fit check reads the drafter geometry from # the same config.json, so for a hub repo id the validation # pass only covered the token budget. Re-run it now that the diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 0a9f5c33c94b..fa8abd0ac607 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy import gc import math import os @@ -302,6 +303,92 @@ def assert_kv_cache_reuse_for_llm( assert warm_output.token_ids == cold_output.token_ids +def gsm8k_fewshot_prompts(num_prompts: int, num_shots: int = 5) -> List[str]: + """Completion prompts of ``num_shots`` GSM8K train examples and one test question. + + Several hundred tokens each, so a repeated prompt spans many KV-cache + blocks. Deterministic: the first ``num_shots`` train examples and the + first ``num_prompts`` test questions. + """ + from datasets import load_dataset + + data_dir = f"{llm_models_root()}/datasets/openai/gsm8k/main" + dataset = load_dataset("parquet", + data_files={ + "train": + f"{data_dir}/train-00000-of-00001.parquet", + "test": f"{data_dir}/test-00000-of-00001.parquet" + }) + shots = "".join(f"Question: {row['question']}\nAnswer: {row['answer']}\n\n" + for row in dataset["train"].select(range(num_shots))) + return [ + f"{shots}Question: {row['question']}\nAnswer:" + for row in dataset["test"].select(range(num_prompts)) + ] + + +def acceptance_length_cold_and_warm( + llm: PyTorchLLM, prompts: List[str], + sampling_params: SamplingParams) -> tuple[float, float]: + """Acceptance length of ``prompts`` run cold, then again over their cached prefix. + + Each prompt gets its own ``cache_salt``, so the cold pass reuses nothing + (not even a prefix the prompts share) and the warm pass reuses each + prompt's own blocks; the warm pass is checked to have hit the cache. The + target's output is the same on both passes, so a drop in the warm + acceptance length can only come from a drafter that lost the reused + prefix's context. Requires ``enable_iter_perf_stats=True`` and block reuse. + + Returns: + ``(cold, warm)`` mean acceptance lengths. + """ + sampling_params = copy.copy(sampling_params) + sampling_params.return_perf_metrics = True + salts = [f"acceptance-under-reuse-{i}" for i in range(len(prompts))] + # Drop whatever iterations an earlier workload left in the stats queue. + llm.get_stats(timeout=2) + cold = llm.generate(prompts, + sampling_params=sampling_params, + cache_salt=salts, + use_tqdm=False) + cold_al = compute_acceptance_length(llm) + warm = llm.generate(prompts, + sampling_params=sampling_params, + cache_salt=salts, + use_tqdm=False) + warm_al = compute_acceptance_length(llm) + + def reused_blocks(outputs): + return sum(output.outputs[0].request_perf_metrics.kv_cache_metrics. + num_reused_blocks for output in outputs) + + assert reused_blocks(cold) == 0, "the cold pass reused KV-cache blocks" + assert reused_blocks(warm) > 0, "the warm pass did not reuse KV cache" + print(f"[AL] cold {cold_al:.4f}, warm {warm_al:.4f} " + f"({100 * (warm_al / cold_al - 1):+.2f}%)") + return cold_al, warm_al + + +def assert_acceptance_length_survives_prefix_reuse( + llm: PyTorchLLM, + prompts: List[str], + sampling_params: SamplingParams, + *, + min_ratio: float = 0.95) -> tuple[float, float]: + """Assert a prefix-cache hit keeps the drafter's acceptance length. + + See ``acceptance_length_cold_and_warm``; the warm acceptance length must + be at least ``min_ratio`` of the cold one. + """ + cold_al, warm_al = acceptance_length_cold_and_warm(llm, prompts, + sampling_params) + assert warm_al >= min_ratio * cold_al, ( + f"[AL] Prefix-cache reuse cost acceptance length: warm {warm_al:.3f} " + f"< {min_ratio} x cold {cold_al:.3f}. The drafter drafts a reused " + "prefix without its context.") + return cold_al, warm_al + + class AccuracyTask: REFERENCE_DIR = f"{os.path.dirname(__file__)}/references" diff --git a/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py b/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py index 93c5a8a34ab9..78a9dff1f611 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py @@ -21,6 +21,7 @@ target's fast paths end to end. """ +import json import os import pytest @@ -28,10 +29,24 @@ from tensorrt_llm import LLM from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV from tensorrt_llm._utils import get_sm_version -from tensorrt_llm.llmapi import CudaGraphConfig, DFlashDecodingConfig, KvCacheConfig, MoeConfig +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + DFlashDecodingConfig, + KvCacheConfig, + MoeConfig, + SamplingParams, +) from ...conftest import llm_models_root -from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness, assert_acceptance_length_for_llm +from ..accuracy_core import ( + GSM8K, + LlmapiAccuracyTestHarness, + assert_acceptance_length, + assert_acceptance_length_for_llm, + assert_acceptance_length_survives_prefix_reuse, + compute_acceptance_length, + gsm8k_fewshot_prompts, +) # The TEP4 compositions are written for sm103 (GB300) and sm107 (Vera Rubin). skip_not_served_sm = pytest.mark.skipif( @@ -42,13 +57,16 @@ # The public opt-ins the compositions are written against: PDL launches, the # MNNVL one-shot attention all-reduce, the fused MoE finalize all-reduce, the # sharded projections, the low-M GEMM dispatcher and the windowed drafter context -# arena. The target refuses to build without them. The KV-cache estimation pass is skipped because its +# arena. The target refuses to build without them. On sm107 the sharded +# projections and the fused finalize also need the SM107 fused-kernel opt-in, +# which sm103 ignores. The KV-cache estimation pass is skipped because its # temporary cache does not fit next to this checkpoint's weights. _SERVED_ENV = { "TRTLLM_ENABLE_PDL": "1", "MUSESPARK_OPEN_ALLREDUCE": "MNNVL", "MUSESPARK_OPEN_FUSED_FINALIZE_AR_MNNVL": "1", "MUSESPARK_OPEN_SHARDED_PROJECTIONS": "1", + "MUSESPARK_OPEN_SM107_FUSED_KERNELS": "1", "TRTLLM_LOW_M_GEMM_BACKEND": "auto", "TRTLLM_DFLASH_WINDOWED_CTX_ARENA": "1", "TRTLLM_SKIP_KV_CACHE_ESTIMATION": "1", @@ -134,3 +152,90 @@ def test_dflash2_gsm8k_and_acceptance(self, monkeypatch): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) assert_acceptance_length_for_llm(self.ACCEPTANCE_KEY, llm) + + # The SemiAnalysis golden protocol for this drafter (InferenceX + # golden_al_distribution/musespark_dflash.yaml): the first user turn of the + # 80 SPEED-Bench qualitative coding prompts, the served chat template, + # T=1.0, top_p=0.95, up to 4096 tokens, EOS honored, rejection sampling. + SPEED_BENCH_CODING = os.environ.get("MUSESPARK_SPEED_BENCH_CODING") or ( + f"{llm_models_root()}/datasets/nvidia/SPEED-Bench/qualitative_coding_turn0.jsonl" + ) + GOLDEN_ACCEPTANCE_KEY = "ModelingV2MuseSparkOpenNvfp4Tep4::dflash2_speed_bench_coding_golden" + + def _speed_bench_coding_prompts(self, llm) -> list: + if not os.path.exists(self.SPEED_BENCH_CODING): + pytest.skip( + f"SPEED-Bench coding prompts not staged at {self.SPEED_BENCH_CODING} " + "(set MUSESPARK_SPEED_BENCH_CODING)" + ) + with open(self.SPEED_BENCH_CODING, encoding="utf-8") as f: + texts = [json.loads(line)["text"] for line in f if line.strip()] + assert len(texts) == 80, f"expected the 80 coding prompts, found {len(texts)}" + return [ + llm.tokenizer.apply_chat_template( + [{"role": "user", "content": text}], tokenize=False, add_generation_prompt=True + ) + for text in texts + ] + + @skip_not_served_sm + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(180000) + def test_dflash2_speed_bench_coding_golden_acceptance(self, monkeypatch): + """Natural acceptance length on the golden protocol, against the golden. + + Runs the default drafter backend (AUTO, TRTLLM here) with rejection + sampling, two passes with different seeds pooled, and checks the + reference's floor: the SemiAnalysis K=7 golden 4.05 less 2%. + """ + with LLM( + self.MODEL_PATH, + speculative_config=DFlashDecodingConfig( + max_draft_len=7, + speculative_model=self.DFLASH2_MODEL_PATH, + use_rejection_sampling=True, + ), + **self.ENGINE, + **_require_target(monkeypatch, **_SERVED_ENV), + ) as llm: + prompts = self._speed_bench_coding_prompts(llm) + llm.get_stats(timeout=2) + for seed in (0, 1): + # The chat template already emits BOS; like the served chat + # endpoint, the tokenizer must not add a second one. + llm.generate( + prompts, + SamplingParams( + max_tokens=4096, + temperature=1.0, + top_p=0.95, + seed=seed, + add_special_tokens=False, + ), + cache_salt=[f"golden-{seed}-{i}" for i in range(len(prompts))], + use_tqdm=False, + ) + acceptance_length = compute_acceptance_length(llm) + print(f"[AL] {self.GOLDEN_ACCEPTANCE_KEY} acceptance_length = {acceptance_length:.4f}") + assert_acceptance_length(self.GOLDEN_ACCEPTANCE_KEY, acceptance_length) + + @skip_not_served_sm + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(180000) + def test_dflash2_acceptance_under_prefix_reuse(self, monkeypatch): + """A prefix-cache hit keeps the drafter's context on the default backend. + + See ``assert_acceptance_length_survives_prefix_reuse``: GSM8K few-shot + prompts run cold and then over their cached prefix, greedy. + """ + with LLM( + self.MODEL_PATH, + speculative_config=DFlashDecodingConfig( + max_draft_len=7, speculative_model=self.DFLASH2_MODEL_PATH + ), + **self.ENGINE, + **_require_target(monkeypatch, **_SERVED_ENV), + ) as llm: + assert_acceptance_length_survives_prefix_reuse( + llm, gsm8k_fewshot_prompts(64), SamplingParams(max_tokens=256, temperature=0.0) + ) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index e41e0346a7a6..9527964b7503 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -166,3 +166,11 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: ModelingV2MuseSparkOpenNvfp4Tep4::dflash2: ref_al: 4.892 min_al: 4.647 +# Read by TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_speed_bench_coding_golden_acceptance. +# Not a measurement of this suite: ref_al is the SemiAnalysis golden for this +# drafter at seven drafts (InferenceX golden_al_distribution/musespark_dflash.yaml, +# SPEED-Bench qualitative coding, T=1.0, top_p 0.95, rejection sampling), and +# min_al is that golden less 2%, the bar every submitted deployment must clear. +ModelingV2MuseSparkOpenNvfp4Tep4::dflash2_speed_bench_coding_golden: + ref_al: 4.05 + min_al: 3.969 diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 63fd1283c9e4..d5aeb0826caf 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -37,8 +37,10 @@ from datasets import load_dataset from defs.common import wait_for_reported_addr +from tensorrt_llm import LLM from tensorrt_llm.executor.result import GenerationResultBase -from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams +from tensorrt_llm.llmapi import (CompletionOutput, KvCacheConfig, RequestOutput, + SamplingParams) from tensorrt_llm.llmapi.llm_args import (DFlashDecodingConfig, DSparkDecodingConfig, Eagle3DecodingConfig, LlmArgs, @@ -52,7 +54,8 @@ from .accuracy_core import (GSM8K, MMLU, GSM8KInferenceX, LlmapiAccuracyTestHarness, acceptance_length_from_iteration_stats, - assert_acceptance_length, get_accuracy_task) + assert_acceptance_length, compute_acceptance_length, + get_accuracy_task, gsm8k_fewshot_prompts) class Result(GenerationResultBase): @@ -1754,6 +1757,96 @@ def test_gsm8k_1p1d_unified_draft_kv(self): "TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv", acceptance_length) + @pytest.mark.skip_less_device(2) + def test_default_1p1d_acceptance_matches_aggregated(self): + """With the default speculative_config, 1P1D drafts as well as aggregated. + + Neither worker names an attention backend or sets + TRTLLM_DISAGG_SHARED_DRAFT_KV, so AUTO must resolve to TRTLLM on both + and share the draft KV, even for this unwindowed drafter (aggregated + AUTO keeps it on VANILLA). The same greedy prompts then run on one + aggregated TRTLLM-drafter LLM and through 1P1D: the generation + worker's acceptance length must be within 5% of the aggregated one. + A generation worker drafting without the prompt loses ~17% here. + """ + prompts = gsm8k_fewshot_prompts(64) + sampling_params = SamplingParams(max_tokens=256, temperature=0.0) + kv_cache_config = { + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.5, + "use_kv_cache_manager_v2": True, + } + with LLM(self.MODEL_PATH, + speculative_config=DFlashDecodingConfig( + max_draft_len=self.MAX_DRAFT_LEN, + speculative_model=self.DRAFT_PATH, + attention_backend="TRTLLM"), + max_batch_size=16, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=KvCacheConfig(**kv_cache_config), + disable_overlap_scheduler=True, + enable_iter_perf_stats=True, + max_stats_len=-1) as agg_llm: + agg_llm.generate(prompts, sampling_params, use_tqdm=False) + aggregated_al = compute_acceptance_length(agg_llm) + + # No attention_backend: AUTO is what is under test. + speculative_config = { + "decoding_type": "DFlash", + "max_draft_len": self.MAX_DRAFT_LEN, + "speculative_model": self.DRAFT_PATH, + } + common_server_config = { + "disable_overlap_scheduler": True, + "max_batch_size": 16, + "max_seq_len": 4096, + "max_num_tokens": 4096, + "kv_cache_config": kv_cache_config, + "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, + "enable_iter_perf_stats": True, + "iter_stats_max_iterations": -1, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + }, + } + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + assert_worker_log_contains= + "target and draft KV layers share one KV cache manager") as llm: + futures = [ + llm.generate_async(prompt, sampling_params) + for prompt in prompts + ] + for future in futures: + future.result() + disaggregated_al = compute_disagg_acceptance_length(llm.serve_url) + print(f"[AL] aggregated {aggregated_al:.4f}, 1P1D " + f"{disaggregated_al:.4f} " + f"({100 * (disaggregated_al / aggregated_al - 1):+.2f}%)") + assert disaggregated_al >= 0.95 * aggregated_al, ( + f"[AL] 1P1D acceptance length {disaggregated_al:.3f} is below 95% " + f"of aggregated {aggregated_al:.3f}: the generation worker drafts " + "without the prompt context.") + @skip_pre_blackwell @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e25a688be596..ede3ed8957c7 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -53,8 +53,9 @@ GSM8K, MMLU, CnnDailymail, GPQADiamond, GSM8KInferenceX, JsonModeEval, LlmapiAccuracyTestHarness, LongBenchV1, LongBenchV2, assert_acceptance_length, assert_acceptance_length_for_llm, + assert_acceptance_length_survives_prefix_reuse, assert_guided_decoding_regex, assert_kv_cache_reuse_for_llm, - compute_acceptance_length) + compute_acceptance_length, gsm8k_fewshot_prompts) # isort: on @@ -3064,6 +3065,50 @@ def test_dflash(self): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_blackwell + @pytest.mark.parametrize("dflash_attention_backend", [ + "TRTLLM", + pytest.param( + "VANILLA", + marks=pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason="VANILLA keeps the drafter's context outside the KV " + "cache manager, so a prefix-cache hit drafts without the " + "reused prefix")), + ]) + def test_dflash_acceptance_under_prefix_reuse(self, + dflash_attention_backend): + """A prefix-cache hit must not cost the drafter its context. + + The same prompts run twice with block reuse: the second pass skips + each prompt's cached prefix, and its acceptance length must stay + within 5% of the first pass's. The target's output is identical on + both passes, so only a drafter that lost the reused prefix's context + can fail this. The drafter is unwindowed, the case AUTO still runs + on VANILLA in aggregated serving. + """ + spec_config = DFlashDecodingConfig( + max_draft_len=15, + speculative_model=f"{llm_models_root()}/Qwen3-8B-DFlash-b16", + attention_backend=dflash_attention_backend) + with LLM(model=f"{llm_models_root()}/Qwen3/Qwen3-8B", + max_batch_size=16, + max_seq_len=4096, + max_num_tokens=8192, + kv_cache_config=KvCacheConfig(enable_block_reuse=True, + free_gpu_memory_fraction=0.6, + tokens_per_block=32, + use_kv_cache_manager_v2=True), + cuda_graph_config=CudaGraphConfig(max_batch_size=16, + enable_padding=True), + speculative_config=spec_config, + enable_iter_perf_stats=True, + max_stats_len=-1) as llm: + assert_acceptance_length_survives_prefix_reuse( + llm, gsm8k_fewshot_prompts(64), + SamplingParams(max_tokens=256, temperature=0.0)) + @skip_pre_hopper @pytest.mark.parametrize( "attention_backend", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 7ac7879b6dc8..456a5268905f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -43,6 +43,8 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 - accuracy/test_llm_api_pytorch.py::TestLagunaXS_2_1::test_bf16_dflash - accuracy/test_llm_api_pytorch.py::TestLagunaXS_2_1::test_nvfp4_dflash + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse[TRTLLM] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse[VANILLA] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] diff --git a/tests/integration/test_lists/test-db/l0_cpu_musespark.yml b/tests/integration/test_lists/test-db/l0_cpu_musespark.yml index fc89e5845f60..3a3282a98158 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_musespark.yml @@ -24,6 +24,8 @@ l0_cpu_musespark: - unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py - unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py - unittest/llmapi/test_disagg_dflash_guard.py + # Cold-page KV-cache compression with a DFlash drafter on every attention backend. + - unittest/_torch/kv_cache_compression/test_cold_page_dflash.py - unittest/disaggregated/test_kv_transfer.py - unittest/disaggregated/test_peer.py - unittest/disaggregated/test_rank_info.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 23efa45be82e..f28646ea7dca 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -21,6 +21,7 @@ l0_dgx_b200: - 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 + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_default_1p1d_acceptance_matches_aggregated # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index d3d551c9419c..451c16b8ad9f 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -18,6 +18,8 @@ l0_gb300_multi_gpus: # One ordered sequence in one job: the cases share module-global rank state. - unittest/_torch/modeling_v2/comm - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_gsm8k_and_acceptance + - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_acceptance_under_prefix_reuse + - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_speed_bench_coding_golden_acceptance - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3 - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2] - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock] diff --git a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml index 96846251b211..4ec22820bc70 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml @@ -21,6 +21,9 @@ l0_gb300_musespark: tests: # Disabled: needs the DFlash2 draft model, which the CI model cache does not have. # - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_gsm8k_and_acceptance + # - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_acceptance_under_prefix_reuse + # Also needs the SPEED-Bench coding prompts staged under llm-models/datasets/nvidia/SPEED-Bench. + # - accuracy/modeling_v2/test_musespark_open.py::TestModelingV2MuseSparkOpenNvfp4Tep4::test_dflash2_speed_bench_coding_golden_acceptance # MuseSparkOpen GSM8K on the llm-models NVFP4 checkpoint; listed nowhere else. - accuracy/test_musespark_open.py::TestMuseSparkOpen::test_nvfp4_4gpus[eager] - accuracy/test_musespark_open.py::TestMuseSparkOpen::test_nvfp4_4gpus[cuda_graph] @@ -50,5 +53,10 @@ l0_gb300_musespark: - unittest/disaggregated/test_kv_transfer.py - unittest/disaggregated/test_bounce.py - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv + # DFlash acceptance length: a prefix-cache hit keeps the drafter's context (TRTLLM), and loses + # it on VANILLA (strict xfail); a default 1P1D config drafts as well as aggregated. + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse[TRTLLM] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash_acceptance_under_prefix_reuse[VANILLA] + - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_default_1p1d_acceptance_matches_aggregated # The Muse modeling_v2 composition over a TRTLLM drafter's paged context. - unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py index 655755ac4cc5..52aac554bf74 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py @@ -419,3 +419,31 @@ def test_rank_pins_auto_to_unify_like_trtllm(tmp_path, monkeypatch, is_disagg): assert creator._unified_drafter_attention_backend() == "TRTLLM" creator._check_shared_draft_kv_request() # TRTLLM_DISAGG_SHARED_DRAFT_KV=1 (autouse) assert creator._should_unify_dflash_kv_cache() + + +@pytest.mark.parametrize(("shared", "warns"), [(None, True), ("0", False)]) +def test_disagg_unwindowed_drafter_that_cannot_share_warns(monkeypatch, shared, warns): + """An unwindowed TRTLLM drafter kept out of the shared manager is named, not silent. + + Its separate pool is not transferred, so the generation worker drafts + without the prompt; TRTLLM_DISAGG_SHARED_DRAFT_KV=0 asks for exactly that. + """ + from tensorrt_llm._torch.pyexecutor import _util + + if shared is None: + monkeypatch.delenv(_SHARED_DRAFT_KV_ENV) + else: + monkeypatch.setenv(_SHARED_DRAFT_KV_ENV, shared) + warnings = [] + monkeypatch.setattr(_util.logger, "warning", lambda msg, *a, **k: warnings.append(msg)) + layout = _model_config(kv_heads=8, head_dim=128) + creator = _creator(is_disagg=True, target=layout, draft=layout, draft_windows=None) + creator._speculative_config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + creator._sparse_attention_config = None + creator._mapping.pp_size = 2 + + creator._check_shared_draft_kv_request() + assert not creator._should_unify_dflash_kv_cache() + assert any("drafts without the prompt" in str(w) for w in warnings) == warns + if warns: + assert any("pipeline_parallel_size is 2" in str(w) for w in warnings) diff --git a/tests/unittest/_torch/kv_cache_compression/test_cold_page_dflash.py b/tests/unittest/_torch/kv_cache_compression/test_cold_page_dflash.py new file mode 100644 index 000000000000..4f44deff4b9a --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_cold_page_dflash.py @@ -0,0 +1,70 @@ +# 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. +"""Cold-page KV-cache compression with a DFlash drafter on every attention backend. + +The codec compresses only the target's cold pages. A TRTLLM drafter, whose +context lives in a separate draft KV cache manager under compression, gets no +codec, so its pages move to and from the Host and Disk tiers losslessly. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm._torch.pyexecutor import _util as util_mod +from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache +from tensorrt_llm.llmapi.llm_args import ( + ColdPageQuantizationCompressionConfig, + DFlashDecodingConfig, + Eagle3DecodingConfig, +) + +pytestmark = pytest.mark.cpu_only + + +def _validate(spec_config) -> None: + with patch.object(util_mod, "is_sm_100f", return_value=True): + util_mod.validate_kv_cache_compression_compatibility( + ColdPageQuantizationCompressionConfig(), + SimpleNamespace(enable_block_reuse=True), + spec_config, + ) + + +@pytest.mark.parametrize("attention_backend", ["VANILLA", "FA4", "TRTLLM", "AUTO"]) +def test_cold_page_compression_admits_dflash_on_every_backend(attention_backend) -> None: + spec_config = DFlashDecodingConfig(max_draft_len=7, attention_backend=attention_backend) + _validate(spec_config) + + +@pytest.mark.parametrize( + ("attention_backend", "has_draft_manager"), + [("VANILLA", False), ("FA4", False), ("TRTLLM", True)], +) +def test_dflash_draft_manager_gets_no_cold_page_codec(attention_backend, has_draft_manager) -> None: + """Only the target's pages reach the codec, whichever backend holds the drafter's context.""" + spec_config = DFlashDecodingConfig(max_draft_len=7, attention_backend=attention_backend) + assert should_use_separate_draft_kv_cache(spec_config) is has_draft_manager + provider = object() + assert util_mod.draft_kv_cold_page_codec_provider(provider, spec_config) is None + + +def test_eagle3_draft_manager_keeps_the_cold_page_codec() -> None: + spec_config = Eagle3DecodingConfig(max_draft_len=3, speculative_model="/tmp/eagle3-draft") + _validate(spec_config) + provider = object() + assert util_mod.draft_kv_cold_page_codec_provider(provider, spec_config) is provider + assert util_mod.draft_kv_cold_page_codec_provider(None, spec_config) is None diff --git a/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py index acc11208d3a8..51678c0057e4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py +++ b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py @@ -1285,12 +1285,8 @@ def test_runtime_admission_is_checked_before_manager_creation(monkeypatch) -> No def test_speculative_admission_accepts_verified_one_model_modes() -> None: _validate_compression(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) _validate_compression(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) - for mode in ( - SpeculativeDecodingMode.MTP, - SpeculativeDecodingMode.DFLASH, - ): - with pytest.raises(ValueError, match="one-model MTP-EAGLE or EAGLE3"): - _validate_compression(mode) + with pytest.raises(ValueError, match="one-model MTP-EAGLE or EAGLE3"): + _validate_compression(SpeculativeDecodingMode.MTP) def test_qwen35_mtp3_resolves_to_supported_one_model_mode() -> None: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index a4a3228ffbaf..465192ff36b7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -4,6 +4,8 @@ TRTLLM borrows the managed paged pool. VANILLA and FA4 own private context buffers and must not allocate a redundant manager or be gated on its capacity. +AUTO resolves to TRTLLM wherever its kernels run the drafter, except an +unwindowed drafter on an aggregated server and under KV cache compression. """ import json @@ -110,6 +112,200 @@ def test_rank_pins_auto_before_choosing_its_draft_cache(tmp_path, monkeypatch) - assert dspark.attention_backend == "AUTO" +@pytest.mark.parametrize( + ("drafter", "unavailable", "disagg_shared", "resolved"), + [ + (_UNWINDOWED, None, True, "TRTLLM"), + (_UNWINDOWED, "no", True, "VANILLA"), + (_WINDOWED, None, True, "TRTLLM"), + (dict(_UNWINDOWED, num_attention_heads=8, head_dim=80), None, True, "VANILLA"), + (dict(_WINDOWED, num_attention_heads=64, num_key_value_heads=1), None, False, "VANILLA"), + (dict(_WINDOWED, num_attention_heads=64, num_key_value_heads=8), None, False, "TRTLLM"), + ], +) +def test_auto_resolves_to_trtllm_wherever_it_runs_the_drafter( + tmp_path, monkeypatch, drafter, unavailable, disagg_shared, resolved +) -> None: + """A sharing disaggregated worker takes TRTLLM for any drafter TRTLLM's kernels take.""" + from tensorrt_llm._torch.speculative import dflash + from tensorrt_llm.llmapi import llm_args + + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: unavailable) + (tmp_path / "config.json").write_text(json.dumps(drafter)) + config = DFlashDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + config.resolve_from_checkpoint(disagg_shared_draft_kv=disagg_shared) + assert config.attention_backend == resolved + + +def _cold_page_compression(): + from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig + + return ColdPageQuantizationCompressionConfig() + + +@pytest.mark.parametrize("drafter", [_WINDOWED, _UNWINDOWED]) +def test_auto_keeps_vanilla_under_kv_cache_compression(tmp_path, monkeypatch, drafter) -> None: + """With compression the drafter cannot share the target's KV cache manager.""" + from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _pin_dflash_auto_attention_backend, + ) + from tensorrt_llm._torch.speculative import dflash + from tensorrt_llm.llmapi import KvCacheConfig, TorchLlmArgs, llm_args + + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: None) + (tmp_path / "config.json").write_text(json.dumps(drafter)) + config = DFlashDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + config.resolve_from_checkpoint(kv_cache_compression=True) + assert config.attention_backend == "VANILLA" + assert should_use_separate_draft_kv_cache(config) is False + + def server(cuda: bool): + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: cuda) + return TorchLlmArgs( + model="dummy", + speculative_config=DFlashDecodingConfig( + max_draft_len=7, speculative_model=str(tmp_path), skip_ctx_buffer_budget_check=True + ), + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=1 << 30), + kv_cache_compression_config=_cold_page_compression(), + max_batch_size=8, + max_num_tokens=8192, + max_seq_len=8192, + ) + + # The frontend pins AUTO when it has a GPU, and a GPU-less frontend leaves + # it to the rank, which resolves it alike. + assert server(cuda=True).speculative_config.attention_backend == "VANILLA" + args = server(cuda=False) + assert args.speculative_config.attention_backend == "AUTO" + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: True) + _pin_dflash_auto_attention_backend(args.speculative_config, args) + assert args.speculative_config.attention_backend == "VANILLA" + + +def _disagg_args(tmp_path, spec_config, **overrides): + from tensorrt_llm.llmapi import CacheTransceiverConfig, KvCacheConfig, TorchLlmArgs + + kwargs = dict( + model="dummy", + speculative_config=spec_config, + cache_transceiver_config=CacheTransceiverConfig(backend="NIXL"), + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), + max_batch_size=8, + max_num_tokens=8192, + max_seq_len=8192, + ) + kwargs.update(overrides) + return TorchLlmArgs(**kwargs) + + +@pytest.mark.parametrize( + ("shared", "overrides", "resolved"), + [ + (None, {}, "TRTLLM"), + ("1", {}, "TRTLLM"), + ("0", {}, "VANILLA"), + (None, {"cache_transceiver_config": None}, "VANILLA"), + (None, {"kv_cache_config": "auto"}, "VANILLA"), + (None, {"attn_backend": "FLASHINFER"}, "VANILLA"), + (None, {"pipeline_parallel_size": 2}, "VANILLA"), + ], +) +def test_disagg_auto_shares_an_unwindowed_drafter_when_its_kv_can_be_shared( + tmp_path, monkeypatch, shared, overrides, resolved +) -> None: + """With VANILLA a disaggregated generation worker drafts without the prompt. + + So a worker that will share its draft KV runs an unwindowed drafter on + TRTLLM too; one that cannot (a reason the KV cache creator would refuse a + shared manager is already in the arguments) keeps VANILLA's faster kernel. + """ + from tensorrt_llm._torch.speculative import dflash + from tensorrt_llm.llmapi import KvCacheConfig, llm_args + + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: None) + if shared is None: + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", shared) + if overrides.get("kv_cache_config") == "auto": + overrides["kv_cache_config"] = KvCacheConfig(use_kv_cache_manager_v2="auto") + (tmp_path / "config.json").write_text(json.dumps(_UNWINDOWED)) + spec_config = DFlashDecodingConfig( + max_draft_len=7, speculative_model=str(tmp_path), skip_ctx_buffer_budget_check=True + ) + args = _disagg_args(tmp_path, spec_config, **overrides) + assert args.speculative_config.attention_backend == resolved + + +def test_rank_pins_an_unwindowed_auto_like_its_disaggregated_frontend( + tmp_path, monkeypatch +) -> None: + """A GPU-less frontend leaves AUTO open; the rank resolves it with the same sharing rule.""" + from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _pin_dflash_auto_attention_backend, + ) + from tensorrt_llm._torch.speculative import dflash + from tensorrt_llm.llmapi import llm_args + + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: False) + (tmp_path / "config.json").write_text(json.dumps(_UNWINDOWED)) + spec_config = DFlashDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + args = _disagg_args(tmp_path, spec_config) + assert args.speculative_config.attention_backend == "AUTO" + + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: None) + _pin_dflash_auto_attention_backend(args.speculative_config, args) + assert args.speculative_config.attention_backend == "TRTLLM" + + +@pytest.mark.parametrize( + ("drafter", "disagg", "overrides", "warns"), + [ + (_WINDOWED, False, {}, True), + (_UNWINDOWED, False, {}, False), + (_UNWINDOWED, True, {}, True), + # A disaggregated TRTLLM drafter could not share its KV with pp=2 either. + (_WINDOWED, True, {"pipeline_parallel_size": 2}, False), + ], +) +def test_explicit_vanilla_warns_where_auto_runs_the_drafter_on_trtllm( + tmp_path, monkeypatch, drafter, disagg, overrides, warns +) -> None: + """VANILLA stays selectable, but pinning it over the TRTLLM default names the lost context.""" + from tensorrt_llm._torch.speculative import dflash + from tensorrt_llm.llmapi import TorchLlmArgs, llm_args + + warnings = [] + monkeypatch.setattr(llm_args.logger, "warning", lambda msg, *a, **k: warnings.append(msg)) + monkeypatch.delenv("TRTLLM_DISAGG_SHARED_DRAFT_KV", raising=False) + monkeypatch.setattr(llm_args.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(dflash, "dflash_trtllm_gen_unavailability_reason", lambda: None) + (tmp_path / "config.json").write_text(json.dumps(drafter)) + spec_config = DFlashDecodingConfig( + max_draft_len=7, + speculative_model=str(tmp_path), + attention_backend="VANILLA", + skip_ctx_buffer_budget_check=True, + ) + if disagg: + args = _disagg_args(tmp_path, spec_config, **overrides) + else: + args = TorchLlmArgs( + model="dummy", + speculative_config=spec_config, + max_batch_size=8, + max_num_tokens=8192, + max_seq_len=8192, + ) + assert args.speculative_config.attention_backend == "VANILLA" + assert any("VANILLA is deprecated" in str(w) for w in warnings) == warns + + @pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM"]) def test_dflash_explicit_cache_opt_out_is_preserved(backend: str) -> None: config = DFlashDecodingConfig(max_draft_len=7, attention_backend=backend)