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 <noreply@anthropic.com>

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)

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