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>
177 lines
6.7 KiB
Python
177 lines
6.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/swiglu.py,
|
|
# commit 30d54d1351c80728fe5685478fac838b44c53d43), modified:
|
|
# adds `swiglu_fused` (gate and up as the two halves of one `[rows, 2 * inter]` tensor, DeepInfer's layout).
|
|
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
|
|
"""The experts' clamped SwiGLU — and the routing weight — as one NKI kernel.
|
|
|
|
gate = nki_linear(x, w1, ...).float()
|
|
up = nki_linear(x, w3, ...).float()
|
|
up = torch.clamp(up, min=-limit, max=limit)
|
|
gate = torch.clamp(gate, max=limit)
|
|
y = F.silu(gate) * up
|
|
y = weights * y # routing weight, routed experts only
|
|
return nki_linear(y.to(dtype), w2, ...)
|
|
|
|
nisa.tensor_scalar(dst=g, data=g_bf16, op0=nl.minimum, operand0=SWIGLU_LIMIT)
|
|
nisa.tensor_scalar(dst=u, data=u_bf16, op0=nl.maximum, operand0=-SWIGLU_LIMIT,
|
|
op1=nl.minimum, operand1=SWIGLU_LIMIT)
|
|
nisa.activation(dst=sg, op=nl.silu, data=g)
|
|
nisa.scalar_tensor_tensor(dst=dst, data=sg, op0=nl.multiply, operand0=wcol,
|
|
op1=nl.multiply, operand1=u)
|
|
"""
|
|
|
|
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(width):
|
|
"""Free-axis chunk: a divisor of `width` small enough that four fp32 copies
|
|
of it are a few KiB per partition. `rmsnorm._chunk`'s rule, widened at the
|
|
bottom because `moe_intermediate_size = 2304` has no divisor above 256 in
|
|
that list.
|
|
"""
|
|
for cand in (2048, 1024, 512, 256, 128, 64, 32):
|
|
if width % cand == 0:
|
|
return cand
|
|
return width
|
|
|
|
|
|
@nki.jit
|
|
def swiglu(gate, up, weights, limit):
|
|
"""The clamped SwiGLU with its routing weight. See the module docstring.
|
|
|
|
`weights` is `(rows, 1)` for a routed expert — one routing weight a row —
|
|
or `(1, 1)` for the shared expert, where every token carries the same
|
|
weight and there is no per-row vector to materialize.
|
|
"""
|
|
rows, inter = gate.shape
|
|
assert up.shape == (rows, inter), f"up {up.shape} wants ({rows}, {inter})"
|
|
assert weights.shape in ((rows, 1), (1, 1)), (
|
|
f"weights {weights.shape} wants ({rows}, 1) — one routing weight a row "
|
|
"— or (1, 1) for a weight shared by every row"
|
|
)
|
|
assert limit.shape == (1, 1), f"limit {limit.shape} wants (1, 1)"
|
|
uniform = weights.shape[0] == 1
|
|
|
|
ch = _chunk(inter)
|
|
nchunks = inter // ch
|
|
|
|
out = nl.ndarray((rows, inter), dtype=gate.dtype, buffer=nl.shared_hbm)
|
|
|
|
def body(r0, nr):
|
|
lim = nl.broadcast_to(nl.load(limit[0:1, 0:1], dtype=nl.float32), (nr, 1))
|
|
nlim = nl.multiply(lim, -1.0)
|
|
if uniform:
|
|
w = nl.broadcast_to(
|
|
nl.load(weights[0:1, 0:1], dtype=nl.float32), (nr, 1)
|
|
)
|
|
else:
|
|
w = nl.load(weights[r0 : r0 + nr, 0:1], dtype=nl.float32)
|
|
|
|
for j in range(nchunks):
|
|
c0 = j * ch
|
|
g = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(
|
|
dst=g,
|
|
data=nl.load(gate[r0 : r0 + nr, c0 : c0 + ch], dtype=nl.float32),
|
|
op0=nl.minimum,
|
|
operand0=lim,
|
|
)
|
|
u = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(
|
|
dst=u,
|
|
data=nl.load(up[r0 : r0 + nr, c0 : c0 + ch], dtype=nl.float32),
|
|
op0=nl.maximum,
|
|
operand0=nlim,
|
|
op1=nl.minimum,
|
|
operand1=lim,
|
|
)
|
|
sg = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.activation(dst=sg, op=nl.silu, data=g)
|
|
y = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_tensor(dst=y, data1=sg, op=nl.multiply, data2=u)
|
|
nisa.tensor_scalar(dst=y, data=y, op0=nl.multiply, operand0=w)
|
|
nl.store(
|
|
out[r0 : r0 + nr, c0 : c0 + ch], value=nl.copy(y, dtype=gate.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
|
|
|
|
|
|
@nki.jit
|
|
def swiglu_fused(x, weights, limit):
|
|
"""`swiglu(x[:, :inter], x[:, inter:], weights, limit)` without slicing `x` outside the kernel.
|
|
|
|
DeepInfer's gate_up GEMM writes gate and up as the two halves of one row; this reads them in place
|
|
(column offsets) instead of taking two strided views as kernel inputs.
|
|
"""
|
|
rows, width = x.shape
|
|
assert width % 2 == 0, f"x width {width} is not 2 * inter"
|
|
inter = width // 2
|
|
assert weights.shape in ((rows, 1), (1, 1)), (
|
|
f"weights {weights.shape} wants ({rows}, 1) or (1, 1)"
|
|
)
|
|
assert limit.shape == (1, 1), f"limit {limit.shape} wants (1, 1)"
|
|
uniform = weights.shape[0] == 1
|
|
|
|
ch = _chunk(inter)
|
|
nchunks = inter // ch
|
|
|
|
out = nl.ndarray((rows, inter), dtype=x.dtype, buffer=nl.shared_hbm)
|
|
|
|
def body(r0, nr):
|
|
lim = nl.broadcast_to(nl.load(limit[0:1, 0:1], dtype=nl.float32), (nr, 1))
|
|
nlim = nl.multiply(lim, -1.0)
|
|
if uniform:
|
|
w = nl.broadcast_to(
|
|
nl.load(weights[0:1, 0:1], dtype=nl.float32), (nr, 1)
|
|
)
|
|
else:
|
|
w = nl.load(weights[r0 : r0 + nr, 0:1], dtype=nl.float32)
|
|
|
|
for j in range(nchunks):
|
|
c0 = j * ch
|
|
g = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(
|
|
dst=g,
|
|
data=nl.load(x[r0 : r0 + nr, c0 : c0 + ch], dtype=nl.float32),
|
|
op0=nl.minimum,
|
|
operand0=lim,
|
|
)
|
|
u = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_scalar(
|
|
dst=u,
|
|
data=nl.load(
|
|
x[r0 : r0 + nr, inter + c0 : inter + c0 + ch], dtype=nl.float32
|
|
),
|
|
op0=nl.maximum,
|
|
operand0=nlim,
|
|
op1=nl.minimum,
|
|
operand1=lim,
|
|
)
|
|
sg = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.activation(dst=sg, op=nl.silu, data=g)
|
|
y = nl.ndarray((nr, ch), dtype=nl.float32, buffer=nl.sbuf)
|
|
nisa.tensor_tensor(dst=y, data1=sg, op=nl.multiply, data2=u)
|
|
nisa.tensor_scalar(dst=y, data=y, op0=nl.multiply, operand0=w)
|
|
nl.store(
|
|
out[r0 : r0 + nr, c0 : c0 + 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
|