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 6, 05:45 UTC
[Test] Only Run MLA model when user explicitly set for batch invariance (#37719)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 22, 2026 at 09:09 UTC
eaf4978621acdff42be80ff7eb0b2818c876ecca
3 files changed
+23
-27
tests/v1/determinism/test_batch_invariance.py
+6
-9
@@ -8,10 +8,10 @@ import pytest
8
import torch
9
from utils import (
10
BACKENDS,
11
+ TEST_MODEL,
12
_extract_step_logprobs,
13
_random_prompt,
14
is_device_capability_below_90,
14
- resolve_model_name,
15
skip_unsupported,
16
)
17
@@ -57,7 +57,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
57
attention_config = {"backend": backend}
58
# Allow overrides from environment (useful for CI tuning)
59
# "facebook/opt-125m" is too small, doesn't reliably test determinism
60
- model = resolve_model_name(backend)
60
+ model = TEST_MODEL
61
num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
62
max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))
63
min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))
@@ -169,7 +169,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
169
):
170
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
171
random.seed(seed)
172
- model_name = resolve_model_name(backend)
172
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
173
174
# For batch invariance, disable custom all-reduce to ensure deterministic
@@ -186,7 +185,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
185
print(f"{'=' * 80}\n")
186
187
llm = LLM(
189
- model=model_name,
188
+ model=TEST_MODEL,
189
tensor_parallel_size=tp_size,
190
max_num_seqs=128,
191
max_model_len=8192,
@@ -395,7 +394,7 @@ def test_simple_generation(backend):
394
Simple test that runs the model with a basic prompt and prints the output.
395
Useful for quick smoke testing and debugging.
396
"""
398
- model = resolve_model_name(backend)
397
+ model = TEST_MODEL
398
399
llm = LLM(
400
model=model,
@@ -458,7 +457,6 @@ def test_logprobs_without_batch_invariance_should_fail(
457
monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)
458
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
459
random.seed(seed)
461
- model_name = resolve_model_name(backend)
460
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
461
462
print(f"\n{'=' * 80}")
@@ -466,7 +464,7 @@ def test_logprobs_without_batch_invariance_should_fail(
464
print(f"{'=' * 80}\n")
465
466
llm = LLM(
469
- model=model_name,
467
+ model=TEST_MODEL,
468
tensor_parallel_size=tp_size,
469
max_num_seqs=32,
470
max_model_len=8192,
@@ -674,7 +672,6 @@ def test_decode_logprobs_match_prefill_logprobs(
672
"""
673
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
674
random.seed(seed)
677
- model_name = resolve_model_name(backend)
675
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
676
677
from vllm.model_executor.layers.batch_invariant import (
@@ -689,7 +686,7 @@ def test_decode_logprobs_match_prefill_logprobs(
686
print(f"{'=' * 80}\n")
687
688
llm = LLM(
692
- model=model_name,
689
+ model=TEST_MODEL,
690
tensor_parallel_size=tp_size,
691
max_num_seqs=32,
692
max_model_len=8192,
tests/v1/determinism/test_online_batch_invariance.py
+3
-4
@@ -17,7 +17,7 @@ from typing import Any
17
18
import openai
19
import pytest
20
-from utils import BACKENDS, _random_prompt, resolve_model_name, skip_unsupported
20
+from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_unsupported
21
22
from tests.utils import RemoteOpenAIServer
23
@@ -139,7 +139,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
139
backend: str,
140
) -> None:
141
random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
142
- model_name = resolve_model_name(backend)
142
prompts_all = [_random_prompt(10, 50) for _ in range(32)]
143
144
sp_kwargs: dict[str, Any] = {
@@ -159,11 +158,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
158
if tp_size:
159
server_args += ["-tp", tp_size]
160
162
- with RemoteOpenAIServer(model_name, server_args) as server:
161
+ with RemoteOpenAIServer(TEST_MODEL, server_args) as server:
162
client = server.get_client()
163
_compare_bs1_vs_bsn_single_process(
164
prompts=prompts_all,
165
sp_kwargs=sp_kwargs,
166
client=client,
168
- model_name=model_name,
167
+ model_name=TEST_MODEL,
168
)
tests/v1/determinism/utils.py
+14
-14
@@ -7,6 +7,10 @@ import pytest
7
import torch
8
9
from vllm.platforms import current_platform
10
+from vllm.transformers_utils.config import get_config
11
+from vllm.transformers_utils.model_arch_config_convertor import (
12
+ ModelArchConfigConvertorBase,
13
+)
14
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
15
16
skip_unsupported = pytest.mark.skipif(
@@ -16,10 +20,12 @@ skip_unsupported = pytest.mark.skipif(
20
reason="Requires CUDA and >= Ampere (SM80)",
21
)
22
23
+DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
24
+TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
25
+
26
BACKENDS: list[str] = [
27
"FLASH_ATTN",
28
"TRITON_ATTN",
22
- "TRITON_MLA",
29
]
30
31
# FlashInfer temporarily disabled due to invariant CTA sizes.
@@ -27,19 +33,13 @@ BACKENDS: list[str] = [
33
# if has_flashinfer():
34
# BACKENDS.append("FLASHINFER")
35
30
-if flash_attn_supports_mla():
31
- BACKENDS.append("FLASH_ATTN_MLA")
32
-
33
-DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
34
-MLA_MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat"
35
-
36
-
37
-def resolve_model_name(backend: str) -> str:
38
- """Resolve the model name for the given backend."""
39
- model = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
40
- if backend.endswith("MLA") and model == DEFAULT_MODEL:
41
- return MLA_MODEL
42
- return model
36
+# only run MLA backends when the requested test model is itself an MLA model.
37
+if os.getenv("VLLM_TEST_MODEL"):
38
+ config = get_config(TEST_MODEL, trust_remote_code=False)
39
+ if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla():
40
+ BACKENDS = ["TRITON_MLA"]
41
+ if flash_attn_supports_mla():
42
+ BACKENDS.append("FLASH_ATTN_MLA")
43
44
45
def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: