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 14, 05:30 UTC
[Bug] Fix batch invariant test `has` to `is` (#27032)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Oct 16, 2025 at 15:45 UTC
2ed8b6b3d0ccc5f577dc0d86845b63abc025468b
2 files changed
+21
-74
tests/v1/generation/test_batch_invariance.py
+10
-32
@@ -10,6 +10,11 @@ import torch
10
from vllm import LLM, SamplingParams
11
from vllm.platforms import current_platform
12
13
+hopper_only = pytest.mark.skipif(
14
+ not (current_platform.is_cuda() and current_platform.is_device_capability(90)),
15
+ reason="Requires CUDA and Hopper (SM90)",
16
+)
17
+
18
19
@pytest.fixture(autouse=True)
20
def enable_batch_invariant_mode():
@@ -66,10 +71,7 @@ def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
71
return base_prompt
72
73
69
-@pytest.mark.skipif(
70
- not current_platform.has_device_capability(90),
71
- reason="Batch invariance tests only supported on Hopper (SM90)",
72
-)
74
+@hopper_only
75
@pytest.mark.timeout(1000)
76
def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
77
"""
@@ -214,14 +216,7 @@ def _extract_step_logprobs(request_output):
216
return None, None
217
218
217
-@pytest.mark.skipif(
218
- not current_platform.has_device_capability(90),
219
- reason="Batch invariance tests only supported on Hopper (SM90)",
220
-)
221
-@pytest.mark.skipif(
222
- not torch.cuda.is_available(),
223
- reason="Requires CUDA to match production inference path.",
224
-)
219
+@hopper_only
220
@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])
221
@pytest.mark.forked
222
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):
@@ -436,10 +431,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):
431
pytest.fail(msg)
432
433
439
-@pytest.mark.skipif(
440
- not current_platform.has_device_capability(90),
441
- reason="Batch invariance tests only supported on Hopper (SM90)",
442
-)
434
+@hopper_only
435
def test_simple_generation():
436
"""
437
Simple test that runs the model with a basic prompt and prints the output.
@@ -485,14 +477,7 @@ def test_simple_generation():
477
llm.shutdown()
478
479
488
-@pytest.mark.skipif(
489
- not current_platform.has_device_capability(90),
490
- reason="Batch invariance tests only supported on Hopper (SM90)",
491
-)
492
-@pytest.mark.skipif(
493
- not torch.cuda.is_available(),
494
- reason="Requires CUDA to match production inference path.",
495
-)
480
+@hopper_only
481
@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])
482
@pytest.mark.forked
483
def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):
@@ -707,14 +692,7 @@ def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):
692
os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = old_value
693
694
710
-@pytest.mark.skipif(
711
- not current_platform.has_device_capability(90),
712
- reason="Batch invariance tests only supported on Hopper (SM90)",
713
-)
714
-@pytest.mark.skipif(
715
- not torch.cuda.is_available(),
716
- reason="Requires CUDA to match production inference path.",
717
-)
695
+@hopper_only
696
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
697
@pytest.mark.forked
698
def test_decode_logprobs_match_prefill_logprobs(backend):
tests/v1/generation/test_rms_norm_batch_invariant.py
+11
-42
@@ -14,14 +14,13 @@ from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_no
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.platforms import current_platform
16
17
-
18
-@pytest.mark.skipif(
19
- not current_platform.has_device_capability(90),
20
- reason="Batch invariance tests only supported on Hopper (SM90)",
21
-)
22
-@pytest.mark.skipif(
23
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
17
+hopper_only = pytest.mark.skipif(
18
+ not (current_platform.is_cuda() and current_platform.is_device_capability(90)),
19
+ reason="Requires CUDA and Hopper (SM90)",
20
)
21
+
22
+
23
+@hopper_only
24
@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
25
@pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
26
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@@ -70,13 +69,7 @@ def test_rms_norm_batch_invariant_vs_standard(
69
)
70
71
73
-@pytest.mark.skipif(
74
- not current_platform.has_device_capability(90),
75
- reason="Batch invariance tests only supported on Hopper (SM90)",
76
-)
77
-@pytest.mark.skipif(
78
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
79
-)
72
+@hopper_only
73
@pytest.mark.parametrize("batch_size", [1, 16, 128])
74
@pytest.mark.parametrize("seq_len", [1, 32, 512])
75
@pytest.mark.parametrize("hidden_size", [2048, 4096])
@@ -118,13 +111,7 @@ def test_rms_norm_3d_input(batch_size: int, seq_len: int, hidden_size: int):
111
)
112
113
121
-@pytest.mark.skipif(
122
- not current_platform.has_device_capability(90),
123
- reason="Batch invariance tests only supported on Hopper (SM90)",
124
-)
125
-@pytest.mark.skipif(
126
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
127
-)
114
+@hopper_only
115
def test_rms_norm_numerical_stability():
116
"""
117
Test RMS norm numerical stability with extreme values.
@@ -184,13 +171,7 @@ def test_rms_norm_numerical_stability():
171
)
172
173
187
-@pytest.mark.skipif(
188
- not current_platform.has_device_capability(90),
189
- reason="Batch invariance tests only supported on Hopper (SM90)",
190
-)
191
-@pytest.mark.skipif(
192
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
193
-)
174
+@hopper_only
175
def test_rms_norm_formula():
176
"""
177
Test that RMS norm follows the correct mathematical formula.
@@ -223,13 +204,7 @@ def test_rms_norm_formula():
204
)
205
206
226
-@pytest.mark.skipif(
227
- not current_platform.has_device_capability(90),
228
- reason="Batch invariance tests only supported on Hopper (SM90)",
229
-)
230
-@pytest.mark.skipif(
231
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
232
-)
207
+@hopper_only
208
@pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
209
def test_rms_norm_different_hidden_sizes(hidden_size: int):
210
"""
@@ -267,13 +242,7 @@ def test_rms_norm_different_hidden_sizes(hidden_size: int):
242
)
243
244
270
-@pytest.mark.skipif(
271
- not current_platform.has_device_capability(90),
272
- reason="Batch invariance tests only supported on Hopper (SM90)",
273
-)
274
-@pytest.mark.skipif(
275
- not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
276
-)
245
+@hopper_only
246
def test_rms_norm_determinism():
247
"""
248
Test that batch-invariant RMS norm produces deterministic results.