mlx-bun 0.3.0 → 0.4.0

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 (150) hide show
  1. package/docs/README.md +4 -1
  2. package/docs/archive/investigations/openjarvis-fit.md +198 -0
  3. package/docs/archive/investigations/qwen38-27b-campaign-log.md +1092 -0
  4. package/docs/design/batching.md +1321 -49
  5. package/docs/design/decode-speed-program.md +331 -1163
  6. package/docs/design/generic-model-support.md +79 -3
  7. package/docs/design/kv-cache.md +198 -54
  8. package/docs/design/speculative-decoding.md +387 -24
  9. package/docs/design/turboquant.md +123 -13
  10. package/docs/design/unified-engine-frontier-plan.md +144 -142
  11. package/docs/planning/PRODUCT_ROADMAP.md +16 -0
  12. package/docs/planning/release-notes-v0.4.0.md +45 -0
  13. package/docs/reference/benchmarks.md +2457 -6
  14. package/docs/reference/cli.md +41 -2
  15. package/docs/reference/distribution.md +16 -9
  16. package/docs/reference/environment.md +55 -5
  17. package/docs/reference/models.md +13 -4
  18. package/docs/reference/server-api.md +79 -25
  19. package/docs/reference/server-config.md +178 -84
  20. package/docs/reference/troubleshooting.md +13 -21
  21. package/package.json +1 -1
  22. package/src/backends/mlx/assistant-target.ts +46 -0
  23. package/src/backends/mlx/autoregressive-draft-rows.ts +120 -0
  24. package/src/backends/mlx/batch-group.ts +368 -224
  25. package/src/backends/mlx/cache-attachment.ts +67 -0
  26. package/src/backends/mlx/cache-layout.ts +31 -0
  27. package/src/backends/mlx/checkpoint-state.ts +53 -0
  28. package/src/backends/mlx/continuation-persistence.ts +106 -0
  29. package/src/backends/mlx/continuation-request.ts +43 -0
  30. package/src/backends/mlx/continuation.ts +77 -0
  31. package/src/backends/mlx/draft-target.ts +7 -22
  32. package/src/backends/mlx/full-kv-row-append.ts +32 -0
  33. package/src/backends/mlx/full-kv-row-donor.ts +54 -0
  34. package/src/backends/mlx/full-transitioning-kv-rows.ts +65 -0
  35. package/src/backends/mlx/gateway-binding.ts +77 -18
  36. package/src/backends/mlx/kv-codec.ts +9 -0
  37. package/src/backends/mlx/kv-maintenance.ts +133 -0
  38. package/src/backends/mlx/kv-tensor-rows.ts +212 -0
  39. package/src/backends/mlx/model-serving.ts +1 -1
  40. package/src/backends/mlx/prefill-cohort.ts +89 -0
  41. package/src/backends/mlx/prefill-rows.ts +200 -0
  42. package/src/backends/mlx/prefill.ts +1 -1
  43. package/src/backends/mlx/projected-context-rows.ts +72 -0
  44. package/src/backends/mlx/request-state-policy.ts +41 -0
  45. package/src/backends/mlx/rollback.ts +41 -3
  46. package/src/backends/mlx/row-sampling.ts +49 -0
  47. package/src/backends/mlx/serial-executor.ts +35 -66
  48. package/src/backends/mlx/speculative-group.ts +359 -0
  49. package/src/backends/mlx/speculative-round.ts +161 -0
  50. package/src/backends/mlx/state-rows.ts +104 -0
  51. package/src/backends/mlx/state-views.ts +21 -1
  52. package/src/backends/mlx/target-layout-capability.ts +35 -0
  53. package/src/backends/mlx/transitioning-kv-rows.ts +165 -0
  54. package/src/cli.ts +11 -36
  55. package/src/contracts/execution.ts +11 -0
  56. package/src/contracts/prefix-cache.ts +23 -0
  57. package/src/contracts/resources.ts +7 -0
  58. package/src/contracts/sampling.ts +23 -0
  59. package/src/contracts/scheduling.ts +18 -0
  60. package/src/engine/execution-plan.ts +18 -16
  61. package/src/engine/resources.ts +12 -1
  62. package/src/engine/scheduler.ts +10 -5
  63. package/src/fit.ts +29 -15
  64. package/src/generate.ts +77 -173
  65. package/src/inference/draft-acceptance.ts +41 -0
  66. package/src/inference/draft-output.ts +49 -0
  67. package/src/inference/rollback.ts +2 -2
  68. package/src/kv-scheme.ts +70 -12
  69. package/src/kv-store.ts +70 -16
  70. package/src/lab/paged-kv/paged-kv-rows.ts +83 -0
  71. package/src/lab/paged-kv/paged-kv.ts +25 -11
  72. package/src/mlx/array.ts +2 -1
  73. package/src/model/artifact-identity.ts +73 -0
  74. package/src/model/batched-kv.ts +235 -0
  75. package/src/model/batched-mask.ts +12 -4
  76. package/src/model/batched-quant.ts +8 -2
  77. package/src/model/batched-quantized-kv.ts +144 -0
  78. package/src/model/batched-rotating-quant.ts +76 -15
  79. package/src/model/batched-rotating-state.ts +76 -0
  80. package/src/model/batched-rotating.ts +64 -22
  81. package/src/model/batched-row-storage.ts +52 -0
  82. package/src/model/batched-ssm.ts +52 -0
  83. package/src/model/batched-turboquant-kv.ts +93 -0
  84. package/src/model/compiled-decode.ts +7 -1
  85. package/src/model/decoded-kv-donor.ts +12 -0
  86. package/src/model/delayed-quantized-kv.ts +72 -0
  87. package/src/model/delayed-rotating-quantized-kv.ts +93 -0
  88. package/src/model/delayed-turboquant-kv.ts +55 -0
  89. package/src/model/full-prefill-padding.ts +56 -0
  90. package/src/model/full-prefill-row.ts +105 -0
  91. package/src/model/gemma4-base.ts +162 -80
  92. package/src/model/gemma4.ts +34 -26
  93. package/src/model/generated/gemma4-12b.ts +4 -1
  94. package/src/model/generated/gemma4-26b.ts +4 -1
  95. package/src/model/generated/gemma4-e4b.ts +4 -1
  96. package/src/model/glm52-cache.ts +80 -4
  97. package/src/model/glm52-mla.ts +12 -6
  98. package/src/model/glm52.ts +19 -40
  99. package/src/model/kv-attention-view.ts +67 -0
  100. package/src/model/minicpm5.ts +11 -7
  101. package/src/model/qwen3-delta-state.ts +74 -0
  102. package/src/model/qwen3-delta.ts +59 -36
  103. package/src/model/qwen3_5.ts +82 -42
  104. package/src/model/rotating-kv-donor.ts +35 -0
  105. package/src/model/rotating-kv-layout.ts +262 -0
  106. package/src/model/rotating-row-transaction.ts +35 -0
  107. package/src/model/speculative-rotating-kv.ts +101 -0
  108. package/src/model/ssm-prefill-padding.ts +39 -0
  109. package/src/model/trellis-linear.ts +2 -2
  110. package/src/model/turboquant-codec.ts +83 -0
  111. package/src/model/universal/dense.ts +14 -13
  112. package/src/preflight.ts +31 -15
  113. package/src/prompt-cache.ts +42 -44
  114. package/src/sampler.ts +36 -25
  115. package/src/serve/chat-stage.ts +2 -2
  116. package/src/serve/generation-gateway.ts +39 -30
  117. package/src/serve/model-host.ts +3 -4
  118. package/src/serve/request-plan.ts +4 -3
  119. package/src/serve/request-prep.ts +2 -4
  120. package/src/serve/text-completion-stage.ts +2 -2
  121. package/src/server.ts +64 -108
  122. package/src/spec/assistant-rows.ts +92 -0
  123. package/src/spec/assistant-source.ts +28 -18
  124. package/src/spec/cache-identity.ts +22 -0
  125. package/src/spec/deepspec-rows.ts +29 -0
  126. package/src/spec/deepspec-source.ts +12 -10
  127. package/src/spec/dflash-source.ts +28 -4
  128. package/src/spec/drafter.ts +153 -142
  129. package/src/spec/dspark/deepspec-module.ts +95 -120
  130. package/src/spec/dspark/module-dflash.ts +126 -170
  131. package/src/spec/glm52-mtp-graph.ts +29 -0
  132. package/src/spec/glm52-mtp-rows.ts +132 -0
  133. package/src/spec/glm52-mtp-source.ts +53 -136
  134. package/src/spec/glm52-mtp-state.ts +33 -0
  135. package/src/spec/ngram-source.ts +92 -20
  136. package/src/spec/projected-draft-rows.ts +73 -0
  137. package/src/spec/qwen-mtp-module.ts +171 -0
  138. package/src/spec/qwen-mtp-rows.ts +219 -0
  139. package/src/spec/qwen-mtp-source.ts +81 -226
  140. package/src/spec/qwen-mtp-state.ts +30 -0
  141. package/src/spec/serve-loop.ts +70 -68
  142. package/src/spec/source.ts +80 -15
  143. package/src/spec/two-model-rows.ts +61 -0
  144. package/src/spec/two-model.ts +43 -101
  145. package/src/ssd-cache.ts +33 -21
  146. package/src/ssd-durability.ts +11 -7
  147. package/src/web/app.js +57 -57
  148. package/src/backends/mlx/memory-guard.ts +0 -40
  149. package/src/spec/prefix-state.ts +0 -45
  150. /package/docs/{planning → archive/planning}/release-notes-v0.3.0.md +0 -0
package/src/kv-store.ts CHANGED
@@ -1,3 +1,4 @@
1
+ import { type CheckpointAttachment, attachmentBytes, disposeAttachments } from "./backends/mlx/checkpoint-state";
1
2
  import { cleanupFailure, disposeResources } from "./engine/resources";
2
3
  // KV-cache persistence: save prompt caches to disk, reload by streamed copy.
3
4
  // The serialization core of the SSD cold tier (docs/design/kv-cache.md).
@@ -86,6 +87,7 @@ interface TensorSlot {
86
87
  export interface CacheHeaderEntry {
87
88
  kind: CacheKind;
88
89
  offset: number;
90
+ minimumReusableOffset?: number;
89
91
  /** rotating variants: ring write index */
90
92
  idx?: number;
91
93
  /** rotating variants: window size */
@@ -114,6 +116,7 @@ export interface CacheHeaderEntry {
114
116
  }
115
117
 
116
118
  export interface KvSaveMeta {
119
+ attachments?: CheckpointAttachment[];
117
120
  modelId?: string;
118
121
  /** configFingerprint(config) — covers every graph-shaping field incl.
119
122
  * the kv-quant scheme, so a scheme flip invalidates naturally. */
@@ -139,10 +142,11 @@ export interface KvSaveMeta {
139
142
  };
140
143
  }
141
144
 
142
- export interface KvFileHeader extends KvSaveMeta {
145
+ export interface KvFileHeader extends Omit<KvSaveMeta, "attachments"> {
146
+ attachments?: Array<Omit<CheckpointAttachment, "tensors"> & { tensors: TensorSlot[] }>;
143
147
  /** Absent in earlier v3 files, which used mlx-cache-v3. */
144
148
  codecProvider?: string;
145
- formatVersion: 3;
149
+ formatVersion: 3 | 4;
146
150
  createdAt: number;
147
151
  tokens: number[];
148
152
  caches: CacheHeaderEntry[];
@@ -273,8 +277,8 @@ const CACHE_CODECS = {
273
277
  },
274
278
  clone: (cache, context) => {
275
279
  const c = cache as KVCache;
276
- if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
277
280
  const clone = new KVCache();
281
+ if (!c.keys || !c.values) return clone;
278
282
  clone.restoreState(
279
283
  context.liveView(c.keys, c.offset),
280
284
  context.liveView(c.values, c.offset),
@@ -305,8 +309,8 @@ const CACHE_CODECS = {
305
309
  },
306
310
  clone: (cache, context) => {
307
311
  const c = cache as RotatingKVCache;
308
- if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
309
312
  const clone = new RotatingKVCache(c.maxSize);
313
+ if (!c.keys || !c.values) return clone;
310
314
  clone.restoreState(context.view(c.keys), context.view(c.values), c.offset, c.ringIdx);
311
315
  return clone;
312
316
  },
@@ -332,8 +336,8 @@ const CACHE_CODECS = {
332
336
  },
333
337
  clone: (cache, context) => {
334
338
  const c = cache as QuantizedKVCache;
335
- if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
336
339
  const clone = new QuantizedKVCache(c.groupSize, c.bits);
340
+ if (!c.keys || !c.values) return clone;
337
341
  clone.restoreState(
338
342
  context.tripleView(c.keys, c.offset),
339
343
  context.tripleView(c.values, c.offset),
@@ -363,8 +367,8 @@ const CACHE_CODECS = {
363
367
  },
364
368
  clone: (cache, context) => {
365
369
  const c = cache as RotatingQuantizedKVCache;
366
- if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
367
370
  const clone = new RotatingQuantizedKVCache(c.maxSize, c.groupSize, c.bits);
371
+ if (!c.keys || !c.values) return clone;
368
372
  clone.restoreState(
369
373
  context.tripleView(c.keys, null), context.tripleView(c.values, null),
370
374
  c.offset, c.ringIdx,
@@ -515,6 +519,17 @@ export function createCacheCodecProvider(
515
519
 
516
520
  export const legacyCacheCodecs = createCacheCodecProvider("mlx-cache-v3", CACHE_CODECS);
517
521
 
522
+ /** Old TQ files lack the conversion boundary. Keep them reusable at their
523
+ * stored offset, but never infer that an earlier precision state is valid. */
524
+ export function cacheHeaderMinimumReusableOffset(entry: CacheHeaderEntry): number {
525
+ const offset = entry.minimumReusableOffset ?? (entry.kind === "turboquant" ? entry.offset : 0);
526
+ if (!Number.isSafeInteger(offset) || offset < 0 || offset > entry.offset)
527
+ throw new Error("invalid cache minimum reusable offset");
528
+ return offset;
529
+ }
530
+ export const cacheHeadersMinimumReusableOffset = (caches: readonly CacheHeaderEntry[]): number =>
531
+ Math.max(0, ...caches.map(cacheHeaderMinimumReusableOffset));
532
+
518
533
  export const cacheHeadersTrimmable = (
519
534
  caches: readonly CacheHeaderEntry[], codecs: CacheCodecProvider = legacyCacheCodecs,
520
535
  ): boolean => caches.every((entry) => codecs.forHeader(entry).headerTrimmable(entry));
@@ -547,6 +562,7 @@ function snapshotCache(c: Cache, dataOffset: number, codecs: CacheCodecProvider)
547
562
 
548
563
  try {
549
564
  const entry = codecs.forCache(c).snapshot(c, { slots, push, liveSlice, liveMlaSlice, pushTriple });
565
+ if (c.minimumReusableOffset !== undefined) entry.minimumReusableOffset = c.minimumReusableOffset;
550
566
  return { entry, sources, next: cursor };
551
567
  } catch (error) {
552
568
  for (const source of sources) if (source.disposeAfter) source.arr.dispose();
@@ -588,7 +604,9 @@ export function cloneKvCaches(caches: Cache[], codecs: CacheCodecProvider = lega
588
604
  const out: Cache[] = [];
589
605
  try {
590
606
  for (const c of caches) {
591
- out.push(codecs.forCache(c).clone(c, { view, liveView, mlaView, tripleView }));
607
+ const clone = codecs.forCache(c).clone(c, { view, liveView, mlaView, tripleView });
608
+ if (c.minimumReusableOffset !== undefined) clone.minimumReusableOffset = c.minimumReusableOffset;
609
+ out.push(clone);
592
610
  pending.length = 0; // transferred to the returned cache
593
611
  }
594
612
  } catch (err) {
@@ -615,10 +633,24 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
615
633
  dataOffset = s.next;
616
634
  }
617
635
 
636
+ const { attachments: liveAttachments, ...identity } = meta;
637
+ const attachments = liveAttachments?.map((attachment) => ({
638
+ schema: attachment.schema,
639
+ metadata: { ...attachment.metadata },
640
+ tensors: attachment.tensors.map((arr) => {
641
+ const slot: TensorSlot = { off: dataOffset, bytes: arr.nbytes,
642
+ shape: [...arr.shape], dtype: arr.dtype, hash: "0000000000000000" };
643
+ dataOffset = alignUp(dataOffset + slot.bytes);
644
+ sources.push({ arr, disposeAfter: false });
645
+ return slot;
646
+ }),
647
+ }));
648
+
618
649
  // Header size is final NOW: hashes are fixed-width placeholders that get
619
650
  // patched in place (same byte length) after the data pass computes them.
620
651
  const header: KvFileHeader = {
621
- formatVersion: 3, createdAt: Date.now(), ...meta, codecProvider: codecs.id, tokens, caches: entries,
652
+ formatVersion: attachments?.length ? 4 : 3, createdAt: Date.now(), ...identity,
653
+ ...(attachments?.length ? { attachments } : {}), codecProvider: codecs.id, tokens, caches: entries,
622
654
  };
623
655
  const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
624
656
  const dataStart = alignUp(PREFIX_LEN + headerLen);
@@ -639,7 +671,7 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
639
671
  // pages, which makes them VISIBLE to ps RSS — accounting, not an
640
672
  // allocation (see bench-serve.ts' per-leg RSS note).
641
673
  let srcIdx = 0;
642
- for (const e of entries) {
674
+ for (const e of [...entries, ...(attachments ?? [])]) {
643
675
  for (const slot of e.tensors) {
644
676
  const src = sources[srcIdx++]!;
645
677
  const c = ops.contiguous(src.arr);
@@ -717,6 +749,7 @@ export async function saveKvCacheAsync(
717
749
  /** One pending write-behind item: an entry's copied tokens + OWNED
718
750
  * zero-copy cache clones (the queue disposes them on every exit path). */
719
751
  export interface SpillItem {
752
+ attachments?: CheckpointAttachment[];
720
753
  tokens: number[];
721
754
  caches: Cache[];
722
755
  ns: string;
@@ -761,6 +794,8 @@ export class SpillQueue {
761
794
  readonly store: (item: SpillItem) => Promise<unknown>,
762
795
  /** Clone disposal — frees the pinned GPU memory. */
763
796
  readonly disposeClones: (caches: Cache[]) => void,
797
+ /** Zero-byte control barriers may need to survive capacity shedding. */
798
+ readonly canDrop: (item: SpillItem) => boolean = () => true,
764
799
  ) {}
765
800
 
766
801
  /** Bytes pinned by queued (not yet flushed) clones. */
@@ -776,7 +811,7 @@ export class SpillQueue {
776
811
  const result = new Promise<boolean>((resolve) => { settle = resolve; });
777
812
  const rec: SpillRec = {
778
813
  ...item,
779
- bytes: this.bytesOf(item.caches),
814
+ bytes: this.bytesOf(item.caches) + attachmentBytes(item.attachments),
780
815
  dropped: false,
781
816
  settle,
782
817
  };
@@ -784,7 +819,7 @@ export class SpillQueue {
784
819
  this.#bytes += rec.bytes;
785
820
  while (this.#bytes > this.capBytes) {
786
821
  const victim = this.#queue.find(
787
- (p) => !p.dropped && p !== this.#inFlight && p !== rec,
822
+ (p) => !p.dropped && p !== this.#inFlight && p !== rec && this.canDrop(p),
788
823
  );
789
824
  if (!victim) break; // only the new item (or in-flight) left — soft cap
790
825
  this.#drop(victim);
@@ -803,12 +838,19 @@ export class SpillQueue {
803
838
  } finally {
804
839
  this.#inFlight = null;
805
840
  this.#remove(rec);
806
- this.disposeClones(rec.caches);
841
+ disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
842
+ { dispose: () => disposeAttachments(rec.attachments) }]);
807
843
  }
808
844
  });
809
845
  return result;
810
846
  }
811
847
 
848
+ /** Cancel queued ownership, excluding the write already in flight. */
849
+ cancelWhere(matches: (item: SpillItem) => boolean): void {
850
+ for (const rec of [...this.#queue])
851
+ if (!rec.dropped && rec !== this.#inFlight && matches(rec)) this.#drop(rec);
852
+ }
853
+
812
854
  /** Settles when everything enqueued so far has flushed or dropped
813
855
  * (tests and graceful shutdown). */
814
856
  drain(): Promise<void> { return this.#chain; }
@@ -818,7 +860,8 @@ export class SpillQueue {
818
860
  this.#dropped++;
819
861
  rec.settle(false);
820
862
  this.#remove(rec);
821
- this.disposeClones(rec.caches);
863
+ disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
864
+ { dispose: () => disposeAttachments(rec.attachments) }]);
822
865
  }
823
866
 
824
867
  #remove(rec: SpillRec): void {
@@ -848,7 +891,7 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
848
891
  if (BigInt(Bun.hash(body)) !== expectHash)
849
892
  throw new Error(`${path}: header hash mismatch (truncated or corrupt)`);
850
893
  const header = JSON.parse(new TextDecoder().decode(body)) as KvFileHeader;
851
- if (header.formatVersion !== 3)
894
+ if (header.formatVersion !== 3 && header.formatVersion !== 4)
852
895
  throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
853
896
  return { ...header, dataStart };
854
897
  } finally {
@@ -857,6 +900,7 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
857
900
  }
858
901
 
859
902
  export interface LoadedKvCache {
903
+ attachments?: CheckpointAttachment[];
860
904
  tokens: number[];
861
905
  header: KvFileHeader;
862
906
  caches: Cache[];
@@ -1026,20 +1070,30 @@ export function loadKvCache(
1026
1070
  ({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) });
1027
1071
 
1028
1072
  const caches: Cache[] = [];
1073
+ const attachments: CheckpointAttachment[] = [];
1029
1074
  try {
1030
1075
  for (const e of header.caches) {
1031
1076
  const codec = codecs.forHeader(e);
1032
- caches.push(codec.load(e, { path, arr, grownArr, triple }));
1077
+ const minimumOffset = cacheHeaderMinimumReusableOffset(e);
1078
+ const cache = codec.load(e, { path, arr, grownArr, triple });
1079
+ cache.minimumReusableOffset = minimumOffset;
1080
+ caches.push(cache);
1033
1081
  // This entry's tensors are now owned by its cache — stop tracking them
1034
1082
  // (the catch must not double-free through both pending AND caches).
1035
1083
  pending.length = 0;
1036
1084
  }
1085
+ for (const attachment of header.attachments ?? []) {
1086
+ attachments.push({ schema: attachment.schema, metadata: { ...attachment.metadata },
1087
+ tensors: attachment.tensors.map(arr) });
1088
+ pending.length = 0;
1089
+ }
1037
1090
  } catch (err) {
1091
+ disposeAttachments(attachments);
1038
1092
  for (const a of pending) a.dispose(); // the mid-entry orphans
1039
1093
  for (const c of caches) c.dispose();
1040
1094
  throw err;
1041
1095
  } finally {
1042
1096
  mmap.unmap(); // every tensor was copied out — nothing aliases the file
1043
1097
  }
1044
- return { tokens: header.tokens, header, caches };
1098
+ return { tokens: header.tokens, header, caches, ...(attachments.length ? { attachments } : {}) };
1045
1099
  }
@@ -0,0 +1,83 @@
1
+ import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
2
+ import { buildBatchedDecodeMask, mergeKVRows } from "../../model/batched-mask";
3
+ import { MlxArray } from "../../mlx/array";
4
+ import { disposeResources } from "../../engine/resources";
5
+ import { PagedKVCache, poolBlocksFor } from "./paged-kv";
6
+
7
+ /** Block storage owns row membership; attention consumes the same gathered
8
+ * K/V representation as the existing paged cache. No scheduler policy here. */
9
+ export class PagedKvRows implements BatchableCache {
10
+ #rows: PagedKVCache[] = [];
11
+ #rope?: MlxArray;
12
+ #ropeOffset = -1;
13
+
14
+ constructor(readonly capacityTokens: number, readonly blockSize: number) {}
15
+ signature(): string { return "kv:paged-rows"; }
16
+ get batchSize(): number { return this.#rows.length; }
17
+ get offset(): number { return Math.max(0, ...this.#rows.map(row => row.offset)); }
18
+ get rowOffsets(): number[] { return this.#rows.map(row => row.offset); }
19
+ get leftPad(): number[] { const width = this.offset; return this.#rows.map(row => width - row.offset); }
20
+ get ropeOffsetArr(): MlxArray | undefined {
21
+ if (this.#rows.length < 2) return undefined;
22
+ if (!this.#rope || this.#ropeOffset !== this.offset) {
23
+ this.releaseRopeArr();
24
+ this.#rope = MlxArray.fromInt32(Int32Array.from(this.rowOffsets), [this.batchSize]);
25
+ this.#ropeOffset = this.offset;
26
+ }
27
+ return this.#rope;
28
+ }
29
+ releaseRopeArr(): void { this.#rope?.dispose(); this.#rope = undefined; this.#ropeOffset = -1; }
30
+ makeEmptyBatch(): PagedKvRows { return new PagedKvRows(this.capacityTokens, this.blockSize); }
31
+
32
+ mergeRows(rows: readonly Cache[]): void {
33
+ const next: PagedKVCache[] = [];
34
+ try {
35
+ for (const source of rows) {
36
+ if (source instanceof PagedKvRows) {
37
+ for (const row of source.#rows) next.push(row.clone());
38
+ } else if (source instanceof PagedKVCache) next.push(source.clone());
39
+ else throw new Error(`Paged rows cannot merge ${source.signature()}`);
40
+ }
41
+ } catch (error) { disposeResources(next); throw error; }
42
+ const previous = this.#rows; this.#rows = next;
43
+ this.releaseRopeArr(); disposeResources(previous);
44
+ }
45
+ extractRow(row: number): PagedKVCache { return this.#rows[row]!.clone(); }
46
+ filterRows(keep: readonly number[]): void {
47
+ const previous = this.#rows;
48
+ this.#rows = keep.map(row => previous[row]!);
49
+ const retained = new Set(keep);
50
+ this.releaseRopeArr();
51
+ disposeResources(previous.filter((_, row) => !retained.has(row)));
52
+ }
53
+
54
+ updateAndFetch(keys: MlxArray, values: MlxArray): [MlxArray, MlxArray] {
55
+ if (this.#rows.length === 1) return this.#rows[0]!.updateAndFetch(keys, values);
56
+ const fetched: Array<{ keys: MlxArray; values: MlxArray }> = [];
57
+ try {
58
+ for (const [row, storage] of this.#rows.entries()) {
59
+ using k = keys.slice([row, 0, 0, 0], [row + 1, ...keys.shape.slice(1)]);
60
+ using v = values.slice([row, 0, 0, 0], [row + 1, ...values.shape.slice(1)]);
61
+ const [nextK, nextV] = storage.updateAndFetch(k, v);
62
+ fetched.push({ keys: nextK, values: nextV });
63
+ }
64
+ const result = mergeKVRows(fetched);
65
+ return [result.keys, result.values];
66
+ } finally { disposeResources(fetched.flatMap(row => [row.keys, row.values])); }
67
+ }
68
+ makeMask(tokens: number, window: number | null): Mask {
69
+ if (this.#rows.length === 1) return this.#rows[0]!.makeMask(tokens, window);
70
+ return { mode: "array", arr: buildBatchedDecodeMask(this.batchSize, tokens,
71
+ this.offset + tokens, this.leftPad, window) };
72
+ }
73
+ state(): MlxArray[] { return this.#rows.flatMap(row => row.state()); }
74
+ isTrimmable(): boolean { return true; }
75
+ trim(tokens: number): void { for (const row of this.#rows) row.trim(tokens); this.releaseRopeArr(); }
76
+ projectedBytes(tokens: number): number {
77
+ const pool = this.#rows.find(row => row.pool)?.pool;
78
+ if (!pool) return 0;
79
+ const bytesPerBlock = (pool.keys.nbytes + pool.values.nbytes) / pool.numBlocks;
80
+ return poolBlocksFor(tokens, this.blockSize) * bytesPerBlock;
81
+ }
82
+ dispose(): void { const rows = this.#rows; this.#rows = []; this.releaseRopeArr(); disposeResources(rows); }
83
+ }
@@ -1,7 +1,7 @@
1
1
  // Paged KV cache — vLLM-style block-pool storage behind the standard Cache
2
2
  // interface (docs/design/kv-cache.md). OPTIONAL and default-off
3
- // (`--paged-kv` / GenerateOptions.pagedKv); v1 scope is serial batch=1
4
- // decode on Gemma4-family plain full-attention layers, bf16 only.
3
+ // (`--paged-kv` / GenerateOptions.pagedKv); Gemma4-family plain full-attention
4
+ // layers, bf16 only. PagedKvRows composes independent pools at B1/B>1.
5
5
  //
6
6
  // Storage: K/V live in fixed-size per-layer pool tensors
7
7
  // [numBlocks, H_kv, blockSize, headDim]; a host-side block table maps the
@@ -17,9 +17,8 @@
17
17
  // Motivation, which says this honestly).
18
18
  //
19
19
  // Deliberately NOT a KVCache subclass (the TurboQuantKVCache reasoning):
20
- // every instanceof gate in the tree — CompiledDecode.supports, the batch
21
- // gateway's #modelCachesBatchable, generated forwards' #matches() — must
22
- // EXCLUDE paged caches and fall back to the monolith/uncompiled path.
20
+ // compiled graph adapters exclude the changing block-list geometry. The
21
+ // cache-layout binding supplies dynamic rows independently of scheduling.
23
22
 
24
23
  import { MlxArray } from "../../mlx/array";
25
24
  import { Dtype } from "../../mlx/ffi";
@@ -58,16 +57,24 @@ export class BlockPool {
58
57
  * reach v1's gemma4 scope, but the pool stays shape-honest). */
59
58
  vHeadDim?: number;
60
59
  dtype: Dtype;
61
- }) {
60
+ }, source?: BlockPool) {
62
61
  this.numBlocks = opts.numBlocks;
63
62
  this.blockSize = opts.blockSize;
64
- this.keys = ops.zeros(
63
+ this.keys = source ? ops.contiguous(source.keys) : ops.zeros(
65
64
  [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.headDim], opts.dtype);
66
- this.values = ops.zeros(
65
+ this.values = source ? ops.contiguous(source.values) : ops.zeros(
67
66
  [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.vHeadDim ?? opts.headDim],
68
67
  opts.dtype);
69
68
  // LIFO free list, low indices first — deterministic layout for tests.
70
- this.#free = Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
69
+ this.#free = source ? [...source.#free] : Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
70
+ }
71
+
72
+ /** Independent ownership and allocation metadata over immutable MLX buffers.
73
+ * Functional writes detach storage while another snapshot retains it. */
74
+ clone(): BlockPool {
75
+ return new BlockPool({ numBlocks: this.numBlocks, blockSize: this.blockSize,
76
+ numKvHeads: this.keys.shape[1]!, headDim: this.keys.shape[3]!,
77
+ vHeadDim: this.values.shape[3]!, dtype: this.keys.dtype }, this);
71
78
  }
72
79
 
73
80
  get freeBlocks(): number {
@@ -134,8 +141,7 @@ export class BlockPool {
134
141
  * shape): head count / head dims / dtype come from the first k/v pair,
135
142
  * so the wiring (maybePageKv) needs no per-model shape plumbing. */
136
143
  export class PagedKVCache implements Cache {
137
- /** Distinct kind: paged layout is serial-only and never merges into a
138
- * batch, so no capability guard should ever match it as plain. */
144
+ /** Distinct storage kind; the row adapter preserves its block pools. */
139
145
  signature(): string { return "kv:paged"; }
140
146
  /** Matches KVCache.STEP: v1's growth granularity is a permutation of
141
147
  * today's 256-token step into fixed reusable slots, not a new tuning
@@ -155,6 +161,14 @@ export class PagedKVCache implements Cache {
155
161
  readonly blockSize: number,
156
162
  ) {}
157
163
 
164
+ clone(): PagedKVCache {
165
+ const copy = new PagedKVCache(this.capacityTokens, this.blockSize);
166
+ copy.offset = this.offset;
167
+ copy.blockTable = [...this.blockTable];
168
+ copy.pool = this.pool?.clone() ?? null;
169
+ return copy;
170
+ }
171
+
158
172
  get #blockSize(): number {
159
173
  return this.blockSize;
160
174
  }
package/src/mlx/array.ts CHANGED
@@ -232,7 +232,7 @@ export class MlxArray {
232
232
  );
233
233
  }
234
234
 
235
- /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
235
+ /** Raw bytes of a supported floating-point or integer array (copy). */
236
236
  rawBytes(): Uint8Array {
237
237
  return this.rawBytesView().slice();
238
238
  }
@@ -253,6 +253,7 @@ export class MlxArray {
253
253
  : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
254
254
  : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
255
255
  : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
256
+ : dt === Dtype.int32 ? C.mlx_array_data_int32(this.handle)
256
257
  : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
257
258
  : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
258
259
  : null;
@@ -0,0 +1,73 @@
1
+ import { createHash, randomUUID } from "node:crypto";
2
+ import { mkdir, readFile, realpath, rename, rm, stat, writeFile } from "node:fs/promises";
3
+ import { homedir } from "node:os";
4
+ import { join } from "node:path";
5
+
6
+ /** Identity calculation owns bytes; this interface owns memo persistence. */
7
+ export interface ArtifactIdentityStore {
8
+ get(key: string): Promise<string | undefined>;
9
+ put(key: string, digest: string): void;
10
+ }
11
+
12
+ export class FileArtifactIdentityStore implements ArtifactIdentityStore {
13
+ readonly #memory = new Map<string, string>();
14
+ readonly #pending = new Set<Promise<void>>();
15
+ constructor(readonly directory: string) {}
16
+
17
+ async get(key: string): Promise<string | undefined> {
18
+ const memory = this.#memory.get(key);
19
+ if (memory) return memory;
20
+ try {
21
+ const digest = await readFile(join(this.directory, key), "utf8");
22
+ if (!/^[0-9a-f]{64}$/.test(digest)) return undefined;
23
+ this.#memory.set(key, digest);
24
+ return digest;
25
+ } catch { return undefined; }
26
+ }
27
+
28
+ put(key: string, digest: string): void {
29
+ this.#memory.set(key, digest);
30
+ const temporary = join(this.directory, `${key}.${process.pid}.${randomUUID()}.tmp`);
31
+ const write = async () => {
32
+ try {
33
+ await mkdir(this.directory, { recursive: true });
34
+ await writeFile(temporary, digest);
35
+ await rename(temporary, join(this.directory, key));
36
+ } catch { await rm(temporary, { force: true }).catch(() => {}); }
37
+ };
38
+ const pending = write().finally(() => this.#pending.delete(pending));
39
+ this.#pending.add(pending);
40
+ }
41
+
42
+ async flush(): Promise<void> { await Promise.all(this.#pending); }
43
+ }
44
+
45
+ const defaultStore = new FileArtifactIdentityStore(join(homedir(), ".cache", "mlx-bun", "artifact-identities"));
46
+ export interface ArtifactIdentityFile { name: string; path: string; }
47
+
48
+ async function revision(files: readonly ArtifactIdentityFile[]): Promise<string> {
49
+ return JSON.stringify(await Promise.all(files.map(async file => {
50
+ const path = await realpath(file.path), info = await stat(path, { bigint: true });
51
+ return [file.name, path, info.dev, info.ino, info.size, info.mtimeNs, info.ctimeNs].map(String);
52
+ })));
53
+ }
54
+
55
+ /** Preserve the existing seed/name/byte digest, reusing it while file revisions
56
+ * are unchanged. ctime catches same-size edits even when mtime is restored.
57
+ * A memo miss computes the full identity; persistence never delays its caller. */
58
+ export async function artifactIdentity(seed: string, input: readonly ArtifactIdentityFile[],
59
+ store: ArtifactIdentityStore = defaultStore): Promise<string> {
60
+ const files = [...input].sort((a, b) => a.name < b.name ? -1 : a.name > b.name ? 1 : 0);
61
+ const before = await revision(files);
62
+ const key = createHash("sha256").update(JSON.stringify([seed, before])).digest("hex");
63
+ const cached = await store.get(key);
64
+ if (cached) return cached;
65
+ const digest = createHash("sha256").update(seed);
66
+ for (const file of files) {
67
+ digest.update(file.name);
68
+ for await (const chunk of Bun.file(file.path).stream()) digest.update(chunk);
69
+ }
70
+ const value = digest.digest("hex");
71
+ if (before === await revision(files)) store.put(key, value);
72
+ return value;
73
+ }