tzhouamandClaude Opus 5
c4e61133fb
Two corrections and the paperwork they imply.
The fused-RMSNorm registration defaulted to *on*, which contradicts the
measurement it was committed with: those kernels are 3.6x and 4.5x faster in
isolation but measured 72.7 tok/s against 74.0 end to end, because inductor
already fuses the native norm into the residual add and an opaque call breaks
that fusion. The implementations stay registered, so the provider remains
selectable through `kernel_config.ir_op_priority` on a machine where the trade
lands differently; only the default ranking is left alone now, and the test
asserts both states rather than just the opt-in one.
The four switches this branch added (VLLM_OMNI_AOT_CACHE_SHIM,
_CPU_FUSED_NORMS, _CPU_RELEASE_MEMORY, _SPARK_GATE_REDUCTION) were not in the
reviewed environment-variable inventory, so
`test_statically_resolvable_environment_accesses_are_classified` failed. They
are classified now, the reviewed count moves 28 -> 32 with a note saying why,
and each has a row in docs/configuration/environment_variables.md carrying the
measurement that justifies its default.
`VLLM_OMNI_CPU_FAST_SAMPLER` and `VLLM_OMNI_STEP_STATS_SYNC` remain
unclassified and that test still fails on them. They arrived in 6ce8de1b, in
files this branch does not touch, so they are left to whoever owns that work
rather than documented from a guess.
Under `-m "core_model and cpu"`: 2593 passed before these fixes with 9
failures, of which two were mine. The rest are CUDA-only tests in a CPU venv
("Torch not compiled with CUDA enabled") plus test_qnn_export.py, which
another session is actively editing.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Returning freed weight-loading memory to the OS."""
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.edge.memory import (
|
|
anon_rss_bytes,
|
|
release_after_load,
|
|
release_free_memory,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_anon_rss_excludes_mapped_files():
|
|
"""RssAnon, not VmRSS: VmRSS counts the mmap'd checkpoint's pages and moved
|
|
1.2-2.0 GB between identical runs on the host this was developed on."""
|
|
anon = anon_rss_bytes()
|
|
assert anon is not None and anon > 0
|
|
with open("/proc/self/status") as f:
|
|
status = f.read()
|
|
total = int([l for l in status.splitlines() if l.startswith("VmRSS:")][0].split()[1])
|
|
assert anon <= total * 1024
|
|
|
|
|
|
def test_release_reports_what_it_called_and_is_repeatable():
|
|
first = release_free_memory()
|
|
second = release_free_memory()
|
|
assert isinstance(first, list)
|
|
# Idempotent: calling it twice must not raise, and glibc is always there.
|
|
assert first == second
|
|
assert any("malloc_trim" in name for name in first)
|
|
|
|
|
|
def test_release_actually_returns_a_large_freed_block():
|
|
"""The real case is gigabytes of freed repack buffers; a smaller version of
|
|
the same thing should still come back."""
|
|
import ctypes
|
|
|
|
before = anon_rss_bytes()
|
|
blocks = [bytearray(64 * 1024 * 1024) for _ in range(4)] # 256 MiB, touched
|
|
for b in blocks:
|
|
b[::4096] = b"\x01" * (len(b) // 4096)
|
|
peak = anon_rss_bytes()
|
|
assert peak - before > 200 * 2**20, "test did not actually allocate"
|
|
del blocks
|
|
release_free_memory()
|
|
after = anon_rss_bytes()
|
|
# Most of it should be back; allocators keep some, so this is deliberately
|
|
# loose -- the point is that the call does something, not how much.
|
|
assert after < peak - 100 * 2**20, (before, peak, after)
|
|
|
|
|
|
def test_release_after_load_is_safe_when_nothing_is_held():
|
|
release_after_load("unit test") # must not raise
|
|
|
|
|
|
def test_patch_installs_the_hook_on_cpu():
|
|
from vllm.platforms import current_platform
|
|
|
|
if not current_platform.is_cpu():
|
|
pytest.skip("CPU-only patch")
|
|
import vllm_omni.patch # noqa: F401
|
|
from vllm.v1.worker.cpu_worker import CPUWorker
|
|
|
|
assert getattr(CPUWorker, "_omni_release_memory", False)
|