Skip to content

fix(gms): cherry-pick support SGLang 0.5.16 memory pool API (#12445) #12492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nv-nmailhot merged 1 commit into from
Aug 3, 2026
Merged

fix(gms): cherry-pick support SGLang 0.5.16 memory pool API (#12445) #12492

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 121 additions & 0 deletions components/src/dynamo/sglang/tests/test_sglang_gms_memory_saver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from __future__ import annotations

from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import Mock

import pytest

Expand All @@ -23,7 +25,9 @@
pytest.mark.pre_merge,
pytest.mark.unit,
pytest.mark.gpu_0,
pytest.mark.profiled_vram_gib(0),
pytest.mark.sglang,
pytest.mark.core,
]


Expand Down Expand Up @@ -161,3 +165,120 @@ def test_region_requires_rw_allocator(build_impl, tag):
with pytest.raises(RuntimeError, match=rf"requires '{tag}' to be RW"):
with impl.region(tag, enable_cpu_backup=False):
pass


def test_write_publication_follows_outermost_region_and_switches_weights_ro(
build_impl, monkeypatch
):
impl, weights, _, _ = build_impl()
model = torch.nn.Module()
impl.preloaded_weights_bytes = 456
finalized_models = []
events = []

@contextmanager
def traced_pool(tag, device):
events.append(f"pool_enter:{tag}")
try:
yield
finally:
events.append(f"pool_exit:{tag}")

def finalize(allocator, pending_model):
events.append("finalize")
finalized_models.append(pending_model)
allocator.granted_lock_type = GrantedLockType.RO
return SimpleNamespace(committed_bytes=123)

monkeypatch.setattr(gms_memory_saver, "gms_use_mem_pool", traced_pool)
monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize)

with impl.region("weights", enable_cpu_backup=False):
with impl.region("kv_cache", enable_cpu_backup=False):
impl.finalize_write_mode(model)
events.append("defer")

assert events == [
"pool_enter:weights",
"pool_enter:kv_cache",
"defer",
"pool_exit:kv_cache",
"pool_exit:weights",
"finalize",
]
assert finalized_models == [model]
assert weights.granted_lock_type == GrantedLockType.RO
assert impl.imported_weights_bytes == 123
assert impl.preloaded_weights_bytes == 0

with impl.region("weights", enable_cpu_backup=False):
impl.finalize_write_mode(torch.nn.Module())

assert events[-1] == "finalize"
assert finalized_models == [model]


@pytest.mark.parametrize("failure_source", ["body", "pool_exit"])
def test_nested_region_failure_discards_pending_publication(
build_impl, monkeypatch, failure_source
):
impl, weights, _, _ = build_impl()
stale_model = torch.nn.Module()
fresh_model = torch.nn.Module()
finalize = Mock(return_value=SimpleNamespace(committed_bytes=1))
monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize)

@contextmanager
def maybe_failing_pool(tag, device):
yield
if failure_source == "pool_exit" and tag == "kv_cache":
raise ValueError("pool exit failed")

monkeypatch.setattr(gms_memory_saver, "gms_use_mem_pool", maybe_failing_pool)

failure_message = failure_source.replace("_", " ")
with impl.region("weights", enable_cpu_backup=False):
impl.finalize_write_mode(stale_model)
with pytest.raises(ValueError, match=f"{failure_message} failed"):
with impl.region("kv_cache", enable_cpu_backup=False):
if failure_source == "body":
raise ValueError("body failed")

finalize.assert_not_called()

with impl.region("weights", enable_cpu_backup=False):
impl.finalize_write_mode(fresh_model)
finalize.assert_called_once_with(weights, fresh_model)


def test_failed_finalization_is_cleared_and_not_retried(build_impl, monkeypatch):
impl, weights, _, _ = build_impl()
model = torch.nn.Module()
finalize = Mock(side_effect=ValueError("finalization failed"))
monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize)

with pytest.raises(ValueError, match="finalization failed"):
with impl.region("weights", enable_cpu_backup=False):
impl.finalize_write_mode(model)

with impl.region("weights", enable_cpu_backup=False):
pass
finalize.assert_called_once_with(weights, model)


def test_invalid_or_duplicate_publication_preserves_first_model(
build_impl, monkeypatch
):
impl, weights, _, _ = build_impl()
first_model = torch.nn.Module()
finalize = Mock(return_value=SimpleNamespace(committed_bytes=1))
monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize)

with impl.region("weights", enable_cpu_backup=False):
impl.finalize_write_mode(first_model)
with pytest.raises(TypeError, match="must not be None"):
impl.finalize_write_mode(None)
with pytest.raises(RuntimeError, match="publication is already pending"):
impl.finalize_write_mode(torch.nn.Module())

finalize.assert_called_once_with(weights, first_model)
39 changes: 35 additions & 4 deletions lib/gpu_memory_service/integrations/sglang/memory_saver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ def __init__(
self.imported_weights_bytes = 0
self.preloaded_weights_bytes = 0
self.ro_connect_timeout_ms = ro_connect_timeout_ms
self._active_region_depth = 0
self._pending_write_model: Optional[torch.nn.Module] = None
requested_mode = mode or RequestedLockType.RW_OR_RO
self.allocators = {
tag: get_or_create_gms_client_memory_manager(
Expand All @@ -95,7 +97,7 @@ def __init__(

@contextmanager
def region(self, tag: str, enable_cpu_backup: bool):
"""Mark allocation region with tag."""
"""Use the tag's RW pool and publish pending weights after clean exit."""
if enable_cpu_backup:
raise ValueError(
"SGLang with GMS does not support CPU backup for allocations."
Expand Down Expand Up @@ -134,8 +136,19 @@ def region(self, tag: str, enable_cpu_backup: bool):
f"SGLang with GMS requires {tag!r} to be RW for allocations; got {mode}"
)

with gms_use_mem_pool(tag, self._device):
yield
self._active_region_depth += 1
clean_exit = False
try:
with gms_use_mem_pool(tag, self._device):
yield
clean_exit = True
finally:
self._active_region_depth -= 1
if not clean_exit:
self._pending_write_model = None

if self._active_region_depth == 0:
self._finalize_pending_write()

@contextmanager
def cuda_graph(
Expand Down Expand Up @@ -184,9 +197,27 @@ def resume(self, tag: Optional[str] = None) -> None:
self.allocators[target_tag].remap_all_vas()

def finalize_write_mode(self, model: torch.nn.Module) -> None:
"""Finalize write mode: register tensors, commit, and switch to read."""
"""Publish write-mode weights after all managed GMS regions exit."""
if model is None:
raise TypeError("GMS weight publication model must not be None")

if self.allocators["weights"].granted_lock_type != GrantedLockType.RW:
# Read-only import mode never republishes weights.
self._pending_write_model = None
return

if self._pending_write_model is not None:
raise RuntimeError("GMS weight publication is already pending")

self._pending_write_model = model
if self._active_region_depth == 0:
self._finalize_pending_write()

def _finalize_pending_write(self) -> None:
"""Consume and publish pending weights outside managed GMS regions."""
model = self._pending_write_model
self._pending_write_model = None
if model is None:
return

stats = finalize_gms_write(self.allocators["weights"], model)
Expand Down
94 changes: 21 additions & 73 deletions lib/gpu_memory_service/integrations/sglang/patches.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@

from __future__ import annotations

import inspect
import logging
from contextlib import contextmanager
from typing import Optional
Expand Down Expand Up @@ -163,84 +162,33 @@ def patch_model_runner() -> None:
if hasattr(ModelRunner, "_gms_patched"):
return

original_init_memory_pool = ModelRunner.init_memory_pool
memory_arg_name = next(
(
name
for name in inspect.signature(original_init_memory_pool).parameters
if name != "self"
),
None,
)
original_alloc_memory_pool = ModelRunner.alloc_memory_pool

def patched_init_memory_pool(self, *args, **kwargs):
"""Patch memory baseline for SGLang old/new init_memory_pool signatures."""
def patched_alloc_memory_pool(self, *args, **kwargs):
impl = get_gms_memory_saver_impl()
preloaded_weights_gib = 0.0
if impl is not None:
if (
impl is not None
and impl.preloaded_weights_bytes > 0
and not self.__dict__.get("_gms_memory_baseline_adjusted", False)
):
preloaded_weights_gib = impl.preloaded_weights_bytes / (1 << 30)
old_value = self.pre_model_load_memory
self.pre_model_load_memory += preloaded_weights_gib
Comment on lines +165 to +176

@devin-ai-integration devin-ai-integration Bot Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Worker startup crashes hard on SGLang builds that lack the new memory-pool entry point

The startup patch grabs the new memory-pool entry point unconditionally (ModelRunner.alloc_memory_pool at lib/gpu_memory_service/integrations/sglang/patches.py:165) without any guard, so on an engine build that does not expose it the worker dies at import time instead of continuing with a warning.
Impact: Users on a slightly older engine release get an unexplained worker crash at start instead of a degraded-but-working start.

Loss of the previous version-tolerant patch path

The old implementation wrapped ModelRunner.init_memory_pool and probed its signature with inspect, tolerating both old and new parameter names, and the only failure mode was the ImportError branch (lib/gpu_memory_service/integrations/sglang/patches.py:156-160) which logs a warning and returns. The new code accesses ModelRunner.alloc_memory_pool directly outside that try, so an AttributeError propagates out of patch_model_runner(), which is executed at module import in the scheduler child process (lib/gpu_memory_service/integrations/sglang/model_loader.py:44). Similarly, patched_alloc_memory_pool reads self.pre_model_load_memory (patches.py:175-176) with no getattr fallback, so a renamed/absent attribute raises during read-mode startup instead of being logged and skipped as the old code did. Wrapping the attribute lookups (e.g. getattr(ModelRunner, "alloc_memory_pool", None) and a hasattr check on the baseline attribute) restores graceful degradation.

Prompt for agents 
In lib/gpu_memory_service/integrations/sglang/patches.py, patch_model_runner() now hard-depends on the SGLang 0.5.16 API: it reads ModelRunner.alloc_memory_pool at module import time (via model_loader.py's top-level patch_model_runner() call) and the patched wrapper reads self.pre_model_load_memory unconditionally. Previously the patch tolerated multiple SGLang signatures and only failed softly via the ImportError branch. Restore graceful degradation so that a missing method or missing baseline attribute logs a warning and leaves SGLang untouched rather than raising AttributeError inside the scheduler child process.
Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

self._gms_memory_baseline_adjusted = True
Comment on lines +175 to +177

@devin-ai-integration devin-ai-integration Bot Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔍 Baseline adjustment now mutates ModelRunner state permanently

Previously the preloaded-weight correction was applied only to the argument passed into init_memory_pool, leaving ModelRunner state untouched. Now self.pre_model_load_memory is permanently increased and guarded by a per-instance _gms_memory_baseline_adjusted flag. That makes the correction survive later re-allocations of the memory pool (e.g. resume-memory-occupation flows), which appears to be the intent, but it also means any other SGLang code that reads pre_model_load_memory (logging, memory accounting, later KV re-sizing) now sees an inflated baseline. Worth confirming against SGLang 0.5.16 that no other consumer of that attribute is affected.

Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

logger.info(
"[GMS] Adjusted pre_model_load_memory for preloaded weights: "
"%.2f GiB + %.2f GiB = %.2f GiB",
old_value,
preloaded_weights_gib,
self.pre_model_load_memory,
)

if preloaded_weights_gib > 0 and memory_arg_name in (
"pre_model_load_memory",
"total_gpu_memory",
):
if args:
old_value = args[0]
new_value = (
old_value + preloaded_weights_gib
if isinstance(old_value, (int, float))
else old_value
)
args = (new_value,) + args[1:]
elif memory_arg_name in kwargs:
old_value = kwargs[memory_arg_name]
new_value = (
old_value + preloaded_weights_gib
if isinstance(old_value, (int, float))
else old_value
)
kwargs = dict(kwargs)
kwargs[memory_arg_name] = new_value
else:
old_value = None
new_value = None

if isinstance(old_value, (int, float)) and isinstance(
new_value, (int, float)
):
logger.info(
"[GMS] Adjusted %s for preloaded weights: "
"%.2f GiB + %.2f GiB = %.2f GiB",
memory_arg_name,
old_value,
preloaded_weights_gib,
new_value,
)
else:
logger.info(
"[GMS] Could not adjust %s for preloaded weights; value=%r",
memory_arg_name,
old_value,
)
elif impl is not None and impl.imported_weights_bytes > 0:
if preloaded_weights_gib > 0:
logger.info(
"[GMS] Leaving %s unchanged; unsupported SGLang "
"init_memory_pool signature for preloaded weights",
memory_arg_name,
)
else:
logger.info(
"[GMS] Leaving %s unchanged; weights were loaded by this process",
memory_arg_name,
)

return original_init_memory_pool(self, *args, **kwargs)

ModelRunner.init_memory_pool = patched_init_memory_pool
return original_alloc_memory_pool(self, *args, **kwargs)

ModelRunner.alloc_memory_pool = patched_alloc_memory_pool
ModelRunner._gms_patched = True
_model_runner_patched = True
logger.info("[GMS] Patched ModelRunner.init_memory_pool")
logger.info("[GMS] Patched ModelRunner.alloc_memory_pool")


def patch_static_state_for_gms() -> None:
Expand Down
Loading
Loading