From 63c992191b75a3b67d1f11c7ebb1ee443b2a5003 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 26 Oct 2026 13:50:09 -0700 Subject: [PATCH] [None][fix] Keep target page size for standalone DSpark/DFlash draft KV pool when drafter kernels accept it A standalone drafter's separate KV pool keeps the target's tokens_per_block when every drafter block-decode kernel accepts it (32 or 64) and falls back to 32 otherwise, as before. 128-page targets (FP4 MLA) still get 32-token draft pages; MTP and Eagle3 draft pools keep the target's page size, unchanged. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/pyexecutor/_util.py | 32 ++++--- .../kv_cache/test_kv_cache_budget_split.py | 84 ++++++++++++++----- 2 files changed, 80 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b404db1c4087..ed8b16f6ae74 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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 @@ -2234,11 +2237,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 @@ -2249,7 +2254,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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 467da9f44a5b..ca7ef6186038 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/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, +) pytestmark = pytest.mark.cpu_only @@ -1161,22 +1167,23 @@ 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 @@ -1184,19 +1191,49 @@ def _creator(self, mocker): 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") @@ -1209,12 +1246,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 @@ -1231,7 +1269,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