graph_ir_ops.h

Go to the documentation of this file.

#pragma once
// CPU kernels for the IR ops of the gated and attention-residual families
// (paperkiln ir.hpp at 176b425: flex highway/plain, attnres, srd's router):
// sub, relu, sigmoid, affine, mul_col, rms_row and depth_attend. Each one is
// the named paperkiln CPU op (src/ops.cpp, nn::AttnResStack) in its order of
// operations, every product and sum rounded to f32 on its own: this file is
// compiled with -ffp-contract=off (CMakeLists.txt), so no a * b + c becomes a
// fused multiply-add. Kept free of ember's other headers so the op-level gate
// can build it alone. Rows are [T, width] row-major.
#include <cstddef>
#include <vector>
 
namespace ember_ir {
namespace cpu {
 
// a - b elementwise. EXACT.
void sub(const float* a, const float* b, float* out, size_t n);
// max(0, x) as std::max(0.0f, x) (ops::relu: -0 and NaN give +0). EXACT.
void relu(const float* x, float* out, size_t n);
// 1 / (1 + exp(-x)) (ops::sigmoid). ROUNDING.
void sigmoid(const float* x, float* out, size_t n);
// (x * scale) + shift: the product rounded, then the sum (ops::scale, then
// ops::add_scalar). EXACT.
void affine(const float* x, float* out, size_t n, float scale, float shift);
// x [T, w] times the column c [T, 1]: out[t, j] = x[t, j] * c[t]. EXACT.
void mul_col(const float* x, const float* c, float* out, int T, int w);
// Per feature: out[t, j] = (x[t, j] * g[j]) + b[j], the product rounded, then
// the sum; b null: the product alone (ops `scale`, a gain without a [d, d]
// diagonal linear). EXACT.
void scale(const float* x, const float* g, const float* b, float* out, int T, int w);
// [T, w] -> [T, 1]: sqrt(ss / w + eps), ss = sum of x[t, j]^2 in ascending j
// (ops::rms_row). ROUNDING.
void rms_row(const float* x, float* out, int T, int w, float eps);
// Attention over depth (nn::AttnResStack::depth_attend, block form): sources
// s_0..s_{S-1}, each [T, d]; w the pseudo-query [d]. S = 1: out = s_0 and w
// is not read (may be null). S >= 2, per row: r_i = 1 / sqrt(ss_i / d + eps)
// (unit-weight rmsnorm), l_i = sum_j (s_i[j] * r_i) * w[j] in ascending j,
// a = max-shifted softmax over i (z summed in ascending i, a_i = e_i / z),
// out = s_0 a_0 + s_1 a_1 + ... in ascending i. ROUNDING.
void depth_attend(const std::vector<const float*>& sources, const float* w, float* out, int T, int d, float eps);
// Mixture-of-Recursions token-choice routing (ir.hpp depth_route) from one token's router logits
// l [R]: depth = 1 + the index of the largest logit (ties to the lower index); gate = p[depth - 1],
// p the max-shifted softmax (e_j = exp(l_j - max), z summed in ascending j, p_j = e_j / z).
// EXACT given the logits for the depth; ROUNDING for the gate.
void mor_route(const float* l, int R, int& depth, float& gate);
// MoR gate "score" after a recursion step: x = x_in + g (x - x_in) per element, the difference
// rounded, then the product, then the sum. EXACT.
void mor_gate(const float* x_in, float* x, size_t n, float g);
// Multi-head latent attention (IR_SPEC 5.20, revision 1.5; DeepSeek-V2's MLA in the absorbed form
// llama.cpp converts to). qn [T, H * dn] and qr [T, H * dr] (head h at columns h * dn and h * dr);
// kv the slot's rows [c | kr], r + dr floats each (row index as in `rows`); rows[t] the kv rows query t
// sees, in ascending position (section 5.15's masking); Kb [H, r, dn] (ne {dn, r, H}) and Vb [H, dv, r]
// (ne {r, dv, H}); out [T, H * dv]. Per head, in the reference's order: qa_i = sum_k Kb[h][i][k] qn_k
// (ascending k); s_j = ((qa . c_j) + (qr_h . kr_j)) * scale, each dot product accumulated in ascending
// order and rounded, then the sum, then the product; a_j = e_j / z with e_j = exp(s_j - max) and z summed
// in ascending j; u_i = sum_j a_j c_j[i] (ascending j); out_o = sum_i Vb[h][o][i] u_i. ROUNDING.
void mla_attention(const float* qn, const float* qr, const float* kv, float* out, int T,
                   const std::vector<std::vector<int>>& rows, int H, int r, int dn, int dr, int dv, const float* Kb,
                   const float* Vb, float scale);
// split with groups g (IR_SPEC 5.5, revision 1.5): x [T, g * sum(sizes)] is g consecutive groups of
// sum(sizes) columns, and output k is the k-th part of every group, the groups in order ([T, g * sizes[k]]).
// g = 1 is the plain split. Copies. EXACT.
void split_groups(const float* x, int T, const std::vector<int>& sizes, int g, const std::vector<float*>& outs);
 
}  // namespace cpu
}  // namespace ember_ir