mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 22:59:44 +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>
48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
"""HyperConnection prenorm GEMM; mirrors ``DeepGEMM-Ascend/deep_gemm/tilelang_ops/hc_prenorm.py``.
|
|
|
|
``D = float(A) @ B.T`` and ``sqr_sum = sum(float(A) ** 2, dim=-1)``, both in fp32
|
|
(reference check: ``DeepGEMM-Ascend/tests/test_hyperconnection.py``).
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
def tf32_hc_prenorm_gemm(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
d: torch.Tensor,
|
|
sqr_sum: torch.Tensor,
|
|
num_splits: int | None = None,
|
|
) -> None:
|
|
"""Fused hyperconnection prenorm forward, written in place into ``d`` / ``sqr_sum``.
|
|
|
|
Args:
|
|
a: [M, K] BF16 input.
|
|
b: [N, K] FP32 weight.
|
|
d: [M, N] FP32 output, or [num_splits, M, N] when ``num_splits`` is given.
|
|
sqr_sum: [M] FP32 output, or [num_splits, M] when ``num_splits`` is given.
|
|
num_splits: K split count. Like the Ascend backend, the split is reduced
|
|
internally, so only ``None`` and ``1`` are accepted.
|
|
"""
|
|
assert a.dim() == 2 and b.dim() == 2, f"{a.shape=}, {b.shape=}"
|
|
assert a.dtype == torch.bfloat16, f"{a.dtype=}"
|
|
assert b.dtype == torch.float32, f"{b.dtype=}"
|
|
m, k = a.shape
|
|
n = b.shape[0]
|
|
assert b.shape[1] == k, f"{a.shape=}, {b.shape=}"
|
|
|
|
if num_splits is None:
|
|
out_mul, out_sqr_sum = d, sqr_sum
|
|
else:
|
|
assert num_splits == 1, f"This backend reduces the K split internally, got {num_splits=}"
|
|
out_mul, out_sqr_sum = d[0], sqr_sum[0]
|
|
|
|
assert out_mul.shape == (m, n) and out_mul.dtype == torch.float32, f"{out_mul.shape=}, {out_mul.dtype=}"
|
|
assert out_sqr_sum.shape == (m,) and out_sqr_sum.dtype == torch.float32, f"{out_sqr_sum.shape=}"
|
|
if m == 0:
|
|
return
|
|
|
|
a_f32 = a.float()
|
|
out_mul.copy_(a_f32 @ b.transpose(0, 1))
|
|
out_sqr_sum.copy_(a_f32.square().sum(dim=-1))
|