diff --git a/deepinfer/neuron_ops/nki_topk.py b/deepinfer/neuron_ops/nki_topk.py new file mode 100644 index 0000000..89a95b8 --- /dev/null +++ b/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) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index 6d9080e..c3d15c9 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/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] diff --git a/deepinfer/neuron_ops/vendor/deep_select/interface.py b/deepinfer/neuron_ops/vendor/deep_select/interface.py index 48b4fbe..47a0ba7 100644 --- a/deepinfer/neuron_ops/vendor/deep_select/interface.py +++ b/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]