From 43fbd6510bf289dc3fedde089d5a710433fb3eed Mon Sep 17 00:00:00 2001 From: MOVZX Date: Sat, 26 Sep 2026 21:14:50 +0000 Subject: [PATCH] [Qwen4Exp] Serve PLE rows from the loader's mapped checkpoint shards (PR #58815) Adds an opt-in backend, --engram-config '{"host_file_gather": true}', that serves Qwen4Exp PLE rows from the checkpoint files through the page cache instead of allocating the 47.7 GiB table on GPU or in pinned host memory. - Qwen4ExpPLEFileGatherEmbedding allocates no table rows; load_weights keeps the default loader's zero-copy safetensors views of each shard. At bind (post-load hook) every view must be file-backed per /proc/self/maps (HF cache symlinks work; heap/pinned/anonymous memory is rejected) and the shards must cover the whole table. - Each step, Qwen4ExpModelState.prepare_inputs computes the PLE n-gram ids, copies them to pinned host memory, syncs the stream once, issues POSIX_FADV_WILLNEED over the coalesced page ranges of the step's rows, gathers the rows on the host, and copies them into a fixed device staging buffer. The forward reads only the staging buffer, so FULL CUDA graph capture is unaffected. - Rows stay in the checkpoint dtype (FP8 table stays FP8, dequantized on the GPU); dummy loads serve zeros; weight reload is rejected; takes precedence over cpu_offload; rejected on ROCm and with elastic EP. Co-Authored-By: Hinata Signed-off-by: MOVZX --- docs/features/engram.md | 44 ++- tests/models/qwen4_exp/test_config.py | 35 ++ tests/models/qwen4_exp/test_ple.py | 351 ++++++++++++++++++ tests/test_config.py | 39 +- vllm/config/engram.py | 11 + .../qwen4_exp/common/ngram_embedding.py | 13 +- vllm/models/qwen4_exp/nvidia/model_state.py | 22 +- .../qwen4_exp/nvidia/ngram_embedding.py | 184 ++++++++- 8 files changed, 684 insertions(+), 15 deletions(-) diff --git a/docs/features/engram.md b/docs/features/engram.md index 73dc05e9f1..b3470ac632 100644 --- a/docs/features/engram.md +++ b/docs/features/engram.md @@ -83,10 +83,7 @@ directly from the Triton kernel through unified virtual addressing (UVA), prefetching on a side CUDA stream so transfers overlap with the compute of preceding layers. The freed GPU memory goes to the KV cache. -`cpu_offload` defaults to `true` (on): it follows the -`VLLM_PLE_CPU_OFFLOAD` environment variable, which also defaults to `true`. -An explicit `--engram-config` value takes precedence over the environment -variable. +`cpu_offload` defaults to `true` and is set only through `--engram-config`. To keep the tables resident on GPU instead (e.g. when host memory is scarce): @@ -96,8 +93,40 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \ --engram-config.cpu_offload false ``` -or equivalently `VLLM_PLE_CPU_OFFLOAD=0`. CPU offload requires a GPU with UVA -support; vLLM fails fast when it is unavailable. +CPU offload requires a GPU with UVA support; vLLM fails fast when it is +unavailable. + +## Host file gather (Qwen4Exp) + +`host_file_gather` serves Qwen4Exp PLE rows from the loader's own file +mapping of the safetensors shards instead of allocating the table. Set +`--engram-config '{"host_file_gather": true}'`. It applies only to Qwen4Exp +and takes precedence over `cpu_offload`. + +Rows stay in the checkpoint dtype: an FP8 table stays FP8 and is dequantized +on the GPU, and an unquantized table must match `--dtype`. CPU offload +converts on load; this mode rejects the mismatch at load. + +Faulted pages are reclaimable page cache, shared across co-located processes. +They appear in each process's RSS as file-backed (`RssFile`), and PSS is +where they are split across processes. A write to the mapping copies that +page into anonymous RSS. Each step pays a host sync, and every TP rank +gathers all rows. Co-located ranks share the cache; ranks on different +machines do not. + +Serving depends on the checkpoint files staying unchanged while the server +runs: an in-place overwrite serves the new bytes, truncation kills the worker +with SIGBUS, and deleting or renaming the files is safe. + +The first touch of a row costs a disk read, and a page evicted under memory +pressure is read again on the next miss. By default no prewarm is needed; for +a fast first request, read the PLE files once after the server has started +and allocated its KV cache. + +Reload is unsupported and fails closed. ROCm and elastic EP reject the key. +Load paths that do not hand over file-backed safetensors tensors fail at bind +(`eager`, `.bin`, runai, `fastsafetensors`, `instanttensor`). The multithread +loader is file-backed and is admitted. `--load-format dummy` serves zeros. ## Data-parallel topologies (DeepSeek V4.1) @@ -118,7 +147,8 @@ vllm serve deepseek-ai/DeepSeek-V4.1-Flash \ ``` `embedding_across_dp` is not supported with elastic expert parallelism yet. -Qwen4Exp PLE tables are ETP-sharded instead and honor `cpu_offload` only. +Qwen4Exp PLE tables are ETP-sharded instead and honor only `cpu_offload` and +`host_file_gather`. ## Sharing host tables across DP replicas diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 9d30365b81..8bc74c5532 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -21,6 +21,7 @@ from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from ...utils import spawn_new_process_for_each_test +from .test_ple import _make_file_gather_embedding def _text_config(**kwargs) -> Qwen4ExpTextConfig: @@ -175,6 +176,7 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None: model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32) model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64) model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32) + model_state.file_gather_modules = [] input_batch = SimpleNamespace( num_reqs=2, @@ -224,6 +226,39 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None: assert model_inputs["ngram_context"].data_ptr() == ngram_context.data_ptr() +def test_qwen4_exp_model_state_stages_file_gather_rows(monkeypatch, tmp_path) -> None: + """Rows are staged from the same PLE inputs the forward receives.""" + module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True) + model_state = object.__new__(Qwen4ExpModelState) + model_state.uses_ngram_embedding = True + model_state.ngram_eos_token_id = 0 + model_state.ngram_context = torch.empty((4, 1), dtype=torch.int32) + model_state.ngram_context_offsets = torch.arange(-1, 0, dtype=torch.int64) + model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32) + model_state.file_gather_modules = [module] + input_batch = SimpleNamespace( + num_reqs=2, + num_reqs_after_padding=2, + idx_mapping=torch.tensor([0, 1]), + query_start_loc=torch.tensor([0, 4, 6], dtype=torch.int32), + input_ids=torch.tensor([5, 9, 3, 7, 2, 11]), + ) + req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([2, 1])), + all_token_ids=SimpleNamespace(gpu=torch.tensor([[1, 8, 6], [4, 3, 3]])), + ) + + with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}): + model_inputs = model_state.prepare_inputs(input_batch, req_states) + + ids = module.compute_ngram_ids( + input_batch.input_ids, + model_inputs["query_start_loc"], + model_inputs["ngram_context"], + ) + assert torch.equal(module.ngram_embedding._staging[:6], table[ids]) + + def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None: model_state = object.__new__(Qwen4ExpModelState) model_state.uses_ngram_embedding = True diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 85928d293b..88aae40af1 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1,19 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap +import os from dataclasses import dataclass from functools import partial from itertools import accumulate +from pathlib import Path from types import SimpleNamespace import pytest import torch +from safetensors import safe_open +from safetensors.torch import load, load_file, save_file from torch import nn from torch.nn import functional as F import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module import vllm.model_executor.parameter as parameter_module import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module +import vllm.models.qwen4_exp.nvidia.ngram_embedding as nvidia_ngram_embedding_module from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.modelopt import ( @@ -35,6 +41,7 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( Qwen4ExpNGramEmbedding, Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, + Qwen4ExpPLEFileGatherEmbedding, Qwen4ExpPLEFp8EmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, @@ -634,6 +641,350 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: ) +def _patch_ple_construction( + monkeypatch: pytest.MonkeyPatch, + *, + engram_config: SimpleNamespace | None = None, + load_format: str = "auto", +) -> None: + _mock_etp_group(monkeypatch) + for module in (embedding_module, parameter_module): + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + vllm_config = SimpleNamespace( + load_config=SimpleNamespace(load_format=load_format), + engram_config=engram_config, + ) + monkeypatch.setattr( + nvidia_ngram_embedding_module, "get_current_vllm_config", lambda: vllm_config + ) + + +def _make_file_gather_embedding( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + *, + fp8: bool = False, + load_format: str = "auto", + bind: bool = False, +) -> tuple[Qwen4ExpNGramEmbedding, torch.Tensor]: + """Build a two-head, 12-token host-file-gather embedding and its saved table. + + The 21-row table is saved to ``tmp_path / "model.safetensors"`` as four + checkpoint shards of stride 6, so the last shard holds 3 rows. + """ + _patch_ple_construction(monkeypatch, load_format=load_format) + method = ( + Qwen4ExpPLEFp8EmbeddingMethod() + if fp8 + else Qwen4ExpPLEUnquantizedEmbeddingMethod() + ) + module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding) + nn.Module.__init__(module) + module.split_ngram_parts = 4 + module.eos_token_id = 0 + module.ngram_size = 2 + module.heads_per_ngram = 2 + module.register_buffer("layer_multipliers", torch.tensor([3, 5])) + module.register_buffer("ngram_heads_vocab_sizes", torch.tensor([13, 7])) + module.register_buffer("ngram_heads_offsets", torch.tensor([0, 13])) + module.ngram_embedding = Qwen4ExpPLEFileGatherEmbedding( + 21, + 2, + params_dtype=torch.bfloat16, + padding_size=1, + prefix="test.ple_embedding", + embedding_method=method, + num_ngram_heads=2, + max_total_tokens=12, + ) + dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + table = torch.arange(42, dtype=torch.float32).reshape(21, 2).to(dtype) + checkpoint = { + f"ngram_embedding.shard_{index}.weight": shard.clone() + for index, shard in enumerate(table.split(6)) + } + if fp8: + checkpoint["ngram_embedding.weight_scale"] = torch.tensor([0.25]).bfloat16() + save_file(checkpoint, tmp_path / "model.safetensors") + if bind: + _load_file_shards(module, tmp_path / "model.safetensors") + module.ngram_embedding.bind_file_shards() + return module, table + + +def _load_file_shards(module: Qwen4ExpNGramEmbedding, path: Path) -> set[str]: + """Load ``path`` through ``load_weights`` as the lazy safetensors iterator does.""" + with safe_open(path, framework="pt") as checkpoint: + return module.load_weights( + (name, checkpoint.get_tensor(name)) + for name in checkpoint.keys() # noqa: SIM118 + ) + + +def _dev_zero_shards() -> list[tuple[str, torch.Tensor]]: + """Return shards in a ``/dev/zero`` mapping, which is how pinned memory maps.""" + with open("/dev/zero", "r+b") as dev_zero: + mapping = mmap.mmap(dev_zero.fileno(), mmap.PAGESIZE) + table = torch.frombuffer(mapping, dtype=torch.uint8)[:84].view(torch.bfloat16) + return [ + (f"ngram_embedding.shard_{index}.weight", shard) + for index, shard in enumerate(table.reshape(21, 2).split(6)) + ] + + +@pytest.mark.parametrize( + ("engram_config", "backend", "rows"), + [ + ( + SimpleNamespace(host_file_gather=True, cpu_offload=True), + Qwen4ExpPLEFileGatherEmbedding, + 0, + ), + ( + SimpleNamespace(host_file_gather=False, cpu_offload=False), + Qwen4ExpPLEDeviceEmbedding, + 24, + ), + (None, Qwen4ExpPLEDeviceEmbedding, 24), + ], +) +def test_ngram_embedding_selects_backend_from_engram_config( + monkeypatch, engram_config, backend, rows +) -> None: + """File gather wins with no rows; ``cpu_offload=False`` keeps rows on device.""" + _patch_ple_construction(monkeypatch, engram_config=engram_config) + config = SimpleNamespace( + ngram_size=2, + heads_per_ngram=2, + eos_token_id=0, + vocab_size=64, + ngram_vocab_size_base=8, + make_ngram_vocab_size_divisible_by=1, + ) + + module = Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t") + + assert tuple(module.ngram_embedding.weight.shape) == (rows, 2) + assert isinstance(module.ngram_embedding, backend) + + +def test_ngram_embedding_default_cpu_offload_selects_pinned_backend( + monkeypatch, +) -> None: + """Default ``cpu_offload=True`` selects pinned storage and requires UVA.""" + _patch_ple_construction( + monkeypatch, + engram_config=SimpleNamespace(host_file_gather=False, cpu_offload=True), + ) + monkeypatch.setattr(ngram_embedding_module, "is_uva_available", lambda: False) + config = SimpleNamespace( + ngram_size=2, + heads_per_ngram=2, + eos_token_id=0, + vocab_size=64, + ngram_vocab_size_base=8, + make_ngram_vocab_size_divisible_by=1, + ) + + with pytest.raises(RuntimeError, match="CPU offload requires UVA"): + Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t") + + +@pytest.mark.parametrize( + ("read_shards", "error"), + [ + pytest.param( + lambda path: [("ngram_embedding.shard_0.weight", torch.zeros(6, 2))], + "dtype torch.float32 must match the embedding dtype torch.bfloat16", + id="wrong_dtype", + ), + pytest.param( + lambda path: [*load_file(path).items()] * 2, + "Duplicate PLE embedding shard 0", + id="duplicate_shard", + ), + pytest.param( + lambda path: [s for s in load_file(path).items() if "_2." not in s[0]], + "PLE shards cover 15 of 21 rows", + id="missing_shard", + ), + pytest.param(lambda path: [], "PLE shards cover 0 of 21 rows", id="no_shards"), + pytest.param( + lambda path: [(n, t.clone()) for n, t in load_file(path).items()], + "not a file-backed view", + id="anonymous", + ), + pytest.param( + lambda path: load(path.read_bytes()).items(), + "not a file-backed view", + id="heap", + ), + pytest.param( + lambda path: _dev_zero_shards(), "not a file-backed view", id="dev_zero" + ), + ], +) +def test_file_gather_rejects_shards_that_cannot_serve_every_row_from_a_file( + monkeypatch, tmp_path, read_shards, error +) -> None: + """A real load must fail instead of serving zeros or a copied table.""" + module, _ = _make_file_gather_embedding(monkeypatch, tmp_path) + + with pytest.raises(ValueError, match=error): + module.load_weights(read_shards(tmp_path / "model.safetensors")) + module.ngram_embedding.bind_file_shards() + + +@pytest.mark.parametrize("fp8", [False, True]) +def test_file_gather_binds_zero_copy_shards_through_loader_hook( + monkeypatch, tmp_path, fp8 +) -> None: + """Fails if safetensors stops returning zero-copy views of the file.""" + module, table = _make_file_gather_embedding(monkeypatch, tmp_path, fp8=fp8) + embedding = module.ngram_embedding + + loaded = _load_file_shards(module, tmp_path / "model.safetensors") + embedding.quant_method.process_weights_after_loading(embedding) + embedding.quant_method.process_weights_after_loading(embedding) + + assert loaded == {"ngram_embedding.weight"} | ( + {"ngram_embedding.weight_scale"} if fp8 else set() + ) + assert embedding._bound + stored = torch.cat([embedding._shards[index] for index in range(4)]) + assert torch.equal(stored, table.view(torch.uint8)) + + +def test_file_gather_binds_shards_opened_through_a_symlink( + monkeypatch, tmp_path +) -> None: + """HF cache snapshots link to suffix-less blobs, which is the mapped path.""" + module, _ = _make_file_gather_embedding(monkeypatch, tmp_path) + blob = (tmp_path / "model.safetensors").rename(tmp_path / "blob") + snapshot = tmp_path / "snapshot.safetensors" + snapshot.symlink_to(blob) + _load_file_shards(module, snapshot) + + module.ngram_embedding.bind_file_shards() + + assert module.ngram_embedding._bound + + +def test_file_gather_rejects_reload_and_stays_poisoned(monkeypatch, tmp_path) -> None: + module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True) + + with pytest.raises(RuntimeError, match="does not support weight reload"): + module.load_weights([("ngram_embedding.shard_0.weight", table[:6])]) + with pytest.raises(RuntimeError, match="unbound or was reloaded"): + module.ngram_embedding.stage_rows(1) + + +@pytest.mark.parametrize( + ("bad_id", "error"), + [(21, IndexError), (24, IndexError), (-1, RuntimeError)], + ids=["past_last_shard_rows", "past_last_shard", "negative"], +) +def test_file_gather_out_of_range_id_keeps_previous_staged_rows( + monkeypatch, tmp_path, bad_id, error +) -> None: + """Out-of-range ids fail before the H2D copy, so no wrong row is served.""" + module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True) + embedding = module.ngram_embedding + ids = torch.tensor( + [ + [20, 0], + [19, 1], + [18, 2], + [17, 3], + [16, 4], + [15, 5], + [14, 6], + [13, 7], + [12, 8], + [11, 9], + [10, 10], + ] + ) + embedding._host_ids[:11] = ids + embedding.stage_rows(11) + embedding._host_ids[5, 1] = bad_id + + with pytest.raises(error): + embedding.stage_rows(11) + + assert torch.equal(embedding._staging[:11], table[ids]) + + +def test_file_gather_reads_ahead_the_coalesced_pages_of_staged_rows( + monkeypatch, tmp_path +) -> None: + """Rows start 34 bytes before page 2 of a file mapped from page 1. + + Rows 0-7 are in page 1, row 8 straddles pages 1 and 2, and the rows of shard 1 + arrive unsorted. + """ + module, table = _make_file_gather_embedding(monkeypatch, tmp_path) + embedding, page = module.ngram_embedding, mmap.PAGESIZE + start = 2 * page - 34 + path = tmp_path / "table.bin" + path.write_bytes(bytes(start) + table.view(torch.uint8).numpy().tobytes()) + with open(path, "rb") as file: + mapping = mmap.mmap(file.fileno(), 0, offset=page, access=mmap.ACCESS_COPY) + rows = torch.frombuffer(mapping, dtype=torch.uint8)[start - page :] + module.load_weights( + (f"ngram_embedding.shard_{index}.weight", shard) + for index, shard in enumerate(rows.view(table.dtype).reshape(21, 2).split(6)) + ) + embedding.bind_file_shards() + calls = [] + monkeypatch.setattr( + nvidia_ngram_embedding_module.os, + "posix_fadvise", + lambda fd, *args: calls.append((os.fstat(fd).st_ino, *args)), + ) + ids = torch.tensor([[0, 1], [8, 7]]) + embedding._host_ids[:2] = ids + + embedding.stage_rows(2) + + inode, willneed = path.stat().st_ino, os.POSIX_FADV_WILLNEED + assert calls == [(inode, page, page, willneed), (inode, page, 2 * page, willneed)] + assert torch.equal(embedding._staging[:2], table[ids]) + + +def test_file_gather_dummy_load_stages_zeros(monkeypatch, tmp_path) -> None: + module, _ = _make_file_gather_embedding(monkeypatch, tmp_path, load_format="dummy") + embedding = module.ngram_embedding + embedding.quant_method.process_weights_after_loading(embedding) + embedding._host_rows.view(torch.uint8).fill_(0xFF) + embedding._staging.view(torch.uint8).fill_(0xFF) + + embedding.stage_rows(4) + + staged = embedding._staging.view(torch.uint8) + assert torch.all(staged[:4] == 0) + assert torch.all(staged[4:] == 0xFF) + + +def test_file_gather_forward_reads_staging_without_host_work( + monkeypatch, tmp_path +) -> None: + """The forward must be a fixed-address read so FULL CUDA graphs replay it.""" + module, _ = _make_file_gather_embedding(monkeypatch, tmp_path) + staging = module.ngram_embedding._staging + staging.view(torch.uint8).fill_(0xFF) + step = (torch.arange(6), torch.tensor([0, 6]), torch.tensor([[1]])) + hidden_states = torch.zeros(6, 8, dtype=torch.bfloat16) + + assert module.start_prefetch(hidden_states, *step) is None + output = module(hidden_states, *step) + + assert torch.all(staging.view(torch.uint8) == 0xFF) + assert output.data_ptr() == staging.data_ptr() + assert output.shape == (6, 4) + + def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None: conv_state_len = 6 dilation = 2 diff --git a/tests/test_config.py b/tests/test_config.py index bbb12e3917..609d59ed5b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1478,7 +1478,9 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in assert config.get_parallel_size(parallel) == expected -@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"]) +@pytest.mark.parametrize( + "option", ["embedding_across_dp", "dp_shared_memory", "host_file_gather"] +) def test_engram_rejects_elastic_cross_dp(option): parallel = ParallelConfig( tensor_parallel_size=4, @@ -1681,12 +1683,47 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, expli assert config.engram_config.cpu_offload is True +@pytest.mark.parametrize( + ("architecture", "layer_field", "rocm", "expect_error"), + [ + pytest.param( + "DeepseekV41ForCausalLM", "engram_layer_ids", False, True, id="deepseek" + ), + pytest.param( + "Qwen4ExpForCausalLM", "ple_layer_ids", True, True, id="qwen4exp_rocm" + ), + pytest.param( + "Qwen4ExpForCausalLM", "ple_layer_ids", False, False, id="qwen4exp_cuda" + ), + ], +) +def test_engram_host_file_gather_requires_cuda_qwen4exp( + monkeypatch, architecture, layer_field, rocm, expect_error +): + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm) + model = cast( + ModelConfig, + SimpleNamespace( + architecture=architecture, + hf_text_config=SimpleNamespace(**{layer_field: [1]}), + ), + ) + config = EngramConfig(host_file_gather=True) + if expect_error: + with pytest.raises(ValueError, match="host_file_gather"): + config.verify_model_config(model) + else: + config.verify_model_config(model) + + def test_engram_hash_tracks_execution_options(): configs = [ EngramConfig(cpu_offload=True), EngramConfig(cpu_offload=False), EngramConfig(cpu_offload=True, embedding_across_dp=True), EngramConfig(cpu_offload=True, dp_shared_memory=True), + EngramConfig(host_file_gather=True), ] assert len({config.compute_hash() for config in configs}) == len(configs) diff --git a/vllm/config/engram.py b/vllm/config/engram.py index fa28d803b5..da2919f97b 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -38,6 +38,11 @@ class EngramConfig: cpu_offload: bool = True """Store embedding weights in pinned CPU memory for UVA lookup.""" + host_file_gather: bool = False + """Serve Qwen4Exp PLE rows by gathering them on the CPU from the loader's + file-backed checkpoint views and staging them on the device each step, + instead of allocating the table; takes precedence over `cpu_offload`.""" + embedding_across_dp: bool = False """Shard embeddings across TP and all DP ranks when enabled. Otherwise, each DP rank has a separate TP-sharded embedding replica.""" @@ -85,6 +90,10 @@ class EngramConfig: "embeddings, non-empty n-gram layer ids, and a CUDA-alike " "device (CUDA or ROCm)." ) + if self.host_file_gather and ( + field != "ple_layer_ids" or current_platform.is_rocm() + ): + raise ValueError("host_file_gather requires a Qwen4Exp model on CUDA") def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None: """Share host tables by default wherever the configuration permits.""" @@ -110,6 +119,8 @@ class EngramConfig: raise ValueError( "Engram embedding_across_dp is not supported with elastic EP yet." ) + if self.host_file_gather and parallel_config.enable_elastic_ep: + raise ValueError("host_file_gather is not supported with elastic EP.") def get_parallel_size(self, parallel_config: "ParallelConfig") -> int: """Derive the embedding group size from the parallel configuration.""" diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 1828b65b73..662a427ff1 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -48,7 +48,7 @@ logger = init_logger(__name__) class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): - """ETP-sharded PLE table shared by device and pinned-host backends.""" + """PLE table base shared by the device, pinned-host, and file-gather backends.""" supports_prefetch: ClassVar[bool] = False @@ -157,9 +157,9 @@ class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): - """Quantization interface shared by resident and pinned PLE tables.""" + """Quantization interface shared by device, pinned, and file-gather PLE tables.""" - # PLE post-load processing only validates scales in their current storage. + # PLE post-load processing validates scales in place and binds host file shards. requires_device_loading: bool = False @staticmethod @@ -216,6 +216,12 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: return F.embedding(input_, layer.weight) + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Bind host-file-gather tables once every checkpoint shard is loaded.""" + bind = getattr(layer, "bind_file_shards", None) + if bind is not None: + bind() + @abstractmethod def dequantize( self, @@ -305,6 +311,7 @@ class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): sentinel = torch.finfo(torch.float32).min if torch.any(layer.weight_scale == sentinel): raise ValueError("FP8 PLE checkpoint is missing its global scale") + super().process_weights_after_loading(layer) def dequantize( self, diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py index 246168509b..c9584d8496 100644 --- a/vllm/models/qwen4_exp/nvidia/model_state.py +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -13,6 +13,12 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from vllm.v1.worker.gpu.states import RequestState +from .ngram_embedding import ( + Qwen4ExpNGramEmbedding, + Qwen4ExpPLEFileGatherEmbedding, + stage_checkpoint_rows, +) + class Qwen4ExpModelState(MambaHybridModelState): """Add rollback-safe PLE n-gram context to the model inputs.""" @@ -63,6 +69,12 @@ class Qwen4ExpModelState(MambaHybridModelState): dtype=torch.int32, device=self.device, ) + self.file_gather_modules = [ + module + for module in model.modules() + if isinstance(module, Qwen4ExpNGramEmbedding) + and isinstance(module.ngram_embedding, Qwen4ExpPLEFileGatherEmbedding) + ] def _prepare_ngram_context( self, @@ -106,9 +118,17 @@ class Qwen4ExpModelState(MambaHybridModelState): query_start_loc[: num_reqs_padded + 1].copy_(input_batch.query_start_loc) # Represent unused capacity as trailing zero-length requests. query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1]) + ngram_context = self._prepare_ngram_context(input_batch, req_states) + if self.file_gather_modules: + stage_checkpoint_rows( + self.file_gather_modules, + input_batch.input_ids, + query_start_loc, + ngram_context, + ) model_inputs.update( query_start_loc=query_start_loc, - ngram_context=self._prepare_ngram_context(input_batch, req_states), + ngram_context=ngram_context, ) return model_inputs diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 0b6fd650b1..eb66fcf06b 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Qwen4Exp n-gram embeddings with device and pinned-host storage.""" +"""Qwen4Exp n-gram embeddings with device, pinned-host and host-file-gather storage.""" +import mmap +import os from collections.abc import Iterable import torch @@ -40,6 +42,155 @@ __all__ = [ ] +def _maps_entry(ptr: int) -> tuple[str, int]: + """Return the path and file offset of ``ptr`` from ``/proc/self/maps``.""" + with open("/proc/self/maps") as maps: + for line in maps: + fields = line.split(maxsplit=5) + low, high = (int(bound, 16) for bound in fields[0].split("-")) + if low <= ptr < high: + path = fields[5].rstrip("\n") if len(fields) > 5 else "" + return path, ptr - low + int(fields[2], 16) + return "", 0 + + +class Qwen4ExpPLEFileGatherEmbedding(Qwen4ExpPLEEmbedding): + """PLE rows gathered on the host from the loader's file-backed shard views.""" + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + params_dtype: torch.dtype, + padding_size: int, + prefix: str, + embedding_method: Qwen4ExpPLEEmbeddingMethod, + num_ngram_heads: int = 1, + max_total_tokens: int = 0, + data_parallel_rank: int = 0, + ) -> None: + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + padding_size=padding_size, + prefix=prefix, + embedding_method=embedding_method, + data_parallel_rank=data_parallel_rank, + ) + self._dummy = get_current_vllm_config().load_config.load_format == "dummy" + shape = (max_total_tokens, num_ngram_heads, embedding_dim) + dtype, pin = self.weight.dtype, self.weight.device.type != "cpu" + self._staging = torch.zeros(shape, dtype=dtype, device=self.weight.device) + self._host_ids = torch.empty( + shape[:2], dtype=torch.int64, device="cpu", pin_memory=pin + ) + # Step N+1 writes these rows only after its stream sync, so after step N's H2D. + self._host_rows = torch.empty(shape, dtype=dtype, device="cpu", pin_memory=pin) + self._shards: dict[int, torch.Tensor] | None = {} + self._shard_size = 0 + self._bound = False + self._fds: dict[str, int] = {} + self._fds_t = self._bases = torch.empty(0, dtype=torch.int64) + + def allocate_embedding_weight( + self, + num_embeddings: int, + embedding_dim: int, + dtype: torch.dtype, + ) -> torch.Tensor: + """Allocate no rows; they stay in the checkpoint files.""" + return torch.empty(0, embedding_dim, dtype=dtype) + + def accept_checkpoint_shard( + self, + shard_index: int, + loaded_weight: torch.Tensor, + shard_size: int, + ) -> None: + """Keep a checkpoint shard as a zero-copy byte view of its storage.""" + if self._bound or self._shards is None: + self._shards = None + raise RuntimeError("PLE host file gather does not support weight reload") + if loaded_weight.dtype != self.weight.dtype: + raise ValueError( + f"PLE shard dtype {loaded_weight.dtype} must match the embedding " + f"dtype {self.weight.dtype}; host file gather serves rows as stored" + ) + if shard_index in self._shards: + raise ValueError(f"Duplicate PLE embedding shard {shard_index}") + self._shards[shard_index] = loaded_weight.view(torch.uint8) + self._shard_size = shard_size + + def bind_file_shards(self) -> None: + """Check that file-backed shards cover every row; dummy loads serve zeros.""" + assert self._shards is not None + if not self._dummy: + rows = sum(shard.shape[0] for shard in self._shards.values()) + if rows != self.org_vocab_size: + raise ValueError( + f"PLE shards cover {rows} of {self.org_vocab_size} rows" + ) + self._fds_t, self._bases = torch.zeros(2, len(self._shards)).long() + for index, shard in self._shards.items(): + path, offset = _maps_entry(shard.data_ptr()) + if not path or not os.path.isfile(path): + raise ValueError( + f"PLE embedding shard {index} is not a file-backed view " + f"({path or 'anonymous'})" + ) + if path not in self._fds: + self._fds[path] = os.open(path, os.O_RDONLY) + self._fds_t[index], self._bases[index] = self._fds[path], offset + self._bound = True + + def stage_rows(self, num_tokens: int) -> None: + """Gather the rows for ``_host_ids[:num_tokens]`` and copy them to staging.""" + shards = self._shards + if shards is None or not self._bound: + raise RuntimeError("PLE host file gather is unbound or was reloaded") + ids = self._host_ids[:num_tokens].flatten() + out = self._host_rows[:num_tokens].flatten(0, 1).view(torch.uint8) + if not shards: + out.zero_() + else: + ids, order = ids.sort() + shard = ids // self._shard_size + local = ids - shard * self._shard_size + counts = torch.bincount(shard, minlength=len(shards)).tolist() + page, row_bytes = mmap.PAGESIZE, out.shape[1] + start = self._bases[shard] + local * row_bytes + first, last = start // page, (start + row_bytes - 1) // page + run = torch.ones_like(first, dtype=torch.bool) + run[1:] = (first[1:] > last[:-1] + 1) | (shard[1:] != shard[:-1]) + fds, lows = self._fds_t[shard[run]].tolist(), first[run].tolist() + # Queue every page read before the gather's first serial fault. + for fd, lo, hi in zip(fds, lows, last[run.roll(-1)].tolist()): + os.posix_fadvise( + fd, lo * page, (hi - lo + 1) * page, os.POSIX_FADV_WILLNEED + ) + groups = zip(order.split(counts), local.split(counts)) + for index, (rows, rows_local) in enumerate(groups): + if rows.numel(): + out.index_copy_(0, rows, shards[index].index_select(0, rows_local)) + self._staging[:num_tokens].copy_( + self._host_rows[:num_tokens], non_blocking=True + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return the rows staged for this step.""" + return self._staging[: hidden_states.shape[0]].flatten(-2) + + def start_prefetch( + self, + hidden_states: torch.Tensor, + ngram_ids: torch.Tensor, + ) -> None: + """Rows are staged in ``prepare_inputs``, so prefetch is a no-op.""" + return None + + class Qwen4ExpNGramEmbedding(nn.Module): _MASK64 = (1 << 64) - 1 _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 @@ -209,7 +360,9 @@ class Qwen4ExpNGramEmbedding(nn.Module): params_dtype = torch.get_default_dtype() engram_config = get_current_vllm_config().engram_config embedding_cls = ( - Qwen4ExpPLEPinnedHostEmbedding + Qwen4ExpPLEFileGatherEmbedding + if engram_config is not None and engram_config.host_file_gather + else Qwen4ExpPLEPinnedHostEmbedding if engram_config is not None and engram_config.cpu_offload else Qwen4ExpPLEDeviceEmbedding ) @@ -358,7 +511,9 @@ class Qwen4ExpNGramEmbedding(nn.Module): ngram_context: torch.Tensor, ) -> torch.Tensor: embedding = self.ngram_embedding - if embedding.supports_prefetch: + if embedding.supports_prefetch or isinstance( + embedding, Qwen4ExpPLEFileGatherEmbedding + ): return embedding(hidden_states) ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) return self.ngram_embedding(ngram_ids).flatten(-2) @@ -434,6 +589,12 @@ class Qwen4ExpNGramEmbedding(nn.Module): f"expected {expected_shape}, got " f"{tuple(loaded_weight.shape)}" ) + if isinstance(embedding, Qwen4ExpPLEFileGatherEmbedding): + embedding.accept_checkpoint_shard( + shard_index, loaded_weight, shard_size + ) + loaded.add("ngram_embedding.weight") + continue embedding.weight.weight_loader( embedding.weight, loaded_weight, @@ -448,6 +609,23 @@ class Qwen4ExpNGramEmbedding(nn.Module): return loaded +def stage_checkpoint_rows( + modules: list[Qwen4ExpNGramEmbedding], + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, +) -> None: + """Stage host-file-gather PLE rows for one step behind a single stream sync.""" + num_tokens = input_ids.shape[0] + for module in modules: + ids = module.compute_ngram_ids(input_ids, query_start_loc, ngram_context) + module.ngram_embedding._host_ids[:num_tokens].copy_(ids, non_blocking=True) + if input_ids.device.type != "cpu": + torch.accelerator.current_stream().synchronize() + for module in modules: + module.ngram_embedding.stage_rows(num_tokens) + + __all__ = [ "Qwen4ExpNGramEmbedding", ]