[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>

tests/distributed/test_engram_dp_shard.py

@@ -685,6 +685,13 @@ def test_engram_dp_shared_memory_runtime_requirements(
         )
 
 
+def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch):
+    """A /dev/shm smaller than the tables must fall back instead of failing startup."""
+    monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: 2)
+    layout = SimpleNamespace(num_embeddings=(1 << 50,), head_dim=DIM)
+    assert not engram_ops.can_share_engram_tables(layout)
+
+
 @pytest.mark.parametrize(
     "node_ids,dp_size,replica_size,expected",
     [

tests/engine/test_arg_utils.py

@@ -114,7 +114,7 @@ def test_engram_config_cli(options):
     "options,provided,dp_shared_memory",
     [
         ([], False, False),
-        (["--engram-config", "{}"], True, False),
+        (["--engram-config", "{}"], True, None),
         (
             ["--engram-config", '{"dp_shared_memory": true}'],
             True,

tests/test_config.py

@@ -1284,15 +1284,27 @@ def test_engram_dp_shared_memory_requires_cpu_offload():
         EngramConfig(cpu_offload=False, dp_shared_memory=True)
 
 
+@pytest.mark.parametrize(
+    "cpu_offload,dp_size,elastic_ep,expected",
+    [(True, 2, False, True), (False, 2, False, False), (True, 1, False, False)],
+)
+def test_engram_dp_shared_memory_defaults_when_supported(
+    cpu_offload, dp_size, elastic_ep, expected
+):
+    """Unset dp_shared_memory enables sharing only for offloaded, non-elastic DP."""
+    parallel = ParallelConfig(data_parallel_size=dp_size)
+    parallel.enable_elastic_ep = elastic_ep
+    config = EngramConfig(cpu_offload=cpu_offload)
+    config.resolve_dp_shared_memory(parallel)
+    assert config.dp_shared_memory is expected
+
+
 @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(

vllm/config/engram.py

@@ -10,7 +10,6 @@ import vllm.envs as envs
 from vllm.config.utils import config, get_hash_factors, hash_factors
 
 if TYPE_CHECKING:
-    from vllm.config.load import LoadConfig
     from vllm.config.model import ModelConfig
     from vllm.config.parallel import ParallelConfig
 
@@ -50,11 +49,13 @@ class EngramConfig:
     """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
+    dp_shared_memory: bool | None = None
     """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."""
+    /dev/shm capacity and a shared IPC namespace. Defaults to enabled whenever
+    the other settings allow it, falling back to per-replica tables when DP
+    replicas are not co-located on one node or /dev/shm cannot hold them."""
 
     @model_validator(mode="after")
     def _validate_shared_memory(self) -> Self:
@@ -82,6 +83,15 @@ class EngramConfig:
                 "embeddings, non-empty n-gram layer ids, and CUDA."
             )
 
+    def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
+        """Share host tables by default wherever the configuration permits."""
+        if self.dp_shared_memory is None:
+            self.dp_shared_memory = (
+                self.cpu_offload
+                and parallel_config.data_parallel_size > 1
+                and not parallel_config.enable_elastic_ep
+            )
+
     def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
         """Reject unsupported embedding parallel topologies."""
         if self.dp_shared_memory:
@@ -98,18 +108,6 @@ 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

@@ -1266,8 +1266,8 @@ class VllmConfig:
                 return
             self.engram_config = EngramConfig()
         self.engram_config.verify_model_config(model_config)
+        self.engram_config.resolve_dp_shared_memory(self.parallel_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/models/deepseek_v41/nvidia/engram.py

@@ -3,6 +3,7 @@
 """NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
 
 import mmap
+import shutil
 import tempfile
 import weakref
 from contextlib import ExitStack
@@ -89,6 +90,27 @@ def gather_engram_hashes(
     return dp_group.all_gather(hash_ids, dim=0)
 
 
+def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
+    """Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
+    if get_engram_dp_size() == 1:
+        logger.warning_once(
+            "Engram DP replicas are not co-located on one node; "
+            "storing the offloaded tables per rank instead of sharing them."
+        )
+        return False
+    num_bytes = sum(layout.num_embeddings) * (
+        layout.head_dim + layout.head_dim // block_size
+    )
+    if shutil.disk_usage("/dev/shm").total < num_bytes:
+        logger.warning_once(
+            "Sharing Engram tables across DP replicas needs %.1f GiB of /dev/shm "
+            "(--shm-size or --ipc=host); sharding them across replicas instead.",
+            num_bytes / 1024**3,
+        )
+        return False
+    return True
+
+
 class DPSharedEngramStorage:
     """Registered host weights shared by a node-local DP group with one writer."""
 
@@ -362,7 +384,7 @@ class Engram(BaseEngram):
             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,
+            dp_shared_memory=bool(engram_config.dp_shared_memory),
         )
 
     def _init_staging(self, max_tokens: int, head_dim: int) -> None:

vllm/models/deepseek_v41/nvidia/model.py

@@ -88,7 +88,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 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
+from .engram import Engram, can_share_engram_tables, gather_engram_hashes
 from .ops.mhc import (
     MHC_OVERLAP_MAX_TOKENS,
     mhc_pre_delayed_overlap,
@@ -589,6 +589,13 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             self.embed_tokens = PPMissingLayer()
 
         self.engram_layout = EngramLayout.from_config(config)
+        engram_config = vllm_config.engram_config
+        if (
+            self.engram_layout is not None
+            and engram_config is not None
+            and engram_config.dp_shared_memory
+        ):
+            engram_config.dp_shared_memory = can_share_engram_tables(self.engram_layout)
 
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,