[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 <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

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
 

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

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

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)
 

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."""

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,

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
 

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",
 ]