Neuron: nkilib rotational_topk for the sampler; single-host shared engram tables (#27)

- deep_select.topk on a NeuronCore with k >= 1024 (the sampler's force_topk
  4096 over the 129280 vocab) uses nkilib's rotational_topk: torch.topk with
  that k does not compile on the device (NCC_IXCG857). 8.4 ms on one core,
  values identical to torch.topk.
- EngramBuffer on Neuron: the ranks of a trn2 instance share one host, so the
  table shards are written into /dev/shm-backed tables (unlinked once mapped)
  and every rank gathers its own rows: no all-to-all and no host sync on split
  sizes per fetch. DEEPINFER_NEURON_ENGRAM_SHARED=0 restores the sharded path.
  Host EP=2 tiny: logits bitwise equal to the sharded path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/neuron_ops/nki_topk.py

@@ -0,0 +1,45 @@
+"""Top-k on a NeuronCore through nkilib's ``rotational_topk`` (issue #27).
+
+``torch.topk`` with a large ``k`` does not compile on the device (neuronx-cc NCC_IXCG857 for k=4096 over the
+129280-token vocab), so the eager sampler fell back to slow paths. nkilib's rotational top-k does k=4096 over
+the vocab in ~8.4 ms with values identical to ``torch.topk`` (equal values may come out in another order).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+
+#: nkilib's tested range; larger widths go through the torch path.
+_MAX_VOCAB = 151936
+_MIN_K = 1024  # the sampler's force_topk; the indexer's k=512 compiles with torch.topk
+
+
+def supported(scores: torch.Tensor, k: int) -> bool:
+    return (
+        scores.device.type == "neuron"
+        and scores.dim() == 2
+        and scores.dtype == torch.float32
+        and _MIN_K <= k <= 4096
+        and k < scores.shape[1] <= _MAX_VOCAB
+    )
+
+
+@functools.cache
+def _kernel_and_config(rows: int, vocab: int, k: int):  # noqa: ANN202
+    import nki.language as nl
+    from nkilib.core.topk.rotational_topk import rotational_topk
+    from nkilib.core.topk.rotational_topk_utils import create_rotational_topk_config, create_topk_config
+    from torch_neuronx.nki_hop import wrap_nki
+
+    tc = create_topk_config(inp_shape=(rows, vocab), inp_dtype=nl.float32, k=k, sorted=True, num_programs=1 if rows == 1 else 2)
+    return wrap_nki(rotational_topk), create_rotational_topk_config(inp_shape=(rows, vocab), topk_config=tc)
+
+
+def topk(scores: torch.Tensor, k: int) -> torch.Tensor:
+    """Indices ``[rows, k]`` int64 of the ``k`` largest values of each row, in descending value order."""
+    rows, vocab = scores.shape
+    kernel, config = _kernel_and_config(rows, vocab, k)
+    _, idx = kernel(scores.contiguous(), config)
+    return idx.reshape(rows, -1)[:, :k].to(torch.int64)

deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py

@@ -25,6 +25,7 @@ Basic mode: the request counts are host values, so `fetch` host-syncs once per l
 from __future__ import annotations
 
 import functools
+import os
 from typing import Callable, List, Optional, Tuple
 
 import torch
@@ -168,13 +169,15 @@ class EngramBuffer(BufferBase):
         assert len(storages) == num_layers
         assert (sfs is not None) == (self.num_sf_packs > 0)
 
-        self._storages = []
         for layer_idx, storage in enumerate(storages):
             assert storage.dim() == 2 and storage.is_contiguous()
             assert storage.size(0) == self.num_entries_per_layer[layer_idx]
             assert storage.size(1) == self.hidden
             assert storage.dtype == self.dtype
-            self._storages.append(storage.clone())
+        if self._use_shared_tables(storages, sfs):
+            self._write_shared(storages, sfs)
+            return
+        self._storages = [storage.clone() for storage in storages]
 
         self._sf_shards = []
         if sfs is not None:
@@ -191,6 +194,61 @@ class EngramBuffer(BufferBase):
 
         comm.barrier(self)
 
+    # --- Single-host shared tables (Neuron) -------------------------------------------
+
+    def _use_shared_tables(self, storages: List[torch.Tensor], sfs: Optional[List[torch.Tensor]]) -> bool:
+        """NEURON: every rank of a trn2 instance shares one host, so the whole table can sit in shared host memory and
+        each rank reads its own rows: no all-to-all, no host sync on split sizes (DEEPINFER_NEURON_ENGRAM_SHARED=0 off)."""
+        return (
+            os.environ.get("DEEPINFER_NEURON_ENGRAM_SHARED", "1") == "1"
+            and self.num_ranks > 1
+            and all(t.device.type == "cpu" for t in storages)
+            and (sfs is None or all(sf.size(0) == n for sf, n in zip(sfs, self.num_entries_per_layer)))
+        )
+
+    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}"
+        if self.rank_idx == 0:
+            with open(path, "wb") as f:
+                f.truncate(rows * row_bytes)
+        comm.barrier(self)
+        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:
+            os.unlink(path)  # the mappings stay valid; the memory goes away with the last process
+        return t
+
+    def _write_shared(self, storages: List[torch.Tensor], sfs: Optional[List[torch.Tensor]]) -> None:
+        self._storages, self._sf_shards = [], []
+        for layer_idx, storage in enumerate(storages):
+            n = self.num_entries_per_layer[layer_idx]
+            begin = self.rank_idx * n
+            table = self._shared(f"t{layer_idx}", self.num_ranks * n, self.hidden * storage.element_size())
+            table[begin:begin + n] = byte_rows(storage)
+            self._storages.append(table)
+            if sfs is not None:
+                sf = sfs[layer_idx]
+                sf_table = self._shared(f"s{layer_idx}", self.num_ranks * n, sf.size(1) * sf.element_size())
+                sf_table[begin:begin + n] = byte_rows(sf.contiguous())
+                self._sf_shards.append((sf_table, sf.dtype))
+        self._shared_tables = True
+        comm.barrier(self)
+
+    def _fetch_layer_shared(self, layer_idx: int, indices: torch.Tensor) -> torch.Tensor | tuple:
+        device = indices.device
+        num_tokens, num_entries_per_token = indices.shape
+        flat = indices.reshape(-1).to(device="cpu", dtype=torch.int64)
+        flat = flat.clamp(0, self._storages[layer_idx].size(0) - 1)
+        rows = self._storages[layer_idx].index_select(0, flat).view(self.dtype)
+        fetched = rows.reshape(num_tokens, num_entries_per_token * self.hidden).to(device)
+        if not self._sf_shards:
+            return fetched
+        sf_table, sf_dtype = self._sf_shards[layer_idx]
+        sf_rows = sf_table.index_select(0, flat).view(sf_dtype)
+        return fetched, sf_rows.reshape(num_tokens, num_entries_per_token * self.num_sf_packs).to(device)
+
     # --- Fetch ----------------------------------------------------------------------
 
     def _fetch_layer(self, layer_idx: int, indices: torch.Tensor,
@@ -276,6 +334,9 @@ class EngramBuffer(BufferBase):
         assert num_entries_per_token == self.num_entries_per_token
         assert self._storages, 'write() must be called before fetch()'
 
-        results = [self._fetch_layer(layer_idx, indices[layer_idx], use_tma_aligned_col_major_sf)
-                   for layer_idx in range(num_layers)]
+        if getattr(self, "_shared_tables", False):
+            results = [self._fetch_layer_shared(layer_idx, indices[layer_idx]) for layer_idx in range(num_layers)]
+        else:
+            results = [self._fetch_layer(layer_idx, indices[layer_idx], use_tma_aligned_col_major_sf)
+                       for layer_idx in range(num_layers)]
         return [(lambda result=result: result) for result in results]

deepinfer/neuron_ops/vendor/deep_select/interface.py

@@ -102,7 +102,12 @@ def topk(
     k_eff = min(topk, vocab_size)
     # torch.topk(sorted=True) returns values descending, which is what `sorted=True` asks for; it
     # is used unconditionally so the result is deterministic for every flag combination.
-    idx = torch.topk(scores, k_eff, dim=-1, sorted=True).indices
+    from deepinfer.neuron_ops import nki_topk  # NEURON: torch.topk with a large k does not compile on the device
+
+    if nki_topk.supported(scores, k_eff):
+        idx = nki_topk.topk(scores, k_eff)
+    else:
+        idx = torch.topk(scores, k_eff, dim=-1, sorted=True).indices
     idx = idx.to(torch.int64)
 
     counts = torch.clamp(valid_len.to(torch.int64), max=topk)  # [b, 1]