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, 07:30 UTC
[ROCm][CI] Fix MLA prefill scale for DeepSeek GSM8K (#41569)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 4, 2026 at 18:32 UTC
1cb08387214fba82c2705429100c5942fab65f05
3 files changed
+99
-4
tests/v1/attention/test_mla_backends.py
+5
-3
@@ -22,6 +22,7 @@ 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,
26
)
27
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -785,7 +786,8 @@ def test_backend_correctness(
786
assert kv_lora_rank + qk_rope_head_dim == head_size, (
787
f"MLA dimensions don't match: {total_head_size} != {head_size}"
788
)
788
- scale = 1.0 / (total_head_size**0.5)
789
+ decode_scale = 1.0 / (total_head_size**0.5)
790
+ prefill_scale = get_mla_prefill_scale(vllm_config.model_config)
791
792
# 2. Generate data and compute SDPA reference output for MLA
793
all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
@@ -902,7 +904,7 @@ def test_backend_correctness(
904
v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
905
906
sdpa_out_i_decode = torch.nn.functional.scaled_dot_product_attention(
905
- q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale
907
+ q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=decode_scale
908
)
909
sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze(
910
0
@@ -938,7 +940,7 @@ def test_backend_correctness(
940
941
# Single attention call with custom mask
942
sdpa_out_i_prefill = torch.nn.functional.scaled_dot_product_attention(
941
- q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale
943
+ q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=prefill_scale
944
)
945
sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)
946
sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)
tests/v1/attention/test_mla_prefill_selector.py
+61
@@ -2,12 +2,17 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Tests for MLA prefill backend selector."""
4
5
+from types import SimpleNamespace
6
from unittest.mock import MagicMock, patch
7
8
import pytest
9
import torch
10
11
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
+)
16
from vllm.platforms.interface import DeviceCapability
17
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
18
from vllm.v1.attention.backends.mla.prefill.selector import (
@@ -53,6 +58,62 @@ def _make_vllm_config(
58
return mock_vllm_config
59
60
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
+
117
class TestGetMLAPrefillBackend:
118
"""Tests for get_mla_prefill_backend (public API)."""
119
vllm/model_executor/layers/attention/mla_attention.py
+33
-1
@@ -238,6 +238,9 @@ 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
+)
244
from vllm.platforms import current_platform
245
from vllm.utils.flashinfer import has_flashinfer
246
from vllm.utils.math_utils import cdiv, round_down
@@ -1327,6 +1330,35 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:
1330
)
1331
1332
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
+
1362
@functools.cache
1363
def backend_supports_prefill_query_quantization() -> bool:
1364
"""Check if the selected MLA prefill backend supports query quantization.
@@ -1527,7 +1559,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1559
prefill_backend_cls = get_mla_prefill_backend(vllm_config)
1560
self._prefill_backend = prefill_backend_cls(
1561
num_heads=self.num_heads,
1530
- scale=self.model_config.get_head_size() ** -0.5,
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,