[Edge][CPU] Return the weight-loading churn to the OS: 7338 -> 3510 MB

Loading a quantized checkpoint allocates far more than the weights occupy.
`_process_gptq_weights_w4a8` unpacks 4-bit codes to int32 -- four bytes per
weight -- permutes, repacks, and drops the intermediates. For a 1.7 B model
that is gigabytes of short-lived allocation, correctly freed, and still
resident afterwards because a caching allocator keeps the arenas.

Measured on Spark-X2.5-1.7B, anonymous RSS:

    path               after load   after release   freed
    W4A8 (GPTQ)          7338 MB       3510 MB     3828 MB   (-52%)
    W4A16 (tinygemm)     5023 MB       3411 MB     1612 MB

The second column is the more useful finding: after the release both paths
land in the same place, so W4A8's apparent 2.3 GB penalty over W4A16 was never
a footprint difference -- it was the repack. W4A8 is therefore faster *and*
level on memory, and fidelity is the only remaining reason to pick W4A16.

Hooked to CPUWorker.load_model and compile_or_warm_up_model, since those churn
separately; `VLLM_OMNI_CPU_RELEASE_MEMORY=0` disables it. CPU only: on GPU the
weights are in device memory and host churn constrains nothing.

`anon_rss_bytes()` ships alongside because the metric matters as much as the
fix. VmRSS counts the mmap'd checkpoint's file pages and moved 1.2-2.0 GB
between *identical* runs on this host; RssAnon repeats to within 0.5 MB. An
earlier claim in this tree that KV sizing saved 3.1 GB came from reading VmRSS
across a comparison that also changed the model file; re-measured on RssAnon,
shrinking the KV cache 32x is worth 188 MB, because its pages are never
touched. The analysis note is corrected accordingly.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

vllm_omni/edge/memory.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Give freed weight-loading memory back to the OS.
+
+Loading a quantized checkpoint churns far more memory than the weights
+occupy. The GPTQ path is the extreme case: `_process_gptq_weights_w4a8`
+unpacks 4-bit codes to **int32 -- four bytes per weight** -- permutes, repacks,
+and drops the intermediates. For a 1.7 B model that is gigabytes of
+short-lived allocation. It is correctly freed, but a caching allocator keeps
+the arenas, so it stays resident for the life of the process.
+
+Measured on Spark-X2.5-1.7B (Xeon 8480C, anonymous RSS, which repeats to
+within 0.5 MB where VmRSS swings 1-2 GB):
+
+    path                after load   after release   freed
+    W4A8 (GPTQ)           7338 MB        3510 MB     3828 MB
+    W4A16 (tinygemm)      5023 MB        3411 MB     1612 MB
+
+Two things worth reading off that table. The release is worth **52%** of
+resident memory on the path we actually want to deploy; and afterwards both
+paths land in the same place, so W4A8's apparent 2.3 GB penalty over W4A16 was
+never a real footprint difference -- it was the repack, held by the allocator.
+
+This is only useful on CPU. On GPU the weights are in device memory and the
+host-side churn is not what constrains anything.
+"""
+
+import ctypes
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# glibc's trim, then tcmalloc's equivalent: vLLM's CPU docs tell you to
+# LD_PRELOAD tcmalloc, so on a correctly configured host the second one is the
+# one that does the work.
+_RELEASERS = (
+    ("libc.so.6", "malloc_trim", (0,)),
+    ("libtcmalloc_minimal.so.4", "MallocExtension_ReleaseFreeMemory", ()),
+    ("libtcmalloc.so.4", "MallocExtension_ReleaseFreeMemory", ()),
+)
+
+
+def anon_rss_bytes() -> int | None:
+    """Resident anonymous memory: allocations, excluding mapped files.
+
+    The right metric for a memory budget. ``VmRSS`` also counts the mmap'd
+    checkpoint's pages, which come and go with page-cache pressure and moved
+    1-2 GB between identical runs on this host.
+    """
+    try:
+        with open("/proc/self/status") as f:
+            for line in f:
+                if line.startswith("RssAnon:"):
+                    return int(line.split()[1]) * 1024
+    except OSError:
+        pass
+    return None
+
+
+def release_free_memory() -> list[str]:
+    """Ask every allocator we can reach to return free arenas to the OS.
+
+    Returns the names of the calls that succeeded. Safe to call repeatedly and
+    safe when none are available -- it simply does nothing.
+    """
+    released: list[str] = []
+    for lib, symbol, args in _RELEASERS:
+        try:
+            fn = getattr(ctypes.CDLL(lib), symbol)
+        except (OSError, AttributeError):
+            continue
+        try:
+            fn(*args)
+            released.append(f"{lib}:{symbol}")
+        except Exception as exc:  # pragma: no cover - defensive
+            logger.debug("%s:%s failed: %s", lib, symbol, exc)
+    return released
+
+
+def release_after_load(context: str = "weight load") -> None:
+    """Release, and say how much came back, at INFO."""
+    before = anon_rss_bytes()
+    released = release_free_memory()
+    if not released:
+        return
+    after = anon_rss_bytes()
+    if before is None or after is None:
+        logger.info("Released allocator arenas after %s via %s", context, released)
+        return
+    logger.info(
+        "Released %.0f MiB after %s (anonymous RSS %.0f -> %.0f MiB) via %s",
+        (before - after) / 2**20,
+        context,
+        before / 2**20,
+        after / 2**20,
+        released,
+    )

vllm_omni/patch.py

@@ -694,3 +694,56 @@ def _patch_cpu_explicit_kv_cache_memory() -> None:
 
 
 _patch_cpu_explicit_kv_cache_memory()
+
+
+def _patch_cpu_release_memory_after_load() -> None:
+    """Return the weight-loading churn to the OS on CPU.
+
+    Loading a quantized checkpoint allocates far more than the weights occupy
+    -- the GPTQ path unpacks 4-bit codes to int32, four bytes per weight, to
+    permute and repack them -- and the intermediates, though freed, stay
+    resident because a caching allocator keeps the arenas. On Spark-X2.5-1.7B
+    that is 3828 MiB held after a W4A8 load and 1612 MiB after W4A16; released,
+    both land at ~3.5 GiB, so W4A8's apparent memory penalty over W4A16 was
+    entirely this and not a real difference in footprint.
+
+    Hooked in two places because they churn separately: after weights are
+    loaded and processed, and after the compile/warm-up pass.
+    """
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_cpu():
+        return
+    if os.environ.get("VLLM_OMNI_CPU_RELEASE_MEMORY", "1") != "1":
+        return
+
+    from vllm.v1.worker.cpu_worker import CPUWorker
+
+    if getattr(CPUWorker, "_omni_release_memory", False):
+        return
+
+    from vllm_omni.edge.memory import release_after_load
+
+    base_load = CPUWorker.load_model
+    base_warm = CPUWorker.compile_or_warm_up_model
+
+    def load_model(self, *args, **kwargs):
+        out = base_load(self, *args, **kwargs)
+        release_after_load("weight load")
+        return out
+
+    def compile_or_warm_up_model(self, *args, **kwargs):
+        out = base_warm(self, *args, **kwargs)
+        release_after_load("warm-up")
+        return out
+
+    CPUWorker.load_model = load_model
+    CPUWorker.compile_or_warm_up_model = compile_or_warm_up_model
+    CPUWorker._omni_release_memory = True
+    _PATCH_LOGGER.info(
+        "CPU weight-load memory release installed (set "
+        "VLLM_OMNI_CPU_RELEASE_MEMORY=0 to disable)."
+    )
+
+
+_patch_cpu_release_memory_after_load()