[None][feat] DFlash: TRTLLM drafter attention by default wherever it runs, cold-page compression with a TRTLLM drafter, and acceptance-length tests

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.
 

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

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],
     *,

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():

tensorrt_llm/_torch/speculative/dflash.py

@@ -81,17 +81,84 @@ 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(
             f"DFlash attention_backend AUTO keeps VANILLA: TRTLLM is unavailable ({reason}).",
             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"
 
 

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 "