mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
9c8c85d83d
- Decode runs eagerly and compiles per-layer regions (attention half, MoE) with stage boundaries (neuron_ops/regions.py, DEEPINFER_NEURON_GRAPH_MODE=regions); layer KV views are sliced from the roots inside the graph (_LazyKVCache). - vllm-neuron#71 NKI kernels (Apache-2.0, vendored under neuron_ops/nki/ with attribution), called via torch_neuronx.nki_hop.wrap_nki: blockquant (fp8/fp4 GEMMs incl. grouped einsum, MoE and shared expert on bf16 activations), sparse_attn / index_score / compress_pool (nki_attention.py), hc_mixes / hc_expand / hc_collapse / hc_attn_out / rmsnorm / swiglu (nki_hc.py, nki_norm.py). - Compile-safe formulations: row-major scales, no dtype bitcasts of graph inputs (uint8 weight tensors), whole-row KV stores, byte assembly in int32 narrowed once at the store, arithmetic int16 scale packing, plain fma/rsqrt and bf16 F.linear under compile; default backend deepinfer_neuron (torch fp8 carried as uint8, avoiding NCC_EOCP001 with the kernels' OCP fp8). - Host: tiny model 9/9 vs golden on all prompts; device: kernel tests pass (13 attention, 88 hc/norm); region decode graph compiles except the attention projection stage (being reduced). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/index_score.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""The indexer's score — per-head `q . k`, rectified, and combined by the
|
|
learned per-head weights — as one NKI kernel.
|
|
|
|
index_score = torch.einsum("bshd,btd->bsht", q.float(), index_k.float())
|
|
index_score = (index_score.relu() * weights.float().unsqueeze(-1)).sum(2)
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
#: Partition width of the tensor engine.
|
|
PT = 128
|
|
#: A psum bank holds 2 KiB per partition = 512 fp32, which is also `nc_matmul`'s
|
|
#: widest moving free axis.
|
|
PSUM_FP32 = 512
|
|
|
|
|
|
def _ntile(n):
|
|
"""Free-axis width for the score tile: as wide as a psum bank allows, and a
|
|
divisor of `n` so no compressed position is dropped.
|
|
"""
|
|
if n <= PSUM_FP32:
|
|
return n
|
|
for cand in (PSUM_FP32, 256, PT):
|
|
if n % cand == 0:
|
|
return cand
|
|
return PT
|
|
|
|
|
|
def _align16(n):
|
|
"""Round an SBUF free extent up to 16 elements = 32 bytes, because a
|
|
`dma_transpose` destination has to be 32-byte aligned. See
|
|
`sparse_attn._align16` for the failure it avoids — a NEFF *load* error
|
|
rather than a compile error.
|
|
"""
|
|
return -(-n // 16) * 16
|
|
|
|
|
|
@nki.jit
|
|
def weighted_scores(q, k, w):
|
|
"""Rectified per-head QK scores combined by the indexer's head weights.
|
|
See the module docstring for the layout and the precision notes.
|
|
"""
|
|
tq, d = q.shape
|
|
n, dk = k.shape
|
|
heads, tokens = w.shape
|
|
|
|
assert dk == d, f"k dim {dk} does not match q dim {d}"
|
|
assert tq == tokens * heads, (
|
|
f"q rows {tq} is not tokens({tokens}) * heads({heads})"
|
|
)
|
|
assert d % PT == 0, f"dim={d} must be a multiple of {PT}"
|
|
assert n % PT == 0, f"n={n} must be a multiple of {PT} — the caller rounds"
|
|
assert heads <= PT, f"heads={heads} exceeds the partition width {PT}"
|
|
|
|
dchunks = d // PT
|
|
ktiles = n // PT
|
|
nt = _ntile(n)
|
|
ntiles = n // nt
|
|
hp = _align16(heads)
|
|
|
|
out = nl.ndarray((tokens, n), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
|
|
kt = nl.ndarray((PT, dchunks, n), dtype=nl.bfloat16, buffer=nl.sbuf)
|
|
for c in nl.static_range(dchunks):
|
|
for kb in nl.static_range(ktiles):
|
|
nisa.dma_transpose(
|
|
dst=kt[0:PT, c, kb * PT : (kb + 1) * PT],
|
|
src=k[kb * PT : (kb + 1) * PT, c * PT : (c + 1) * PT],
|
|
)
|
|
|
|
for t in nl.affine_range(tokens):
|
|
qt = nl.ndarray((PT, dchunks, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
|
|
for c in nl.static_range(dchunks):
|
|
nisa.dma_transpose(
|
|
dst=qt[0:PT, c, 0:heads],
|
|
src=q[t * heads : t * heads + heads, c * PT : (c + 1) * PT],
|
|
)
|
|
wt = nl.load(w[0:heads, t : t + 1], dtype=nl.float32)
|
|
|
|
row = nl.ndarray((1, n), dtype=nl.float32, buffer=nl.sbuf)
|
|
for nb in nl.static_range(ntiles):
|
|
acc = nl.ndarray((heads, nt), dtype=nl.float32, buffer=nl.psum)
|
|
for c in nl.static_range(dchunks):
|
|
nisa.nc_matmul(
|
|
dst=acc,
|
|
stationary=qt[0:PT, c, 0:heads],
|
|
moving=kt[0:PT, c, nb * nt : (nb + 1) * nt],
|
|
accumulate=(c > 0),
|
|
)
|
|
rect = nl.ndarray((heads, nt), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(
|
|
dst=rect, data=acc, op0=nl.maximum, operand0=0.0
|
|
)
|
|
wsum = nl.ndarray((1, nt), dtype=nl.float32, buffer=nl.psum)
|
|
nisa.nc_matmul(
|
|
dst=wsum, stationary=wt[0:heads, 0:1], moving=rect[0:heads, 0:nt]
|
|
)
|
|
nisa.tensor_copy(dst=row[0:1, nb * nt : (nb + 1) * nt], src=wsum)
|
|
nl.store(out[t : t + 1, 0:n], value=row)
|
|
return out
|