From 25efcfa7887c4a9541b6328af69dbd5fee4e8173 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Tue, 1 Sep 2026 15:43:28 +0200 Subject: [PATCH] [Attention] Enable adaptive verification for FLASHINFER_MLA_SPARSE_DSV4 (#52724) Signed-off-by: Markov Ilya --- .../deepseek_v4/nvidia/flashinfer_sparse.py | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index bca71abbad..4ac362c95b 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -21,13 +21,18 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( from vllm.models.deepseek_v4.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, + DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 -from vllm.v1.attention.backend import MultipleOf +from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) +from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWABackend, + DeepseekSparseSWAMetadataBuilder, +) if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -164,11 +169,34 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): return None return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x" + @staticmethod + def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]: + return DeepseekV4FlashInferSparseMLAMetadataBuilder + + +class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): + """Varlen-capable metadata builder for the FlashInfer sparse MLA backend.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + +class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + """SWA metadata for the FlashInfer sparse decode path (varlen decode).""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + +class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): + @staticmethod + def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]: + return DeepseekSparseSWAFlashInferMetadataBuilder + class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4.""" backend_cls = DeepseekV4FlashInferMLASparseBackend + swa_backend_cls = DeepseekSparseSWAFlashInferBackend use_fp8_ds_mla_layout: ClassVar[bool] = False @classmethod @@ -530,6 +558,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention): """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels.""" backend_cls = DeepseekV4FlashInferMLASparseBackend + swa_backend_cls = DeepseekSparseSWAFlashInferBackend use_fp8_ds_mla_layout: ClassVar[bool] = True @staticmethod