mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
42e05844d1
- 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>
46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
"""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)
|