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
[CI/Build] Fix test_torch_utils in AMD CI (#27317)
Signed-off-by: zhewenli <zhewenli@meta.com>
Zhewen Li committed
Oct 24, 2025 at 12:26 UTC
fc168c33f35e0610d41206e864b6bf90fe613f19
2 files changed
+35
-11
.buildkite/test-amd.yaml
+1
-1
@@ -50,7 +50,7 @@ steps:
50
51
- label: Async Engine, Inputs, Utils, Worker Test # 36min
52
timeout_in_minutes: 50
53
- mirror_hardwares: [amdexperimental]
53
+ mirror_hardwares: [amdexperimental, amdproduction]
54
agent_pool: mi325_1
55
# grade: Blocking
56
source_file_dependencies:
tests/utils_/test_torch_utils.py
+34
-10
@@ -60,15 +60,10 @@ def test_common_broadcastable_dtype(dtypes, expected_result):
60
assert common_broadcastable_dtype(dtypes) == expected_result
61
62
63
-def test_current_stream_multithread():
63
+def _test_stream_thread(main_expected_stream: torch.cuda.Stream):
64
import threading
65
66
- if not torch.cuda.is_available():
67
- pytest.skip("CUDA not available")
68
-
69
- main_default_stream = torch.cuda.current_stream()
66
child_stream = torch.cuda.Stream()
71
-
67
thread_stream_ready = threading.Event()
68
thread_can_exit = threading.Event()
69
@@ -90,15 +85,44 @@ def test_current_stream_multithread():
85
assert main_current_stream != child_stream, (
86
"Main thread's current_stream was contaminated by child thread"
87
)
93
- assert main_current_stream == main_default_stream, (
94
- "Main thread's current_stream is not the default stream"
88
+ assert main_current_stream == main_expected_stream, (
89
+ f"Main thread's stream changed unexpectedly. "
90
+ f"Expected {main_expected_stream}, got {main_current_stream}"
91
)
92
97
- # Notify child thread it can exit
93
thread_can_exit.set()
94
95
finally:
101
- # Ensure child thread exits properly
96
child_thread.join(timeout=5)
97
if child_thread.is_alive():
98
pytest.fail("Child thread failed to exit properly")
99
+
100
+
101
+def test_current_stream_multithread():
102
+ from vllm.platforms import current_platform
103
+
104
+ if not torch.cuda.is_available():
105
+ pytest.skip("CUDA not available")
106
+
107
+ if current_platform.is_rocm():
108
+ main_dedicated_stream = current_stream()
109
+
110
+ assert main_dedicated_stream.cuda_stream != 0, (
111
+ "ROCm should create a dedicated stream, not use default stream (0x0)"
112
+ )
113
+
114
+ main_stream_again = current_stream()
115
+ assert main_stream_again == main_dedicated_stream, (
116
+ "Multiple calls to current_stream should return the same dedicated stream"
117
+ )
118
+
119
+ _test_stream_thread(main_dedicated_stream)
120
+ else:
121
+ main_default_stream = torch.cuda.default_stream()
122
+ main_initial_stream = current_stream()
123
+
124
+ assert main_initial_stream == main_default_stream, (
125
+ "First call to current_stream should return default stream on CUDA"
126
+ )
127
+
128
+ _test_stream_thread(main_default_stream)