[Bugfix][Engram] Fall back when /dev/shm is absent before sharing tables (#57914)
Signed-off-by: Juntian Liu <juntianl@inferact.ai> Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Juntian Liu committed
Sep 22, 2026 at 09:19 UTC
a6131b0e656e18eeb3797b8856aa3dbc69efd8e5
1 file changed
+50
-58
vllm/models/deepseek_v41/nvidia/engram.py
+50
-58
@@ -3,6 +3,7 @@
3
"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
4
5
import mmap
6
+import os
7
import shutil
8
import tempfile
9
import weakref
@@ -20,6 +21,10 @@ from vllm.distributed import (
21
get_tensor_model_parallel_rank,
22
tensor_model_parallel_all_gather,
23
)
24
+from vllm.distributed.device_communicators.shm_broadcast import (
25
+ SHM_PATH,
26
+ check_shm_free_space,
27
+)
28
from vllm.distributed.parallel_state import GroupCoordinator
29
from vllm.forward_context import get_forward_context
30
from vllm.logger import init_logger
@@ -92,6 +97,13 @@ def gather_engram_hashes(
97
98
def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
99
"""Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
100
+ if not os.path.isdir(SHM_PATH):
101
+ logger.warning_once(
102
+ "%s is not mounted; storing the offloaded Engram tables per rank "
103
+ "instead of sharing them across DP replicas.",
104
+ SHM_PATH,
105
+ )
106
+ return False
107
if get_engram_dp_size() == 1:
108
logger.warning_once(
109
"Engram DP replicas are not co-located on one node; "
@@ -101,11 +113,12 @@ def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
113
num_bytes = sum(layout.num_embeddings) * (
114
layout.head_dim + layout.head_dim // block_size
115
)
104
- if shutil.disk_usage("/dev/shm").total < num_bytes:
116
+ if shutil.disk_usage(SHM_PATH).total < num_bytes:
117
logger.warning_once(
106
- "Sharing Engram tables across DP replicas needs %.1f GiB of /dev/shm "
118
+ "Sharing Engram tables across DP replicas needs %.1f GiB of %s "
119
"(--shm-size or --ipc=host); sharding them across replicas instead.",
120
num_bytes / 1024**3,
121
+ SHM_PATH,
122
)
123
return False
124
return True
@@ -126,20 +139,14 @@ class DPSharedEngramStorage:
139
140
def _allocate(self, num_bytes: int) -> torch.Tensor:
141
"""Map and register one physical allocation across a node-local DP group."""
129
- from vllm.distributed.device_communicators.shm_broadcast import (
130
- check_shm_free_space,
131
- )
132
-
142
group = self.group
143
with ExitStack() as stack:
135
- path, error = None, None
144
+ path = error = None
145
if group.rank_in_group == 0:
146
try:
147
check_shm_free_space(num_bytes)
148
backing_file = stack.enter_context(
140
- tempfile.NamedTemporaryFile(
141
- prefix="vllm_engram_", dir="/dev/shm"
142
- )
149
+ tempfile.NamedTemporaryFile(prefix="vllm_engram_", dir=SHM_PATH)
150
)
151
backing_file.truncate(num_bytes)
152
path = backing_file.name
@@ -151,57 +158,42 @@ class DPSharedEngramStorage:
158
"Engram shared-memory creation failed on EDP rank 0: " + error
159
)
160
154
- mapping = owner = tensor = finalizer = None
155
- stage = "open"
161
+ mapping = owner = tensor = None
162
try:
157
- try:
158
- with open(path, "r+b") as file:
159
- stage = "mmap"
160
- mapping = mmap.mmap(
161
- file.fileno(), num_bytes, flags=mmap.MAP_SHARED
162
- )
163
- stage = "cudaHostRegister"
164
- owner = np.frombuffer(mapping, dtype=np.uint8)
165
- pointer = owner.ctypes.data
166
- tensor = torch.from_numpy(owner)
167
- result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
168
- if result.value != 0:
169
- raise RuntimeError(f"cudaHostRegister failed: {result}")
170
- finalizer = weakref.finalize(
171
- owner, self._unregister, mapping, pointer
172
- )
173
- finalizer.atexit = False # type: ignore[misc]
174
- # The UVA helper otherwise allocates a private pinned copy.
175
- if not tensor.is_pinned():
176
- raise RuntimeError(
177
- "CUDA did not recognize the shared Engram registration"
178
- )
179
- except Exception as exc:
180
- error = f"{stage}: {type(exc).__name__}: {exc}"
181
-
182
- errors: list[str | None] = [None] * group.world_size
183
- # Also fences peer mappings before the leader unlinks the file.
184
- torch.distributed.all_gather_object(
185
- errors, error, group=group.cpu_group
186
- )
187
- failures = "; ".join(
188
- f"EDP rank {rank}: {error}"
189
- for rank, error in enumerate(errors)
190
- if error is not None
191
- )
192
- if failures:
163
+ with open(path, "r+b") as file:
164
+ mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED)
165
+ owner = np.frombuffer(mapping, dtype=np.uint8)
166
+ pointer = owner.ctypes.data
167
+ tensor = torch.from_numpy(owner)
168
+ result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
169
+ if result.value != 0:
170
+ raise RuntimeError(f"cudaHostRegister failed: {result}")
171
+ finalizer = weakref.finalize(owner, self._unregister, mapping, pointer)
172
+ finalizer.atexit = False # type: ignore[misc]
173
+ # The UVA helper otherwise allocates a private pinned copy.
174
+ if not tensor.is_pinned():
175
raise RuntimeError(
194
- "Engram shared-memory initialization failed: " + failures
176
+ "cudaHostRegister did not pin the shared Engram mapping"
177
)
196
- assert tensor is not None
197
- return tensor
198
- except Exception:
199
- if finalizer is not None:
200
- finalizer()
201
- tensor = owner = None
202
- if mapping is not None:
203
- mapping.close()
204
- raise
178
+ except Exception as exc:
179
+ error = f"{type(exc).__name__}: {exc}"
180
+
181
+ # Also fences peer mappings before the leader unlinks the file.
182
+ errors: list[str | None] = [None] * group.world_size
183
+ torch.distributed.all_gather_object(errors, error, group=group.cpu_group)
184
+ failures = "; ".join(
185
+ f"EDP rank {rank}: {error}"
186
+ for rank, error in enumerate(errors)
187
+ if error is not None
188
+ )
189
+ if failures:
190
+ # Dropping the owner runs the finalizer; the mapping then unmaps.
191
+ del tensor, owner, mapping
192
+ raise RuntimeError(
193
+ "Engram shared-memory initialization failed: " + failures
194
+ )
195
+ assert tensor is not None
196
+ return tensor
197
198
@staticmethod
199
def _unregister(mapping: mmap.mmap, pointer: int) -> None: