[DS V4.1][Engram] Support async prefetch for offloaded engram lookups and engram DP sharding (#56512)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: ywang96 <>
Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

.buildkite/test_areas/distributed.yaml

@@ -206,6 +206,9 @@ steps:
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
   - tests/distributed/test_utils
+  - tests/distributed/test_engram_dp_shard.py
+  - vllm/config/engram.py
+  - vllm/models/deepseek_v4_1/
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
@@ -217,6 +220,7 @@ steps:
   - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
   - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
   - pytest -v -s distributed/test_utils.py
+  - pytest -v -s distributed/test_engram_dp_shard.py
 
 - label: ":nvidia: (L4) Distributed Compile + Comm"
   key: distributed-compile-comm-4-gpus

tests/distributed/test_engram_dp_shard.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Engram tables shared by co-located DP replicas.
+
+Four GPUs on one node split the hash heads four ways, whether as DEP
+(TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns
+the right rows if the ids of every replica are gathered, each rank looks up
+its own heads, and each replica trades those tokens back for the heads its
+peers own.
+
+With shared host storage, DP replicas instead map the same TP slice and
+prefetch only their own tokens without DP gathers.
+"""
+
+from itertools import product
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from threading import Event
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+
+from vllm.config import (
+    EngramConfig,
+    VllmConfig,
+    get_current_vllm_config_or_none,
+    set_current_vllm_config,
+)
+from vllm.config.load import LoadConfig
+from vllm.config.parallel import ParallelConfig
+from vllm.config.scheduler import SchedulerConfig
+from vllm.distributed import cleanup_dist_env_and_memory, parallel_state
+from vllm.distributed.parallel_state import (
+    get_engram_dp_group,
+    get_engram_dp_size,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.forward_context import set_forward_context
+from vllm.models.deepseek_v4_1.common.engram import EngramLayout
+from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops
+from vllm.models.deepseek_v4_1.nvidia.engram import (
+    Engram,
+    ParallelEngramEmbedding,
+    engram_head_shard_rank,
+    gather_engram_hashes,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.system_utils import update_environment_variables
+
+WORLD_SIZE = 4
+DIM, BLOCK = 64, 32
+# Prime-sized buckets, as the model's layout builds them.
+HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131)
+
+
+def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]:
+    """Even loads, then uneven ones with an idle replica to pad around."""
+    return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size)
+
+
+def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """The unsharded table, identical on every rank."""
+    torch.manual_seed(0)
+    weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn)
+    scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8)
+    return weight, scale_inv
+
+
+def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor:
+    """Ids of each head, drawn from that head's own bucket range."""
+    torch.manual_seed(seed)
+    offset = 0
+    columns = []
+    for size in head_sizes:
+        columns.append(torch.randint(offset, offset + size, (num_tokens, 1)))
+        offset += size
+    return torch.cat(columns, dim=1).to(torch.int32).cuda()
+
+
+def _reference(weight, scale_inv, ids) -> torch.Tensor:
+    """Dequantized lookup against the whole table."""
+    valid = ids >= 0
+    safe_ids = ids.clamp_min(0).long()
+    values = torch.nn.functional.embedding(safe_ids, weight).float()
+    scales = torch.nn.functional.embedding(safe_ids, scale_inv)
+    scales = (scales.to(torch.int32) << 23).view(torch.float32)
+    values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1)
+    return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0)
+
+
+@pytest.mark.parametrize(
+    "architecture,options,edp_ranks,etp_ranks",
+    [
+        ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]),
+        ("Qwen4ExpForCausalLM", {}, None, [0, 1]),
+        (
+            "Qwen4ExpForCausalLM",
+            {"embedding_across_dp": True},
+            None,
+            [0, 1, 2, 3],
+        ),
+        ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]),
+    ],
+)
+def test_engram_group_creation_honors_model_and_options(
+    monkeypatch, architecture, options, edp_ranks, etp_ranks
+):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(architecture=architecture, is_moe=False),
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=2,
+            data_parallel_size=2,
+            distributed_executor_backend="mp",
+        ),
+        engram_config=EngramConfig(**options),
+    )
+    for name in (
+        "_TP",
+        "_ETP",
+        "_ENGRAM_DP",
+        "_DCP",
+        "_PCP",
+        "_PP",
+        "_DP",
+        "_EP",
+        "_EPLB",
+    ):
+        monkeypatch.setattr(parallel_state, name, None)
+    monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0))
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0)
+    monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1)
+
+    def make_group(group_ranks, *args, **kwargs):
+        ranks = next(ranks for ranks in group_ranks if 0 in ranks)
+        return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks))
+
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group)
+    with set_current_vllm_config(config):
+        initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl")
+    group = get_engram_dp_group()
+    assert (group.ranks if group is not None else None) == edp_ranks
+    assert parallel_state.get_etp_group().ranks == etp_ranks
+
+
+def _worker(rank: int, tp_size: int, port: int) -> None:
+    torch.accelerator.set_device_index(rank)
+    update_environment_variables(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(WORLD_SIZE),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(port),
+        }
+    )
+    dp_size = WORLD_SIZE // tp_size
+    dp_rank, tp_rank = divmod(rank, tp_size)
+    vllm_config = VllmConfig(
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=tp_size,
+            data_parallel_size=dp_size,
+            data_parallel_rank=dp_rank,
+        ),
+        scheduler_config=SchedulerConfig(
+            max_model_len=8,
+            is_encoder_decoder=False,
+            max_num_batched_tokens=8,
+            max_num_seqs=8,
+        ),
+    )
+    # Supply the resolved Engram settings without constructing a full model.
+    vllm_config.engram_config = EngramConfig()
+    vllm_config.model_config = SimpleNamespace(
+        architecture="DeepseekV41ForCausalLM", is_moe=True
+    )
+    try:
+        init_distributed_environment()
+        with set_current_vllm_config(vllm_config):
+            initialize_model_parallel(tensor_model_parallel_size=tp_size)
+        assert get_engram_dp_size() == dp_size
+        assert get_engram_dp_group().ranks == [
+            tp_rank + replica * tp_size for replica in range(dp_size)
+        ]
+        assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank
+
+        for failure in ("create", "open", "mmap", "register", "pinned", None):
+            _check_shared_storage_lifetime_and_failures(failure)
+
+        # Reuse the processes and communication groups across storage modes.
+        for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product(
+            [(False, False), (True, False), (True, True)],
+            [8, 7],
+            [False, True] if tp_size > 1 else [False],
+        ):
+            _check_table(
+                vllm_config,
+                cpu_offload,
+                dp_shared_memory,
+                n_heads,
+                sequence_parallel,
+            )
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+def _check_shared_storage_lifetime_and_failures(failure):
+    """One peer's failure must roll back every mapping/registration, without hangs."""
+    import gc
+    import weakref
+
+    group = get_engram_dp_group()
+    runtime = torch.cuda.cudart()
+    make_file = engram_ops.tempfile.NamedTemporaryFile
+    map_file = engram_ops.mmap.mmap
+    paths, mappings, registrations, unregistrations = [], [], [], []
+    failed_rank = 0 if failure == "create" else 1
+
+    def fail():
+        raise OSError(f"injected {failure} failure")
+
+    def temporary_file(*args, **kwargs):
+        if failure == "create":
+            fail()
+        file = make_file(*args, **kwargs)
+        paths.append(Path(file.name))
+        return file
+
+    def open_file(*args, **kwargs):
+        if failure == "open" and group.rank_in_group == failed_rank:
+            fail()
+        return open(*args, **kwargs)
+
+    def mapping(*args, **kwargs):
+        if failure == "mmap" and group.rank_in_group == failed_rank:
+            fail()
+        result = map_file(*args, **kwargs)
+        mappings.append(weakref.ref(result))
+        return result
+
+    def register(pointer, size, flags):
+        if failure == "register" and group.rank_in_group == failed_rank:
+            return SimpleNamespace(value=1)
+        result = runtime.cudaHostRegister(pointer, size, flags)
+        assert result.value == 0
+        registrations.append(pointer)
+        return result
+
+    def unregister(pointer):
+        unregistrations.append(pointer)
+        result = runtime.cudaHostUnregister(pointer)
+        assert result.value == 0
+        return result
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(
+            engram_ops,
+            "tempfile",
+            SimpleNamespace(NamedTemporaryFile=temporary_file),
+        )
+        patch.setattr(engram_ops, "open", open_file, raising=False)
+        patch.setattr(
+            engram_ops,
+            "mmap",
+            SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED),
+        )
+        patch.setattr(
+            torch.cuda,
+            "cudart",
+            lambda: SimpleNamespace(
+                cudaHostRegister=register,
+                cudaHostUnregister=unregister,
+            ),
+        )
+        if failure == "pinned" and group.rank_in_group == failed_rank:
+            patch.setattr(torch.Tensor, "is_pinned", lambda self: False)
+        if failure is not None:
+            stage = "cudaHostRegister" if failure in ("register", "pinned") else failure
+            with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"):
+                engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
+            assert unregistrations == registrations
+        else:
+            storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group)
+            assert all(not path.exists() for path in paths)
+            storage_ref = weakref.ref(storage)
+            parameter = torch.nn.Parameter(storage.weight, requires_grad=False)
+            views = storage.get_views(parameter, storage.weight_scale_inv)
+            del storage
+            gc.collect()
+            assert storage_ref() is None
+            assert not unregistrations
+            del parameter
+            gc.collect()
+            assert not unregistrations  # UVA views still retain the CPU allocation.
+            del views
+        gc.collect()
+        assert unregistrations == registrations
+        for ref in mappings:
+            mapped = ref()
+            assert mapped is None or mapped.closed
+        assert all(not path.exists() for path in paths)
+
+
+def _check_table(
+    vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel
+):
+    parallel = vllm_config.parallel_config
+    dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size
+    dp_rank = parallel.data_parallel_rank
+    tp_rank = engram_ops.get_tensor_model_parallel_rank()
+    multithread = (
+        dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1)
+    )
+    load_config = (
+        LoadConfig(
+            load_format="safetensors",
+            use_tqdm_on_load=False,
+            model_loader_extra_config={
+                "enable_multithread_load": True,
+                "num_threads": 2,
+            },
+        )
+        if multithread
+        else vllm_config.load_config
+    )
+    with pytest.MonkeyPatch.context() as m:
+        head_sizes = HEAD_SIZES[:n_heads]
+        rows = sum(head_sizes)
+        config = SimpleNamespace(
+            hidden_size=DIM,
+            hc_mult=2,
+            rms_norm_eps=1e-6,
+            engram_layer_ids=[0],
+            engram_num_embeddings=[rows],
+            engram_max_ngram_size=2,
+            engram_n_heads=n_heads,
+            engram_head_dim=DIM,
+            engram_compressed_vocab_size=32,
+            engram_pad_token_id=0,
+            engram_vocab_size=HEAD_SIZES[0],
+        )
+        with set_current_vllm_config(vllm_config):
+            weight, scale_inv = _full_table(rows)
+            layout = EngramLayout(config)
+            # Stub configuration inputs, not Engram's initialization or execution.
+            component_config = SimpleNamespace(
+                engram_config=EngramConfig(
+                    cpu_offload=cpu_offload,
+                    dp_shared_memory=dp_shared_memory,
+                ),
+                scheduler_config=vllm_config.scheduler_config,
+                load_config=load_config,
+            )
+            with m.context() as init_patch, torch.device("cuda"):
+                init_patch.setattr(
+                    engram_ops, "get_current_vllm_config", lambda: component_config
+                )
+                engram = Engram(
+                    config,
+                    quant_config=None,
+                    layout=layout,
+                    layer_hash_index=0,
+                    use_sequence_parallel=sequence_parallel,
+                    prefix="model.layers.0.engram",
+                )
+            layer = engram.embed_tokens
+            if multithread:
+                _load_shared_multithread(layer, weight, scale_inv, load_config)
+            else:
+                # Only the leader has valid checkpoint payload in shared mode.
+                loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None
+                loader_scale = (
+                    scale_inv if not dp_shared_memory or dp_rank == 0 else None
+                )
+                layer.weight.weight_loader(layer.weight, loader_weight)
+                layer.weight_scale_inv.weight_loader(
+                    layer.weight_scale_inv, loader_scale
+                )
+            if dp_shared_memory:
+                assert layer.dp_size == 1
+                assert layer.head_start == tp_rank * layer.part_n_hash_cols
+                assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned()
+
+                def unexpected_collective(*args, **kwargs):
+                    raise AssertionError("shared lookup must not use DP collectives")
+
+                m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective)
+
+            token_counts = _token_counts(dp_size)
+
+        # Production forward runs after the construction-time config has exited.
+        assert get_current_vllm_config_or_none() is None
+        for batch_idx, counts in enumerate(token_counts):
+            num_tokens = counts[dp_rank]
+            # Ids differ per replica but must agree across its TP ranks.
+            ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank)
+            if batch_idx == 2 and dp_rank == 1:
+                ids.fill_(engram_ops.DEAD_ID)
+            with set_forward_context(
+                None,
+                vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32),
+            ):
+                gathered = gather_engram_hashes(
+                    ids, dp_shared_memory=layer.dp_shared_memory
+                )
+                expected_tokens = (
+                    num_tokens if dp_shared_memory else max(counts) * dp_size
+                )
+                assert gathered.shape[0] == expected_tokens
+                engram.prepare_embeddings(gathered)
+                looked_up = engram.embed(ids)
+                standalone = layer(ids)
+            expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
+            torch.testing.assert_close(standalone, expected, atol=0, rtol=0)
+            if sequence_parallel:
+                # SP hands back only this rank's token shard, zero-padded.
+                chunk = -(-num_tokens // tp_size)
+                expected = torch.nn.functional.pad(
+                    expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens)
+                )
+                expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk]
+            torch.testing.assert_close(looked_up, expected, atol=0, rtol=0)
+
+        if (
+            cpu_offload
+            and n_heads == 7
+            and not sequence_parallel
+            and (not dp_shared_memory or tp_size == 1)
+        ):
+            _check_dummy_hash_model_forward(
+                vllm_config,
+                engram,
+                head_sizes,
+                weight,
+                scale_inv,
+                dp_shared_memory,
+            )
+
+        replay = n_heads == 7 and (
+            (dp_shared_memory and sequence_parallel == (tp_size > 1))
+            or (
+                cpu_offload
+                and not dp_shared_memory
+                and sequence_parallel == (tp_size > 1)
+            )
+        )
+        if replay:
+            captures = ("full", "breakable")
+            for capture in captures:
+                _check_prefetch_replay(
+                    engram,
+                    head_sizes,
+                    weight,
+                    scale_inv,
+                    dp_shared_memory,
+                    capture,
+                )
+        if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1):
+            with m.context() as storage_patch:
+                storage_patch.setattr(layer.weight, "data", layer.weight.data.clone())
+                with pytest.raises(RuntimeError, match="storage must not be replaced"):
+                    layer(ids)
+
+
+def _check_dummy_hash_model_forward(
+    vllm_config, engram, head_sizes, weight, scales, dp_shared_memory
+):
+    """A metadata-less replica joins DP sharding but skips a shared lookup."""
+    from vllm.models.deepseek_v4_1.common.engram import NgramHashState
+    from vllm.models.deepseek_v4_1.nvidia import model as model_ops
+
+    dp_rank = vllm_config.parallel_config.data_parallel_rank
+    if dp_shared_memory and dp_rank != 1:
+        return
+    tokens = 4
+    ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank)
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda")
+    state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda")
+    state.ensure_cache = lambda: True
+    # Hash arithmetic is covered separately; keep the real model's branch and
+    # NgramHashState.dummy_hashes, plus real preparation/lookup collectives.
+    state.forward = lambda *args: ids.unsqueeze(1)
+
+    class Decoder(SimpleNamespace):
+        def __call__(self, hidden, positions, input_ids, *args):
+            hashes, keep = args[-2:]
+            if dp_shared_memory and dp_rank == 1:
+                assert hashes is None and keep is None
+                output = hidden.new_zeros(
+                    tokens, len(head_sizes), DIM, dtype=torch.bfloat16
+                )
+            else:
+                assert hashes is not None and keep is not None
+                assert bool(keep.all()) == (dp_rank != 1)
+                if dp_rank == 1:
+                    assert torch.all(hashes == engram_ops.DEAD_ID)
+                output = engram.embed(hashes[:, 0])
+            return output, None, None, None, None
+
+    model = SimpleNamespace(
+        use_mega_moe=False,
+        use_sequence_parallel=False,
+        engram_hash=state,
+        engram_swa_prefix="swa",
+        engram_dp_shared_memory=dp_shared_memory,
+        layers=[Decoder(engram=engram)],
+        start_layer=0,
+        end_layer=1,
+        aux_hidden_state_layers=(),
+    )
+    metadata = (
+        None
+        if dp_rank == 1
+        else {
+            "swa": SimpleNamespace(
+                query_start_loc=None, slot_mapping=None, block_table=None
+            )
+        }
+    )
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(
+            model_ops,
+            "get_pp_group",
+            lambda: SimpleNamespace(
+                is_first_rank=True,
+                is_last_rank=False,
+            ),
+        )
+        patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden)
+        with set_forward_context(
+            metadata,
+            vllm_config,
+            num_tokens=tokens,
+            num_tokens_across_dp=torch.full(
+                (vllm_config.parallel_config.data_parallel_size,),
+                tokens,
+                dtype=torch.int32,
+            ),
+        ):
+            output = model_ops.DeepseekV4Model.forward(
+                model,
+                torch.arange(tokens, device="cuda"),
+                torch.arange(tokens, device="cuda"),
+                intermediate_tensors=None,
+                inputs_embeds=torch.zeros(tokens, DIM, device="cuda"),
+                lookback_token_ids=torch.full((1, 1), -1, device="cuda"),
+            )["hidden_states"]
+    expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids
+    expected = _reference(weight.cuda(), scales.cuda(), expected_ids)
+    torch.testing.assert_close(output, expected, atol=0, rtol=0)
+
+
+def _load_shared_multithread(layer, weight, scale_inv, load_config):
+    """Opposite shard completion orders must leave every shared parameter ready."""
+    from safetensors.torch import save_file
+
+    from vllm.model_executor.model_loader import weight_utils
+    from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+
+    tensors = {
+        "weight": weight,
+        "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu),
+    }
+    order = list(tensors)
+    if get_engram_dp_group().rank_in_group % 2:
+        order.reverse()
+    first_loaded = Event()
+    original_load = weight_utils.load_file
+
+    def load_file(path, **kwargs):
+        if Path(path).stem != order[0]:
+            assert first_loaded.wait(timeout=30), "First shared weight load stalled"
+        return original_load(path, **kwargs)
+
+    with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m:
+        for name, tensor in tensors.items():
+            save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors"))
+        m.setattr(weight_utils, "load_file", load_file)
+        loader = DefaultModelLoader(load_config)
+        source = DefaultModelLoader.Source(directory, revision=None)
+        loaded = []
+        for name, tensor in loader._get_weights_iterator(source):
+            param = getattr(layer, name)
+            param.weight_loader(param, tensor)
+            loaded.append(name)
+            first_loaded.set()
+        assert loaded == order
+
+    # Check before any further collective could mask incomplete shared writes.
+    for name, tensor in tensors.items():
+        param = getattr(layer, name)
+        expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0])
+        torch.testing.assert_close(
+            param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0
+        )
+
+
+def _check_prefetch_replay(
+    engram, head_sizes, weight, scale_inv, dp_shared_memory, capture
+):
+    """DP-sharded and shared prefetches must replay with new IDs."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    ids = _make_ids(head_sizes, 8, seed=501)
+    tp_rank = engram_ops.get_tensor_model_parallel_rank()
+    tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8
+    output = torch.empty(
+        tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16
+    )
+
+    def step():
+        gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory)
+        engram.prepare_embeddings(gathered)
+        output.copy_(engram.embed(ids))
+
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8)
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            step()
+        if capture == "breakable":
+            graph = BreakableCUDAGraphCapture()
+            with torch.cuda.stream(warmup), graph:
+                step()
+            assert graph.num_eager_breaks == 2
+        else:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=warmup):
+                step()
+        torch.cuda.current_stream().wait_stream(warmup)
+        for iteration in range(2):
+            ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration))
+            graph.replay()
+            expected = _reference(weight.cuda(), scale_inv.cuda(), ids)
+            if engram.use_sequence_parallel:
+                expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens]
+            torch.testing.assert_close(output, expected, atol=0, rtol=0)
+        torch.accelerator.synchronize()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch):
+    """Check all storage modes with one process launch per TP/DP topology."""
+    if torch.accelerator.device_count() < WORLD_SIZE:
+        pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.")
+    monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
+    mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE)
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.parametrize(
+    "cpu_offload,dp_size,error",
+    [
+        (False, 4, "requires cpu_offload=True"),
+        (True, 1, "effective Engram DP size is 1"),
+    ],
+)
+def test_engram_dp_shared_memory_runtime_requirements(
+    monkeypatch, cpu_offload, dp_size, error
+):
+    """Report unmet storage/topology requirements before allocating CUDA memory."""
+    monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size)
+    with pytest.raises(ValueError, match=error):
+        ParallelEngramEmbedding(
+            sum(HEAD_SIZES),
+            DIM,
+            HEAD_SIZES,
+            cpu_offload=cpu_offload,
+            dp_shared_memory=True,
+        )
+
+
+@pytest.mark.parametrize(
+    "node_ids,dp_size,replica_size,expected",
+    [
+        ([0] * 4, 1, 4, 1),
+        ([0] * 4, 4, 1, 4),
+        ([0] * 4, 2, 2, 2),
+        ([0] * 8 + [1] * 8, 8, 2, 4),
+        ([0] * 6 + [1] * 6, 4, 1, 2),
+        ([0] * 3 + [1] * 4, 7, 1, 1),
+        ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1),
+        ([0, 1] * 4, 4, 2, 1),
+        ([0] * 3 + [1] * 5, 4, 2, 1),
+    ],
+)
+def test_engram_dp_shard_size_respects_node_boundaries(
+    node_ids, dp_size, replica_size, expected, monkeypatch
+):
+    """Sharing must fall back when replicas or node rank ranges do not align."""
+    monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids)))
+    monkeypatch.setattr(
+        parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None)
+    )
+    monkeypatch.setattr(
+        parallel_state,
+        "in_the_same_node_as",
+        lambda group, src: [node == node_ids[src] for node in node_ids],
+    )
+    assert (
+        parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size)
+        == expected
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [0, 2, 4, 5])
+def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch):
+    """Idle/padded slots must not query row zero; oversized batches are rejected."""
+    monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4)
+    monkeypatch.setattr(
+        engram_ops,
+        "get_engram_dp_group",
+        lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)),
+    )
+    ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32)
+    if num_tokens > 4:
+        with pytest.raises(ValueError, match="exceeds the DP token slot"):
+            gather_engram_hashes(ids)
+        return
+    gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3)
+    for replica in gathered:
+        torch.testing.assert_close(replica[:num_tokens], ids)
+        assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID)
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+@pytest.mark.parametrize("uniform", [False, True])
+def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch):
+    """Two four-GPU nodes pad only to their own maximum, including TP x DP."""
+    group_size = 4 // tp_size
+    counts = (
+        [4] * (8 // tp_size)
+        if uniform
+        else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2)
+    )
+    monkeypatch.setattr(
+        engram_ops,
+        "get_forward_context",
+        lambda: SimpleNamespace(
+            dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts))
+        ),
+    )
+    for dp_rank in range(8 // tp_size):
+        slot = 4 if uniform else (4096 if dp_rank < group_size else 3)
+        ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32)
+        monkeypatch.setattr(
+            engram_ops,
+            "get_dp_group",
+            lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank),
+        )
+        monkeypatch.setattr(
+            engram_ops,
+            "get_engram_dp_group",
+            lambda rank=dp_rank: SimpleNamespace(
+                rank_in_group=rank % group_size,
+                world_size=group_size,
+                all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim),
+            ),
+        )
+        gathered = gather_engram_hashes(ids)
+        expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID)
+        expected[: ids.shape[0]] = ids
+        torch.testing.assert_close(
+            gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}"
+        )

tests/engine/test_arg_utils.py

@@ -89,7 +89,10 @@ def test_watermark_config_cli():
 @pytest.mark.parametrize(
     "options",
     [
-        ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'],
+        [
+            "--engram-config",
+            '{"cpu_offload": false, "embedding_across_dp": true}',
+        ],
         [
             "--engram-config.cpu_offload",
             "false",
@@ -98,9 +101,8 @@ def test_watermark_config_cli():
         ],
     ],
 )
-def test_engram_config_cli(options, monkeypatch):
-    """CLI settings take precedence over the legacy offload environment."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
+def test_engram_config_cli(options):
+    """JSON and dotted CLI options independently control Engram settings."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = EngineArgs.from_cli_args(parser.parse_args(options))
     assert args.engram_config is not None
@@ -109,14 +111,26 @@ def test_engram_config_cli(options, monkeypatch):
 
 
 @pytest.mark.parametrize(
-    "options, provided",
-    [([], False), (["--engram-config", "{}"], True)],
+    "options,provided,dp_shared_memory",
+    [
+        ([], False, False),
+        (["--engram-config", "{}"], True, False),
+        (
+            ["--engram-config", '{"dp_shared_memory": true}'],
+            True,
+            True,
+        ),
+        (["--engram-config.dp_shared_memory", "true"], True, True),
+    ],
 )
-def test_engram_config_cli_optional(options, provided):
-    """An explicit empty config must remain distinct from an omitted config."""
+def test_engram_config_cli_optional(options, provided, dp_shared_memory):
+    """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag."""
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
     args = EngineArgs.from_cli_args(parser.parse_args(options))
     assert (args.engram_config is not None) == provided
+    if provided:
+        assert args.engram_config.cpu_offload is True
+        assert args.engram_config.dp_shared_memory is dp_shared_memory
 
 
 @pytest.mark.parametrize(

tests/kernels/test_engram.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import bisect
+import inspect
 from types import SimpleNamespace
 
 import pytest
@@ -8,10 +9,13 @@ import torch
 
 from vllm.models.deepseek_v4_1.common import engram as engram_ops
 from vllm.models.deepseek_v4_1.common.engram import (
-    Engram,
-    NgramHashState,
-    ParallelEngramEmbedding,
+    Engram as CommonEngram,
 )
+from vllm.models.deepseek_v4_1.common.engram import (
+    NgramHashState,
+)
+from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops
+from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding
 from vllm.platforms import current_platform
 
 
@@ -107,7 +111,7 @@ def test_fused_engram_post_wkv_matches_reference(
     )
     monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
 
-    module = Engram.__new__(Engram)
+    module = CommonEngram.__new__(CommonEngram)
     torch.nn.Module.__init__(module)
     module.dim = dim
     module.hc_mult = hc_mult
@@ -116,7 +120,7 @@ def test_fused_engram_post_wkv_matches_reference(
     module.use_sequence_parallel = use_sequence_parallel
     module.embed_tokens = torch.nn.Identity()
     # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
-    module.embed_tokens.tp_size = 1
+    module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1
     module.staged_rows = kv.unsqueeze(1)
     if use_sequence_parallel:
         padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
@@ -151,6 +155,25 @@ def _hash_state(use_slot_cache: bool) -> NgramHashState:
     return state
 
 
+@pytest.mark.parametrize("num_tokens", [0, 7])
+def test_engram_dummy_hashes_leave_history_untouched(num_tokens):
+    """Dummy forwards have no valid table IDs and do not alter cached history."""
+    state = _hash_state(use_slot_cache=True)
+    state.multipliers = torch.empty(2, 4, dtype=torch.int64)
+    state.primes = torch.empty(2, 3, 8, dtype=torch.int64)
+    state._cache = torch.arange(16, dtype=torch.int32)
+    history = state._cache.clone()
+
+    hashes, keep = state.dummy_hashes(torch.arange(num_tokens))
+
+    assert hashes.shape == (num_tokens, 2, 24)
+    assert hashes.dtype == torch.int32
+    assert torch.all(hashes == engram_ops.DEAD_ID)
+    assert keep.shape == (num_tokens,) and keep.dtype == torch.bool
+    assert not keep.any()
+    torch.testing.assert_close(state._cache, history)
+
+
 @pytest.mark.parametrize("num_blocks", [4, 100])
 def test_engram_cache_follows_kv_cache_binding(num_blocks):
     """Discard profiling history when the real KV cache replaces its storage."""
@@ -533,9 +556,11 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
     layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding)
     torch.nn.Module.__init__(layer)
     layer.dim, layer.block_size, layer.tp_size = dim, block, 1
+    layer.dp_size = 1
     layer.n_hash_cols = layer.part_n_hash_cols = 24
     layer.head_start = 0
     layer.cpu_offload = cpu_offload
+    layer.dp_shared_memory = False
     layer.part_num_embeddings = rows
     # A window strictly inside the table, so unowned rows are exercised too.
     layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2
@@ -559,12 +584,27 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
     return layer
 
 
+@pytest.mark.parametrize(
+    "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)]
+)
+def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch):
+    """Reject empty owners before allocating weights or accessing CUDA."""
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size)
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    with pytest.raises(AssertionError, match="ranks without hash heads"):
+        ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads)
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
 @pytest.mark.parametrize("cpu_offload", [False, True])
 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
 def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
     """Keep complete buckets and reconstruct head order, including TP padding."""
-    head_sizes = (17, 19, 23, 29, 31, 37)
+    head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73)
     num_rows, dim = sum(head_sizes), 64
     torch.manual_seed(0)
     weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn)
@@ -623,7 +663,7 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp
     module.staged_rows = torch.empty_like(shards[0])
     module.prepare_embeddings(ids)
     torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0)
-    # Slicing before head reordering must preserve padded heads and empty owners.
+    # Slicing before head reordering must preserve padded heads and tokens.
     module.use_sequence_parallel = True
     chunk = (len(ids) + tp_size - 1) // tp_size
     padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size))
@@ -636,6 +676,31 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp
         )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_engram_lookup_reuses_jit_across_token_shapes():
+    """Runtime token counts and launch grids must share one JIT variant."""
+    layer = _make_embedding(cpu_offload=False)
+    kernel = engram_ops._engram_lookup_kernel
+    kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0]
+    kernel_cache.clear()
+
+    cache_sizes = []
+    for background in (False, True):
+        for num_tokens in (1, 7, 256):
+            ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda")
+            out = torch.empty(
+                num_tokens,
+                24,
+                layer.dim,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+            layer.lookup(ids, out, background=background)
+            cache_sizes.append(len(kernel_cache))
+
+    assert cache_sizes == [1] * len(cache_sizes)
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
 @pytest.mark.parametrize("cpu_offload", [False, True])
 @pytest.mark.parametrize("background", [False, True])
@@ -660,19 +725,147 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
-@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize(
+    "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
+)
+def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
+    """Default offload uses pinned storage and a side stream; False uses HBM."""
+    from vllm.config import EngramConfig
+
+    config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6)
+    layout = SimpleNamespace(
+        num_embeddings=(72,),
+        head_dim=32,
+        max_ngram_size=2,
+        n_heads=24,
+        primes=(((3,) * 24,),),
+    )
+    vllm_config = SimpleNamespace(
+        engram_config=EngramConfig()
+        if cpu_offload is None
+        else EngramConfig(cpu_offload=cpu_offload),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+    )
+    offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload
+    if backend == "common":
+        vllm_config.engram_config = None
+    monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config)
+    monkeypatch.setattr(
+        nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config
+    )
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity()
+    )
+    with torch.device("cuda"):
+        cls = Engram if backend == "nvidia" else CommonEngram
+        module = cls(config, None, layout, 0, False, "engram")
+    layer = module.embed_tokens
+    if offloaded:
+        assert layer.weight.device.type == "cpu" and layer.weight.is_pinned()
+    else:
+        assert layer.weight.device.type == "cuda"
+    layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn))
+    layer.weight_scale_inv.data.fill_(127)
+    streams = []
+    lookup = layer.lookup
+
+    def track_lookup(ids, out, background=False):
+        streams.append(torch.cuda.current_stream())
+        lookup(ids, out, background=background)
+
+    layer.lookup = track_lookup
+    main = torch.cuda.current_stream()
+    ids = torch.randint(
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+        (8, 24),
+        device="cuda",
+        dtype=torch.int32,
+    )
+    module.prepare_embeddings(ids)
+    expected = _reference_lookup(
+        layer.weight.cuda(),
+        layer.weight_scale_inv.cuda(),
+        ids,
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+    )
+    torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0)
+    if offloaded:
+        assert streams == [module._prefetch_stream]
+        assert streams[0] != main
+    else:
+        assert streams == [main]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "cpu_offload,delay",
+    [(False, None), (True, None), (True, "producer"), (True, "lookup")],
+)
 @pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
-def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture):
-    """Consume early lookup results after a break, with fresh IDs each replay."""
-    _run_engram_prepared_rows(cpu_offload, capture)
+def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay):
+    """Temporary lookup IDs survive allocator reuse and graph replay."""
+    _run_engram_prepared_rows(cpu_offload, capture, delay=delay)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"]
+)
+def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency):
+    """Negative controls must expose incorrect rows when a dependency is removed."""
+    if (
+        missing_dependency == "record_stream"
+        and torch.cuda.memory.get_allocator_backend() != "native"
+    ):
+        pytest.skip(
+            "The allocator reuse negative control requires the native allocator"
+        )
+
+    def start(self, hash_ids, rows, stream):
+        if missing_dependency != "producer_wait":
+            stream.wait_stream(torch.cuda.current_stream())
+        if missing_dependency != "record_stream":
+            hash_ids.record_stream(stream)
+        with torch.cuda.stream(stream):
+            self.embed_tokens.lookup(hash_ids, rows, background=True)
+
+    def finish(self, stream):
+        pass
+
+    monkeypatch.setattr(Engram, "_start_prefetch", start)
+    if missing_dependency == "consumer_wait":
+        monkeypatch.setattr(Engram, "_finish_prefetch", finish)
+    delay = "producer" if missing_dependency == "producer_wait" else "lookup"
+    try:
+        with pytest.raises(
+            AssertionError, match=r"Tensor-likes are not (?:equal|close)!"
+        ):
+            _run_engram_prepared_rows(True, "eager", delay=delay)
+    finally:
+        torch.accelerator.synchronize()
 
 
 def _run_engram_prepared_rows(
-    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False
+    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None
 ):
-    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
 
     layer = _make_embedding(cpu_offload)
+    if delay == "lookup":
+        lookup = layer.lookup
+
+        def delayed_lookup(indices, out, background=False):
+            torch.cuda._sleep(2_000_000)
+            lookup(indices, out, background=background)
+
+        layer.lookup = delayed_lookup
     cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64)
     layer.n_hash_cols = cols
     layer.tp_size = tp_size
@@ -681,6 +874,17 @@ def _run_engram_prepared_rows(
     engram = Engram.__new__(Engram)
     torch.nn.Module.__init__(engram)
     engram.embed_tokens = layer
+    engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None
+    # Exercise the production eager boundaries even when the test process
+    # imported Engram before breakable graphs were enabled.
+    if capture == "breakable":
+        import vllm.envs as envs
+
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True)
+            for name in ("_start_prefetch", "_finish_prefetch"):
+                fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name)))
+                setattr(engram, name, fn.__get__(engram, Engram))
     engram.use_sequence_parallel = use_sequence_parallel
     engram.staged_rows = torch.empty(
         num_tokens,
@@ -709,7 +913,12 @@ def _run_engram_prepared_rows(
         embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True)
 
     def step(cap=None):
-        engram.prepare_embeddings(src)
+        if delay == "producer":
+            torch.cuda._sleep(2_000_000)
+        # Drop the last reference to a non-contiguous input after launching lookup.
+        engram.prepare_embeddings(hashes.clone()[:, 1])
+        # Exercise same-size allocator reuse before consuming the prefetched rows.
+        torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1)
         if cap is not None:
             cap.add_eager(lambda: None)
         out.copy_(embed(src))
@@ -730,9 +939,12 @@ def _run_engram_prepared_rows(
         with torch.cuda.stream(warmup), graph:
             step(graph)
         torch.cuda.current_stream().wait_stream(warmup)
-        assert graph.num_graphs == 2
+        assert graph.num_eager_breaks == (3 if cpu_offload else 1)
 
     for _ in range(3):
+        # Eager prefill can overwrite staging rows between decode replays.
+        hashes.random_(0, layer.part_num_embeddings)
+        step()
         hashes.random_(0, layer.part_num_embeddings)
         if graph is None:
             step()

tests/test_config.py

@@ -1158,31 +1158,68 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in
     assert config.get_parallel_size(parallel) == expected
 
 
-def test_engram_rejects_elastic_cross_dp():
+@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"])
+def test_engram_rejects_elastic_cross_dp(option):
     parallel = ParallelConfig(
         tensor_parallel_size=4,
         data_parallel_size=2,
         distributed_executor_backend="mp",
     )
     parallel.enable_elastic_ep = True
-    with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"):
-        EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel)
+    with pytest.raises(ValueError, match=f"{option}.*elastic EP"):
+        EngramConfig(**{option: True}).verify_parallel_config(parallel)
 
 
-@pytest.mark.parametrize("legacy", [None, "0", "1"])
-def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy):
-    """Explicit settings must override the legacy environment fallback."""
-    monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
-    if legacy is not None:
-        monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
-    assert EngramConfig().cpu_offload == (legacy == "1")
-    assert EngramConfig(cpu_offload=False).cpu_offload is False
-    assert EngramConfig(cpu_offload=True).cpu_offload is True
+def test_engram_dp_shared_memory_requires_cpu_offload():
+    with pytest.raises(ValueError, match="requires cpu_offload"):
+        EngramConfig(cpu_offload=False, dp_shared_memory=True)
+
+
+@pytest.mark.parametrize(
+    "dp_size,load_format,multithread,error",
+    [
+        (1, "auto", False, "requires data_parallel_size > 1"),
+        (2, "dummy", False, "requires load_format"),
+        (2, "sharded_state", False, "requires load_format"),
+        (2, "auto", False, None),
+        (2, "safetensors", True, None),
+        (2, "pt", True, None),
+    ],
+)
+def test_engram_dp_shared_memory_config_validation(
+    monkeypatch, dp_size, load_format, multithread, error
+):
+    """Reject invalid shared configs before distributed init; allow threaded loads."""
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    config = cast(
+        VllmConfig,
+        SimpleNamespace(
+            model_config=SimpleNamespace(
+                architecture="DeepseekV41ForCausalLM",
+                hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
+            ),
+            speculative_config=None,
+            engram_config=EngramConfig(cpu_offload=True, dp_shared_memory=True),
+            parallel_config=ParallelConfig(data_parallel_size=dp_size),
+            load_config=LoadConfig(
+                load_format=load_format,
+                model_loader_extra_config={"enable_multithread_load": multithread},
+            ),
+        ),
+    )
+    if error:
+        with pytest.raises(ValueError, match=error):
+            VllmConfig._resolve_and_verify_engram_config(config)
+    else:
+        VllmConfig._resolve_and_verify_engram_config(config)
 
 
 @pytest.mark.parametrize(
     "architecture, ple_layers, cuda, supported",
     [
+        ("DeepseekV41ForCausalLM", [1], True, True),
+        ("DeepseekV41ForCausalLM", [], True, False),
+        ("DeepseekV41ForCausalLM", [1], False, False),
         ("Qwen4ExpForCausalLM", [1], True, True),
         ("Qwen4ExpForConditionalGeneration", [1], True, True),
         ("Qwen4ExpForCausalLM", [], True, False),
@@ -1201,7 +1238,9 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo
             ModelConfig,
             SimpleNamespace(
                 architecture=architecture,
-                hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers),
+                hf_text_config=SimpleNamespace(
+                    ple_layer_ids=ple_layers, engram_layer_ids=ple_layers
+                ),
             ),
         )
         if architecture is not None
@@ -1214,55 +1253,73 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo
         with pytest.raises(ValueError, match="requires a model with supported Engram"):
             config.verify_model_config(model)
 
+    resolved = cast(
+        VllmConfig,
+        SimpleNamespace(
+            model_config=model,
+            speculative_config=None,
+            parallel_config=ParallelConfig(),
+            load_config=LoadConfig(load_format="dummy"),
+            engram_config=None,
+        ),
+    )
+    VllmConfig._resolve_and_verify_engram_config(resolved)
+    assert (resolved.engram_config is not None) == supported
+    if supported:
+        assert resolved.engram_config.cpu_offload is True
 
-def test_engram_config_defaults_to_none(monkeypatch):
+
+@pytest.mark.parametrize(
+    ("value", "expected"), [(None, True), ("0", False), ("1", True)]
+)
+def test_engram_cpu_offload_environment_default(monkeypatch, value, expected):
     monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
+    if value is not None:
+        monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", value)
+    assert EngramConfig().cpu_offload is expected
+    assert EngramConfig(cpu_offload=False).cpu_offload is False
+    assert EngramConfig(cpu_offload=True).cpu_offload is True
+
+
+def test_engram_config_defaults_to_none():
     config = VllmConfig()
     assert config.engram_config is None
     assert config.compute_hash()
 
 
-@pytest.mark.parametrize("legacy", ["0", "1"])
-def test_engram_none_resolves_legacy_offload(monkeypatch, legacy):
-    """Legacy enablement materializes a config before model validation."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
-    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+@pytest.mark.parametrize("explicit", [False, True])
+@pytest.mark.parametrize("enable_dbo,ubatch_size", [(True, 0), (False, 2)])
+def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size):
+    """Reject overlapping Engram staging even without an explicit EngramConfig."""
     config = cast(
         VllmConfig,
         SimpleNamespace(
             model_config=SimpleNamespace(
-                architecture="Qwen4ExpForCausalLM",
-                hf_text_config=SimpleNamespace(ple_layer_ids=[1]),
+                architecture="DeepseekV41ForCausalLM",
+                hf_text_config=SimpleNamespace(engram_layer_ids=[1]),
             ),
             speculative_config=None,
-            engram_config=None,
-            parallel_config=ParallelConfig(),
+            engram_config=EngramConfig(cpu_offload=False) if explicit else None,
+            parallel_config=ParallelConfig(
+                enable_dbo=enable_dbo, ubatch_size=ubatch_size
+            ),
         ),
     )
-    VllmConfig._resolve_and_verify_engram_config(config)
-    if legacy == "1":
-        assert config.engram_config is not None
-        assert config.engram_config.cpu_offload is True
-    else:
-        assert config.engram_config is None
+    with pytest.raises(
+        ValueError, match="Engram does not support DBO or microbatching"
+    ):
+        VllmConfig._resolve_and_verify_engram_config(config)
 
 
-@pytest.mark.parametrize("legacy", ["0", "1"])
-def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy):
+def test_engram_explicit_config_requires_supported_model():
     """Explicit all-false settings still opt into model validation."""
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
     with pytest.raises(ValueError, match="requires a model with supported Engram"):
         VllmConfig(engram_config=EngramConfig(cpu_offload=False))
 
 
-def test_engram_legacy_offload_requires_supported_model(monkeypatch):
-    monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
-    with pytest.raises(ValueError, match="requires a model with supported Engram"):
-        VllmConfig()
-
-
 @pytest.mark.parametrize("target_has_ple", [False, True])
-def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
+@pytest.mark.parametrize("explicit", [False, True])
+def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit):
     """MTP may inherit cross-DP sharding without having its own PLE layers."""
     monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
     target = SimpleNamespace(
@@ -1277,24 +1334,29 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
             speculative_config=SimpleNamespace(
                 draft_model_config=draft, target_model_config=target
             ),
-            engram_config=EngramConfig(embedding_across_dp=True),
+            engram_config=EngramConfig(embedding_across_dp=True) if explicit else None,
             parallel_config=ParallelConfig(),
+            load_config=LoadConfig(load_format="dummy"),
         ),
     )
-    if target_has_ple:
-        VllmConfig._resolve_and_verify_engram_config(config)
-    else:
+    if explicit and not target_has_ple:
         with pytest.raises(ValueError, match="requires a model with supported Engram"):
             VllmConfig._resolve_and_verify_engram_config(config)
+    else:
+        VllmConfig._resolve_and_verify_engram_config(config)
+        assert (config.engram_config is not None) == target_has_ple
+        if target_has_ple:
+            assert config.engram_config.cpu_offload is True
 
 
-def test_engram_hash_tracks_storage_and_sharding():
+def test_engram_hash_tracks_execution_options():
     configs = [
-        EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp)
-        for offload in (False, True)
-        for across_dp in (False, True)
+        EngramConfig(cpu_offload=True),
+        EngramConfig(cpu_offload=False),
+        EngramConfig(cpu_offload=True, embedding_across_dp=True),
+        EngramConfig(cpu_offload=True, dp_shared_memory=True),
     ]
-    assert len({config.compute_hash() for config in configs}) == 4
+    assert len({config.compute_hash() for config in configs}) == len(configs)
 
 
 @pytest.mark.parametrize("port", [1, 29550, 65535])

vllm/config/engram.py

@@ -3,18 +3,17 @@
 
 from typing import TYPE_CHECKING
 
-from pydantic import Field
+from pydantic import Field, model_validator
+from typing_extensions import Self
 
 import vllm.envs as envs
 from vllm.config.utils import config, get_hash_factors, hash_factors
-from vllm.logger import init_logger
 
 if TYPE_CHECKING:
+    from vllm.config.load import LoadConfig
     from vllm.config.model import ModelConfig
     from vllm.config.parallel import ParallelConfig
 
-logger = init_logger(__name__)
-
 # Architecture -> the hf_text_config field naming its n-gram layers. A model is
 # only configurable here if it actually has such layers to store.
 _NGRAM_LAYER_FIELDS = {
@@ -25,28 +24,44 @@ _NGRAM_LAYER_FIELDS = {
 
 
 def _default_cpu_offload() -> bool:
-    """Honor the legacy environment variable only when the field is omitted."""
-    if envs.is_set("VLLM_PLE_CPU_OFFLOAD"):
-        logger.warning_once(
-            "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a "
-            "future release. Use --engram-config.cpu_offload instead."
-        )
     return envs.VLLM_PLE_CPU_OFFLOAD
 
 
+def model_has_engram_layers(model_config: "ModelConfig | None") -> bool:
+    """Whether the model carries n-gram embedding layers."""
+    if model_config is None:
+        return False
+    field = _NGRAM_LAYER_FIELDS.get(model_config.architecture)
+    if field is None:
+        return False
+    return bool(getattr(model_config.hf_text_config, field, None))
+
+
 @config
 class EngramConfig:
     """Configuration for Engram embedding storage and sharding."""
 
     cpu_offload: bool = Field(default_factory=_default_cpu_offload)
     """Store embedding weights in pinned CPU memory for UVA lookup.
-    Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility.
-    An explicit value takes precedence over the legacy environment variable."""
+    Defaults to VLLM_PLE_CPU_OFFLOAD, which is enabled by default. An explicit
+    value takes precedence over the environment variable."""
 
     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."""
 
+    dp_shared_memory: bool = False
+    """Share CPU-offloaded embedding weights between co-located
+    DP replicas. Each node stores one copy of every TP shard, reducing host
+    memory without per-step Engram DP collectives. Requires sufficient
+    /dev/shm capacity and a shared IPC namespace."""
+
+    @model_validator(mode="after")
+    def _validate_shared_memory(self) -> Self:
+        if self.dp_shared_memory and not self.cpu_offload:
+            raise ValueError("dp_shared_memory requires cpu_offload=True")
+        return self
+
     def verify_model_config(self, model_config: "ModelConfig | None") -> None:
         """Reject Engram configuration for models without n-gram embeddings."""
         from vllm.platforms import current_platform
@@ -64,13 +79,16 @@ class EngramConfig:
         ):
             raise ValueError(
                 "EngramConfig requires a model with supported Engram "
-                "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek "
-                "V4.1 implementations with non-empty n-gram layer ids are "
-                "supported."
+                "embeddings, non-empty n-gram layer ids, and CUDA."
             )
 
     def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
         """Reject unsupported embedding parallel topologies."""
+        if self.dp_shared_memory:
+            if parallel_config.data_parallel_size <= 1:
+                raise ValueError("dp_shared_memory requires data_parallel_size > 1.")
+            if parallel_config.enable_elastic_ep:
+                raise ValueError("dp_shared_memory is not supported with elastic EP.")
         if (
             self.embedding_across_dp
             and parallel_config.data_parallel_size > 1
@@ -80,6 +98,18 @@ class EngramConfig:
                 "Engram embedding_across_dp is not supported with elastic EP yet."
             )
 
+    def verify_load_config(self, load_config: "LoadConfig") -> None:
+        """Shared tables require a loader that invokes parameter weight callbacks."""
+        if self.dp_shared_memory and load_config.load_format not in (
+            "auto",
+            "safetensors",
+            "pt",
+        ):
+            raise ValueError(
+                "dp_shared_memory requires load_format 'auto', "
+                f"'safetensors' or 'pt'; got {load_config.load_format!r}."
+            )
+
     def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
         """Derive the embedding group size from the parallel configuration."""
         size = parallel_config.tensor_parallel_size

vllm/config/vllm.py

@@ -34,7 +34,7 @@ from .device import DeviceConfig
 from .diffusion import DiffusionConfig
 from .ec_manager_config import EncoderCacheManagerConfig
 from .ec_transfer import ECTransferConfig
-from .engram import EngramConfig
+from .engram import EngramConfig, model_has_engram_layers
 from .kernel import KernelConfig
 from .kv_events import KVEventsConfig
 from .kv_transfer import KVTransferConfig
@@ -381,7 +381,7 @@ class VllmConfig:
     attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
     """Attention configuration."""
     engram_config: EngramConfig | None = None
-    """Optional Engram configuration, only valid for supported PLE models."""
+    """N-gram embedding storage and sharding settings."""
     mamba_config: MambaConfig = Field(default_factory=MambaConfig)
     """Mamba configuration."""
     kernel_config: KernelConfig = Field(default_factory=KernelConfig)
@@ -1218,22 +1218,37 @@ class VllmConfig:
             )
 
     def _resolve_and_verify_engram_config(self) -> None:
-        """Resolve legacy offload settings and validate model and parallel configs."""
-        if self.engram_config is None:
-            if not envs.VLLM_PLE_CPU_OFFLOAD:
-                return
-            self.engram_config = EngramConfig()
+        """Resolve defaults and validate n-gram embedding settings."""
+        from vllm.platforms import current_platform
+
         model_config = self.model_config
         speculative_config = self.speculative_config
         # Draft configs inherit the target's communication groups and settings.
-        # Qwen4Exp MTP itself disables PLE, so validate its target instead.
+        # Validate the target because the draft may disable n-gram embeddings.
         if (
             speculative_config is not None
             and model_config is speculative_config.draft_model_config
         ):
             model_config = speculative_config.target_model_config
+        if (
+            model_config is not None
+            and model_config.architecture == "DeepseekV41ForCausalLM"
+            and getattr(model_config.hf_text_config, "engram_layer_ids", None)
+            and self.parallel_config.use_ubatching
+        ):
+            raise ValueError(
+                "DeepSeek V4.1 Engram does not support DBO or microbatching. "
+                "Disable --enable-dbo and set --ubatch-size to 0."
+            )
+        if self.engram_config is None:
+            if not current_platform.is_cuda() or not model_has_engram_layers(
+                model_config
+            ):
+                return
+            self.engram_config = EngramConfig()
         self.engram_config.verify_model_config(model_config)
         self.engram_config.verify_parallel_config(self.parallel_config)
+        self.engram_config.verify_load_config(self.load_config)
         logger.info_once("Resolved Engram configuration: %s", str(self.engram_config))
 
     def __post_init__(self):

vllm/distributed/parallel_state.py

@@ -32,6 +32,7 @@ from collections.abc import Callable
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
+from math import gcd
 from multiprocessing import shared_memory
 from typing import TYPE_CHECKING, Any, Protocol
 from unittest.mock import patch
@@ -1558,6 +1559,23 @@ def get_etp_group() -> GroupCoordinator:
     return _ETP
 
 
+_ENGRAM_DP: GroupCoordinator | None = None
+
+
+def get_engram_dp_group() -> GroupCoordinator | None:
+    """Return the DP replicas that share one engram embedding table.
+
+    None when every replica holds a full (TP-sharded) copy, which is the case
+    for models without engram layers and for replicas that span nodes.
+    """
+    return _ENGRAM_DP
+
+
+def get_engram_dp_size() -> int:
+    """Number of DP replicas one engram embedding table is sharded over."""
+    return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1
+
+
 _DCP: GroupCoordinator | None = None
 
 
@@ -1933,6 +1951,29 @@ def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
     return nullcontext()
 
 
+def _engram_dp_shard_size(
+    world_size: int,
+    data_parallel_size: int,
+    replica_size: int,
+) -> int:
+    """Size of DP shards aligned with complete replicas and node boundaries."""
+    node_count = get_node_count()
+    replicas_per_node, remainder = divmod(world_size, node_count * replica_size)
+    if remainder:
+        return 1
+    # Shards must divide both the DP group and each node's replica count.
+    shard_size = gcd(data_parallel_size, replicas_per_node)
+    if shard_size == 1 or node_count == 1:
+        return shard_size
+    ranks_per_node = replicas_per_node * replica_size
+    for start in range(0, world_size, ranks_per_node):
+        same_node = in_the_same_node_as(get_world_group().cpu_group, start)
+        node_ranks = [rank for rank, local in enumerate(same_node) if local]
+        if node_ranks != list(range(start, start + ranks_per_node)):
+            return 1
+    return shard_size
+
+
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     pipeline_model_parallel_size: int = 1,
@@ -2055,6 +2096,40 @@ def initialize_model_parallel(
             group_name="etp",
         )
 
+    # Build the node-local DP group used to shard or share an Engram table.
+    global _ENGRAM_DP
+    assert _ENGRAM_DP is None, "engram data parallel group is already initialized"
+    engram_dp_size = 1
+    engram_config = config.engram_config
+    if (
+        engram_config is not None
+        and config.model_config is not None
+        and config.model_config.architecture == "DeepseekV41ForCausalLM"
+        and not enable_elastic_ep
+    ):
+        engram_dp_size = _engram_dp_shard_size(
+            world_size,
+            data_parallel_size,
+            tensor_model_parallel_size
+            * pipeline_model_parallel_size
+            * prefill_context_model_parallel_size,
+        )
+    if engram_dp_size > 1:
+        group_ranks = (
+            all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0)
+        )
+        _ENGRAM_DP = init_model_parallel_group(
+            [x.tolist() for x in group_ranks],
+            get_world_group().local_rank,
+            backend,
+            group_name="edp",
+        )
+        logger.info_once(
+            "Engram node-local DP group size: %d (TP=%d)",
+            engram_dp_size,
+            tensor_model_parallel_size,
+        )
+
     # Build the DCP model-parallel groups.
     global _DCP
     assert _DCP is None, "decode context model parallel group is already initialized"
@@ -2304,6 +2379,11 @@ def destroy_model_parallel():
         _TP.destroy()
     _TP = None
 
+    global _ENGRAM_DP
+    if _ENGRAM_DP:
+        _ENGRAM_DP.destroy()
+    _ENGRAM_DP = None
+
     global _DCP
     if _DCP:
         _DCP.destroy()

vllm/envs.py

@@ -154,7 +154,7 @@ if TYPE_CHECKING:
     VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False
     VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
     VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
-    VLLM_PLE_CPU_OFFLOAD: bool = False
+    VLLM_PLE_CPU_OFFLOAD: bool = True
     VLLM_DISABLE_COMPILE_CACHE: bool = False
     VLLM_REPLICATE_EMBED: bool = False
     VLLM_USE_LAYERNAME: bool = True
@@ -2059,11 +2059,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_LOG_MODEL_INSPECTION": lambda: bool(
         int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0"))
     ),
-    # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their
-    # rows through UVA on a dedicated CUDA stream.
-    # Legacy fallback for EngramConfig.cpu_offload, which takes precedence.
-    # This environment variable may be removed in a future release.
-    "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))),
+    # Store n-gram embedding tables in pinned CPU memory for UVA lookup.
+    "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "1"))),
     # Debug logging for --enable-mfu-metrics
     "VLLM_DEBUG_MFU_METRICS": lambda: bool(
         int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))

vllm/models/deepseek_v4_1/common/engram.py

@@ -51,8 +51,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.triton_utils import tl, triton
-from vllm.utils.platform_utils import is_uva_available
-from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
 
 logger = init_logger(__name__)
 
@@ -236,6 +234,7 @@ def _write_hash_cache_kernel(
     tl.store(cache + slot, value, valid)
 
 
+# Keep request shapes out of the cache key so one warmed variant covers runtime batches.
 @triton.jit(
     do_not_specialize=[
         "num_tokens",
@@ -462,6 +461,21 @@ class NgramHashState(nn.Module):
         self._kv_cache_ref = weakref.ref(kv_cache)
         return True
 
+    def dummy_hashes(
+        self, input_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Participate in DP lookups without valid rows or hash-cache updates."""
+        num_tokens = input_ids.shape[0]
+        num_layers, max_ngram = self.multipliers.shape
+        num_heads = self.primes.shape[-1]
+        hashes = input_ids.new_full(
+            (num_tokens, num_layers, (max_ngram - 1) * num_heads),
+            DEAD_ID,
+            dtype=torch.int32,
+        )
+        keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device)
+        return hashes, keep
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -561,7 +575,16 @@ def _engram_head_shard_weight_loader(
     param.data.copy_(shard)
 
 
-@triton.jit
+@triton.jit(
+    do_not_specialize=[
+        "vocab_start",
+        "vocab_end",
+        "num_rows",
+        "ids_stride_t",
+        "ids_stride_h",
+        "GRID",
+    ]
+)
 def _engram_lookup_kernel(
     weight,
     scales,
@@ -578,7 +601,7 @@ def _engram_lookup_kernel(
     DIM: tl.constexpr,
     QUANT_BLOCK: tl.constexpr,
     BLOCK_R: tl.constexpr,
-    GRID: tl.constexpr,
+    GRID,
 ):
     """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
 
@@ -620,12 +643,9 @@ def _engram_lookup_kernel(
 
 
 class ParallelEngramEmbedding(nn.Module):
-    """The n-gram hash table, sharded by complete hash heads over TP ranks.
-    Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup.
+    """TP-sharded hash heads with FP8 rows and per-block E8M0 scales."""
 
-    With `cpu_offload` the shard lives in pinned host memory and is read over
-    UVA instead of HBM; the TP sharding is unchanged either way.
-    """
+    _weight_loader = staticmethod(_engram_head_shard_weight_loader)
 
     def __init__(
         self,
@@ -633,89 +653,55 @@ class ParallelEngramEmbedding(nn.Module):
         dim: int,
         head_sizes: tuple[int, ...],
         block_size: int = 32,
-        cpu_offload: bool = False,
-    ):
+    ) -> None:
         super().__init__()
-        tp_size = get_tensor_model_parallel_world_size()
-        tp_rank = get_tensor_model_parallel_rank()
         assert head_sizes and all(size > 0 for size in head_sizes)
         assert sum(head_sizes) <= num_embeddings
-        if cpu_offload and not is_uva_available():
-            raise RuntimeError("Engram CPU offload requires UVA support")
         self.num_embeddings = num_embeddings
         self.dim = dim
         self.block_size = block_size
         self.n_hash_cols = len(head_sizes)
-        self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size)
-        self.head_start = tp_rank * self.part_n_hash_cols
+        self.tp_size = get_tensor_model_parallel_world_size()
+        num_shards, head_rank = self._get_shard_info()
+        self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards)
+        # TODO: Support row-wise sharding when there are too few hash heads.
+        assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, (
+            f"Engram sharding leaves ranks without hash heads: "
+            f"{self.n_hash_cols} heads over {num_shards} shards"
+        )
+        self.head_start = head_rank * self.part_n_hash_cols
         head_end = self.head_start + self.part_n_hash_cols
         self.vocab_start_idx = sum(head_sizes[: self.head_start])
         self.vocab_end_idx = sum(head_sizes[:head_end])
         self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx
-        self.tp_size = tp_size
-        self.cpu_offload = cpu_offload
-        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
-        self._view_src: tuple[int, int] | None = None
         self._num_sms = torch.cuda.get_device_properties(
             torch.accelerator.current_device_index()
         ).multi_processor_count
-
-        # Explicit device: model init runs under a `torch.device("cuda")`
-        # context, which would otherwise put the shard in HBM.
-        kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {}
-        self.weight = nn.Parameter(
-            torch.empty(
-                self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs
-            ),
-            requires_grad=False,
-        )
-        self.weight_scale_inv = nn.Parameter(
-            torch.empty(
-                self.part_num_embeddings,
-                dim // block_size,
-                dtype=torch.uint8,
-                **kwargs,
-            ),
-            requires_grad=False,
-        )
+        weight, scales = self._allocate_weights()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.weight_scale_inv = nn.Parameter(scales, requires_grad=False)
         for param in (self.weight, self.weight_scale_inv):
             set_weight_attrs(
                 param,
                 {
-                    "weight_loader": _engram_head_shard_weight_loader,
+                    "weight_loader": self._weight_loader,
                     "engram_vocab_start": self.vocab_start_idx,
                 },
             )
-        if cpu_offload:
-            # Constant dummy values avoid randomizing huge CPU lookup tables.
-            set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
-            # The ue8m0 encoding of scale 1.0 is exponent byte 127.
-            set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
-            logger.info(
-                "Engram table offloaded to pinned host memory: %d rows x %d, "
-                "%.2f GiB per rank",
-                self.part_num_embeddings,
-                dim,
-                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
-            )
+
+    def _get_shard_info(self) -> tuple[int, int]:
+        return self.tp_size, get_tensor_model_parallel_rank()
+
+    def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return (
+            torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn),
+            torch.empty(
+                self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8
+            ),
+        )
 
     def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
-        """Parameters when resident, else cached UVA views of the pinned shard.
-
-        Rebuilt if anything swaps `.data`, so a stale device pointer cannot
-        survive silently.
-        """
-        if not self.cpu_offload:
-            return self.weight.data, self.weight_scale_inv.data
-        src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
-        if self._view_src != src:
-            self._views = (
-                get_accelerator_view_from_cpu_tensor(self.weight.data),
-                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
-            )
-            self._view_src = src
-        assert self._views is not None
-        return self._views
+        return self.weight.data, self.weight_scale_inv.data
 
     def lookup(
         self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
@@ -753,7 +739,7 @@ class ParallelEngramEmbedding(nn.Module):
 
     def forward(self, indices: torch.Tensor) -> torch.Tensor:
         """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
-        bf16, gathered from all TP shards."""
+        bf16, gathered from all shards for this replica's tokens."""
         out = torch.empty(
             (indices.shape[0], self.part_n_hash_cols, self.dim),
             dtype=torch.bfloat16,
@@ -762,8 +748,7 @@ class ParallelEngramEmbedding(nn.Module):
         self.lookup(indices, out)
         if self.tp_size > 1:
             out = tensor_model_parallel_all_gather(out, dim=1)
-            out = out[:, : self.n_hash_cols]
-        return out
+        return out[:, : self.n_hash_cols]
 
 
 @triton.jit(do_not_specialize=["num_kv_tokens"])
@@ -859,7 +844,7 @@ def _fused_engram_post_wkv_kernel(
 
 
 @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"])
-def _engram_sp_rows_kernel(
+def _engram_select_rows_kernel(
     gathered,
     output,
     num_tokens,
@@ -880,6 +865,32 @@ def _engram_sp_rows_kernel(
     tl.store(output + offsets, values, offsets < num_elements)
 
 
+def _engram_select_rows(
+    gathered: torch.Tensor,
+    output: torch.Tensor,
+    source_tokens: int,
+    token_start: int,
+    local_width: int,
+) -> None:
+    """Copy one token window out of a rank-major gathered buffer.
+
+    Both gathers land rank-major ([rank][token][local width]); this walks the
+    window the rank keeps and lays its ranks out side by side as width.
+    """
+    if output.numel() == 0:
+        return
+    _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)](
+        gathered,
+        output,
+        source_tokens,
+        token_start,
+        output.numel(),
+        local_width,
+        output.shape[1] * output.shape[2],
+        BLOCK_SIZE=1024,
+    )
+
+
 class Engram(nn.Module):
     """Writes an n-gram lookup into the residual stream, gated by how well it
     matches that stream.
@@ -910,13 +921,7 @@ class Engram(nn.Module):
         # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight``
         # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight``
         # suffix rule.
-        engram_config = get_current_vllm_config().engram_config
-        self.embed_tokens = ParallelEngramEmbedding(
-            layout.num_embeddings[layer_hash_index],
-            layout.head_dim,
-            tuple(size for order in layout.primes[layer_hash_index] for size in order),
-            cpu_offload=engram_config.cpu_offload if engram_config else True,
-        )
+        self.embed_tokens = self._create_embedding(layout, layer_hash_index)
         n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads
         self.wkv = ReplicatedLinear(
             n_hash_cols * layout.head_dim,
@@ -936,40 +941,54 @@ class Engram(nn.Module):
         )
 
         max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens
-        # Keep lookup results alive across breakable graph segments.
+        self._init_staging(max_tokens, layout.head_dim)
+
+    def _create_embedding(
+        self, layout: EngramLayout, layer_hash_index: int
+    ) -> ParallelEngramEmbedding:
+        return ParallelEngramEmbedding(
+            layout.num_embeddings[layer_hash_index],
+            layout.head_dim,
+            tuple(size for order in layout.primes[layer_hash_index] for size in order),
+        )
+
+    def _init_staging(self, max_tokens: int, head_dim: int) -> None:
+        # Persistent storage keeps lookup addresses stable across graph replays.
         self.staged_rows = torch.empty(
             max_tokens,
             self.embed_tokens.part_n_hash_cols,
-            layout.head_dim,
+            head_dim,
             dtype=torch.bfloat16,
         )
 
     def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
-        """Gather this layer's rows on the main stream before decoder layers."""
-        self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]])
+        """Look up rows before the decoder layers consume them."""
+        rows = self.staged_rows[: hash_ids.shape[0]]
+        assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
+        self.embed_tokens.lookup(hash_ids, rows)
+
+    def _ready_rows(self, num_tokens: int) -> torch.Tensor:
+        return self.staged_rows[:num_tokens]
 
     def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
         """Gather heads, returning only local tokens when SP is enabled."""
-        rows = self.staged_rows[: hash_ids.shape[0]]
+        rows = self._ready_rows(hash_ids.shape[0])
         if self.embed_tokens.tp_size == 1:
-            return rows
+            return rows[:, : self.embed_tokens.n_hash_cols]
         if self.use_sequence_parallel:
             tp_size = self.embed_tokens.tp_size
             num_tokens, local_heads, dim = rows.shape
             gathered = tensor_model_parallel_all_gather(rows, dim=0)
             chunk = (num_tokens + tp_size - 1) // tp_size
-            rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
-            _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)](
+            out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim))
+            _engram_select_rows(
                 gathered,
-                rows,
+                out,
                 num_tokens,
                 get_tensor_model_parallel_rank() * chunk,
-                rows.numel(),
                 local_heads * dim,
-                self.embed_tokens.n_hash_cols * dim,
-                BLOCK_SIZE=1024,
             )
-            return rows
+            return out
         rows = tensor_model_parallel_all_gather(rows, dim=1)
         return rows[:, : self.embed_tokens.n_hash_cols]
 

vllm/models/deepseek_v4_1/nvidia/engram.py

@@ -0,0 +1,403 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
+
+import mmap
+import tempfile
+import weakref
+from contextlib import ExitStack
+
+import numpy as np
+import torch
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import get_current_vllm_config
+from vllm.distributed import (
+    get_dp_group,
+    get_engram_dp_group,
+    get_engram_dp_size,
+    get_tensor_model_parallel_rank,
+    tensor_model_parallel_all_gather,
+)
+from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.deepseek_v4_1.common.engram import (
+    DEAD_ID,
+    EngramLayout,
+    _engram_head_shard_weight_loader,
+    _engram_select_rows,
+)
+from vllm.models.deepseek_v4_1.common.engram import (
+    Engram as BaseEngram,
+)
+from vllm.models.deepseek_v4_1.common.engram import (
+    ParallelEngramEmbedding as BaseParallelEngramEmbedding,
+)
+from vllm.utils.platform_utils import is_uva_available
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
+
+logger = init_logger(__name__)
+
+
+def engram_head_shard_rank() -> int:
+    """This rank's slot among the hash-head shards of one engram table.
+
+    TP-major, so the shards a DP gather brings in are contiguous heads and
+    the following TP gather completes the head order.
+    """
+    dp_group = get_engram_dp_group()
+    dp_size = dp_group.world_size if dp_group is not None else 1
+    dp_rank = dp_group.rank_in_group if dp_group is not None else 0
+    return get_tensor_model_parallel_rank() * dp_size + dp_rank
+
+
+def engram_gathered_num_tokens() -> int:
+    """Per-replica token slot for the node-local Engram DP group."""
+    dp_metadata = get_forward_context().dp_metadata
+    if dp_metadata is None:
+        raise RuntimeError("a DP-shared engram table needs DP token metadata")
+    group = get_engram_dp_group()
+    assert group is not None
+    # Engram groups are contiguous slices of the full DP group.
+    start = get_dp_group().rank_in_group - group.rank_in_group
+    return int(
+        dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max()
+    )
+
+
+def gather_engram_hashes(
+    hash_ids: torch.Tensor, *, dp_shared_memory: bool = False
+) -> torch.Tensor:
+    """Collect the n-gram ids of every DP replica sharing one table.
+
+    Replicas are padded to a common token slot, so the gathered shape is
+    static under CUDA graph capture (where DP already pads alike).
+    """
+    dp_group = get_engram_dp_group()
+    if dp_group is None or dp_shared_memory:
+        return hash_ids
+    slot = engram_gathered_num_tokens()
+    if hash_ids.shape[0] > slot:
+        raise ValueError("Engram token count exceeds the DP token slot")
+    if hash_ids.shape[0] < slot:
+        pad = hash_ids.new_full(
+            (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID
+        )
+        hash_ids = torch.cat((hash_ids, pad))
+    return dp_group.all_gather(hash_ids, dim=0)
+
+
+class DPSharedEngramStorage:
+    """Registered host weights shared by a node-local DP group with one writer."""
+
+    def __init__(
+        self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator
+    ) -> None:
+        self.group = group
+        weight_bytes = num_rows * dim
+        storage = self._allocate(weight_bytes + weight_bytes // block_size)
+        self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
+        self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size)
+        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
+
+    def _allocate(self, num_bytes: int) -> torch.Tensor:
+        """Map and register one physical allocation across a node-local DP group."""
+        from vllm.distributed.device_communicators.shm_broadcast import (
+            check_shm_free_space,
+        )
+
+        group = self.group
+        with ExitStack() as stack:
+            path, error = None, None
+            if group.rank_in_group == 0:
+                try:
+                    check_shm_free_space(num_bytes)
+                    backing_file = stack.enter_context(
+                        tempfile.NamedTemporaryFile(
+                            prefix="vllm_engram_", dir="/dev/shm"
+                        )
+                    )
+                    backing_file.truncate(num_bytes)
+                    path = backing_file.name
+                except Exception as exc:
+                    error = f"{type(exc).__name__}: {exc}"
+            path, error = group.broadcast_object((path, error))
+            if error is not None:
+                raise RuntimeError(
+                    "Engram shared-memory creation failed on EDP rank 0: " + error
+                )
+
+            mapping = owner = tensor = finalizer = None
+            stage = "open"
+            try:
+                try:
+                    with open(path, "r+b") as file:
+                        stage = "mmap"
+                        mapping = mmap.mmap(
+                            file.fileno(), num_bytes, flags=mmap.MAP_SHARED
+                        )
+                    stage = "cudaHostRegister"
+                    owner = np.frombuffer(mapping, dtype=np.uint8)
+                    pointer = owner.ctypes.data
+                    tensor = torch.from_numpy(owner)
+                    result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
+                    if result.value != 0:
+                        raise RuntimeError(f"cudaHostRegister failed: {result}")
+                    finalizer = weakref.finalize(
+                        owner, self._unregister, mapping, pointer
+                    )
+                    finalizer.atexit = False  # type: ignore[misc]
+                    # The UVA helper otherwise allocates a private pinned copy.
+                    if not tensor.is_pinned():
+                        raise RuntimeError(
+                            "CUDA did not recognize the shared Engram registration"
+                        )
+                except Exception as exc:
+                    error = f"{stage}: {type(exc).__name__}: {exc}"
+
+                errors: list[str | None] = [None] * group.world_size
+                # Also fences peer mappings before the leader unlinks the file.
+                torch.distributed.all_gather_object(
+                    errors, error, group=group.cpu_group
+                )
+                failures = "; ".join(
+                    f"EDP rank {rank}: {error}"
+                    for rank, error in enumerate(errors)
+                    if error is not None
+                )
+                if failures:
+                    raise RuntimeError(
+                        "Engram shared-memory initialization failed: " + failures
+                    )
+                assert tensor is not None
+                return tensor
+            except Exception:
+                if finalizer is not None:
+                    finalizer()
+                tensor = owner = None
+                if mapping is not None:
+                    mapping.close()
+                raise
+
+    @staticmethod
+    def _unregister(mapping: mmap.mmap, pointer: int) -> None:
+        # Torch storage retains the numpy owner, including through cached UVA views.
+        # Keep its mmap alive until CUDA has released the registration.
+        result = torch.cuda.cudart().cudaHostUnregister(pointer)
+        if result.value != 0:
+            logger.warning("Engram cudaHostUnregister failed: %s", result)
+
+    def load_weight(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        if self.group.rank_in_group == 0:
+            _engram_head_shard_weight_loader(param, loaded_weight)
+        # Read order may differ across ranks. Equal load counts ensure all shared
+        # weights are ready after the last weight-loader call returns.
+        torch.distributed.barrier(group=self.group.cpu_group)
+
+    def get_views(
+        self, weight: torch.Tensor, scales: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (weight.data_ptr(), scales.data_ptr()) != (
+            self.weight.data_ptr(),
+            self.weight_scale_inv.data_ptr(),
+        ):
+            raise RuntimeError("Shared Engram parameter storage must not be replaced")
+        if self._views is None:
+            self._views = (
+                get_accelerator_view_from_cpu_tensor(self.weight),
+                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv),
+            )
+        return self._views
+
+
+class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
+    """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices."""
+
+    _shared_memory: DPSharedEngramStorage | None = None
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        dim: int,
+        head_sizes: tuple[int, ...],
+        block_size: int = 32,
+        cpu_offload: bool = False,
+        dp_shared_memory: bool = False,
+    ) -> None:
+        self.cpu_offload = cpu_offload
+        self.dp_shared_memory = dp_shared_memory
+        self.dp_size = get_engram_dp_size()
+        if dp_shared_memory:
+            if not cpu_offload:
+                raise ValueError("dp_shared_memory requires cpu_offload=True")
+            if self.dp_size <= 1:
+                raise ValueError(
+                    "dp_shared_memory requires a node-local Engram DP "
+                    f"group with size > 1; effective Engram DP size is {self.dp_size}. "
+                    "Check that the node layout and rank placement allow complete "
+                    "DP replicas to be co-located."
+                )
+            self.dp_size = 1
+        if cpu_offload and not is_uva_available():
+            raise RuntimeError("Engram CPU offload requires UVA support")
+        self._views: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._view_src: tuple[int, int] | None = None
+        super().__init__(num_embeddings, dim, head_sizes, block_size)
+        if cpu_offload:
+            # Constant dummy values avoid randomizing huge CPU lookup tables.
+            set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
+            # The ue8m0 encoding of scale 1.0 is exponent byte 127.
+            set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
+            logger.info(
+                "Engram table offloaded to pinned host memory: %d rows x %d, "
+                "%.2f GiB %s",
+                self.part_num_embeddings,
+                dim,
+                self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
+                "shared across DP replicas" if dp_shared_memory else "per rank",
+            )
+
+    def _get_shard_info(self) -> tuple[int, int]:
+        if self.dp_size == 1:
+            return super()._get_shard_info()
+        return self.tp_size * self.dp_size, engram_head_shard_rank()
+
+    def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.dp_shared_memory:
+            group = get_engram_dp_group()
+            assert group is not None
+            storage = DPSharedEngramStorage(
+                self.part_num_embeddings, self.dim, self.block_size, group
+            )
+            self._shared_memory = storage
+            self._weight_loader = storage.load_weight
+            return storage.weight, storage.weight_scale_inv
+        if not self.cpu_offload:
+            return super()._allocate_weights()
+        # Model initialization may be inside a CUDA device context.
+        return (
+            torch.empty(
+                self.part_num_embeddings,
+                self.dim,
+                dtype=torch.float8_e4m3fn,
+                device="cpu",
+                pin_memory=True,
+            ),
+            torch.empty(
+                self.part_num_embeddings,
+                self.dim // self.block_size,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=True,
+            ),
+        )
+
+    def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._shared_memory is not None:
+            return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
+        if not self.cpu_offload:
+            return super()._storage()
+        src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
+        if self._view_src != src:
+            self._views = (
+                get_accelerator_view_from_cpu_tensor(self.weight.data),
+                get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
+            )
+            self._view_src = src
+        assert self._views is not None
+        return self._views
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        if self.dp_size == 1:
+            return super().forward(indices)
+        num_tokens = indices.shape[0]
+        indices = gather_engram_hashes(indices)
+        out = torch.empty(
+            (indices.shape[0], self.part_n_hash_cols, self.dim),
+            dtype=torch.bfloat16,
+            device=indices.device,
+        )
+        self.lookup(indices, out)
+        out = _gather_engram_rows(out, num_tokens)
+        if self.tp_size > 1:
+            out = tensor_model_parallel_all_gather(out, dim=1)
+        return out[:, : self.n_hash_cols]
+
+
+def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
+    """Exchange DP tokens for heads, retaining only this replica's tokens."""
+    dp_group = get_engram_dp_group()
+    assert dp_group is not None
+    slot, remainder = divmod(staged.shape[0], dp_group.world_size)
+    assert remainder == 0 and 0 <= num_tokens <= slot
+    gathered = dp_group.all_gather(staged, dim=0)
+    local_heads, dim = staged.shape[1:]
+    rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim))
+    _engram_select_rows(
+        gathered,
+        rows,
+        staged.shape[0],
+        dp_group.rank_in_group * slot,
+        local_heads * dim,
+    )
+    return rows
+
+
+class Engram(BaseEngram):
+    """NVIDIA Engram with asynchronous offload and node-local DP lookup."""
+
+    _prefetch_stream: torch.cuda.Stream | None = None
+
+    def _create_embedding(
+        self, layout: EngramLayout, layer_hash_index: int
+    ) -> ParallelEngramEmbedding:
+        engram_config = get_current_vllm_config().engram_config
+        assert engram_config is not None
+        return ParallelEngramEmbedding(
+            layout.num_embeddings[layer_hash_index],
+            layout.head_dim,
+            tuple(size for order in layout.primes[layer_hash_index] for size in order),
+            cpu_offload=engram_config.cpu_offload,
+            dp_shared_memory=engram_config.dp_shared_memory,
+        )
+
+    def _init_staging(self, max_tokens: int, head_dim: int) -> None:
+        super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
+        if self.embed_tokens.cpu_offload:
+            self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
+
+    def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
+        """Prefetch local shared rows or the DP group's gathered hash IDs."""
+        if self._prefetch_stream is None:
+            return super().prepare_embeddings(hash_ids)
+        rows = self.staged_rows[: hash_ids.shape[0]]
+        assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
+        self._start_prefetch(hash_ids, rows, self._prefetch_stream)
+
+    @eager_break_during_capture
+    def _start_prefetch(
+        self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream
+    ) -> None:
+        # Eager boundaries let the lookup span piecewise graph segments.
+        stream.wait_stream(torch.cuda.current_stream())
+        # Keep temporary hash storage alive until lookup finishes reading it.
+        hash_ids.record_stream(stream)
+        with torch.cuda.stream(stream):
+            self.embed_tokens.lookup(hash_ids, rows, background=True)
+
+    @eager_break_during_capture
+    def _finish_prefetch(self, stream: torch.cuda.Stream) -> None:
+        torch.cuda.current_stream().wait_stream(stream)
+
+    def _ready_rows(self, num_tokens: int) -> torch.Tensor:
+        if self._prefetch_stream is not None:
+            self._finish_prefetch(self._prefetch_stream)
+        if self.embed_tokens.dp_size > 1:
+            slot = engram_gathered_num_tokens()
+            staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
+            return _gather_engram_rows(staged, num_tokens)
+        return super()._ready_rows(num_tokens)

vllm/models/deepseek_v4_1/nvidia/model.py

@@ -12,6 +12,7 @@ import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.config.kernel import MEGA_MOE_BACKENDS
 from vllm.distributed import (
+    get_engram_dp_size,
     get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -74,8 +75,9 @@ from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
-from ..common.engram import Engram, EngramLayout, NgramHashState
+from ..common.engram import EngramLayout, NgramHashState
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
+from .engram import Engram, gather_engram_hashes
 
 if typing.TYPE_CHECKING:
     from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -466,6 +468,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         # layer's sliding-window KV cache. Only PP ranks owning an engram
         # layer need it.
         self.engram_hash: NgramHashState | None = None
+        self.engram_dp_shared_memory = bool(
+            vllm_config.engram_config and vllm_config.engram_config.dp_shared_memory
+        )
         self.engram_swa_prefix: str | None = None
         if self.engram_layout is not None:
             local_engram = any(
@@ -594,12 +599,21 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                     swa_metadata.slot_mapping,
                     swa_metadata.block_table,
                 )
+            elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1:
+                # DP-sharded lookups are collective, so a replica skipping the
+                # hash still has to reach them.
+                engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids)
+            if engram_hashes is not None:
                 # Gather all Engram rows before entering the decoder layers.
+                # One gather feeds every layer sharing the DP-split table.
+                gathered_hashes = gather_engram_hashes(
+                    engram_hashes, dp_shared_memory=self.engram_dp_shared_memory
+                )
                 for layer in islice(self.layers, self.start_layer, self.end_layer):
                     engram = getattr(layer, "engram", None)
                     if engram is not None:
                         engram.prepare_embeddings(
-                            engram_hashes[:, engram.layer_hash_index]
+                            gathered_hashes[:, engram.layer_hash_index]
                         )
 
         full_num_tokens = positions.shape[0]