mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 14:49:45 +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>
360 lines
13 KiB
Python
360 lines
13 KiB
Python
"""Dense / grouped GEMMs; mirrors ``DeepGEMM-Ascend/csrc/apis/gemm.hpp``.
|
|
|
|
Everything is a straightforward dequantize-then-matmul reference with fp32
|
|
accumulation. Outputs are written in place into the caller's ``d`` tensor, and
|
|
``c`` follows the upstream contract:
|
|
|
|
* ``c is d`` (same storage) -> accumulate: ``D = A @ B^T + D``
|
|
* separate contiguous ``[1, N]`` tensor -> bias
|
|
* ``None`` -> plain ``D = A @ B^T``
|
|
|
|
Operand majorness does not change the logical shapes, so the ``nt`` / ``nn`` /
|
|
``tn`` / ``tt`` variants only differ by which operands are logically transposed.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from ._quant import dequantize, is_packed_fp4
|
|
from .layout import transform_sf_into_required_layout
|
|
|
|
|
|
def _is_alias(a: torch.Tensor, b: torch.Tensor) -> bool:
|
|
return a.data_ptr() == b.data_ptr() and a.shape == b.shape
|
|
|
|
|
|
def _logical_k(data: torch.Tensor) -> int:
|
|
"""Logical K of a low-precision operand whose K axis is the last dim."""
|
|
return data.shape[-1] * (2 if is_packed_fp4(data.dtype) else 1)
|
|
|
|
|
|
def _write_output(out: torch.Tensor, d: torch.Tensor, c: torch.Tensor | None, alpha: float | None) -> None:
|
|
if alpha is not None:
|
|
out = out * alpha
|
|
if c is not None:
|
|
if _is_alias(c, d):
|
|
out = out + d.float()
|
|
else:
|
|
# Separate bias tensor, [1, N] (or broadcastable).
|
|
assert c.shape[-1] == d.shape[-1], f"{c.shape=}, {d.shape=}"
|
|
out = out + c.float()
|
|
d.copy_(out.to(d.dtype))
|
|
|
|
|
|
def _dequant_operand(
|
|
data: torch.Tensor,
|
|
sf: torch.Tensor,
|
|
mn: int,
|
|
k: int,
|
|
recipe: tuple[int, int, int],
|
|
is_sfa: bool,
|
|
num_groups: int | None = None,
|
|
disable_ue8m0_cast: bool = False,
|
|
) -> torch.Tensor:
|
|
sf = transform_sf_into_required_layout(
|
|
sf, mn, k, recipe, num_groups=num_groups, is_sfa=is_sfa, disable_ue8m0_cast=disable_ue8m0_cast
|
|
)
|
|
return dequantize(data, sf, gran_k=recipe[2])
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# BF16
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _bf16_gemm(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
d: torch.Tensor,
|
|
c: torch.Tensor | None,
|
|
trans_a: bool,
|
|
trans_b: bool,
|
|
alpha: float | None,
|
|
) -> None:
|
|
if trans_a:
|
|
a = a.transpose(0, 1)
|
|
if trans_b:
|
|
b = b.transpose(0, 1)
|
|
assert a.dtype == b.dtype == torch.bfloat16, f"{a.dtype=}, {b.dtype=}"
|
|
assert d.dtype in (torch.bfloat16, torch.float32), f"{d.dtype=}"
|
|
m, k = a.shape
|
|
n = b.shape[0]
|
|
assert b.shape[1] == k and d.shape == (m, n), f"{a.shape=}, {b.shape=}, {d.shape=}"
|
|
if m == 0 or n == 0:
|
|
return
|
|
if k == 0:
|
|
_write_output(torch.zeros((m, n), dtype=torch.float32, device=d.device), d, c, alpha)
|
|
return
|
|
_write_output(a.float() @ b.float().transpose(0, 1), d, c, alpha)
|
|
|
|
|
|
def bf16_gemm_nt(a, b, d, c=None, compiled_dims="nk", alpha=None, epilogue=None) -> None:
|
|
"""D = A[M, K] @ B[N, K]^T (+ C)."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_bf16_gemm(a, b, d, c, False, False, alpha)
|
|
|
|
|
|
def bf16_gemm_nn(a, b, d, c=None, compiled_dims="nk", alpha=None, epilogue=None) -> None:
|
|
"""D = A[M, K] @ B[K, N] (+ C)."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_bf16_gemm(a, b, d, c, False, True, alpha)
|
|
|
|
|
|
def bf16_gemm_tn(a, b, d, c=None, compiled_dims="mn", alpha=None, epilogue=None) -> None:
|
|
"""D = A[K, M]^T @ B[K, N] (+ C)."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_bf16_gemm(a, b, d, c, True, True, alpha)
|
|
|
|
|
|
def bf16_gemm_tt(a, b, d, c=None, compiled_dims="mn", alpha=None, epilogue=None) -> None:
|
|
"""D = A[K, M]^T @ B[N, K]^T (+ C)."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_bf16_gemm(a, b, d, c, True, False, alpha)
|
|
|
|
|
|
# The Ascend build aliases the cuBLASLt entry points onto its ACLNN bf16 GEMMs
|
|
# (``csrc/apis/aclnn.hpp``), so they are the same kernel here.
|
|
def cublaslt_gemm_nt(a, b, d, c=None) -> None:
|
|
_bf16_gemm(a, b, d, c, False, False, None)
|
|
|
|
|
|
def cublaslt_gemm_nn(a, b, d, c=None) -> None:
|
|
_bf16_gemm(a, b, d, c, False, True, None)
|
|
|
|
|
|
def cublaslt_gemm_tn(a, b, d, c=None) -> None:
|
|
_bf16_gemm(a, b, d, c, True, True, None)
|
|
|
|
|
|
def cublaslt_gemm_tt(a, b, d, c=None) -> None:
|
|
_bf16_gemm(a, b, d, c, True, False, None)
|
|
|
|
|
|
aclnn_bf16_gemm_nt = cublaslt_gemm_nt
|
|
aclnn_bf16_gemm_nn = cublaslt_gemm_nn
|
|
aclnn_bf16_gemm_tn = cublaslt_gemm_tn
|
|
aclnn_bf16_gemm_tt = cublaslt_gemm_tt
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# FP8 / FP4
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _transposed_pair(
|
|
operand: tuple[torch.Tensor, torch.Tensor], trans: bool, dim0: int, dim1: int
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
data, sf = operand
|
|
if not trans:
|
|
return data, sf
|
|
assert not is_packed_fp4(data.dtype), "packed FP4 operands cannot be logically transposed"
|
|
return data.transpose(dim0, dim1), sf.transpose(dim0, dim1)
|
|
|
|
|
|
def _fp8_fp4_gemm(
|
|
a_: tuple[torch.Tensor, torch.Tensor],
|
|
b_: tuple[torch.Tensor, torch.Tensor],
|
|
d: torch.Tensor,
|
|
c: torch.Tensor | None,
|
|
recipe,
|
|
recipe_a,
|
|
recipe_b,
|
|
disable_ue8m0_cast: bool,
|
|
alpha: float | None,
|
|
trans_a: bool,
|
|
trans_b: bool,
|
|
) -> None:
|
|
from .layout import get_recipe
|
|
|
|
a, sfa = _transposed_pair(a_, trans_a, 0, 1)
|
|
b, sfb = _transposed_pair(b_, trans_b, 0, 1)
|
|
assert d.dtype in (torch.bfloat16, torch.float32), f"{d.dtype=}"
|
|
m = a.shape[0]
|
|
n = b.shape[0]
|
|
k = _logical_k(a)
|
|
assert _logical_k(b) == k, f"{a.shape=}, {b.shape=}"
|
|
assert d.shape == (m, n), f"{d.shape=}, {m=}, {n=}"
|
|
if m == 0 or n == 0:
|
|
return
|
|
full_recipe = get_recipe(recipe, recipe_a, recipe_b)
|
|
a_f32 = _dequant_operand(a, sfa, m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast)
|
|
b_f32 = _dequant_operand(b, sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast)
|
|
_write_output(a_f32 @ b_f32.transpose(0, 1), d, c, alpha)
|
|
|
|
|
|
def fp8_fp4_gemm_nt(
|
|
a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None,
|
|
compiled_dims="nk", disable_ue8m0_cast=False, alpha=None, epilogue=None,
|
|
) -> None:
|
|
"""D = A[M, K] @ B[N, K]^T (+ C), with FP8/FP4 operands and UE8M0 scales."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_fp8_fp4_gemm(a, b, d, c, recipe, recipe_a, recipe_b, disable_ue8m0_cast, alpha, False, False)
|
|
|
|
|
|
def fp8_fp4_gemm_nn(
|
|
a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None,
|
|
compiled_dims="nk", disable_ue8m0_cast=False, alpha=None, epilogue=None,
|
|
) -> None:
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_fp8_fp4_gemm(a, b, d, c, recipe, recipe_a, recipe_b, disable_ue8m0_cast, alpha, False, True)
|
|
|
|
|
|
def fp8_fp4_gemm_tn(
|
|
a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None,
|
|
compiled_dims="mn", disable_ue8m0_cast=False, alpha=None, epilogue=None,
|
|
) -> None:
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_fp8_fp4_gemm(a, b, d, c, recipe, recipe_a, recipe_b, disable_ue8m0_cast, alpha, True, True)
|
|
|
|
|
|
def fp8_fp4_gemm_tt(
|
|
a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None,
|
|
compiled_dims="mn", disable_ue8m0_cast=False, alpha=None, epilogue=None,
|
|
) -> None:
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
_fp8_fp4_gemm(a, b, d, c, recipe, recipe_a, recipe_b, disable_ue8m0_cast, alpha, True, False)
|
|
|
|
|
|
# Upstream keeps separate FP8-only entry points with the same semantics.
|
|
fp8_gemm_nt = fp8_fp4_gemm_nt
|
|
fp8_gemm_nn = fp8_fp4_gemm_nn
|
|
fp8_gemm_tn = fp8_fp4_gemm_tn
|
|
fp8_gemm_tt = fp8_fp4_gemm_tt
|
|
aclnn_fp8_fp4_gemm_nt = fp8_fp4_gemm_nt
|
|
aclnn_fp8_fp4_gemm_nn = fp8_fp4_gemm_nn
|
|
aclnn_fp8_fp4_gemm_tn = fp8_fp4_gemm_tn
|
|
aclnn_fp8_fp4_gemm_tt = fp8_fp4_gemm_tt
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# M-grouped contiguous ("psum") layout
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _psum_group_spans(grouped_layout: torch.Tensor, alignment: int, total_m: int) -> list[tuple[int, int]]:
|
|
"""Decode the contiguous psum layout into per-group ``[start, end)`` row spans.
|
|
|
|
``grouped_layout`` is the cumulative end offset of each group's valid rows; the
|
|
next group starts at ``align(end, alignment)``
|
|
(``DeepGEMM-Ascend/tests/generators.py::generate_psum_layout``).
|
|
"""
|
|
assert grouped_layout.dim() == 1 and grouped_layout.dtype == torch.int32, f"{grouped_layout.shape=}, {grouped_layout.dtype=}"
|
|
# NEURON: host sync - the group boundaries are data dependent, and the number of
|
|
# groups is small (one per local expert), so the loop runs on the host.
|
|
ends = grouped_layout.tolist()
|
|
spans = []
|
|
start = 0
|
|
for end in ends:
|
|
end = max(start, min(int(end), total_m))
|
|
spans.append((start, end))
|
|
start = min((int(end) + alignment - 1) // alignment * alignment, total_m)
|
|
return spans
|
|
|
|
|
|
def m_grouped_bf16_gemm_nt_contiguous(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
d: torch.Tensor,
|
|
grouped_layout: torch.Tensor,
|
|
compiled_dims="nk",
|
|
use_psum_layout: bool = True,
|
|
ensure_zero_padding: bool = True,
|
|
expected_m_for_psum_layout: int | None = None,
|
|
epilogue=None,
|
|
) -> None:
|
|
"""D[rows of group i] = A[rows of group i] @ B[i]^T, contiguous psum layout."""
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
assert use_psum_layout, "only the psum layout is supported"
|
|
assert a.dtype == b.dtype == torch.bfloat16, f"{a.dtype=}, {b.dtype=}"
|
|
num_groups, n, k = b.shape
|
|
total_m = a.shape[0]
|
|
assert a.shape[1] == k and d.shape == (total_m, n), f"{a.shape=}, {b.shape=}, {d.shape=}"
|
|
if total_m == 0 or n == 0:
|
|
return
|
|
from .runtime import get_mk_alignment_for_contiguous_layout
|
|
|
|
spans = _psum_group_spans(grouped_layout, get_mk_alignment_for_contiguous_layout(), total_m)
|
|
assert len(spans) == num_groups, f"{len(spans)=}, {num_groups=}"
|
|
d.zero_()
|
|
for group_idx, (start, end) in enumerate(spans):
|
|
if end <= start:
|
|
continue
|
|
out = a[start:end].float() @ b[group_idx].float().transpose(0, 1)
|
|
d[start:end].copy_(out.to(d.dtype))
|
|
|
|
|
|
def m_grouped_bf16_gemm_nn_contiguous(
|
|
a, b, d, grouped_layout, compiled_dims="nk", use_psum_layout=True,
|
|
ensure_zero_padding=True, expected_m_for_psum_layout=None, epilogue=None,
|
|
) -> None:
|
|
m_grouped_bf16_gemm_nt_contiguous(
|
|
a, b.transpose(1, 2), d, grouped_layout, compiled_dims, use_psum_layout,
|
|
ensure_zero_padding, expected_m_for_psum_layout, epilogue,
|
|
)
|
|
|
|
|
|
def m_grouped_fp8_fp4_gemm_nt_contiguous(
|
|
a_: tuple[torch.Tensor, torch.Tensor],
|
|
b_: tuple[torch.Tensor, torch.Tensor],
|
|
d: torch.Tensor,
|
|
grouped_layout: torch.Tensor,
|
|
recipe=None,
|
|
recipe_a=None,
|
|
recipe_b=None,
|
|
compiled_dims="nk",
|
|
disable_ue8m0_cast: bool = False,
|
|
use_psum_layout: bool = True,
|
|
ensure_zero_padding: bool = True,
|
|
expected_m_for_psum_layout: int | None = None,
|
|
epilogue=None,
|
|
) -> None:
|
|
"""D[rows of group i] = A[rows of group i] @ B[i]^T for FP8/FP4 operands."""
|
|
from .layout import get_recipe
|
|
from .runtime import get_mk_alignment_for_contiguous_layout
|
|
|
|
assert epilogue is None, "epilogue classes are not supported"
|
|
assert use_psum_layout, "only the psum layout is supported"
|
|
a, sfa = a_
|
|
b, sfb = b_
|
|
total_m = a.shape[0]
|
|
num_groups, n = b.shape[0], b.shape[1]
|
|
k = _logical_k(a)
|
|
assert _logical_k(b) == k, f"{a.shape=}, {b.shape=}"
|
|
assert d.shape == (total_m, n), f"{d.shape=}, {total_m=}, {n=}"
|
|
if total_m == 0 or n == 0:
|
|
return
|
|
full_recipe = get_recipe(recipe, recipe_a, recipe_b)
|
|
a_f32 = _dequant_operand(a, sfa, total_m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast)
|
|
b_f32 = _dequant_operand(
|
|
b, sfb, n, k, full_recipe, False, num_groups=num_groups, disable_ue8m0_cast=disable_ue8m0_cast
|
|
)
|
|
|
|
spans = _psum_group_spans(grouped_layout, get_mk_alignment_for_contiguous_layout(), total_m)
|
|
assert len(spans) == num_groups, f"{len(spans)=}, {num_groups=}"
|
|
d.zero_()
|
|
for group_idx, (start, end) in enumerate(spans):
|
|
if end <= start:
|
|
continue
|
|
out = a_f32[start:end] @ b_f32[group_idx].transpose(0, 1)
|
|
d[start:end].copy_(out.to(d.dtype))
|
|
|
|
|
|
def m_grouped_fp8_fp4_gemm_nn_contiguous(
|
|
a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None,
|
|
compiled_dims="nk", disable_ue8m0_cast=False, use_psum_layout=True,
|
|
ensure_zero_padding=True, expected_m_for_psum_layout=None, epilogue=None,
|
|
) -> None:
|
|
data, sf = b
|
|
assert not is_packed_fp4(data.dtype), "packed FP4 B cannot be logically transposed"
|
|
m_grouped_fp8_fp4_gemm_nt_contiguous(
|
|
a, (data.transpose(1, 2), sf.transpose(1, 2)), d, grouped_layout, recipe, recipe_a, recipe_b,
|
|
compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding,
|
|
expected_m_for_psum_layout, epilogue,
|
|
)
|
|
|
|
|
|
m_grouped_fp8_gemm_nt_contiguous = m_grouped_fp8_fp4_gemm_nt_contiguous
|
|
m_grouped_fp8_gemm_nn_contiguous = m_grouped_fp8_fp4_gemm_nn_contiguous
|
|
|
|
|
|
__all__ = [name for name in globals() if "gemm" in name]
|