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 12, 22:15 UTC
[Bug] Fix Batch Invariant MLA test (#28967)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Nov 19, 2025 at 16:18 UTC
1607e664f0de4b7eb113c0259b889edbe73c4341
2 files changed
+33
-10
tests/v1/determinism/test_batch_invariance.py
+32
-9
@@ -9,13 +9,33 @@ import torch
9
from utils import _extract_step_logprobs, _random_prompt, skip_unsupported
10
11
from vllm import LLM, SamplingParams
12
+from vllm.platforms import current_platform
13
+
14
+BACKENDS: list[str] = [
15
+ "FLASH_ATTN",
16
+ "FLASHINFER",
17
+]
18
+
19
+if current_platform.is_cuda() and current_platform.is_device_capability(90):
20
+ BACKENDS.append("FLASH_ATTN_MLA")
21
+
22
+DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
23
+MLA_MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat"
24
+
25
+
26
+def resolve_model_name(backend: str) -> str:
27
+ """Resolve the model name for the given backend, respecting env overrides."""
28
+ model = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
29
+ if backend.endswith("MLA") and model == DEFAULT_MODEL:
30
+ return MLA_MODEL
31
+ return model
32
33
34
@skip_unsupported
35
@pytest.mark.timeout(1000)
36
@pytest.mark.parametrize(
37
"backend",
18
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
38
+ BACKENDS,
39
)
40
def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
41
backend, monkeypatch: pytest.MonkeyPatch
@@ -47,7 +67,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
67
monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
68
# Allow overrides from environment (useful for CI tuning)
69
# "facebook/opt-125m" is too small, doesn't reliably test determinism
50
- model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
70
+ model = resolve_model_name(backend)
71
num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
72
max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))
73
min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))
@@ -150,7 +170,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
170
@skip_unsupported
171
@pytest.mark.parametrize(
172
"backend",
153
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
173
+ BACKENDS,
174
)
175
@pytest.mark.forked
176
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
@@ -160,7 +180,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
180
181
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
182
random.seed(seed)
163
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
183
+ model_name = resolve_model_name(backend)
184
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
185
186
# For batch invariance, disable custom all-reduce to ensure deterministic
@@ -369,7 +389,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
389
@skip_unsupported
390
@pytest.mark.parametrize(
391
"backend",
372
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
392
+ BACKENDS,
393
)
394
def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
395
"""
@@ -377,7 +397,7 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
397
Useful for quick smoke testing and debugging.
398
"""
399
monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
380
- model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
400
+ model = resolve_model_name(backend)
401
402
llm = LLM(
403
model=model,
@@ -419,7 +439,7 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
439
@skip_unsupported
440
@pytest.mark.parametrize(
441
"backend",
422
- ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],
442
+ BACKENDS,
443
)
444
@pytest.mark.forked
445
def test_logprobs_without_batch_invariance_should_fail(
@@ -434,6 +454,9 @@ def test_logprobs_without_batch_invariance_should_fail(
454
The test will PASS if we detect differences (proving batch invariance matters).
455
The test will FAIL if everything matches (suggesting batch invariance isn't needed).
456
"""
457
+ from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
458
+
459
+ vllm_is_batch_invariant.cache_clear()
460
monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)
461
462
# CRITICAL: Disable batch invariance for this test
@@ -441,7 +464,7 @@ def test_logprobs_without_batch_invariance_should_fail(
464
465
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
466
random.seed(seed)
444
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
467
+ model_name = resolve_model_name(backend)
468
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
469
470
print(f"\n{'=' * 80}")
@@ -659,7 +682,7 @@ def test_decode_logprobs_match_prefill_logprobs(
682
683
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
684
random.seed(seed)
662
- model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
685
+ model_name = resolve_model_name(backend)
686
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
687
688
from vllm.model_executor.layers.batch_invariant import (
vllm/model_executor/layers/batch_invariant.py
+1
-1
@@ -803,11 +803,11 @@ def override_envs_for_invariance():
803
"FLASH_ATTN", # best supported backend
804
"FLASHINFER",
805
"FLASH_ATTN_MLA",
806
- "FLASHINFER_MLA",
806
"TRITON_MLA",
807
# Not yet supported MLA backends
808
# "FLASHMLA",
809
# "FLEX_ATTENTION", # IMA issue even if we disable batch invariance
810
+ # "FLASHINFER_MLA", https://github.com/vllm-project/vllm/pull/28967
811
]
812
if curr_attn_backend not in supported_backends:
813
warning = (