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(