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>
181 lines
7.0 KiB
Python
181 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/hc_mixes.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""The hyper-connection mixing coefficients — the stream statistic, the three
|
|
nonlinearities and the 20-iteration Sinkhorn — as one NKI kernel.
|
|
|
|
rms one statistic per token over the whole flattened `hc * dim` stream
|
|
pre sigmoid(raw * scale0 + base0) + eps
|
|
post 2 * sigmoid(raw * scale1 + base1)
|
|
comb softmax(raw * scale2 + base2) + eps, then Sinkhorn
|
|
|
|
AssertionError: ap() pattern has invalid partition stride. Partition
|
|
step 1 must equal tensor free dimension size 3.
|
|
|
|
A freshly `nl.load`ed `(1, n)` tile broadcast to `(rows, n)` is fine, which
|
|
is why every coefficient here is loaded from its own HBM slice rather than
|
|
sliced out of one staged tile.
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
#: Partition width.
|
|
PT = 128
|
|
#: `hc_sinkhorn_iters`. Not a convergence tolerance — the reference runs exactly
|
|
#: this many passes, so matching the count matters bit-for-bit.
|
|
SINKHORN_ITERS = 20
|
|
#: `hc_eps`.
|
|
HC_EPS = 1e-06
|
|
#: `rms_norm_eps`.
|
|
NORM_EPS = 1e-20
|
|
|
|
|
|
def _hc_from_width(width):
|
|
"""`width = (2 + hc) * hc` inverted. 24 -> 4."""
|
|
hc = 1
|
|
while (2 + hc) * hc < width:
|
|
hc += 1
|
|
assert (2 + hc) * hc == width, f"width={width} is not (2 + hc) * hc"
|
|
return hc
|
|
|
|
|
|
def _sqchunk(fd):
|
|
"""Free-axis chunk for the sum of squares: a divisor of the stream width,
|
|
small enough that one fp32 copy of it is a few KiB per partition.
|
|
"""
|
|
for cand in (2048, 1024, 512, 256, 128):
|
|
if fd % cand == 0:
|
|
return cand
|
|
return fd
|
|
|
|
|
|
def _rms_recip(flat, r0, nr, fd):
|
|
"""`1 / sqrt(mean(flat^2) + eps)` for each of `nr` rows, `[nr, 1]` fp32."""
|
|
ch = _sqchunk(fd)
|
|
total = None
|
|
for c in range(fd // ch):
|
|
x = nl.load(flat[r0 : r0 + nr, c * ch : (c + 1) * ch], dtype=nl.float32)
|
|
part = nl.sum(nl.square(x), axis=1, keepdims=True)
|
|
total = part if total is None else nl.add(total, part)
|
|
m = nl.add(nl.multiply(total, 1.0 / fd), NORM_EPS)
|
|
y = nl.rsqrt(m)
|
|
corr = nl.add(nl.multiply(nl.multiply(m, nl.multiply(y, y)), -0.5), 1.5)
|
|
return nl.multiply(y, corr)
|
|
|
|
|
|
@nki.jit
|
|
def kernel(raw, flat, base, scale):
|
|
"""See the module docstring. Returns `(pre, post, comb)`."""
|
|
rows, width = raw.shape
|
|
frows, fd = flat.shape
|
|
hc = _hc_from_width(width)
|
|
|
|
assert frows == rows, f"flat has {frows} rows, raw has {rows}"
|
|
assert fd % hc == 0, f"stream width {fd} is not a multiple of hc={hc}"
|
|
assert base.shape == (1, width), f"base {base.shape} wants (1, {width})"
|
|
assert scale.shape == (1, 3), f"scale {scale.shape} wants (1, 3)"
|
|
|
|
pre = nl.ndarray((rows, hc), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
post = nl.ndarray((rows, hc), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
comb = nl.ndarray((rows, hc * hc), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
|
|
def body(r0, nr):
|
|
rs = _rms_recip(flat, r0, nr, fd)
|
|
rawt = nl.load(raw[r0 : r0 + nr, 0:width], dtype=nl.float32)
|
|
mix = nl.ndarray((nr, width), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(dst=mix, data=rawt, op0=nl.multiply, operand0=rs)
|
|
|
|
s0 = nl.broadcast_to(nl.load(scale[0:1, 0:1]), (nr, 1))
|
|
s1 = nl.broadcast_to(nl.load(scale[0:1, 1:2]), (nr, 1))
|
|
s2 = nl.broadcast_to(nl.load(scale[0:1, 2:3]), (nr, 1))
|
|
b0 = nl.broadcast_to(nl.load(base[0:1, 0:hc]), (nr, hc))
|
|
b1 = nl.broadcast_to(nl.load(base[0:1, hc : 2 * hc]), (nr, hc))
|
|
b2 = nl.broadcast_to(
|
|
nl.load(base[0:1, 2 * hc : 2 * hc + hc * hc]), (nr, hc * hc)
|
|
)
|
|
|
|
t0 = nl.ndarray((nr, hc), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.scalar_tensor_tensor(
|
|
dst=t0, data=mix[0:nr, 0:hc], op0=nl.multiply, operand0=s0,
|
|
op1=nl.add, operand1=b0,
|
|
)
|
|
nl.store(pre[r0 : r0 + nr, 0:hc], value=nl.add(nl.sigmoid(t0), HC_EPS))
|
|
|
|
t1 = nl.ndarray((nr, hc), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.scalar_tensor_tensor(
|
|
dst=t1, data=mix[0:nr, hc : 2 * hc], op0=nl.multiply, operand0=s1,
|
|
op1=nl.add, operand1=b1,
|
|
)
|
|
nl.store(post[r0 : r0 + nr, 0:hc], value=nl.multiply(nl.sigmoid(t1), 2.0))
|
|
|
|
t2 = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.scalar_tensor_tensor(
|
|
dst=t2,
|
|
data=mix[0:nr, 2 * hc : 2 * hc + hc * hc],
|
|
op0=nl.multiply, operand0=s2, op1=nl.add, operand1=b2,
|
|
)
|
|
def rowsum_div(src, dst):
|
|
"""Divide each `hc x hc` row by its own sum (the reference's
|
|
`dim=-1` normalization). The divisor is one value per token, so it
|
|
rides the per-partition operand and needs no broadcast.
|
|
"""
|
|
for i in range(hc):
|
|
lo, hi = i * hc, (i + 1) * hc
|
|
s = nl.sum(src[0:nr, lo:hi], axis=1, keepdims=True)
|
|
nisa.tensor_scalar(
|
|
dst=dst[0:nr, lo:hi],
|
|
data=src[0:nr, lo:hi],
|
|
op0=nl.multiply,
|
|
operand0=nl.reciprocal(nl.add(s, HC_EPS)),
|
|
)
|
|
|
|
def colsum_div(src, dst):
|
|
"""Divide each `hc x hc` column by its own sum (the reference's
|
|
`dim=-2` normalization). The sum of the row blocks *is* the divisor
|
|
laid out the way each block wants it.
|
|
"""
|
|
s = nl.add(src[0:nr, 0:hc], src[0:nr, hc : 2 * hc])
|
|
for i in range(2, hc):
|
|
s = nl.add(s, src[0:nr, i * hc : (i + 1) * hc])
|
|
d = nl.reciprocal(nl.add(s, HC_EPS))
|
|
for i in range(hc):
|
|
lo, hi = i * hc, (i + 1) * hc
|
|
nisa.tensor_tensor(
|
|
dst=dst[0:nr, lo:hi], data1=src[0:nr, lo:hi], data2=d,
|
|
op=nl.multiply,
|
|
)
|
|
|
|
a = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf)
|
|
b = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf)
|
|
for i in range(hc):
|
|
lo, hi = i * hc, (i + 1) * hc
|
|
mx = nl.max(t2[0:nr, lo:hi], axis=1, keepdims=True)
|
|
e = nl.exp(
|
|
nl.subtract(t2[0:nr, lo:hi], nl.broadcast_to(mx, (nr, hc)))
|
|
)
|
|
s = nl.sum(e, axis=1, keepdims=True)
|
|
nisa.tensor_scalar(
|
|
dst=a[0:nr, lo:hi],
|
|
data=e,
|
|
op0=nl.multiply,
|
|
operand0=nl.reciprocal(s),
|
|
op1=nl.add,
|
|
operand1=HC_EPS,
|
|
)
|
|
|
|
colsum_div(a, b)
|
|
for _ in range(SINKHORN_ITERS - 1):
|
|
rowsum_div(b, a)
|
|
colsum_div(a, b)
|
|
nl.store(comb[r0 : r0 + nr, 0 : hc * hc], value=b[0:nr, 0 : hc * hc])
|
|
|
|
full, tail = rows // PT, rows % PT
|
|
for t in nl.affine_range(full):
|
|
body(t * PT, PT)
|
|
if tail:
|
|
body(full * PT, tail)
|
|
return pre, post, comb
|