A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][CI] Defer AITER sampler import and isolate server test PYTHONPATH (#44823)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 10, 2026 at 03:56 UTC
82a42234be76edb07b509006de7809475d878b71
3 files changed
+117
-17
tests/utils.py
+47
-3
@@ -18,7 +18,7 @@ import tempfile
18
import threading
19
import time
20
import warnings
21
-from collections.abc import Callable, Iterable, Sequence
21
+from collections.abc import Callable, Iterable, MutableMapping, Sequence
22
from contextlib import ExitStack, contextmanager
23
from multiprocessing import Process, get_context
24
from pathlib import Path
@@ -149,6 +149,46 @@ ROCM_ENGINE_KWARGS: dict = (
149
if current_platform.is_rocm()
150
else {}
151
)
152
+_TILELANG_TVM_PYTHONPATH_FRAGMENT = os.path.join(
153
+ "tilelang", "3rdparty", "tvm", "python"
154
+)
155
+
156
+
157
+def _sanitize_pythonpath_value(pythonpath: str | None) -> str:
158
+ if not pythonpath:
159
+ return ""
160
+ entries = []
161
+ for entry in pythonpath.split(os.pathsep):
162
+ normalized = entry.replace(os.sep, "/")
163
+ if _TILELANG_TVM_PYTHONPATH_FRAGMENT.replace(os.sep, "/") in normalized:
164
+ continue
165
+ entries.append(entry)
166
+ return os.pathsep.join(entries)
167
+
168
+
169
+def _sanitize_pythonpath_env(env: MutableMapping[str, str]) -> None:
170
+ cleaned = _sanitize_pythonpath_value(env.get("PYTHONPATH"))
171
+ if cleaned:
172
+ env["PYTHONPATH"] = cleaned
173
+ else:
174
+ env.pop("PYTHONPATH", None)
175
+
176
+
177
+def _sanitize_current_pythonpath_env() -> None:
178
+ _sanitize_pythonpath_env(os.environ)
179
+
180
+
181
+@contextmanager
182
+def _temporarily_sanitized_pythonpath_env():
183
+ original = os.environ.get("PYTHONPATH")
184
+ _sanitize_current_pythonpath_env()
185
+ try:
186
+ yield
187
+ finally:
188
+ if original is None:
189
+ os.environ.pop("PYTHONPATH", None)
190
+ else:
191
+ os.environ["PYTHONPATH"] = original
192
193
194
def requires_spawn_multiprocessing() -> bool:
@@ -253,7 +293,8 @@ class RemoteVLLMServer:
293
getattr(args, "show_hidden_metrics_for_version", None) is not None
294
)
295
256
- self._pre_download_model(model, args)
296
+ with _temporarily_sanitized_pythonpath_env():
297
+ self._pre_download_model(model, args)
298
self._shutdown_complete = False
299
300
# Record GPU memory before server start so we know what
@@ -727,6 +768,7 @@ class RemoteOpenAIServer(RemoteVLLMServer):
768
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
769
if env_dict is not None:
770
env.update(env_dict)
771
+ _sanitize_pythonpath_env(env)
772
serve_cmd = ["vllm", "serve", model, *vllm_serve_args]
773
print(f"Launching RemoteOpenAIServer with: {' '.join(serve_cmd)}")
774
print(f"Environment variables: {env}")
@@ -754,6 +796,7 @@ class RemoteLaunchRenderServer(RemoteVLLMServer):
796
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
797
if env_dict is not None:
798
env.update(env_dict)
799
+ _sanitize_pythonpath_env(env)
800
serve_cmd = ["vllm", "launch", "render", model, *vllm_serve_args]
801
print(f"Launching RemoteLaunchRenderServer with: {' '.join(serve_cmd)}")
802
self.proc: subprocess.Popen = subprocess.Popen(
@@ -795,7 +838,8 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
838
target=_run_in_new_process_group,
839
args=(self.child_process_fxn, env_dict, model, vllm_serve_args),
840
) # type: ignore[assignment]
798
- self.proc.start()
841
+ with _temporarily_sanitized_pythonpath_env():
842
+ self.proc.start()
843
844
def __init__(
845
self,
tests/v1/sample/test_topk_topp_sampler.py
+45
@@ -64,6 +64,51 @@ def test_sampler_threads_fp64_gumbel_to_topk_topp_sampler():
64
assert sampler.topk_topp_sampler.use_fp64_gumbel
65
66
67
+def test_rocm_aiter_sampler_defers_import_when_generators_force_native(
68
+ monkeypatch: pytest.MonkeyPatch,
69
+):
70
+ from vllm.v1.sample.ops import topk_topp_sampler
71
+
72
+ class MockPlatform:
73
+ @staticmethod
74
+ def is_cuda():
75
+ return False
76
+
77
+ @staticmethod
78
+ def is_cpu():
79
+ return False
80
+
81
+ @staticmethod
82
+ def is_xpu():
83
+ return False
84
+
85
+ class MockRocmAiterOps:
86
+ @staticmethod
87
+ def is_enabled():
88
+ return True
89
+
90
+ real_import = __import__
91
+
92
+ def guard_aiter_sampling_import(name, *args, **kwargs):
93
+ if name == "aiter.ops.sampling":
94
+ raise AssertionError("aiter sampling import should be deferred")
95
+ return real_import(name, *args, **kwargs)
96
+
97
+ monkeypatch.setattr(topk_topp_sampler, "current_platform", MockPlatform())
98
+ monkeypatch.setattr(topk_topp_sampler, "rocm_aiter_ops", MockRocmAiterOps())
99
+ monkeypatch.setattr("builtins.__import__", guard_aiter_sampling_import)
100
+
101
+ sampler = topk_topp_sampler.TopKTopPSampler()
102
+ logits = torch.randn(2, 8)
103
+ k = torch.full((2,), 2, dtype=torch.int32)
104
+ generators = {0: torch.Generator(device=logits.device).manual_seed(0)}
105
+
106
+ token_ids, logits_to_return = sampler(logits, generators, k, None)
107
+
108
+ assert token_ids.shape == (2,)
109
+ assert logits_to_return is None
110
+
111
+
112
def test_random_sample_uses_fp64_exponential_race_when_requested():
113
torch.set_default_device(DEVICE_TYPE)
114
probs = torch.tensor(
vllm/v1/sample/ops/topk_topp_sampler.py
+25
-14
@@ -111,20 +111,12 @@ class TopKTopPSampler(nn.Module):
111
logprobs_mode not in ("processed_logits", "processed_logprobs")
112
and rocm_aiter_ops.is_enabled()
113
):
114
- try:
115
- import aiter.ops.sampling # noqa: F401
116
-
117
- self.aiter_ops = torch.ops.aiter
118
- logger.info_once(
119
- "Using aiter sampler on ROCm (lazy import, sampling-only)."
120
- )
121
- self.forward = self.forward_hip
122
- except ImportError:
123
- logger.warning_once(
124
- "aiter.ops.sampling is not available on ROCm. "
125
- "Falling back to forward_native implementation."
126
- )
127
- self.forward = self.forward_native
114
+ self.aiter_ops = None
115
+ self._aiter_ops_import_failed = False
116
+ logger.info_once(
117
+ "Using aiter sampler on ROCm (lazy import, sampling-only)."
118
+ )
119
+ self.forward = self.forward_hip
120
else:
121
self.forward = self.forward_native
122
@@ -211,6 +203,22 @@ class TopKTopPSampler(nn.Module):
203
204
return sample_with_exponential_noise(probs, q), logits_to_return
205
206
+ def _init_aiter_ops(self) -> bool:
207
+ if self._aiter_ops_import_failed:
208
+ return False
209
+ try:
210
+ import aiter.ops.sampling # noqa: F401
211
+ except ImportError:
212
+ self._aiter_ops_import_failed = True
213
+ self.forward = self.forward_native
214
+ logger.warning_once(
215
+ "aiter.ops.sampling is not available on ROCm. "
216
+ "Falling back to PyTorch-native implementation."
217
+ )
218
+ return False
219
+ self.aiter_ops = torch.ops.aiter
220
+ return True
221
+
222
def forward_hip(
223
self,
224
logits: torch.Tensor,
@@ -232,6 +240,8 @@ class TopKTopPSampler(nn.Module):
240
"processed_logits",
241
"processed_logprobs",
242
), "aiter sampler does not support returning logits/logprobs."
243
+ if self.aiter_ops is None and not self._init_aiter_ops():
244
+ return self.forward_native(logits, generators, k, p)
245
return self.aiter_sample(logits, k, p, generators), None
246
247
def aiter_sample(
@@ -242,6 +252,7 @@ class TopKTopPSampler(nn.Module):
252
generators: dict[int, torch.Generator],
253
) -> torch.Tensor:
254
"""Sample from logits using aiter ops."""
255
+ assert self.aiter_ops is not None
256
use_top_k = k is not None
257
use_top_p = p is not None
258
# Joint k+p path