From d660bed89af14ab181e0173aef6b25845ed54a21 Mon Sep 17 00:00:00 2001 From: Thiago Mota Martins Date: Mon, 21 Sep 2026 18:45:20 -0300 Subject: [PATCH] fix(rocm): preserve config during GPU memory profiling Signed-off-by: Thiago Mota Martins Assisted-by: OpenAI Codex --- tests/v1/worker/test_gpu_worker.py | 40 ++++++++++++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 6 ++++- 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index b63681deca2c..90adce187448 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -152,6 +152,46 @@ def test_scoped_max_split_ignores_async_allocator(monkeypatch): set_settings.assert_not_called() +@pytest.mark.parametrize("kv_cache_memory_bytes", [None, 1]) +def test_memory_profile_bounds_decode_logits_rows(monkeypatch, kv_cache_memory_bytes): + """ROCm decode workspace uses max_num_seqs during startup profiling.""" + from vllm.config import get_current_vllm_config_or_none + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _max_decode_logits_rows + + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=SimpleNamespace(num_speculative_tokens=5), + ) + + class Profiled(Exception): + pass + + def profile_run(*, randomize_inputs=False): + assert randomize_inputs is True + # Without the config context, the helper falls back to all 32768 + # batched tokens, which would reserve 128 GiB at 1M context. + assert _max_decode_logits_rows(32768) == 16 * (1 + 5) + raise Profiled + + worker = SimpleNamespace( + vllm_config=config, + randomize_dummy_inputs=True, + cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_cache_memory_bytes), + model_runner=SimpleNamespace(profile_run=profile_run, model_memory_usage=0), + init_snapshot=SimpleNamespace(free_memory=2), + _scoped_allocator_max_split=lambda **kwargs: nullcontext(), + ) + monkeypatch.setattr(gpu_worker, "maybe_apply_startup_plan", lambda _: None) + monkeypatch.setattr( + gpu_worker, "memory_profiling", lambda *args, **kwargs: nullcontext() + ) + + assert get_current_vllm_config_or_none() is None + with pytest.raises(Profiled): + gpu_worker.Worker.determine_available_memory(worker) + assert get_current_vllm_config_or_none() is None + + # Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and # saved by Worker.determine_available_memory / compile_or_warm_up_model. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index e9021bd142c0..4ca4214f011f 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -613,7 +613,10 @@ def determine_available_memory(self) -> int: if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: # still need a profile run which compiles the model for # max_num_batched_tokens - self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs) + with set_current_vllm_config(self.vllm_config): + self.model_runner.profile_run( + randomize_inputs=self.randomize_dummy_inputs + ) msg = ( f"Initial free memory {format_gib(self.init_snapshot.free_memory)} " @@ -649,6 +652,7 @@ def determine_available_memory(self) -> int: # limit are never split, so they stay releasable. Exits before # memory_profiling measures, restoring the original limit. self._scoped_allocator_max_split(max_split_size_mb=20), + set_current_vllm_config(self.vllm_config), ): self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)