mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 14:49:45 +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>
89 lines
3.4 KiB
Python
89 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/hc_attn_out.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""`hc_attn_out` — `HyperConnectSublayerOut`, the mHC `hc_post`: expand a sublayer's output back
|
|
to `hc` copies of the residual stream, with the residual stream mixed back in through the
|
|
doubly-stochastic `comb`. Broadcast plus an `hc x hc` stream mix, no parameters. Whole module on
|
|
one core; no collective.
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
PMAX = 128
|
|
|
|
CHUNKS = (1024, 512, 256, 128, 64, 32)
|
|
|
|
|
|
def _chunk(dim):
|
|
for cand in CHUNKS:
|
|
if dim % cand == 0:
|
|
return cand
|
|
return dim
|
|
|
|
|
|
@nki.jit
|
|
def kernel(input_0, input_1, input_2, input_3):
|
|
"""`post[..., j, None] * x + sum(comb[..., i, j] * residual[..., i, :] for i)`."""
|
|
batch, seqlen, dim = input_0.shape
|
|
hc = input_2.shape[-1]
|
|
rows = batch * seqlen
|
|
chunk = _chunk(dim)
|
|
|
|
out_hbm = nl.ndarray((batch, seqlen, hc, dim), dtype=input_0.dtype,
|
|
buffer=nl.shared_hbm)
|
|
|
|
x_hbm = input_0.reshape((rows, dim))
|
|
res_hbm = input_1.reshape((rows, hc, dim))
|
|
post_hbm = input_2.reshape((rows, hc))
|
|
comb_hbm = input_3.reshape((rows, hc * hc))
|
|
y_hbm = out_hbm.reshape((rows, hc, dim))
|
|
|
|
def body(row0, nr):
|
|
post = nl.load(post_hbm[nl.ds(row0, nr), 0:hc])
|
|
comb = nl.load(comb_hbm[nl.ds(row0, nr), 0:hc * hc])
|
|
|
|
for c in nl.static_range(dim // chunk):
|
|
col0 = c * chunk
|
|
|
|
xg = nl.ndarray((nr, chunk), dtype=input_0.dtype, buffer=nl.sbuf)
|
|
nisa.dma_copy(dst=xg, src=x_hbm[nl.ds(row0, nr), nl.ds(col0, chunk)],
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
rg = nl.ndarray((nr, hc, chunk), dtype=input_1.dtype, buffer=nl.sbuf)
|
|
nisa.dma_copy(dst=rg,
|
|
src=res_hbm[nl.ds(row0, nr), 0:hc, nl.ds(col0, chunk)],
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
og = nl.ndarray((nr, hc, chunk), dtype=input_0.dtype, buffer=nl.sbuf)
|
|
acc = nl.ndarray((nr, hc, chunk), dtype=nl.float32, buffer=nl.sbuf)
|
|
|
|
for j in nl.static_range(hc):
|
|
nisa.tensor_scalar(dst=acc[:, j, :], data=rg[:, 0, :],
|
|
op0=nl.multiply, operand0=comb[:, j:j + 1],
|
|
engine=nisa.scalar_engine)
|
|
for i in nl.static_range(1, hc):
|
|
col = i * hc + j
|
|
nisa.scalar_tensor_tensor(dst=acc[:, j, :], data=rg[:, i, :],
|
|
op0=nl.multiply,
|
|
operand0=comb[:, col:col + 1],
|
|
op1=nl.add, operand1=acc[:, j, :])
|
|
|
|
nisa.scalar_tensor_tensor(dst=og[:, j, :], data=xg,
|
|
op0=nl.multiply, operand0=post[:, j:j + 1],
|
|
op1=nl.add, operand1=acc[:, j, :])
|
|
|
|
nisa.dma_copy(dst=y_hbm[nl.ds(row0, nr), 0:hc, nl.ds(col0, chunk)], src=og,
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
full = rows // PMAX
|
|
tail = rows - full * PMAX
|
|
for t in nl.affine_range(full):
|
|
body(t * PMAX, PMAX)
|
|
if tail:
|
|
body(full * PMAX, tail)
|
|
|
|
return out_hbm
|