From c4e61133fbb5587e0cf8972ce32ffac4e37fa739 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 12 Sep 2026 14:55:24 +0800 Subject: [PATCH] [Edge][CPU] Fused norms off by default, and register the switches this branch added 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) --- docs/configuration/environment_variables.md | 4 ++ tests/config/test_environment_variables.py | 5 +- tests/edge/test_memory.py | 67 +++++++++++++++++++ .../model_executor/layers/test_cpu_ir_ops.py | 16 ++++- .../config/environment_variable_inventory.py | 4 ++ vllm_omni/model_executor/layers/cpu_ir_ops.py | 18 +++-- 6 files changed, 107 insertions(+), 7 deletions(-) create mode 100644 tests/edge/test_memory.py diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 41e6e456d..1420e5406 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -84,6 +84,10 @@ depends on the installed kernels and model path. | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental | +| `VLLM_OMNI_AOT_CACHE_SHIM` | `1` enables (default); `0` disables | CPU only; installed at `import vllm_omni` (`vllm_omni/patch.py`) | Environment-only setting. vLLM's CPU backend cannot load its own AOT-compiled model: `_try_load_aot_compiled_fn` calls `finalize_loading` on an artifact that is a plain function there, so the `AttributeError` is treated as a cache miss and the model is recompiled on every start. The shim attaches the missing no-op; measured warm start 49.8 s -> 14.9 s on Spark-X2.5-1.7B. Set `0` to restore stock behaviour. | Experimental | +| `VLLM_OMNI_CPU_FUSED_NORMS` | `1` enables; default off | CPU only; read by `vllm_omni/model_executor/layers/cpu_ir_ops.py` | Environment-only setting. Ranks vLLM's own `_C::rms_norm` / `_C::fused_add_rms_norm` kernels ahead of the native path (vLLM gates them on `GPGPU_DEVICE`, so CPU never uses them). Off by default because it lost in situ: 3.6x and 4.5x faster in isolation, but 72.7 tok/s against 74.0 end to end, since inductor already fuses the native norm into the residual add. | Experimental | +| `VLLM_OMNI_CPU_RELEASE_MEMORY` | `1` enables (default); `0` disables | CPU only; hooks `CPUWorker.load_model` and `compile_or_warm_up_model` (`vllm_omni/edge/memory.py`) | Environment-only setting. Returns freed weight-loading arenas to the OS. Loading a quantized checkpoint unpacks 4-bit codes to int32 (four bytes per weight) to repack them; the intermediates are freed but the allocator keeps them. Measured 7338 -> 3510 MB of anonymous RSS on Spark-X2.5-1.7B W4A8. | Experimental | +| `VLLM_OMNI_SPARK_GATE_REDUCTION` | `auto` (default), `1`, `0` | Spark-X2.5 attention (`vllm_omni/model_executor/models/spark2_5/spark2_5.py`) | Environment-only setting. `auto` uses a broadcast-and-reduce for the head-wise gate on CPU and a matmul elsewhere. A 2048x8 projection cannot amortize a thread fan-out: as `aten::mm` it costs 1.155 ms of a decode step, as a reduction 0.444 ms. Override to pin either form when measuring on new hardware. | Experimental | | `VLLM_OMNI_CPU_INDUCTOR` | `1` enables; default off | CPU OmniPlatform (`VLLM_TARGET_DEVICE=cpu`); read when the platform reports `supports_torch_inductor` | Environment-only setting. Turns on torch.compile/inductor for the CPU AR model runner (measured variant, see `benchmarks/tts/cpu_thread_sweep.py`). | Experimental | | `VLLM_OMNI_HW_PROFILE` | Filesystem path to a hardware-profile JSON | `deploy_profile="auto"` hardware adaptation (`vllm_omni/edge/hardware_probe.py`); read when the profile is loaded | Environment-only override of the live probe, used for emulation and tests; an unreadable file raises. | Experimental | | `VLLM_OMNI_INIT_TIMELINE` | Filesystem path (JSON) | Engine init timeline (`vllm_omni/engine/init_timeline.py`); read when `AsyncOmniEngine` starts | Environment-only setting. When set, init phases from the orchestrator and every stage worker are merged and written to this path; a partial timeline is written on init timeout. | Diagnostic | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index ea253aaf8..a9bc766fc 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 28, + # 28 -> 32: the edge branch's four CPU switches + # (VLLM_OMNI_AOT_CACHE_SHIM, _CPU_FUSED_NORMS, _CPU_RELEASE_MEMORY, + # _SPARK_GATE_REDUCTION), each an on/off for a measured change. + EnvironmentVariableCategory.PUBLIC_OMNI: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 62, diff --git a/tests/edge/test_memory.py b/tests/edge/test_memory.py new file mode 100644 index 000000000..bb205dcee --- /dev/null +++ b/tests/edge/test_memory.py @@ -0,0 +1,67 @@ +# 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) diff --git a/tests/model_executor/layers/test_cpu_ir_ops.py b/tests/model_executor/layers/test_cpu_ir_ops.py index d89436a65..f37d763cb 100644 --- a/tests/model_executor/layers/test_cpu_ir_ops.py +++ b/tests/model_executor/layers/test_cpu_ir_ops.py @@ -79,11 +79,25 @@ def test_declines_arguments_the_kernel_cannot_take(): assert not impl.supports_args(x, w.float(), 1e-6) -def test_priority_patch_is_idempotent_and_puts_us_first(): +def test_priority_patch_is_opt_in_and_idempotent(monkeypatch): + """Off by default, because that is what the measurement says. + + In isolation these kernels are 3.6x and 4.5x faster than the native path; + in the compiled model they measured 72.7 tok/s against 74.0, because + inductor already fuses the native norm into the residual add. So the + registration always happens -- the provider stays selectable through + `kernel_config.ir_op_priority` -- but the default ranking does not change + unless it is asked for. + """ if not cpu_ir_ops.CPU_FUSED_NORMS: pytest.skip(_REASON) from vllm.platforms.cpu import CpuPlatform + monkeypatch.delenv("VLLM_OMNI_CPU_FUSED_NORMS", raising=False) + cpu_ir_ops.prefer_cpu_fused_norms() + assert CpuPlatform.get_default_ir_op_priority(None).rms_norm == ["native"] + + monkeypatch.setenv("VLLM_OMNI_CPU_FUSED_NORMS", "1") cpu_ir_ops.prefer_cpu_fused_norms() cpu_ir_ops.prefer_cpu_fused_norms() # second call must not double-wrap priority = CpuPlatform.get_default_ir_op_priority(None) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 400d0b823..b5ce2872b 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -65,7 +65,10 @@ _PUBLIC_OMNI = ( "SPEAKER_SAMPLES_DIR", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", # edge branch (vllm_omni/edge, init timeline, CPU platform) + "VLLM_OMNI_AOT_CACHE_SHIM", + "VLLM_OMNI_CPU_FUSED_NORMS", "VLLM_OMNI_CPU_INDUCTOR", + "VLLM_OMNI_CPU_RELEASE_MEMORY", "VLLM_OMNI_HW_PROFILE", "VLLM_OMNI_INIT_TIMELINE", "VLLM_OMNI_STEP_STATS_DIR", @@ -77,6 +80,7 @@ _PUBLIC_OMNI = ( "VLLM_OMNI_SERVER_STORAGE__PATH", "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", "VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP", + "VLLM_OMNI_SPARK_GATE_REDUCTION", "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "VLLM_OMNI_STORAGE_PATH", "VLLM_OMNI_USE_QUACK_FP8", diff --git a/vllm_omni/model_executor/layers/cpu_ir_ops.py b/vllm_omni/model_executor/layers/cpu_ir_ops.py index 8fe76107c..fb03224c0 100644 --- a/vllm_omni/model_executor/layers/cpu_ir_ops.py +++ b/vllm_omni/model_executor/layers/cpu_ir_ops.py @@ -109,15 +109,23 @@ def fused_add_rms_norm( def prefer_cpu_fused_norms() -> None: - """Make the CPU platform rank these kernels ahead of the native path. + """Rank these kernels ahead of the native path -- opt-in only. - The platform hook only fills in ops the user left unset, so an explicit - ``kernel_config.ir_op_priority`` still wins. + The implementations are always registered (so they can be selected through + ``kernel_config.ir_op_priority`` like any other provider); this function + only changes the *default* ranking, and it does nothing unless + ``VLLM_OMNI_CPU_FUSED_NORMS=1``. The platform hook fills in ops the user + left unset, so an explicit priority still wins either way. """ if not CPU_FUSED_NORMS: return - if os.environ.get("VLLM_OMNI_CPU_FUSED_NORMS", "1") != "1": - logger.info("CPU fused RMSNorm kernels disabled by environment.") + # Off by default, and the default is the measurement: in isolation these + # kernels are 3.6x and 4.5x faster than the native path, but in the + # compiled model they measured 72.7 tok/s against 74.0, because inductor + # was already fusing the native norm into the residual add and an opaque + # call breaks that fusion. Set VLLM_OMNI_CPU_FUSED_NORMS=1 to try them on + # a machine where that trade may land differently. + if os.environ.get("VLLM_OMNI_CPU_FUSED_NORMS", "0") != "1": return from vllm.platforms.cpu import CpuPlatform