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 4, 06:15 UTC
[Attention] Minor refactor: layer takes ownership of the MLA prefill backend (#41744)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
May 5, 2026 at 19:22 UTC
01b9b5af67b77715fd1c33f19598f213ae4fec6a
7 files changed
+74
-150
tests/v1/attention/test_mla_backends.py
+16
-2
@@ -22,7 +22,6 @@ from vllm.config.vllm import set_current_vllm_config
22
from vllm.model_executor.layers.attention.mla_attention import (
23
QueryLenSupport,
24
_DecodeConcatQuantFP8,
25
- get_mla_prefill_scale,
25
)
26
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
27
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -31,6 +30,7 @@ from vllm.utils.math_utils import cdiv
30
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
31
from vllm.v1.attention.backend import CommonAttentionMetadata
32
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
33
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
34
from vllm.v1.attention.backends.registry import AttentionBackendEnum
35
from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
36
from vllm.v1.kv_cache_interface import MLAAttentionSpec
@@ -622,6 +622,19 @@ def run_attention_backend(
622
k_scale=k_scale,
623
)
624
625
+ # Attach prefill backend (normally created by MLAAttention.__init__)
626
+ prefill_scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
627
+ prefill_backend_cls = get_mla_prefill_backend(vllm_config)
628
+ mock_layer.prefill_backend = prefill_backend_cls(
629
+ num_heads=num_heads,
630
+ scale=prefill_scale,
631
+ kv_lora_rank=kv_lora_rank,
632
+ qk_nope_head_dim=qk_nope_head_dim,
633
+ qk_rope_head_dim=qk_rope_head_dim,
634
+ v_head_dim=v_head_dim,
635
+ vllm_config=vllm_config,
636
+ )
637
+
638
# Populate static_forward_context with mock attention layers
639
for layer_name in layer_names:
640
vllm_config.compilation_config.static_forward_context[layer_name] = (
@@ -787,7 +800,8 @@ def test_backend_correctness(
800
f"MLA dimensions don't match: {total_head_size} != {head_size}"
801
)
802
decode_scale = 1.0 / (total_head_size**0.5)
790
- prefill_scale = get_mla_prefill_scale(vllm_config.model_config)
803
+ qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
804
+ prefill_scale = qk_head_dim**-0.5
805
806
# 2. Generate data and compute SDPA reference output for MLA
807
all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
tests/v1/attention/test_mla_prefill_selector.py
-61
@@ -2,17 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Tests for MLA prefill backend selector."""
4
5
-from types import SimpleNamespace
5
from unittest.mock import MagicMock, patch
6
7
import pytest
8
import torch
9
10
from vllm.config import AttentionConfig, ModelConfig, VllmConfig
12
-from vllm.model_executor.layers.attention.mla_attention import get_mla_prefill_scale
13
-from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
14
- yarn_get_mscale,
15
-)
11
from vllm.platforms.interface import DeviceCapability
12
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
13
from vllm.v1.attention.backends.mla.prefill.selector import (
@@ -58,62 +53,6 @@ def _make_vllm_config(
53
return mock_vllm_config
54
55
61
-class TestMLAPrefillScale:
62
- """Tests for the MLA prefill softmax scale."""
63
-
64
- def test_uses_qk_head_dim_for_deepseek_v2_style_mla(self):
65
- model_config = SimpleNamespace(
66
- hf_text_config=SimpleNamespace(
67
- q_lora_rank=None,
68
- kv_lora_rank=512,
69
- qk_nope_head_dim=128,
70
- qk_rope_head_dim=64,
71
- v_head_dim=128,
72
- rope_parameters={"rope_type": "default"},
73
- )
74
- )
75
-
76
- assert get_mla_prefill_scale(model_config) == pytest.approx(192**-0.5)
77
-
78
- def test_applies_deepseek_yarn_mscale(self):
79
- model_config = SimpleNamespace(
80
- hf_text_config=SimpleNamespace(
81
- q_lora_rank=None,
82
- kv_lora_rank=512,
83
- qk_nope_head_dim=128,
84
- qk_rope_head_dim=64,
85
- v_head_dim=128,
86
- rope_parameters={
87
- "rope_type": "yarn",
88
- "factor": 40,
89
- "mscale_all_dim": 0.707,
90
- },
91
- )
92
- )
93
-
94
- mscale = yarn_get_mscale(40, 0.707)
95
- assert get_mla_prefill_scale(model_config) == pytest.approx(
96
- 192**-0.5 * mscale * mscale
97
- )
98
-
99
- def test_deepseek_v4_style_mla_does_not_apply_yarn_mscale(self):
100
- model_config = SimpleNamespace(
101
- hf_text_config=SimpleNamespace(
102
- compress_ratios=[4],
103
- q_lora_rank=1536,
104
- head_dim=128,
105
- qk_rope_head_dim=64,
106
- rope_parameters={
107
- "rope_type": "yarn",
108
- "factor": 40,
109
- "mscale_all_dim": 0.707,
110
- },
111
- )
112
- )
113
-
114
- assert get_mla_prefill_scale(model_config) == pytest.approx(128**-0.5)
115
-
116
-
56
class TestGetMLAPrefillBackend:
57
"""Tests for get_mla_prefill_backend (public API)."""
58
vllm/model_executor/layers/attention/mla_attention.py
+24
-52
@@ -238,9 +238,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
238
kFp8StaticTensorSym,
239
kNvfp4Dynamic,
240
)
241
-from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
242
- yarn_get_mscale,
243
-)
241
from vllm.platforms import current_platform
242
from vllm.utils.flashinfer import has_flashinfer
243
from vllm.utils.math_utils import cdiv, round_down
@@ -262,7 +259,10 @@ from vllm.v1.attention.backend import (
259
MLAAttentionImpl,
260
SparseMLAAttentionImpl,
261
)
265
-from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackend
262
+from vllm.v1.attention.backends.mla.prefill import (
263
+ MLAPrefillBackend,
264
+ get_mla_prefill_backend,
265
+)
266
from vllm.v1.attention.backends.utils import (
267
get_dcp_local_seq_lens,
268
split_decodes_and_prefills,
@@ -454,20 +454,32 @@ class MLAAttention(nn.Module, AttentionLayerBase):
454
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
455
self.use_direct_call = not current_platform.opaque_attention_op()
456
457
- compilation_config = get_current_vllm_config().compilation_config
457
+ vllm_config = get_current_vllm_config()
458
+ compilation_config = vllm_config.compilation_config
459
if prefix in compilation_config.static_forward_context:
460
raise ValueError(f"Duplicate layer name: {prefix}")
461
compilation_config.static_forward_context[prefix] = self
462
463
+ prefill_backend_cls = get_mla_prefill_backend(vllm_config)
464
+ self.prefill_backend = prefill_backend_cls(
465
+ num_heads=self.num_heads,
466
+ scale=self.scale,
467
+ kv_lora_rank=self.kv_lora_rank,
468
+ qk_nope_head_dim=self.qk_nope_head_dim,
469
+ qk_rope_head_dim=self.qk_rope_head_dim,
470
+ v_head_dim=self.v_head_dim,
471
+ vllm_config=vllm_config,
472
+ )
473
+
474
self.kv_cache = torch.tensor([])
475
476
self.use_sparse = use_sparse
477
466
- vllm_config = get_current_vllm_config_or_none()
478
+ _vllm_config = get_current_vllm_config_or_none()
479
self.dcp_a2a = (
468
- vllm_config is not None
469
- and vllm_config.parallel_config.decode_context_parallel_size > 1
470
- and vllm_config.parallel_config.dcp_comm_backend == "a2a"
480
+ _vllm_config is not None
481
+ and _vllm_config.parallel_config.decode_context_parallel_size > 1
482
+ and _vllm_config.parallel_config.dcp_comm_backend == "a2a"
483
)
484
485
# Initialize q/k/v range constants.
@@ -1330,35 +1342,6 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:
1342
)
1343
1344
1333
-def get_mla_prefill_scale(model_config: ModelConfig) -> float:
1334
- hf_text_config = model_config.hf_text_config
1335
- mla_dims = get_mla_dims(model_config)
1336
- qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
1337
- scale = qk_head_dim**-0.5
1338
-
1339
- # Deepseek V4 disables YaRN mscale for attention; Deepseek V2/V3 applies
1340
- # the same mscale correction when constructing the MLA attention module.
1341
- if hasattr(hf_text_config, "compress_ratios"):
1342
- return scale
1343
-
1344
- rope_parameters = getattr(hf_text_config, "rope_parameters", None)
1345
- if rope_parameters is None:
1346
- rope_parameters = getattr(hf_text_config, "rope_scaling", None)
1347
-
1348
- if rope_parameters is None:
1349
- return scale
1350
-
1351
- rope_type = rope_parameters.get("rope_type", rope_parameters.get("type"))
1352
- apply_yarn_scaling = rope_parameters.get("apply_yarn_scaling", True)
1353
- if rope_type != "default" and apply_yarn_scaling:
1354
- mscale_all_dim = rope_parameters.get("mscale_all_dim", False)
1355
- scaling_factor = rope_parameters["factor"]
1356
- mscale = yarn_get_mscale(float(scaling_factor), float(mscale_all_dim))
1357
- scale *= mscale * mscale
1358
-
1359
- return scale
1360
-
1361
-
1345
@functools.cache
1346
def backend_supports_prefill_query_quantization() -> bool:
1347
"""Check if the selected MLA prefill backend supports query quantization.
@@ -1554,20 +1537,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1537
device=device,
1538
)
1539
1557
- from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
1558
-
1559
- prefill_backend_cls = get_mla_prefill_backend(vllm_config)
1560
- self._prefill_backend = prefill_backend_cls(
1561
- num_heads=self.num_heads,
1562
- scale=get_mla_prefill_scale(self.model_config),
1563
- kv_lora_rank=self.mla_dims.kv_lora_rank,
1564
- qk_nope_head_dim=self.mla_dims.qk_nope_head_dim,
1565
- qk_rope_head_dim=self.mla_dims.qk_rope_head_dim,
1566
- v_head_dim=self.mla_dims.v_head_dim,
1567
- vllm_config=vllm_config,
1568
- device=device,
1569
- layer_names=layer_names,
1570
- )
1540
+ self._prefill_backend = self.compilation_config.static_forward_context[
1541
+ layer_names[0]
1542
+ ].prefill_backend
1543
1544
supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
1545
self._init_reorder_batch_threshold(
vllm/v1/attention/backends/mla/prefill/base.py
-4
@@ -81,8 +81,6 @@ class MLAPrefillBackend(ABC):
81
qk_rope_head_dim: int,
82
v_head_dim: int,
83
vllm_config: "VllmConfig",
84
- device: torch.device,
85
- layer_names: list[str] | None = None,
84
) -> None:
85
self.num_heads = num_heads
86
self.scale = scale
@@ -91,8 +89,6 @@ class MLAPrefillBackend(ABC):
89
self.qk_rope_head_dim = qk_rope_head_dim
90
self.v_head_dim = v_head_dim
91
self.vllm_config = vllm_config
94
- self.device = device
95
- self.layer_names = layer_names
92
93
def prepare_metadata( # noqa: B027
94
self,
vllm/v1/attention/backends/mla/prefill/flash_attn.py
-4
@@ -44,8 +44,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
44
qk_rope_head_dim: int,
45
v_head_dim: int,
46
vllm_config: "VllmConfig",
47
- device: torch.device,
48
- layer_names: list[str] | None = None,
47
) -> None:
48
super().__init__(
49
num_heads=num_heads,
@@ -55,8 +53,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
53
qk_rope_head_dim=qk_rope_head_dim,
54
v_head_dim=v_head_dim,
55
vllm_config=vllm_config,
58
- device=device,
59
- layer_names=layer_names,
56
)
57
58
# Handle the differences between the flash_attn_varlen from
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+34
-23
@@ -9,6 +9,7 @@ import torch
9
import vllm.envs as envs
10
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
11
from vllm.v1.attention.backends.utils import (
12
+ PerLayerParameters,
13
get_per_layer_parameters,
14
infer_global_hyperparameters,
15
)
@@ -62,8 +63,6 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
63
qk_rope_head_dim: int,
64
v_head_dim: int,
65
vllm_config: "VllmConfig",
65
- device: torch.device,
66
- layer_names: list[str] | None = None,
66
) -> None:
67
super().__init__(
68
num_heads=num_heads,
@@ -73,25 +72,11 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
72
qk_rope_head_dim=qk_rope_head_dim,
73
v_head_dim=v_head_dim,
74
vllm_config=vllm_config,
76
- device=device,
77
- layer_names=layer_names,
75
)
76
77
self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
78
self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
82
- if layer_names is None:
83
- raise ValueError(
84
- "FlashInferPrefillBackend requires layer_names to "
85
- "initialize global hyperparameters."
86
- )
87
-
88
- from vllm.model_executor.layers.attention.mla_attention import (
89
- MLACommonImpl,
90
- )
91
-
92
- self._global_hyperparameters = infer_global_hyperparameters(
93
- get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]
94
- )
79
+ self._global_hyperparameters: PerLayerParameters | None = None
80
81
def _ensure_chunks(
82
self,
@@ -106,10 +91,36 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
91
)
92
)
93
94
+ def _resolve_global_hyperparameters(self) -> PerLayerParameters:
95
+ if self._global_hyperparameters is not None:
96
+ return self._global_hyperparameters
97
+
98
+ from vllm.model_executor.layers.attention.mla_attention import (
99
+ MLAAttention,
100
+ MLACommonImpl,
101
+ )
102
+
103
+ forward_context = self.vllm_config.compilation_config.static_forward_context
104
+ layer_names = [
105
+ name
106
+ for name, layer in forward_context.items()
107
+ if isinstance(layer, MLAAttention)
108
+ ]
109
+
110
+ self._global_hyperparameters = infer_global_hyperparameters(
111
+ get_per_layer_parameters(
112
+ self.vllm_config,
113
+ layer_names,
114
+ MLACommonImpl, # type: ignore[type-abstract]
115
+ )
116
+ )
117
+ return self._global_hyperparameters
118
+
119
def prepare_metadata(
120
self,
121
prefill_metadata: "MLACommonPrefillMetadata",
122
) -> None:
123
+ global_hyperparameters = self._resolve_global_hyperparameters()
124
qo_indptr = prefill_metadata.query_start_loc
125
has_context = prefill_metadata.chunked_context is not None
126
(workspace_buffer,) = current_workspace_manager().get_simultaneous(
@@ -144,9 +155,9 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
155
head_dim_qk=head_dim_qk,
156
head_dim_vo=head_dim_vo,
157
causal=True,
147
- sm_scale=self._global_hyperparameters.sm_scale,
148
- window_left=self._global_hyperparameters.window_left,
149
- logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
158
+ sm_scale=global_hyperparameters.sm_scale,
159
+ window_left=global_hyperparameters.window_left,
160
+ logits_soft_cap=global_hyperparameters.logits_soft_cap,
161
q_data_type=prefill_metadata.q_data_type,
162
o_data_type=prefill_metadata.output_dtype,
163
)
@@ -165,9 +176,9 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
176
head_dim_qk=head_dim_qk,
177
head_dim_vo=head_dim_vo,
178
causal=False,
168
- sm_scale=self._global_hyperparameters.sm_scale,
169
- window_left=self._global_hyperparameters.window_left,
170
- logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
179
+ sm_scale=global_hyperparameters.sm_scale,
180
+ window_left=global_hyperparameters.window_left,
181
+ logits_soft_cap=global_hyperparameters.logits_soft_cap,
182
q_data_type=prefill_metadata.q_data_type,
183
o_data_type=prefill_metadata.output_dtype,
184
)
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
-4
@@ -51,8 +51,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
51
qk_rope_head_dim: int,
52
v_head_dim: int,
53
vllm_config: "VllmConfig",
54
- device: torch.device,
55
- layer_names: list[str] | None = None,
54
) -> None:
55
super().__init__(
56
num_heads=num_heads,
@@ -62,8 +60,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
60
qk_rope_head_dim=qk_rope_head_dim,
61
v_head_dim=v_head_dim,
62
vllm_config=vllm_config,
65
- device=device,
66
- layer_names=layer_names,
63
)
64
65
def _get_workspace_buffer(self) -> torch.Tensor: