[Bugfix][ROCm] Preserve config during GPU memory profiling

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.
 

vllm/v1/worker/gpu_worker.py

@@ -615,7 +615,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)} "
@@ -651,6 +654,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)