From dc062c170dfcb0dd3cc9b63bfacd780f32e983af Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:27:35 +0800 Subject: [PATCH] [Perf][SM70] Qwen3.8 dual-compile lane: hybrid PLE by default when the host can pin it (#1340) The dual-compile lane auto-enabled disk-mmap PLE. Every MTP4 decode round then waited about 1 ms for the CPU worker's disk lookup before the target graph could continue (and the automatic default also blocked the PLE cascade as an "existing explicit placement"). When physical host memory, less an 8 GiB reserve, holds three quarters of the checkpoint's PLE n-gram table, the default now selects the existing hybrid placement instead: decode gathers PLE rows from local device and pinned-host shards on the GPU, prefill keeps the CPU/disk path. The per-rank host budget is fixed from physical memory ((total - PLE host reserve) / TP - 2 GiB, capped by the table share). Smaller hosts keep the disk default. Co-authored-by: Claude Opus 5.5 --- .../qwen4_exp/test_ple_hybrid_default.py | 60 ++++++++++++++++++ vllm/models/qwen4_exp/common/ple.py | 62 +++++++++++++++++++ vllm/platforms/runtime_defaults.py | 30 ++++++--- 3 files changed, 144 insertions(+), 8 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_hybrid_default.py 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