From d2389f8a9d2f41c16df8f7392519cf826b8b009f Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 20 Oct 2026 21:35:48 +0800 Subject: [PATCH] Neuron deep_ep: on-device fixes from the first NeuronCore run - Engram: gather rows where the table lives (host by default, read via UVA on GPU) and move them to the compute device. - EPHandle: the topk_idx version guard tolerates inference-mode tensors on the neuron device that do not report is_inference(). Co-Authored-By: Claude Opus 5.5 --- .../neuron_ops/vendor/deep_ep/buffers/engram.py | 14 ++++++++++---- deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py | 13 +++++++++++-- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index a8a3fe5..8ddc2fa 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -47,6 +47,12 @@ def _accept_legacy_storage_size(init): return wrapper + +def _take_rows_to(table: torch.Tensor, index: torch.Tensor, device: torch.device) -> torch.Tensor: + """Gather rows where the table lives (Engram tables default to host memory, which the GPU reads + through UVA), then move the result to the compute device.""" + return take_rows(table, index.to(table.device)).to(device) + class EngramBuffer(BufferBase): """Distributed n-gram embedding lookup over a sharded table.""" @@ -193,8 +199,8 @@ class EngramBuffer(BufferBase): local = flat - owner * num_entries if self.num_ranks == 1: - rows = take_rows(storage, local) - sf_rows = take_rows(sf_shard, local) if sf_shard is not None else None + rows = _take_rows_to(storage, local, device) + sf_rows = _take_rows_to(sf_shard, local, device) if sf_shard is not None else None else: # Group the requests by owner rank (stable: ascending request position). order = torch.argsort(owner, stable=True) @@ -214,7 +220,7 @@ class EngramBuffer(BufferBase): comm.all_to_all_single(served_idx, send_idx, recv_counts, send_counts, self.group) served_idx = served_idx.clamp(0, max(num_entries - 1, 0)) - served_rows = take_rows(storage, served_idx).contiguous() + served_rows = _take_rows_to(storage, served_idx, device).contiguous() answer_rows = empty_rows(num_requests, self.hidden, storage.dtype, device) # Split sizes count rows: `all_to_all_single` splits along dim 0. comm.all_to_all_single(byte_rows(answer_rows), byte_rows(served_rows), @@ -224,7 +230,7 @@ class EngramBuffer(BufferBase): sf_rows = None if sf_shard is not None: - served_sf = take_rows(sf_shard, served_idx).contiguous() + served_sf = _take_rows_to(sf_shard, served_idx, device).contiguous() answer_sf = empty_rows(num_requests, self.num_sf_packs, sf_shard.dtype, device) comm.all_to_all_single(byte_rows(answer_sf), byte_rows(served_sf), send_counts, recv_counts, self.group) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py index 04e8626..a8862ae 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py @@ -38,6 +38,15 @@ def _value_or(value, default): return default if value is None else value + +def _version_or_none(t: torch.Tensor) -> Optional[int]: + """Version counter for the modified-after-dispatch guard; None when the tensor has none (inference + tensors, and on the neuron device some inference-mode outputs that do not report is_inference()).""" + try: + return None if t.is_inference() else t._version + except RuntimeError: + return None + class EPHandle: """Communication handle returned by `EPBuffer.dispatch`. @@ -91,11 +100,11 @@ class EPHandle: self._plan: Optional['_RoutingPlan'] = None self._topk_idx = topk_idx - self._topk_idx_version = None if topk_idx.is_inference() else topk_idx._version + self._topk_idx_version = _version_or_none(topk_idx) @property def topk_idx(self) -> torch.Tensor: - assert self._topk_idx_version is None or self._topk_idx._version == self._topk_idx_version, \ + assert self._topk_idx_version is None or _version_or_none(self._topk_idx) == self._topk_idx_version, \ '`topk_idx` must not be modified while the EP handle is in use' return self._topk_idx