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>
43 lines
1.8 KiB
C++
43 lines
1.8 KiB
C++
#include "cuda_utils.hpp"
|
|
|
|
#include <cstdint>
|
|
#include <optional>
|
|
#include <torch/extension.h>
|
|
|
|
#include "utils/log.hpp"
|
|
|
|
// Neuron (native PrivateUse1 device). Device memory is not addressable through raw
|
|
// pointers from the host, so allocation goes through the torch allocator and slot
|
|
// copies are done on tensor views (see clone_and_zero_states_and_blocks).
|
|
namespace deepinfer::cuda::utils {
|
|
|
|
std::optional<torch::Tensor> alloc_tensor(ssize_t size) {
|
|
try {
|
|
auto t = torch::empty({(int64_t)size}, at::TensorOptions().device(get_default_device()).dtype(torch::kByte));
|
|
LOG_INFO("alloc_tensor: allocated neuron memory of size {}", size);
|
|
return {t};
|
|
} catch (const c10::Error &e) {
|
|
LOG_ERROR("alloc_tensor: failed to allocate neuron memory of size {}: {}", size, e.what());
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<torch::Tensor> import_ipc_tensor(py::bytes data) { TORCH_CHECK(false, "IPC tensors are not supported on Neuron"); }
|
|
|
|
py::bytes export_ipc_tensor(size_t ptr, ssize_t nbytes) { TORCH_CHECK(false, "IPC tensors are not supported on Neuron"); }
|
|
|
|
void cuda_copy_async(void *dst, void *src, size_t size, uint64_t stream) {
|
|
TORCH_CHECK(false, "raw device-pointer copies are not supported on Neuron; copy tensor views instead");
|
|
}
|
|
|
|
uint64_t get_current_stream() { return 0; }
|
|
|
|
void pymodule_init(pybind11::module &m) {
|
|
m.def("alloc_tensor", &deepinfer::cuda::utils::alloc_tensor, "Allocate a 1D byte tensor on the Neuron device.", py::arg("size"))
|
|
.def("export_ipc_tensor", &deepinfer::cuda::utils::export_ipc_tensor, "Not supported on Neuron.", py::arg("ptr"),
|
|
py::arg("nbytes"))
|
|
.def("import_ipc_tensor", &deepinfer::cuda::utils::import_ipc_tensor, "Not supported on Neuron.", py::arg("data"));
|
|
}
|
|
|
|
} // namespace deepinfer::cuda::utils
|