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>
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/compress_pool.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""The KV compressor's learned softmax pooling, as one NKI kernel.
|
|
|
|
g0 = nl.load(cgt[r0:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
|
|
g1 = nl.load(cgt[r0 + 1:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
|
|
v0 = nl.load(ckv[r0:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
|
|
v1 = nl.load(ckv[r0 + 1:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
|
|
mx = nl.maximum(g0, g1)
|
|
e0 = nl.exp(nl.subtract(g0, mx))
|
|
e1 = nl.exp(nl.subtract(g1, mx))
|
|
nisa.reciprocal(dst=rz, data=nl.add(e0, e1))
|
|
p0 = nl.multiply(e0, rz)
|
|
p1 = nl.multiply(e1, rz)
|
|
acc = nl.add(nl.multiply(v0, p0), nl.multiply(v1, p1))
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
#: Partition width; one *group* per partition, 128 groups per tile.
|
|
PT = 128
|
|
|
|
|
|
def _chunk(dim):
|
|
"""Free-axis chunk for the pooling: a divisor of `dim` small enough that
|
|
`ratio + 3` fp32 copies of it fit a partition comfortably.
|
|
"""
|
|
for cand in (2048, 1024, 512, 256, 128):
|
|
if dim % cand == 0:
|
|
return cand
|
|
return dim
|
|
|
|
|
|
def _recip(d, shape):
|
|
"""`1 / d` with one Newton step: `y * (2 - d * y)`."""
|
|
y = nl.ndarray(shape, dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.reciprocal(dst=y, data=d)
|
|
return nl.multiply(y, nl.add(nl.multiply(nl.multiply(d, y), -1.0), 2.0))
|
|
|
|
|
|
@nki.jit
|
|
def softmax_pool(kv, score):
|
|
"""`(kv * score.softmax(dim=1)).sum(dim=1)`. See the module docstring."""
|
|
groups, ratio, dim = kv.shape
|
|
assert score.shape == (groups, ratio, dim), (
|
|
f"score {score.shape} wants ({groups}, {ratio}, {dim})"
|
|
)
|
|
assert ratio >= 2, f"ratio={ratio} — a single-member group needs no pooling"
|
|
|
|
kv2 = kv.reshape((groups * ratio, dim))
|
|
sc2 = score.reshape((groups * ratio, dim))
|
|
|
|
ch = _chunk(dim)
|
|
nchunks = dim // ch
|
|
|
|
out = nl.ndarray((groups, dim), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
|
|
def body(g0, ng):
|
|
r0, rn = g0 * ratio, ng * ratio
|
|
for j in range(nchunks):
|
|
c0 = j * ch
|
|
g = [
|
|
nl.load(
|
|
sc2[r0 + m : r0 + rn : ratio, c0 : c0 + ch], dtype=nl.float32
|
|
)
|
|
for m in range(ratio)
|
|
]
|
|
mx = g[0]
|
|
for m in range(1, ratio):
|
|
mx = nl.maximum(mx, g[m])
|
|
e = [nl.exp(nl.subtract(g[m], mx)) for m in range(ratio)]
|
|
den = e[0]
|
|
for m in range(1, ratio):
|
|
den = nl.add(den, e[m])
|
|
rz = _recip(den, (ng, ch))
|
|
|
|
acc = None
|
|
for m in range(ratio):
|
|
v = nl.load(
|
|
kv2[r0 + m : r0 + rn : ratio, c0 : c0 + ch], dtype=nl.float32
|
|
)
|
|
term = nl.multiply(v, nl.multiply(e[m], rz))
|
|
acc = term if acc is None else nl.add(acc, term)
|
|
nl.store(out[g0 : g0 + ng, c0 : c0 + ch], value=acc)
|
|
|
|
full, tail = groups // PT, groups % PT
|
|
for t in nl.affine_range(full):
|
|
body(t * PT, PT)
|
|
if tail:
|
|
body(full * PT, tail)
|
|
return out
|