# SPDX-License-Identifier: Apache-2.0 # From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/hc_mixes.py, # commit 30d54d1351c80728fe5685478fac838b44c53d43), unmodified. # Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE. """The hyper-connection mixing coefficients — the stream statistic, the three nonlinearities and the 20-iteration Sinkhorn — as one NKI kernel. rms one statistic per token over the whole flattened `hc * dim` stream pre sigmoid(raw * scale0 + base0) + eps post 2 * sigmoid(raw * scale1 + base1) comb softmax(raw * scale2 + base2) + eps, then Sinkhorn AssertionError: ap() pattern has invalid partition stride. Partition step 1 must equal tensor free dimension size 3. A freshly `nl.load`ed `(1, n)` tile broadcast to `(rows, n)` is fine, which is why every coefficient here is loaded from its own HBM slice rather than sliced out of one staged tile. """ import nki import nki.isa as nisa import nki.language as nl #: Partition width. PT = 128 #: `hc_sinkhorn_iters`. Not a convergence tolerance — the reference runs exactly #: this many passes, so matching the count matters bit-for-bit. SINKHORN_ITERS = 20 #: `hc_eps`. HC_EPS = 1e-06 #: `rms_norm_eps`. NORM_EPS = 1e-20 def _hc_from_width(width): """`width = (2 + hc) * hc` inverted. 24 -> 4.""" hc = 1 while (2 + hc) * hc < width: hc += 1 assert (2 + hc) * hc == width, f"width={width} is not (2 + hc) * hc" return hc def _sqchunk(fd): """Free-axis chunk for the sum of squares: a divisor of the stream width, small enough that one fp32 copy of it is a few KiB per partition. """ for cand in (2048, 1024, 512, 256, 128): if fd % cand == 0: return cand return fd def _rms_recip(flat, r0, nr, fd): """`1 / sqrt(mean(flat^2) + eps)` for each of `nr` rows, `[nr, 1]` fp32.""" ch = _sqchunk(fd) total = None for c in range(fd // ch): x = nl.load(flat[r0 : r0 + nr, c * ch : (c + 1) * ch], dtype=nl.float32) part = nl.sum(nl.square(x), axis=1, keepdims=True) total = part if total is None else nl.add(total, part) m = nl.add(nl.multiply(total, 1.0 / fd), NORM_EPS) y = nl.rsqrt(m) corr = nl.add(nl.multiply(nl.multiply(m, nl.multiply(y, y)), -0.5), 1.5) return nl.multiply(y, corr) @nki.jit def kernel(raw, flat, base, scale): """See the module docstring. Returns `(pre, post, comb)`.""" rows, width = raw.shape frows, fd = flat.shape hc = _hc_from_width(width) assert frows == rows, f"flat has {frows} rows, raw has {rows}" assert fd % hc == 0, f"stream width {fd} is not a multiple of hc={hc}" assert base.shape == (1, width), f"base {base.shape} wants (1, {width})" assert scale.shape == (1, 3), f"scale {scale.shape} wants (1, 3)" pre = nl.ndarray((rows, hc), dtype=nl.float32, buffer=nl.shared_hbm) post = nl.ndarray((rows, hc), dtype=nl.float32, buffer=nl.shared_hbm) comb = nl.ndarray((rows, hc * hc), dtype=nl.float32, buffer=nl.shared_hbm) def body(r0, nr): rs = _rms_recip(flat, r0, nr, fd) rawt = nl.load(raw[r0 : r0 + nr, 0:width], dtype=nl.float32) mix = nl.ndarray((nr, width), dtype=nl.float32, buffer=nl.sbuf) nisa.tensor_scalar(dst=mix, data=rawt, op0=nl.multiply, operand0=rs) s0 = nl.broadcast_to(nl.load(scale[0:1, 0:1]), (nr, 1)) s1 = nl.broadcast_to(nl.load(scale[0:1, 1:2]), (nr, 1)) s2 = nl.broadcast_to(nl.load(scale[0:1, 2:3]), (nr, 1)) b0 = nl.broadcast_to(nl.load(base[0:1, 0:hc]), (nr, hc)) b1 = nl.broadcast_to(nl.load(base[0:1, hc : 2 * hc]), (nr, hc)) b2 = nl.broadcast_to( nl.load(base[0:1, 2 * hc : 2 * hc + hc * hc]), (nr, hc * hc) ) t0 = nl.ndarray((nr, hc), dtype=nl.float32, buffer=nl.sbuf) nisa.scalar_tensor_tensor( dst=t0, data=mix[0:nr, 0:hc], op0=nl.multiply, operand0=s0, op1=nl.add, operand1=b0, ) nl.store(pre[r0 : r0 + nr, 0:hc], value=nl.add(nl.sigmoid(t0), HC_EPS)) t1 = nl.ndarray((nr, hc), dtype=nl.float32, buffer=nl.sbuf) nisa.scalar_tensor_tensor( dst=t1, data=mix[0:nr, hc : 2 * hc], op0=nl.multiply, operand0=s1, op1=nl.add, operand1=b1, ) nl.store(post[r0 : r0 + nr, 0:hc], value=nl.multiply(nl.sigmoid(t1), 2.0)) t2 = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf) nisa.scalar_tensor_tensor( dst=t2, data=mix[0:nr, 2 * hc : 2 * hc + hc * hc], op0=nl.multiply, operand0=s2, op1=nl.add, operand1=b2, ) def rowsum_div(src, dst): """Divide each `hc x hc` row by its own sum (the reference's `dim=-1` normalization). The divisor is one value per token, so it rides the per-partition operand and needs no broadcast. """ for i in range(hc): lo, hi = i * hc, (i + 1) * hc s = nl.sum(src[0:nr, lo:hi], axis=1, keepdims=True) nisa.tensor_scalar( dst=dst[0:nr, lo:hi], data=src[0:nr, lo:hi], op0=nl.multiply, operand0=nl.reciprocal(nl.add(s, HC_EPS)), ) def colsum_div(src, dst): """Divide each `hc x hc` column by its own sum (the reference's `dim=-2` normalization). The sum of the row blocks *is* the divisor laid out the way each block wants it. """ s = nl.add(src[0:nr, 0:hc], src[0:nr, hc : 2 * hc]) for i in range(2, hc): s = nl.add(s, src[0:nr, i * hc : (i + 1) * hc]) d = nl.reciprocal(nl.add(s, HC_EPS)) for i in range(hc): lo, hi = i * hc, (i + 1) * hc nisa.tensor_tensor( dst=dst[0:nr, lo:hi], data1=src[0:nr, lo:hi], data2=d, op=nl.multiply, ) a = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf) b = nl.ndarray((nr, hc * hc), dtype=nl.float32, buffer=nl.sbuf) for i in range(hc): lo, hi = i * hc, (i + 1) * hc mx = nl.max(t2[0:nr, lo:hi], axis=1, keepdims=True) e = nl.exp( nl.subtract(t2[0:nr, lo:hi], nl.broadcast_to(mx, (nr, hc))) ) s = nl.sum(e, axis=1, keepdims=True) nisa.tensor_scalar( dst=a[0:nr, lo:hi], data=e, op0=nl.multiply, operand0=nl.reciprocal(s), op1=nl.add, operand1=HC_EPS, ) colsum_div(a, b) for _ in range(SINKHORN_ITERS - 1): rowsum_div(b, a) colsum_div(a, b) nl.store(comb[r0 : r0 + nr, 0 : hc * hc], value=b[0:nr, 0 : hc * hc]) full, tail = rows // PT, rows % PT for t in nl.affine_range(full): body(t * PT, PT) if tail: body(full * PT, tail) return pre, post, comb