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 1, 21:15 UTC
[DSV4] Refactor DeepseekV4Attention (#44569)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 4, 2026 at 20:23 UTC
4efd6ffde09477800294a8ed9cc752017812c3b1
8 files changed
+521
-918
vllm/models/deepseek_v4/amd/model.py
+3
-161
@@ -18,7 +18,6 @@ from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClam
18
from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (
21
- ColumnParallelLinear,
21
MergedColumnParallelLinear,
22
RowParallelLinear,
23
)
@@ -45,11 +44,7 @@ from vllm.model_executor.models.utils import (
44
make_layers,
45
maybe_prefix,
46
)
48
-from vllm.models.deepseek_v4.attention import (
49
- DeepseekV4Indexer,
50
- DeepseekV4MLA,
51
-)
52
-from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
47
+from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
48
from vllm.platforms import current_platform
49
from vllm.sequence import IntermediateTensors
50
from vllm.utils.import_utils import has_tilelang
@@ -225,158 +220,6 @@ class DeepseekV4MoE(nn.Module):
220
return final_hidden_states.view(org_shape)
221
222
228
-class DeepseekV4Attention(nn.Module):
229
- def __init__(
230
- self,
231
- vllm_config: VllmConfig,
232
- prefix: str,
233
- topk_indices_buffer: torch.Tensor | None = None,
234
- aux_stream_list: list[torch.cuda.Stream] | None = None,
235
- ):
236
- super().__init__()
237
- config = vllm_config.model_config.hf_config
238
- quant_config = vllm_config.quant_config
239
- layer_id = extract_layer_index(prefix)
240
-
241
- self.layer_id = layer_id
242
- self.hidden_size = config.hidden_size
243
- self.n_heads = config.num_attention_heads
244
- tp_size = get_tensor_model_parallel_world_size()
245
- assert self.n_heads % tp_size == 0
246
-
247
- self.n_local_heads = self.n_heads // tp_size
248
- self.q_lora_rank = config.q_lora_rank
249
- self.o_lora_rank = config.o_lora_rank
250
- self.head_dim = config.head_dim
251
- self.rope_head_dim = config.qk_rope_head_dim
252
- self.nope_head_dim = self.head_dim - self.rope_head_dim
253
- self.n_groups = config.o_groups
254
- self.n_local_groups = self.n_groups // tp_size
255
- self.window_size = config.sliding_window
256
- # NOTE(zyongye) Compress ratio can't be 0
257
- # we do this for because MTP layer is not included
258
- # in the compress ratio list
259
- if layer_id < config.num_hidden_layers:
260
- self.compress_ratio = max(1, config.compress_ratios[layer_id])
261
- else:
262
- self.compress_ratio = 1
263
- self.eps = config.rms_norm_eps
264
- self.max_position_embeddings = config.max_position_embeddings
265
-
266
- # Padded to min 64 heads for FlashMLA, initialized to -inf
267
- # (no sink effect). Weight loading fills the first n_local_heads slots.
268
- padded_heads = max(self.n_local_heads, 64)
269
- self.attn_sink = nn.Parameter(
270
- torch.full((padded_heads,), -float("inf"), dtype=torch.float32),
271
- requires_grad=False,
272
- )
273
-
274
- self.fused_wqa_wkv = MergedColumnParallelLinear(
275
- self.hidden_size,
276
- [self.q_lora_rank, self.head_dim],
277
- bias=False,
278
- quant_config=quant_config,
279
- prefix=f"{prefix}.fused_wqa_wkv",
280
- disable_tp=True, # fused ReplicatedLinear
281
- )
282
- self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
283
- self.wq_b = ColumnParallelLinear(
284
- self.q_lora_rank,
285
- self.n_heads * self.head_dim,
286
- bias=False,
287
- quant_config=quant_config,
288
- return_bias=False,
289
- prefix=f"{prefix}.wq_b",
290
- )
291
-
292
- self.kv_norm = RMSNorm(self.head_dim, self.eps)
293
- self.wo_a = ColumnParallelLinear(
294
- self.n_heads * self.head_dim // self.n_groups,
295
- self.n_groups * self.o_lora_rank,
296
- bias=False,
297
- quant_config=quant_config,
298
- return_bias=False,
299
- prefix=f"{prefix}.wo_a",
300
- )
301
- self.wo_a.is_bmm = True
302
- self.wo_a.bmm_batch_size = self.n_local_groups
303
- self.wo_b = RowParallelLinear(
304
- self.n_groups * self.o_lora_rank,
305
- self.hidden_size,
306
- bias=False,
307
- quant_config=quant_config,
308
- return_bias=False,
309
- prefix=f"{prefix}.wo_b",
310
- )
311
- self.softmax_scale = self.head_dim**-0.5
312
- self.scale_fmt = config.quantization_config["scale_fmt"]
313
-
314
- self.rope_parameters = config.rope_scaling
315
-
316
- # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)
317
- self.rotary_emb = build_deepseek_v4_rope(
318
- config,
319
- head_dim=self.head_dim,
320
- rope_head_dim=self.rope_head_dim,
321
- max_position_embeddings=self.max_position_embeddings,
322
- compress_ratio=self.compress_ratio,
323
- )
324
-
325
- self.indexer = None
326
- if self.compress_ratio == 4:
327
- # Only C4A uses sparse attention and hence has indexer.
328
- self.indexer = DeepseekV4Indexer(
329
- vllm_config,
330
- config=config,
331
- hidden_size=self.hidden_size,
332
- q_lora_rank=self.q_lora_rank,
333
- quant_config=quant_config,
334
- cache_config=vllm_config.cache_config,
335
- topk_indices_buffer=topk_indices_buffer,
336
- compress_ratio=self.compress_ratio,
337
- prefix=f"{prefix}.indexer",
338
- )
339
-
340
- self.mla_attn = DeepseekV4MLA(
341
- hidden_size=self.hidden_size,
342
- num_heads=self.n_local_heads,
343
- head_dim=self.head_dim,
344
- scale=self.softmax_scale,
345
- qk_nope_head_dim=self.nope_head_dim,
346
- qk_rope_head_dim=self.rope_head_dim,
347
- v_head_dim=self.head_dim,
348
- q_lora_rank=self.q_lora_rank,
349
- kv_lora_rank=self.head_dim,
350
- o_lora_rank=self.o_lora_rank,
351
- vllm_config=vllm_config,
352
- fused_wqa_wkv=self.fused_wqa_wkv,
353
- q_norm=self.q_norm,
354
- wq_b=self.wq_b,
355
- kv_norm=self.kv_norm,
356
- wo_a=self.wo_a,
357
- wo_b=self.wo_b,
358
- attn_sink=self.attn_sink,
359
- rotary_emb=self.rotary_emb,
360
- indexer=self.indexer,
361
- indexer_rotary_emb=self.rotary_emb,
362
- topk_indices_buffer=topk_indices_buffer,
363
- aux_stream_list=aux_stream_list,
364
- window_size=self.window_size,
365
- compress_ratio=self.compress_ratio,
366
- cache_config=vllm_config.cache_config,
367
- quant_config=quant_config,
368
- prefix=prefix,
369
- )
370
-
371
- def forward(
372
- self,
373
- positions: torch.Tensor,
374
- hidden_states: torch.Tensor,
375
- llama_4_scaling: torch.Tensor | None,
376
- ):
377
- return self.mla_attn(positions, hidden_states, llama_4_scaling)
378
-
379
-
223
class DeepseekV4DecoderLayer(nn.Module):
224
def __init__(
225
self,
@@ -395,7 +238,7 @@ class DeepseekV4DecoderLayer(nn.Module):
238
self.hidden_size = config.hidden_size
239
240
self.rms_norm_eps = config.rms_norm_eps
398
- self.attn = DeepseekV4Attention(
241
+ self.attn = DeepseekV4ROCMAiterMLAAttention(
242
vllm_config,
243
prefix=f"{prefix}.attn",
244
topk_indices_buffer=topk_indices_buffer,
@@ -601,7 +444,7 @@ class DeepseekV4Model(nn.Module):
444
self.rms_norm_eps = config.rms_norm_eps
445
446
# Three aux streams: one per non-default input GEMM in
604
- # DeepseekV4MLA.attn_gemm_parallel_execute
447
+ # DeepseekV4Attention.attn_gemm_parallel_execute
448
# (compressor kv_score, indexer.weights_proj, indexer.compressor
449
# kv_score). fused_wqa_wkv stays on the default stream.
450
# Disable them on ROCm because of hang issues.
@@ -897,7 +740,6 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
740
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
741
},
742
orig_to_new_substr={
900
- ".attn.compressor.": ".attn.mla_attn.compressor.",
743
".shared_experts.w2": ".shared_experts.down_proj",
744
},
745
)
vllm/models/deepseek_v4/amd/rocm.py
+65
-68
@@ -2,15 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, cast
5
+from typing import cast
6
7
import torch
8
9
from vllm.forward_context import get_forward_context
10
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
11
from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
12
from vllm.models.deepseek_v4.nvidia.flashmla import (
13
DeepseekV4FlashMLASparseBackend,
13
- DeepseekV4SparseMLAAttentionImpl,
14
)
15
from vllm.triton_utils import tl, triton
16
from vllm.v1.attention.backend import (
@@ -26,16 +26,12 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
26
)
27
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
28
build_ragged_indices_from_dense,
29
+ rocm_inv_rope_einsum,
30
rocm_sparse_attn_decode,
31
rocm_sparse_attn_prefill,
32
)
33
from vllm.v1.worker.workspace import current_workspace_manager
34
34
-if TYPE_CHECKING:
35
- from vllm.models.deepseek_v4.attention import (
36
- DeepseekV4MLAAttention,
37
- )
38
-
35
36
def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
37
lengths = lengths.to(dtype=torch.int32).contiguous()
@@ -582,13 +578,9 @@ class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend):
578
def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]:
579
return DeepseekV4ROCMAiterMLASparseMetadataBuilder
580
585
- @staticmethod
586
- def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]:
587
- return DeepseekV4ROCMAiterMLASparseImpl
588
-
581
590
-class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
591
- """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer."""
582
+class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
583
+ """ROCm sparse MLA attention layer for DeepSeek V4."""
584
585
backend_cls = DeepseekV4ROCMAiterMLASparseBackend
586
@@ -596,10 +588,21 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
588
def get_padded_num_q_heads(cls, num_heads: int) -> int:
589
return num_heads
590
599
- @classmethod
600
- def forward_mqa( # type: ignore[override]
601
- cls,
602
- layer: "DeepseekV4MLAAttention",
591
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
592
+ # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
593
+ z = rocm_inv_rope_einsum(
594
+ self.rotary_emb,
595
+ o,
596
+ positions,
597
+ self.rope_head_dim,
598
+ self.n_local_groups,
599
+ self.o_lora_rank,
600
+ self.wo_a,
601
+ )
602
+ return self.wo_b(z.flatten(1))
603
+
604
+ def forward_mqa(
605
+ self,
606
q: torch.Tensor,
607
kv: torch.Tensor,
608
positions: torch.Tensor,
@@ -619,16 +622,16 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
622
# Warmup dummy run: no real metadata. Reserve the same bf16
623
# gather workspace _forward_prefill would; the dequantize / topk
624
# / sparse_fwd kernels are skipped this step.
622
- swa_only = layer.compress_ratio <= 1
625
+ swa_only = self.compress_ratio <= 1
626
N = (
627
0
628
if swa_only
626
- else (layer.max_model_len + layer.compress_ratio - 1)
627
- // layer.compress_ratio
629
+ else (self.max_model_len + self.compress_ratio - 1)
630
+ // self.compress_ratio
631
)
629
- M = N + layer.window_size + layer.max_num_batched_tokens
632
+ M = N + self.window_size + self.max_num_batched_tokens
633
current_workspace_manager().get_simultaneous(
631
- ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
634
+ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
635
)
636
output.zero_()
637
return
@@ -636,25 +639,24 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
639
assert isinstance(attn_metadata, dict)
640
rocm_metadata = cast(
641
DeepseekV4ROCMAiterMLASparseMetadata | None,
639
- attn_metadata.get(layer.prefix),
642
+ attn_metadata.get(self.prefix),
643
)
644
swa_metadata = cast(
645
DeepseekV4ROCMAiterSparseSWAMetadata | None,
643
- attn_metadata.get(layer.swa_cache_layer.prefix),
646
+ attn_metadata.get(self.swa_cache_layer.prefix),
647
)
648
assert swa_metadata is not None
649
647
- swa_only = layer.compress_ratio <= 1
648
- self_kv_cache = layer.kv_cache if not swa_only else None
649
- swa_kv_cache = layer.swa_cache_layer.kv_cache
650
+ swa_only = self.compress_ratio <= 1
651
+ self_kv_cache = self.kv_cache if not swa_only else None
652
+ swa_kv_cache = self.swa_cache_layer.kv_cache
653
654
num_decodes = swa_metadata.num_decodes
655
num_prefills = swa_metadata.num_prefills
656
num_decode_tokens = swa_metadata.num_decode_tokens
657
658
if num_prefills > 0:
656
- cls._forward_prefill(
657
- layer=layer,
659
+ self._forward_prefill(
660
q=q[num_decode_tokens:],
661
positions=positions[num_decode_tokens:],
662
compressed_k_cache=self_kv_cache,
@@ -664,8 +666,7 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
666
swa_metadata=swa_metadata,
667
)
668
if num_decodes > 0:
667
- cls._forward_decode(
668
- layer=layer,
669
+ self._forward_decode(
670
q=q[:num_decode_tokens],
671
kv_cache=self_kv_cache,
672
swa_metadata=swa_metadata,
@@ -674,10 +675,8 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
675
output=output[:num_decode_tokens],
676
)
677
677
- @classmethod
678
def _forward_decode(
679
- cls,
680
- layer: "DeepseekV4MLAAttention",
679
+ self,
680
q: torch.Tensor,
681
kv_cache: torch.Tensor | None,
682
swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
@@ -695,16 +694,16 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
694
if not swa_only:
695
assert attn_metadata is not None
696
assert swa_metadata.is_valid_token is not None
698
- block_size = attn_metadata.block_size // layer.compress_ratio
697
+ block_size = attn_metadata.block_size // self.compress_ratio
698
is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
700
- if layer.compress_ratio == 4:
701
- assert layer.topk_indices_buffer is not None
699
+ if self.compress_ratio == 4:
700
+ assert self.topk_indices_buffer is not None
701
(
702
topk_ragged_indices,
703
topk_ragged_indptr,
704
topk_lens,
705
) = compute_global_topk_ragged_indices_and_indptr(
707
- layer.topk_indices_buffer[:num_decode_tokens],
706
+ self.topk_indices_buffer[:num_decode_tokens],
707
swa_metadata.token_to_req_indices,
708
attn_metadata.block_table[:num_decodes],
709
block_size,
@@ -719,7 +718,7 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
718
rocm_sparse_attn_decode(
719
q=q,
720
kv_cache=kv_cache,
722
- swa_k_cache=layer.swa_cache_layer.kv_cache,
721
+ swa_k_cache=self.swa_cache_layer.kv_cache,
722
swa_only=swa_only,
723
topk_indices=topk_indices,
724
topk_lens=topk_lens,
@@ -729,18 +728,16 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
728
swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr,
729
topk_ragged_indices=topk_ragged_indices,
730
topk_ragged_indptr=topk_ragged_indptr,
732
- attn_sink=layer.attn_sink,
733
- scale=layer.scale,
734
- head_dim=layer.head_dim,
735
- nope_head_dim=layer.nope_head_dim,
736
- rope_head_dim=layer.rope_head_dim,
731
+ attn_sink=self.attn_sink,
732
+ scale=self.scale,
733
+ head_dim=self.head_dim,
734
+ nope_head_dim=self.nope_head_dim,
735
+ rope_head_dim=self.rope_head_dim,
736
output=output,
737
)
738
740
- @classmethod
739
def _forward_prefill(
742
- cls,
743
- layer: "DeepseekV4MLAAttention",
740
+ self,
741
q: torch.Tensor,
742
positions: torch.Tensor,
743
compressed_k_cache: torch.Tensor | None,
@@ -768,34 +765,34 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
765
prefill_token_base = query_start_loc_cpu[num_decodes]
766
767
if not swa_only:
771
- if layer.compress_ratio == 4:
772
- assert layer.topk_indices_buffer is not None
773
- topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
768
+ if self.compress_ratio == 4:
769
+ assert self.topk_indices_buffer is not None
770
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
771
topk_indices = topk_indices[:num_prefill_tokens]
772
else:
773
assert attn_metadata is not None
774
topk_indices = attn_metadata.c128a_prefill_topk_indices
775
assert topk_indices is not None
776
top_k = topk_indices.shape[-1]
780
- N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio
777
+ N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
778
else:
782
- assert layer.topk_indices_buffer is not None
783
- topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
779
+ assert self.topk_indices_buffer is not None
780
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
781
top_k = 0
782
N = 0
783
787
- M = N + layer.window_size + layer.max_num_batched_tokens
788
- num_chunks = (num_prefills + cls.PREFILL_CHUNK_SIZE - 1) // (
789
- cls.PREFILL_CHUNK_SIZE
784
+ M = N + self.window_size + self.max_num_batched_tokens
785
+ num_chunks = (num_prefills + self.PREFILL_CHUNK_SIZE - 1) // (
786
+ self.PREFILL_CHUNK_SIZE
787
)
788
789
workspace_manager = current_workspace_manager()
790
kv = workspace_manager.get_simultaneous(
794
- ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
791
+ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
792
)[0]
793
for chunk_idx in range(num_chunks):
797
- chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE
798
- chunk_end = min(chunk_start + cls.PREFILL_CHUNK_SIZE, num_prefills)
794
+ chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
795
+ chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
796
chunk_size = chunk_end - chunk_start
797
if not swa_only:
798
assert attn_metadata is not None
@@ -804,10 +801,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
801
dequantize_and_gather_k_cache(
802
kv[:chunk_size],
803
compressed_k_cache,
807
- seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio,
804
+ seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
805
gather_lens=None,
806
block_table=block_table[chunk_start:chunk_end],
810
- block_size=attn_metadata.block_size // layer.compress_ratio,
807
+ block_size=attn_metadata.block_size // self.compress_ratio,
808
offset=0,
809
)
810
@@ -836,8 +833,8 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
833
],
834
seq_lens[chunk_start:chunk_end],
835
gather_lens[chunk_start:chunk_end],
839
- layer.window_size,
840
- layer.compress_ratio,
836
+ self.window_size,
837
+ self.compress_ratio,
838
top_k,
839
M,
840
N,
@@ -847,10 +844,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
844
kv=kv.view(-1, 1, q.shape[-1]),
845
indices=combined_indices,
846
topk_length=combined_lens,
850
- scale=layer.scale,
851
- head_dim=layer.head_dim,
852
- nope_head_dim=layer.nope_head_dim,
853
- rope_head_dim=layer.rope_head_dim,
854
- attn_sink=layer.attn_sink,
847
+ scale=self.scale,
848
+ head_dim=self.head_dim,
849
+ nope_head_dim=self.nope_head_dim,
850
+ rope_head_dim=self.rope_head_dim,
851
+ attn_sink=self.attn_sink,
852
output=output[query_start:query_end],
853
)
vllm/models/deepseek_v4/attention.py
+224
-345
@@ -4,8 +4,9 @@
4
DeepseekV4 MLA Attention Layer
5
"""
6
7
+from abc import ABC, abstractmethod
8
from collections.abc import Callable
8
-from typing import TYPE_CHECKING, Any, cast
9
+from typing import TYPE_CHECKING, Any, ClassVar, cast
10
11
import torch
12
import torch.nn as nn
@@ -15,16 +16,16 @@ from transformers import DeepseekV2Config, DeepseekV3Config
16
import vllm.envs as envs
17
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
18
from vllm.model_executor.layers.linear import (
19
+ ColumnParallelLinear,
20
+ MergedColumnParallelLinear,
21
ReplicatedLinear,
22
+ RowParallelLinear,
23
)
24
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
25
from vllm.models.deepseek_v4.common.ops import (
26
fused_indexer_q_rope_quant,
23
- fused_inv_rope_fp8_quant,
27
fused_q_kv_rmsnorm,
28
)
26
-from vllm.utils.deep_gemm import fp8_einsum
27
-from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
29
30
if TYPE_CHECKING:
31
from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -42,14 +43,9 @@ from vllm.logger import init_logger
43
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.quantization import QuantizationConfig
45
-from vllm.model_executor.layers.quantization.input_quant_fp8 import (
46
- QuantFP8,
47
-)
48
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
49
- GroupShape,
50
-)
46
+from vllm.model_executor.models.utils import extract_layer_index
47
+from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
48
from vllm.models.deepseek_v4.compressor import DeepseekCompressor
52
-from vllm.platforms import current_platform
49
from vllm.utils.multi_stream_utils import (
50
execute_in_parallel,
51
maybe_execute_in_parallel,
@@ -62,187 +58,209 @@ from vllm.v1.attention.backends.mla.indexer import (
58
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
59
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
60
65
-if TYPE_CHECKING:
66
- from vllm.models.deepseek_v4.nvidia.flashmla import (
67
- DeepseekV4SparseMLAAttentionImpl,
68
- )
69
-
61
logger = init_logger(__name__)
62
63
73
-def _resolve_dsv4_backend(vllm_config: VllmConfig | None):
74
- """Return the explicitly-requested DSv4 sparse backend enum, or None."""
75
- if vllm_config is None:
76
- return None
77
- attn_config = getattr(vllm_config, "attention_config", None)
78
- return getattr(attn_config, "backend", None) if attn_config is not None else None
79
-
80
-
81
-def _select_v4_sparse_impl(
82
- vllm_config: VllmConfig | None = None,
83
-) -> "type[DeepseekV4SparseMLAAttentionImpl]":
84
- """Pick the V4 sparse MLA impl class.
85
-
86
- An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
87
- FlashInfer TRTLLM-gen path; otherwise the platform default (FlashMLA on
88
- NVIDIA, ROCm Aiter on AMD) is used.
89
- """
90
- from vllm.v1.attention.backends.registry import AttentionBackendEnum
91
-
92
- backend = _resolve_dsv4_backend(vllm_config)
93
- if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
94
- from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
95
- DeepseekV4FlashInferMLASparseImpl,
96
- )
97
-
98
- logger.info_once("Using FLASHINFER_MLA_SPARSE_DSV4 backend.")
99
- return DeepseekV4FlashInferMLASparseImpl
100
- if current_platform.is_rocm():
101
- from vllm.models.deepseek_v4.amd.rocm import (
102
- DeepseekV4ROCMAiterMLASparseImpl,
103
- )
104
-
105
- logger.info_once("Using ROCM_FLASHMLA_SPARSE_DSV4 backend.")
106
- return DeepseekV4ROCMAiterMLASparseImpl
107
- from vllm.models.deepseek_v4.nvidia.flashmla import (
108
- DeepseekV4FlashMLASparseImpl,
109
- )
110
-
111
- logger.info_once("Using FLASHMLA_SPARSE_DSV4 backend.")
112
- return DeepseekV4FlashMLASparseImpl
113
-
114
-
64
def _resolve_dsv4_kv_cache_dtype(
116
- backend,
65
+ use_flashmla_fp8_layout: bool,
66
kv_cache_dtype: str,
67
cache_config: CacheConfig | None,
68
) -> tuple[str, torch.dtype]:
120
- """Map ``(backend, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
121
-
122
- FlashInfer V4 reads a contiguous 512-wide KV row (bf16 or per-tensor FP8
123
- E4M3); FlashMLA V4 reads the legacy UE8M0 paged layout (uint8 /
124
- ``fp8_ds_mla``). For FlashMLA the canonical ``fp8_ds_mla`` string is
125
- written back onto ``cache_config`` so the page-size specs pick the 576B
126
- layout.
69
+ """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
70
+
71
+ Both layouts are paged; they differ in the per-token block format. The
72
+ FlashMLA fp8 layout (FlashMLA / ROCm Aiter) is the ``fp8_ds_mla`` format:
73
+ UE8M0 block-scaled fp8 packed as ``uint8`` (the canonical ``fp8_ds_mla``
74
+ string is written back onto ``cache_config`` so the page-size specs pick
75
+ the 576B per-token slot). Otherwise (FlashInfer) each token's KV row is
76
+ stored in its plain element dtype — bf16 or per-tensor FP8 E4M3.
77
"""
128
- from vllm.v1.attention.backends.registry import AttentionBackendEnum
129
-
130
- if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
131
- if kv_cache_dtype.startswith("fp8"):
132
- return kv_cache_dtype, torch.float8_e4m3fn
133
- # auto / bfloat16 -> contiguous BF16 cache.
134
- return kv_cache_dtype, torch.bfloat16
135
-
136
- # FlashMLA (and ROCm Aiter): legacy UE8M0 paged uint8 cache.
137
- assert kv_cache_dtype.startswith("fp8"), (
138
- f"DeepseekV4 FlashMLA sparse backend only supports fp8 kv-cache, "
139
- f"got {kv_cache_dtype}"
140
- )
141
- if kv_cache_dtype != "fp8_ds_mla":
142
- if cache_config is not None:
143
- cache_config.cache_dtype = "fp8_ds_mla"
144
- kv_cache_dtype = "fp8_ds_mla"
145
- logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
146
- return kv_cache_dtype, torch.uint8
78
+ if use_flashmla_fp8_layout:
79
+ # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
80
+ assert kv_cache_dtype.startswith("fp8"), (
81
+ f"DeepseekV4 FlashMLA fp8 layout only supports fp8 kv-cache, "
82
+ f"got {kv_cache_dtype}"
83
+ )
84
+ if kv_cache_dtype != "fp8_ds_mla":
85
+ if cache_config is not None:
86
+ cache_config.cache_dtype = "fp8_ds_mla"
87
+ kv_cache_dtype = "fp8_ds_mla"
88
+ logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
89
+ return kv_cache_dtype, torch.uint8
90
+
91
+ # Plain bf16 / per-tensor fp8 KV row (FlashInfer).
92
+ if kv_cache_dtype.startswith("fp8"):
93
+ return kv_cache_dtype, torch.float8_e4m3fn
94
+ # auto / bfloat16 -> plain bf16 KV row.
95
+ return kv_cache_dtype, torch.bfloat16
96
+
97
+
98
+class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
99
+ """DeepseekV4 MLA attention layer.
100
+
101
+ The platform-specific sparse-MLA forward (``forward_mqa`` /
102
+ ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
103
+ subclass — ``DeepseekV4FlashMLAAttention`` / ``DeepseekV4FlashInferMLAAttention``
104
+ (CUDA) or ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the
105
+ platform-specific deepseek_v4 model module. The base is never instantiated
106
+ directly.
107
+ """
108
+
109
+ # Provided by the platform subclass.
110
+ backend_cls: ClassVar[type[AttentionBackend]]
111
+ # KV-cache per-token block format (both layouts are paged). True (default)
112
+ # = FlashMLA / ROCm fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8);
113
+ # False = FlashInfer plain bf16 / per-tensor fp8 KV row.
114
+ use_flashmla_fp8_layout: ClassVar[bool] = True
115
+ # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
116
+ # workspace allocated in _forward_prefill and is also read by the dummy-run
117
+ # path to pre-reserve that workspace.
118
+ PREFILL_CHUNK_SIZE: ClassVar[int] = 4
119
+
120
+ @classmethod
121
+ @abstractmethod
122
+ def get_padded_num_q_heads(cls, num_heads: int) -> int:
123
+ """Q head count the q/output buffers are allocated at.
124
+
125
+ The layer allocates the q/output buffers at
126
+ ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy
127
+ ``result >= num_heads``. Backends with no padding constraint return
128
+ ``num_heads``.
129
+ """
130
+ raise NotImplementedError
131
+
132
+ @abstractmethod
133
+ def forward_mqa(
134
+ self,
135
+ q: torch.Tensor,
136
+ kv: torch.Tensor,
137
+ positions: torch.Tensor,
138
+ output: torch.Tensor,
139
+ ) -> None:
140
+ """Platform-specific sparse MLA forward; writes attention into ``output``."""
141
+ raise NotImplementedError
142
143
+ @abstractmethod
144
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
145
+ """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
146
+ raise NotImplementedError
147
149
-class DeepseekV4MLA(nn.Module):
148
def __init__(
149
self,
152
- hidden_size: int,
153
- num_heads: int,
154
- head_dim: int,
155
- scale: float,
156
- qk_nope_head_dim: int,
157
- qk_rope_head_dim: int,
158
- v_head_dim: int,
159
- q_lora_rank: int | None,
160
- kv_lora_rank: int,
161
- o_lora_rank: int | None,
150
vllm_config: VllmConfig,
163
- fused_wqa_wkv: torch.nn.Module,
164
- q_norm: torch.nn.Module,
165
- wq_b: torch.nn.Module,
166
- kv_norm: torch.nn.Module,
167
- wo_a: torch.nn.Module,
168
- wo_b: torch.nn.Module,
169
- attn_sink: torch.nn.Module,
170
- rotary_emb: torch.nn.Module,
171
- indexer: torch.nn.Module | None,
172
- indexer_rotary_emb: torch.nn.Module,
173
- topk_indices_buffer: torch.Tensor | None,
174
- aux_stream_list: list[torch.cuda.Stream] | None,
175
- window_size: int,
176
- compress_ratio: int | None,
177
- cache_config: CacheConfig | None = None,
178
- quant_config: QuantizationConfig | None = None,
179
- prefix: str = "",
151
+ prefix: str,
152
+ topk_indices_buffer: torch.Tensor | None = None,
153
+ aux_stream_list: list[torch.cuda.Stream] | None = None,
154
) -> None:
155
super().__init__()
182
- self.hidden_size = hidden_size
183
- self.n_local_heads = num_heads
184
- self.head_dim = head_dim
185
- self.scale = scale
186
-
187
- self.q_lora_rank = q_lora_rank
188
- self.kv_lora_rank = kv_lora_rank
189
- self.window_size = window_size
190
- self.compress_ratio = compress_ratio if compress_ratio is not None else 1
191
- self.prefix = prefix
192
-
193
- # Extract config from vllm_config
156
config = vllm_config.model_config.hf_config
157
+ quant_config = vllm_config.quant_config
158
+ cache_config = vllm_config.cache_config
159
tp_size = get_tensor_model_parallel_world_size()
160
+ layer_id = extract_layer_index(prefix)
161
197
- # DeepseekV4-specific attributes (num_heads is already TP-adjusted)
198
- self.eps = config.rms_norm_eps
199
- self.rope_head_dim = config.qk_rope_head_dim
200
- self.nope_head_dim = head_dim - self.rope_head_dim
201
- self.n_local_groups = config.o_groups // tp_size
162
+ self.prefix = prefix # Alias for compatibility with compressor
163
+ self.hidden_size = config.hidden_size
164
+ self.n_heads = config.num_attention_heads
165
+ assert self.n_heads % tp_size == 0
166
+ self.n_local_heads = self.n_heads // tp_size
167
+ self.q_lora_rank = config.q_lora_rank
168
self.o_lora_rank = config.o_lora_rank
169
+ self.head_dim = config.head_dim
170
+ self.rope_head_dim = config.qk_rope_head_dim
171
+ self.nope_head_dim = self.head_dim - self.rope_head_dim
172
+ self.n_groups = config.o_groups
173
+ self.n_local_groups = self.n_groups // tp_size
174
+ self.window_size = config.sliding_window
175
+ # NOTE(zyongye) Compress ratio can't be 0
176
+ # we do this for because MTP layer is not included
177
+ # in the compress ratio list
178
+ if layer_id < config.num_hidden_layers:
179
+ self.compress_ratio = max(1, config.compress_ratios[layer_id])
180
+ else:
181
+ self.compress_ratio = 1
182
+ self.eps = config.rms_norm_eps
183
+ self.scale = self.head_dim**-0.5
184
+
185
+ # Padded Q head count is dictated by the platform subclass.
186
+ self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads)
187
+ # Sink padded to the same head count, initialized to -inf (no sink
188
+ # effect). Weight loading fills the first n_local_heads slots.
189
+ self.attn_sink = nn.Parameter(
190
+ torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32),
191
+ requires_grad=False,
192
+ )
193
204
- # Store projection modules
205
- self.fused_wqa_wkv = fused_wqa_wkv
206
- self.q_norm = q_norm
207
- self.wq_b = wq_b
194
+ self.fused_wqa_wkv = MergedColumnParallelLinear(
195
+ self.hidden_size,
196
+ [self.q_lora_rank, self.head_dim],
197
+ bias=False,
198
+ quant_config=quant_config,
199
+ prefix=f"{prefix}.fused_wqa_wkv",
200
+ disable_tp=True, # fused ReplicatedLinear
201
+ )
202
+ self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
203
+ self.wq_b = ColumnParallelLinear(
204
+ self.q_lora_rank,
205
+ self.n_heads * self.head_dim,
206
+ bias=False,
207
+ quant_config=quant_config,
208
+ return_bias=False,
209
+ prefix=f"{prefix}.wq_b",
210
+ )
211
209
- self.kv_norm = kv_norm
210
- self.wo_a = wo_a
212
+ self.kv_norm = RMSNorm(self.head_dim, self.eps)
213
+ self.wo_a = ColumnParallelLinear(
214
+ self.n_heads * self.head_dim // self.n_groups,
215
+ self.n_groups * self.o_lora_rank,
216
+ bias=False,
217
+ quant_config=quant_config,
218
+ return_bias=False,
219
+ prefix=f"{prefix}.wo_a",
220
+ )
221
+ self.wo_a.is_bmm = True
222
+ self.wo_a.bmm_batch_size = self.n_local_groups
223
+ self.wo_b = RowParallelLinear(
224
+ self.n_groups * self.o_lora_rank,
225
+ self.hidden_size,
226
+ bias=False,
227
+ quant_config=quant_config,
228
+ return_bias=False,
229
+ prefix=f"{prefix}.wo_b",
230
+ )
231
212
- self._wo_a_act_quant = QuantFP8(
213
- static=False,
214
- group_shape=GroupShape(1, 128),
215
- use_ue8m0=True,
232
+ # Initialize rotary embedding before the indexer/compressor consume it.
233
+ self.rotary_emb = build_deepseek_v4_rope(
234
+ config,
235
+ head_dim=self.head_dim,
236
+ rope_head_dim=self.rope_head_dim,
237
+ max_position_embeddings=config.max_position_embeddings,
238
+ compress_ratio=self.compress_ratio,
239
)
217
- # Bypass packed-for-deepgemm path — we need FP32 scales (not packed
218
- # INT32) so fp8_einsum can handle layout transform internally.
219
- self._wo_a_act_quant.use_deep_gemm_supported = False
220
- self.wo_b = wo_b
221
-
222
- # Pick fp8_einsum recipe based on GPU arch:
223
- # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128
224
- # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1
225
- cap = current_platform.get_device_capability()
226
- assert cap is not None, "DeepseekV4 attention requires a CUDA device"
227
- self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
228
- self._tma_aligned_scales = cap.major >= 10
229
-
230
- self.rotary_emb = rotary_emb
231
- self.indexer_rotary_emb = indexer_rotary_emb
240
+ self.indexer_rotary_emb = self.rotary_emb
241
self.topk_indices_buffer = topk_indices_buffer
242
234
- self.indexer = indexer
235
-
236
- # Per-head RMS normalization for Q (no learnable weights)
237
- self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False)
238
-
239
- # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic
240
- head_bytes = (
241
- self.nope_head_dim # 448 fp8 NoPE
242
- + self.rope_head_dim * 2 # 64 bf16 RoPE
243
- + self.nope_head_dim // 64 # 7B scale factors
244
- + 1 # 1B pad
245
- )
243
+ self.indexer = None
244
+ if self.compress_ratio == 4:
245
+ # Only C4A uses sparse attention and hence has indexer.
246
+ # aux_stream_list[2] is free here (outer GEMMs joined) for the inner
247
+ # overlap of wq_b+fused_indexer_q_rope_quant vs compressor. None on
248
+ # ROCm, where aux_stream_list is None.
249
+ indexer_aux_stream = (
250
+ aux_stream_list[2] if aux_stream_list is not None else None
251
+ )
252
+ self.indexer = DeepseekV4Indexer(
253
+ vllm_config,
254
+ config=config,
255
+ hidden_size=self.hidden_size,
256
+ q_lora_rank=self.q_lora_rank,
257
+ quant_config=quant_config,
258
+ cache_config=cache_config,
259
+ topk_indices_buffer=topk_indices_buffer,
260
+ compress_ratio=self.compress_ratio,
261
+ prefix=f"{prefix}.indexer",
262
+ aux_stream=indexer_aux_stream,
263
+ )
264
265
# Will be None on ROCm for now.
266
self.aux_stream_list = aux_stream_list
@@ -252,45 +270,39 @@ class DeepseekV4MLA(nn.Module):
270
self.ln_events = [torch.cuda.Event() for _ in range(4)]
271
272
assert cache_config is not None, "DeepseekV4 attention requires cache_config"
255
- # Resolve the SWA cache tensor dtype from the selected backend: FlashMLA
256
- # uses the legacy UE8M0 paged uint8 layout; FlashInfer uses a contiguous
257
- # bf16 / per-tensor fp8 row.
258
- backend = _resolve_dsv4_backend(vllm_config)
259
- _, swa_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
260
- backend, cache_config.cache_dtype, cache_config
273
+ # ---- Attention / KV-cache setup ----
274
+ self.max_num_batched_tokens = (
275
+ vllm_config.scheduler_config.max_num_batched_tokens
276
)
277
+ self.max_model_len = vllm_config.model_config.max_model_len
278
+
279
+ # Resolve the kv-cache dtype from this backend's block format (a
280
+ # ClassVar set by the subclass): fp8_ds_mla (UE8M0 block-scaled fp8 as
281
+ # uint8) for FlashMLA / ROCm, vs a plain bf16 / per-tensor fp8 row for
282
+ # FlashInfer. The same resolution drives the SWA cache tensor dtype
283
+ # below.
284
+ self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
285
+ self.use_flashmla_fp8_layout, cache_config.cache_dtype, cache_config
286
+ )
287
+
288
self.swa_cache_layer = DeepseekV4SWACache(
289
head_dim=self.head_dim,
290
window_size=self.window_size,
265
- dtype=swa_cache_torch_dtype,
291
+ dtype=self.kv_cache_torch_dtype,
292
prefix=f"{prefix}.swa_cache",
293
cache_config=cache_config,
294
)
295
270
- self.mla_attn = DeepseekV4MLAAttention(
271
- num_heads=self.n_local_heads,
272
- head_dim=self.head_dim,
273
- scale=self.scale,
274
- qk_nope_head_dim=self.nope_head_dim,
275
- qk_rope_head_dim=self.rope_head_dim,
276
- q_lora_rank=self.q_lora_rank,
277
- kv_lora_rank=self.kv_lora_rank,
278
- compress_ratio=self.compress_ratio,
279
- window_size=self.window_size,
280
- head_bytes=head_bytes,
281
- swa_cache_layer=self.swa_cache_layer,
282
- attn_sink=attn_sink, # already padded with -inf
283
- cache_config=cache_config,
284
- quant_config=quant_config,
285
- prefix=prefix,
286
- indexer=self.indexer,
287
- topk_indices_buffer=self.topk_indices_buffer,
288
- )
289
- # Mirror the inner layer's padded head count (single source of truth).
290
- self.padded_heads = self.mla_attn.padded_heads
296
+ # Register with compilation context for metadata lookup.
297
+ compilation_config = vllm_config.compilation_config
298
+ if prefix and prefix in compilation_config.static_forward_context:
299
+ raise ValueError(f"Duplicate layer name: {prefix}")
300
+ if prefix:
301
+ compilation_config.static_forward_context[prefix] = self
302
+ self.kv_cache = torch.tensor([])
303
292
- # Create the compressor for layers with compress_ratio > 1; after
293
- # creating the DeepseekV4MLAAttention layer to get its cache.
304
+ # Create the compressor for layers with compress_ratio > 1; after the
305
+ # attention setup above so its KV-cache prefix (self.prefix) is set.
306
self.compressor = None
307
if self.compress_ratio > 1:
308
self.compressor = DeepseekCompressor(
@@ -300,7 +312,7 @@ class DeepseekV4MLA(nn.Module):
312
head_dim=self.head_dim,
313
rotate=True,
314
prefix=f"{prefix}.compressor",
303
- k_cache_prefix=self.mla_attn.prefix,
315
+ k_cache_prefix=self.prefix,
316
)
317
318
def forward(
@@ -324,48 +336,8 @@ class DeepseekV4MLA(nn.Module):
336
self.attention_impl(hidden_states, positions, o_padded)
337
o = o_padded[:, : self.n_local_heads, :]
338
327
- # Keep ROCm on the BF16 reference wo_a path util kernel ready.
328
- if current_platform.is_rocm():
329
- z = rocm_inv_rope_einsum(
330
- self.rotary_emb,
331
- o,
332
- positions,
333
- self.rope_head_dim,
334
- self.n_local_groups,
335
- self.o_lora_rank,
336
- self.wo_a,
337
- )
338
- return self.wo_b(z.flatten(1))
339
-
340
- # O projection: inverse RoPE + FP8 quant + einsum + wo_b
341
- o_fp8, o_scale = fused_inv_rope_fp8_quant(
342
- o,
343
- positions,
344
- self.rotary_emb.cos_sin_cache,
345
- n_groups=self.n_local_groups,
346
- heads_per_group=self.n_local_heads // self.n_local_groups,
347
- nope_dim=self.nope_head_dim,
348
- rope_dim=self.rope_head_dim,
349
- tma_aligned_scales=self._tma_aligned_scales,
350
- )
351
-
352
- wo_a_fp8 = self.wo_a.weight
353
- wo_a_scale = self.wo_a.weight_scale_inv
354
-
355
- z = torch.empty(
356
- (num_tokens, self.n_local_groups, self.o_lora_rank),
357
- device=o.device,
358
- dtype=torch.bfloat16,
359
- )
360
- fp8_einsum(
361
- "bhr,hdr->bhd",
362
- (o_fp8, o_scale),
363
- (wo_a_fp8, wo_a_scale),
364
- z,
365
- recipe=self._einsum_recipe,
366
- )
367
-
368
- return self.wo_b(z.flatten(1))
339
+ # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
340
+ return self._o_proj(o, positions)
341
342
def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
343
aux_streams = self.aux_stream_list
@@ -451,7 +423,7 @@ class DeepseekV4MLA(nn.Module):
423
)
424
425
# wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
454
- # on the default stream so q stays on its consumer stream (mla_attn
426
+ # on the default stream so q stays on its consumer stream (forward_mqa
427
# downstream reads q on default). Indexer/compressor go on aux for
428
# overlap with default's GEMM + cache write.
429
if self.indexer is not None:
@@ -514,7 +486,7 @@ class DeepseekV4MLA(nn.Module):
486
487
# MLA attention writes into the pre-allocated `out` buffer
488
# ([num_tokens, padded_heads, head_dim]).
517
- self.mla_attn(q, kv, positions, output=out)
489
+ self.forward_mqa(q, kv, positions, out)
490
491
def _fused_qnorm_rope_kv_insert(
492
self,
@@ -549,7 +521,7 @@ class DeepseekV4MLA(nn.Module):
521
cos_sin_cache = self.rotary_emb.cos_sin_cache
522
cache_dtype = swa_kv_cache.dtype
523
552
- # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
524
+ # kv is unchanged; attention reads kv solely via swa_kv_cache.
525
if cache_dtype == torch.uint8:
526
# Legacy FlashMLA UE8M0 paged path. Horizontally fused:
527
# Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
@@ -569,9 +541,10 @@ class DeepseekV4MLA(nn.Module):
541
swa_metadata.block_size,
542
)
543
572
- # FlashInfer full-cache path: contiguous [num_blocks, block_size, 512]
573
- # cache (no Q padding). bf16 rewrites q in place; per-tensor fp8 writes a
574
- # separately-allocated fp8 q and quantizes the KV row.
544
+ # FlashInfer full-cache path: the [num_blocks, block_size, 512] cache
545
+ # stores the KV row in its plain dtype (no Q padding). bf16 rewrites q
546
+ # in place; per-tensor fp8 writes a separately-allocated fp8 q and
547
+ # quantizes the KV row.
548
block_size = swa_metadata.block_size
549
swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
550
if cache_dtype == torch.bfloat16:
@@ -597,99 +570,13 @@ class DeepseekV4MLA(nn.Module):
570
swa_metadata.slot_mapping,
571
positions,
572
cos_sin_cache,
600
- self.mla_attn._flashinfer_fp8_kv_scale,
601
- self.mla_attn._flashinfer_fp8_q_scale_inv,
573
+ self._flashinfer_fp8_kv_scale,
574
+ self._flashinfer_fp8_q_scale_inv,
575
self.eps,
576
block_size,
577
)
578
return q_fp8
579
607
-
608
-class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
609
- def __init__(
610
- self,
611
- num_heads: int,
612
- head_dim: int,
613
- scale: float,
614
- qk_nope_head_dim: int,
615
- qk_rope_head_dim: int,
616
- q_lora_rank: int | None,
617
- kv_lora_rank: int,
618
- compress_ratio: int,
619
- window_size: int,
620
- head_bytes: int,
621
- swa_cache_layer: DeepseekV4SWACache,
622
- attn_sink: torch.Tensor,
623
- cache_config: CacheConfig | None = None,
624
- quant_config: QuantizationConfig | None = None,
625
- prefix: str = "",
626
- # Sparse MLA Args
627
- indexer: object | None = None,
628
- topk_indices_buffer: torch.Tensor | None = None,
629
- aux_stream: torch.cuda.Stream | None = None,
630
- **extra_impl_args,
631
- ) -> None:
632
- super().__init__()
633
- vllm_config = get_current_vllm_config()
634
- self.impl_cls = _select_v4_sparse_impl(vllm_config)
635
- self.backend_cls = self.impl_cls.backend_cls
636
- self.num_heads = num_heads
637
- self.num_kv_heads = 1
638
- self.head_dim = head_dim
639
- self.scale = scale
640
- self.window_size = window_size
641
- self.head_bytes = head_bytes
642
- self.compress_ratio = compress_ratio
643
- self.q_lora_rank = q_lora_rank
644
- self.kv_lora_rank = kv_lora_rank
645
- self.nope_head_dim = qk_nope_head_dim
646
- self.rope_head_dim = qk_rope_head_dim
647
- self.indexer = indexer
648
- self.topk_indices_buffer = topk_indices_buffer
649
-
650
- self.prefix = prefix # Alias for compatibility with compressor
651
-
652
- self.aux_stream = aux_stream
653
- self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
654
-
655
- # Padded Q head count is dictated by the selected impl.
656
- self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads)
657
-
658
- # Store attention sink
659
- assert attn_sink is not None
660
- self.attn_sink: torch.Tensor = attn_sink
661
- # Store SWA cache
662
- assert swa_cache_layer is not None
663
- self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer
664
-
665
- # Get vllm config for cache setup
666
- self.max_num_batched_tokens = (
667
- vllm_config.scheduler_config.max_num_batched_tokens
668
- )
669
- self.max_model_len = vllm_config.model_config.max_model_len
670
-
671
- # Resolve the kv-cache dtype from the selected backend. FlashMLA uses
672
- # the legacy UE8M0 paged uint8 (fp8_ds_mla) layout; FlashInfer uses a
673
- # contiguous bf16 / per-tensor fp8 row.
674
- backend = _resolve_dsv4_backend(vllm_config)
675
- kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"
676
- self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
677
- backend, kv_cache_dtype, cache_config
678
- )
679
-
680
- # Per-impl layer buffers (e.g. FlashInfer FP8 scale buffers). No-op for
681
- # the FlashMLA / ROCm impls.
682
- self.impl_cls.init_layer_buffers(self)
683
-
684
- # Register with compilation context for metadata lookup
685
- compilation_config = vllm_config.compilation_config
686
- if prefix and prefix in compilation_config.static_forward_context:
687
- raise ValueError(f"Duplicate layer name: {prefix}")
688
- if prefix:
689
- compilation_config.static_forward_context[prefix] = self
690
-
691
- self.kv_cache = torch.tensor([])
692
-
580
def get_attn_backend(self) -> type[AttentionBackend]:
581
return self.backend_cls
582
@@ -698,8 +585,9 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
585
self.compress_ratio <= 1
586
): # SWA part. Allocated separately as DeepseekV4SWACache.
587
return None
701
- # FlashMLA uses the UE8M0 paged uint8 layout (576B aligned); FlashInfer
702
- # uses a contiguous bf16 / per-tensor fp8 cache with no extra alignment.
588
+ # FlashMLA uses the fp8_ds_mla block format (UE8M0 block-scaled fp8 as
589
+ # uint8, 576B aligned); FlashInfer stores a plain bf16 / per-tensor fp8
590
+ # row with no extra alignment.
591
is_flashmla = self.kv_cache_dtype == "fp8_ds_mla"
592
return MLAAttentionSpec(
593
block_size=vllm_config.cache_config.block_size,
@@ -712,15 +600,6 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
600
model_version="deepseek_v4",
601
)
602
715
- def forward(
716
- self,
717
- q: torch.Tensor,
718
- kv: torch.Tensor,
719
- positions: torch.Tensor,
720
- output: torch.Tensor,
721
- ) -> None:
722
- self.impl_cls.forward_mqa(self, q, kv, positions, output)
723
-
603
604
class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
605
def __init__(
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+71
-62
@@ -3,28 +3,30 @@
3
"""DeepSeek V4 FlashInfer TRTLLM-gen sparse MLA backend.
4
5
Uses FlashInfer's public ``trtllm_batch_decode_sparse_mla_dsv4`` launcher with a
6
-contiguous bf16 / per-tensor FP8 KV cache. Shares the V4 sparse-index pipeline
7
-(SWA cache + compressor + indexer, 256-token blocks, head_size 512) with the
8
-FlashMLA V4 backend; only the attention forward differs.
6
+plain bf16 / per-tensor FP8 KV row (vs FlashMLA's packed ``fp8_ds_mla`` block
7
+format). Shares the V4 sparse-index pipeline (SWA cache + compressor + indexer,
8
+256-token blocks, head_size 512) with the FlashMLA V4 backend; only the
9
+attention forward differs.
10
"""
11
11
-from typing import TYPE_CHECKING, cast
12
+from typing import TYPE_CHECKING, ClassVar, cast
13
14
import torch
15
16
from vllm.forward_context import get_forward_context
17
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
18
from vllm.models.deepseek_v4.common.ops import (
19
build_flashinfer_mixed_sparse_indices,
20
)
19
-from vllm.models.deepseek_v4.nvidia.flashmla import (
20
- DeepseekV4FlashMLASparseBackend,
21
- DeepseekV4SparseMLAAttentionImpl,
21
+from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLASparseBackend
22
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
23
+ compute_fp8_einsum_recipe,
24
+ deep_gemm_fp8_o_proj,
25
)
26
from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
27
from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseMetadata
28
29
if TYPE_CHECKING:
27
- from vllm.models.deepseek_v4.attention import DeepseekV4MLAAttention
30
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
31
32
# 128 MB TRTLLM-gen workspace, allocated once per device and zero-initialized
@@ -51,23 +53,21 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLASparseBackend):
53
Inheriting from the FlashMLA V4 backend reuses its ``FlashMLASparseMetadata``
54
builder (which the V4 sparse-index pipeline needs — the V3.2 FlashInfer
55
builder lacks the ``c128a_*`` fields), 256-token blocks, head_size 512, and
54
- the contiguous (num_blocks, block_size, 512) cache shape for non-``fp8_ds_mla``
55
- dtypes.
56
+ the (num_blocks, block_size, 512) cache shape for non-``fp8_ds_mla`` dtypes.
57
"""
58
59
@staticmethod
60
def get_name() -> str:
61
return "FLASHINFER_MLA_SPARSE_DSV4"
62
62
- @staticmethod
63
- def get_impl_cls() -> type["DeepseekV4FlashInferMLASparseImpl"]:
64
- return DeepseekV4FlashInferMLASparseImpl
65
-
63
67
-class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
68
- """FlashInfer TRTLLM-gen sparse MLA implementation for DeepSeek V4."""
64
+class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
65
+ """FlashInfer TRTLLM-gen sparse MLA attention layer for DeepSeek V4."""
66
67
backend_cls = DeepseekV4FlashInferMLASparseBackend
68
+ # FlashInfer stores a plain bf16 / per-tensor fp8 KV row, not the FlashMLA
69
+ # packed fp8_ds_mla block format (UE8M0 block-scaled fp8 as uint8).
70
+ use_flashmla_fp8_layout: ClassVar[bool] = False
71
72
@classmethod
73
def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -79,27 +79,44 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
79
)
80
return 64 if num_heads <= 64 else 128
81
82
- @classmethod
83
- def init_layer_buffers(cls, layer: "DeepseekV4MLAAttention") -> None:
82
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
83
+ return deep_gemm_fp8_o_proj(
84
+ o,
85
+ positions,
86
+ self.rotary_emb.cos_sin_cache,
87
+ self.wo_a,
88
+ self.wo_b,
89
+ n_groups=self.n_local_groups,
90
+ heads_per_group=self.n_local_heads // self.n_local_groups,
91
+ nope_dim=self.nope_head_dim,
92
+ rope_dim=self.rope_head_dim,
93
+ o_lora_rank=self.o_lora_rank,
94
+ einsum_recipe=self._einsum_recipe,
95
+ tma_aligned_scales=self._tma_aligned_scales,
96
+ )
97
+
98
+ def __init__(self, *args, **kwargs) -> None:
99
+ super().__init__(*args, **kwargs)
100
+ self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
101
# Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
85
- # per-tensor FP8 cache path consumes these; bf16 reads ``layer.scale``.
86
- if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
102
+ # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
103
+ if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
104
return
105
# TODO: load real per-tensor Q/KV scales from the checkpoint; unit
106
# scales until the scale tensor names are wired.
107
fp8_q_scale = 1.0
108
fp8_kv_scale = 1.0
92
- layer.register_buffer(
109
+ self.register_buffer(
110
"_flashinfer_fp8_q_scale",
111
torch.tensor([fp8_q_scale], dtype=torch.float32),
112
persistent=False,
113
)
97
- layer.register_buffer(
114
+ self.register_buffer(
115
"_flashinfer_fp8_q_scale_inv",
116
torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
117
persistent=False,
118
)
102
- layer.register_buffer(
119
+ self.register_buffer(
120
"_flashinfer_fp8_kv_scale",
121
torch.tensor([fp8_kv_scale], dtype=torch.float32),
122
persistent=False,
@@ -107,13 +124,11 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
124
# TRTLLM-gen takes scalar scale args on a distinct (correct) C++ path
125
# vs 1-elem tensors, so these are Python floats. bmm1 folds the softmax
126
# scale and the Q/KV per-tensor scales; bmm2 is the KV scale.
110
- layer._flashinfer_fp8_bmm1_scale = layer.scale * fp8_q_scale * fp8_kv_scale
111
- layer._flashinfer_fp8_bmm2_scale = fp8_kv_scale
127
+ self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
128
+ self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
129
113
- @classmethod
114
- def forward_mqa( # type: ignore[override]
115
- cls,
116
- layer: "DeepseekV4MLAAttention",
130
+ def forward_mqa(
131
+ self,
132
q: torch.Tensor,
133
kv: torch.Tensor,
134
positions: torch.Tensor,
@@ -147,21 +162,20 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
162
163
assert isinstance(attn_metadata, dict)
164
flashmla_metadata = cast(
150
- FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix)
165
+ FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)
166
)
167
swa_metadata = cast(
168
"DeepseekSparseSWAMetadata | None",
154
- attn_metadata.get(layer.swa_cache_layer.prefix),
169
+ attn_metadata.get(self.swa_cache_layer.prefix),
170
)
171
assert swa_metadata is not None
172
158
- swa_only = layer.compress_ratio <= 1
173
+ swa_only = self.compress_ratio <= 1
174
# SWA-only layers don't allocate their own compressed KV cache.
160
- self_kv_cache = layer.kv_cache if not swa_only else None
161
- swa_kv_cache = layer.swa_cache_layer.kv_cache
175
+ self_kv_cache = self.kv_cache if not swa_only else None
176
+ swa_kv_cache = self.swa_cache_layer.kv_cache
177
163
- cls._forward(
164
- layer=layer,
178
+ self._forward(
179
q=q,
180
kv_cache=self_kv_cache,
181
swa_k_cache=swa_kv_cache,
@@ -171,10 +185,8 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
185
output=output,
186
)
187
174
- @classmethod
188
def _build_sparse_index_metadata(
176
- cls,
177
- layer: "DeepseekV4MLAAttention",
189
+ self,
190
kv_cache: torch.Tensor | None,
191
swa_k_cache: torch.Tensor,
192
swa_metadata: "DeepseekSparseSWAMetadata",
@@ -200,17 +212,17 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
212
assert swa_metadata.block_table is not None
213
214
decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
203
- num_decode_tokens, layer.window_size
215
+ num_decode_tokens, self.window_size
216
)
217
decode_compressed_topk_lens = None
218
decode_compressed_indices_are_local = False
219
decode_is_valid_token = None
220
221
if swa_only:
210
- assert layer.topk_indices_buffer is not None
222
+ assert self.topk_indices_buffer is not None
223
compressed_kv_cache = swa_k_cache
224
decode_compressed_indices = None
213
- prefill_topk_indices = layer.topk_indices_buffer[
225
+ prefill_topk_indices = self.topk_indices_buffer[
226
num_decode_tokens:num_tokens, :0
227
]
228
compressed_block_table = None
@@ -221,24 +233,24 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
233
assert attn_metadata is not None
234
compressed_kv_cache = kv_cache
235
compressed_block_table = attn_metadata.block_table[:num_reqs]
224
- compressed_block_size = attn_metadata.block_size // layer.compress_ratio
236
+ compressed_block_size = attn_metadata.block_size // self.compress_ratio
237
226
- if layer.compress_ratio == 4:
227
- assert layer.topk_indices_buffer is not None
238
+ if self.compress_ratio == 4:
239
+ assert self.topk_indices_buffer is not None
240
if num_prefill_tokens > 0:
229
- prefill_topk_indices = layer.topk_indices_buffer[
241
+ prefill_topk_indices = self.topk_indices_buffer[
242
num_decode_tokens:num_tokens
243
]
244
top_k = prefill_topk_indices.shape[-1]
245
else:
234
- prefill_topk_indices = layer.topk_indices_buffer[:0, :0]
246
+ prefill_topk_indices = self.topk_indices_buffer[:0, :0]
247
top_k = 0
248
249
decode_compressed_indices_are_local = True
250
assert swa_metadata.is_valid_token is not None
251
decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
252
if num_decode_tokens > 0:
241
- decode_compressed_indices = layer.topk_indices_buffer[
253
+ decode_compressed_indices = self.topk_indices_buffer[
254
:num_decode_tokens
255
]
256
else:
@@ -284,18 +296,16 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
296
swa_metadata.block_size,
297
compressed_block_table,
298
compressed_block_size,
287
- layer.window_size,
288
- layer.compress_ratio,
299
+ self.window_size,
300
+ self.compress_ratio,
301
top_k,
302
decode_compressed_indices_are_local=decode_compressed_indices_are_local,
303
decode_is_valid_token=decode_is_valid_token,
304
)
305
return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
306
295
- @classmethod
307
def _forward(
297
- cls,
298
- layer: "DeepseekV4MLAAttention",
308
+ self,
309
q: torch.Tensor,
310
kv_cache: torch.Tensor | None,
311
swa_k_cache: torch.Tensor,
@@ -304,7 +314,7 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
314
swa_only: bool,
315
output: torch.Tensor,
316
) -> None:
307
- assert layer.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
317
+ assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
318
num_decodes = swa_metadata.num_decodes
319
num_prefills = swa_metadata.num_prefills
320
num_decode_tokens = swa_metadata.num_decode_tokens
@@ -319,8 +329,7 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
329
seq_lens,
330
sparse_indices,
331
sparse_topk_lens,
322
- ) = cls._build_sparse_index_metadata(
323
- layer=layer,
332
+ ) = self._build_sparse_index_metadata(
333
kv_cache=kv_cache,
334
swa_k_cache=swa_k_cache,
335
swa_metadata=swa_metadata,
@@ -332,12 +341,12 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
341
# restrict to the real tokens (the launcher validates sparse indices).
342
query = q[:num_tokens]
343
output = output[:num_tokens]
335
- bmm1_scale: float | torch.Tensor = layer.scale
344
+ bmm1_scale: float | torch.Tensor = self.scale
345
bmm2_scale: float | torch.Tensor = 1.0
337
- if layer.kv_cache_torch_dtype == torch.float8_e4m3fn:
346
+ if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
347
assert query.dtype == torch.float8_e4m3fn
339
- bmm1_scale = layer._flashinfer_fp8_bmm1_scale
340
- bmm2_scale = layer._flashinfer_fp8_bmm2_scale
348
+ bmm1_scale = self._flashinfer_fp8_bmm1_scale
349
+ bmm2_scale = self._flashinfer_fp8_bmm2_scale
350
else:
351
assert query.dtype == torch.bfloat16
352
query = query.contiguous()
@@ -376,7 +385,7 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
385
out=output[:num_decode_tokens],
386
bmm1_scale=bmm1_scale,
387
bmm2_scale=bmm2_scale,
379
- sinks=layer.attn_sink,
388
+ sinks=self.attn_sink,
389
cum_seq_lens_q=decode_cu,
390
max_q_len=int(decode_lens_cpu.max().item()),
391
)
@@ -401,7 +410,7 @@ class DeepseekV4FlashInferMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
410
out=output[num_decode_tokens:num_tokens],
411
bmm1_scale=bmm1_scale,
412
bmm2_scale=bmm2_scale,
404
- sinks=layer.attn_sink,
413
+ sinks=self.attn_sink,
414
cum_seq_lens_q=prefill_cu,
415
max_q_len=int(prefill_lens_cpu.max().item()),
416
)
vllm/models/deepseek_v4/nvidia/flashmla.py
+72
-118
@@ -1,22 +1,22 @@
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
4
+from typing import TYPE_CHECKING, cast
5
6
import torch
7
8
from vllm.forward_context import get_forward_context
9
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
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,
15
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
16
+ compute_fp8_einsum_recipe,
17
+ deep_gemm_fp8_o_proj,
18
)
19
+from vllm.v1.attention.backend import MultipleOf
20
from vllm.v1.attention.backends.mla.flashmla_sparse import (
21
FlashMLASparseBackend,
22
FlashMLASparseMetadata,
@@ -28,63 +28,9 @@ from vllm.v1.attention.ops.flashmla import (
28
from vllm.v1.worker.workspace import current_workspace_manager
29
30
if TYPE_CHECKING:
31
- from vllm.models.deepseek_v4.attention import (
32
- DeepseekV4MLAAttention,
33
- )
31
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
32
33
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
- @classmethod
67
- @abstractmethod
68
- def get_padded_num_q_heads(cls, num_heads: int) -> int:
69
- """Q head count the backend wants q allocated at.
70
-
71
- The MLA wrapper allocates the q/output buffers at
72
- ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must
73
- satisfy ``result >= num_heads``. Backends with no padding constraint
74
- return ``num_heads``.
75
- """
76
- raise NotImplementedError
77
-
78
- @classmethod
79
- def init_layer_buffers(cls, layer: "DeepseekV4MLAAttention") -> None:
80
- """Register impl-specific buffers on the layer at construction.
81
-
82
- No-op by default; FlashInfer overrides this to register its per-tensor
83
- FP8 scale buffers.
84
- """
85
- return None
86
-
87
-
34
class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
35
@staticmethod
36
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@@ -94,10 +40,6 @@ class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
40
def get_name() -> str:
41
return "FLASHMLA_SPARSE_DSV4"
42
97
- @staticmethod
98
- def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]:
99
- return DeepseekV4FlashMLASparseImpl
100
-
43
@classmethod
44
def get_supported_head_sizes(cls) -> list[int]:
45
# DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512 (overrides the
@@ -120,11 +62,31 @@ class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
62
return (num_blocks, block_size, head_size)
63
64
123
-class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
124
- """FlashMLA sparse MLA implementation for DeepSeek V4's custom MLA layer."""
65
+class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
66
+ """FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA)."""
67
68
backend_cls = DeepseekV4FlashMLASparseBackend
69
70
+ def __init__(self, *args, **kwargs) -> None:
71
+ super().__init__(*args, **kwargs)
72
+ self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
73
+
74
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
75
+ return deep_gemm_fp8_o_proj(
76
+ o,
77
+ positions,
78
+ self.rotary_emb.cos_sin_cache,
79
+ self.wo_a,
80
+ self.wo_b,
81
+ n_groups=self.n_local_groups,
82
+ heads_per_group=self.n_local_heads // self.n_local_groups,
83
+ nope_dim=self.nope_head_dim,
84
+ rope_dim=self.rope_head_dim,
85
+ o_lora_rank=self.o_lora_rank,
86
+ einsum_recipe=self._einsum_recipe,
87
+ tma_aligned_scales=self._tma_aligned_scales,
88
+ )
89
+
90
@classmethod
91
def get_padded_num_q_heads(cls, num_heads: int) -> int:
92
# FP8 decode kernel only supports h_q = 64 or 128.
@@ -135,10 +97,8 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
97
)
98
return 64 if num_heads <= 64 else 128
99
138
- @classmethod
139
- def forward_mqa( # type: ignore[override]
140
- cls,
141
- layer: "DeepseekV4MLAAttention",
100
+ def forward_mqa(
101
+ self,
102
q: torch.Tensor,
103
kv: torch.Tensor,
104
positions: torch.Tensor,
@@ -159,35 +119,35 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
119
# Warmup dummy run: no real metadata. Reserve the same bf16
120
# gather workspace _forward_prefill would; the dequantize / topk
121
# / sparse_fwd kernels are skipped this step.
162
- swa_only = layer.compress_ratio <= 1
122
+ swa_only = self.compress_ratio <= 1
123
N = (
124
0
125
if swa_only
166
- else (layer.max_model_len + layer.compress_ratio - 1)
167
- // layer.compress_ratio
126
+ else (self.max_model_len + self.compress_ratio - 1)
127
+ // self.compress_ratio
128
)
169
- M = N + layer.window_size + layer.max_num_batched_tokens
129
+ M = N + self.window_size + self.max_num_batched_tokens
130
current_workspace_manager().get_simultaneous(
171
- ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
131
+ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
132
)
133
output.zero_()
134
return
135
136
assert isinstance(attn_metadata, dict)
137
flashmla_metadata = cast(
178
- FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix)
138
+ FlashMLASparseMetadata | None, attn_metadata.get(self.prefix)
139
)
140
swa_metadata = cast(
141
"DeepseekSparseSWAMetadata | None",
182
- attn_metadata.get(layer.swa_cache_layer.prefix),
142
+ attn_metadata.get(self.swa_cache_layer.prefix),
143
)
144
assert swa_metadata is not None
145
186
- swa_only = layer.compress_ratio <= 1
146
+ swa_only = self.compress_ratio <= 1
147
# SWA-only layers (compress_ratio <= 1) don't have their own KV cache
188
- # allocation, so layer.kv_cache may be empty after profiling cleanup.
189
- self_kv_cache = layer.kv_cache if not swa_only else None
190
- swa_kv_cache = layer.swa_cache_layer.kv_cache
148
+ # allocation, so self.kv_cache may be empty after profiling cleanup.
149
+ self_kv_cache = self.kv_cache if not swa_only else None
150
+ swa_kv_cache = self.swa_cache_layer.kv_cache
151
152
# Split prefill and decode
153
num_decodes = swa_metadata.num_decodes
@@ -195,8 +155,7 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
155
num_decode_tokens = swa_metadata.num_decode_tokens
156
157
if num_prefills > 0:
198
- cls._forward_prefill(
199
- layer=layer,
158
+ self._forward_prefill(
159
q=q[num_decode_tokens:],
160
positions=positions[num_decode_tokens:],
161
compressed_k_cache=self_kv_cache,
@@ -206,8 +165,7 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
165
swa_metadata=swa_metadata,
166
)
167
if num_decodes > 0:
209
- cls._forward_decode(
210
- layer=layer,
168
+ self._forward_decode(
169
q=q[:num_decode_tokens],
170
kv_cache=self_kv_cache,
171
swa_metadata=swa_metadata,
@@ -216,10 +174,8 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
174
output=output[:num_decode_tokens],
175
)
176
219
- @classmethod
177
def _forward_decode(
221
- cls,
222
- layer: "DeepseekV4MLAAttention",
178
+ self,
179
q: torch.Tensor,
180
kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
181
swa_metadata: "DeepseekSparseSWAMetadata",
@@ -235,13 +191,13 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
191
if not swa_only:
192
assert attn_metadata is not None
193
assert swa_metadata.is_valid_token is not None
238
- block_size = attn_metadata.block_size // layer.compress_ratio
194
+ block_size = attn_metadata.block_size // self.compress_ratio
195
is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
240
- if layer.compress_ratio == 4:
196
+ if self.compress_ratio == 4:
197
# C4A: local indices differ per layer (filled by Indexer).
242
- assert layer.topk_indices_buffer is not None
198
+ assert self.topk_indices_buffer is not None
199
global_indices, topk_lens = compute_global_topk_indices_and_lens(
244
- layer.topk_indices_buffer[:num_decode_tokens],
200
+ self.topk_indices_buffer[:num_decode_tokens],
201
swa_metadata.token_to_req_indices,
202
attn_metadata.block_table[:num_decodes],
203
block_size,
@@ -258,12 +214,12 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
214
215
# We treat queries in the same seq as different queries
216
# and later we only attend by generated indices.
261
- # q arrives pre-padded to layer.padded_heads by the outer wrapper.
217
+ # q arrives pre-padded to self.padded_heads by the outer wrapper.
218
q = q.unsqueeze(1)
219
220
# Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
221
# Use unsqueeze to preserve strides (handles padded blocks correctly)
266
- swa_cache = layer.swa_cache_layer.kv_cache.unsqueeze(-2)
222
+ swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2)
223
# Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
224
if kv_cache is not None:
225
kv_cache = kv_cache.unsqueeze(-2)
@@ -274,20 +230,20 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
230
# and num_splits via PyTorch's graph-aware allocator so CUDA graph
231
# capture reuses the same addresses on replay); subsequent same-type
232
# layers see have_initialized=True and skip the planner.
277
- if layer.compress_ratio <= 1:
233
+ if self.compress_ratio <= 1:
234
tile_metadata = swa_metadata.tile_sched_swaonly
279
- elif layer.compress_ratio == 4:
235
+ elif self.compress_ratio == 4:
236
tile_metadata = swa_metadata.tile_sched_c4a
281
- elif layer.compress_ratio == 128:
237
+ elif self.compress_ratio == 128:
238
tile_metadata = swa_metadata.tile_sched_c128a
239
else:
240
raise ValueError(
285
- f"Unsupported compress_ratio={layer.compress_ratio}; "
241
+ f"Unsupported compress_ratio={self.compress_ratio}; "
242
"expected 1, 4, or 128."
243
)
244
assert tile_metadata is not None, (
245
"swa_metadata missing tile_sched entry for "
290
- f"compress_ratio={layer.compress_ratio}; "
246
+ f"compress_ratio={self.compress_ratio}; "
247
"DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not "
248
"allocate one for this layer type."
249
)
@@ -302,18 +258,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
258
is_fp8_kvcache=True,
259
indices=swa_indices,
260
topk_length=swa_lens,
305
- softmax_scale=layer.scale,
306
- attn_sink=layer.attn_sink,
261
+ softmax_scale=self.scale,
262
+ attn_sink=self.attn_sink,
263
extra_k_cache=kv_cache if not swa_only else None,
264
extra_indices_in_kvcache=topk_indices,
265
extra_topk_length=topk_lens,
266
out=output.unsqueeze(1),
267
)
268
313
- @classmethod
269
def _forward_prefill(
315
- cls,
316
- layer: "DeepseekV4MLAAttention",
270
+ self,
271
q: torch.Tensor,
272
positions: torch.Tensor,
273
compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1
@@ -343,9 +297,9 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
297
prefill_token_base = query_start_loc_cpu[num_decodes]
298
299
if not swa_only:
346
- if layer.compress_ratio == 4:
347
- assert layer.topk_indices_buffer is not None
348
- topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
300
+ if self.compress_ratio == 4:
301
+ assert self.topk_indices_buffer is not None
302
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
303
topk_indices = topk_indices[:num_prefill_tokens]
304
else:
305
# C128A: pre-computed during metadata build.
@@ -355,16 +309,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
309
# Compressed region must fit the full compressed pool (seq_len //
310
# compress_ratio), not just top_k. top_k bounds how many indices
311
# the indexer selects, not the pool size it indexes into.
358
- N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio
312
+ N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
313
else:
314
# NOTE(woosuk): topk_indices will not be used for SWA-only layers.
361
- assert layer.topk_indices_buffer is not None
362
- topk_indices = layer.topk_indices_buffer[num_decode_tokens:]
315
+ assert self.topk_indices_buffer is not None
316
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
317
top_k = 0
318
N = 0
319
366
- M = N + layer.window_size + layer.max_num_batched_tokens
367
- chunk_size_const = cls.PREFILL_CHUNK_SIZE
320
+ M = N + self.window_size + self.max_num_batched_tokens
321
+ chunk_size_const = self.PREFILL_CHUNK_SIZE
322
num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const
323
324
workspace_manager = current_workspace_manager()
@@ -382,10 +336,10 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
336
dequantize_and_gather_k_cache(
337
kv[:chunk_size],
338
compressed_k_cache,
385
- seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio,
339
+ seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
340
gather_lens=None,
341
block_table=block_table[chunk_start:chunk_end],
388
- block_size=attn_metadata.block_size // layer.compress_ratio,
342
+ block_size=attn_metadata.block_size // self.compress_ratio,
343
offset=0,
344
)
345
@@ -416,8 +370,8 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
370
],
371
seq_lens[chunk_start:chunk_end],
372
gather_lens[chunk_start:chunk_end],
419
- layer.window_size,
420
- layer.compress_ratio,
373
+ self.window_size,
374
+ self.compress_ratio,
375
top_k,
376
M,
377
N,
@@ -426,8 +380,8 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):
380
q=q[query_start:query_end],
381
kv=kv.view(-1, 1, q.shape[-1]),
382
indices=combined_indices.unsqueeze(1),
429
- sm_scale=layer.scale,
430
- attn_sink=layer.attn_sink,
383
+ sm_scale=self.scale,
384
+ attn_sink=self.attn_sink,
385
topk_length=combined_lens,
386
out=output[query_start:query_end],
387
)
vllm/models/deepseek_v4/nvidia/model.py
+17
-163
@@ -33,7 +33,6 @@ from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
33
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
34
from vllm.model_executor.layers.layernorm import RMSNorm
35
from vllm.model_executor.layers.linear import (
36
- ColumnParallelLinear,
36
MergedColumnParallelLinear,
37
RowParallelLinear,
38
)
@@ -55,13 +54,14 @@ from vllm.model_executor.models.utils import (
54
maybe_prefix,
55
)
56
from vllm.model_executor.utils import set_weight_attrs
58
-from vllm.models.deepseek_v4.attention import (
59
- DeepseekV4Indexer,
60
- DeepseekV4MLA,
57
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
58
+from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
59
+ DeepseekV4FlashInferMLAAttention,
60
)
62
-from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
61
+from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
62
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
63
from vllm.sequence import IntermediateTensors
64
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
65
66
67
class DeepseekV4MLP(nn.Module):
@@ -713,163 +713,18 @@ class DeepseekV4MoE(nn.Module):
713
self.experts.finalize_weights()
714
715
716
-class DeepseekV4Attention(nn.Module):
717
- def __init__(
718
- self,
719
- vllm_config: VllmConfig,
720
- prefix: str,
721
- topk_indices_buffer: torch.Tensor | None = None,
722
- aux_stream_list: list[torch.cuda.Stream] | None = None,
723
- ):
724
- super().__init__()
725
- config = vllm_config.model_config.hf_config
726
- quant_config = vllm_config.quant_config
727
- layer_id = extract_layer_index(prefix)
728
-
729
- self.layer_id = layer_id
730
- self.hidden_size = config.hidden_size
731
- self.n_heads = config.num_attention_heads
732
- tp_size = get_tensor_model_parallel_world_size()
733
- assert self.n_heads % tp_size == 0
734
-
735
- self.n_local_heads = self.n_heads // tp_size
736
- self.q_lora_rank = config.q_lora_rank
737
- self.o_lora_rank = config.o_lora_rank
738
- self.head_dim = config.head_dim
739
- self.rope_head_dim = config.qk_rope_head_dim
740
- self.nope_head_dim = self.head_dim - self.rope_head_dim
741
- self.n_groups = config.o_groups
742
- self.n_local_groups = self.n_groups // tp_size
743
- self.window_size = config.sliding_window
744
- # NOTE(zyongye) Compress ratio can't be 0
745
- # we do this for because MTP layer is not included
746
- # in the compress ratio list
747
- if layer_id < config.num_hidden_layers:
748
- self.compress_ratio = max(1, config.compress_ratios[layer_id])
749
- else:
750
- self.compress_ratio = 1
751
- self.eps = config.rms_norm_eps
752
- self.max_position_embeddings = config.max_position_embeddings
753
-
754
- # Padded to min 64 heads for FlashMLA, initialized to -inf
755
- # (no sink effect). Weight loading fills the first n_local_heads slots.
756
- padded_heads = max(self.n_local_heads, 64)
757
- self.attn_sink = nn.Parameter(
758
- torch.full((padded_heads,), -float("inf"), dtype=torch.float32),
759
- requires_grad=False,
760
- )
761
-
762
- self.fused_wqa_wkv = MergedColumnParallelLinear(
763
- self.hidden_size,
764
- [self.q_lora_rank, self.head_dim],
765
- bias=False,
766
- quant_config=quant_config,
767
- prefix=f"{prefix}.fused_wqa_wkv",
768
- disable_tp=True, # fused ReplicatedLinear
769
- )
770
- self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
771
- self.wq_b = ColumnParallelLinear(
772
- self.q_lora_rank,
773
- self.n_heads * self.head_dim,
774
- bias=False,
775
- quant_config=quant_config,
776
- return_bias=False,
777
- prefix=f"{prefix}.wq_b",
778
- )
716
+def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
717
+ """Pick the CUDA sparse-MLA attention class for the configured backend.
718
780
- self.kv_norm = RMSNorm(self.head_dim, self.eps)
781
- self.wo_a = ColumnParallelLinear(
782
- self.n_heads * self.head_dim // self.n_groups,
783
- self.n_groups * self.o_lora_rank,
784
- bias=False,
785
- quant_config=quant_config,
786
- return_bias=False,
787
- prefix=f"{prefix}.wo_a",
788
- )
789
- self.wo_a.is_bmm = True
790
- self.wo_a.bmm_batch_size = self.n_local_groups
791
- self.wo_b = RowParallelLinear(
792
- self.n_groups * self.o_lora_rank,
793
- self.hidden_size,
794
- bias=False,
795
- quant_config=quant_config,
796
- return_bias=False,
797
- prefix=f"{prefix}.wo_b",
798
- )
799
- self.softmax_scale = self.head_dim**-0.5
800
- self.scale_fmt = config.quantization_config["scale_fmt"]
801
-
802
- self.rope_parameters = config.rope_scaling
803
-
804
- # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)
805
- self.rotary_emb = build_deepseek_v4_rope(
806
- config,
807
- head_dim=self.head_dim,
808
- rope_head_dim=self.rope_head_dim,
809
- max_position_embeddings=self.max_position_embeddings,
810
- compress_ratio=self.compress_ratio,
811
- )
812
-
813
- self.indexer = None
814
- if self.compress_ratio == 4:
815
- # Only C4A uses sparse attention and hence has indexer.
816
- # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is
817
- # free here (outer GEMMs joined) for the inner overlap of
818
- # wq_b+fused_indexer_q_rope_quant vs compressor.
819
- indexer_aux_stream = (
820
- aux_stream_list[2] if aux_stream_list is not None else None
821
- )
822
- self.indexer = DeepseekV4Indexer(
823
- vllm_config,
824
- config=config,
825
- hidden_size=self.hidden_size,
826
- q_lora_rank=self.q_lora_rank,
827
- quant_config=quant_config,
828
- cache_config=vllm_config.cache_config,
829
- topk_indices_buffer=topk_indices_buffer,
830
- compress_ratio=self.compress_ratio,
831
- prefix=f"{prefix}.indexer",
832
- aux_stream=indexer_aux_stream,
833
- )
834
-
835
- self.mla_attn = DeepseekV4MLA(
836
- hidden_size=self.hidden_size,
837
- num_heads=self.n_local_heads,
838
- head_dim=self.head_dim,
839
- scale=self.softmax_scale,
840
- qk_nope_head_dim=self.nope_head_dim,
841
- qk_rope_head_dim=self.rope_head_dim,
842
- v_head_dim=self.head_dim,
843
- q_lora_rank=self.q_lora_rank,
844
- kv_lora_rank=self.head_dim,
845
- o_lora_rank=self.o_lora_rank,
846
- vllm_config=vllm_config,
847
- fused_wqa_wkv=self.fused_wqa_wkv,
848
- q_norm=self.q_norm,
849
- wq_b=self.wq_b,
850
- kv_norm=self.kv_norm,
851
- wo_a=self.wo_a,
852
- wo_b=self.wo_b,
853
- attn_sink=self.attn_sink,
854
- rotary_emb=self.rotary_emb,
855
- indexer=self.indexer,
856
- indexer_rotary_emb=self.rotary_emb,
857
- topk_indices_buffer=topk_indices_buffer,
858
- aux_stream_list=aux_stream_list,
859
- window_size=self.window_size,
860
- compress_ratio=self.compress_ratio,
861
- cache_config=vllm_config.cache_config,
862
- quant_config=quant_config,
863
- prefix=prefix,
864
- )
865
-
866
- def forward(
867
- self,
868
- positions: torch.Tensor,
869
- hidden_states: torch.Tensor,
870
- llama_4_scaling: torch.Tensor | None,
719
+ An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
720
+ FlashInfer TRTLLM-gen path; otherwise the FlashMLA path is used.
721
+ """
722
+ if (
723
+ vllm_config.attention_config.backend
724
+ == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4
725
):
872
- return self.mla_attn(positions, hidden_states, llama_4_scaling)
726
+ return DeepseekV4FlashInferMLAAttention
727
+ return DeepseekV4FlashMLAAttention
728
729
730
class DeepseekV4DecoderLayer(nn.Module):
@@ -886,7 +741,7 @@ class DeepseekV4DecoderLayer(nn.Module):
741
self.hidden_size = config.hidden_size
742
743
self.rms_norm_eps = config.rms_norm_eps
889
- self.attn = DeepseekV4Attention(
744
+ self.attn = _select_dsv4_attn_cls(vllm_config)(
745
vllm_config,
746
prefix=f"{prefix}.attn",
747
topk_indices_buffer=topk_indices_buffer,
@@ -1043,7 +898,7 @@ class DeepseekV4Model(nn.Module):
898
self.rms_norm_eps = config.rms_norm_eps
899
900
# Three aux streams: one per non-default input GEMM in
1046
- # DeepseekV4MLA.attn_gemm_parallel_execute
901
+ # DeepseekV4Attention.attn_gemm_parallel_execute
902
# (compressor kv_score, indexer.weights_proj, indexer.compressor
903
# kv_score). fused_wqa_wkv stays on the default stream.
904
aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
@@ -1341,7 +1196,6 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1196
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1197
},
1198
orig_to_new_substr={
1344
- ".attn.compressor.": ".attn.mla_attn.compressor.",
1199
".shared_experts.w2": ".shared_experts.down_proj",
1200
},
1201
)
vllm/models/deepseek_v4/nvidia/ops/o_proj.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.nn as nn
5
+
6
+from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant
7
+from vllm.platforms import current_platform
8
+from vllm.utils.deep_gemm import fp8_einsum
9
+
10
+
11
+def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]:
12
+ """fp8_einsum recipe + scale layout for the current GPU arch.
13
+
14
+ SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128.
15
+ SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1.
16
+
17
+ Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``.
18
+ """
19
+ cap = current_platform.get_device_capability()
20
+ assert cap is not None, "DeepseekV4 attention requires a CUDA device"
21
+ einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
22
+ tma_aligned_scales = cap.major >= 10
23
+ return einsum_recipe, tma_aligned_scales
24
+
25
+
26
+def deep_gemm_fp8_o_proj(
27
+ o: torch.Tensor,
28
+ positions: torch.Tensor,
29
+ cos_sin_cache: torch.Tensor,
30
+ wo_a: nn.Module,
31
+ wo_b: nn.Module,
32
+ *,
33
+ n_groups: int,
34
+ heads_per_group: int,
35
+ nope_dim: int,
36
+ rope_dim: int,
37
+ o_lora_rank: int,
38
+ einsum_recipe: tuple[int, int, int],
39
+ tma_aligned_scales: bool,
40
+) -> torch.Tensor:
41
+ """O projection: inverse RoPE + FP8 quant + einsum + wo_b.
42
+
43
+ Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
44
+ ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
45
+ """
46
+ o_fp8, o_scale = fused_inv_rope_fp8_quant(
47
+ o,
48
+ positions,
49
+ cos_sin_cache,
50
+ n_groups=n_groups,
51
+ heads_per_group=heads_per_group,
52
+ nope_dim=nope_dim,
53
+ rope_dim=rope_dim,
54
+ tma_aligned_scales=tma_aligned_scales,
55
+ )
56
+ z = torch.empty(
57
+ (o.shape[0], n_groups, o_lora_rank),
58
+ device=o.device,
59
+ dtype=torch.bfloat16,
60
+ )
61
+ fp8_einsum(
62
+ "bhr,hdr->bhd",
63
+ (o_fp8, o_scale),
64
+ (wo_a.weight, wo_a.weight_scale_inv),
65
+ z,
66
+ recipe=einsum_recipe,
67
+ )
68
+ return wo_b(z.flatten(1))
vllm/v1/attention/backends/mla/flashmla_sparse.py
+1
-1
@@ -602,7 +602,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
602
fp8_use_mixed_batch = (
603
self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.is_deepseek_v4
604
)
605
- # DeepseekV4 has its own attention impl (DeepseekV4MLAAttention) that does not
605
+ # DeepseekV4 has its own attention impl (DeepseekV4Attention) that does not
606
# consume fp8_extra_metadata. Skipping the build here avoids a
607
# forced D2H sync on seq_lens that would otherwise fire on every
608
# prefill-bearing step, lifting GPU utilization on long-prefill