diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index dd7105a23b..ba03387aac 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/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", [ diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 68fe250677..bded621fad 100644 --- a/tests/engine/test_arg_utils.py +++ b/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, diff --git a/tests/test_config.py b/tests/test_config.py index f31adfdb8d..cad5ffd627 100644 --- a/tests/test_config.py +++ b/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( diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 773ca20911..c9f05c9382 100644 --- a/vllm/config/engram.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6974d1fbb3..1ff675a8f0 100644 --- a/vllm/config/vllm.py +++ b/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): diff --git a/vllm/models/deepseek_v41/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py index 624a36c214..7dd820fee4 100644 --- a/vllm/models/deepseek_v41/nvidia/engram.py +++ b/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: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b91ad1d7ba..d525aff3c7 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/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,