KIRI on Hardware

// Processing one token through the Pulse atom (2 layers, 32-dim, 4 heads):

1. Embedding lookup         2 reads × 32 dims          = 64 ops
2. RMSNorm                  32 multiply + 32 add + sqrt  = ~70 ops
3. Attention Q,K,V          3 × (32×32 MAC)             = 3,072 ops
4. Score + softmax          4 heads × T scores           = ~128 ops
5. Weighted sum + O proj    32×32 MAC                    = 1,024 ops
6. Residual add             32 adds                       = 32 ops
7. MLP (expand + ReLU)      128×32 MAC + 128 relu         = 4,224 ops
8. MLP (compress)           32×128 MAC                    = 4,096 ops
9. Residual add             32 adds                       = 32 ops
10. ×2 layers              (steps 2-9) × 2              = ×2
11. lm_head + softmax       43×32 MAC + 43 exp           = 1,419 ops
────────────────────────────────────────────────────────
Total: ~27,000 operations per token, ~160,000 for a 6-token sequence
At 240 MHz (ESP32-S3): <1 ms inference time

// kiri_atom.h — Fixed-point transformer forward pass
// Weights are int16_t, scaled by 2^10 (10 fractional bits)

#include <stdint.h>

#define N_EMBD    32
#define N_HEAD    4
#define N_LAYER   2
#define VOCAB     43
#define BLOCK_SZ  16
#define HEAD_DIM  (N_EMBD / N_HEAD)  // 8
#define FRAC_BITS 10
#define SCALE     (1 << FRAC_BITS)    // 1024

// Weight matrices (stored in flash, loaded to SRAM on boot)
typedef struct {
    int16_t wte[VOCAB][N_EMBD];          // 43 × 32 = 1,376
    int16_t wpe[BLOCK_SZ][N_EMBD];      // 16 × 32 = 512
    int16_t lm_head[VOCAB][N_EMBD];     // 43 × 32 = 1,376
    // Per layer: wq, wk, wv, wo, f1, f2
    int16_t wq[N_LAYER][N_EMBD][N_EMBD];
    int16_t wk[N_LAYER][N_EMBD][N_EMBD];
    int16_t wv[N_LAYER][N_EMBD][N_EMBD];
    int16_t wo[N_LAYER][N_EMBD][N_EMBD];
    int16_t f1[N_LAYER][4*N_EMBD][N_EMBD];
    int16_t f2[N_LAYER][N_EMBD][4*N_EMBD];
} AtomWeights;

// Fixed-point linear layer: out[nout] = W[nout][nin] @ x[nin]
static void linear(int32_t *out, const int16_t *W,
                   const int32_t *x, int nout, int nin) {
    for (int i = 0; i < nout; i++) {
        int32_t acc = 0;
        for (int j = 0; j < nin; j++) {
            acc += (int32_t)W[i * nin + j] * (x[j] >> (FRAC_BITS/2));
        }
        out[i] = acc >> (FRAC_BITS/2);  // keep in Q10 range
    }
}

// ReLU: max(0, x)
static void relu(int32_t *x, int n) {
    for (int i = 0; i < n; i++)
        if (x[i] < 0) x[i] = 0;
}

// Score one observation: returns anomaly score (fixed-point)
int32_t atom_score(const AtomWeights *w, const uint8_t *tokens, int n_tokens) {
    int32_t x[N_EMBD], xr[N_EMBD], tmp[4*N_EMBD];
    int32_t total_score = 0;

    // KV cache for attention
    int32_t k_cache[N_LAYER][BLOCK_SZ][N_EMBD];
    int32_t v_cache[N_LAYER][BLOCK_SZ][N_EMBD];

    for (int pos = 0; pos < n_tokens - 1; pos++) {
        // Embedding: x = wte[token] + wpe[pos]
        for (int j = 0; j < N_EMBD; j++)
            x[j] = ((int32_t)w->wte[tokens[pos]][j]
                  + (int32_t)w->wpe[pos][j]) << (FRAC_BITS/2);

        // rmsnorm(x) — normalize
        rmsnorm(x, N_EMBD);

        // Transformer layers
        for (int li = 0; li < N_LAYER; li++) {
            memcpy(xr, x, sizeof(x));  // save for residual
            rmsnorm(x, N_EMBD);

            // Attention: Q*K^T/sqrt(d) -> softmax -> *V
            // ... (multi-head attention with KV cache) ...

            // MLP: expand -> relu -> compress
            linear(tmp, (int16_t*)w->f1[li], x, 4*N_EMBD, N_EMBD);
            relu(tmp, 4*N_EMBD);
            linear(x, (int16_t*)w->f2[li], tmp, N_EMBD, 4*N_EMBD);

            // Residual connection
            for (int j = 0; j < N_EMBD; j++) x[j] += xr[j];
        }

        // lm_head -> score
        int32_t logits[VOCAB];
        linear(logits, (int16_t*)w->lm_head, x, VOCAB, N_EMBD);

        // -log(softmax(logits)[target]) = surprise
        total_score += neg_log_softmax(logits, tokens[pos+1], VOCAB);
    }

    return total_score / (n_tokens - 1);
}

# Python: export weights as binary for ESP32
import struct, json

with open('pulse_weights.json') as f:
    data = json.load(f)

with open('pulse_weights.bin', 'wb') as f:
    for name in ['wte','wpe','lm_head','l0.wq','l0.wk', ...]:
        for row in data['weights'][name]:
            for val in row:
                # Quantize float to int16 (Q10 fixed-point)
                q = int(round(val * 1024))
                q = max(-32768, min(32767, q))
                f.write(struct.pack('<h', q))

# Flash to ESP32:
# esptool.py write_flash 0x100000 pulse_weights.bin