@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4013f3ad6c

[ROCm] Fix CI runtime and tests for MI355 DPX (#58244)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Mahesh Kunreddi <mahesh.kunreddi@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

mkunredd committed Sep 24, 2026 at 23:15 UTC 4013f3ad6ceafb063807a62c75bd6828e66913bf

12 files changed +175 -47

.buildkite/scripts/hardware_ci/run-amd-test.sh

+36 -5

@@ -605,6 +605,33 @@ initialize_native_environment() {

605 fi

606 }

607

608 +check_dpx_gpu_exclusivity() {

609 + local devices=(/dev/dri/renderD*)

610 + local device_id lock_status

611 +

612 + if [[ "${#devices[@]}" -ne 1 || ! -c "${devices[0]}" ]]; then

613 + echo "DPX guard requires exactly one render device; refusing to start tests." >&2

614 + return 1

615 + fi

616 + # This queue mounts the same host-local HF cache into every pod on the node.

617 + if ! mountpoint -q "${HF_HOME}"; then

618 + echo "DPX guard requires the shared HF cache mount; refusing to start tests." >&2

619 + return 1

620 + fi

621 + device_id=$(stat -Lc '%t-%T' "${devices[0]}") || return 1

622 + mkdir -p "${HF_HOME}/.vllm-dpx-locks" || return 1

623 + # Hold the descriptor through workload and teardown; never delete the file.

624 + exec {dpx_gpu_lock_fd}>>"${HF_HOME}/.vllm-dpx-locks/${device_id}.lock" || return 1

625 + flock -n -E 75 "${dpx_gpu_lock_fd}" && return 0

626 + lock_status=$?

627 + if [[ "${lock_status}" -eq 75 ]]; then

628 + echo "DPX GPU collision: ${devices[0]} (${device_id}) is already locked by another CI job; refusing to start tests." >&2

629 + else

630 + echo "DPX GPU lock failed (status ${lock_status}); refusing to start tests." >&2

631 + fi

632 + return 1

633 +}

634 +

635 run_native_preflight() {

636 local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"

637

@@ -1166,15 +1193,15 @@ collect_rocm_failure_diagnostics() {

1193 # Bus data identifies a card within the public node without publishing its

1194 # persistent UUID, serial number, or process list.

1195 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \

1169 - amd-smi static -b -g all

1196 + amd-smi static --bus --gpu all

1197 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \

1171 - amd-smi metric -e -k -P -x -g all

1198 + amd-smi metric --ecc --ecc-blocks --pcie --gpu all

1199 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \

1173 - amd-smi bad-pages -p -r -u -g all

1200 + amd-smi static --ras --gpu all

1201 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \

1175 - amd-smi metric -p -t -u -m -v -g all

1202 + amd-smi metric --power --temperature --usage --mem-usage --gpu all

1203 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \

1177 - amd-smi xgmi -l -g all

1204 + amd-smi xgmi --link-status --gpu all

1205 elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \

1206 && command -v rocm-smi >/dev/null 2>&1; then

1207 run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" rocm-smi \

@@ -1485,6 +1512,10 @@ if is_native_runtime; then

1512 echo "Failed to initialize the native test environment"

1513 exit 1

1514 fi

1515 + if [[ "${BUILDKITE_AGENT_META_DATA_QUEUE:-}" == *dpx* \

1516 + && "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" != "0" ]]; then

1517 + check_dpx_gpu_exclusivity || exit 1

1518 + fi

1519 if [[ "${commands}" == *python_only_compile.sh* ]]; then

1520 # This no-GPU job validates the ROCm precompiled/editable install path,

1521 # rather than CPU runtime platform selection.

tests/basic_correctness/test_mem.py

+6 -4

@@ -364,7 +364,7 @@ def test_end_to_end(model: str):

364

365

366 @create_new_process_for_each_test()

367 -def test_deep_sleep():

367 +def test_deep_sleep(gpu_memory_cleared):

368 model = "hmellor/tiny-random-LlamaForCausalLM"

369 free, total = torch.accelerator.get_memory_info()

370 used_bytes_baseline = total - free # in case other process is running

@@ -395,7 +395,7 @@ def test_deep_sleep():

395

396

397 @create_new_process_for_each_test()

398 -def test_deep_sleep_lora():

398 +def test_deep_sleep_lora(gpu_memory_cleared):

399 """Level-2 sleep/wake/reload with enable_lora=True.

400

401 LoRA wrapping moves parameters under base_layer and adds LoRA

@@ -488,7 +488,7 @@ def test_deep_sleep_lora_tp2(num_gpus_available, monkeypatch):

488

489

490 @create_new_process_for_each_test()

491 -def test_deep_sleep_async():

491 +def test_deep_sleep_async(gpu_memory_cleared):

492 async def test():

493 model = "hmellor/tiny-random-LlamaForCausalLM"

494 free, total = torch.accelerator.get_memory_info()

@@ -527,7 +527,9 @@ def test_deep_sleep_async():

527

528

529 @requires_fp8

530 -def test_deep_sleep_fp8_kvcache_mrv1(monkeypatch: pytest.MonkeyPatch):

530 +def test_deep_sleep_fp8_kvcache_mrv1(

531 + monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared

532 +):

533 # Regression test for https://github.com/vllm-project/vllm/pull/28783.

534 # In particular, verify that MRV1 does not rely on post_kv_cache_wake_up()

535 # to restore correct output after level-2 sleep.

tests/conftest.py

+18

@@ -1790,6 +1790,24 @@ def disable_deepgemm_ue8m0(monkeypatch):

1790 is_deep_gemm_e8m0_used.cache_clear()

1791

1792

1793 +@pytest.fixture

1794 +def gpu_memory_cleared():

1795 + """Wait for prior tests to release GPU memory on gfx950."""

1796 + if not current_platform.is_rocm():

1797 + return

1798 +

1799 + from tests.utils import wait_for_gpu_memory_to_clear

1800 + from vllm.platforms.rocm import on_gfx950

1801 +

1802 + if on_gfx950():

1803 + wait_for_gpu_memory_to_clear(

1804 + devices=[0],

1805 + threshold_ratio=0.08,

1806 + timeout_s=30,

1807 + stable_duration_s=1,

1808 + )

1809 +

1810 +

1811 def _should_clean_gpu_memory_between_tests() -> bool:

1812 # This must stay opt-in: a function-scoped fixture cannot distinguish

1813 # stale VRAM from allocations owned by longer-lived module/session fixtures.

tests/evals/gsm8k/configs/models-gfx950-small.txt new

@@ -0,0 +1,2 @@

1 +Qwen3-0.6B-FP8.yaml

2 +Qwen3-1.7B-MXFP4.yaml

tests/kernels/moe/test_block_int8.py

+31 -2

@@ -12,6 +12,9 @@ from tests.kernels.quant_utils import (

12 from vllm.config import VllmConfig, set_current_vllm_config

13 from vllm.model_executor.layers.activation import SiluAndMul

14 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk

15 +from vllm.model_executor.layers.quantization.utils.int8_utils import (

16 + per_token_group_quant_int8,

17 +)

18 from vllm.platforms import current_platform

20 if current_platform.get_device_capability() < (7, 0):

@@ -64,7 +67,7 @@ def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):

67 topk_ids = topk_ids.view(-1)

69 _, block_k = block_shape[0], block_shape[1]

67 - a_q, a_s = native_per_token_group_quant_int8(a, block_k)

70 + a_q, a_s = native_per_token_group_quant_int8(a, block_k, round_to_nearest=False)

71 for i in range(w1.shape[0]):

72 mask = topk_ids == i

73 if mask.sum():

@@ -72,7 +75,9 @@ def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):

75 a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype

76 )

77 act_out = SiluAndMul().forward_native(inter_out)

75 - act_out_q, act_out_s = native_per_token_group_quant_int8(act_out, block_k)

78 + act_out_q, act_out_s = native_per_token_group_quant_int8(

79 + act_out, block_k, round_to_nearest=False

80 + )

81 act_out = act_out.to(torch.float32)

82 out[mask] = native_w8a8_block_matmul(

83 act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype

@@ -132,3 +137,27 @@ def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):

137

138 # Check results

139 torch.testing.assert_close(out, ref_out, atol=0.065, rtol=0.065)

140 +

141 +

142 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])

143 +def test_block_int8_quantization_truncates_fractional_values(dtype):

144 + """Block INT8 quantization truncates, unlike per-token round-to-nearest."""

145 + device = current_platform.device_type

146 + x = torch.tensor(

147 + [[-127, -1.75, -0.75, 0.75, 1.75, 127] + [0] * 122],

148 + dtype=dtype,

149 + device=device,

150 + )

151 + expected = torch.tensor(

152 + [[-127, -1, 0, 0, 1, 127] + [0] * 122], dtype=torch.int8, device=device

153 + )

154 +

155 + out_q, scale = per_token_group_quant_int8(x, group_size=128)

156 + ref_q, ref_s = native_per_token_group_quant_int8(

157 + x, group_size=128, round_to_nearest=False

158 + )

159 +

160 + torch.testing.assert_close(out_q, expected, atol=0, rtol=0)

161 + torch.testing.assert_close(ref_q, expected, atol=0, rtol=0)

162 + torch.testing.assert_close(scale, torch.ones_like(scale), atol=0, rtol=0)

163 + torch.testing.assert_close(ref_s, scale, atol=0, rtol=0)

tests/kernels/quant_utils.py

+8 -4

@@ -181,12 +181,15 @@ def native_per_token_group_quant_fp8(

181 return x_q, x_s

182

183

184 -def native_per_token_group_quant_int8(x, group_size, eps=1e-10, dtype=torch.int8):

184 +def native_per_token_group_quant_int8(

185 + x, group_size, eps=1e-10, dtype=torch.int8, *, round_to_nearest=True

186 +):

187 """Function to perform per-token-group quantization on an input tensor

188 `x` using native torch.

189

190 It converts the tensor values into int8 values and returns the

191 quantized tensor along with the scaling factor used for quantization.

192 + Per-token quantization rounds to nearest; block quantization truncates.

193 """

194 assert x.shape[-1] % group_size == 0, (

195 "the last dimension of `x` must be divisible by `group_size`"

@@ -201,9 +204,10 @@ def native_per_token_group_quant_int8(x, group_size, eps=1e-10, dtype=torch.int8

204 # Use float32 for scale calculation for stability

205 amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)

206 x_s = amax / int8_max

204 - x_q = (

205 - (x_.to(torch.float32) / x_s).round().clamp(min=int8_min, max=int8_max).to(dtype)

206 - ) # Round before clamping

207 + x_q = x_.to(torch.float32) / x_s

208 + if round_to_nearest:

209 + x_q = x_q.round()

210 + x_q = x_q.clamp(min=int8_min, max=int8_max).to(dtype)

211 x_q = x_q.reshape(x.shape)

212 x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))

213

tests/kernels/quantization/test_rocm_mxfp4.py

+3 -1

@@ -614,11 +614,13 @@ def test_aiter_fp4_gemm_skinny_shapes(M, N, K):

614 out = gemm_afp4wfp4(A_fp4, B_fp4, A_scale, B_scale)

615

616 assert out.shape == (M, N)

617 + assert out.dtype == torch.bfloat16

618 assert not torch.any(torch.isnan(out))

619

620 + # Accumulate in FP32, then match the GEMM's BF16 output rounding.

621 A_dq = quant_dequant_mxfp4(A)

622 B_dq = quant_dequant_mxfp4(B)

621 - ref = torch.matmul(A_dq.float(), B_dq.t().float())

623 + ref = torch.matmul(A_dq.float(), B_dq.t().float()).to(torch.bfloat16).float()

624

625 _print_close_stats(

626 f"skinny_gemm M={M} N={N} K={K}",

tests/kernels/quantization/test_rocm_skinny_gemms.py

+2 -1

@@ -132,7 +132,8 @@ M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]

132 NKM_FACTORS_WVSPLITKRC_LARGE_K = [

133 (128, 6144, 128),

134 (96, 8192, 128),

135 - (128, 12288, 128),

135 + # Keep 24 K-shards (batched readback) within the DPX budget of 128 CUs.

136 + (128, 12288, 64),

137 (32, 12288, 128),

138 ]

139

tests/models/language/generation/test_hybrid.py

@@ -73,6 +73,7 @@ def _set_conv_state_layout(monkeypatch, layout: str) -> None:

73 def test_models(

74 hf_runner,

75 vllm_runner,

76 + gpu_memory_cleared,

77 example_prompts,

78 monkeypatch,

79 model: str,

tests/v1/kv_connector/unit/test_nixl_rocm_gpu_mem_diag.py

+21 -4

@@ -12,6 +12,7 @@ import gc

12 import pytest

13 import torch

15 +from tests.utils import wait_for_gpu_memory_to_clear

16 from vllm.platforms import current_platform

18 pytestmark = pytest.mark.skipif(

@@ -19,6 +20,8 @@ pytestmark = pytest.mark.skipif(

20 reason="ROCm platform required",

21 )

23 +MAX_LEAK_PCT = 10

24 +

26 def _mb(b: int) -> float:

27 return b / (1024 * 1024)

@@ -61,6 +64,18 @@ def _full_gpu_cleanup():

64 torch.accelerator.empty_cache()

67 +def _wait_for_driver_memory_release(baseline: dict, peak: dict) -> None:

68 + # Process exit can precede ROCm's release of the backing allocation.

69 + peak_delta_mb = max(0.0, peak["drv_used_mb"] - baseline["drv_used_mb"])

70 + target_mb = baseline["drv_used_mb"] + peak_delta_mb * MAX_LEAK_PCT / 100

71 + wait_for_gpu_memory_to_clear(

72 + devices=[torch.accelerator.current_device_index()],

73 + threshold_bytes=int(target_mb * 1024**2),

74 + timeout_s=120,

75 + poll_interval_s=1,

76 + )

77 +

78 +

79 @pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])

80 def test_gpu_memory_nixl_hma(model_name, sw_size):

81 """Track GPU memory through NixlConnector create/infer/shutdown cycle."""

@@ -116,7 +131,7 @@ def test_gpu_memory_nixl_hma(model_name, sw_size):

131

132 # shutdown + cleanup

133 print("\n--- shutdown ---")

119 - llm.llm_engine.engine_core.shutdown()

134 + llm.llm_engine.engine_core.shutdown(timeout=60)

135 _gpu_snapshot("3. after shutdown()", snap2["alloc_mb"])

136

137 del llm

@@ -126,6 +141,7 @@ def test_gpu_memory_nixl_hma(model_name, sw_size):

141 torch._dynamo.reset()

142 gc.collect()

143 torch.accelerator.empty_cache()

144 + _wait_for_driver_memory_release(snap0, snap2)

145 snap_final = _gpu_snapshot("4. final", snap2["alloc_mb"])

146

147 # summary

@@ -160,7 +176,7 @@ def test_gpu_memory_nixl_hma(model_name, sw_size):

176 # Peak driver memory used above baseline

177 drv_peak = snap2["drv_used_mb"] - drv_base

178 leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0

163 - max_leak_pct = 10

179 + max_leak_pct = MAX_LEAK_PCT

180 assert leak_pct <= max_leak_pct, (

181 f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "

182 f"not freed after NixlConnector shutdown "

@@ -196,7 +212,7 @@ def test_gpu_memory_no_nixl_baseline(model_name):

212 llm.generate(["hi " * 500], SamplingParams(max_tokens=1))

213 snap_peak = _gpu_snapshot("after generate()", snap0["alloc_mb"])

214

199 - llm.llm_engine.engine_core.shutdown()

215 + llm.llm_engine.engine_core.shutdown(timeout=60)

216 del llm

217 _full_gpu_cleanup()

218 cleanup_dist_env_and_memory()

@@ -204,6 +220,7 @@ def test_gpu_memory_no_nixl_baseline(model_name):

220 torch._dynamo.reset()

221 gc.collect()

222 torch.accelerator.empty_cache()

223 + _wait_for_driver_memory_release(snap0, snap_peak)

224 snap_final = _gpu_snapshot("final", snap0["alloc_mb"])

225

226 drv_base = snap0["drv_used_mb"]

@@ -213,7 +230,7 @@ def test_gpu_memory_no_nixl_baseline(model_name):

230 print("=" * 90)

231

232 leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0

216 - max_leak_pct = 10

233 + max_leak_pct = MAX_LEAK_PCT

234 assert leak_pct <= max_leak_pct, (

235 f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "

236 f"not freed after baseline shutdown "

tests/v1/metrics/test_engine_logger_apis.py

+24 -12

@@ -7,6 +7,7 @@ import pytest

7 from tests.plugins.vllm_add_dummy_stat_logger.dummy_stat_logger.dummy_stat_logger import ( # noqa E501

8 DummyStatLogger,

9 )

10 +from tests.utils import wait_for_memory_to_settle

11 from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM

12 from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger

@@ -30,8 +31,15 @@ async def test_async_llm_replace_default_loggers(log_stats_enabled_engine_args):

31 engine = AsyncLLM.from_engine_args(

32 log_stats_enabled_engine_args, stat_loggers=[RayPrometheusStatLogger]

33 )

33 - assert isinstance(engine.logger_manager.stat_loggers[0], RayPrometheusStatLogger)

34 - engine.shutdown()

34 + try:

35 + assert isinstance(

36 + engine.logger_manager.stat_loggers[0], RayPrometheusStatLogger

37 + )

38 + finally:

39 + engine.shutdown()

40 + wait_for_memory_to_settle(

41 + threshold_ratio=1.0 - log_stats_enabled_engine_args.gpu_memory_utilization

42 + )

45 @pytest.mark.asyncio

@@ -48,14 +56,18 @@ async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args):

56 disabled_log_engine_args, stat_loggers=[DummyStatLogger]

57 )

51 - assert len(engine.logger_manager.stat_loggers) == 2

52 - assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1

53 - assert isinstance(

54 - engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0],

55 - DummyStatLogger,

56 - )

57 -

58 - # log_stats is still True, since custom stat loggers are used

59 - assert engine.log_stats

59 + try:

60 + assert len(engine.logger_manager.stat_loggers) == 2

61 + assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1

62 + assert isinstance(

63 + engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0],

64 + DummyStatLogger,

65 + )

61 - engine.shutdown()

67 + # log_stats is still True, since custom stat loggers are used

68 + assert engine.log_stats

69 + finally:

70 + engine.shutdown()

71 + wait_for_memory_to_settle(

72 + threshold_ratio=1.0 - disabled_log_engine_args.gpu_memory_utilization

73 + )

tests/v1/metrics/test_ray_metrics.py

+23 -14

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock

6 import pytest

7 import ray

9 +from tests.utils import wait_for_memory_to_settle

10 from vllm.config.model import ModelDType

11 from vllm.platforms import current_platform

12 from vllm.sampling_params import SamplingParams

@@ -46,31 +47,39 @@ def test_engine_log_metrics_ray(

47 ) -> None:

48 """Simple smoke test, verifying this can be used without exceptions.

49 Need to start a Ray cluster in order to verify outputs."""

50 + engine_args = AsyncEngineArgs(

51 + model=model, dtype=dtype, disable_log_stats=False, enforce_eager=True

52 + )

54 @ray.remote(num_gpus=1)

55 class EngineTestActor:

56 async def run(self):

53 - engine_args = AsyncEngineArgs(

54 - model=model, dtype=dtype, disable_log_stats=False, enforce_eager=True

55 - )

56 -

57 engine = AsyncLLM.from_engine_args(

58 engine_args, stat_loggers=[RayPrometheusStatLogger]

59 )

61 - for i, prompt in enumerate(example_prompts):

62 - results = engine.generate(

63 - request_id=f"request-id-{i}",

64 - prompt=prompt,

65 - sampling_params=SamplingParams(max_tokens=max_tokens),

66 - )

61 + try:

62 + for i, prompt in enumerate(example_prompts):

63 + results = engine.generate(

64 + request_id=f"request-id-{i}",

65 + prompt=prompt,

66 + sampling_params=SamplingParams(max_tokens=max_tokens),

67 + )

68 - async for _ in results:

69 - pass

69 + async for _ in results:

70 + pass

71 + finally:

72 + engine.shutdown()

74 # Create the actor and call the async method

72 - actor = EngineTestActor.remote() # type: ignore[attr-defined]

73 - ray.get(actor.run.remote())

75 + try:

76 + actor = EngineTestActor.remote() # type: ignore[attr-defined]

77 + ray.get(actor.run.remote())

78 + finally:

79 + ray.shutdown()

80 + wait_for_memory_to_settle(

81 + threshold_ratio=1.0 - engine_args.gpu_memory_utilization

82 + )

85 def test_sanitized_opentelemetry_name():