diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index c3d15c9..1555921 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -208,12 +208,18 @@ class EngramBuffer(BufferBase): def _shared(self, name: str, rows: int, row_bytes: int) -> torch.Tensor: """A ``[rows, row_bytes]`` uint8 tensor mapping one /dev/shm file shared by the group's ranks.""" - # Ranks of one job share their launcher: its pid names the job. - path = f"/dev/shm/deepinfer_engram_{os.getppid()}_{name}" + # One name for the whole group: rank 0 publishes its pid through the distributed store. (The parent pid is not + # shared: CP followers are children of their engine's driver, not of the launcher.) + store = dist.distributed_c10d._get_default_store() + key = f"deepinfer_engram_shm_{name}" if self.rank_idx == 0: + path = f"/dev/shm/deepinfer_engram_{os.getpid()}_{name}" with open(path, "wb") as f: f.truncate(rows * row_bytes) + store.set(key, path) comm.barrier(self) + path = store.get(key).decode() + assert os.path.getsize(path) == rows * row_bytes, path # opening must never create a second table t = torch.from_file(path, shared=True, size=rows * row_bytes, dtype=torch.uint8).view(rows, row_bytes) comm.barrier(self) if self.rank_idx == 0: