From 2836e365eaff9fe8aaeda786e254252a229ad9a4 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 07:08:37 +0800 Subject: [PATCH] Neuron engram shared tables: one file for the whole group (#28) The name came from the parent pid, which CP followers do not share with their driver (they are its children): followers created a second table, never unlinked (622 GB of /dev/shm leaked over five runs, then a host OOM), and each side read zeros for the other's shards. Rank 0 now publishes the path through the distributed store, and opening asserts the size so it can never create a file. Co-Authored-By: Claude Opus 5.5 --- deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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: