mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 15:27:26 +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>
28 lines
965 B
Python
28 lines
965 B
Python
"""Mirrors ``tilelang_ops/asc/gather.py`` (``gather_tl``): bytewise row gather with zero-fill."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
|
|
|
|
def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
|
|
"""Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.
|
|
|
|
``y[i] = x[index[i]]`` if ``0 <= index[i] < n`` else 0; written in place.
|
|
"""
|
|
|
|
def _gather(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
|
|
n = x.shape[0]
|
|
m = index.shape[0]
|
|
if m == 0:
|
|
return
|
|
if n == 0:
|
|
y.zero_()
|
|
return
|
|
valid = (index >= 0) & (index < n)
|
|
safe = torch.where(valid, index, torch.zeros_like(index))
|
|
rows = torch.index_select(x[:, :D], 0, safe)
|
|
y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))
|
|
|
|
return _gather
|