mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
ed49792a0c
Run DeepInfer on the native torch-neuronx device (PrivateUse1 `neuron`, torch-neuronx 2.14, no torch-xla), following the Ascend port's structure. - deep_cluster: Neuron discovery (one DeepInfer device = one logical NeuronCore, NEURON_RT_VISIBLE_CORES, NUMA from neuron-ls). - PlatformNeuron: takes the Ascend (is_npu) kernel layouts, adds is_neuron; DEEPINFER_NEURON_DEVICE=cpu runs the same code on the host for correctness work. - neuron_ops/transfer_to_neuron: map "cuda" devices, torch.device and torch.cuda.* onto neuron:0 (counterpart of torch_npu transfer_to_npu). - C++ extension: -DWITH_NEURON host-only build; torch-allocated buffers and tensor-view slot copies instead of raw device-pointer copies. - Pure-torch drop-in vendor packages (Ascend API semantics): deep_gemm, flash_mla, deep_select, tile_kernels, tilelang (constants only). - tilelang_ops/neuron backend: KV cache / metadata / norm / compress kernels; CP, ZOS and VL entry points raise NotImplementedError. - Basic mode defaults on Neuron: eager, no DSpark/MTP drafting, no zero-overhead scheduling; MoE avoids the GPU-only mega kernel; mHC prenorm GEMM uses the Ascend single-split reduction. Not yet included: deep_ep, the sampling kernels, end-to-end harness. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
"""AWS Neuron (Trainium) discovery using device files and ``neuron-ls``.
|
|
|
|
One DeepInfer device is one logical NeuronCore; visibility is controlled by
|
|
``NEURON_RT_VISIBLE_CORES``. A trn2.48xlarge exposes 16 devices x 4 logical
|
|
cores (LNC=2) = 64 devices.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import json
|
|
import re
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
_DEVICE_PATTERN = re.compile(r"^neuron(\d+)$")
|
|
|
|
|
|
def _path(root: Path, absolute: str) -> Path:
|
|
return root / absolute.lstrip("/")
|
|
|
|
|
|
def _device_node_indices(*, _root: Path = Path("/")) -> tuple[int, ...]:
|
|
device_dir = _path(_root, "/dev")
|
|
if not device_dir.is_dir():
|
|
return ()
|
|
try:
|
|
return tuple(sorted(int(m.group(1)) for p in device_dir.iterdir() if (m := _DEVICE_PATTERN.match(p.name)) is not None))
|
|
except OSError as error:
|
|
raise OSError(f"cannot enumerate {device_dir}: {error}") from error
|
|
|
|
|
|
def is_available(*, _root: Path = Path("/")) -> bool:
|
|
return bool(_device_node_indices(_root=_root))
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _neuron_ls() -> tuple[dict, ...]:
|
|
try:
|
|
output = subprocess.check_output(["neuron-ls", "-j"], text=True, stderr=subprocess.DEVNULL)
|
|
except (FileNotFoundError, subprocess.CalledProcessError):
|
|
return ()
|
|
return tuple(json.loads(output))
|
|
|
|
|
|
def discover_cores(*, _root: Path = Path("/")) -> list[tuple[int, int, int]]:
|
|
"""Logical NeuronCores as ``(core_id, device_id, numa_node)``; NUMA is -1 when unknown."""
|
|
cores: list[tuple[int, int, int]] = []
|
|
if _root == Path("/"):
|
|
for device in _neuron_ls():
|
|
numa = int(device.get("numa_node", -1))
|
|
cores.extend((int(core), int(device["neuron_device"]), numa) for core in device["neuroncore_ids"])
|
|
if cores:
|
|
return sorted(cores)
|
|
# Fallback without neuron-ls: sysfs core_count is physical cores; trn2 defaults to LNC=2.
|
|
lnc = 2
|
|
for index in _device_node_indices(_root=_root):
|
|
count_path = _path(_root, f"/sys/class/neuron_device/neuron{index}/core_count")
|
|
physical = int(count_path.read_text().strip()) if count_path.is_file() else 8
|
|
base = len(cores)
|
|
cores.extend((base + i, index, -1) for i in range(physical // lnc))
|
|
return cores
|
|
|
|
|
|
def physical_device_count(*, _root: Path = Path("/")) -> int:
|
|
return len(discover_cores(_root=_root))
|
|
|
|
|
|
def device_numa_node(index: int, *, _root: Path = Path("/")) -> int:
|
|
cores = discover_cores(_root=_root)
|
|
if index >= len(cores):
|
|
raise IndexError(f"no neuron core at index {index}")
|
|
numa = cores[index][2]
|
|
if numa < 0:
|
|
raise RuntimeError(f"NUMA metadata is unavailable for neuron core {index}")
|
|
return numa
|
|
|
|
|
|
__all__ = ["device_numa_node", "discover_cores", "is_available", "physical_device_count"]
|