From 7492ac12a3b13d6be1af44cffa79577371ec8b1b Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 23:42:55 -0700 Subject: [PATCH] [None][fix] Use 128-token DSpark draft pages only where the dynamic-page kernels are verified (SM100/SM103) dflash_trtllm_supports_page_size allowed power-of-two pages from 128 tokens wherever the DFlash TRTLLM ops run (SM100-SM109). On SM107 the DSpark draft pool therefore took the target's 128-token pages and paired with it for joint block reuse, which the disaggregated draft-window transfer refuses at executor creation; flashinfer's dynamic page-size kernels are also verified on SM100/SM103 only. Allow the dynamic sizes there only; elsewhere the pool keeps the static sizes (64 tokens, unpaired) as before. Signed-off-by: Zheyu Fu --- .../_torch/speculative/dflash_attention.py | 5 +++- .../kv_cache/test_kv_cache_estimation.py | 24 ++++++++++++++----- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash_attention.py b/tensorrt_llm/_torch/speculative/dflash_attention.py index 63274115deff..3927de7a647d 100644 --- a/tensorrt_llm/_torch/speculative/dflash_attention.py +++ b/tensorrt_llm/_torch/speculative/dflash_attention.py @@ -168,6 +168,8 @@ def dflash_trtllm_supports_page_size(tokens_per_block: int, *, has_context_atten not has_context_attention and is_power_of_two and tokens_per_block >= _TRTLLM_GEN_DYNAMIC_PAGE_MIN_TOKENS + # The dynamic page-size kernels are verified on SM100 and SM103 only. + and get_sm_version() in (100, 103) ) @@ -206,7 +208,8 @@ def validate_dflash_trtllm_gen_runtime( raise RuntimeError( "DFlash TRTLLM attention does not support context-cache page size " f"{tokens_per_block}; supported sizes are {supported}, plus powers of two " - f">= {_TRTLLM_GEN_DYNAMIC_PAGE_MIN_TOKENS} without the non-causal context kernel." + f">= {_TRTLLM_GEN_DYNAMIC_PAGE_MIN_TOKENS} on SM100/SM103 without the non-causal " + "context kernel." ) if has_context_attention and head_dim in FlashInferTrtllmGenFmha.UNSUPPORTED_HEAD_SIZES_CONTEXT: raise RuntimeError( diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index cf7cac06cace..003b33219835 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -1514,17 +1514,29 @@ def test_premise_trtllm_gen_page_sizes(self): assert max(size for size in supported if size < 128) == 64 @pytest.mark.parametrize( - "tokens_per_block, has_context_attention, expected", - [(64, True, True), (128, False, True), (128, True, False), (96, False, False)], + "tokens_per_block, has_context_attention, sm_version, expected", + [ + (64, True, 103, True), + (128, False, 103, True), + (128, False, 107, False), + (128, True, 103, False), + (96, False, 103, False), + ], ) - def test_trtllm_page_size_capability(self, tokens_per_block, has_context_attention, expected): + def test_trtllm_page_size_capability( + self, tokens_per_block, has_context_attention, sm_version, expected + ): from tensorrt_llm._torch.speculative.dflash_attention import ( dflash_trtllm_supports_page_size, ) - supported = dflash_trtllm_supports_page_size( - tokens_per_block, has_context_attention=has_context_attention - ) + with patch( + "tensorrt_llm._torch.speculative.dflash_attention.get_sm_version", + return_value=sm_version, + ): + supported = dflash_trtllm_supports_page_size( + tokens_per_block, has_context_attention=has_context_attention + ) assert supported is expected @pytest.mark.parametrize(