mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 22:59:44 +08:00
3629d113fb
- PrepareArgs.prompt_as_decode: decode-layout metadata for a prompt chunk; input_ids are the prompt's own tokens instead of last token + MTP drafts. - DEEPINFER_NEURON_PREFILL_BUCKET adds one graph size; the model runner pads a prompt chunk to it with the decode placeholders (every EP rank takes the same bucket, so the static collectives match). Sampling still sees a prompt step. - nki_attention.paged_mqa_logits: one gather + one kernel call for a chunk (one sequence) instead of a T-way loop. - bf16 linear for the LM head eagerly on a NeuronCore (b.float() was a 2.6 GiB copy a call); DEEPINFER_NEURON_REGION_TIMING; --num_prompts. Host mode, tiny V4.1, chunk 128: same tokens as eager prefill (short 9/9, long 715 tokens in 6 chunks 8/9 against golden). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
428 lines
16 KiB
Python
428 lines
16 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
|
|
|
|
try: # the NKI kernels live in deepinfer; standalone shim tests run without it
|
|
from deepinfer.neuron_ops import nki_gemm
|
|
except ImportError: # pragma: no cover
|
|
nki_gemm = None
|
|
|
|
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])
|
|
|
|
|
|
# Dequantized weight rows per chunk (fp32 elements). The device is a 24 GB NeuronCore that already holds the
|
|
# packed weights: dequantizing a whole projection (or all local experts) to fp32 at once ran the full model out
|
|
# of HBM. Chunking along N keeps every output element's K-reduction intact, so results are unchanged.
|
|
_CHUNK_ELEMS = 1 << 25
|
|
|
|
|
|
def _row_chunks(n: int, k: int) -> list[tuple[int, int]]:
|
|
step = max(1, _CHUNK_ELEMS // max(1, k))
|
|
return [(n0, min(n, n0 + step)) for n0 in range(0, n, step)]
|
|
|
|
|
|
def _required_sf(sf, mn, k, recipe, is_sfa, num_groups=None, disable_ue8m0_cast=False):
|
|
return transform_sf_into_required_layout(
|
|
sf, mn, k, recipe, num_groups=num_groups, is_sfa=is_sfa, disable_ue8m0_cast=disable_ue8m0_cast
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 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
|
|
if torch.compiler.is_compiling() or b.device.type == "neuron":
|
|
# NEURON graphs: one bf16 dot with fp32 accumulation and no materialized fp32 transpose (vllm-neuron#71 runs
|
|
# these as bf16 linears; an fp32 transpose feeding a dot is a neuronx-cc InferNonlocalTensors trigger).
|
|
# Eagerly on a NeuronCore too: ``b.float()`` of the LM head is a 2.6 GiB copy a call.
|
|
_write_output(torch.nn.functional.linear(a, b).float(), 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
|
|
if c is None and alpha is None and not trans_a and not trans_b and nki_gemm is not None and nki_gemm.supported(a, b, k, n):
|
|
nki_gemm.gemm(a, sfa, b, d)
|
|
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)
|
|
sfb_req = _required_sf(sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast)
|
|
c_is_d = c is not None and _is_alias(c, d)
|
|
for n0, n1 in _row_chunks(n, k):
|
|
b_f32 = dequantize(b[n0:n1], sfb_req[n0:n1], gran_k=full_recipe[2])
|
|
c_chunk = None if c is None else (d[:, n0:n1] if c_is_d else c[..., n0:n1])
|
|
_write_output(a_f32 @ b_f32.transpose(0, 1), d[:, n0:n1], c_chunk, 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)
|
|
sfb_req = _required_sf(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
|
|
# One expert (and one N chunk) at a time: never all local experts dequantized at once.
|
|
for n0, n1 in _row_chunks(n, k):
|
|
b_f32 = dequantize(b[group_idx, n0:n1], sfb_req[group_idx, n0:n1], gran_k=full_recipe[2])
|
|
out = a_f32[start:end] @ b_f32.transpose(0, 1)
|
|
d[start:end, n0:n1].copy_(out.to(d.dtype))
|
|
|
|
|
|
def m_grouped_fp8_fp4_gemm_nt_shared_a(
|
|
a_: tuple[torch.Tensor, torch.Tensor],
|
|
b_: tuple[torch.Tensor, torch.Tensor],
|
|
d: torch.Tensor,
|
|
group_idx: int,
|
|
recipe=None,
|
|
recipe_a=None,
|
|
recipe_b=None,
|
|
disable_ue8m0_cast: bool = False,
|
|
) -> None:
|
|
"""NEURON extension, static shapes: ``D = A @ B[group_idx]^T`` with every row of A (no psum layout, no host
|
|
sync). The static decode MoE runs each local expert over all gathered rows."""
|
|
from .layout import get_recipe
|
|
|
|
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 nki_gemm is not None and nki_gemm.supported(a, b, k, n):
|
|
nki_gemm.gemm(a, sfa, b, d, group=group_idx)
|
|
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)
|
|
sfb_req = _required_sf(sfb, n, k, full_recipe, False, num_groups=num_groups, disable_ue8m0_cast=disable_ue8m0_cast)
|
|
for n0, n1 in _row_chunks(n, k):
|
|
b_f32 = dequantize(b[group_idx, n0:n1], sfb_req[group_idx, n0:n1], gran_k=full_recipe[2])
|
|
d[:, n0:n1].copy_((a_f32 @ b_f32.transpose(0, 1)).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]
|