diff --git a/tests/v1/attention/test_dsv4_kernel_block_size.py b/tests/v1/attention/test_dsv4_kernel_block_size.py new file mode 100644 index 000000000000..1cc0f0239a8d --- /dev/null +++ b/tests/v1/attention/test_dsv4_kernel_block_size.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( + DeepseekV4FlashInferMLASparseBackend, +) +from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4SparseMLABackend, + dsv4_supported_kernel_block_sizes, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.indexer import DeepseekV4IndexerBackend + + +def test_dsv4_kernel_block_size_sm12x(monkeypatch): + monkeypatch.setattr( + current_platform, "is_device_capability_family", lambda family: family == 120 + ) + assert dsv4_supported_kernel_block_sizes() == [64] + assert DeepseekV4SparseMLABackend.get_supported_kernel_block_sizes() == [64] + assert DeepseekV4FlashInferMLASparseBackend.get_supported_kernel_block_sizes() == [ + 64 + ] + assert DeepseekV4IndexerBackend.get_supported_kernel_block_sizes() == [64] + + +def test_dsv4_kernel_block_size_not_sm12x(monkeypatch): + monkeypatch.setattr( + current_platform, "is_device_capability_family", lambda family: False + ) + assert dsv4_supported_kernel_block_sizes() == [256] + assert DeepseekV4SparseMLABackend.get_supported_kernel_block_sizes() == [256] + assert DeepseekV4FlashInferMLASparseBackend.get_supported_kernel_block_sizes() == [ + 256 + ] + assert DeepseekV4IndexerBackend.get_supported_kernel_block_sizes() == [256] diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 883ef5ecafd4..81e5ea8e3303 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -25,7 +25,7 @@ ) from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 -from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf +from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) @@ -110,10 +110,6 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): "fp8_ds_mla", ] - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [256] - @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE_DSV4" diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 4953bdc9256f..5f44c1465ae4 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -39,6 +39,19 @@ _C128A_TOPK_ALIGNMENT = 128 +def dsv4_supported_kernel_block_sizes() -> list[int | MultipleOf]: + """Kernel page size for DeepSeek-V4 sparse MLA. + + FlashInfer SM120 DSV4 decode is compiled for 64-token pages + (``_DECODE_DSV4_PAGE_BLOCK_SIZE = 64``). Keep manager ``--block-size 256`` + so C128 storage is 2. Returning ``[64]`` lets ``select_common_block_size`` + split each manager block into four kernel pages. SM100 stays at 256. + """ + if current_platform.is_device_capability_family(120): + return [64] + return [256] + + class DeepseekV4SparseMLABackend(AttentionBackend): """DeepSeek-V4 sparse-MLA backend base. @@ -58,7 +71,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [256] + return dsv4_supported_kernel_block_sizes() @staticmethod def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]: diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2321cf8d13d4..5729a9a2fbf5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -251,7 +251,15 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows. - return [256] + # Imported lazily: vllm.models.deepseek_v4 transitively imports + # vllm._aiter_ops (via fused_moe), which would deadlock cold start + # with an import cycle when vllm._aiter_ops pulls in this module + # first (see vllm/v1/attention/ops/rocm_aiter_mla_sparse.py). + from vllm.models.deepseek_v4.sparse_mla import ( + dsv4_supported_kernel_block_sizes, + ) + + return dsv4_supported_kernel_block_sizes() class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):