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 10, 10:15 UTC
[ROCm][CI][V1] Fix `nixl_connector` test failure and achieve CUDA parity in `test_async_scheduling` (#32000)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jan 9, 2026 at 06:48 UTC
e02706d2d27c9af429adf89e7dec2b37e3ec39c1
3 files changed
+28
-39
tests/v1/e2e/test_async_scheduling.py
+3
-23
@@ -163,14 +163,7 @@ def run_tests(
163
uni/multiproc executor with spec decoding."""
164
165
# Determine attention config based on platform
166
- if current_platform.is_rocm():
167
- if is_testing_with_spec_decoding:
168
- # Use TRITON_ATTN for spec decoding test for consistency
169
- attention_config = {"backend": "TRITON_ATTN"}
170
- else:
171
- attention_config = {"backend": "ROCM_ATTN"}
172
- else:
173
- attention_config = {"backend": "FLEX_ATTENTION"}
166
+ attention_config = {"backend": "FLEX_ATTENTION"}
167
168
with monkeypatch.context() as m:
169
# lock matmul precision to full FP32 (IEEE)
@@ -226,15 +219,7 @@ def run_tests(
219
name_1=f"config=[{test_config}], params={params}",
220
)
221
229
- # On ROCm with TRITON_ATTN (spec decoding test), skip strict
230
- # logprobs comparison when logprobs are requested
231
- skip_logprobs_check = (
232
- current_platform.is_rocm()
233
- and params.get("logprobs")
234
- and is_testing_with_spec_decoding
235
- )
236
- if not skip_logprobs_check:
237
- assert _all_logprobs_match(base_logprobs, test_logprobs)
222
+ assert _all_logprobs_match(base_logprobs, test_logprobs)
223
224
if (
225
base_acceptance_rate is not None
@@ -374,12 +359,7 @@ def _all_logprobs_match(req_a, req_b) -> bool:
359
360
361
def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
377
- if current_platform.is_rocm():
378
- # ROCm has higher numerical variance
379
- # due to use of float16.
380
- rel_tol, abs_tol = 5e-2, 1e-5
381
- else:
382
- rel_tol, abs_tol = 1e-3, 1e-6
362
+ rel_tol, abs_tol = 1e-3, 1e-6
363
return (
364
len(lps_a) == len(lps_b)
365
and lps_a.keys() == lps_b.keys()
tests/v1/kv_connector/unit/test_nixl_connector.py
+20
-16
@@ -185,18 +185,21 @@ class FakeNixlWrapper:
185
def _make_fake_nixl_pkg():
186
"""Context manager that creates a temporary package making
187
`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
188
+ Also creates rixl package for ROCm compatibility.
189
190
Automatically cleans up the temporary directory when done.
191
"""
192
with tempfile.TemporaryDirectory() as td:
192
- pkg_root = os.path.join(td, "nixl", "_api")
193
- os.makedirs(pkg_root, exist_ok=True)
193
+ # Create both nixl and rixl packages for cross-platform compatibility
194
+ for pkg_name in ["nixl", "rixl"]:
195
+ pkg_root = os.path.join(td, pkg_name, "_api")
196
+ os.makedirs(pkg_root, exist_ok=True)
197
195
- # Get the source code of FakeNixlWrapper class and dedent it
196
- fake_nixl_source = inspect.getsource(FakeNixlWrapper)
197
- fake_nixl_source = textwrap.dedent(fake_nixl_source)
198
+ # Get the source code of FakeNixlWrapper class and dedent it
199
+ fake_nixl_source = inspect.getsource(FakeNixlWrapper)
200
+ fake_nixl_source = textwrap.dedent(fake_nixl_source)
201
199
- stub = f"""\
202
+ stub = f"""\
203
# Copy of FakeNixlWrapper implementation for Ray workers
204
import uuid
205
from collections import defaultdict
@@ -206,16 +209,17 @@ from collections import defaultdict
209
# Export as nixl_agent
210
nixl_agent = FakeNixlWrapper
211
"""
209
- with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
210
- f.write(stub)
211
-
212
- # Mock nixlXferTelemetry class
213
- pkg_root2 = os.path.join(td, "nixl", "_bindings")
214
- os.makedirs(pkg_root2, exist_ok=True)
215
- with open(os.path.join(pkg_root2, "__init__.py"), "w") as f:
216
- f.write("class nixlXferTelemetry: pass")
217
- # touch parent package
218
- open(os.path.join(td, "nixl", "__init__.py"), "w").close()
212
+ with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
213
+ f.write(stub)
214
+
215
+ # Mock nixlXferTelemetry class
216
+ pkg_root2 = os.path.join(td, pkg_name, "_bindings")
217
+ os.makedirs(pkg_root2, exist_ok=True)
218
+ with open(os.path.join(pkg_root2, "__init__.py"), "w") as f:
219
+ f.write("class nixlXferTelemetry: pass")
220
+ # touch parent package
221
+ open(os.path.join(td, pkg_name, "__init__.py"), "w").close()
222
+
223
yield td
224
225
vllm/v1/spec_decode/eagle.py
+5
@@ -187,6 +187,11 @@ class EagleProposer:
187
188
rocm_types.append(MLACommonMetadata)
189
190
+ # FlexAttention backend support
191
+ from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadata
192
+
193
+ rocm_types.append(FlexAttentionMetadata)
194
+
195
self.allowed_attn_types = tuple(rocm_types)
196
197
# Parse the speculative token tree.