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