From fb6f96c468d9b7685e5ec948ae6ff0d0f5a3838d Mon Sep 17 00:00:00 2001 From: Shai Shmuel Date: Mon, 2 Nov 2026 13:56:34 +0200 Subject: [PATCH] fix(server): reserve declared SGLang runtime headroom at load staging SGLang memory admission reserves the static pool plus a fixed allowance: 1 GiB with an explicit mem_fraction_static, 2 GiB with automatic sizing. A vision encoder can need several GiB more outside the static pool for one full-size image. A model that loads onto a shared device therefore passed admission without room for its own peak, and its first large page could run out of memory. Profiles can now declare the measured excess as the loadtime option runtime_headroom_gib. The memory budget adds it to the load reservation in both modes, capped at the device total for an explicit fraction. Load staging evicts idle models until that memory is free, and an automatically sized engine does not launch without it. The headroom is never given to the KV pool. The value defaults to 0, so existing profiles keep their reservations. Co-Authored-By: Claude Opus 5.5 --- packages/sie_server/README.md | 10 +++++ .../src/sie_server/adapters/sglang/_memory.py | 9 +++- .../sie_server/adapters/sglang/generation.py | 5 +++ .../tests/adapters/test_sglang_memory.py | 44 +++++++++++++++++++ .../config/test_sglang_memory_profiles.py | 13 ++++++ 5 files changed, 80 insertions(+), 1 deletion(-) diff --git a/packages/sie_server/README.md b/packages/sie_server/README.md index f3d5c76ee..eb75f379c 100644 --- a/packages/sie_server/README.md +++ b/packages/sie_server/README.md @@ -1220,6 +1220,16 @@ attention. Calibrated profiles for these configurations retain their overrides. Automatic sizing is an estimate; unusually large runtime workspaces can still invoke the existing OOM recovery path. +A profile whose engine needs more memory outside its static pool than these +allowances can declare the measured excess as +`adapter_options.loadtime.runtime_headroom_gib`. An example is a vision encoder +processing one full-size page. SIE adds the value to the reservation in both +modes. Load staging then evicts idle eligible models until that much memory is +free, and an automatically sized engine does not launch without it. The +headroom stays outside the engine's static pool. A model loaded later is still +admitted on actual device usage, so it can take headroom that the engine has not +allocated yet. + | Env var | Default | Effect | |--|--|--| | `SIE_MEMORY_PRESSURE_THRESHOLD_PERCENT` | `95` | VRAM utilisation that triggers reactive LRU eviction by the pressure monitor. | diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_memory.py b/packages/sie_server/src/sie_server/adapters/sglang/_memory.py index 08ad6d6ca..f670dffb4 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_memory.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_memory.py @@ -116,6 +116,7 @@ class SGLangMemoryBudget: mem_fraction_static: float | None = None tensor_parallel_size: int = 1 has_auxiliary_weights: bool = False + runtime_headroom_gib: float = 0 reserved_bytes: int = 0 def __post_init__(self) -> None: @@ -126,6 +127,9 @@ def __post_init__(self) -> None: isinstance(fraction, bool) or not isinstance(fraction, int | float) or not 0 < fraction <= 1 ): raise ValueError("mem_fraction_static must be a finite number in (0, 1]") + headroom = self.runtime_headroom_gib + if isinstance(headroom, bool) or not isinstance(headroom, int | float) or not 0 <= headroom < math.inf: + raise ValueError("runtime_headroom_gib must be a finite non-negative number") @property def token_capacity(self) -> int: @@ -150,15 +154,18 @@ def static_bytes(self) -> int: def required_bytes(self, *, device_type: str, device_total_bytes: int) -> int | None: if device_type != "cuda" or device_total_bytes <= 0: return None + headroom = math.ceil(self.runtime_headroom_gib * GIB) if self.mem_fraction_static is not None: required = _server.estimate_load_required_memory_bytes( device_type=device_type, device_total_bytes=device_total_bytes, mem_fraction_static=self.mem_fraction_static, ) + if required is not None: + required = min(required + headroom, device_total_bytes) self.reserved_bytes = (required or 0) * self.tensor_parallel_size return required - required = self.static_bytes + RUNTIME_HEADROOM_BYTES + required = self.static_bytes + RUNTIME_HEADROOM_BYTES + headroom self.reserved_bytes = required return required diff --git a/packages/sie_server/src/sie_server/adapters/sglang/generation.py b/packages/sie_server/src/sie_server/adapters/sglang/generation.py index 2cf2f9d34..a97721a1e 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/generation.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/generation.py @@ -552,6 +552,10 @@ def __init__( # it, then follows as usual. None leaves reasoning unbounded and the # launch unchanged. Needs a reasoning parser with known control tokens. thinking_budget: int | None = None, + # GiB the engine allocates outside its static pool beyond the default + # allowance at peak, such as a vision encoder's activations for one + # full-size image. Load staging requires it free before launch. + runtime_headroom_gib: float = 0, **kwargs: Any, # accept extra args from loader for compatibility ) -> None: _ = kwargs @@ -571,6 +575,7 @@ def __init__( mem_fraction_static=mem_fraction_static, tensor_parallel_size=self._tensor_parallel_size, has_auxiliary_weights=bool(lora_paths or (speculative and speculative.get("enabled"))), + runtime_headroom_gib=runtime_headroom_gib, ) self._watchdog_timeout_s = _server.validate_optional_positive(watchdog_timeout_s, field="watchdog_timeout_s") self._declared_nccl_port = _server.validate_optional_port(nccl_port) diff --git a/packages/sie_server/tests/adapters/test_sglang_memory.py b/packages/sie_server/tests/adapters/test_sglang_memory.py index b0f6022ad..07b7747a0 100644 --- a/packages/sie_server/tests/adapters/test_sglang_memory.py +++ b/packages/sie_server/tests/adapters/test_sglang_memory.py @@ -93,6 +93,50 @@ def test_explicit_fraction_keeps_admission_and_unbounded_kv_behavior(monkeypatch assert budget.launch_args(0) == ["--mem-fraction-static", "0.85"] +def test_runtime_headroom_is_reserved_in_both_admission_paths(model_dir): + automatic = _budget(model_dir).required_bytes(device_type="cuda", device_total_bytes=80 * GIB) + budget = _budget(model_dir, runtime_headroom_gib=4) + assert budget.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == automatic + 4 * GIB + assert budget.reserved_bytes == automatic + 4 * GIB + + explicit = _budget(Path("missing-model"), mem_fraction_static=0.85, runtime_headroom_gib=4) + assert explicit.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == 73 * GIB + assert explicit.reserved_bytes == 73 * GIB + oversized = _budget(Path("missing-model"), mem_fraction_static=0.85, runtime_headroom_gib=20) + assert oversized.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == 80 * GIB + + +def test_runtime_headroom_stays_outside_the_static_pool(model_dir, monkeypatch): + budget = _budget(model_dir, runtime_headroom_gib=4) + without_headroom = _budget(model_dir).required_bytes(device_type="cuda", device_total_bytes=80 * GIB) + monkeypatch.setattr(_memory, "device_memory", lambda _: (without_headroom, 80 * GIB)) + with pytest.raises(RuntimeError, match="CUDA out of memory"): + budget.launch_args(0) + monkeypatch.setattr(_memory, "device_memory", lambda _: (60 * GIB, 80 * GIB)) + assert float(budget.launch_args(0)[1]) * 60 * GIB == pytest.approx(budget.static_bytes) + + +@pytest.mark.parametrize("headroom", [-1, True, float("nan"), float("inf"), "4"]) +def test_invalid_runtime_headroom_rejected(headroom): + with pytest.raises(ValueError, match="runtime_headroom_gib"): + _budget(Path("missing"), runtime_headroom_gib=headroom) + + +def test_load_staging_evicts_until_the_runtime_headroom_is_free(model_dir): + plain = SGLangVisionExtractAdapter(str(model_dir), max_seq_length=32768, kv_budget_tokens=32768) + vision = SGLangVisionExtractAdapter( + str(model_dir), max_seq_length=32768, kv_budget_tokens=32768, runtime_headroom_gib=4 + ) + plain_required = plain.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + assert plain_required is not None + manager = MemoryManager(device="cuda:0") + free = plain_required + GIB + with patch.object(manager, "get_stats", return_value=MemoryStats(80 * GIB - free, 80 * GIB, "cuda")): + assert not manager.should_evict_for_load(_adapter_load_required_bytes(plain, manager)) + assert _adapter_load_required_bytes(vision, manager) == plain_required + 4 * GIB + assert manager.should_evict_for_load(_adapter_load_required_bytes(vision, manager)) + + def test_sharded_weights_use_index_and_serving_precision(model_dir): (model_dir / "model.safetensors").unlink() _write_weights(model_dir / "model-1.safetensors", 1000, "F32", 4) diff --git a/packages/sie_server/tests/config/test_sglang_memory_profiles.py b/packages/sie_server/tests/config/test_sglang_memory_profiles.py index c3175e34a..c5a445b5d 100644 --- a/packages/sie_server/tests/config/test_sglang_memory_profiles.py +++ b/packages/sie_server/tests/config/test_sglang_memory_profiles.py @@ -107,6 +107,19 @@ def test_glm_ocr_launches_with_its_explicit_fraction(): assert budget.launch_args(0) == ["--mem-fraction-static", "0.85"] +def test_profile_runtime_headroom_reaches_the_load_reservation(): + config = load_model_config(MODELS / "zai-org__GLM-OCR.yaml") + baseline = load_adapter(config, MODELS, device="cuda:0") + assert baseline._memory_budget.runtime_headroom_gib == 0 + values = config.model_dump() + values["profiles"]["default"]["adapter_options"]["loadtime"]["runtime_headroom_gib"] = 4 + adapter = load_adapter(type(config).model_validate(values), MODELS, device="cuda:0") + assert adapter._memory_budget.runtime_headroom_gib == 4 + required = adapter.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + assert required == baseline.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + 4 * GIB + assert required == 73 * GIB + + def test_guard_calibrated_profile_retains_explicit_override(): config = load_model_config(MODELS / "Qwen__Qwen3Guard-Gen-4B.yaml") assert config.resolve_profile("default").kv_budget_tokens == 32768