@seanhogg/builderforce-memory-engine 2026.7.14 → 2026.9.27

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (186) hide show
  1. package/README.md +14 -0
  2. package/dist/bench/adaptation.d.ts +57 -0
  3. package/dist/bench/adaptation.d.ts.map +1 -0
  4. package/dist/bench/adaptation.js +167 -0
  5. package/dist/bench/adaptation.js.map +1 -0
  6. package/dist/bench/harness.d.ts +2 -0
  7. package/dist/bench/harness.d.ts.map +1 -1
  8. package/dist/bench/harness.js +1 -1
  9. package/dist/bench/harness.js.map +1 -1
  10. package/dist/bench/index.d.ts +5 -3
  11. package/dist/bench/index.d.ts.map +1 -1
  12. package/dist/bench/index.js +4 -2
  13. package/dist/bench/index.js.map +1 -1
  14. package/dist/bench/types.d.ts +88 -0
  15. package/dist/bench/types.d.ts.map +1 -1
  16. package/dist/evermind/adapt.d.ts +40 -0
  17. package/dist/evermind/adapt.d.ts.map +1 -0
  18. package/dist/evermind/adapt.js +25 -0
  19. package/dist/evermind/adapt.js.map +1 -0
  20. package/dist/evermind/delta_wire.d.ts +43 -0
  21. package/dist/evermind/delta_wire.d.ts.map +1 -0
  22. package/dist/evermind/delta_wire.js +94 -0
  23. package/dist/evermind/delta_wire.js.map +1 -0
  24. package/dist/evermind/eval.d.ts +31 -0
  25. package/dist/evermind/eval.d.ts.map +1 -0
  26. package/dist/evermind/eval.js +36 -0
  27. package/dist/evermind/eval.js.map +1 -0
  28. package/dist/evermind/index.d.ts +16 -0
  29. package/dist/evermind/index.d.ts.map +1 -0
  30. package/dist/evermind/index.js +11 -0
  31. package/dist/evermind/index.js.map +1 -0
  32. package/dist/evermind/merge.d.ts +21 -0
  33. package/dist/evermind/merge.d.ts.map +1 -0
  34. package/dist/evermind/merge.js +79 -0
  35. package/dist/evermind/merge.js.map +1 -0
  36. package/dist/{moe/moe_package.d.ts → evermind/package.d.ts} +55 -5
  37. package/dist/evermind/package.d.ts.map +1 -0
  38. package/dist/{moe/moe_package.js → evermind/package.js} +118 -19
  39. package/dist/evermind/package.js.map +1 -0
  40. package/dist/import/foreign/adapter.d.ts +59 -0
  41. package/dist/import/foreign/adapter.d.ts.map +1 -0
  42. package/dist/import/foreign/adapter.js +76 -0
  43. package/dist/import/foreign/adapter.js.map +1 -0
  44. package/dist/import/foreign/apply.d.ts +23 -0
  45. package/dist/import/foreign/apply.d.ts.map +1 -0
  46. package/dist/import/foreign/apply.js +43 -0
  47. package/dist/import/foreign/apply.js.map +1 -0
  48. package/dist/import/foreign/codestral_mamba.d.ts +57 -0
  49. package/dist/import/foreign/codestral_mamba.d.ts.map +1 -0
  50. package/dist/import/foreign/codestral_mamba.js +174 -0
  51. package/dist/import/foreign/codestral_mamba.js.map +1 -0
  52. package/dist/import/foreign/falcon_mamba.d.ts +51 -0
  53. package/dist/import/foreign/falcon_mamba.d.ts.map +1 -0
  54. package/dist/import/foreign/falcon_mamba.js +131 -0
  55. package/dist/import/foreign/falcon_mamba.js.map +1 -0
  56. package/dist/import/foreign/index.d.ts +21 -0
  57. package/dist/import/foreign/index.d.ts.map +1 -0
  58. package/dist/import/foreign/index.js +18 -0
  59. package/dist/import/foreign/index.js.map +1 -0
  60. package/dist/import/foreign/plan.d.ts +122 -0
  61. package/dist/import/foreign/plan.d.ts.map +1 -0
  62. package/dist/import/foreign/plan.js +171 -0
  63. package/dist/import/foreign/plan.js.map +1 -0
  64. package/dist/import/foreign/registry.d.ts +35 -0
  65. package/dist/import/foreign/registry.d.ts.map +1 -0
  66. package/dist/import/foreign/registry.js +89 -0
  67. package/dist/import/foreign/registry.js.map +1 -0
  68. package/dist/import/index.d.ts +7 -0
  69. package/dist/import/index.d.ts.map +1 -1
  70. package/dist/import/index.js +7 -0
  71. package/dist/import/index.js.map +1 -1
  72. package/dist/index.d.ts +25 -12
  73. package/dist/index.d.ts.map +1 -1
  74. package/dist/index.js +21 -9
  75. package/dist/index.js.map +1 -1
  76. package/dist/kernels/selective_scan.d.ts.map +1 -1
  77. package/dist/kernels/selective_scan.js +15 -9
  78. package/dist/kernels/selective_scan.js.map +1 -1
  79. package/dist/kernels/slice.d.ts +26 -0
  80. package/dist/kernels/slice.d.ts.map +1 -0
  81. package/dist/kernels/slice.js +52 -0
  82. package/dist/kernels/slice.js.map +1 -0
  83. package/dist/kernels/weight_update.d.ts.map +1 -1
  84. package/dist/kernels/weight_update.js +30 -4
  85. package/dist/kernels/weight_update.js.map +1 -1
  86. package/dist/lm/evermind_lm.d.ts +110 -1
  87. package/dist/lm/evermind_lm.d.ts.map +1 -1
  88. package/dist/lm/evermind_lm.js +220 -35
  89. package/dist/lm/evermind_lm.js.map +1 -1
  90. package/dist/lm/index.d.ts +4 -2
  91. package/dist/lm/index.d.ts.map +1 -1
  92. package/dist/lm/index.js +3 -1
  93. package/dist/lm/index.js.map +1 -1
  94. package/dist/lm/text_embedder.d.ts +55 -0
  95. package/dist/lm/text_embedder.d.ts.map +1 -0
  96. package/dist/lm/text_embedder.js +62 -0
  97. package/dist/lm/text_embedder.js.map +1 -0
  98. package/dist/model/attention_block.d.ts.map +1 -1
  99. package/dist/model/attention_block.js +9 -5
  100. package/dist/model/attention_block.js.map +1 -1
  101. package/dist/model/mamba1_block.d.ts.map +1 -1
  102. package/dist/model/mamba1_block.js +28 -12
  103. package/dist/model/mamba1_block.js.map +1 -1
  104. package/dist/model/mamba2_block.d.ts.map +1 -1
  105. package/dist/model/mamba2_block.js +14 -9
  106. package/dist/model/mamba2_block.js.map +1 -1
  107. package/dist/model/mamba3_block.d.ts.map +1 -1
  108. package/dist/model/mamba3_block.js +14 -9
  109. package/dist/model/mamba3_block.js.map +1 -1
  110. package/dist/moe/index.d.ts +0 -2
  111. package/dist/moe/index.d.ts.map +1 -1
  112. package/dist/moe/index.js +0 -1
  113. package/dist/moe/index.js.map +1 -1
  114. package/dist/optim/adamw.d.ts +32 -0
  115. package/dist/optim/adamw.d.ts.map +1 -1
  116. package/dist/optim/adamw.js +34 -0
  117. package/dist/optim/adamw.js.map +1 -1
  118. package/dist/tokenizer/bpe.d.ts +28 -0
  119. package/dist/tokenizer/bpe.d.ts.map +1 -1
  120. package/dist/tokenizer/bpe.js +47 -0
  121. package/dist/tokenizer/bpe.js.map +1 -1
  122. package/dist/training/cpu_ops.d.ts +44 -0
  123. package/dist/training/cpu_ops.d.ts.map +1 -0
  124. package/dist/training/cpu_ops.js +111 -0
  125. package/dist/training/cpu_ops.js.map +1 -0
  126. package/dist/training/mamba1_cpu.d.ts +109 -0
  127. package/dist/training/mamba1_cpu.d.ts.map +1 -0
  128. package/dist/training/mamba1_cpu.js +316 -0
  129. package/dist/training/mamba1_cpu.js.map +1 -0
  130. package/dist/training/model_cpu.d.ts +100 -0
  131. package/dist/training/model_cpu.d.ts.map +1 -0
  132. package/dist/training/model_cpu.js +252 -0
  133. package/dist/training/model_cpu.js.map +1 -0
  134. package/dist/training/trainer.d.ts +156 -16
  135. package/dist/training/trainer.d.ts.map +1 -1
  136. package/dist/training/trainer.js +327 -101
  137. package/dist/training/trainer.js.map +1 -1
  138. package/dist/utils/base64.d.ts +12 -0
  139. package/dist/utils/base64.d.ts.map +1 -0
  140. package/dist/utils/base64.js +27 -0
  141. package/dist/utils/base64.js.map +1 -0
  142. package/package.json +3 -1
  143. package/src/bench/adaptation.ts +180 -0
  144. package/src/bench/harness.ts +1 -1
  145. package/src/bench/index.ts +13 -2
  146. package/src/bench/types.ts +93 -0
  147. package/src/evermind/adapt.ts +61 -0
  148. package/src/evermind/delta_wire.ts +114 -0
  149. package/src/evermind/eval.ts +62 -0
  150. package/src/evermind/index.ts +33 -0
  151. package/src/evermind/merge.ts +106 -0
  152. package/src/{moe/moe_package.ts → evermind/package.ts} +148 -17
  153. package/src/import/foreign/adapter.ts +106 -0
  154. package/src/import/foreign/apply.ts +54 -0
  155. package/src/import/foreign/codestral_mamba.ts +224 -0
  156. package/src/import/foreign/falcon_mamba.ts +173 -0
  157. package/src/import/foreign/index.ts +37 -0
  158. package/src/import/foreign/plan.ts +294 -0
  159. package/src/import/foreign/registry.ts +126 -0
  160. package/src/import/index.ts +30 -0
  161. package/src/index.ts +98 -13
  162. package/src/kernels/selective_scan.ts +15 -9
  163. package/src/kernels/slice.ts +64 -0
  164. package/src/kernels/weight_update.ts +30 -4
  165. package/src/lm/evermind_lm.ts +244 -34
  166. package/src/lm/index.ts +6 -2
  167. package/src/lm/text_embedder.ts +84 -0
  168. package/src/model/attention_block.ts +9 -5
  169. package/src/model/mamba1_block.ts +28 -12
  170. package/src/model/mamba2_block.ts +14 -9
  171. package/src/model/mamba3_block.ts +14 -9
  172. package/src/moe/index.ts +0 -9
  173. package/src/optim/adamw.ts +60 -0
  174. package/src/tokenizer/bpe.ts +55 -0
  175. package/src/training/cpu_ops.ts +127 -0
  176. package/src/training/mamba1_cpu.ts +421 -0
  177. package/src/training/model_cpu.ts +360 -0
  178. package/src/training/trainer.ts +397 -154
  179. package/src/utils/base64.ts +28 -0
  180. package/dist/model/mamba_block.d.ts +0 -64
  181. package/dist/model/mamba_block.d.ts.map +0 -1
  182. package/dist/model/mamba_block.js +0 -303
  183. package/dist/model/mamba_block.js.map +0 -1
  184. package/dist/moe/moe_package.d.ts.map +0 -1
  185. package/dist/moe/moe_package.js.map +0 -1
  186. package/src/model/mamba_block.ts +0 -401
@@ -0,0 +1,64 @@
1
+ /**
2
+ * slice.ts — strided COLUMN slice of a row-major (rows, stride) tensor.
3
+ *
4
+ * Every block in this engine projects into one fused tensor and then splits it:
5
+ * Mamba-1 into (x, z) and (Δ, B, C), Mamba-2/3 into (x, Δ) and (x, B, C),
6
+ * attention into (Q, K, V). `LINEAR_FORWARD_WGSL` writes its output `(M, N)`
7
+ * ROW-major, so each of those parts is a contiguous range of COLUMNS **within
8
+ * every row** — not a contiguous range of the buffer.
9
+ *
10
+ * Splitting them with `copyBufferToBuffer` therefore took whole ROWS instead: it
11
+ * is only the right answer when `M == 1` (batch × seq_len), and for any real
12
+ * sequence it silently handed the block the wrong tensor. This kernel is the one
13
+ * correct implementation; every block uses it.
14
+ */
15
+
16
+ import { createBindGroup, createUniformBuffer, dispatchKernel, cdiv } from '../utils/gpu_utils.js';
17
+
18
+ export const COL_SLICE_WGSL: string = /* wgsl */`
19
+ struct SliceParams {
20
+ rows : u32,
21
+ stride : u32, // source row width
22
+ offset : u32, // first column to take
23
+ width : u32, // number of columns to take
24
+ };
25
+ @group(0) @binding(0) var<uniform> p : SliceParams;
26
+ @group(0) @binding(1) var<storage, read> src : array<f32>;
27
+ @group(0) @binding(2) var<storage, read_write> dst : array<f32>;
28
+
29
+ // Dispatch: (ceil(rows * width / 256), 1, 1)
30
+ @compute @workgroup_size(256)
31
+ fn col_slice(@builtin(global_invocation_id) gid: vec3<u32>) {
32
+ let i = gid.x;
33
+ if (i >= p.rows * p.width) { return; }
34
+ let r = i / p.width;
35
+ let c = i % p.width;
36
+ dst[i] = src[r * p.stride + p.offset + c];
37
+ }
38
+ `;
39
+
40
+ /** Entry point name of {@link COL_SLICE_WGSL}. */
41
+ export const COL_SLICE_ENTRY = 'col_slice';
42
+
43
+ /**
44
+ * `dst[r, 0..width) = src[r, offset..offset+width)` for every row.
45
+ *
46
+ * @param pipeline a pipeline built from {@link COL_SLICE_WGSL} / {@link COL_SLICE_ENTRY}
47
+ * @param rows batch * seqLen
48
+ * @param stride source row width (the fused projection's output features)
49
+ */
50
+ export function dispatchColumnSlice(
51
+ device: GPUDevice,
52
+ pipeline: GPUComputePipeline,
53
+ src: GPUBuffer,
54
+ dst: GPUBuffer,
55
+ rows: number,
56
+ stride: number,
57
+ offset: number,
58
+ width: number,
59
+ ): void {
60
+ const pBuf = createUniformBuffer(device, new Uint32Array([rows, stride, offset, width]).buffer);
61
+ const bg = createBindGroup(device, pipeline, [pBuf, src, dst]);
62
+ dispatchKernel(device, pipeline, bg, [cdiv(rows * width, 256), 1, 1]);
63
+ pBuf.destroy();
64
+ }
@@ -79,6 +79,20 @@ fn adamw_update(
79
79
 
80
80
  // Gradient clipping kernel – clips global gradient norm to max_norm.
81
81
  // Run before weight updates. Two-pass: first compute squared norm, then scale.
82
+ //
83
+ // The accumulator is a REAL atomic. It used to be a plain `array<f32>` written as
84
+ // `norm_sq[0] = norm_sq[0] + local_sq[0]` by the lane-0 invocation of every
85
+ // workgroup — a read-modify-write with no synchronisation, so with more than one
86
+ // workgroup (i.e. any tensor above 256 elements) concurrent workgroups lost each
87
+ // other's contributions and the norm came out too SMALL, under-clipping exactly
88
+ // when clipping matters most. WGSL has no `atomic<f32>`, so the accumulator is an
89
+ // `atomic<u32>` holding the f32 bit pattern and updated with a
90
+ // compare-exchange loop — the standard portable float-atomic-add.
91
+ //
92
+ // Summation ORDER across workgroups is not deterministic (floating-point addition
93
+ // is not associative), so the norm can differ in the last ulp or two between runs.
94
+ // That is fine for a clip threshold and is the accepted trade for correctness;
95
+ // nothing downstream depends on a bit-exact norm.
82
96
  export const GRAD_CLIP_WGSL: string = /* wgsl */`
83
97
 
84
98
  struct ClipParams {
@@ -88,10 +102,23 @@ struct ClipParams {
88
102
 
89
103
  @group(0) @binding(0) var<uniform> clip_p : ClipParams;
90
104
  @group(0) @binding(1) var<storage, read_write> grad : array<f32>;
91
- @group(0) @binding(2) var<storage, read_write> norm_sq : array<f32>; // size 1, atomic accumulator
105
+ // size 1 – the f32 bit pattern of the accumulated sum of squares, updated
106
+ // atomically. Zero-initialise it by writing 0.0f (bit pattern 0x00000000).
107
+ @group(0) @binding(2) var<storage, read_write> norm_sq : array<atomic<u32>>;
92
108
 
93
109
  var<workgroup> local_sq : array<f32, 256>;
94
110
 
111
+ // Portable atomic add on an f32 stored as its u32 bit pattern.
112
+ fn atomic_add_f32(value: f32) {
113
+ var old_bits : u32 = atomicLoad(&norm_sq[0]);
114
+ loop {
115
+ let new_bits = bitcast<u32>(bitcast<f32>(old_bits) + value);
116
+ let res = atomicCompareExchangeWeak(&norm_sq[0], old_bits, new_bits);
117
+ if (res.exchanged) { break; }
118
+ old_bits = res.old_value;
119
+ }
120
+ }
121
+
95
122
  // Pass 1: reduce sum of squares into norm_sq[0]
96
123
  @compute @workgroup_size(256, 1, 1)
97
124
  fn grad_norm_reduce(
@@ -117,8 +144,7 @@ fn grad_norm_reduce(
117
144
  }
118
145
 
119
146
  if (lid == 0u) {
120
- // Non-atomic accumulation (single workgroup assumption for small models)
121
- norm_sq[0] = norm_sq[0] + local_sq[0];
147
+ atomic_add_f32(local_sq[0]);
122
148
  }
123
149
  }
124
150
 
@@ -130,7 +156,7 @@ fn grad_clip_scale(
130
156
  let i = gid.x;
131
157
  if (i >= clip_p.num_elements) { return; }
132
158
 
133
- let ns = norm_sq[0];
159
+ let ns = bitcast<f32>(atomicLoad(&norm_sq[0]));
134
160
  if (ns > clip_p.max_norm_sq) {
135
161
  let scale = sqrt(clip_p.max_norm_sq / ns);
136
162
  grad[i] = grad[i] * scale;
@@ -102,9 +102,37 @@ export interface LMGenerateOptions {
102
102
  stopToken?: number;
103
103
  }
104
104
 
105
+ /**
106
+ * Incremental decode state — everything a continuation needs from the prefix.
107
+ *
108
+ * The model's only cross-position dependency is the depthwise causal conv, whose
109
+ * receptive field is `convKernel` positions, so the entire "KV cache" is the
110
+ * trailing `convKernel - 1` normalised conv inputs per layer. That makes an exact
111
+ * prefix cache cheap: process a prompt once, then extend it with each candidate
112
+ * continuation for the cost of the candidate alone.
113
+ */
114
+ export interface EvermindLMDecodeState {
115
+ /** [layer][slot] the trailing `convKernel - 1` RMSNormed conv inputs, oldest first. */
116
+ convWindow: Float32Array[][];
117
+ /** Positions consumed so far. */
118
+ length: number;
119
+ /** Logits at the LAST consumed position — what predicts the next token. */
120
+ lastLogits: Float32Array | null;
121
+ }
122
+
105
123
  export class EvermindLM {
106
124
  readonly config: Required<Omit<EvermindLMConfig, "seed">>;
107
125
 
126
+ /**
127
+ * Token-positions pushed through a block since the last {@link resetStats}.
128
+ *
129
+ * The cost metric that matters for the tool-calling path: scoring N candidate
130
+ * tool names used to replay the WHOLE prompt N times, so this counted
131
+ * N x (prompt + candidate) x layers. With {@link stepDecode} the prompt is paid
132
+ * for once. Exposed so a benchmark can assert the saving rather than assume it.
133
+ */
134
+ private _positionsEvaluated = 0;
135
+
108
136
  /** Tied token embedding / output head: vocabSize × dModel (row-major). */
109
137
  emb: Float32Array;
110
138
  private gEmb: Float32Array;
@@ -201,48 +229,81 @@ export class EvermindLM {
201
229
  * instead of retaining every layer's cache at once.
202
230
  */
203
231
  private _forwardLayer(l: number, layerIn: Float32Array[]): { afterMoe: Float32Array[]; cache: LayerCache } {
204
- const { dModel, convKernel } = this.config;
205
232
  const T = layerIn.length;
206
- const ker = this.conv[l]!;
207
- const nConv = this.nConv[l]!;
208
- const nMoe = this.nMoe[l]!;
233
+ const window: Float32Array[] = [];
209
234
 
210
235
  const normedConv: Float32Array[] = [];
211
236
  const rmsConv: number[] = [];
212
- for (let t = 0; t < T; t++) {
213
- const { y, r } = rmsNorm(layerIn[t]!, nConv);
214
- normedConv.push(y);
215
- rmsConv.push(r);
216
- }
217
237
  const afterConv: Float32Array[] = [];
218
- for (let t = 0; t < T; t++) {
219
- const out = Float32Array.from(layerIn[t]!); // residual base
220
- for (let c = 0; c < dModel; c++) {
221
- let acc = 0;
222
- for (let j = 0; j < convKernel; j++) {
223
- const ti = t - j;
224
- if (ti >= 0) acc += ker[c * convKernel + j]! * normedConv[ti]![c]!;
225
- }
226
- out[c] = out[c]! + acc;
227
- }
228
- afterConv.push(out);
229
- }
230
-
231
238
  const rmsMoe: number[] = [];
232
239
  const moeCache: MoECacheLike[] = [];
233
240
  const afterMoe: Float32Array[] = [];
241
+
234
242
  for (let t = 0; t < T; t++) {
235
- const { y, r } = rmsNorm(afterConv[t]!, nMoe);
236
- rmsMoe.push(r);
237
- const out = Float32Array.from(afterConv[t]!); // residual base
238
- const mr = this.moe[l]!.forward(y);
239
- for (let c = 0; c < dModel; c++) out[c] = out[c]! + mr.output[c]!;
240
- afterMoe.push(out);
241
- moeCache.push(mr.cache as unknown as MoECacheLike);
243
+ const st = this._stepLayerPosition(l, layerIn[t]!, window);
244
+ normedConv.push(st.normed);
245
+ rmsConv.push(st.rmsConv);
246
+ afterConv.push(st.afterConv);
247
+ rmsMoe.push(st.rmsMoe);
248
+ moeCache.push(st.moeCache);
249
+ afterMoe.push(st.out);
242
250
  }
243
251
  return { afterMoe, cache: { layerIn, normedConv, rmsConv, afterConv, rmsMoe, moeCache } };
244
252
  }
245
253
 
254
+ /**
255
+ * ONE position through one (conv + MoE) block, given the trailing conv window.
256
+ *
257
+ * The single place the block's maths lives: the batch path
258
+ * ({@link _forwardLayer}, which needs the full activation cache for backward)
259
+ * and the incremental path ({@link stepDecode}, which needs none of it) both
260
+ * run this, so they cannot drift apart. `window` holds the previous
261
+ * `convKernel - 1` RMSNormed conv inputs, OLDEST FIRST, and is advanced in
262
+ * place — a shorter window is exactly the causal zero-padding at the start of
263
+ * a sequence.
264
+ */
265
+ private _stepLayerPosition(
266
+ l: number,
267
+ xt: Float32Array,
268
+ window: Float32Array[],
269
+ ): { normed: Float32Array; rmsConv: number; afterConv: Float32Array; rmsMoe: number; moeCache: MoECacheLike; out: Float32Array } {
270
+ const { dModel, convKernel } = this.config;
271
+ const ker = this.conv[l]!;
272
+
273
+ const { y: normed, r: rConv } = rmsNorm(xt, this.nConv[l]!);
274
+ const afterConv = Float32Array.from(xt); // residual base
275
+ for (let c = 0; c < dModel; c++) {
276
+ let acc = ker[c * convKernel]! * normed[c]!;
277
+ for (let j = 1; j < convKernel; j++) {
278
+ const prev = window[window.length - j]; // position t - j, or undefined before the start
279
+ if (prev) acc += ker[c * convKernel + j]! * prev[c]!;
280
+ }
281
+ afterConv[c] = afterConv[c]! + acc;
282
+ }
283
+ window.push(normed);
284
+ while (window.length > convKernel - 1) window.shift();
285
+
286
+ const { y: normedMoe, r: rMoe } = rmsNorm(afterConv, this.nMoe[l]!);
287
+ const out = Float32Array.from(afterConv); // residual base
288
+ const mr = this.moe[l]!.forward(normedMoe);
289
+ for (let c = 0; c < dModel; c++) out[c] = out[c]! + mr.output[c]!;
290
+
291
+ this._positionsEvaluated++;
292
+ return {
293
+ normed, rmsConv: rConv, afterConv, rmsMoe: rMoe,
294
+ moeCache: mr.cache as unknown as MoECacheLike, out,
295
+ };
296
+ }
297
+
298
+ /** One token's embedding row (a copy — callers write into it). */
299
+ private _embedOne(tok: number): Float32Array {
300
+ const { dModel } = this.config;
301
+ const row = new Float32Array(dModel);
302
+ const off = tok * dModel;
303
+ for (let c = 0; c < dModel; c++) row[c] = this.emb[off + c]!;
304
+ return row;
305
+ }
306
+
246
307
  /** Tied output head: logits_t[v] = x_t · emb[v]. */
247
308
  private _head(x: Float32Array[]): Float32Array[] {
248
309
  const { dModel, vocabSize } = this.config;
@@ -270,6 +331,61 @@ export class EvermindLM {
270
331
  return { logits: this._head(x), cache: { tokens, layers, finalX: x } };
271
332
  }
272
333
 
334
+ /**
335
+ * Produces a single fixed-length embedding vector for a token sequence — the
336
+ * CPU counterpart of `HybridMambaModel.embed()`, to the same contract.
337
+ *
338
+ * Runs the full block stack and takes the final hidden state — i.e. exactly the
339
+ * state the tied LM head consumes in {@link forward} — then mean-pools across
340
+ * sequence positions and L2-normalises. The result has length `dModel`, so
341
+ * cosine similarity between two embeddings reduces to a dot product.
342
+ *
343
+ * Like the GPU path it skips the (expensive) head projection: it only needs the
344
+ * `dModel`-wide hidden state, not `vocabSize` logits. Unlike the GPU path it
345
+ * needs no WebGPU device at all, which is what lets a headless host (the stdio
346
+ * MCP server) rank recall by SSM embedding instead of word overlap.
347
+ *
348
+ * Streamed one position at a time through the SHARED block implementation
349
+ * ({@link _stepLayerPosition}, the same routine {@link forward} and
350
+ * {@link stepDecode} run), so it can never drift from them and it retains no
351
+ * per-position activation cache.
352
+ *
353
+ * The embedding reflects whatever the model currently knows — an untrained
354
+ * model behaves like a random projection of the token embeddings (still
355
+ * lexically discriminative), and the representation sharpens as the model is
356
+ * adapted/distilled.
357
+ */
358
+ embed(tokenIds: number[] | Uint32Array): Float32Array {
359
+ const { dModel, numLayers } = this.config;
360
+ const out = new Float32Array(dModel);
361
+ const seqLen = tokenIds.length;
362
+ if (seqLen === 0) return out;
363
+
364
+ // One rolling conv window per layer — a shorter window at the start of the
365
+ // sequence IS the causal zero-padding, exactly as in stepDecode.
366
+ const windows: Float32Array[][] = Array.from({ length: numLayers }, () => [] as Float32Array[]);
367
+ for (let t = 0; t < seqLen; t++) {
368
+ let x = this._embedOne(tokenIds[t]!);
369
+ for (let l = 0; l < numLayers; l++) x = this._stepLayerPosition(l, x, windows[l]!).out;
370
+ for (let c = 0; c < dModel; c++) out[c] = out[c]! + x[c]!;
371
+ }
372
+
373
+ // Mean-pool across positions.
374
+ for (let c = 0; c < dModel; c++) out[c] = out[c]! / seqLen;
375
+
376
+ // L2-normalise so cosine similarity reduces to a dot product.
377
+ let norm = 0;
378
+ for (let c = 0; c < dModel; c++) norm += out[c]! * out[c]!;
379
+ norm = Math.sqrt(norm) || 1;
380
+ for (let c = 0; c < dModel; c++) out[c] = out[c]! / norm;
381
+ return out;
382
+ }
383
+
384
+ /** Text-level embedding: encode with `codec`, then {@link embed}. */
385
+ embedText(text: string, codec: TextCodec): Float32Array {
386
+ return this.embed(codec.encode(text));
387
+ }
388
+
273
389
  // ── Loss + backward ──────────────────────────────────────────────────────────
274
390
 
275
391
  /**
@@ -405,6 +521,80 @@ export class EvermindLM {
405
521
  return loss;
406
522
  }
407
523
 
524
+ // ── Incremental decoding (prefix cache) ──────────────────────────────────────
525
+
526
+ /** Positions pushed through a block since the last {@link resetStats}. */
527
+ get positionsEvaluated(): number { return this._positionsEvaluated; }
528
+ resetStats(): void { this._positionsEvaluated = 0; }
529
+
530
+ /** An empty decode state — the start of a sequence. */
531
+ newDecodeState(): EvermindLMDecodeState {
532
+ return {
533
+ convWindow: Array.from({ length: this.config.numLayers }, () => [] as Float32Array[]),
534
+ length: 0,
535
+ lastLogits: null,
536
+ };
537
+ }
538
+
539
+ /** Deep copy, so one prefix state can seed many independent continuations. */
540
+ cloneDecodeState(state: EvermindLMDecodeState): EvermindLMDecodeState {
541
+ return {
542
+ convWindow: state.convWindow.map((w) => w.map((v) => Float32Array.from(v))),
543
+ length: state.length,
544
+ lastLogits: state.lastLogits ? Float32Array.from(state.lastLogits) : null,
545
+ };
546
+ }
547
+
548
+ /**
549
+ * Consume `tokens` from `state`, returning the per-position logits and the
550
+ * ADVANCED state. `state` is not mutated — pass {@link cloneDecodeState} output
551
+ * or a fresh state; this clones internally so the caller's prefix stays reusable.
552
+ *
553
+ * Numerically identical to {@link forward} over the concatenated sequence: the
554
+ * shared {@link _stepLayerPosition} is the only implementation of the block.
555
+ */
556
+ stepDecode(state: EvermindLMDecodeState, tokens: number[]): { logits: Float32Array[]; state: EvermindLMDecodeState } {
557
+ const next = this.cloneDecodeState(state);
558
+ const logits: Float32Array[] = [];
559
+ for (const tok of tokens) {
560
+ let x = this._embedOne(tok);
561
+ for (let l = 0; l < this.config.numLayers; l++) {
562
+ x = this._stepLayerPosition(l, x, next.convWindow[l]!).out;
563
+ }
564
+ const lg = this._head([x])[0]!;
565
+ logits.push(lg);
566
+ next.lastLogits = lg;
567
+ next.length++;
568
+ }
569
+ return { logits, state: next };
570
+ }
571
+
572
+ /** Run a prompt once and keep the state, so continuations are cheap. */
573
+ prefixState(tokens: number[]): EvermindLMDecodeState {
574
+ return this.stepDecode(this.newDecodeState(), tokens).state;
575
+ }
576
+
577
+ /**
578
+ * Mean log-probability the model assigns to `continuation` following a cached
579
+ * prefix — the tool-decision vote, without replaying the prompt.
580
+ *
581
+ * Returns `-Infinity` for an empty continuation, and requires a prefix with at
582
+ * least one position (there must be a position to predict the first
583
+ * continuation token FROM).
584
+ */
585
+ scoreContinuation(prefix: EvermindLMDecodeState, continuation: number[]): number {
586
+ if (continuation.length === 0) return -Infinity;
587
+ if (!prefix.lastLogits) return -Infinity;
588
+ // The last prefix position predicts continuation[0]; each continuation
589
+ // position then predicts the next one, so the final one is never scored.
590
+ const { logits } = this.stepDecode(prefix, continuation.slice(0, -1));
591
+ let total = logProbOfToken(prefix.lastLogits, continuation[0]!);
592
+ for (let i = 1; i < continuation.length; i++) {
593
+ total += logProbOfToken(logits[i - 1]!, continuation[i]!);
594
+ }
595
+ return total / continuation.length;
596
+ }
597
+
408
598
  // ── Generation ───────────────────────────────────────────────────────────────
409
599
 
410
600
  /**
@@ -417,19 +607,24 @@ export class EvermindLM {
417
607
  return codec.decode(this.generate(codec.encode(prompt), opts));
418
608
  }
419
609
 
420
- /** Greedy / temperature-sampled autoregressive generation. Returns NEW token ids. */
610
+ /**
611
+ * Greedy / temperature-sampled autoregressive generation. Returns NEW token ids.
612
+ *
613
+ * Carries an incremental decode state, so the prompt is processed ONCE and each
614
+ * new token costs one position. The previous implementation re-ran the whole
615
+ * sequence per token, making generation quadratic in the prompt length.
616
+ */
421
617
  generate(prompt: number[], opts: LMGenerateOptions): number[] {
422
618
  const temperature = opts.temperature ?? 0;
423
619
  const rng = temperature > 0 ? new SeededRng((opts.seed ?? 1) >>> 0 || 1) : null;
424
- const tokens = [...prompt];
620
+ let state = this.stepDecode(this.newDecodeState(), prompt.length > 0 ? prompt : [0]).state;
425
621
  const produced: number[] = [];
426
622
  for (let n = 0; n < opts.maxNewTokens; n++) {
427
- const { logits } = this.forward(tokens.length > 0 ? tokens : [0]);
428
- const last = logits[logits.length - 1]!;
623
+ const last = state.lastLogits!;
429
624
  const next = rng ? sampleTemperature(last, temperature, rng) : argmax(last);
430
625
  produced.push(next);
431
- tokens.push(next);
432
626
  if (opts.stopToken !== undefined && next === opts.stopToken) break;
627
+ state = this.stepDecode(state, [next]).state;
433
628
  }
434
629
  return produced;
435
630
  }
@@ -745,3 +940,18 @@ function sampleTemperature(logits: Float32Array, temperature: number, rng: Seede
745
940
  }
746
941
  return probs.length - 1;
747
942
  }
943
+
944
+ /**
945
+ * Log-probability the model assigned to `id` at a position, from that position's
946
+ * raw logits. Log-softmax computed max-shifted so a long-tail logit cannot
947
+ * overflow `exp`.
948
+ */
949
+ export function logProbOfToken(row: Float32Array, id: number): number {
950
+ let max = -Infinity;
951
+ for (let i = 0; i < row.length; i++) { const v = row[i]!; if (v > max) max = v; }
952
+ if (!Number.isFinite(max)) return -Infinity;
953
+ let sum = 0;
954
+ for (let i = 0; i < row.length; i++) sum += Math.exp(row[i]! - max);
955
+ const logit = id >= 0 && id < row.length ? row[id]! : -Infinity;
956
+ return logit - (max + Math.log(sum));
957
+ }
package/src/lm/index.ts CHANGED
@@ -2,8 +2,12 @@
2
2
  * EvermindLM — the generative language model (the runnable "AI").
3
3
  */
4
4
 
5
- export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED } from "./evermind_lm.js";
6
- export type { EvermindLMConfig, LMGenerateOptions, TextCodec, EvermindLMTrainOptions } from "./evermind_lm.js";
5
+ export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED, logProbOfToken } from "./evermind_lm.js";
6
+ export type { EvermindLMConfig, LMGenerateOptions, TextCodec, EvermindLMTrainOptions, EvermindLMDecodeState } from "./evermind_lm.js";
7
+
8
+ // CPU text→vector embedder (semantic recall with no GPU).
9
+ export { EvermindTextEmbedder } from "./text_embedder.js";
10
+ export type { TextEmbedder, EmbedderCodec } from "./text_embedder.js";
7
11
 
8
12
  // PEFT / efficient-training toolkit (LoRA, QLoRA, mixed precision).
9
13
  export { LoRAAdapter, EvermindLMLoRA, quantizeBase } from "../training/lora.js";
@@ -0,0 +1,84 @@
1
+ /**
2
+ * lm/text_embedder.ts — text → vector, on the CPU, from an `EvermindLM`.
3
+ *
4
+ * `EvermindLM.embed()` turns TOKENS into a `dModel` vector. Every consumer that
5
+ * wants SEMANTIC RECALL wants it over TEXT, and needs the vector to be stable
6
+ * across processes, so it also needs (a) the tokenizer that produced those ids
7
+ * and (b) a way to tell whether a cached vector came from this exact model. This
8
+ * class is that pairing, in one place, so no consumer re-derives it.
9
+ *
10
+ * It deliberately satisfies the SAME structural shape the memory runtime's
11
+ * `embed(text)` seam expects (`MemoryStore.recallRanked(query, k, runtime)`
12
+ * accepts anything with an async `embed(text): Float32Array`), so a headless host
13
+ * can hand one of these straight to the memory layer and get embedding-ranked
14
+ * recall with NO GPU — the CPU LM is the whole dependency.
15
+ */
16
+
17
+ import type { EvermindLM, TextCodec } from "./evermind_lm.js";
18
+ import { crc32 } from "../utils/crc32.js";
19
+
20
+ /** The minimal text-embedder seam the memory/recall layers consume. */
21
+ export interface TextEmbedder {
22
+ /** Embed `text` as an L2-normalised vector of length {@link dimensions}. */
23
+ embed(text: string): Promise<Float32Array>;
24
+ /** Vector width (the model's `dModel`). */
25
+ readonly dimensions: number;
26
+ /**
27
+ * Identity of the model+tokenizer that produces these vectors. Vectors are only
28
+ * comparable within one fingerprint, so a persistent cache MUST discard entries
29
+ * stamped with a different one.
30
+ */
31
+ readonly fingerprint: string;
32
+ }
33
+
34
+ /** A tokenizer this embedder can encode through (the engine's `BPETokenizer` fits). */
35
+ export type EmbedderCodec = Pick<TextCodec, "encode">;
36
+
37
+ /**
38
+ * Pairs an `EvermindLM` with the tokenizer whose ids it was trained on.
39
+ *
40
+ * Pure compute and stateless: caching belongs to whoever owns a lifetime (the
41
+ * memory store's in-process LRU, or a host's on-disk vector cache), not here.
42
+ */
43
+ export class EvermindTextEmbedder implements TextEmbedder {
44
+ readonly dimensions: number;
45
+ readonly fingerprint: string;
46
+
47
+ constructor(
48
+ private readonly model: EvermindLM,
49
+ private readonly codec: EmbedderCodec,
50
+ /** Override the derived fingerprint (e.g. with a published model version). */
51
+ fingerprint?: string,
52
+ ) {
53
+ this.dimensions = model.config.dModel;
54
+ this.fingerprint = fingerprint ?? derivedFingerprint(model);
55
+ }
56
+
57
+ /**
58
+ * Embed `text`. Async only to match the seam — the work is synchronous CPU, so
59
+ * a caller that wants the vector without a microtask can use
60
+ * {@link embedSync}.
61
+ */
62
+ async embed(text: string): Promise<Float32Array> {
63
+ return this.embedSync(text);
64
+ }
65
+
66
+ /** The synchronous form — encode with the codec, then run the CPU LM. */
67
+ embedSync(text: string): Float32Array {
68
+ return this.model.embedText(text, this.codec as TextCodec);
69
+ }
70
+ }
71
+
72
+ /**
73
+ * Model identity: architecture plus a CRC-32 over the tied embedding table.
74
+ *
75
+ * The embedding table is the one tensor every position's vector flows through, so
76
+ * an adapted checkpoint (which is what a delta merge produces) always changes it —
77
+ * making this a cheap but honest "are these vectors still comparable?" stamp.
78
+ */
79
+ function derivedFingerprint(model: EvermindLM): string {
80
+ const c = model.config;
81
+ const bytes = new Uint8Array(model.emb.buffer, model.emb.byteOffset, model.emb.byteLength);
82
+ const arch = [c.vocabSize, c.dModel, c.numLayers, c.convKernel, c.hiddenDim, c.numExperts, c.topK].join("-");
83
+ return `evermind-lm/${arch}/${crc32(bytes).toString(16)}`;
84
+ }
@@ -36,6 +36,7 @@ import {
36
36
  import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
37
37
  import { gaussianArray } from '../utils/rng.js';
38
38
  import { ACTIVATIONS_WGSL } from '../kernels/activations.js';
39
+ import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
39
40
 
40
41
  import type { SequenceLayer, LayerForwardResult, LayerParam } from './sequence_layer.js';
41
42
 
@@ -147,6 +148,7 @@ export class AttentionBlock implements SequenceLayer {
147
148
  attn_fwd: createComputePipeline(d, ATTENTION_FORWARD_WGSL, 'attention_forward'),
148
149
  attn_val: createComputePipeline(d, ATTENTION_FORWARD_WGSL, 'attention_value'),
149
150
  softmax : createComputePipeline(d, SOFTMAX_WGSL, 'softmax_forward'),
151
+ colSlice: createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
150
152
  elAdd : createComputePipeline(d, ADD_SHADER, 'main'),
151
153
  };
152
154
 
@@ -194,11 +196,13 @@ export class AttentionBlock implements SequenceLayer {
194
196
  const KBuf = createEmptyStorageBuffer(d, M * dModel * 4, true);
195
197
  const VBuf = createEmptyStorageBuffer(d, M * dModel * 4, true);
196
198
  {
197
- const enc = d.createCommandEncoder();
198
- enc.copyBufferToBuffer(qkvOut, 0, QBuf, 0, M * dModel * 4);
199
- enc.copyBufferToBuffer(qkvOut, M * dModel * 4, KBuf, 0, M * dModel * 4);
200
- enc.copyBufferToBuffer(qkvOut, 2 * M * dModel * 4, VBuf, 0, M * dModel * 4);
201
- d.queue.submit([enc.finish()]);
199
+ // COLUMN split. `linear_forward` writes (M, 3*dModel) ROW-major, so Q, K
200
+ // and V are ranges of columns WITHIN every row — a strided gather, not a
201
+ // contiguous copy (which only coincides for M == 1).
202
+ const slice = this.pipelines['colSlice']!;
203
+ dispatchColumnSlice(d, slice, qkvOut, QBuf, M, 3 * dModel, 0, dModel);
204
+ dispatchColumnSlice(d, slice, qkvOut, KBuf, M, 3 * dModel, dModel, dModel);
205
+ dispatchColumnSlice(d, slice, qkvOut, VBuf, M, 3 * dModel, 2 * dModel, dModel);
202
206
  }
203
207
  qkvOut.destroy();
204
208