[None][fix] Keep target page size for standalone DSpark/DFlash draft KV pool when drafter kernels accept it

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -88,13 +88,16 @@
 
 GB = 1 << 30
 
-# Page size of a standalone drafter's separate KV pool (DFlash, standalone
-# DSpark). Deliberately not kv_cache_config.tokens_per_block: that describes
-# the target pool, and target attention backends and deployment configs set it
-# to values the drafter's block-decode kernels reject (FP4 MLA forces 128, which
-# the trtllm-gen MLA decode does not take). The drafter adopts whatever page
-# size its pool is carved with, so the pool picks one every drafter kernel
-# accepts. Matches the KvCacheConfig.tokens_per_block default.
+# Page sizes of a standalone drafter's separate KV pool (DFlash, standalone
+# DSpark). The pool keeps the target's kv_cache_config.tokens_per_block when it
+# is in _DRAFTER_KERNEL_TOKENS_PER_BLOCK, the page sizes every drafter
+# block-decode kernel accepts: trtllm-gen MLA decode takes 32 or 64, the CuTe
+# DSL MLA decode a divisor of 128 above 1, and the trtllm-gen attention of a GQA
+# drafter 16, 32 or 64. Target attention backends and deployment configs set
+# other values (FP4 MLA forces 128), and those get
+# _STANDALONE_DRAFT_TOKENS_PER_BLOCK, the KvCacheConfig.tokens_per_block
+# default. The drafter adopts whatever page size its pool is carved with.
+_DRAFTER_KERNEL_TOKENS_PER_BLOCK = frozenset((32, 64))
 _STANDALONE_DRAFT_TOKENS_PER_BLOCK = 32
 
 # Aggregated helix skips memory profiling, so its default KV fraction leaves
@@ -2281,11 +2284,13 @@ def _get_draft_kv_model_config(self) -> ModelConfig:
     def _get_draft_tokens_per_block(self) -> int:
         """Page size of the separate one-model draft KV pool.
 
-        A standalone drafter (DFlash, standalone DSpark) gets its own page
-        size, ``_STANDALONE_DRAFT_TOKENS_PER_BLOCK``, NOT self._tokens_per_block:
-        the target's value can be one the drafter's kernels reject (FP4 MLA
-        forces 128; trtllm-gen MLA decode takes only 32 or 64). Those drafters
-        read their pool geometry from the manager, so nothing else has to agree.
+        A standalone drafter (DFlash, standalone DSpark) keeps
+        self._tokens_per_block when every drafter block-decode kernel accepts
+        it (``_DRAFTER_KERNEL_TOKENS_PER_BLOCK``), and otherwise gets its own
+        page size, ``_STANDALONE_DRAFT_TOKENS_PER_BLOCK``: the target's value
+        can be one the drafter's kernels reject (FP4 MLA forces 128;
+        trtllm-gen MLA decode takes only 32 or 64). Those drafters read their
+        pool geometry from the manager, so nothing else has to agree.
 
         Every other separate draft pool (MTP, Eagle3, ...) keeps the target's
         page size: its draft layers run through the target's attention
@@ -2296,7 +2301,8 @@ def _get_draft_tokens_per_block(self) -> int:
         the split charges one page granularity and the pool uses another.
         """
         spec_dec_mode = self._speculative_config.spec_dec_mode
-        if spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark():
+        if ((spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark()) and
+                self._tokens_per_block not in _DRAFTER_KERNEL_TOKENS_PER_BLOCK):
             return _STANDALONE_DRAFT_TOKENS_PER_BLOCK
         return self._tokens_per_block
 

tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py

@@ -29,7 +29,13 @@
 from tensorrt_llm._torch.pyexecutor.config_utils import uses_vswa_kv_cache_layout
 from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode
-from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, KvCacheConfig, MTPDecodingConfig
+from tensorrt_llm.llmapi.llm_args import (
+    DFlashDecodingConfig,
+    DSparkDecodingConfig,
+    Eagle3DecodingConfig,
+    KvCacheConfig,
+    MTPDecodingConfig,
+)
 from tensorrt_llm.mapping import CpType, Mapping
 
 pytestmark = pytest.mark.cpu_only
@@ -1257,42 +1263,73 @@ def test_non_external_drafter_is_untouched(self, mocker):
 
 
 class TestStandaloneDrafterPageSize:
-    """A standalone drafter's separate pool has its own page size.
+    """A standalone drafter's pool keeps the target's page size only where its kernels accept it.
 
     The target's tokens_per_block is whatever its attention backend or the
     deployment config chose (FP4 MLA forces 128), and the drafter's block-decode
     kernels need not accept it: trtllm-gen MLA decode takes 32 or 64 only. The
-    pool is built with the drafter's page size, and the budget split must charge
-    the same one, since per-request quotas round up to whole pages.
+    pool keeps a target page size every drafter kernel accepts and is built with
+    the drafter's page size otherwise, and the budget split must charge the same
+    one, since per-request quotas round up to whole pages.
     """
 
-    TARGET_TOKENS_PER_BLOCK = 128
+    KEPT_AND_REPLACED = [(64, 64), (128, _STANDALONE_DRAFT_TOKENS_PER_BLOCK)]
 
-    def _creator(self, mocker):
+    def _creator(self, mocker, target_tokens_per_block):
         from tensorrt_llm.models.modeling_utils import QuantConfig
 
         c, draft_model_config, _ = TestExternalDrafterKvDtype()._creator(mocker, QuantConfig())
-        c._tokens_per_block = self.TARGET_TOKENS_PER_BLOCK
+        c._tokens_per_block = target_tokens_per_block
         # Read by the draft manager-class resolution on both the cost and the
         # allocation path; the resolution itself is not under test here.
         c._cache_transceiver_config = None
         c._max_beam_width = 1
         c._validate_or_fallback_kv_cache_manager_v2 = Mock(side_effect=lambda cls, *_: cls)
         return c, draft_model_config
 
-    def test_standalone_drafter_page_size_is_independent_of_the_target(self, mocker):
-        c, _ = self._creator(mocker)
-        assert c._get_draft_tokens_per_block() == _STANDALONE_DRAFT_TOKENS_PER_BLOCK
-        assert c._get_draft_tokens_per_block() != c._tokens_per_block
+    @pytest.mark.parametrize(
+        "spec_config",
+        [DSparkDecodingConfig(max_draft_len=4), DFlashDecodingConfig(max_draft_len=4)],
+        ids=["dspark", "dflash"],
+    )
+    @pytest.mark.parametrize(
+        "target_tokens_per_block, draft_tokens_per_block",
+        [
+            (64, 64),
+            (32, 32),
+            (128, _STANDALONE_DRAFT_TOKENS_PER_BLOCK),
+            (16, _STANDALONE_DRAFT_TOKENS_PER_BLOCK),
+        ],
+    )
+    def test_standalone_drafter_keeps_a_target_page_size_its_kernels_accept(
+        self, mocker, spec_config, target_tokens_per_block, draft_tokens_per_block
+    ):
+        c, _ = self._creator(mocker, target_tokens_per_block)
+        c._speculative_config = spec_config
+        assert c._get_draft_tokens_per_block() == draft_tokens_per_block
 
-    def test_draft_layers_of_the_target_keep_the_target_page_size(self, mocker):
+    @pytest.mark.parametrize(
+        "spec_config",
+        [
+            MTPDecodingConfig(num_nextn_predict_layers=1, max_draft_len=4),
+            Eagle3DecodingConfig(max_draft_len=4, speculative_model="draft-model"),
+        ],
+        ids=["mtp", "eagle3"],
+    )
+    @pytest.mark.parametrize("target_tokens_per_block", [64, 128])
+    def test_draft_layers_of_the_target_keep_the_target_page_size(
+        self, mocker, spec_config, target_tokens_per_block
+    ):
         """MTP/Eagle3 draft layers run on the target's attention metadata."""
-        c, _ = self._creator(mocker)
-        c._speculative_config = MTPDecodingConfig(num_nextn_predict_layers=1, max_draft_len=4)
-        assert c._get_draft_tokens_per_block() == self.TARGET_TOKENS_PER_BLOCK
+        c, _ = self._creator(mocker, target_tokens_per_block)
+        c._speculative_config = spec_config
+        assert c._get_draft_tokens_per_block() == target_tokens_per_block
 
-    def test_cost_charges_the_draft_page_size(self, mocker):
-        c, _ = self._creator(mocker)
+    @pytest.mark.parametrize("target_tokens_per_block, draft_tokens_per_block", KEPT_AND_REPLACED)
+    def test_cost_charges_the_draft_page_size(
+        self, mocker, target_tokens_per_block, draft_tokens_per_block
+    ):
+        c, _ = self._creator(mocker, target_tokens_per_block)
         size_spy = mocker.spy(KVCacheManagerV2, "get_cache_size_per_token")
         bytes_spy = mocker.spy(KVCacheManagerV2, "get_cache_bytes_per_request")
         quota_spy = mocker.spy(KVCacheManagerV2, "get_cache_quota_per_request")
@@ -1305,12 +1342,13 @@ def test_cost_charges_the_draft_page_size(self, mocker):
 
         for spy in (size_spy, bytes_spy, quota_spy):
             assert spy.call_count == 1, spy
-            assert spy.call_args.kwargs["tokens_per_block"] == (
-                _STANDALONE_DRAFT_TOKENS_PER_BLOCK
-            ), spy
+            assert spy.call_args.kwargs["tokens_per_block"] == draft_tokens_per_block, spy
 
-    def test_allocation_uses_the_page_size_the_cost_charged(self, mocker):
-        c, draft_model_config = self._creator(mocker)
+    @pytest.mark.parametrize("target_tokens_per_block, draft_tokens_per_block", KEPT_AND_REPLACED)
+    def test_allocation_uses_the_page_size_the_cost_charged(
+        self, mocker, target_tokens_per_block, draft_tokens_per_block
+    ):
+        c, draft_model_config = self._creator(mocker, target_tokens_per_block)
         draft_model_config.sparse_attention_config = None
         c._model_engine.max_cuda_graph_batch_size = 1
         c._skip_est = False
@@ -1327,7 +1365,7 @@ def test_allocation_uses_the_page_size_the_cost_charged(self, mocker):
 
         assert create.call_count == 1
         allocated = create.call_args.kwargs["tokens_per_block"]
-        assert allocated == _STANDALONE_DRAFT_TOKENS_PER_BLOCK
+        assert allocated == draft_tokens_per_block
         assert size_spy.call_args.kwargs["tokens_per_block"] == allocated