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>
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/rmsnorm.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""`RMSNorm` as one NKI kernel, with a **row-count-independent** reduction."""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
#: Partition width; one row per partition, 128 rows per tile.
|
|
PT = 128
|
|
|
|
|
|
def _chunk(dim):
|
|
"""Free-axis chunk for the reduction: a divisor of `dim`, small enough that
|
|
one fp32 copy is a few KiB per partition. `engram_gate._chunk`'s rule, and
|
|
the reason this kernel is row-count invariant — it reads `dim` only.
|
|
"""
|
|
for cand in (2048, 1024, 512, 256, 128):
|
|
if dim % cand == 0:
|
|
return cand
|
|
return dim
|
|
|
|
|
|
def _rsqrt(m):
|
|
"""`1 / sqrt(m)` with one Newton step: `y * (1.5 - 0.5 * m * y^2)`."""
|
|
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 rms_norm(x, weight, eps):
|
|
"""`x * rsqrt(mean(x^2, -1) + eps) * weight`. See the module docstring."""
|
|
rows, dim = x.shape
|
|
assert weight.shape == (1, dim), f"weight {weight.shape} wants (1, {dim})"
|
|
assert eps.shape == (1, 1), f"eps {eps.shape} wants (1, 1)"
|
|
|
|
ch = _chunk(dim)
|
|
nchunks = dim // ch
|
|
inv_dim = 1.0 / dim
|
|
|
|
out = nl.ndarray((rows, dim), dtype=x.dtype, buffer=nl.shared_hbm)
|
|
|
|
def body(r0, nr):
|
|
xf = nl.ndarray((nr, dim), dtype=nl.float32, buffer=nl.sbuf)
|
|
acc = None
|
|
for j in range(nchunks):
|
|
xj = nl.load(x[r0 : r0 + nr, j * ch : (j + 1) * ch], dtype=nl.float32)
|
|
nisa.tensor_copy(dst=xf[0:nr, j * ch : (j + 1) * ch], src=xj)
|
|
sq = nl.sum(nl.square(xj), axis=1, keepdims=True)
|
|
acc = sq if acc is None else nl.add(acc, sq)
|
|
|
|
rstd = nl.ndarray((nr, 1), dtype=nl.float32, buffer=nl.sbuf)
|
|
rstd[...] = _rsqrt(nl.add(nl.multiply(acc, inv_dim), nl.broadcast_to(
|
|
nl.load(eps[0:1, 0:1], dtype=nl.float32), (nr, 1)
|
|
)))
|
|
|
|
for j in range(nchunks):
|
|
wj = nl.broadcast_to(
|
|
nl.load(weight[0:1, j * ch : (j + 1) * ch], dtype=nl.float32),
|
|
(nr, ch),
|
|
)
|
|
y = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.activation(
|
|
dst=y, op=nl.copy, data=xf[0:nr, j * ch : (j + 1) * ch], scale=rstd
|
|
)
|
|
nisa.tensor_tensor(dst=y, data1=y, op=nl.multiply, data2=wj)
|
|
nl.store(
|
|
out[r0 : r0 + nr, j * ch : (j + 1) * ch],
|
|
value=nl.copy(y, dtype=x.dtype),
|
|
)
|
|
|
|
full, tail = rows // PT, rows % PT
|
|
for t in nl.affine_range(full):
|
|
body(t * PT, PT)
|
|
if tail:
|
|
body(full * PT, tail)
|
|
return out
|