[None][fix] DFlash buffer sizing/position fixes and an accurate FlashInfer speculation gate

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -35,7 +35,8 @@
 from ..attention.backends.interface import AttentionRuntimeFeatures
 from ..distributed import Distributed
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
-                           get_spec_resource_manager)
+                           get_spec_resource_manager,
+                           should_use_separate_draft_kv_cache)
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     compute_max_num_sequences, create_py_executor_instance,
@@ -190,6 +191,24 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType):
                 ))
 
 
+def _flashinfer_one_engine_spec_supported(attn_backend: str,
+                                          spec_config) -> bool:
+    """Whether this speculation is qualified on the FlashInfer target backend.
+
+    DFlash remains unqualified regardless of its draft attention backend or
+    cache ownership. Its VANILLA and FA4 backends own private context buffers,
+    but that does not establish FlashInfer serving-scale qualification.
+
+    Other speculative modes are admitted when they do not need a separate
+    draft KV cache manager.
+    """
+    if spec_config is None or attn_backend != "FLASHINFER":
+        return True
+    if spec_config.spec_dec_mode.is_dflash():
+        return False
+    return not should_use_separate_draft_kv_cache(spec_config)
+
+
 def _set_model_engines_cache_reuse(model_engines, cache_reuse: bool):
     for engine in model_engines:
         if engine is None:
@@ -439,12 +458,18 @@ def _create_py_executor_impl(
             )
             llm_args.disable_overlap_scheduler = True
 
-    if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
-            and spec_config.spec_dec_mode.use_one_engine()
-            and not spec_config._use_shared_kv_cache):
+    if not _flashinfer_one_engine_spec_supported(llm_args.attn_backend,
+                                                 spec_config):
+        if spec_config.spec_dec_mode.is_dflash():
+            raise ValueError(
+                "FLASHINFER target attention is not qualified for DFlash, "
+                "regardless of the draft attention backend or cache ownership. "
+                "Use TRTLLM target attention for DFlash.")
         raise ValueError(
-            "FLASHINFER attention backend supports one-engine speculative "
-            "decoding only when the draft model shares the target KV cache.")
+            f"FLASHINFER target attention is not qualified for "
+            f"{spec_config.spec_dec_mode.name}: this one-engine speculative "
+            "mode needs a separate draft KV cache manager, which FLASHINFER "
+            "does not support. Use TRTLLM target attention.")
 
     if mm_encoder_only:
         llm_args.mm_encoder_only = True

tensorrt_llm/_torch/speculative/interface.py

@@ -136,6 +136,11 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
     if (spec_config.spec_dec_mode.is_dspark()
             and spec_config.draft_is_embedded_in_target):
         return False
+    # Suffix automaton (SA) drafts from a dedicated suffix-automaton state pool
+    # and never reads a paged draft KV cache manager, so it needs no separate
+    # draft KV cache despite reaching this one-engine path.
+    if spec_config.spec_dec_mode.is_sa():
+        return False
     return spec_config._allow_separate_draft_kv_cache
 
 

tensorrt_llm/llmapi/llm_args.py

@@ -2984,6 +2984,14 @@ class DFlashDecodingConfig(DecodingBaseConfig):
         "and supports SM100/SM103 only. FA4 uses the flash-attn CuTe DSL kernels "
         "on the same paged cache and supports SM90 only.")
 
+    skip_ctx_buffer_budget_check: bool = Field(
+        default=False,
+        description=
+        "Skip the config-time check that the pooled-context K/V buffers fit in "
+        "the memory left free after the KV-cache pool commits. The estimate is "
+        "deliberately conservative; set this to proceed when you know the "
+        "device has the headroom. The token-budget check is not affected.")
+
     @model_validator(mode="after")
     def set_max_total_draft_tokens(self):
         self.max_total_draft_tokens = self.max_draft_len
@@ -6287,6 +6295,128 @@ def normalize_prefill_cuda_graph_config(self) -> 'TorchLlmArgs':
 
         return self
 
+    def _kv_cache_estimation_runs(self) -> bool:
+        """Whether the executor will profile a forward to size the KV pool.
+
+        Mirrors the config-time-knowable conditions of
+        ``KvCacheCreator.try_prepare_estimation``: the
+        ``TRTLLM_SKIP_KV_CACHE_ESTIMATION`` env var, a VANILLA target
+        attention backend, and context parallelism all skip estimation.
+        Encoder-decoder targets also skip it but are only known from the
+        model config at load time; they are treated as estimating here,
+        which errs toward the looser arena budget.
+        """
+        if os.environ.get("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "0") == "1":
+            return False
+        if self.attn_backend == "VANILLA":
+            return False
+        if self.cp_config is not None:
+            return False
+        return True
+
+    def _validate_dflash_ctx_budget(self,
+                                    memory_budget_bytes: Optional[int] = None
+                                    ) -> None:
+        """Fail at config time on DFlash setups that would die late.
+
+        Delegates to ``validate_dflash_ctx_buffer_budget``: the token-budget
+        rule (``max_batch_size * (1 + max_draft_len) <= max_num_tokens``) and
+        the pooled-context K/V buffer fit, both of which otherwise fail only
+        once the model is loaded and the first forward runs.
+
+        The buffer budget depends on when the arena is allocated relative to
+        the KV-cache pool (``_kv_cache_estimation_runs``). With estimation
+        (the default) the arena is allocated inside the estimation forward,
+        before the pool is sized, so it must fit the device total less the
+        runtime overhead, the estimated per-rank weight footprint (checkpoint
+        shard bytes on disk divided across TP x PP ranks) and a reserve for
+        activations and CUDA graphs; ``free_gpu_memory_fraction`` bounds the
+        pool that is sized afterwards, not the arena. With estimation skipped
+        the pool commits first as ``free_gpu_memory_fraction`` of the
+        post-load memory, and the arena gets ``(1 - fraction)`` of that. See
+        ``derive_dflash_ctx_memory_budget_bytes`` for both forms. The check
+        is only enforced when the KV pool is sized by fraction (an explicit
+        ``kv_cache_config.max_tokens`` cap can leave more headroom than the
+        fraction implies) and when a CUDA device is visible.
+        ``memory_budget_bytes`` overrides the derivation (tests).
+        """
+        from tensorrt_llm._torch.speculative.dflash import (
+            derive_dflash_ctx_memory_budget_bytes,
+            estimate_checkpoint_weight_bytes, validate_dflash_ctx_buffer_budget)
+
+        spec_cfg = self.speculative_config
+
+        # Token budget: max_batch_size * (1 + max_draft_len) must fit
+        # max_num_tokens or the config corrupts memory at engine init. Under
+        # the stock defaults (max_batch_size=2048, max_num_tokens=8192) this is
+        # violated for any max_draft_len >= 4, so a plain
+        # DFlashDecodingConfig(max_draft_len=7) would fail at construction.
+        # When max_batch_size was not set explicitly, clamp it to what fits and
+        # warn; an explicitly-set max_batch_size is left to hard-fail in the
+        # validator below rather than silently overridden.
+        K = spec_cfg.max_draft_len
+        if (self.max_num_tokens is not None and self.max_batch_size is not None
+                and "max_batch_size" not in self.model_fields_set):
+            tokens_per_req = 1 + K
+            if self.max_batch_size * tokens_per_req > self.max_num_tokens:
+                clamped = self.max_num_tokens // tokens_per_req
+                if clamped >= 1:
+                    logger.warning(
+                        f"DFlash: max_batch_size ({self.max_batch_size}) x (1 "
+                        f"+ max_draft_len ({K})) exceeds max_num_tokens "
+                        f"({self.max_num_tokens}); clamping max_batch_size to "
+                        f"{clamped}. Set max_num_tokens explicitly to raise "
+                        "it.")
+                    self.max_batch_size = clamped
+                # clamped < 1 means not even one request fits; leave
+                # max_batch_size so the validator below raises with the
+                # actionable token-budget message.
+
+        draft_config = None
+        if spec_cfg.speculative_model is not None:
+            draft_config_path = os.path.join(str(spec_cfg.speculative_model),
+                                             "config.json")
+            if os.path.exists(draft_config_path):
+                with open(draft_config_path) as f:
+                    draft_config = json.load(f)
+
+        arena_before_pool = self._kv_cache_estimation_runs()
+        if spec_cfg.skip_ctx_buffer_budget_check:
+            # Opt out of the pooled-context buffer-fit check (the estimate is
+            # conservative); the token-budget check above still applies.
+            memory_budget_bytes = None
+        elif memory_budget_bytes is None:
+            kv_fraction = self.kv_cache_config.free_gpu_memory_fraction
+            if (kv_fraction is not None
+                    and self.kv_cache_config.max_tokens is None
+                    and torch.cuda.is_available()):
+                total = torch.cuda.get_device_properties(0).total_memory
+                per_rank_weight_bytes = None
+                checkpoint_bytes = estimate_checkpoint_weight_bytes(
+                    str(self.model))
+                if checkpoint_bytes is not None:
+                    weight_shards = max(
+                        1,
+                        self.tensor_parallel_size * self.pipeline_parallel_size)
+                    per_rank_weight_bytes = checkpoint_bytes // weight_shards
+                memory_budget_bytes = derive_dflash_ctx_memory_budget_bytes(
+                    total,
+                    kv_fraction,
+                    per_rank_weight_bytes,
+                    arena_before_pool=arena_before_pool)
+
+        validate_dflash_ctx_buffer_budget(
+            max_batch_size=self.max_batch_size,
+            max_num_tokens=self.max_num_tokens,
+            max_seq_len=self.max_seq_len,
+            max_draft_len=spec_cfg.max_draft_len,
+            attention_backend=spec_cfg.attention_backend,
+            draft_config=draft_config,
+            tp_size=self.tensor_parallel_size,
+            memory_budget_bytes=memory_budget_bytes,
+            arena_before_pool=arena_before_pool,
+        )
+
     @model_validator(mode="after")
     def validate_speculative_config(self):
         if self.speculative_config:
@@ -6457,9 +6587,12 @@ def validate_speculative_config(self):
                         "unaffected; expect a lower acceptance rate than the "
                         "same configuration run aggregated.")
                 assert self.speculative_config.max_draft_len > 0, "DFlash max_draft_len must be > 0"
-                # A Hugging Face repo id is not readable yet; CachedModelLoader
-                # calls this again after the drafter is downloaded.
+                # A Hugging Face repo id is not readable yet: both calls below
+                # then run without the drafter's config.json (the budget check
+                # covers only the token budget), and CachedModelLoader repeats
+                # both after the drafter is downloaded.
                 self.speculative_config.resolve_from_checkpoint()
+                self._validate_dflash_ctx_budget()
 
             if isinstance(self.speculative_config, DSparkDecodingConfig):
                 spec_cfg = self.speculative_config

tensorrt_llm/llmapi/llm_utils.py

@@ -25,11 +25,12 @@
                                             warn_if_inline_diverges)
 # yapf: disable
 from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig,
-                       DraftTargetDecodingConfig, Eagle3DecodingConfig,
-                       EagleDecodingConfig, EncodeCudaGraphConfig,
-                       KvCacheConfig, KVEventsConfig, LlmArgs,
-                       MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig,
-                       TorchLlmArgs, UserProvidedDecodingConfig, _ModelWrapper,
+                       DFlashDecodingConfig, DraftTargetDecodingConfig,
+                       Eagle3DecodingConfig, EagleDecodingConfig,
+                       EncodeCudaGraphConfig, KvCacheConfig, KVEventsConfig,
+                       LlmArgs, MTPDecodingConfig, NGramDecodingConfig,
+                       SchedulerConfig, TorchLlmArgs,
+                       UserProvidedDecodingConfig, _ModelWrapper,
                        _ParallelConfig, update_llm_args_with_extra_dict,
                        update_llm_args_with_extra_options)
 # yapf: enable
@@ -432,6 +433,15 @@ def __call__(self) -> Tuple[Path, Union[Path, None]]:
                     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
+                # checkpoint is local; a local path was fully checked at
+                # validation time and is not re-run.
+                if (isinstance(self.llm_args, TorchLlmArgs)
+                        and isinstance(self.llm_args.speculative_config,
+                                       DFlashDecodingConfig)):
+                    self.llm_args._validate_dflash_ctx_budget()
 
         # AutoDeploy doesn't use ModelLoader
         if self.llm_args.backend == "_autodeploy":

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -1707,6 +1707,11 @@
       "kind": "value",
       "path": "speculative_config.sa_config.threshold"
     },
+    {
+      "capture_policy": "bool",
+      "kind": "value",
+      "path": "speculative_config.skip_ctx_buffer_budget_check"
+    },
     {
       "capture_policy": "list[int]|none",
       "kind": "value",

tests/unittest/_torch/executor/test_py_executor_creator_flashinfer_spec_gate.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer qualification is independent of DFlash cache ownership.
+
+DFlash remains refused for every draft backend, including those using private
+buffers. Other modes retain the existing separate-manager qualification gate.
+"""
+
+import pytest
+
+from tensorrt_llm._torch.pyexecutor.py_executor_creator import _flashinfer_one_engine_spec_supported
+from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
+from tensorrt_llm.llmapi import (
+    DFlashDecodingConfig,
+    MTPDecodingConfig,
+    NGramDecodingConfig,
+    SADecodingConfig,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+
+def test_no_speculation_is_admitted():
+    assert _flashinfer_one_engine_spec_supported("FLASHINFER", None)
+
+
+def test_suffix_automaton_is_admitted_on_flashinfer():
+    config = SADecodingConfig(max_draft_len=4)
+    assert config.spec_dec_mode.use_one_engine()
+    assert not config._use_shared_kv_cache
+    assert _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+@pytest.mark.parametrize(
+    ("backend", "uses_managed_pool"),
+    [("VANILLA", False), ("FA4", False), ("TRTLLM", True)],
+)
+def test_dflash_is_refused_on_flashinfer(backend: str, uses_managed_pool: bool) -> None:
+    config = DFlashDecodingConfig(max_draft_len=7, attention_backend=backend)
+    assert should_use_separate_draft_kv_cache(config) is uses_managed_pool
+    assert not _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+    assert _flashinfer_one_engine_spec_supported("TRTLLM", config)
+
+
+def test_default_dflash_is_refused_without_a_separate_manager() -> None:
+    config = DFlashDecodingConfig(max_draft_len=7)
+    assert not should_use_separate_draft_kv_cache(config)
+    assert not _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+@pytest.mark.parametrize("backend", ["VANILLA", "FA4", "TRTLLM"])
+def test_dflash_cache_opt_out_does_not_lift_flashinfer_gate(backend: str) -> None:
+    config = DFlashDecodingConfig(max_draft_len=7, attention_backend=backend)
+    config._allow_separate_draft_kv_cache = False
+    assert not should_use_separate_draft_kv_cache(config)
+    assert not _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+def test_external_drafters_are_admitted_on_flashinfer():
+    config = NGramDecodingConfig(max_draft_len=4, max_matching_ngram_size=2)
+    assert not config.spec_dec_mode.use_one_engine()
+    assert _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+def test_one_engine_drafters_with_a_draft_model_are_refused_on_flashinfer():
+    config = MTPDecodingConfig(num_nextn_predict_layers=1)
+    assert config.spec_dec_mode.use_one_engine()
+    assert not _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+def test_shared_kv_cache_is_admitted_on_flashinfer():
+    config = MTPDecodingConfig(num_nextn_predict_layers=1)
+    config._use_shared_kv_cache = True
+    assert _flashinfer_one_engine_spec_supported("FLASHINFER", config)
+
+
+def test_other_backends_are_not_gated():
+    config = MTPDecodingConfig(num_nextn_predict_layers=1)
+    assert _flashinfer_one_engine_spec_supported("TRTLLM", config)