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>
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/hc_collapse.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""`hc_collapse` — `HyperConnectCollapse`, the mHC `hc_pre`: the `hc` parallel copies of the
|
|
residual stream are mixed back down into the single hidden state a sublayer reads. One fp32
|
|
multiply-accumulate over the copies, no parameters. Whole module on one core; no collective.
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
PMAX = 128
|
|
|
|
#: Free-axis chunk candidates, widest first. 1024 bf16 is 2 KiB per partition.
|
|
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):
|
|
"""`sum(pre_mix[..., c, None] * x[..., c, :] for c in range(hc))`."""
|
|
batch, n_tok, hc, dim = input_0.shape
|
|
rows = batch * n_tok
|
|
chunk = _chunk(dim)
|
|
|
|
y_out = nl.ndarray((batch, n_tok, dim), dtype=input_0.dtype, buffer=nl.shared_hbm)
|
|
|
|
x3 = input_0.reshape((rows, hc, dim))
|
|
pm2 = input_1.reshape((rows, hc))
|
|
y2 = y_out.reshape((rows, dim))
|
|
|
|
def body(row0, nr):
|
|
pm = nl.load(pm2[nl.ds(row0, nr), 0:hc])
|
|
|
|
for c in nl.static_range(dim // chunk):
|
|
f0 = c * chunk
|
|
|
|
acc = nl.ndarray((nr, chunk), dtype=nl.float32, buffer=nl.sbuf)
|
|
xg = nl.ndarray((nr, hc, chunk), dtype=input_0.dtype, buffer=nl.sbuf)
|
|
|
|
nisa.dma_copy(dst=xg, src=x3[nl.ds(row0, nr), 0:hc, nl.ds(f0, chunk)],
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
for j in nl.static_range(hc):
|
|
if j == 0:
|
|
nisa.tensor_scalar(dst=acc, data=xg[:, j, :], op0=nl.multiply,
|
|
operand0=pm[:, j:j + 1])
|
|
else:
|
|
nisa.scalar_tensor_tensor(dst=acc, data=xg[:, j, :], op0=nl.multiply,
|
|
operand0=pm[:, j:j + 1],
|
|
op1=nl.add, operand1=acc)
|
|
|
|
nisa.dma_copy(dst=y2[nl.ds(row0, nr), nl.ds(f0, chunk)], src=acc,
|
|
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 y_out
|