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 9, 17:06 UTC
[Tests] Standardize RNG seed utility across test files (#32982)
Signed-off-by: 7. Sun <jhao.sun@gmail.com>
7. Sun committed
Jan 24, 2026 at 06:47 UTC
0ccecf8833fbfb0cb4353fc66316a67c6aa6a5b5
4 files changed
+13
-20
tests/kernels/test_flex_attention.py
+3
-14
@@ -2,13 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Integration tests for FlexAttention backend vs default backend"""
4
5
-import random
6
-
7
-import numpy as np
5
import pytest
6
import torch
7
from packaging import version
8
9
+from tests.utils import set_random_seed
10
from tests.v1.attention.utils import (
11
BatchSpec,
12
create_common_attn_metadata,
@@ -27,15 +25,6 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0")
25
DIRECT_BUILD_VERSION = version.parse("2.9.dev0")
26
27
30
-def set_seed(seed):
31
- """Set seeds for reproducibility"""
32
- random.seed(seed)
33
- np.random.seed(seed)
34
- torch.manual_seed(seed)
35
- if torch.cuda.is_available():
36
- torch.cuda.manual_seed_all(seed)
37
-
38
-
28
@pytest.mark.skipif(
29
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
30
reason="CUDA not available or PyTorch version < 2.7",
@@ -57,7 +46,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
46
]
47
48
# Run with flex attention
60
- set_seed(seed)
49
+ set_random_seed(seed)
50
with vllm_runner(
51
model_name,
52
runner="generate",
@@ -71,7 +60,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
60
)
61
62
# Run with default backend
74
- set_seed(seed)
63
+ set_random_seed(seed)
64
with vllm_runner(
65
model_name,
66
runner="generate",
tests/utils.py
+4
-1
@@ -59,7 +59,10 @@ from vllm.tokenizers import get_tokenizer
59
from vllm.utils.argparse_utils import FlexibleArgumentParser
60
from vllm.utils.mem_constants import GB_bytes
61
from vllm.utils.network_utils import get_open_port
62
-from vllm.utils.torch_utils import cuda_device_count_stateless
62
+from vllm.utils.torch_utils import (
63
+ cuda_device_count_stateless,
64
+ set_random_seed, # noqa: F401 - re-exported for use in test files
65
+)
66
67
FP8_DTYPE = current_platform.fp8_dtype()
68
tests/v1/logits_processors/test_custom_offline.py
+4
-5
@@ -1,11 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import random
3
from typing import Any
4
5
import pytest
6
8
-from tests.utils import create_new_process_for_each_test
7
+from tests.utils import create_new_process_for_each_test, set_random_seed
8
from tests.v1.logits_processors.utils import (
9
DUMMY_LOGITPROC_ARG,
10
DUMMY_LOGITPROC_FQCN,
@@ -135,7 +134,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
134
135
# Test that logitproc info is passed to workers
136
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
138
- random.seed(40)
137
+ set_random_seed(40)
138
139
# Choose LLM args based on logitproc source
140
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_NONE:
@@ -194,7 +193,7 @@ def test_custom_logitsprocs_req(monkeypatch):
193
194
# Test that logitproc info is passed to workers
195
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
197
- random.seed(40)
196
+ set_random_seed(40)
197
_run_test(
198
{"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True
199
)
@@ -237,7 +236,7 @@ def test_rejects_custom_logitsprocs(
236
logitproc from
237
"""
238
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
240
- random.seed(40)
239
+ set_random_seed(40)
240
241
test_params: dict[str, dict[str, Any]] = {
242
"pooling": {
vllm/utils/torch_utils.py
+2
@@ -333,6 +333,8 @@ def set_random_seed(seed: int | None) -> None:
333
random.seed(seed)
334
np.random.seed(seed)
335
torch.manual_seed(seed)
336
+ if torch.cuda.is_available():
337
+ torch.cuda.manual_seed_all(seed)
338
339
340
def create_kv_caches_with_random_flash(