esp32: 2.4x faster on the ESP32-S3 (quick-start 2.52 s -> 1.06 s on a Cardputer)

Measured on an M5Stack Cardputer: 67 tokens 2514 -> 1056 ms, 126 tokens 5687 -> 2281 ms, 12 tokens 528 -> 245 ms; same heap. Bit-identical: 4-token PIE kernel with the tile loop in assembly, attention loop restructuring, lrintf-free rounding, LayerNorm/quantisation on both cores, 64-byte-aligned embedded model. Within float rounding (outputs move <= 0.0035): fc2 in 256-neuron groups, Cephes exp in the attention softmax, heads read each matrix once (u = W^T q). Host conformance builds with -ffp-contract=fast like the device GCC: max|dp| 0.0163, no pick or span flips.

README.md

@@ -178,7 +178,7 @@ const char* msg = "Book a table for 4 at an Italian place near the station on Fr
 td::answer(msg, qs, 2, a);                // a[0].pick == 3, msg[a[1].start, a[1].end) == "7:30"
 ```
 
-On the ESP32-S3, the PIE vector unit runs the 4-bit matmuls on both cores. On an M5Stack Cardputer the quick-start request takes 2.5 s. It has four questions and 67 tokens, and the answers match the table above. That timing comes from the ESP-IDF project in `esp32/examples/basic`. The Arduino sketch in `esp32/tinydecide/examples/Basic` builds, but it has not run on a device yet. Each folder's README has the details.
+On the ESP32-S3, the PIE vector unit runs the 4-bit matmuls on both cores. On an M5Stack Cardputer the quick-start request takes 1.1 s. It has four questions and 67 tokens, and the answers match the table above. That timing comes from the ESP-IDF project in `esp32/examples/basic`. The Arduino sketch in `esp32/tinydecide/examples/Basic` builds, but it has not run on a device yet. Each folder's README has the details.
 
 ## Files
 
@@ -228,7 +228,7 @@ Calibration error (ECE) is 0.02 to 0.15 on the tiers with true labels.
 - English only.
 - Every engine reads the first 127 tokens of a message and sets `truncated` when it cuts.
 - It answers best on short, everyday messages. Questions that need two-step reasoning, such as natural language inference, are close to chance.
-- On a Cardputer the quick-start request takes 2.5 s on both cores. [Pocket Notes](https://github.com/therezor/pocket-notes) asks one choice question about a short note and takes about 1.8 s. A request needs about 2.2 KB of heap per token. The bare ESP32 example has 371 KB free, enough for about 150 tokens. An app with a display and Wi-Fi has less.
+- On a Cardputer the quick-start request takes 1.1 s on both cores. [Pocket Notes](https://github.com/therezor/pocket-notes) asks one choice question about a short note and takes about 1.8 s with its own, older kernels. A request needs about 2.2 KB of heap per token. The bare ESP32 example has 371 KB free, enough for about 150 tokens. An app with a display and Wi-Fi has less.
 - The Decision Index numbers are self-scored with the kit. The board's maintainers have not verified them.
 - The 13.8M build's 6.31 is not a clean out-of-sample number. It had two board runs, and one data change came after the first run. The 10.4M build's 5.77 came from a single run.
 

esp32/README.md

@@ -17,7 +17,7 @@ the conformance test below runs it.
 | host build, 84 requests | same token ids, errors, picks and spans; largest probability difference about 0.02 |
 | ESP-IDF 5.5 build for ESP32-S3 with the model embedded (pioarduino 55.03.312) | compiles; the app image is 6.6 MB |
 | Arduino core 3.x build for ESP32-S3 with the model in a partition (same platform) | compiles; the app is 329 KB |
-| `examples/basic` on an M5Stack Cardputer (ESP32-S3, 8 MB flash, no PSRAM) | the four answers match the reference (restaurants 0.98 vs 0.97, 0.22, 0.52, "7:30"); 2.52 s per request; 371 KB heap free |
+| `examples/basic` on an M5Stack Cardputer (ESP32-S3, 8 MB flash, no PSRAM) | the four answers match the reference (restaurants 0.98 vs 0.97, 0.22, 0.52, "7:30"); 1.06 s per request; 371 KB heap free |
 | the Arduino sketch on a device | not run |
 
 The probabilities differ by up to about 0.02 because the PIE kernel needs activations rounded to

esp32/host/conformance.mjs

@@ -21,8 +21,9 @@ const exp = JSON.parse(readFileSync(conf + "expected.S768.json", "utf8"));
 
 const tmp = mkdtempSync(tmpdir() + "/td_host_"), bin = tmp + "/td_host";
 process.on("exit", () => rmSync(tmp, { recursive: true, force: true }));
-// TD_CXXFLAGS adds flags, e.g. TD_CXXFLAGS="-fsanitize=address,undefined -g"
-execFileSync("c++", ["-O2", "-std=c++17", "-Wall", "-Wextra", "-DTD_TEST", ...(process.env.TD_CXXFLAGS || "").split(" ").filter(Boolean), "-I" + here + "tinydecide/src",
+// -ffp-contract=fast: fuse a*b+c as the device GCC does, so results do not depend on how the host
+// compiler happens to contract each loop. TD_CXXFLAGS adds flags, e.g. TD_CXXFLAGS="-fsanitize=address,undefined -g"
+execFileSync("c++", ["-O2", "-std=c++17", "-ffp-contract=fast", "-Wall", "-Wextra", "-DTD_TEST", ...(process.env.TD_CXXFLAGS || "").split(" ").filter(Boolean), "-I" + here + "tinydecide/src",
   here + "host/td_host.cpp", here + "tinydecide/src/engine.cpp", "-o", bin], { stdio: "inherit" });
 
 const hex = (s) => (s === "" ? "-" : Buffer.from(s, "utf8").toString("hex"));

esp32/tinydecide/src/dot_q4_pie.S

@@ -71,4 +71,143 @@ dot_q4q8_pie:
     rfr     a2, f0
     retw.n
     .size   dot_q4q8_pie, . - dot_q4q8_pie
+
+// The same dot product for tiles of four tokens: each weight block is unpacked once and applied to
+// all four, and the loop over tiles runs here rather than in C.
+//
+//   void dot4_q4q8_pie(const uint8_t* nib,     // a2: 16B-aligned nibble plane
+//                      const uint8_t* scales,  // a3: bf16 weight scales
+//                      const int8_t*  xq,      // a4: token 0's int8 activations; 16B-aligned
+//                      const float*   xs,      // a5: token 0's block scales
+//                      int            nb,      // a6: number of 32-weight blocks
+//                      int            xqs,     // a7: bytes from one token's xq to the next;
+//                                              //     its xs are xqs / 32 floats further on
+//                      float*         out,     // stack: 4 * ntile sums, token order
+//                      int            ntile);  // stack: tiles of four tokens, >= 1
+//
+// Each token's sum is accumulated exactly as dot_q4q8_pie does it, so the results are identical.
+//
+// Frame: a1+0 tiles left, +4 nib, +8 scales, +12 nb, +16 xs stride (bytes), +20 xq of the tile,
+// +24 xs of the tile, +28 xqs.
+
+    .text
+    .align  4
+    .global dot4_q4q8_pie
+    .type   dot4_q4q8_pie, @function
+dot4_q4q8_pie:
+    entry   a1, 48
+    l32i    a15, a1, 48         // out (7th argument, on the caller's stack)
+    l32i    a8, a1, 52          // ntile (8th)
+    s32i    a8, a1, 0
+    s32i    a2, a1, 4
+    s32i    a3, a1, 8
+    s32i    a6, a1, 12
+    srai    a8, a7, 3           // xs stride in bytes: xqs / 32 floats
+    s32i    a8, a1, 16
+    s32i    a4, a1, 20
+    s32i    a5, a1, 24
+    s32i    a7, a1, 28
+    movi    a8, .Lconst
+    ee.vldbc.8 q7, a8           // q7 = 0x0F splat
+    addi    a8, a8, 1
+    ee.vldbc.8 q6, a8           // q6 = 8 splat
+    ssai    4
+.Ltile:
+    l32i    a4, a1, 20          // this tile's token pointers
+    l32i    a7, a1, 28
+    add     a9, a4, a7
+    add     a10, a9, a7
+    add     a11, a10, a7
+    l32i    a5, a1, 24
+    l32i    a8, a1, 16
+    add     a12, a5, a8
+    add     a13, a12, a8
+    add     a14, a13, a8
+    l32i    a2, a1, 4           // back to the start of the weight row
+    l32i    a3, a1, 8
+    l32i    a6, a1, 12
+    movi    a8, 0
+    wfr     f0, a8
+    wfr     f1, a8
+    wfr     f2, a8
+    wfr     f3, a8
+    // Per block: unpack the weights once, then token k's MACs overlap token k-1's scalar tail.
+    // Token k's sum is still acc_k += (float)isum_k * (wscale * xs_k[b]), as in dot_q4q8_pie.
+    loopnez a6, .Lloop4_end
+    ee.vld.128.ip q0, a2, 16    // 16 packed nibble bytes
+    l16ui       a8, a3, 0       // bf16 weight scale bits
+    lsip        f5, a5, 4       // xs[b] of tokens 0..3
+    lsip        f6, a12, 4
+    lsip        f7, a13, 4
+    lsip        f8, a14, 4
+    ee.vsr.32   q1, q0
+    ee.andq     q0, q0, q7      // low nibbles  -> q[0..15]
+    slli        a8, a8, 16
+    ee.andq     q1, q1, q7      // high nibbles -> q[16..31]
+    wfr         f4, a8          // wscale
+    ee.vsubs.s8 q0, q0, q6
+    ee.vsubs.s8 q1, q1, q6
+    mul.s       f5, f4, f5      // wscale * xscale, per token
+    mul.s       f6, f4, f6
+    mul.s       f7, f4, f7
+    mul.s       f8, f4, f8
+    addi        a3, a3, 2
+    // token 0
+    ee.vld.128.ip q2, a4, 16
+    ee.vld.128.ip q3, a4, 16
+    ee.zero.accx
+    ee.vmulas.s8.accx q0, q2
+    ee.vld.128.ip q4, a9, 16    // token 1's activations
+    ee.vmulas.s8.accx q1, q3
+    ee.vld.128.ip q5, a9, 16
+    rur.accx_0  a7
+    // token 1
+    ee.zero.accx
+    ee.vmulas.s8.accx q0, q4
+    ee.vld.128.ip q2, a10, 16   // token 2's
+    ee.vmulas.s8.accx q1, q5
+    ee.vld.128.ip q3, a10, 16
+    float.s     f9, a7, 0
+    rur.accx_0  a8
+    // token 2
+    ee.zero.accx
+    ee.vmulas.s8.accx q0, q2
+    ee.vld.128.ip q4, a11, 16   // token 3's
+    ee.vmulas.s8.accx q1, q3
+    ee.vld.128.ip q5, a11, 16
+    madd.s      f0, f9, f5
+    float.s     f10, a8, 0
+    rur.accx_0  a7
+    // token 3
+    ee.zero.accx
+    ee.vmulas.s8.accx q0, q4
+    ee.vmulas.s8.accx q1, q5
+    madd.s      f1, f10, f6
+    float.s     f11, a7, 0
+    rur.accx_0  a8
+    madd.s      f2, f11, f7
+    float.s     f12, a8, 0
+    madd.s      f3, f12, f8
+.Lloop4_end:
+    ssi     f0, a15, 0
+    ssi     f1, a15, 4
+    ssi     f2, a15, 8
+    ssi     f3, a15, 12
+    addi    a15, a15, 16
+    l32i    a7, a1, 28          // next tile: four tokens on
+    slli    a7, a7, 2
+    l32i    a8, a1, 20
+    add     a8, a8, a7
+    s32i    a8, a1, 20
+    l32i    a7, a1, 16
+    slli    a7, a7, 2
+    l32i    a8, a1, 24
+    add     a8, a8, a7
+    s32i    a8, a1, 24
+    l32i    a8, a1, 0
+    addi    a8, a8, -1
+    s32i    a8, a1, 0
+    bnez    a8, .Ltile
+    retw.n
+    .size   dot4_q4q8_pie, . - dot4_q4q8_pie
 #endif  // CONFIG_IDF_TARGET_ESP32S3

esp32/tinydecide/src/engine.cpp

@@ -2,9 +2,11 @@
 //
 // Memory plan ("streamed kernel"): one scratch area per call, ~2.2 KB per token. Each layer keeps
 // the fp32 residual x plus an int8 copy of the layer input; attention runs one head at a time with
-// W_o accumulated straight into x, and the FFN runs in 64-neuron chunks with fc2 accumulated into x.
-// Every weight row is read once per layer and applied to all tokens while it is hot in cache; rows
-// are split across both cores on dual-core chips.
+// W_o accumulated straight into x, and the FFN runs fc in 64-neuron chunks into 256-neuron int8
+// groups (in the then idle K and V buffers), each group's fc2 slice accumulated into x. Every
+// weight row is read once per layer and applied to all tokens while it is hot in cache, four
+// tokens per unpacked weight block on the ESP32-S3; rows are split across both cores on dual-core
+// chips.
 #include "tinydecide.h"
 #include "td_meta.h"
 #include "td_unicode.h"
@@ -276,14 +278,25 @@ static inline float dotq(const uint8_t* nib, const int8_t* xq, const uint8_t* sc
 #endif
 }
 
+// lrintf for |v| < 2^22 under round-to-nearest-even, without the libm call: adding 1.5 * 2^23
+// leaves no fraction bits, so the FPU rounds v to an integer exactly as lrintf does.
+static inline int roundEven(float v) { return (int)((v + 12582912.0f) - 12582912.0f); }
+
 // Symmetric int8 per 32-element block, one fp32 scale each (llama.cpp Q8_0).
+// No a*b+c fusion here: x * inv must be rounded to float before roundEven, as lrintf would see it.
+#if defined(__GNUC__) && !defined(__clang__)
+__attribute__((optimize("fp-contract=off")))
+#endif
 static void quant8(const float* x, int8_t* xq, float* xs, int n) {
+#if defined(__clang__)
+#pragma clang fp contract(off)
+#endif
   for (int b = 0; b < n / 32; b++) {
     float mx = 0.0f;
     for (int i = 0; i < 32; i++) { float a = fabsf(x[i]); if (a > mx) mx = a; }
     float s = mx / 127.0f, inv = s > 0 ? 1.0f / s : 0.0f;
     xs[b] = s;
-    for (int i = 0; i < 32; i++) xq[i] = (int8_t)lrintf(x[i] * inv);
+    for (int i = 0; i < 32; i++) xq[i] = (int8_t)roundEven(x[i] * inv);
     x += 32; xq += 32;
   }
 }
@@ -308,7 +321,22 @@ static void mmRows(const MM& a, int ra, int rb) {
     const uint8_t* sc  = g_m + a.w->sc + ((size_t)r * nbRow + a.blk0) * 2;
     const float b = a.bias ? a.bias[r] : 0.0f;
     float* o = a.out + (r - a.r0);
-    for (int t = 0; t < a.T; t++) {
+    int t = 0;
+#if defined(TD_PIE)
+    // tiles of four tokens, up to eight tiles per call: each weight block is unpacked once per
+    // tile (xss == xqs / 32)
+    float vt[32];
+    while (t + 4 <= a.T) {
+      const int nt = (a.T - t) / 4 < 8 ? (a.T - t) / 4 : 8;
+      dot4_q4q8_pie(nib, sc, a.xq + (size_t)t * a.xqs, a.xs + (size_t)t * a.xss, a.nbk, a.xqs, vt, nt);
+      for (int j = 0; j < 4 * nt; j++) {
+        const float v = vt[j] + b;
+        if (a.acc) o[(size_t)(t + j) * a.os] += v; else o[(size_t)(t + j) * a.os] = v;
+      }
+      t += 4 * nt;
+    }
+#endif
+    for (; t < a.T; t++) {
       float v = dotq(nib, a.xq + (size_t)t * a.xqs, sc, a.xs + (size_t)t * a.xss, a.nbk) + b;
       if (a.acc) o[(size_t)t * a.os] += v; else o[(size_t)t * a.os] = v;
     }
@@ -415,6 +443,38 @@ static void mvI8(const tdm::I8& w, const float* v, float* y) {
   }
 }
 
+// u = W^T v for an int8 head matrix: q . (W h) == u . h, so a head that scores many tokens against
+// one query reads W once instead of once per token. u has w.cols entries.
+static void mtvI8(const tdm::I8& w, const float* v, float* u) {
+  const int8_t* q = (const int8_t*)(g_m + w.q);
+  const float* s = F(w.sc);
+  for (int c = 0; c < w.cols; c++) u[c] = 0;
+  for (int r = 0; r < w.rows; r++) {
+    const float k = v[r] * s[r];
+    const int8_t* row = q + (size_t)r * w.cols;
+    for (int c = 0; c < w.cols; c++) u[c] += k * (float)row[c];
+  }
+}
+
+// exp(x) for the attention softmax, x <= 0: Cephes expf (within about 1 ulp of expf) at a tenth of
+// newlib's cost. Returns 0 below -87, where expf is denormal.
+static inline float expNeg(float x) {
+  if (x < -87.0f) return 0.0f;
+  const float n = (x * 1.44269504089f + 12582912.0f) - 12582912.0f;   // round(x / ln 2)
+  const float r = (x - n * 0.693359375f) - n * -2.12194440e-4f;
+  float p = 1.9875691500e-4f;
+  p = p * r + 1.3981999507e-3f;
+  p = p * r + 8.3334519073e-3f;
+  p = p * r + 4.1665795894e-2f;
+  p = p * r + 1.6666665459e-1f;
+  p = p * r + 5.0000001201e-1f;
+  p = p * r * r + r + 1.0f;
+  const uint32_t bits = (uint32_t)((int)n + 127) << 23;
+  float sc;
+  memcpy(&sc, &bits, 4);
+  return p * sc;
+}
+
 // ============================================================================
 //  Request layout (encodeRequest in tinydecide.js)
 // ============================================================================
@@ -516,15 +576,28 @@ static bool allocScratch(Scratch& m, int T) {
   return false;
 }
 
-// GELU over one FFN chunk for tokens [a, b), then int8 for the fc2 slice.
-struct GeluCtx { float* Hc; int8_t* hq; float* hs; };
+// GELU over one 64-neuron FFN chunk for tokens [a, b), then int8 into columns of hq (row stride
+// qs bytes, scales qs / 32 floats) for fc2.
+struct GeluCtx { float* Hc; int8_t* hq; float* hs; int qs; };
 static void geluRange(void* ctx, int a, int b) {
   const GeluCtx& g = *(const GeluCtx*)ctx;
   const int DHd = tdm::DHEAD;
   for (int t = a; t < b; t++) {
     float* hrow = g.Hc + (size_t)t * DHd;
     for (int c = 0; c < DHd; c++) hrow[c] = gelu(hrow[c]);
-    quant8(hrow, g.hq + (size_t)t * DHd, g.hs + (size_t)t * (DHd / 32), DHd);
+    quant8(hrow, g.hq + (size_t)t * g.qs, g.hs + (size_t)t * (g.qs / 32), DHd);
+  }
+}
+
+// Optional LayerNorm (w != nullptr), then int8 of x for the next matmuls, for tokens [a, b).
+struct NormCtx { float* x; int8_t* xq; float* xs; const float* w; const float* b; };
+static void normRange(void* ctx, int a, int b) {
+  const NormCtx& n = *(const NormCtx*)ctx;
+  const int D = tdm::D;
+  for (int t = a; t < b; t++) {
+    float* row = n.x + (size_t)t * D;
+    if (n.w) layernorm(row, D, n.w, n.b, tdm::LN_EPS);
+    quant8(row, n.xq + (size_t)t * D, n.xs + (size_t)t * (D / 32), D);
   }
 }
 
@@ -551,8 +624,24 @@ static void attnRange(void* ctx, int a, int b) {
     const float* q = c.Qh + (size_t)i * DHd;
     float mx = -1e30f;
     int n = 0;
-    for (const auto& rg : ranges)
-      for (int j = rg[0]; j < rg[1]; j++) {
+    for (const auto& rg : ranges) {
+      int j = rg[0];
+      // four keys at a time: four independent sums, each in the order of the one-key loop
+      for (; j + 4 <= rg[1]; j += 4) {
+        const float* k0 = c.Kh + (size_t)j * DHd;
+        const float *k1 = k0 + DHd, *k2 = k1 + DHd, *k3 = k2 + DHd;
+        float s0 = 0, s1 = 0, s2 = 0, s3 = 0;
+        for (int d = 0; d < DHd; d++) {
+          const float qd = q[d];
+          s0 += qd * k0[d]; s1 += qd * k1[d]; s2 += qd * k2[d]; s3 += qd * k3[d];
+        }
+        for (float s : {s0, s1, s2, s3}) {
+          s *= c.scale;
+          sc[n++] = s;
+          if (s > mx) mx = s;
+        }
+      }
+      for (; j < rg[1]; j++) {
         const float* kk = c.Kh + (size_t)j * DHd;
         float s = 0;
         for (int d = 0; d < DHd; d++) s += q[d] * kk[d];
@@ -560,17 +649,25 @@ static void attnRange(void* ctx, int a, int b) {
         sc[n++] = s;
         if (s > mx) mx = s;
       }
+    }
     float z = 0;
-    for (int j = 0; j < n; j++) { sc[j] = expf(sc[j] - mx); z += sc[j]; }
+    for (int j = 0; j < n; j++) { sc[j] = expNeg(sc[j] - mx); z += sc[j]; }
     const float iz = 1.0f / z;
-    for (int d = 0; d < DHd; d++) out[d] = 0;
-    n = 0;
-    for (const auto& rg : ranges)
-      for (int j = rg[0]; j < rg[1]; j++) {
-        const float p = sc[n++] * iz;
-        const float* v = c.Vh + (size_t)j * DHd;
-        for (int d = 0; d < DHd; d++) out[d] += p * v[d];
-      }
+    for (int j = 0; j < n; j++) sc[j] *= iz;
+    // out = sum_j p_j v_j, eight dimensions at a time so the sums stay in registers
+    for (int d0 = 0; d0 < DHd; d0 += 8) {
+      float o0 = 0, o1 = 0, o2 = 0, o3 = 0, o4 = 0, o5 = 0, o6 = 0, o7 = 0;
+      n = 0;
+      for (const auto& rg : ranges)
+        for (int j = rg[0]; j < rg[1]; j++) {
+          const float p = sc[n++];
+          const float* v = c.Vh + (size_t)j * DHd + d0;
+          o0 += p * v[0]; o1 += p * v[1]; o2 += p * v[2]; o3 += p * v[3];
+          o4 += p * v[4]; o5 += p * v[5]; o6 += p * v[6]; o7 += p * v[7];
+        }
+      out[d0] = o0; out[d0 + 1] = o1; out[d0 + 2] = o2; out[d0 + 3] = o3;
+      out[d0 + 4] = o4; out[d0 + 5] = o5; out[d0 + 6] = o6; out[d0 + 7] = o7;
+    }
     quant8(out, c.hq + (size_t)i * DHd, c.hs + (size_t)i * (DHd / 32), DHd);
   }
 }
@@ -615,7 +712,10 @@ static void encode(const Req& r, const Scratch& m) {
   const float scale = 1.0f / sqrtf((float)DHd);
   for (int l = 0; l < tdm::LAYERS; l++) {
     const tdm::Block& B = tdm::BLOCKS[l];
-    for (int t = 0; t < T; t++) quant8(x + (size_t)t * D, xq + (size_t)t * D, xs + (size_t)t * (D / 32), D);
+    if (l == 0) {
+      NormCtx nc{x, xq, xs, nullptr, nullptr};
+      parallel(normRange, &nc, T);
+    }
 
     // ---- attention, one head at a time; W_o slice accumulated into the residual
     for (int h = 0; h < tdm::HEADS; h++) {
@@ -627,18 +727,26 @@ static void encode(const Req& r, const Scratch& m) {
       parallel(attnRange, &ac, T);
       mm({&B.o, r0 / 32, DHd / 32, hq, DHd, hs, DHd / 32, T, x, D, h == 0 ? F(B.ob) : nullptr, true, 0}, D);
     }
-    for (int t = 0; t < T; t++) layernorm(x + (size_t)t * D, D, F(B.ln1w), F(B.ln1b), tdm::LN_EPS);
+    NormCtx n1{x, xq, xs, F(B.ln1w), F(B.ln1b)};
+    parallel(normRange, &n1, T);
 
-    // ---- FFN in 64-neuron chunks; fc2 slice accumulated into the residual
-    for (int t = 0; t < T; t++) quant8(x + (size_t)t * D, xq + (size_t)t * D, xs + (size_t)t * (D / 32), D);
+    // ---- FFN: fc in 64-neuron chunks through GELU into int8 groups of G neurons, held in the
+    // K and V buffers (free until the next layer); each group's fc2 slice is accumulated into x.
+    constexpr int G = 256;
+    static_assert(tdm::FFN % G == 0 && G <= tdm::DHEAD * 4, "an FFN group must fit the K buffer");
     float* Hc = Qh;
-    for (int c0 = 0; c0 < tdm::FFN; c0 += DHd) {
-      mm({&B.fc, 0, D / 32, xq, D, xs, D / 32, T, Hc, DHd, F(B.fcb), false, c0}, c0 + DHd);
-      GeluCtx gc{Hc, hq, hs};
-      parallel(geluRange, &gc, T);
-      mm({&B.fc2, c0 / 32, DHd / 32, hq, DHd, hs, DHd / 32, T, x, D, c0 == 0 ? F(B.fc2b) : nullptr, true, 0}, D);
+    int8_t* gq = (int8_t*)Kh;                               // [T, G] int8
+    float* gs = Vh;                                         // [T, G / 32]
+    for (int g0 = 0; g0 < tdm::FFN; g0 += G) {
+      for (int c0 = g0; c0 < g0 + G; c0 += DHd) {
+        mm({&B.fc, 0, D / 32, xq, D, xs, D / 32, T, Hc, DHd, F(B.fcb), false, c0}, c0 + DHd);
+        GeluCtx gc{Hc, gq + (c0 - g0), gs + (c0 - g0) / 32, G};
+        parallel(geluRange, &gc, T);
+      }
+      mm({&B.fc2, g0 / 32, G / 32, gq, G, gs, G / 32, T, x, D, g0 == 0 ? F(B.fc2b) : nullptr, true, 0}, D);
     }
-    for (int t = 0; t < T; t++) layernorm(x + (size_t)t * D, D, F(B.ln2w), F(B.ln2b), tdm::LN_EPS);
+    NormCtx n2{x, xq, xs, F(B.ln2w), F(B.ln2b)};   // the int8 copy is the next layer's input
+    parallel(normRange, &n2, T);
   }
 }
 
@@ -688,16 +796,16 @@ static float protoTerm(const Protos& p, const float* qvec, const float* c, int i
 static void headChoice(const Req& r, const float* x, int k, const Question& q, Answer& out) {
   const int sel = q.type == SCORE ? 1 : 0, n = q.n_options, DH = tdm::DH;
   const float Tt = tdm::TEMP[sel ? 2 : 0], s = expf(F(tdm::H_SCALE)[sel]);
-  float h[tdm::D], qa[tdm::DH], ov[tdm::DH], L[MAX_OPTIONS];
+  float h[tdm::D], qa[tdm::DH], u[tdm::D], L[MAX_OPTIONS];
   headRow(x, r.ans[k], h);
   mvI8(tdm::H_A[sel], h, qa);
   mvI8(tdm::H_P, h, out.qvec);
+  mtvI8(tdm::H_O[sel], qa, u);                     // qa . (H_O h) == u . h
   float cbuf[QDIM];
   const float* c = q.protos ? centreOf(*q.protos, n, cbuf) : nullptr;
   for (int i = 0; i < n; i++) {
     headRow(x, r.opt[k][i], h);
-    mvI8(tdm::H_O[sel], h, ov);
-    out.z0[i] = s * dotf(qa, ov, DH) / sqrtf((float)DH) / Tt;
+    out.z0[i] = s * dotf(u, h, tdm::D) / sqrtf((float)DH) / Tt;
     L[i] = out.z0[i] + (q.bias ? q.bias[i] : 0.0f);
     if (q.protos) L[i] += protoTerm(*q.protos, out.qvec, c, i) / Tt;
   }
@@ -742,18 +850,18 @@ static void headNoul(const Req& r, const float* x, int k, const Question& q, Ans
 static void headSpan(const Req& r, const float* x, int k, const char* state, Answer& out) {
   const int S = r.S - 1, DH = tdm::DH;           // state text tokens at rows 1..S
   const float Tt = tdm::TEMP[3], isq = 1.0f / sqrtf((float)DH);
-  float h[tdm::D], hAns[tdm::D], qS[tdm::DH], kk[tdm::DH], qe[tdm::DH];
+  float h[tdm::D], hAns[tdm::D], qS[tdm::DH], qe[tdm::DH], u[tdm::D];
   headRow(x, r.ans[k], hAns);
   mvI8(tdm::H_P, hAns, out.qvec);
   mvI8(tdm::H_SQ, hAns, qS);
+  mtvI8(tdm::H_SK, qS, u);                         // qS . (H_SK h) == u . h
   // start logits: null, then each state token; log-softmax at temperature Tt
   const float zNull = dotf(F(tdm::H_SNULL_W), hAns, tdm::D) + F(tdm::H_SNULL_B)[0];
   static float zs[tdm::TS_MAX];
   float mx = zNull;
   for (int t = 0; t < S; t++) {
     headRow(x, 1 + t, h);
-    mvI8(tdm::H_SK, h, kk);
-    zs[t] = dotf(qS, kk, DH) * isq;
+    zs[t] = dotf(u, h, tdm::D) * isq;
     if (zs[t] > mx) mx = zs[t];
   }
   float Z = expf((zNull - mx) / Tt);
@@ -774,12 +882,12 @@ static void headSpan(const Req& r, const float* x, int k, const char* state, Ans
   headRow(x, 1 + best, h);
   mvI8(tdm::H_ES, h, es);
   for (int i = 0; i < DH; i++) qe[i] += es[i];
+  mtvI8(tdm::H_EK, qe, u);                         // qe . (H_EK h) == u . h
   const int e1 = best + tdm::SPAN_MAX < S ? best + tdm::SPAN_MAX : S;
   float en[tdm::SPAN_MAX], emx = -1e30f;
   for (int t = best; t < e1; t++) {
     headRow(x, 1 + t, h);
-    mvI8(tdm::H_EK, h, kk);
-    en[t - best] = dotf(qe, kk, DH) * isq;
+    en[t - best] = dotf(u, h, tdm::D) * isq;
     if (en[t - best] > emx) emx = en[t - best];
   }
   float EZ = 0;

esp32/tinydecide/src/model_data.S

@@ -1,10 +1,12 @@
 // Embeds model.bin and vocab.bin in flash (.rodata, MMU-mapped) when the ESP-IDF component is built
 // with CONFIG_TINYDECIDE_EMBED_MODEL (CMakeLists.txt defines TINYDECIDE_EMBED and the include path).
-// The PIE kernel loads 128-bit words, so the model starts on a 16-byte boundary; all Q4 offsets
-// inside it are multiples of 16 (checked by tools/pack.mjs).
+// The PIE kernel loads 128-bit words, so the model must start on a 16-byte boundary; all Q4 offsets
+// inside it are multiples of 16 (checked by tools/pack.mjs). It starts on a 64-byte boundary so the
+// weight rows, at 64-byte offsets in model.bin, line up with the flash cache lines: unaligned, each
+// row touches one more line and the matmuls run up to ~10% slower.
 #if defined(TINYDECIDE_EMBED)
     .section .rodata.td_model, "a"
-    .balign 16
+    .balign 64
     .global td_model_bin
 td_model_bin:
     .incbin "model.bin"

esp32/tinydecide/src/tinydecide.h

@@ -117,4 +117,7 @@ int requestTokens(const Question* qs, int nq, int state_tokens);
 // Q4_0 row-planar dot product against Q8_0 activations (dot_q4_pie.S, ESP32-S3 only).
 extern "C" float dot_q4q8_pie(const uint8_t* nib, const int8_t* xq, const uint8_t* scales,
                               const float* xs, int nb);
+// The same for ntile * 4 tokens xqs bytes apart (their scales xqs / 32 floats apart); out[4 * ntile].
+extern "C" void dot4_q4q8_pie(const uint8_t* nib, const uint8_t* scales, const int8_t* xq, const float* xs,
+                              int nb, int xqs, float* out, int ntile);
 #endif