[CI/Build][AMD] Fix ROCm OOM in eagle_correctness_heavy by reserving CUDA graph memory (#47366)
Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
peizhang56 committed
Jul 9, 2026 at 19:14 UTC
88e5e2c57be8ce6e25510c1249352a23b8a85ec4
2 files changed
+14
-5
tests/v1/e2e/spec_decode/test_spec_decode.py
+7
@@ -17,6 +17,7 @@ from tests.utils import (
17
multi_gpu_marks,
18
multi_gpu_only,
19
single_gpu_only,
20
+ wait_for_rocm_memory_to_settle,
21
)
22
from vllm import LLM, SamplingParams
23
from vllm.assets.base import VLLM_S3_BUCKET_URL
@@ -448,6 +449,9 @@ def _run_eagle_correctness(
449
del ref_llm
450
torch.accelerator.empty_cache()
451
cleanup_dist_env_and_memory()
452
+ # ROCm frees VRAM lazily; wait so the spec engine started right after
453
+ # does not OOM on its startup memory guard.
454
+ wait_for_rocm_memory_to_settle()
455
456
spec_llm = LLM(
457
model=model_name,
@@ -485,6 +489,9 @@ def _run_eagle_correctness(
489
del spec_llm
490
torch.accelerator.empty_cache()
491
cleanup_dist_env_and_memory()
492
+ # ROCm frees VRAM lazily; wait so the next parametrization's engine does
493
+ # not OOM on its startup memory guard.
494
+ wait_for_rocm_memory_to_settle()
495
496
497
@single_gpu_only
vllm/v1/worker/gpu_worker.py
+7
-5
@@ -478,11 +478,14 @@ class Worker(WorkerBase):
478
)
479
480
# Profile CUDA graph memory if graphs will be captured.
481
- # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info
482
- # behave differently and can produce incorrect/negative estimates.
481
+ # ROCm is included: #44825 moved the profiler to
482
+ # torch.accelerator.get_memory_info (reliable on ROCm, as used by
483
+ # the AMD-CI mem tests), and graph_pool_handle resolves to the same
484
+ # torch.cuda handle the live capture path already uses on ROCm.
485
+ # XPU stays excluded (see #39977).
486
cudagraph_memory_estimate = 0
487
if (
485
- current_platform.is_cuda()
488
+ current_platform.is_cuda_alike()
489
and self.vllm_config.compilation_config.cudagraph_mode
490
!= CUDAGraphMode.NONE
491
):
@@ -498,8 +501,7 @@ class Worker(WorkerBase):
501
+ profile_result.weights_memory
502
)
503
501
- # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op.
502
- # On CUDA, respect the opt-in flag as originally designed.
504
+ # Respect the opt-in flag as originally designed.
505
cudagraph_memory_estimate_applied = (
506
cudagraph_memory_estimate
507
if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS