diff --git a/tests/models/qwen4_exp/test_ple_hybrid_default.py b/tests/models/qwen4_exp/test_ple_hybrid_default.py new file mode 100644 index 000000000..ef17fd3cd --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_hybrid_default.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + +import vllm.models.qwen4_exp.common.ple as ple + +GIB = 1 << 30 + + +def _checkpoint(tmp_path, table_gib: int): + weight_map = { + "model.language_model.layers.0.ple.ple_embedding.ngram_embedding.weight": ( + "model-plefp8-00000.safetensors" + ), + "model.language_model.layers.0.mlp.gate.weight": "model-bf16-00001.safetensors", + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + with open(tmp_path / "model-plefp8-00000.safetensors", "wb") as f: + f.truncate(table_gib * GIB) # sparse: only the size matters + return str(tmp_path) + + +def test_table_bytes_counts_only_ple_shards(tmp_path): + path = _checkpoint(tmp_path, 3) + assert ple.ple_table_bytes(path) == 3 * GIB + assert ple.ple_table_bytes("org/remote-model") is None + + +def test_hybrid_default_uses_physical_memory(monkeypatch, tmp_path): + path = _checkpoint(tmp_path, 40) + # MemAvailable is never consulted: processes resolving the defaults after + # others pinned their shares must reach the same placement. + monkeypatch.setattr(ple, "available_host_bytes", lambda: 0) + # Reserve 64/4 = 16 GiB, fair share (64 - 16) / 4 - 2 = 10 GiB per rank, + # table share 40 / 4 = 10 GiB. + monkeypatch.setattr(ple, "total_host_bytes", lambda: 64 * GIB) + assert ple.default_hybrid_host_gib(path, 4) == 10.0 + # 0.75 * 40 GiB + 8 GiB reserve = 38 GiB is the smallest host for hybrid. + monkeypatch.setattr(ple, "total_host_bytes", lambda: 37 * GIB) + assert ple.default_hybrid_host_gib(path, 4) is None + assert ple.default_hybrid_host_gib("org/remote-model", 4) is None + + +def test_apply_keeps_an_explicit_host_budget(monkeypatch, tmp_path): + path = _checkpoint(tmp_path, 40) + monkeypatch.setattr(ple, "total_host_bytes", lambda: 64 * GIB) + defaults = {} + assert ple.apply_hybrid_ple_defaults(defaults, path, 4) == 10.0 + assert defaults == { + "VLLM_SM70_QWEN38_HYBRID_PLE": "1", + "VLLM_QWEN4EXP_PLE_HOST_GIB": "10.0", + } + defaults = {"VLLM_QWEN4EXP_PLE_HOST_GIB": "7"} + ple.apply_hybrid_ple_defaults(defaults, path, 4) + assert defaults["VLLM_QWEN4EXP_PLE_HOST_GIB"] == "7" + defaults = {} + assert ple.apply_hybrid_ple_defaults(defaults, None, 4) is None + assert defaults == {"VLLM_SM70_QWEN38_HYBRID_PLE": "0"} diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py index a4c832dc8..79377d319 100644 --- a/vllm/models/qwen4_exp/common/ple.py +++ b/vllm/models/qwen4_exp/common/ple.py @@ -414,6 +414,68 @@ def cap_host_budget_bytes( return min(budget_bytes, share) +# Host memory kept free beside the hybrid default's pinned PLE share (engine +# processes, checkpoint loading, page cache of the disk-backed prefill path). +HYBRID_DEFAULT_HOST_RESERVE_BYTES = 8 << 30 +# Fraction of the PLE table that must fit in pinned host memory; the hybrid +# placement keeps the rest in device memory beside the weights. +HYBRID_DEFAULT_HOST_FRACTION = 0.75 +# Per-rank margin below the fair host share when the default fixes the budget. +HYBRID_DEFAULT_RANK_MARGIN_BYTES = 2 << 30 + + +def ple_table_bytes(model_path: str | None) -> int | None: + """Checkpoint bytes of the PLE n-gram tables, from a local safetensors index.""" + import json + import os + + try: + with open(os.path.join(str(model_path), "model.safetensors.index.json")) as f: + weight_map = json.load(f)["weight_map"] + files = {f for k, f in weight_map.items() if "ngram_embedding" in k} + return sum(os.path.getsize(os.path.join(str(model_path), f)) for f in files) + except (OSError, KeyError, ValueError, TypeError): + return None + + +def default_hybrid_host_gib(model_path: str | None, tp_size: int) -> float | None: + """Per-rank pinned-host PLE budget for the hybrid default, or None for disk. + + Hybrid decode gathers PLE rows on the GPU from local device and pinned-host + shards instead of waiting each step for the CPU worker's disk lookup + (Flash-Next NVFP4 MTP4 on 4x V100-32GB: 18.3 -> 15.7 ms per round). It is + chosen when physical memory less a reserve holds three quarters of the + table. + + Both the choice and the budget come from physical memory. Every engine + process resolves these defaults, some after others pinned their shares, + and a budget derived per process from MemAvailable left the hybrid lane + crashing in its first decode graph replay; every startup with a fixed + budget was stable. + """ + table = ple_table_bytes(model_path) + total = total_host_bytes() + if not table or total is None: + return None + if total - HYBRID_DEFAULT_HOST_RESERVE_BYTES < table * HYBRID_DEFAULT_HOST_FRACTION: + return None + ranks = max(int(tp_size), 1) + share = (total - ple_host_reserve_bytes(total)) // ranks + budget = min(table // ranks, share - HYBRID_DEFAULT_RANK_MARGIN_BYTES) + return round(budget / (1 << 30), 2) if budget > 0 else None + + +def apply_hybrid_ple_defaults( + defaults, model_path: str | None, tp_size: int +) -> float | None: + """Select hybrid placement when the host fits it; return its budget.""" + host_gib = default_hybrid_host_gib(model_path, tp_size) + defaults["VLLM_SM70_QWEN38_HYBRID_PLE"] = "0" if host_gib is None else "1" + if host_gib is not None and "VLLM_QWEN4EXP_PLE_HOST_GIB" not in defaults: + defaults["VLLM_QWEN4EXP_PLE_HOST_GIB"] = str(host_gib) + return host_gib + + def env_gib_bytes(name: str) -> int | None: """Bytes of a GiB-valued PLE variable, or None when it is unset.""" diff --git a/vllm/platforms/runtime_defaults.py b/vllm/platforms/runtime_defaults.py index 8263dff2b..2843b9268 100644 --- a/vllm/platforms/runtime_defaults.py +++ b/vllm/platforms/runtime_defaults.py @@ -130,16 +130,26 @@ def checkpoint_kv_quant_allowed(cfg: VllmConfig) -> bool: def _apply_sm70_qwen38_disk_ple_defaults( - parallel_config: ParallelConfig, *, defaults -) -> None: - """Keep PLE disk-backed and complete its late-bound parallel config.""" - defaults["VLLM_SM70_QWEN38_HYBRID_PLE"] = "0" + parallel_config: ParallelConfig, *, defaults, model: str | None = None +) -> float | None: + """Keep PLE disk-backed and complete its late-bound parallel config. + + When the host can pin the table of ``model``, decode reads local + device/host shards (hybrid placement) while prefill keeps the CPU/disk + path. Returns the per-rank pinned-host budget, or None for disk only. + """ + from vllm.models.qwen4_exp.common.ple import apply_hybrid_ple_defaults + + host_gib = apply_hybrid_ple_defaults( + defaults, model, parallel_config.tensor_parallel_size + ) defaults["VLLM_PLE_CPU_OFFLOAD"] = "1" defaults["VLLM_PLE_DISK_OFFLOAD"] = "1" # ParallelConfig is validated before these model-aware defaults are # applied, so initialize the endpoint that its validator would have # created for an explicit PLE configuration. parallel_config.ensure_ple_offload_ipc_path() + return host_gib def _ple_disk_cascade_cuda_supported() -> bool: @@ -486,12 +496,16 @@ def apply_runtime_policy_defaults(self): ) ) ): - _apply_sm70_qwen38_disk_ple_defaults( - self.parallel_config, defaults=defaults + host_gib = _apply_sm70_qwen38_disk_ple_defaults( + self.parallel_config, defaults=defaults, model=self.model_config.model ) logger.info_once( - "Auto-enabling disk-mmap PLE for the SM70 Qwen3.8 " - "dual-compile lane: bounded result staging, no resident table." + "Auto-enabling %s PLE for the SM70 Qwen3.8 dual-compile lane: %s.", + "disk-mmap" if host_gib is None else "hybrid", + "bounded result staging, no resident table" + if host_gib is None + else f"decode gathers {host_gib} GiB pinned-host/device shards " + "per rank, prefill reads the disk-backed table", ) if self.speculative_config is not None: policy = self.speculative_config.sm70_dflash2