diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_compat/sie_lora_compat.py b/packages/sie_server/src/sie_server/adapters/sglang/_compat/sie_lora_compat.py index bca828aa2..7eda52c4d 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_compat/sie_lora_compat.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_compat/sie_lora_compat.py @@ -1,9 +1,11 @@ -"""Release LoRA GPU slots when the pinned SGLang 0.5.10 engine unloads an adapter. +"""Load partial QKV LoRAs and release their GPU slots on pinned SGLang engines. Upstream issue https://github.com/sgl-project/sglang/issues/21380 describes the missing cleanup: CPU weights disappear, but the GPU pool keeps the old UID. Repeated load/unload cycles then exhaust the pool. The CUDA 13 bundle's -SGLang 0.5.20 already implements this cleanup and must keep its native path. +SGLang 0.5.20 already implements this cleanup and keeps its native unload path. +Both pinned engines assume split attention adapters include q_proj and v_proj. +Complete missing projections with zero factors before their normal packing. Only the SGLang child opts into this hook. It waits for the engine's normal import rather than importing CUDA or SGLang during Python site initialization. @@ -15,6 +17,7 @@ import builtins import importlib.metadata import logging +import re import sys from collections.abc import Mapping, Sequence from functools import wraps @@ -22,13 +25,98 @@ from typing import Any _AFFECTED_VERSIONS = frozenset({"0.5.10", "0.5.10.post1"}) +_PARTIAL_QKV_VERSIONS = _AFFECTED_VERSIONS | {"0.5.20"} +_ADAPTER_MODULE = "sglang.srt.lora.lora" _MANAGER_MODULE = "sglang.srt.lora.lora_manager" _POOL_MODULE = "sglang.srt.lora.mem_pool" _CLASS_MARKER = "_sie_lora_slot_release_compat" +_ADAPTER_MARKER = "_sie_lora_partial_qkv_compat" _IMPORT_MARKER = "_sie_lora_slot_release_deferred" +_SPLIT_QKV_WEIGHT = re.compile(r"((?:.*\.)?)([qkv])_proj\.(lora_[AB]\.weight)") logger = logging.getLogger(__name__) +def _positive_attention_dimension(value: Any) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError("Partial QKV LoRA requires standard positive attention dimensions") + return value + + +def _complete_partial_qkv(adapter: Any, weight_names: list[str], weights: dict[str, Any]) -> list[str]: + """Complete standard attention factors on CPU, before the first GPU use. + + Keep every supplied factor unchanged. A missing projection gets zero A and + B factors, with B sized from attention heads rather than hidden_size (GQA + and explicit head_dim can make those sizes different). Unknown layouts fail + during adapter loading instead of reaching the scheduler with invalid keys. + """ + groups: dict[str, dict[str, dict[str, Any]]] = {} + for name in weight_names: + match = _SPLIT_QKV_WEIGHT.fullmatch(name) + if match is not None: + prefix, projection, factor = match.groups() + groups.setdefault(prefix, {}).setdefault(projection, {})[factor] = weights[name] + + additions: dict[str, Any] = {} + for prefix, projections in groups.items(): + if any(f"{prefix}qkv_proj.{factor}" in weights for factor in ("lora_A.weight", "lora_B.weight")): + raise ValueError("LoRA cannot mix fused and split QKV factors for one attention projection") + if any(set(factors) != {"lora_A.weight", "lora_B.weight"} for factors in projections.values()): + raise ValueError("Split QKV LoRA projections require both lora_A.weight and lora_B.weight") + if len(projections) == 3: + continue + + config = adapter.base_hf_config + if hasattr(config, "get_text_config"): + config = config.get_text_config() + hidden_size = _positive_attention_dimension(getattr(config, "hidden_size", None)) + query_heads = _positive_attention_dimension(getattr(config, "num_attention_heads", None)) + kv_heads = _positive_attention_dimension(getattr(config, "num_key_value_heads", query_heads)) + rank = _positive_attention_dimension(adapter.config.r) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + if hidden_size % query_heads: + raise ValueError("Partial QKV LoRA requires an explicit head_dim for this attention layout") + head_dim = hidden_size // query_heads + head_dim = _positive_attention_dimension(head_dim) + output_dims = {"q": query_heads * head_dim, "k": kv_heads * head_dim, "v": kv_heads * head_dim} + + for projection, factors in projections.items(): + a, b = factors["lora_A.weight"], factors["lora_B.weight"] + if tuple(a.shape) != (rank, hidden_size) or tuple(b.shape) != (output_dims[projection], rank): + raise ValueError("Partial QKV LoRA tensor shapes do not match the base attention dimensions") + reference = next(iter(projections.values())) + for projection in ("q", "k", "v"): + if projection not in projections: + additions[f"{prefix}{projection}_proj.lora_A.weight"] = reference["lora_A.weight"].new_zeros( + (rank, hidden_size) + ) + additions[f"{prefix}{projection}_proj.lora_B.weight"] = reference["lora_B.weight"].new_zeros( + (output_dims[projection], rank) + ) + + # Validate every group before changing the input mapping. + weights.update(additions) + return [*weight_names, *additions] + + +def _patch_adapter_module(module: ModuleType) -> None: + adapter = getattr(module, "LoRAAdapter", None) + if adapter is None: + raise RuntimeError("SGLang does not expose the expected LoRAAdapter") + if getattr(adapter, _ADAPTER_MARKER, False): + return + original_normalize = adapter.normalize_qkv_proj + + @wraps(original_normalize) + def normalize_with_partial_qkv(self: Any, weight_names: list[str], weights: dict[str, Any]) -> Any: + names = _complete_partial_qkv(self, weight_names, weights) + return original_normalize(self, names, weights) + + adapter.normalize_qkv_proj = normalize_with_partial_qkv + setattr(adapter, _ADAPTER_MARKER, True) + + def _release_slot(pool: Any, lora_id: str, empty_slot: Any) -> None: """Forget one drained adapter, preserving the base model's ``None`` UID.""" if lora_id is None: @@ -75,17 +163,29 @@ def unload_with_slot_release(self: Any, lora_ref: Any) -> Any: setattr(manager, _CLASS_MARKER, True) +def _patch_ready_modules(*, release_slots: bool) -> bool: + adapter = sys.modules.get(_ADAPTER_MODULE) + if not isinstance(adapter, ModuleType) or not hasattr(adapter, "LoRAAdapter"): + return False + _patch_adapter_module(adapter) + if release_slots: + manager = sys.modules.get(_MANAGER_MODULE) + if not isinstance(manager, ModuleType) or not hasattr(manager, "LoRAManager"): + return False + _patch_manager_module(manager) + return True + + def install_lora_compat() -> bool: - """Install the exact-version backport; return False for unaffected engines.""" + """Install only reviewed version backports; leave unknown engines native.""" try: version = importlib.metadata.version("sglang") except importlib.metadata.PackageNotFoundError: return False - if version not in _AFFECTED_VERSIONS: + if version not in _PARTIAL_QKV_VERSIONS: return False - loaded = sys.modules.get(_MANAGER_MODULE) - if isinstance(loaded, ModuleType) and hasattr(loaded, "LoRAManager"): - _patch_manager_module(loaded) + release_slots = version in _AFFECTED_VERSIONS + if _patch_ready_modules(release_slots=release_slots): return True current_import = builtins.__import__ @@ -100,14 +200,16 @@ def deferred_import( level: int = 0, ) -> ModuleType: module = current_import(name, globals, locals, fromlist, level) - loaded_module = sys.modules.get(_MANAGER_MODULE) - if not isinstance(loaded_module, ModuleType) or not hasattr(loaded_module, "LoRAManager"): - return module + complete = False try: - _patch_manager_module(loaded_module) - finally: + complete = _patch_ready_modules(release_slots=release_slots) + except Exception: if builtins.__import__ is deferred_import: setattr(builtins, "__import__", current_import) # noqa: B010 + raise + finally: + if complete and builtins.__import__ is deferred_import: + setattr(builtins, "__import__", current_import) # noqa: B010 return module setattr(deferred_import, _IMPORT_MARKER, True) diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_compat/sitecustomize.py b/packages/sie_server/src/sie_server/adapters/sglang/_compat/sitecustomize.py index 05cb6af9b..83b0e2a9a 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_compat/sitecustomize.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_compat/sitecustomize.py @@ -23,8 +23,9 @@ bounds the digit runs of numbers in XGrammar's JSON-schema grammars (see the section at the end of this module). -``SIE_SGLANG_LORA_COMPAT`` enables the exact-version LoRA GPU slot release -backport needed for runtime eviction on the pinned SGLang 0.5.10 engine. +``SIE_SGLANG_LORA_COMPAT`` enables exact-version partial QKV LoRA normalization +on both pinned engines and the GPU slot release backport needed for runtime +eviction on SGLang 0.5.10. ``SIE_SGLANG_CHAT_TOOLS_AS_SENT`` makes the chat route render a request's tools as sent (see ``sie_chat_tools``). diff --git a/packages/sie_server/tests/adapters/test_sglang_lora_compat.py b/packages/sie_server/tests/adapters/test_sglang_lora_compat.py index cb3b73d46..faf5f52a1 100644 --- a/packages/sie_server/tests/adapters/test_sglang_lora_compat.py +++ b/packages/sie_server/tests/adapters/test_sglang_lora_compat.py @@ -1,9 +1,10 @@ -"""The pinned engine must reuse GPU slots across repeated LoRA replacements.""" +"""Pinned LoRA loading preserves PEFT factors and reuses drained GPU slots.""" from __future__ import annotations import builtins import importlib.util +import itertools import os import subprocess import sys @@ -12,6 +13,7 @@ from typing import Any import pytest +import torch _HOOK_PATH = Path(__file__).parents[2] / "src/sie_server/adapters/sglang/_compat/sie_lora_compat.py" @@ -65,7 +67,37 @@ def unload_lora_adapter(self, ref: Any) -> SimpleNamespace: pool_module.__dict__["EMPTY_SLOT"] = empty_slot monkeypatch.setitem(sys.modules, manager_module.__name__, manager_module) monkeypatch.setitem(sys.modules, pool_module.__name__, pool_module) - return SimpleNamespace(manager=Manager, module=manager_module, empty=empty_slot) + adapter_module = ModuleType("sglang.srt.lora.lora") + + class Adapter: + """Pinned engine's packing contract, without importing its CUDA runtime. + + The exact 0.5.10.post1 and 0.5.20 wheel implementations were separately + reproduced: split weights are packed only when q_proj is present, and + q_proj also requires v_proj. The hook must supply every absent slice. + """ + + def __init__(self, config: Any, rank: int = 4) -> None: + self.base_hf_config = config + self.config = SimpleNamespace(r=rank) + self.calls = 0 + + def normalize_qkv_proj(self, names: list[str], weights: dict[str, torch.Tensor]) -> None: + self.calls += 1 + for name in names: + if ".q_proj." in name: + parts = [name.replace(".q_proj.", f".{projection}_proj.") for projection in ("q", "k", "v")] + weights[name.replace(".q_proj.", ".qkv_proj.")] = torch.cat([weights[part] for part in parts]) + for part in parts: + weights.pop(part) + elif ".qkv_proj." in name and "lora_A" in name: + weights[name] = weights[name].repeat(3, 1) + + adapter_module.__dict__["LoRAAdapter"] = Adapter + monkeypatch.setitem(sys.modules, adapter_module.__name__, adapter_module) + return SimpleNamespace( + manager=Manager, module=manager_module, empty=empty_slot, adapter=Adapter, adapter_module=adapter_module + ) def test_repeated_unload_reuses_slot_without_evicting_base(hook: ModuleType, engine: SimpleNamespace) -> None: @@ -124,14 +156,26 @@ def test_base_slot_is_never_released(hook: ModuleType, engine: SimpleNamespace) assert manager.memory_pool.uid_to_buffer_id == {None: 1} -@pytest.mark.parametrize("version", ["0.5.20", "0.6.0", "0.5.11"]) +@pytest.mark.parametrize("version", ["0.6.0", "0.5.11"]) def test_other_engine_versions_keep_native_method( hook: ModuleType, engine: SimpleNamespace, monkeypatch: pytest.MonkeyPatch, version: str ) -> None: monkeypatch.setattr(hook.importlib.metadata, "version", lambda _: version) original = engine.manager.unload_lora_adapter + original_normalize = engine.adapter.normalize_qkv_proj assert hook.install_lora_compat() is False assert engine.manager.unload_lora_adapter is original + assert engine.adapter.normalize_qkv_proj is original_normalize + + +def test_cuda13_normalization_keeps_native_unload( + hook: ModuleType, engine: SimpleNamespace, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(hook.importlib.metadata, "version", lambda _: "0.5.20") + original = engine.manager.unload_lora_adapter + assert hook.install_lora_compat() + assert engine.manager.unload_lora_adapter is original + assert getattr(engine.adapter, hook._ADAPTER_MARKER) def test_exact_version_install_is_idempotent( @@ -140,8 +184,10 @@ def test_exact_version_install_is_idempotent( monkeypatch.setattr(hook.importlib.metadata, "version", lambda _: "0.5.10.post1") assert hook.install_lora_compat() first = engine.manager.unload_lora_adapter + first_normalize = engine.adapter.normalize_qkv_proj assert hook.install_lora_compat() assert engine.manager.unload_lora_adapter is first + assert engine.adapter.normalize_qkv_proj is first_normalize def test_deferred_import_restores_existing_hook( @@ -161,6 +207,7 @@ def test_deferred_import_restores_existing_hook( builtins.__import__("math") assert builtins.__import__ is previous_import assert getattr(engine.manager, hook._CLASS_MARKER) + assert getattr(engine.adapter, hook._ADAPTER_MARKER) def test_child_sitecustomize_installs_the_hook(tmp_path: Path) -> None: @@ -172,8 +219,12 @@ def test_child_sitecustomize_installs_the_hook(tmp_path: Path) -> None: metadata.mkdir() (metadata / "METADATA").write_text("Name: sglang\nVersion: 0.5.10.post1\n") (tmp_path / "sglang/srt/lora/mem_pool.py").write_text("EMPTY_SLOT = object()\n") + (tmp_path / "sglang/srt/lora/lora.py").write_text( + "class LoRAAdapter:\n def normalize_qkv_proj(self, names, weights):\n raise NotImplementedError\n" + ) (tmp_path / "sglang/srt/lora/lora_manager.py").write_text( "from sglang.srt.lora.mem_pool import EMPTY_SLOT\n" + "from sglang.srt.lora.lora import LoRAAdapter\n" "class LoRAManager:\n" " def unload_lora_adapter(self, ref):\n" " raise NotImplementedError\n" @@ -188,7 +239,9 @@ def test_child_sitecustomize_installs_the_hook(tmp_path: Path) -> None: "import sys\n" "assert 'sglang.srt.lora.lora_manager' not in sys.modules\n" "from sglang.srt.lora.lora_manager import LoRAManager\n" - "assert LoRAManager._sie_lora_slot_release_compat\n", + "from sglang.srt.lora.lora import LoRAAdapter\n" + "assert LoRAManager._sie_lora_slot_release_compat\n" + "assert LoRAAdapter._sie_lora_partial_qkv_compat\n", ], env=environment, capture_output=True, @@ -198,3 +251,169 @@ def test_child_sitecustomize_installs_the_hook(tmp_path: Path) -> None: ) assert result.returncode == 0, result.stderr assert "Error in sitecustomize" not in result.stderr + + +_SUBSETS = [subset for size in (1, 2, 3) for subset in itertools.combinations(("q", "k", "v"), size)] + + +def _weights( + projections: tuple[str, ...], *, query_dim: int, kv_dim: int, prefix: str = "base.model.layers.0.attn." +) -> dict[str, torch.Tensor]: + generator = torch.Generator().manual_seed(20261024) + result = {} + for projection in projections: + result[f"{prefix}{projection}_proj.lora_A.weight"] = torch.randn( + 4, 12, generator=generator, dtype=torch.float64 + ) + result[f"{prefix}{projection}_proj.lora_B.weight"] = torch.randn( + query_dim if projection == "q" else kv_dim, 4, generator=generator, dtype=torch.float64 + ) + return result + + +@pytest.mark.parametrize("projections", _SUBSETS) +@pytest.mark.parametrize(("kv_heads", "head_dim"), [(3, None), (1, None), (1, 8)]) +def test_packed_qkv_matches_independent_peft_products( + hook: ModuleType, engine: SimpleNamespace, projections: tuple[str, ...], kv_heads: int, head_dim: int | None +) -> None: + config = SimpleNamespace(hidden_size=12, num_attention_heads=3, num_key_value_heads=kv_heads, head_dim=head_dim) + query_dim, kv_dim = 3 * (head_dim or 4), kv_heads * (head_dim or 4) + weights = _weights(projections, query_dim=query_dim, kv_dim=kv_dim) + original = {name: value.clone() for name, value in weights.items()} + prefix = "base.model.layers.0.attn." + x = torch.arange(36, dtype=torch.float64).reshape(3, 12) / 36 + expected = [] + for projection, dim in zip(("q", "k", "v"), (query_dim, kv_dim, kv_dim), strict=True): + expected.append( + x + @ original[f"{prefix}{projection}_proj.lora_A.weight"].T + @ original[f"{prefix}{projection}_proj.lora_B.weight"].T + if projection in projections + else torch.zeros(3, dim, dtype=torch.float64) + ) + + hook._patch_adapter_module(engine.adapter_module) + adapter = engine.adapter(config) + adapter.normalize_qkv_proj(list(weights), weights) + assert adapter.calls == 1 + assert set(weights) == {f"{prefix}qkv_proj.lora_A.weight", f"{prefix}qkv_proj.lora_B.weight"} + a = weights[f"{prefix}qkv_proj.lora_A.weight"] + b = weights[f"{prefix}qkv_proj.lora_B.weight"] + assert a.shape == (12, 12) + assert b.shape == (query_dim + 2 * kv_dim, 4) + projected = x @ a.T + offset = 0 + actual = [] + for index, dim in enumerate((query_dim, kv_dim, kv_dim)): + actual.append(projected[:, index * 4 : (index + 1) * 4] @ b[offset : offset + dim].T) + offset += dim + torch.testing.assert_close(torch.cat(actual, dim=1), torch.cat(expected, dim=1), rtol=1e-12, atol=1e-12) + + +def test_partial_groups_are_completed_independently(hook: ModuleType, engine: SimpleNamespace) -> None: + config = SimpleNamespace(hidden_size=12, num_attention_heads=3, num_key_value_heads=1) + weights = _weights(("q",), query_dim=12, kv_dim=4) + weights.update(_weights(("v",), query_dim=12, kv_dim=4, prefix="base.model.layers.1.attn.")) + hook._patch_adapter_module(engine.adapter_module) + engine.adapter(config).normalize_qkv_proj(list(weights), weights) + assert len(weights) == 4 + for layer, active_slice in ((0, 0), (1, 2)): + a = weights[f"base.model.layers.{layer}.attn.qkv_proj.lora_A.weight"] + for index in range(3): + assert bool(torch.count_nonzero(a[index * 4 : (index + 1) * 4])) is (index == active_slice) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_padding_preserves_factor_dtype_device_and_values( + hook: ModuleType, engine: SimpleNamespace, dtype: torch.dtype +) -> None: + config = SimpleNamespace(hidden_size=12, num_attention_heads=3, num_key_value_heads=1, head_dim=8) + weights = {name: value.to(dtype) for name, value in _weights(("v",), query_dim=24, kv_dim=8).items()} + original = dict(weights) + hook._complete_partial_qkv(engine.adapter(config), list(weights), weights) + assert all(value.dtype == dtype and value.device.type == "cpu" for value in weights.values()) + assert all(weights[name] is value for name, value in original.items()) + assert all(torch.count_nonzero(value) == 0 for name, value in weights.items() if name not in original) + + +@pytest.mark.parametrize( + "problem", ["missing_a", "missing_b", "wrong_rank", "wrong_input", "wrong_output", "fused_collision"] +) +def test_invalid_partial_factors_fail_before_mutating_weights( + hook: ModuleType, engine: SimpleNamespace, problem: str +) -> None: + config = SimpleNamespace(hidden_size=12, num_attention_heads=3, num_key_value_heads=1) + weights = _weights(("v",), query_dim=12, kv_dim=4) + a_name, b_name = sorted(weights) + if problem == "missing_a": + weights.pop(a_name) + elif problem == "missing_b": + weights.pop(b_name) + elif problem == "wrong_rank": + weights[a_name] = torch.zeros(8, 12) + elif problem == "wrong_input": + weights[a_name] = torch.zeros(4, 13) + elif problem == "wrong_output": + weights[b_name] = torch.zeros(12, 4) + else: + weights[a_name.replace("v_proj", "qkv_proj")] = torch.zeros(4, 12) + before = dict(weights) + with pytest.raises(ValueError, match=r"QKV LoRA|fused and split QKV"): + hook._complete_partial_qkv(engine.adapter(config), list(weights), weights) + assert weights.keys() == before.keys() + assert all(weights[name] is value for name, value in before.items()) + + +@pytest.mark.parametrize( + "config", + [ + SimpleNamespace(), + SimpleNamespace(hidden_size=13, num_attention_heads=3), + SimpleNamespace(hidden_size=12, num_attention_heads=3, head_dim=False), + ], +) +def test_unknown_attention_dimensions_reject_partial_adapters( + hook: ModuleType, engine: SimpleNamespace, config: SimpleNamespace +) -> None: + weights = _weights(("v",), query_dim=12, kv_dim=4) + with pytest.raises(ValueError, match="Partial QKV LoRA"): + hook._complete_partial_qkv(engine.adapter(config), list(weights), weights) + + +def test_full_and_fused_qkv_keep_native_packing(hook: ModuleType, engine: SimpleNamespace) -> None: + hook._patch_adapter_module(engine.adapter_module) + adapter = engine.adapter(SimpleNamespace()) # Full factors do not need dimension inference. + weights = _weights(("q", "k", "v"), query_dim=24, kv_dim=8) + adapter.normalize_qkv_proj(list(weights), weights) + assert len(weights) == 2 + prefix = "base.model.layers.0.attn.qkv_proj." + fused_a, fused_b = torch.ones(4, 12), torch.ones(40, 4) + fused = {prefix + "lora_A.weight": fused_a, prefix + "lora_B.weight": fused_b} + adapter.normalize_qkv_proj(list(fused), fused) + assert torch.equal(fused[prefix + "lora_A.weight"], fused_a.repeat(3, 1)) + assert fused[prefix + "lora_B.weight"] is fused_b + + +def test_no_qkv_weights_do_not_require_attention_config(hook: ModuleType, engine: SimpleNamespace) -> None: + weights = {"base.model.layers.0.attn.o_proj.lora_A.weight": torch.ones(4, 12)} + before = dict(weights) + assert hook._complete_partial_qkv(engine.adapter(SimpleNamespace()), list(weights), weights) == list(weights) + assert weights == before + + +def test_unconfigured_child_does_not_import_lora_hook() -> None: + environment = {key: value for key, value in os.environ.items() if not key.startswith("SIE_SGLANG_")} + environment["PYTHONPATH"] = str(_HOOK_PATH.parent) + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; assert 'sie_lora_compat' not in sys.modules; assert 'torch' not in sys.modules", + ], + env=environment, + capture_output=True, + text=True, + check=False, + timeout=15, + ) + assert result.returncode == 0, result.stderr