A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 00:30 UTC
[Refactor] Extract DeepSeek V4 sparse MLA impl into model folder (#43149)
Yongye Zhu committed
May 22, 2026 at 13:06 UTC
843715739b7b555c61dd6190cafb5ab7a44c41f1
9 files changed
+485
-402
docs/design/attention_backends.md
+1
-1
@@ -219,7 +219,7 @@ MLA decode backends are selected using the standard
219
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
220
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
221
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
222
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
222
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
223
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
224
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
225
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+2
-2
@@ -198,7 +198,7 @@ def _ref_combine_topk_swa_ragged(
198
199
@torch.inference_mode()
200
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
201
- from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
201
+ from vllm.models.deepseek_v4.amd.rocm import (
202
compute_global_topk_ragged_indices_and_indptr,
203
)
204
@@ -330,7 +330,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
330
331
@torch.inference_mode()
332
def test_combine_topk_swa_indices_ragged() -> None:
333
- from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
333
+ from vllm.models.deepseek_v4.amd.rocm import (
334
combine_topk_swa_indices_ragged,
335
)
336
tests/kernels/test_fused_inv_rope_fp8_quant.py
+3
-2
@@ -7,7 +7,8 @@ Tests compare the fused kernel against a reference implementation built from
7
the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
8
via per_token_group_quant_fp8).
9
10
-The reference faithfully reproduces the exact flow in deepseek_v4_attention.py:295-310:
10
+The reference faithfully reproduces the exact flow in
11
+deepseek_v4/nvidia/ops/attention.py:295-310:
12
1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
13
2. Reshape [T, H, head_dim] -> [T, G, D]
14
3. Transpose+flatten to [G*T, D], quantize, reshape back
@@ -667,7 +668,7 @@ def _unfused_inv_rope_fp8_quant(
668
nope_dim: int = NOPE_DIM,
669
rope_dim: int = ROPE_DIM,
670
) -> tuple[torch.Tensor, torch.Tensor]:
670
- """Unfused path matching deepseek_v4_attention.py:295-310.
671
+ """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
672
673
Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
674
"""
vllm/models/deepseek_v4/amd/rocm.py
renamed
+47
-59
@@ -8,14 +8,15 @@ import torch
8
9
from vllm.forward_context import get_forward_context
10
from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
11
+from vllm.models.deepseek_v4.nvidia.flashmla import (
12
+ DeepseekV4FlashMLASparseBackend,
13
+ DeepseekV4SparseMLAAttentionImpl,
14
+)
15
from vllm.triton_utils import tl, triton
16
from vllm.v1.attention.backend import (
13
- AttentionLayer,
17
CommonAttentionMetadata,
15
- SparseMLAAttentionImpl,
18
)
19
from vllm.v1.attention.backends.mla.flashmla_sparse import (
18
- DeepseekV4FlashMLASparseBackend,
20
FlashMLASparseMetadata,
21
FlashMLASparseMetadataBuilder,
22
)
@@ -31,7 +32,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
32
from vllm.v1.worker.workspace import current_workspace_manager
33
34
if TYPE_CHECKING:
34
- from vllm.models.deepseek_v4.attention import DeepseekV4MLAAttention
35
+ from vllm.models.deepseek_v4.nvidia.ops.attention import (
36
+ DeepseekV4MLAAttention,
37
+ )
38
39
40
def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
@@ -449,47 +452,27 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
452
)
453
454
452
-class DeepseekV4ROCMAiterMLASparseImpl(
453
- SparseMLAAttentionImpl[DeepseekV4ROCMAiterMLASparseMetadata]
454
-):
455
- """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer."""
455
+class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):
456
+ @staticmethod
457
+ def get_name() -> str:
458
+ return "ROCM_V4_FLASHMLA_SPARSE"
459
457
- _PREFILL_CHUNK_SIZE = 4
460
+ @staticmethod
461
+ def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]:
462
+ return DeepseekV4ROCMAiterMLASparseMetadataBuilder
463
459
- def __init__(
460
- self,
461
- num_heads: int,
462
- head_size: int,
463
- scale: float,
464
- num_kv_heads: int,
465
- alibi_slopes: list[float] | None,
466
- sliding_window: int | None,
467
- kv_cache_dtype: str,
468
- logits_soft_cap: float | None,
469
- attn_type: str,
470
- kv_sharing_target_layer_name: str | None,
471
- **_: object,
472
- ) -> None:
473
- self.num_heads = num_heads
474
- self.head_size = head_size
475
- self.scale = float(scale)
476
- self.num_kv_heads = num_kv_heads
477
- self.kv_cache_dtype = kv_cache_dtype
464
+ @staticmethod
465
+ def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]:
466
+ return DeepseekV4ROCMAiterMLASparseImpl
467
479
- def forward_mqa(
480
- self,
481
- q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
482
- kv_c_and_k_pe_cache: torch.Tensor,
483
- attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata,
484
- layer: AttentionLayer,
485
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
486
- raise NotImplementedError(
487
- "DeepseekV4ROCMAiterMLASparseImpl is driven by "
488
- "DeepseekV4MLAAttention.forward."
489
- )
468
+
469
+class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
470
+ """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer."""
471
+
472
+ backend_cls = DeepseekV4ROCMAiterMLASparseBackend
473
474
@classmethod
492
- def forward(
475
+ def forward_mqa( # type: ignore[override]
476
cls,
477
layer: "DeepseekV4MLAAttention",
478
q: torch.Tensor,
@@ -506,6 +489,25 @@ class DeepseekV4ROCMAiterMLASparseImpl(
489
490
forward_context = get_forward_context()
491
attn_metadata = forward_context.attn_metadata
492
+
493
+ if attn_metadata is None:
494
+ # Warmup dummy run: no real metadata. Reserve the same bf16
495
+ # gather workspace _forward_prefill would; the dequantize / topk
496
+ # / sparse_fwd kernels are skipped this step.
497
+ swa_only = layer.compress_ratio <= 1
498
+ N = (
499
+ 0
500
+ if swa_only
501
+ else (layer.max_model_len + layer.compress_ratio - 1)
502
+ // layer.compress_ratio
503
+ )
504
+ M = N + layer.window_size + layer.max_num_batched_tokens
505
+ current_workspace_manager().get_simultaneous(
506
+ ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
507
+ )
508
+ output.zero_()
509
+ return
510
+
511
assert isinstance(attn_metadata, dict)
512
rocm_metadata = cast(
513
DeepseekV4ROCMAiterMLASparseMetadata | None,
@@ -658,17 +660,17 @@ class DeepseekV4ROCMAiterMLASparseImpl(
660
N = 0
661
662
M = N + layer.window_size + layer.max_num_batched_tokens
661
- num_chunks = (num_prefills + cls._PREFILL_CHUNK_SIZE - 1) // (
662
- cls._PREFILL_CHUNK_SIZE
663
+ num_chunks = (num_prefills + cls.PREFILL_CHUNK_SIZE - 1) // (
664
+ cls.PREFILL_CHUNK_SIZE
665
)
666
667
workspace_manager = current_workspace_manager()
668
kv = workspace_manager.get_simultaneous(
667
- ((cls._PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
669
+ ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
670
)[0]
671
for chunk_idx in range(num_chunks):
670
- chunk_start = chunk_idx * cls._PREFILL_CHUNK_SIZE
671
- chunk_end = min(chunk_start + cls._PREFILL_CHUNK_SIZE, num_prefills)
672
+ chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE
673
+ chunk_end = min(chunk_start + cls.PREFILL_CHUNK_SIZE, num_prefills)
674
chunk_size = chunk_end - chunk_start
675
if not swa_only:
676
assert attn_metadata is not None
@@ -737,17 +739,3 @@ class DeepseekV4ROCMAiterMLASparseImpl(
739
ragged_indices=combined_ragged_indices,
740
ragged_indptr=combined_ragged_indptr,
741
)
740
-
741
-
742
-class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):
743
- @staticmethod
744
- def get_name() -> str:
745
- return "ROCM_V4_FLASHMLA_SPARSE"
746
-
747
- @staticmethod
748
- def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]:
749
- return DeepseekV4ROCMAiterMLASparseMetadataBuilder
750
-
751
- @staticmethod
752
- def get_impl_cls() -> type["DeepseekV4ROCMAiterMLASparseImpl"]:
753
- return DeepseekV4ROCMAiterMLASparseImpl
vllm/models/deepseek_v4/nvidia/flashmla.py
new
+402
@@ -0,0 +1,402 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import abstractmethod
5
+from typing import TYPE_CHECKING, ClassVar, cast
6
+
7
+import torch
8
+
9
+from vllm.forward_context import get_forward_context
10
+from vllm.models.deepseek_v4.common.ops import (
11
+ combine_topk_swa_indices,
12
+ compute_global_topk_indices_and_lens,
13
+ dequantize_and_gather_k_cache,
14
+)
15
+from vllm.v1.attention.backend import (
16
+ AttentionBackend,
17
+ MultipleOf,
18
+ SparseMLAAttentionImpl,
19
+)
20
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
21
+ FlashMLASparseBackend,
22
+ FlashMLASparseMetadata,
23
+)
24
+from vllm.v1.attention.ops.flashmla import (
25
+ flash_mla_sparse_fwd,
26
+ flash_mla_with_kvcache,
27
+)
28
+from vllm.v1.worker.workspace import current_workspace_manager
29
+
30
+if TYPE_CHECKING:
31
+ from vllm.models.deepseek_v4.nvidia.ops.attention import (
32
+ DeepseekV4MLAAttention,
33
+ )
34
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
35
+
36
+
37
+class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
38
+ """Abstract parent for DeepseekV4 sparse MLA impls.
39
+
40
+ V4 sparse MLA is driven by the layer (``DeepseekV4MLAAttention.forward``)
41
+ rather than the v1 framework, so ``forward_mqa`` is overridden with a
42
+ classmethod that takes the layer as its first argument. This Liskov-broken
43
+ override is intentional: the grandparent's instance-method ``forward_mqa``
44
+ is never called on V4 layers.
45
+ """
46
+
47
+ backend_cls: ClassVar[type[AttentionBackend]]
48
+
49
+ # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
50
+ # workspace allocated in _forward_prefill and is also read by the V4 layer's
51
+ # dummy-run path to pre-reserve that workspace.
52
+ PREFILL_CHUNK_SIZE: ClassVar[int] = 4
53
+
54
+ @classmethod
55
+ @abstractmethod
56
+ def forward_mqa( # type: ignore[override]
57
+ cls,
58
+ layer: "DeepseekV4MLAAttention",
59
+ q: torch.Tensor,
60
+ kv: torch.Tensor,
61
+ positions: torch.Tensor,
62
+ output: torch.Tensor,
63
+ ) -> None:
64
+ raise NotImplementedError
65
+
66
+
67
+class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
68
+ @staticmethod
69
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
70
+ return [256]
71
+
72
+ @staticmethod
73
+ def get_name() -> str:
74
+ return "V4_FLASHMLA_SPARSE"
75
+
76
+ @staticmethod
77
+ def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]:
78
+ return DeepseekV4FlashMLASparseImpl
79
+
80
+ @classmethod
81
+ def get_supported_head_sizes(cls) -> list[int]:
82
+ # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512 (overrides the
83
+ # V3.2 default of 576 from FlashMLASparseBackend).
84
+ return [512]
85
+
86
+ @staticmethod
87
+ def get_kv_cache_shape(
88
+ num_blocks: int,
89
+ block_size: int,
90
+ num_kv_heads: int,
91
+ head_size: int,
92
+ cache_dtype_str: str = "auto",
93
+ ) -> tuple[int, ...]:
94
+ if cache_dtype_str == "fp8_ds_mla":
95
+ # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
96
+ # head_size passed in is the semantic head_dim (512).
97
+ return (num_blocks, block_size, 584)
98
+ else:
99
+ return (num_blocks, block_size, head_size)
100
+
101
+
102
+class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
103
+ """FlashMLA sparse MLA implementation for DeepSeek V4's custom MLA layer."""
104
+
105
+ backend_cls = DeepseekV4FlashMLASparseBackend
106
+
107
+ @classmethod
108
+ def forward_mqa( # type: ignore[override]
109
+ cls,
110
+ layer: "DeepseekV4MLAAttention",
111
+ q: torch.Tensor,
112
+ kv: torch.Tensor,
113
+ positions: torch.Tensor,
114
+ output: torch.Tensor,
115
+ ) -> None:
116
+ assert output.shape == q.shape, (
117
+ f"output buffer shape {output.shape} must match q shape {q.shape}"
118
+ )
119
+ assert output.dtype == q.dtype, (
120
+ f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
121
+ )
122
+
123
+ # Get SWA and indexer metadata from forward context
124
+ forward_context = get_forward_context()
125
+ attn_metadata = forward_context.attn_metadata
126
+
127
+ if attn_metadata is None:
128
+ # Warmup dummy run: no real metadata. Reserve the same bf16
129
+ # gather workspace _forward_prefill would; the dequantize / topk
130
+ # / sparse_fwd kernels are skipped this step.
131
+ swa_only = layer.compress_ratio <= 1
132
+ N = (
133
+ 0
134
+ if swa_only
135
+ else (layer.max_model_len + layer.compress_ratio - 1)
136
+ // layer.compress_ratio
137
+ )
138
+ M = N + layer.window_size + layer.max_num_batched_tokens
139
+ current_workspace_manager().get_simultaneous(
140
+ ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
141
+ )
142
+ output.zero_()
143
+ return
144
+
145
+ assert isinstance(attn_metadata, dict)
146
+ flashmla_metadata = cast(
147
+ FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix)
148
+ )
149
+ swa_metadata = cast(
150
+ "DeepseekSparseSWAMetadata | None",
151
+ attn_metadata.get(layer.swa_cache_layer.prefix),
152
+ )
153
+ assert swa_metadata is not None
154
+
155
+ swa_only = layer.compress_ratio <= 1
156
+ # SWA-only layers (compress_ratio <= 1) don't have their own KV cache
157
+ # allocation, so layer.kv_cache may be empty after profiling cleanup.
158
+ self_kv_cache = layer.kv_cache if not swa_only else None
159
+ swa_kv_cache = layer.swa_cache_layer.kv_cache
160
+
161
+ # Split prefill and decode
162
+ num_decodes = swa_metadata.num_decodes
163
+ num_prefills = swa_metadata.num_prefills
164
+ num_decode_tokens = swa_metadata.num_decode_tokens
165
+
166
+ if num_prefills > 0:
167
+ cls._forward_prefill(
168
+ layer=layer,
169
+ q=q[num_decode_tokens:],
170
+ positions=positions[num_decode_tokens:],
171
+ compressed_k_cache=self_kv_cache,
172
+ swa_k_cache=swa_kv_cache,
173
+ output=output[num_decode_tokens:],
174
+ attn_metadata=flashmla_metadata,
175
+ swa_metadata=swa_metadata,
176
+ )
177
+ if num_decodes > 0:
178
+ cls._forward_decode(
179
+ layer=layer,
180
+ q=q[:num_decode_tokens],
181
+ kv_cache=self_kv_cache,
182
+ swa_metadata=swa_metadata,
183
+ attn_metadata=flashmla_metadata,
184
+ swa_only=swa_only,
185
+ output=output[:num_decode_tokens],
186
+ )
187
+
188
+ @classmethod
189
+ def _forward_decode(
190
+ cls,
191
+ layer: "DeepseekV4MLAAttention",
192
+ q: torch.Tensor,
193
+ kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
194
+ swa_metadata: "DeepseekSparseSWAMetadata",
195
+ attn_metadata: FlashMLASparseMetadata | None,
196
+ swa_only: bool,
197
+ output: torch.Tensor,
198
+ ) -> None:
199
+ num_decodes = swa_metadata.num_decodes
200
+ num_decode_tokens = swa_metadata.num_decode_tokens
201
+
202
+ topk_indices = None
203
+ topk_lens = None
204
+ if not swa_only:
205
+ assert attn_metadata is not None
206
+ assert swa_metadata.is_valid_token is not None
207
+ block_size = attn_metadata.block_size // layer.compress_ratio
208
+ is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
209
+ if layer.compress_ratio == 4:
210
+ # C4A: local indices differ per layer (filled by Indexer).
211
+ assert layer.topk_indices_buffer is not None
212
+ global_indices, topk_lens = compute_global_topk_indices_and_lens(
213
+ layer.topk_indices_buffer[:num_decode_tokens],
214
+ swa_metadata.token_to_req_indices,
215
+ attn_metadata.block_table[:num_decodes],
216
+ block_size,
217
+ is_valid,
218
+ )
219
+ topk_indices = global_indices.view(num_decode_tokens, 1, -1)
220
+ else:
221
+ # C128A: pre-computed during metadata build.
222
+ topk_indices = attn_metadata.c128a_global_decode_topk_indices
223
+ topk_lens = attn_metadata.c128a_decode_topk_lens
224
+
225
+ swa_indices = swa_metadata.decode_swa_indices
226
+ swa_lens = swa_metadata.decode_swa_lens
227
+
228
+ # We treat queries in the same seq as different queries
229
+ # and later we only attend by generated indices.
230
+ # q arrives pre-padded to layer.padded_heads by the outer wrapper.
231
+ q = q.unsqueeze(1)
232
+
233
+ # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
234
+ # Use unsqueeze to preserve strides (handles padded blocks correctly)
235
+ swa_cache = layer.swa_cache_layer.kv_cache.unsqueeze(-2)
236
+ # Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
237
+ if kv_cache is not None:
238
+ kv_cache = kv_cache.unsqueeze(-2)
239
+
240
+ # One FlashMLASchedMeta per layer type, shared across all same-type
241
+ # layers within this decode step. The first forward call per type
242
+ # triggers the in-kernel planner (allocating tile_scheduler_metadata
243
+ # and num_splits via PyTorch's graph-aware allocator so CUDA graph
244
+ # capture reuses the same addresses on replay); subsequent same-type
245
+ # layers see have_initialized=True and skip the planner.
246
+ if layer.compress_ratio <= 1:
247
+ tile_metadata = swa_metadata.tile_sched_swaonly
248
+ elif layer.compress_ratio == 4:
249
+ tile_metadata = swa_metadata.tile_sched_c4a
250
+ elif layer.compress_ratio == 128:
251
+ tile_metadata = swa_metadata.tile_sched_c128a
252
+ else:
253
+ raise ValueError(
254
+ f"Unsupported compress_ratio={layer.compress_ratio}; "
255
+ "expected 1, 4, or 128."
256
+ )
257
+ assert tile_metadata is not None, (
258
+ "swa_metadata missing tile_sched entry for "
259
+ f"compress_ratio={layer.compress_ratio}; "
260
+ "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not "
261
+ "allocate one for this layer type."
262
+ )
263
+
264
+ out, _ = flash_mla_with_kvcache(
265
+ q=q,
266
+ k_cache=swa_cache,
267
+ block_table=None,
268
+ head_dim_v=512,
269
+ tile_scheduler_metadata=tile_metadata,
270
+ cache_seqlens=None,
271
+ is_fp8_kvcache=True,
272
+ indices=swa_indices,
273
+ topk_length=swa_lens,
274
+ softmax_scale=layer.scale,
275
+ attn_sink=layer.attn_sink,
276
+ extra_k_cache=kv_cache if not swa_only else None,
277
+ extra_indices_in_kvcache=topk_indices,
278
+ extra_topk_length=topk_lens,
279
+ out=output.unsqueeze(1),
280
+ )
281
+
282
+ @classmethod
283
+ def _forward_prefill(
284
+ cls,
285
+ layer: "DeepseekV4MLAAttention",
286
+ q: torch.Tensor,
287
+ positions: torch.Tensor,
288
+ compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1
289
+ swa_k_cache: torch.Tensor,
290
+ output: torch.Tensor,
291
+ attn_metadata: FlashMLASparseMetadata | None,
292
+ swa_metadata: "DeepseekSparseSWAMetadata",
293
+ ) -> None:
294
+ swa_only = attn_metadata is None
295
+
296
+ num_prefills = swa_metadata.num_prefills
297
+ num_prefill_tokens = swa_metadata.num_prefill_tokens
298
+ num_decodes = swa_metadata.num_decodes
299
+ num_decode_tokens = swa_metadata.num_decode_tokens
300
+
301
+ # Use pre-computed prefill metadata.
302
+ seq_lens = swa_metadata.prefill_seq_lens
303
+ gather_lens = swa_metadata.prefill_gather_lens
304
+ assert seq_lens is not None
305
+ assert gather_lens is not None
306
+
307
+ # Derive prefill-local token offsets from the full query_start_loc_cpu.
308
+ query_start_loc_cpu = swa_metadata.query_start_loc_cpu
309
+ query_start_loc = swa_metadata.query_start_loc
310
+ assert query_start_loc_cpu is not None
311
+ assert query_start_loc is not None
312
+ prefill_token_base = query_start_loc_cpu[num_decodes]
313
+
314
+ if not swa_only:
315
+ if layer.compress_ratio == 4:
316
+ assert layer.topk_indices_buffer is not None
317
+ topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
318
+ topk_indices = topk_indices[:num_prefill_tokens]
319
+ else:
320
+ # C128A: pre-computed during metadata build.
321
+ assert attn_metadata is not None
322
+ topk_indices = attn_metadata.c128a_prefill_topk_indices
323
+ top_k = topk_indices.shape[-1]
324
+ # Compressed region must fit the full compressed pool (seq_len //
325
+ # compress_ratio), not just top_k. top_k bounds how many indices
326
+ # the indexer selects, not the pool size it indexes into.
327
+ N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio
328
+ else:
329
+ # NOTE(woosuk): topk_indices will not be used for SWA-only layers.
330
+ assert layer.topk_indices_buffer is not None
331
+ topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
332
+ top_k = 0
333
+ N = 0
334
+
335
+ M = N + layer.window_size + layer.max_num_batched_tokens
336
+ chunk_size_const = cls.PREFILL_CHUNK_SIZE
337
+ num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const
338
+
339
+ workspace_manager = current_workspace_manager()
340
+ kv = workspace_manager.get_simultaneous(
341
+ ((chunk_size_const, M, q.shape[-1]), torch.bfloat16),
342
+ )[0]
343
+ for chunk_idx in range(num_chunks):
344
+ chunk_start = chunk_idx * chunk_size_const
345
+ chunk_end = min(chunk_start + chunk_size_const, num_prefills)
346
+ chunk_size = chunk_end - chunk_start
347
+ if not swa_only:
348
+ # Gather compressed KV
349
+ assert attn_metadata is not None
350
+ block_table = attn_metadata.block_table[num_decodes:]
351
+ dequantize_and_gather_k_cache(
352
+ kv[:chunk_size],
353
+ compressed_k_cache,
354
+ seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio,
355
+ gather_lens=None,
356
+ block_table=block_table[chunk_start:chunk_end],
357
+ block_size=attn_metadata.block_size // layer.compress_ratio,
358
+ offset=0,
359
+ )
360
+
361
+ # Gather SWA KV
362
+ swa_block_table = swa_metadata.block_table[num_decodes:]
363
+ dequantize_and_gather_k_cache(
364
+ kv[:chunk_size],
365
+ swa_k_cache,
366
+ seq_lens=seq_lens[chunk_start:chunk_end],
367
+ gather_lens=gather_lens[chunk_start:chunk_end],
368
+ block_table=swa_block_table[chunk_start:chunk_end],
369
+ block_size=swa_metadata.block_size,
370
+ offset=N,
371
+ )
372
+
373
+ # Combine the topk indices and SWA indices for gathered KV cache
374
+ query_start = (
375
+ query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
376
+ )
377
+ query_end = (
378
+ query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
379
+ )
380
+
381
+ combined_indices, combined_lens = combine_topk_swa_indices(
382
+ topk_indices[query_start:query_end],
383
+ query_start_loc[
384
+ num_decodes + chunk_start : num_decodes + chunk_end + 1
385
+ ],
386
+ seq_lens[chunk_start:chunk_end],
387
+ gather_lens[chunk_start:chunk_end],
388
+ layer.window_size,
389
+ layer.compress_ratio,
390
+ top_k,
391
+ M,
392
+ N,
393
+ )
394
+ flash_mla_sparse_fwd(
395
+ q=q[query_start:query_end],
396
+ kv=kv.view(-1, 1, q.shape[-1]),
397
+ indices=combined_indices.unsqueeze(1),
398
+ sm_scale=layer.scale,
399
+ attn_sink=layer.attn_sink,
400
+ topk_length=combined_lens,
401
+ out=output[query_start:query_end],
402
+ )
vllm/models/deepseek_v4/nvidia/model.py
+1
-1
@@ -56,7 +56,7 @@ from vllm.model_executor.models.utils import (
56
maybe_prefix,
57
)
58
from vllm.model_executor.utils import set_weight_attrs
59
-from vllm.models.deepseek_v4.attention import (
59
+from vllm.models.deepseek_v4.nvidia.ops.attention import (
60
DeepseekV4Indexer,
61
DeepseekV4MLAModules,
62
DeepseekV4MultiHeadLatentAttentionWrapper,
vllm/models/deepseek_v4/nvidia/ops/attention.py
renamed
+23
-309
@@ -20,9 +20,6 @@ from vllm.model_executor.layers.linear import (
20
)
21
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
22
from vllm.models.deepseek_v4.common.ops import (
23
- combine_topk_swa_indices,
24
- compute_global_topk_indices_and_lens,
25
- dequantize_and_gather_k_cache,
23
fused_indexer_q_rope_quant,
24
fused_inv_rope_fp8_quant,
25
fused_q_kv_rmsnorm,
@@ -62,28 +59,36 @@ from vllm.utils.multi_stream_utils import (
59
)
60
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
61
from vllm.v1.attention.backends.mla.flashmla_sparse import (
65
- DeepseekV4FlashMLASparseBackend,
62
FlashMLASparseBackend,
67
- FlashMLASparseMetadata,
63
)
64
from vllm.v1.attention.backends.mla.indexer import (
65
DeepseekV4IndexerBackend,
66
get_max_prefill_buffer_size,
67
)
68
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
74
-from vllm.v1.attention.ops.flashmla import (
75
- flash_mla_sparse_fwd,
76
- flash_mla_with_kvcache,
77
-)
69
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
79
-from vllm.v1.worker.workspace import current_workspace_manager
70
+
71
+if TYPE_CHECKING:
72
+ from vllm.models.deepseek_v4.nvidia.flashmla import (
73
+ DeepseekV4SparseMLAAttentionImpl,
74
+ )
75
76
logger = init_logger(__name__)
77
83
-# Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
84
-# workspace allocated at _forward_prefill (and the matching profile-time
85
-# reservation in attention_impl's dummy-run branch).
86
-PREFILL_CHUNK_SIZE = 4
78
+
79
+def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]":
80
+ """Pick the platform-specific V4 sparse MLA impl class. Sole platform check."""
81
+ if current_platform.is_rocm():
82
+ from vllm.models.deepseek_v4.amd.rocm import (
83
+ DeepseekV4ROCMAiterMLASparseImpl,
84
+ )
85
+
86
+ return DeepseekV4ROCMAiterMLASparseImpl
87
+ from vllm.models.deepseek_v4.nvidia.flashmla import (
88
+ DeepseekV4FlashMLASparseImpl,
89
+ )
90
+
91
+ return DeepseekV4FlashMLASparseImpl
92
93
94
@dataclass
@@ -489,25 +494,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
494
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
495
self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
496
492
- # Handle dummy run (no metadata).
493
- if not isinstance(attn_metadata, dict):
494
- # Reserve _forward_prefill's bf16-gather workspace; the dummy
495
- # run returns before mla_attn runs, so without this the shared
496
- # workspace locks below the real prefill size.
497
- sub = self.mla_attn
498
- swa_only = sub.compress_ratio <= 1
499
- N = (
500
- 0
501
- if swa_only
502
- else (sub.max_model_len + sub.compress_ratio - 1) // sub.compress_ratio
503
- )
504
- M = N + sub.window_size + sub.max_num_batched_tokens
505
- current_workspace_manager().get_simultaneous(
506
- ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
507
- )
508
- out.zero_()
509
- return
510
-
497
# Pad q to FlashMLA-required head count (64 or 128)
498
if self.n_local_heads < self.padded_heads:
499
pad_size = self.padded_heads - self.n_local_heads
@@ -643,6 +629,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
629
**extra_impl_args,
630
) -> None:
631
super().__init__()
632
+ self.impl_cls = _select_v4_sparse_impl()
633
+ self.backend_cls = self.impl_cls.backend_cls
634
self.num_heads = num_heads
635
self.num_kv_heads = 1
636
self.head_dim = head_dim
@@ -723,13 +711,7 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
711
self.kv_cache = torch.tensor([])
712
713
def get_attn_backend(self) -> type[AttentionBackend]:
726
- if current_platform.is_rocm():
727
- from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
728
- DeepseekV4ROCMAiterMLASparseBackend,
729
- )
730
-
731
- return DeepseekV4ROCMAiterMLASparseBackend
732
- return DeepseekV4FlashMLASparseBackend
714
+ return self.backend_cls
715
716
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
717
if (
@@ -754,275 +736,7 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
736
positions: torch.Tensor,
737
output: torch.Tensor,
738
) -> None:
757
- assert output.shape == q.shape, (
758
- f"output buffer shape {output.shape} must match q shape {q.shape}"
759
- )
760
- assert output.dtype == q.dtype, (
761
- f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
762
- )
763
-
764
- if current_platform.is_rocm():
765
- from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
766
- DeepseekV4ROCMAiterMLASparseImpl,
767
- )
768
-
769
- DeepseekV4ROCMAiterMLASparseImpl.forward(self, q, kv, positions, output)
770
- return
771
-
772
- # Get SWA and indexer metadata from forward context
773
- forward_context = get_forward_context()
774
- attn_metadata = forward_context.attn_metadata
775
- assert isinstance(attn_metadata, dict)
776
- flashmla_metadata = cast(
777
- FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)
778
- )
779
- swa_metadata = cast(
780
- "DeepseekSparseSWAMetadata | None",
781
- attn_metadata.get(self.swa_cache_layer.prefix),
782
- )
783
- assert swa_metadata is not None
784
-
785
- swa_only = self.compress_ratio <= 1
786
- # SWA-only layers (compress_ratio <= 1) don't have their own KV cache
787
- # allocation, so self.kv_cache may be empty after profiling cleanup.
788
- self_kv_cache = self.kv_cache if not swa_only else None
789
- swa_kv_cache = self.swa_cache_layer.kv_cache
790
-
791
- # Split prefill and decode
792
- num_decodes = swa_metadata.num_decodes
793
- num_prefills = swa_metadata.num_prefills
794
- num_decode_tokens = swa_metadata.num_decode_tokens
795
-
796
- if num_prefills > 0:
797
- self._forward_prefill(
798
- q=q[num_decode_tokens:],
799
- positions=positions[num_decode_tokens:],
800
- compressed_k_cache=self_kv_cache,
801
- swa_k_cache=swa_kv_cache,
802
- output=output[num_decode_tokens:],
803
- attn_metadata=flashmla_metadata,
804
- swa_metadata=swa_metadata,
805
- )
806
- if num_decodes > 0:
807
- self._forward_decode(
808
- q=q[:num_decode_tokens],
809
- kv_cache=self_kv_cache,
810
- swa_metadata=swa_metadata,
811
- attn_metadata=flashmla_metadata,
812
- swa_only=swa_only,
813
- output=output[:num_decode_tokens],
814
- )
815
-
816
- def _forward_decode(
817
- self,
818
- q: torch.Tensor,
819
- kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
820
- swa_metadata: "DeepseekSparseSWAMetadata",
821
- attn_metadata: FlashMLASparseMetadata | None,
822
- swa_only: bool,
823
- output: torch.Tensor,
824
- ) -> None:
825
- num_decodes = swa_metadata.num_decodes
826
- num_decode_tokens = swa_metadata.num_decode_tokens
827
-
828
- topk_indices = None
829
- topk_lens = None
830
- if not swa_only:
831
- assert attn_metadata is not None
832
- assert swa_metadata.is_valid_token is not None
833
- block_size = attn_metadata.block_size // self.compress_ratio
834
- is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
835
- if self.compress_ratio == 4:
836
- # C4A: local indices differ per layer (filled by Indexer).
837
- assert self.topk_indices_buffer is not None
838
- global_indices, topk_lens = compute_global_topk_indices_and_lens(
839
- self.topk_indices_buffer[:num_decode_tokens],
840
- swa_metadata.token_to_req_indices,
841
- attn_metadata.block_table[:num_decodes],
842
- block_size,
843
- is_valid,
844
- )
845
- topk_indices = global_indices.view(num_decode_tokens, 1, -1)
846
- else:
847
- # C128A: pre-computed during metadata build.
848
- topk_indices = attn_metadata.c128a_global_decode_topk_indices
849
- topk_lens = attn_metadata.c128a_decode_topk_lens
850
-
851
- swa_indices = swa_metadata.decode_swa_indices
852
- swa_lens = swa_metadata.decode_swa_lens
853
-
854
- # We treat queries in the same seq as different queries
855
- # and later we only attend by generated indices.
856
- # q arrives pre-padded to self.padded_heads by the outer wrapper.
857
- q = q.unsqueeze(1)
858
-
859
- # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
860
- # Use unsqueeze to preserve strides (handles padded blocks correctly)
861
- swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2)
862
- # Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
863
- if kv_cache is not None:
864
- kv_cache = kv_cache.unsqueeze(-2)
865
-
866
- # One FlashMLASchedMeta per layer type, shared across all same-type
867
- # layers within this decode step. The first forward call per type
868
- # triggers the in-kernel planner (allocating tile_scheduler_metadata
869
- # and num_splits via PyTorch's graph-aware allocator so CUDA graph
870
- # capture reuses the same addresses on replay); subsequent same-type
871
- # layers see have_initialized=True and skip the planner.
872
- if self.compress_ratio <= 1:
873
- tile_metadata = swa_metadata.tile_sched_swaonly
874
- elif self.compress_ratio == 4:
875
- tile_metadata = swa_metadata.tile_sched_c4a
876
- elif self.compress_ratio == 128:
877
- tile_metadata = swa_metadata.tile_sched_c128a
878
- else:
879
- raise ValueError(
880
- f"Unsupported compress_ratio={self.compress_ratio}; "
881
- "expected 1, 4, or 128."
882
- )
883
- assert tile_metadata is not None, (
884
- "swa_metadata missing tile_sched entry for "
885
- f"compress_ratio={self.compress_ratio}; "
886
- "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not "
887
- "allocate one for this layer type."
888
- )
889
-
890
- out, _ = flash_mla_with_kvcache(
891
- q=q,
892
- k_cache=swa_cache,
893
- block_table=None,
894
- head_dim_v=512,
895
- tile_scheduler_metadata=tile_metadata,
896
- cache_seqlens=None,
897
- is_fp8_kvcache=True,
898
- indices=swa_indices,
899
- topk_length=swa_lens,
900
- softmax_scale=self.scale,
901
- attn_sink=self.attn_sink,
902
- extra_k_cache=kv_cache if not swa_only else None,
903
- extra_indices_in_kvcache=topk_indices,
904
- extra_topk_length=topk_lens,
905
- out=output.unsqueeze(1),
906
- )
907
-
908
- def _forward_prefill(
909
- self,
910
- q: torch.Tensor,
911
- positions: torch.Tensor,
912
- compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1
913
- swa_k_cache: torch.Tensor,
914
- output: torch.Tensor,
915
- attn_metadata: FlashMLASparseMetadata | None,
916
- swa_metadata: "DeepseekSparseSWAMetadata",
917
- ) -> None:
918
- swa_only = attn_metadata is None
919
-
920
- num_prefills = swa_metadata.num_prefills
921
- num_prefill_tokens = swa_metadata.num_prefill_tokens
922
- num_decodes = swa_metadata.num_decodes
923
- num_decode_tokens = swa_metadata.num_decode_tokens
924
-
925
- # Use pre-computed prefill metadata.
926
- seq_lens = swa_metadata.prefill_seq_lens
927
- gather_lens = swa_metadata.prefill_gather_lens
928
- assert seq_lens is not None
929
- assert gather_lens is not None
930
-
931
- # Derive prefill-local token offsets from the full query_start_loc_cpu.
932
- query_start_loc_cpu = swa_metadata.query_start_loc_cpu
933
- query_start_loc = swa_metadata.query_start_loc
934
- assert query_start_loc_cpu is not None
935
- assert query_start_loc is not None
936
- prefill_token_base = query_start_loc_cpu[num_decodes]
937
-
938
- if not swa_only:
939
- if self.compress_ratio == 4:
940
- assert self.topk_indices_buffer is not None
941
- topk_indices = self.topk_indices_buffer[num_decode_tokens:]
942
- topk_indices = topk_indices[:num_prefill_tokens]
943
- else:
944
- # C128A: pre-computed during metadata build.
945
- assert attn_metadata is not None
946
- topk_indices = attn_metadata.c128a_prefill_topk_indices
947
- top_k = topk_indices.shape[-1]
948
- # Compressed region must fit the full compressed pool (seq_len //
949
- # compress_ratio), not just top_k. top_k bounds how many indices
950
- # the indexer selects, not the pool size it indexes into.
951
- N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
952
- else:
953
- # NOTE(woosuk): topk_indices will not be used for SWA-only layers.
954
- assert self.topk_indices_buffer is not None
955
- topk_indices = self.topk_indices_buffer[num_decode_tokens:]
956
- top_k = 0
957
- N = 0
958
-
959
- M = N + self.window_size + self.max_num_batched_tokens
960
- num_chunks = (num_prefills + PREFILL_CHUNK_SIZE - 1) // PREFILL_CHUNK_SIZE
961
-
962
- workspace_manager = current_workspace_manager()
963
- kv = workspace_manager.get_simultaneous(
964
- ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
965
- )[0]
966
- for chunk_idx in range(num_chunks):
967
- chunk_start = chunk_idx * PREFILL_CHUNK_SIZE
968
- chunk_end = min(chunk_start + PREFILL_CHUNK_SIZE, num_prefills)
969
- chunk_size = chunk_end - chunk_start
970
- if not swa_only:
971
- # Gather compressed KV
972
- assert attn_metadata is not None
973
- block_table = attn_metadata.block_table[num_decodes:]
974
- dequantize_and_gather_k_cache(
975
- kv[:chunk_size],
976
- compressed_k_cache,
977
- seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
978
- gather_lens=None,
979
- block_table=block_table[chunk_start:chunk_end],
980
- block_size=attn_metadata.block_size // self.compress_ratio,
981
- offset=0,
982
- )
983
-
984
- # Gather SWA KV
985
- swa_block_table = swa_metadata.block_table[num_decodes:]
986
- dequantize_and_gather_k_cache(
987
- kv[:chunk_size],
988
- swa_k_cache,
989
- seq_lens=seq_lens[chunk_start:chunk_end],
990
- gather_lens=gather_lens[chunk_start:chunk_end],
991
- block_table=swa_block_table[chunk_start:chunk_end],
992
- block_size=swa_metadata.block_size,
993
- offset=N,
994
- )
995
-
996
- # Combine the topk indices and SWA indices for gathered KV cache
997
- query_start = (
998
- query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
999
- )
1000
- query_end = (
1001
- query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
1002
- )
1003
-
1004
- combined_indices, combined_lens = combine_topk_swa_indices(
1005
- topk_indices[query_start:query_end],
1006
- query_start_loc[
1007
- num_decodes + chunk_start : num_decodes + chunk_end + 1
1008
- ],
1009
- seq_lens[chunk_start:chunk_end],
1010
- gather_lens[chunk_start:chunk_end],
1011
- self.window_size,
1012
- self.compress_ratio,
1013
- top_k,
1014
- M,
1015
- N,
1016
- )
1017
- flash_mla_sparse_fwd(
1018
- q=q[query_start:query_end],
1019
- kv=kv.view(-1, 1, q.shape[-1]),
1020
- indices=combined_indices.unsqueeze(1),
1021
- sm_scale=self.scale,
1022
- attn_sink=self.attn_sink,
1023
- topk_length=combined_lens,
1024
- out=output[query_start:query_end],
1025
- )
739
+ self.impl_cls.forward_mqa(self, q, kv, positions, output)
740
741
742
class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
vllm/v1/attention/backends/mla/flashmla_sparse.py
+5
-27
@@ -117,8 +117,11 @@ class FlashMLASparseBackend(AttentionBackend):
117
118
@classmethod
119
def get_supported_head_sizes(cls) -> list[int]:
120
- # V3.2: 576 (512 NoPE + 64 RoPE); DeepseekV4: 512 (448 NoPE + 64 RoPE)
121
- return [512, 576]
120
+ # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576.
121
+ # DeepSeek V4 uses 448 NoPE + 64 RoPE = 512 and overrides this in
122
+ # vllm/models/deepseek_v4/nvidia/flashmla.py:
123
+ # DeepseekV4FlashMLASparseBackend.get_supported_head_sizes.
124
+ return [576]
125
126
@classmethod
127
def is_mla(cls) -> bool:
@@ -147,31 +150,6 @@ class FlashMLASparseBackend(AttentionBackend):
150
return (num_blocks, block_size, head_size)
151
152
150
-class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
151
- @staticmethod
152
- def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
153
- return [256]
154
-
155
- @staticmethod
156
- def get_name() -> str:
157
- return "V4_FLASHMLA_SPARSE"
158
-
159
- @staticmethod
160
- def get_kv_cache_shape(
161
- num_blocks: int,
162
- block_size: int,
163
- num_kv_heads: int,
164
- head_size: int,
165
- cache_dtype_str: str = "auto",
166
- ) -> tuple[int, ...]:
167
- if cache_dtype_str == "fp8_ds_mla":
168
- # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
169
- # head_size passed in is the semantic head_dim (512).
170
- return (num_blocks, block_size, 584)
171
- else:
172
- return (num_blocks, block_size, head_size)
173
-
174
-
153
@dataclass
154
class FlashMLASparseMetadata(AttentionMetadata):
155
num_reqs: int
vllm/v1/attention/backends/mla/sparse_swa.py
+1
-1
@@ -113,7 +113,7 @@ class DeepseekSparseSWABackend(AttentionBackend):
113
@staticmethod
114
def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]:
115
if current_platform.is_rocm():
116
- from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
116
+ from vllm.models.deepseek_v4.amd.rocm import (
117
DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
118
)
119