[Attention] Enable adaptive verification for FLASHINFER_MLA_SPARSE_DSV4 (#52724)
Signed-off-by: Markov Ilya <markovilya197@gmail.com>
Ilya Markov committed
Sep 1, 2026 at 15:43 UTC
25efcfa7887c4a9541b6328af69dbd5fee4e8173
1 file changed
+30
-1
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+30
-1
@@ -21,13 +21,18 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
21
from vllm.models.deepseek_v4.sparse_mla import (
22
DeepseekV4FlashMLAMetadata,
23
DeepseekV4SparseMLABackend,
24
+ DeepseekV4SparseMLAMetadataBuilder,
25
)
26
from vllm.platforms.interface import DeviceCapability
27
from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
27
-from vllm.v1.attention.backend import MultipleOf
28
+from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
29
from vllm.v1.attention.backends.mla.compressor_utils import (
30
get_dspark_swa_index_width,
31
)
32
+from vllm.v1.attention.backends.mla.sparse_swa import (
33
+ DeepseekSparseSWABackend,
34
+ DeepseekSparseSWAMetadataBuilder,
35
+)
36
37
if TYPE_CHECKING:
38
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -164,11 +169,34 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
169
return None
170
return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
171
172
+ @staticmethod
173
+ def get_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]:
174
+ return DeepseekV4FlashInferSparseMLAMetadataBuilder
175
+
176
+
177
+class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
178
+ """Varlen-capable metadata builder for the FlashInfer sparse MLA backend."""
179
+
180
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
181
+
182
+
183
+class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
184
+ """SWA metadata for the FlashInfer sparse decode path (varlen decode)."""
185
+
186
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
187
+
188
+
189
+class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
190
+ @staticmethod
191
+ def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
192
+ return DeepseekSparseSWAFlashInferMetadataBuilder
193
+
194
195
class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
196
"""FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
197
198
backend_cls = DeepseekV4FlashInferMLASparseBackend
199
+ swa_backend_cls = DeepseekSparseSWAFlashInferBackend
200
use_fp8_ds_mla_layout: ClassVar[bool] = False
201
202
@classmethod
@@ -530,6 +558,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
558
"""DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""
559
560
backend_cls = DeepseekV4FlashInferMLASparseBackend
561
+ swa_backend_cls = DeepseekSparseSWAFlashInferBackend
562
use_fp8_ds_mla_layout: ClassVar[bool] = True
563
564
@staticmethod