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>
70 lines
2.4 KiB
Python
70 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_expand.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified.
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""`hc_expand` — `HyperConnectExpand`, the head of the mHC rank-4 residual stream: one hidden
|
|
state becomes `hc_mult` = 4 copies of itself, plus the one-hot `pre_mix` that makes the first
|
|
layer's `hc_pre` return copy 0 unchanged. No arithmetic at all — pure DMA, so both outputs come
|
|
out bit-identical to the reference. Whole module on one core; no collective.
|
|
"""
|
|
|
|
import nki
|
|
import nki.isa as nisa
|
|
import nki.language as nl
|
|
|
|
TILE = 128
|
|
|
|
#: `hc_mult`. Baked, because the one-hot `pre_mix` and the broadcast width are
|
|
#: part of the instruction stream; `model.py` asserts the config agrees.
|
|
HC = 4
|
|
|
|
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(hidden_states):
|
|
"""`[batch, seqlen, dim]` -> (`[batch, seqlen, HC, dim]`, `[batch, seqlen, HC]`)."""
|
|
batch, seqlen, dim = hidden_states.shape
|
|
rows = batch * seqlen
|
|
chunk = _chunk(dim)
|
|
|
|
copies_hbm = nl.ndarray((batch, seqlen, HC, dim), dtype=hidden_states.dtype,
|
|
buffer=nl.shared_hbm)
|
|
mix_hbm = nl.ndarray((batch, seqlen, HC), dtype=nl.float32, buffer=nl.shared_hbm)
|
|
|
|
x3 = hidden_states.reshape((rows, 1, dim))
|
|
copies = copies_hbm.reshape((rows, HC, dim))
|
|
mix = mix_hbm.reshape((rows, HC))
|
|
|
|
n_chunks = dim // chunk
|
|
|
|
def body(row0, nr):
|
|
pm = nl.ndarray((nr, HC), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.memset(pm, 0.0)
|
|
nisa.memset(pm[:, 0:1], 1.0)
|
|
nisa.dma_copy(dst=mix[nl.ds(row0, nr), 0:HC], src=pm,
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
for chunk_idx in nl.static_range(n_chunks):
|
|
col0 = chunk_idx * chunk
|
|
|
|
nisa.dma_copy(dst=copies[nl.ds(row0, nr), 0:HC, nl.ds(col0, chunk)],
|
|
src=x3[nl.ds(row0, nr), 0:1, nl.ds(col0, chunk)].broadcast(1, HC),
|
|
dge_mode=nisa.dge_mode.none)
|
|
|
|
full = rows // TILE
|
|
tail = rows - full * TILE
|
|
for tile_idx in nl.affine_range(full):
|
|
body(tile_idx * TILE, TILE)
|
|
if tail:
|
|
body(full * TILE, tail)
|
|
|
|
return copies_hbm, mix_hbm
|