From 9db79ddcd5cb7234ad6ad908483e7a2ceae7ee07 Mon Sep 17 00:00:00 2001 From: zheyuf Date: Mon, 2 Nov 2026 23:21:28 -0800 Subject: [PATCH 1/5] [None][feat] DFlash: TRTLLM drafter attention by default where it runs, cold-page compression with a TRTLLM drafter, and acceptance-length tests - AUTO resolves to TRTLLM on a disaggregated server that shares the draft KV for unwindowed drafters too, when nothing known from the arguments keeps the target and draft KV out of one manager; it also checks the drafter's head geometry against the TRTLLM kernels. - The frontend, the rank and the hub-download path resolve AUTO alike. - An explicit VANILLA where AUTO runs the drafter on TRTLLM warns that it is deprecated and why (prefix-cache hits and disaggregated generation draft without the reused or transferred context). - A disaggregated TRTLLM drafter kept in a separate pool is named instead of silently drafting without the prompt. - Cold-page KV-cache compression admits DFlash on every backend; a TRTLLM drafter's separate draft KV cache manager gets no cold-page codec. - Tests: acceptance length under prefix-cache reuse (Qwen3-8B DFlash, TRTLLM and a strict-xfail VANILLA), default-config 1P1D vs aggregated acceptance length, Muse modeling_v2 golden SPEED-Bench acceptance and prefix reuse, and unit tests for the resolution, warnings and compression admission. Signed-off-by: zheyuf Co-Authored-By: Claude Opus 5.5 --- docs/source/features/speculative-decoding.md | 8 +- .../kv_cache_compression/nvfp4_cold_page.md | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 54 ++++++- .../_torch/pyexecutor/py_executor_creator.py | 16 +- tensorrt_llm/_torch/speculative/dflash.py | 61 +++++++- tensorrt_llm/commands/serve.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 113 ++++++++++++-- tensorrt_llm/llmapi/llm_utils.py | 18 ++- .../defs/accuracy/accuracy_core.py | 87 +++++++++++ .../modeling_v2/test_musespark_open.py | 98 +++++++++++- .../references/acceptance_length.yaml | 8 + .../accuracy/test_disaggregated_serving.py | 97 +++++++++++- .../defs/accuracy/test_llm_api_pytorch.py | 46 +++++- .../test_lists/test-db/l0_b200.yml | 2 + .../test_lists/test-db/l0_cpu_musespark.yml | 2 + .../test_lists/test-db/l0_dgx_b200.yml | 1 + .../test-db/l0_gb300_multi_gpus.yml | 2 + .../test_lists/test-db/l0_gb300_musespark.yml | 8 + .../kv_cache/test_dflash_unified_kv_gate.py | 28 ++++ .../test_cold_page_dflash.py | 70 +++++++++ .../test_quantization_for_cold_page.py | 8 +- .../hw_agnostic/test_dflash_draft_kv.py | 144 ++++++++++++++++++ 22 files changed, 828 insertions(+), 53 deletions(-) create mode 100644 tests/unittest/_torch/kv_cache_compression/test_cold_page_dflash.py diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index b5f585450493..ed6687ceba2e 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) 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..ae1785408375 100644 --- a/examples/kv_cache_compression/nvfp4_cold_page.md +++ b/examples/kv_cache_compression/nvfp4_cold_page.md @@ -68,8 +68,10 @@ 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 -parallelism is not currently supported. +unchanged. One-model MTP-EAGLE and EAGLE3 are supported. DFlash is supported +on an aggregated server: only the target's cold Pages are compressed, and a +TRTLLM drafter's separate draft KV cache 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 or Disk cache is also required for Pages to cross a compression boundary. 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..398b618e059f 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,16 @@ 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``). + """ 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)) def _flashinfer_one_engine_spec_supported(attn_backend: str, @@ -620,7 +626,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..014010fcc876 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -81,10 +81,57 @@ 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 +) -> 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 + one exception: an unwindowed drafter on an aggregated worker, where TRTLLM + runs the non-causal layer on its context kernel, 2-9x slower than VANILLA. + + 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. + + Returns: + ``"TRTLLM"`` or ``"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 +139,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..2c4b88ebd158 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3309,17 +3309,22 @@ 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 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 +3339,21 @@ 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) -> 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. """ if self.speculative_model is None: return @@ -3357,7 +3370,7 @@ 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) logger.info( f"DFlash attention_backend AUTO resolved to {self.attention_backend}." ) @@ -6023,6 +6036,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 +6872,40 @@ 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 resolve_dflash_auto_attention_backend( + draft_config, disagg_shared_draft_kv=auto_shares) != "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 +7175,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 +7188,10 @@ 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)) + 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..be8221c76138 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,16 @@ 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)) + 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..e5b9bd2c8554 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( @@ -134,3 +149,82 @@ 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): + llm.generate( + prompts, + SamplingParams(max_tokens=4096, temperature=1.0, top_p=0.95, seed=seed), + 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..61ff7e598890 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,49 @@ 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, + 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..62b294a49948 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. """ import json @@ -110,6 +112,148 @@ 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 _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", "warns"), + [(_WINDOWED, False, True), (_UNWINDOWED, False, False), (_UNWINDOWED, True, True)], +) +def test_explicit_vanilla_warns_where_auto_runs_the_drafter_on_trtllm( + tmp_path, monkeypatch, drafter, disagg, 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) + 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) From 29da6fb1cb3a23c754cee597d0482d0c1b5e65d1 Mon Sep 17 00:00:00 2001 From: zheyuf Date: Mon, 2 Nov 2026 23:30:04 -0800 Subject: [PATCH 2/5] [None][feat] DFlash: no VANILLA deprecation warning where a TRTLLM drafter could not share its KV On a disaggregated server whose arguments already rule out a shared target and draft KV cache manager, a TRTLLM drafter cannot carry the prompt context either, so pinning VANILLA there is not discouraged. The VANILLA prefix-reuse acceptance-length case expects only the assertion to fail. Signed-off-by: zheyuf Co-Authored-By: Claude Opus 5.5 --- tensorrt_llm/llmapi/llm_args.py | 3 +++ .../defs/accuracy/test_llm_api_pytorch.py | 1 + .../hw_agnostic/test_dflash_draft_kv.py | 14 ++++++++++---- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2c4b88ebd158..b71a8dd3edb1 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6894,6 +6894,9 @@ def _warn_if_dflash_vanilla_overrides_trtllm(self, disagg: bool) -> None: 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) != "TRTLLM": return diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 61ff7e598890..ede3ed8957c7 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3072,6 +3072,7 @@ def test_dflash(self): "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")), 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 62b294a49948..4efeb0139fdc 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 @@ -218,11 +218,17 @@ def test_rank_pins_an_unwindowed_auto_like_its_disaggregated_frontend( @pytest.mark.parametrize( - ("drafter", "disagg", "warns"), - [(_WINDOWED, False, True), (_UNWINDOWED, False, False), (_UNWINDOWED, True, True)], + ("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, warns + 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 @@ -241,7 +247,7 @@ def test_explicit_vanilla_warns_where_auto_runs_the_drafter_on_trtllm( skip_ctx_buffer_budget_check=True, ) if disagg: - args = _disagg_args(tmp_path, spec_config) + args = _disagg_args(tmp_path, spec_config, **overrides) else: args = TorchLlmArgs( model="dummy", From 386e5ebe5b493f532dd25ab4a78aa4f3f39c8cba Mon Sep 17 00:00:00 2001 From: zheyuf Date: Tue, 3 Nov 2026 01:04:07 -0800 Subject: [PATCH 3/5] [None][test] Muse modeling_v2 golden AL test: no second BOS on the chat-templated prompts The Muse chat template emits <|begin_of_text|> and tokenizer_config.json sets add_bos_token, so LLM.generate's default add_special_tokens=True can prefix a second BOS. The served chat endpoint, which the SemiAnalysis golden protocol goes through, tokenizes the templated prompt with add_special_tokens=False. Signed-off-by: zheyuf --- .../defs/accuracy/modeling_v2/test_musespark_open.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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 e5b9bd2c8554..59b20e4fa2a2 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py @@ -198,9 +198,17 @@ def test_dflash2_speed_bench_coding_golden_acceptance(self, monkeypatch): 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), + 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, ) From 844043028a065ef46125755eff2ffb1e1f50ea53 Mon Sep 17 00:00:00 2001 From: zheyuf Date: Tue, 3 Nov 2026 04:29:10 -0800 Subject: [PATCH 4/5] [None][fix] DFlash: AUTO keeps VANILLA under KV cache compression With KV cache compression the drafter cannot share the target's KV cache manager, so a TRTLLM drafter keeps a separate draft KV cache that prefix reuse does not reach. That pool is sized from the drafter's declared attention windows while its mirror reserves the full prompt, so on long agentic contexts it admits only a few requests at a time. AUTO now resolves to VANILLA there; an explicit TRTLLM still runs. Signed-off-by: zheyuf --- docs/source/features/speculative-decoding.md | 2 +- .../kv_cache_compression/nvfp4_cold_page.md | 8 ++-- .../_torch/pyexecutor/py_executor_creator.py | 8 +++- tensorrt_llm/_torch/speculative/dflash.py | 20 ++++++-- tensorrt_llm/llmapi/llm_args.py | 25 +++++++--- tensorrt_llm/llmapi/llm_utils.py | 5 +- .../hw_agnostic/test_dflash_draft_kv.py | 48 ++++++++++++++++++- 7 files changed, 98 insertions(+), 18 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index ed6687ceba2e..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. `"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) 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. +* `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 diff --git a/examples/kv_cache_compression/nvfp4_cold_page.md b/examples/kv_cache_compression/nvfp4_cold_page.md index ae1785408375..93830839a685 100644 --- a/examples/kv_cache_compression/nvfp4_cold_page.md +++ b/examples/kv_cache_compression/nvfp4_cold_page.md @@ -69,9 +69,11 @@ KVCacheManagerV2, and an SM100, SM103, or SM107 GPU. Hot Attention KV can use FP 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. DFlash is supported -on an aggregated server: only the target's cold Pages are compressed, and a -TRTLLM drafter's separate draft KV cache moves its Pages losslessly. HELIX -context parallelism is not currently 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 or Disk cache is also required for Pages to cross a compression boundary. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 398b618e059f..0d06ed04e739 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -228,13 +228,17 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType): 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``). + 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( disagg_shared_draft_kv=llm_args is not None - and dflash_auto_shares_draft_kv(llm_args, spec_config)) + 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, diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 014010fcc876..1ec6c7c6590d 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -110,7 +110,10 @@ def _dflash_trtllm_gen_shape_unsupported_reason(draft_config: Optional[dict]) -> def resolve_dflash_auto_attention_backend( - draft_config: Optional[dict], *, disagg_shared_draft_kv: bool = False + 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. @@ -118,8 +121,11 @@ def resolve_dflash_auto_attention_backend( 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 - one exception: an unwindowed drafter on an aggregated worker, where TRTLLM - runs the non-causal layer on its context kernel, 2-9x slower than VANILLA. + 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. @@ -127,10 +133,18 @@ def resolve_dflash_auto_attention_backend( 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: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b71a8dd3edb1..811d6fd03d13 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3319,9 +3319,10 @@ class DFlashDecodingConfig(DecodingBaseConfig): "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 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; " + "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.") @@ -3341,7 +3342,8 @@ def set_max_total_draft_tokens(self): def resolve_from_checkpoint(self, *, - disagg_shared_draft_kv: bool = False) -> None: + 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: @@ -3354,6 +3356,8 @@ def resolve_from_checkpoint(self, 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 @@ -3370,7 +3374,9 @@ def resolve_from_checkpoint(self, from tensorrt_llm._torch.speculative.dflash import \ resolve_dflash_auto_attention_backend self.attention_backend = resolve_dflash_auto_attention_backend( - draft_cfg, disagg_shared_draft_kv=disagg_shared_draft_kv) + 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}." ) @@ -6898,7 +6904,10 @@ def _warn_if_dflash_vanilla_overrides_trtllm(self, disagg: bool) -> None: # 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) != "TRTLLM": + 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") @@ -7192,7 +7201,9 @@ def validate_speculative_config(self): # covers only the token budget), and CachedModelLoader repeats # both after the drafter is downloaded. self.speculative_config.resolve_from_checkpoint( - disagg_shared_draft_kv=dflash_auto_shares_draft_kv(self)) + 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; diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index be8221c76138..2c3fbc6db39d 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -417,7 +417,10 @@ def __call__(self) -> Optional[Path]: if isinstance(spec_config, DFlashDecodingConfig): spec_config.resolve_from_checkpoint( disagg_shared_draft_kv=dflash_auto_shares_draft_kv( - self.llm_args)) + 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) 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 4efeb0139fdc..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 @@ -5,7 +5,7 @@ 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. +unwindowed drafter on an aggregated server and under KV cache compression. """ import json @@ -138,6 +138,52 @@ def test_auto_resolves_to_trtllm_wherever_it_runs_the_drafter( 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 From 03b61b526ed6f4e8c9472c1a4948165ca3b59b0c Mon Sep 17 00:00:00 2001 From: zheyuf Date: Tue, 3 Nov 2026 05:15:35 -0800 Subject: [PATCH 5/5] [None][test] Muse modeling_v2 tests: SM107 fused-kernel opt-in on Vera Rubin On sm107 the sharded projections and the fused MoE finalize all-reduce also need MUSESPARK_OPEN_SM107_FUSED_KERNELS=1, so without it the target refuses to build on Vera Rubin. sm103 ignores it. Signed-off-by: zheyuf --- .../defs/accuracy/modeling_v2/test_musespark_open.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 59b20e4fa2a2..78a9dff1f611 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_musespark_open.py @@ -57,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",