coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/neuron_ops/nki/sparse_attn.py
whn09 9c8c85d83d Neuron: per-layer region graphs + vllm-neuron#71 NKI kernels (issues #10, #11, #13, #25) 
- 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>
2026-10-22 17:36:54 +08:00

328 lines
13 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/sparse_attn.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""The attention arithmetic — QK score, masked softmax with the learned sink,
and the probability-weighted sum — as one NKI kernel.
* the QK product takes bf16 operands into an fp32 psum accumulator, and the
scale multiplies the *logits* rather than `q`, so it never passes through a
bf16 rounding;
* the denominator sums **fp32** probabilities while the PV product consumes
**bf16** ones — that is the reference's own split (its staging buffer is
bf16), which is why `_softmax` writes the exponentials twice;
* the sink joins the denominator after the row max is final.
"""
import nki
import nki.isa as nisa
import nki.language as nl
#: Partition width of the tensor engine.
PT = 128
#: A psum bank holds 2 KiB per partition = 512 fp32, which is also `nc_matmul`'s
#: widest moving free axis.
PSUM_FP32 = 512
#: `gathered_attn_idx` takes the softmax scale in column 0 of its bias tensor
#: rather than as its own operand, and this is how much of the row comes before
#: the bias proper. 16 fp32 = 64 bytes, so the bias slice the score product
BIAS_PAD = 16
def _ktile(k):
"""Free-axis width for the score tile: as wide as a psum bank allows, and a
divisor of `k` so no slots are dropped.
"""
if k <= PSUM_FP32:
return k
for cand in (PSUM_FP32, 256, PT):
if k % cand == 0:
return cand
return PT
def _align16(n):
"""Round an SBUF free extent up to 16 elements = 32 bytes.
TDRV:dma_util_expand_desc_transpose transpose dest offset ... must be
32B aligned
"""
return -(-n // 16) * 16
@nki.jit
def gathered_attn(q, g, bias, sink):
"""Multi-head attention over per-query *pre-gathered* slots.
* `softmax_scale` arrives in column 0 of `bias` rather than as its own
`(1, 1)` operand — see `BIAS_PAD`. As two separate broadcast tiles the
scale and the bias hit `[NCC_IBIR297]` (two SBUF inputs of one
instruction at different base partitions) at `tokens=512, topk=128`,
in *this* kernel too, so the shipped form cannot be compiled standalone
at every live shape and could not serve as the reference unchanged. The
instruction and both operand *values* are untouched, and the tiny-shape
run of `devcmp.py` compared the shipped form against the combined one
and got identical bytes;
* nothing else.
"""
tq, d = q.shape
tokens, kb_ = bias.shape
k = kb_ - BIAS_PAD
heads = tq // tokens
assert tq == tokens * heads, f"q rows {tq} is not tokens({tokens}) * heads"
assert g.shape == (tokens * k, d), (
f"g {g.shape} does not match tokens({tokens}) * topk({k}) x head_dim({d})"
)
assert d % PT == 0, f"head_dim={d} must be a multiple of {PT}"
assert k % PT == 0, f"topk={k} must be a multiple of {PT} — the caller pads"
assert heads <= PT, f"heads={heads} exceeds the partition width {PT}"
assert sink.shape == (heads, 1), f"sink {sink.shape} wants ({heads}, 1)"
dchunks = d // PT
ktiles = k // PT
kt = _ktile(k)
ntiles = k // kt
hp = _align16(heads)
out = nl.ndarray((tq, d), dtype=nl.bfloat16, buffer=nl.shared_hbm)
sink_t = nl.load(sink[0:heads, 0:1], dtype=nl.float32)
for t in nl.affine_range(tokens):
qt = nl.ndarray((PT, dchunks, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
gt = nl.ndarray((PT, dchunks, k), dtype=nl.bfloat16, buffer=nl.sbuf)
gn = nl.ndarray((PT, ktiles, d), dtype=nl.bfloat16, buffer=nl.sbuf)
for c in nl.static_range(dchunks):
nisa.dma_transpose(
dst=qt[0:PT, c, 0:heads],
src=q[t * heads : t * heads + heads, c * PT : (c + 1) * PT],
)
nisa.dma_transpose(
dst=gt[0:PT, c, 0:k],
src=g[t * k : (t + 1) * k, c * PT : (c + 1) * PT],
)
for kb in nl.static_range(ktiles):
nisa.tensor_copy(
dst=gn[0:PT, kb, 0:d],
src=nl.load(
g[t * k + kb * PT : t * k + (kb + 1) * PT, 0:d],
dtype=nl.bfloat16,
),
)
bias_b = nl.broadcast_to(
nl.load(bias[t : t + 1, 0 : BIAS_PAD + k]), (heads, BIAS_PAD + k)
)
scale_p = bias_b[0:heads, 0:1]
sc = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
for nb in nl.static_range(ntiles):
acc = nl.ndarray((heads, kt), dtype=nl.float32, buffer=nl.psum)
for c in nl.static_range(dchunks):
nisa.nc_matmul(
dst=acc,
stationary=qt[0:PT, c, 0:heads],
moving=gt[0:PT, c, nb * kt : (nb + 1) * kt],
accumulate=(c > 0),
)
nisa.scalar_tensor_tensor(
dst=sc[0:heads, nb * kt : (nb + 1) * kt],
data=acc,
op0=nl.multiply,
operand0=scale_p,
op1=nl.add,
operand1=bias_b[
0:heads, BIAS_PAD + nb * kt : BIAS_PAD + (nb + 1) * kt
],
)
neg_max = nl.negative(nl.max(sc[0:heads, 0:k], axis=1, keepdims=True))
p32 = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
den = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.activation(
dst=p32,
op=nl.exp,
data=sc[0:heads, 0:k],
bias=neg_max,
scale=1.0,
reduce_op=nl.add,
reduce_res=den,
)
sk = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.activation(dst=sk, op=nl.exp, data=sink_t, bias=neg_max)
rcp = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.reciprocal(dst=rcp, data=nl.add(den, sk))
p16 = nl.copy(p32, dtype=nl.bfloat16)
pt = nl.ndarray((PT, ktiles, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
for kb in nl.static_range(ktiles):
ptps = nl.ndarray((PT, hp), dtype=nl.bfloat16, buffer=nl.psum)
nisa.nc_transpose(
dst=ptps[0:PT, 0:heads], data=p16[0:heads, kb * PT : (kb + 1) * PT]
)
nisa.tensor_copy(
dst=pt[0:PT, kb, 0:heads], src=ptps[0:PT, 0:heads]
)
acc_o = nl.ndarray((heads, d), dtype=nl.float32, buffer=nl.psum)
for kb in nl.static_range(ktiles):
nisa.nc_matmul(
dst=acc_o,
stationary=pt[0:PT, kb, 0:heads],
moving=gn[0:PT, kb, 0:d],
accumulate=(kb > 0),
)
o16 = nl.ndarray((heads, d), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=o16, data=acc_o, op0=nl.multiply, operand0=rcp
)
nl.store(out[t * heads : t * heads + heads, 0:d], value=o16)
return out
@nki.jit
def gathered_attn_idx(q, kv, idx, bias, sink):
"""`gathered_attn` with the gather inside it. **This is the live path.**
* the logit of an empty slot is `q·kv[0] * scale + (-1e30)`. `|q·kv|` is
bounded by a few times `head_dim` and one ulp of 1e30 is ~1.2e23, so
the sum is **exactly** -1e30, which is what the zeroed row gives. The
row max, and so every probability, is therefore identical;
* the value product weights that row by `p = exp(-1e30 - max) = 0`, and
`0 * finite` is exactly 0. `kv` is fp8/fp4-rounded activation, never
inf or NaN, so the product cannot be NaN;
* in the all-empty case (`p = 1` for every slot) the accumulator holds
`sum(kv[0])` instead of 0, and the reciprocal of the denominator is
exactly 0 because `exp(sink + 1e30)` is `+inf` — so the row is exactly
zero either way, which is the reference's convention.
"""
tq, d = q.shape
tokens, kb_ = bias.shape
k = kb_ - BIAS_PAD
heads = tq // tokens
nrows, kvd = kv.shape
assert tq == tokens * heads, f"q rows {tq} is not tokens({tokens}) * heads"
assert kvd == d, f"kv head_dim {kvd} does not match q head_dim {d}"
assert idx.shape == (tokens, k), (
f"idx {idx.shape} does not match tokens({tokens}) x topk({k})"
)
assert nrows > 0, "kv is empty"
assert d % PT == 0, f"head_dim={d} must be a multiple of {PT}"
assert k % PT == 0, f"topk={k} must be a multiple of {PT} — the caller pads"
assert heads <= PT, f"heads={heads} exceeds the partition width {PT}"
assert sink.shape == (heads, 1), f"sink {sink.shape} wants ({heads}, 1)"
dchunks = d // PT
ktiles = k // PT
kt = _ktile(k)
ntiles = k // kt
hp = _align16(heads)
out = nl.ndarray((tq, d), dtype=nl.bfloat16, buffer=nl.shared_hbm)
sink_t = nl.load(sink[0:heads, 0:1], dtype=nl.float32)
for t in nl.affine_range(tokens):
qt = nl.ndarray((PT, dchunks, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
gt = nl.ndarray((PT, dchunks, k), dtype=nl.bfloat16, buffer=nl.sbuf)
gn = nl.ndarray((PT, ktiles, d), dtype=nl.bfloat16, buffer=nl.sbuf)
ids = nl.ndarray((PT, ktiles), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(
dst=ids, src=idx[t : t + 1, 0:k].ap(pattern=[[1, PT], [PT, ktiles]])
)
for c in nl.static_range(dchunks):
nisa.dma_transpose(
dst=qt[0:PT, c, 0:heads],
src=q[t * heads : t * heads + heads, c * PT : (c + 1) * PT],
)
for kb in nl.static_range(ktiles):
nisa.dma_copy(
dst=gn[0:PT, kb, 0:d],
src=kv.ap(
pattern=[[d, PT], [1, d]],
offset=0,
vector_offset=ids[0:PT, kb : kb + 1],
indirect_dim=0,
),
)
tps = nl.ndarray((PT, dchunks, PT), dtype=nl.bfloat16, buffer=nl.psum)
for c in nl.static_range(dchunks):
nisa.nc_transpose(
dst=tps[0:PT, c, 0:PT],
data=gn[0:PT, kb, c * PT : (c + 1) * PT],
)
nisa.tensor_copy(
dst=gt[0:PT, 0:dchunks, kb * PT : (kb + 1) * PT], src=tps
)
bias_b = nl.broadcast_to(
nl.load(bias[t : t + 1, 0 : BIAS_PAD + k]), (heads, BIAS_PAD + k)
)
scale_p = bias_b[0:heads, 0:1]
sc = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
for nb in nl.static_range(ntiles):
acc = nl.ndarray((heads, kt), dtype=nl.float32, buffer=nl.psum)
for c in nl.static_range(dchunks):
nisa.nc_matmul(
dst=acc,
stationary=qt[0:PT, c, 0:heads],
moving=gt[0:PT, c, nb * kt : (nb + 1) * kt],
accumulate=(c > 0),
)
nisa.scalar_tensor_tensor(
dst=sc[0:heads, nb * kt : (nb + 1) * kt],
data=acc,
op0=nl.multiply,
operand0=scale_p,
op1=nl.add,
operand1=bias_b[
0:heads, BIAS_PAD + nb * kt : BIAS_PAD + (nb + 1) * kt
],
)
neg_max = nl.negative(nl.max(sc[0:heads, 0:k], axis=1, keepdims=True))
p32 = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
den = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.activation(
dst=p32,
op=nl.exp,
data=sc[0:heads, 0:k],
bias=neg_max,
scale=1.0,
reduce_op=nl.add,
reduce_res=den,
)
sk = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.activation(dst=sk, op=nl.exp, data=sink_t, bias=neg_max)
rcp = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.reciprocal(dst=rcp, data=nl.add(den, sk))
p16 = nl.copy(p32, dtype=nl.bfloat16)
pt = nl.ndarray((PT, ktiles, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
for kb in nl.static_range(ktiles):
ptps = nl.ndarray((PT, hp), dtype=nl.bfloat16, buffer=nl.psum)
nisa.nc_transpose(
dst=ptps[0:PT, 0:heads], data=p16[0:heads, kb * PT : (kb + 1) * PT]
)
nisa.tensor_copy(
dst=pt[0:PT, kb, 0:heads], src=ptps[0:PT, 0:heads]
)
acc_o = nl.ndarray((heads, d), dtype=nl.float32, buffer=nl.psum)
for kb in nl.static_range(ktiles):
nisa.nc_matmul(
dst=acc_o,
stationary=pt[0:PT, kb, 0:heads],
moving=gn[0:PT, kb, 0:d],
accumulate=(kb > 0),
)
o16 = nl.ndarray((heads, d), dtype=nl.bfloat16, buffer=nl.sbuf)
nisa.tensor_scalar(
dst=o16, data=acc_o, op0=nl.multiply, operand0=rcp
)
nl.store(out[t * heads : t * heads + heads, 0:d], value=o16)
return out
Reference in New Issue View Git Blame Copy Permalink