[Model][Engram] Share host tables across co-located DP replicas by default (#57651)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Juntian Liu committed
Sep 21, 2026 at 01:26 UTC
76ffb0d37488988ec158cab609d5569e9f9ca34c
7 files changed
+68
-22
tests/distributed/test_engram_dp_shard.py
+7
@@ -685,6 +685,13 @@ def test_engram_dp_shared_memory_runtime_requirements(
685
)
686
687
688
+def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch):
689
+ """A /dev/shm smaller than the tables must fall back instead of failing startup."""
690
+ monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: 2)
691
+ layout = SimpleNamespace(num_embeddings=(1 << 50,), head_dim=DIM)
692
+ assert not engram_ops.can_share_engram_tables(layout)
693
+
694
+
695
@pytest.mark.parametrize(
696
"node_ids,dp_size,replica_size,expected",
697
[
tests/engine/test_arg_utils.py
+1
-1
@@ -114,7 +114,7 @@ def test_engram_config_cli(options):
114
"options,provided,dp_shared_memory",
115
[
116
([], False, False),
117
- (["--engram-config", "{}"], True, False),
117
+ (["--engram-config", "{}"], True, None),
118
(
119
["--engram-config", '{"dp_shared_memory": true}'],
120
True,
tests/test_config.py
+15
-3
@@ -1284,15 +1284,27 @@ def test_engram_dp_shared_memory_requires_cpu_offload():
1284
EngramConfig(cpu_offload=False, dp_shared_memory=True)
1285
1286
1287
+@pytest.mark.parametrize(
1288
+ "cpu_offload,dp_size,elastic_ep,expected",
1289
+ [(True, 2, False, True), (False, 2, False, False), (True, 1, False, False)],
1290
+)
1291
+def test_engram_dp_shared_memory_defaults_when_supported(
1292
+ cpu_offload, dp_size, elastic_ep, expected
1293
+):
1294
+ """Unset dp_shared_memory enables sharing only for offloaded, non-elastic DP."""
1295
+ parallel = ParallelConfig(data_parallel_size=dp_size)
1296
+ parallel.enable_elastic_ep = elastic_ep
1297
+ config = EngramConfig(cpu_offload=cpu_offload)
1298
+ config.resolve_dp_shared_memory(parallel)
1299
+ assert config.dp_shared_memory is expected
1300
+
1301
+
1302
@pytest.mark.parametrize(
1303
"dp_size,load_format,multithread,error",
1304
[
1305
(1, "auto", False, "requires data_parallel_size > 1"),
1291
- (2, "dummy", False, "requires load_format"),
1292
- (2, "sharded_state", False, "requires load_format"),
1306
(2, "auto", False, None),
1307
(2, "safetensors", True, None),
1295
- (2, "pt", True, None),
1308
],
1309
)
1310
def test_engram_dp_shared_memory_config_validation(
vllm/config/engram.py
+13
-15
@@ -10,7 +10,6 @@ import vllm.envs as envs
10
from vllm.config.utils import config, get_hash_factors, hash_factors
11
12
if TYPE_CHECKING:
13
- from vllm.config.load import LoadConfig
13
from vllm.config.model import ModelConfig
14
from vllm.config.parallel import ParallelConfig
15
@@ -50,11 +49,13 @@ class EngramConfig:
49
"""Shard embeddings across TP and all DP ranks when enabled.
50
Otherwise, each DP rank has a separate TP-sharded embedding replica."""
51
53
- dp_shared_memory: bool = False
52
+ dp_shared_memory: bool | None = None
53
"""Share CPU-offloaded embedding weights between co-located
54
DP replicas. Each node stores one copy of every TP shard, reducing host
55
memory without per-step Engram DP collectives. Requires sufficient
57
- /dev/shm capacity and a shared IPC namespace."""
56
+ /dev/shm capacity and a shared IPC namespace. Defaults to enabled whenever
57
+ the other settings allow it, falling back to per-replica tables when DP
58
+ replicas are not co-located on one node or /dev/shm cannot hold them."""
59
60
@model_validator(mode="after")
61
def _validate_shared_memory(self) -> Self:
@@ -82,6 +83,15 @@ class EngramConfig:
83
"embeddings, non-empty n-gram layer ids, and CUDA."
84
)
85
86
+ def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
87
+ """Share host tables by default wherever the configuration permits."""
88
+ if self.dp_shared_memory is None:
89
+ self.dp_shared_memory = (
90
+ self.cpu_offload
91
+ and parallel_config.data_parallel_size > 1
92
+ and not parallel_config.enable_elastic_ep
93
+ )
94
+
95
def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
96
"""Reject unsupported embedding parallel topologies."""
97
if self.dp_shared_memory:
@@ -98,18 +108,6 @@ class EngramConfig:
108
"Engram embedding_across_dp is not supported with elastic EP yet."
109
)
110
101
- def verify_load_config(self, load_config: "LoadConfig") -> None:
102
- """Shared tables require a loader that invokes parameter weight callbacks."""
103
- if self.dp_shared_memory and load_config.load_format not in (
104
- "auto",
105
- "safetensors",
106
- "pt",
107
- ):
108
- raise ValueError(
109
- "dp_shared_memory requires load_format 'auto', "
110
- f"'safetensors' or 'pt'; got {load_config.load_format!r}."
111
- )
112
-
111
def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
112
"""Derive the embedding group size from the parallel configuration."""
113
size = parallel_config.tensor_parallel_size
vllm/config/vllm.py
+1
-1
@@ -1266,8 +1266,8 @@ class VllmConfig:
1266
return
1267
self.engram_config = EngramConfig()
1268
self.engram_config.verify_model_config(model_config)
1269
+ self.engram_config.resolve_dp_shared_memory(self.parallel_config)
1270
self.engram_config.verify_parallel_config(self.parallel_config)
1270
- self.engram_config.verify_load_config(self.load_config)
1271
logger.info_once("Resolved Engram configuration: %s", str(self.engram_config))
1272
1273
def __post_init__(self):
vllm/models/deepseek_v41/nvidia/engram.py
+23
-1
@@ -3,6 +3,7 @@
3
"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
4
5
import mmap
6
+import shutil
7
import tempfile
8
import weakref
9
from contextlib import ExitStack
@@ -89,6 +90,27 @@ def gather_engram_hashes(
90
return dp_group.all_gather(hash_ids, dim=0)
91
92
93
+def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
94
+ """Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
95
+ if get_engram_dp_size() == 1:
96
+ logger.warning_once(
97
+ "Engram DP replicas are not co-located on one node; "
98
+ "storing the offloaded tables per rank instead of sharing them."
99
+ )
100
+ return False
101
+ num_bytes = sum(layout.num_embeddings) * (
102
+ layout.head_dim + layout.head_dim // block_size
103
+ )
104
+ if shutil.disk_usage("/dev/shm").total < num_bytes:
105
+ logger.warning_once(
106
+ "Sharing Engram tables across DP replicas needs %.1f GiB of /dev/shm "
107
+ "(--shm-size or --ipc=host); sharding them across replicas instead.",
108
+ num_bytes / 1024**3,
109
+ )
110
+ return False
111
+ return True
112
+
113
+
114
class DPSharedEngramStorage:
115
"""Registered host weights shared by a node-local DP group with one writer."""
116
@@ -362,7 +384,7 @@ class Engram(BaseEngram):
384
layout.head_dim,
385
tuple(size for order in layout.primes[layer_hash_index] for size in order),
386
cpu_offload=engram_config.cpu_offload,
365
- dp_shared_memory=engram_config.dp_shared_memory,
387
+ dp_shared_memory=bool(engram_config.dp_shared_memory),
388
)
389
390
def _init_staging(self, max_tokens: int, head_dim: int) -> None:
vllm/models/deepseek_v41/nvidia/model.py
+8
-1
@@ -88,7 +88,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id
88
89
from ..common.engram import EngramLayout, NgramHashState
90
from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
91
-from .engram import Engram, gather_engram_hashes
91
+from .engram import Engram, can_share_engram_tables, gather_engram_hashes
92
from .ops.mhc import (
93
MHC_OVERLAP_MAX_TOKENS,
94
mhc_pre_delayed_overlap,
@@ -589,6 +589,13 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
589
self.embed_tokens = PPMissingLayer()
590
591
self.engram_layout = EngramLayout.from_config(config)
592
+ engram_config = vllm_config.engram_config
593
+ if (
594
+ self.engram_layout is not None
595
+ and engram_config is not None
596
+ and engram_config.dp_shared_memory
597
+ ):
598
+ engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
599
600
self.start_layer, self.end_layer, self.layers = make_layers(
601
config.num_hidden_layers,