Skip to content

Commit 2a75e9e

Browse files
devalshahamddevalshahamd
authored
Merge branch 'main' into main
2 parents 67f3143 + e1e4646 commit 2a75e9e

17 files changed

Lines changed: 542 additions & 189 deletions

File tree

‎docs/design/attention_backends.md‎

Lines changed: 30 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ Priority is **1 = highest** (tried first).
155155
| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
156156
| **Head Sizes** | Supported attention head sizes |
157157
| **Sink** | Attention sink support (for StreamingLLM) |
158+
| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models |
158159
| **Sparse** | Sparse attention support (MLA only) |
159160
| **MM Prefix** | Multimodal prefix full attention support |
160161
| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
@@ -165,22 +166,22 @@ Priority is **1 = highest** (tried first).
165166

166167
## Standard Attention (MHA, MQA, GQA) Backends
167168

168-
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
169-
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
170-
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
171-
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
172-
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
173-
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
174-
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
175-
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
176-
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
177-
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
178-
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
179-
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
180-
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
181-
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
182-
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
183-
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
169+
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
170+
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
171+
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
172+
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
173+
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
174+
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
175+
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
176+
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
177+
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
178+
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
179+
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
180+
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
181+
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
182+
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
183+
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
184+
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
184185

185186
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
186187
>
@@ -211,16 +212,16 @@ hardware and configuration.
211212
MLA decode backends are selected using the standard
212213
`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).
213214

214-
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
215-
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
216-
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
217-
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
218-
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
219-
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
220-
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
221-
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
222-
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
223-
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
224-
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
225-
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
226-
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
215+
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
216+
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
217+
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
218+
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
219+
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
220+
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
221+
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
222+
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
223+
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
224+
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
225+
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
226+
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
227+
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |

‎docs/serving/data_parallel_deployment.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ For larger scale deployments especially, it can make sense to handle the orchest
9898

9999
In this case, it's more convenient to treat each DP rank like a separate vLLM deployment, with its own endpoint, and have an external router balance HTTP requests between them, making use of appropriate real-time telemetry from each server for routing decisions.
100100

101-
This can already be done trivially for non-MoE models, since each deployed server is fully independent. No data parallel CLI options need to be used for this.
101+
This can already be done trivially for non-MoE models, since each deployed server is fully independent. In that case, launch independent vLLM instances without any `--data-parallel-*` arguments; external DP CLI options are only supported for MoE deployments.
102102

103103
We support an equivalent topology for MoE DP+EP which can be configured via the following CLI arguments.
104104

‎tests/quantization/test_turboquant.py‎

Lines changed: 82 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -182,22 +182,100 @@ def test_all_presets_all_head_dims(self, preset, head_dim):
182182

183183
# ---- Boundary skip layers ----
184184

185+
@staticmethod
186+
def _dense_model_config(num_layers):
187+
from types import SimpleNamespace
188+
189+
return SimpleNamespace(
190+
is_hybrid=False,
191+
hf_text_config=SimpleNamespace(num_hidden_layers=num_layers),
192+
)
193+
185194
def test_boundary_skip_layers_basic(self):
186-
layers = TurboQuantConfig.get_boundary_skip_layers(32)
195+
mc = self._dense_model_config(32)
196+
layers = TurboQuantConfig.get_boundary_skip_layers(mc)
187197
assert layers == ["0", "1", "30", "31"]
188198

189199
def test_boundary_skip_layers_zero(self):
190-
assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == []
200+
mc = self._dense_model_config(32)
201+
assert TurboQuantConfig.get_boundary_skip_layers(mc, 0) == []
191202

192203
def test_boundary_skip_layers_small_model(self):
193-
layers = TurboQuantConfig.get_boundary_skip_layers(4)
204+
mc = self._dense_model_config(4)
205+
layers = TurboQuantConfig.get_boundary_skip_layers(mc)
194206
assert layers == ["0", "1", "2", "3"]
195207

196208
def test_boundary_skip_layers_cap_at_half(self):
197-
layers = TurboQuantConfig.get_boundary_skip_layers(8, 10)
209+
mc = self._dense_model_config(8)
210+
layers = TurboQuantConfig.get_boundary_skip_layers(mc, 10)
198211
assert len(layers) == 8
199212

200213

214+
class TestHybridAttentionIndices:
215+
"""Regression tests for boundary protection on hybrid models.
216+
217+
Hybrid models (attention + Mamba / linear-attention) identify KV-carrying
218+
layers via layer_types / layers_block_type / attn_type_list. The helper
219+
must return the *global* layer indices of the full-attention layers so
220+
that kv_cache_dtype_skip_layers matches what extract_layer_index(prefix)
221+
reports on the Attention layers at runtime.
222+
"""
223+
224+
@staticmethod
225+
def _fake_model_config(text_cfg=None, hf_cfg=None):
226+
from types import SimpleNamespace
227+
228+
return SimpleNamespace(
229+
hf_text_config=text_cfg if text_cfg is not None else SimpleNamespace(),
230+
hf_config=hf_cfg if hf_cfg is not None else SimpleNamespace(),
231+
)
232+
233+
def test_layer_types_full_attention(self):
234+
from vllm.model_executor.layers.quantization.turboquant.config import (
235+
_get_full_attention_layer_indices,
236+
)
237+
238+
cfg = type("C", (), {})()
239+
cfg.layer_types = [
240+
"linear_attention",
241+
"linear_attention",
242+
"full_attention",
243+
"linear_attention",
244+
"full_attention",
245+
"full_attention",
246+
]
247+
mc = self._fake_model_config(text_cfg=cfg)
248+
assert _get_full_attention_layer_indices(mc) == [2, 4, 5]
249+
250+
def test_layers_block_type_jamba(self):
251+
from vllm.model_executor.layers.quantization.turboquant.config import (
252+
_get_full_attention_layer_indices,
253+
)
254+
255+
cfg = type("C", (), {})()
256+
cfg.layers_block_type = ["mamba", "attention", "mamba", "attention"]
257+
mc = self._fake_model_config(text_cfg=cfg)
258+
assert _get_full_attention_layer_indices(mc) == [1, 3]
259+
260+
def test_attn_type_list_minimax(self):
261+
from vllm.model_executor.layers.quantization.turboquant.config import (
262+
_get_full_attention_layer_indices,
263+
)
264+
265+
hf = type("C", (), {})()
266+
hf.attn_type_list = [0, 1, 0, 1, 1]
267+
mc = self._fake_model_config(hf_cfg=hf)
268+
assert _get_full_attention_layer_indices(mc) == [1, 3, 4]
269+
270+
def test_no_hybrid_hints_returns_empty(self):
271+
from vllm.model_executor.layers.quantization.turboquant.config import (
272+
_get_full_attention_layer_indices,
273+
)
274+
275+
mc = self._fake_model_config()
276+
assert _get_full_attention_layer_indices(mc) == []
277+
278+
201279
# ============================================================================
202280
# Centroids tests (CPU-only)
203281
# ============================================================================

‎tests/test_config.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1215,8 +1215,6 @@ def test_scheduler_config_init():
12151215
("facebook/opt-125m", 1, False, False),
12161216
# Non-MoE model with DP>1 internal LB should need coordinator
12171217
("facebook/opt-125m", 2, False, True),
1218-
# Non-MoE model with DP>1 external LB should not need coordinator
1219-
("facebook/opt-125m", 2, True, False),
12201218
# MoE model with DP=1 should not need coordinator
12211219
("mistralai/Mixtral-8x7B-Instruct-v0.1", 1, False, False),
12221220
# MoE model with DP>1 internal LB should need both coordinator

‎tests/v1/attention/test_mla_backends.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from vllm.model_executor.layers.attention.mla_attention import (
2323
QueryLenSupport,
2424
_DecodeConcatQuantFP8,
25+
get_mla_prefill_scale,
2526
)
2627
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
2728
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -785,7 +786,8 @@ def test_backend_correctness(
785786
assert kv_lora_rank + qk_rope_head_dim == head_size, (
786787
f"MLA dimensions don't match: {total_head_size} != {head_size}"
787788
)
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)
789791

790792
# 2. Generate data and compute SDPA reference output for MLA
791793
all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
@@ -902,7 +904,7 @@ def test_backend_correctness(
902904
v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
903905

904906
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
906908
)
907909
sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze(
908910
0
@@ -938,7 +940,7 @@ def test_backend_correctness(
938940

939941
# Single attention call with custom mask
940942
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
942944
)
943945
sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)
944946
sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)

‎tests/v1/attention/test_mla_prefill_selector.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,17 @@
22
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
33
"""Tests for MLA prefill backend selector."""
44

5+
from types import SimpleNamespace
56
from unittest.mock import MagicMock, patch
67

78
import pytest
89
import torch
910

1011
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+
)
1116
from vllm.platforms.interface import DeviceCapability
1217
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
1318
from vllm.v1.attention.backends.mla.prefill.selector import (
@@ -53,6 +58,62 @@ def _make_vllm_config(
5358
return mock_vllm_config
5459

5560

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+
56117
class TestGetMLAPrefillBackend:
57118
"""Tests for get_mla_prefill_backend (public API)."""
58119

‎tests/v1/distributed/test_external_lb_dp.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from tests.utils import RemoteOpenAIServer
1515
from vllm.platforms import current_platform
1616

17-
MODEL_NAME = "ibm-research/PowerMoE-3b"
17+
MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b")
1818

1919
# Number of data parallel ranks for external LB testing
2020
DP_SIZE = int(os.getenv("DP_SIZE", "2"))

0 commit comments

Comments
 (0)