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.
- package/docs/README.md +4 -1
- package/docs/archive/investigations/openjarvis-fit.md +198 -0
- package/docs/archive/investigations/qwen38-27b-campaign-log.md +1092 -0
- package/docs/design/batching.md +1321 -49
- package/docs/design/decode-speed-program.md +331 -1163
- package/docs/design/generic-model-support.md +79 -3
- package/docs/design/kv-cache.md +198 -54
- package/docs/design/speculative-decoding.md +387 -24
- package/docs/design/turboquant.md +123 -13
- package/docs/design/unified-engine-frontier-plan.md +144 -142
- package/docs/planning/PRODUCT_ROADMAP.md +16 -0
- package/docs/planning/release-notes-v0.4.0.md +45 -0
- package/docs/reference/benchmarks.md +2457 -6
- package/docs/reference/cli.md +41 -2
- package/docs/reference/distribution.md +16 -9
- package/docs/reference/environment.md +55 -5
- package/docs/reference/models.md +13 -4
- package/docs/reference/server-api.md +79 -25
- package/docs/reference/server-config.md +178 -84
- package/docs/reference/troubleshooting.md +13 -21
- package/package.json +1 -1
- package/src/backends/mlx/assistant-target.ts +46 -0
- package/src/backends/mlx/autoregressive-draft-rows.ts +120 -0
- package/src/backends/mlx/batch-group.ts +368 -224
- package/src/backends/mlx/cache-attachment.ts +67 -0
- package/src/backends/mlx/cache-layout.ts +31 -0
- package/src/backends/mlx/checkpoint-state.ts +53 -0
- package/src/backends/mlx/continuation-persistence.ts +106 -0
- package/src/backends/mlx/continuation-request.ts +43 -0
- package/src/backends/mlx/continuation.ts +77 -0
- package/src/backends/mlx/draft-target.ts +7 -22
- package/src/backends/mlx/full-kv-row-append.ts +32 -0
- package/src/backends/mlx/full-kv-row-donor.ts +54 -0
- package/src/backends/mlx/full-transitioning-kv-rows.ts +65 -0
- package/src/backends/mlx/gateway-binding.ts +77 -18
- package/src/backends/mlx/kv-codec.ts +9 -0
- package/src/backends/mlx/kv-maintenance.ts +133 -0
- package/src/backends/mlx/kv-tensor-rows.ts +212 -0
- package/src/backends/mlx/model-serving.ts +1 -1
- package/src/backends/mlx/prefill-cohort.ts +89 -0
- package/src/backends/mlx/prefill-rows.ts +200 -0
- package/src/backends/mlx/prefill.ts +1 -1
- package/src/backends/mlx/projected-context-rows.ts +72 -0
- package/src/backends/mlx/request-state-policy.ts +41 -0
- package/src/backends/mlx/rollback.ts +41 -3
- package/src/backends/mlx/row-sampling.ts +49 -0
- package/src/backends/mlx/serial-executor.ts +35 -66
- package/src/backends/mlx/speculative-group.ts +359 -0
- package/src/backends/mlx/speculative-round.ts +161 -0
- package/src/backends/mlx/state-rows.ts +104 -0
- package/src/backends/mlx/state-views.ts +21 -1
- package/src/backends/mlx/target-layout-capability.ts +35 -0
- package/src/backends/mlx/transitioning-kv-rows.ts +165 -0
- package/src/cli.ts +11 -36
- package/src/contracts/execution.ts +11 -0
- package/src/contracts/prefix-cache.ts +23 -0
- package/src/contracts/resources.ts +7 -0
- package/src/contracts/sampling.ts +23 -0
- package/src/contracts/scheduling.ts +18 -0
- package/src/engine/execution-plan.ts +18 -16
- package/src/engine/resources.ts +12 -1
- package/src/engine/scheduler.ts +10 -5
- package/src/fit.ts +29 -15
- package/src/generate.ts +77 -173
- package/src/inference/draft-acceptance.ts +41 -0
- package/src/inference/draft-output.ts +49 -0
- package/src/inference/rollback.ts +2 -2
- package/src/kv-scheme.ts +70 -12
- package/src/kv-store.ts +70 -16
- package/src/lab/paged-kv/paged-kv-rows.ts +83 -0
- package/src/lab/paged-kv/paged-kv.ts +25 -11
- package/src/mlx/array.ts +2 -1
- package/src/model/artifact-identity.ts +73 -0
- package/src/model/batched-kv.ts +235 -0
- package/src/model/batched-mask.ts +12 -4
- package/src/model/batched-quant.ts +8 -2
- package/src/model/batched-quantized-kv.ts +144 -0
- package/src/model/batched-rotating-quant.ts +76 -15
- package/src/model/batched-rotating-state.ts +76 -0
- package/src/model/batched-rotating.ts +64 -22
- package/src/model/batched-row-storage.ts +52 -0
- package/src/model/batched-ssm.ts +52 -0
- package/src/model/batched-turboquant-kv.ts +93 -0
- package/src/model/compiled-decode.ts +7 -1
- package/src/model/decoded-kv-donor.ts +12 -0
- package/src/model/delayed-quantized-kv.ts +72 -0
- package/src/model/delayed-rotating-quantized-kv.ts +93 -0
- package/src/model/delayed-turboquant-kv.ts +55 -0
- package/src/model/full-prefill-padding.ts +56 -0
- package/src/model/full-prefill-row.ts +105 -0
- package/src/model/gemma4-base.ts +162 -80
- package/src/model/gemma4.ts +34 -26
- package/src/model/generated/gemma4-12b.ts +4 -1
- package/src/model/generated/gemma4-26b.ts +4 -1
- package/src/model/generated/gemma4-e4b.ts +4 -1
- package/src/model/glm52-cache.ts +80 -4
- package/src/model/glm52-mla.ts +12 -6
- package/src/model/glm52.ts +19 -40
- package/src/model/kv-attention-view.ts +67 -0
- package/src/model/minicpm5.ts +11 -7
- package/src/model/qwen3-delta-state.ts +74 -0
- package/src/model/qwen3-delta.ts +59 -36
- package/src/model/qwen3_5.ts +82 -42
- package/src/model/rotating-kv-donor.ts +35 -0
- package/src/model/rotating-kv-layout.ts +262 -0
- package/src/model/rotating-row-transaction.ts +35 -0
- package/src/model/speculative-rotating-kv.ts +101 -0
- package/src/model/ssm-prefill-padding.ts +39 -0
- package/src/model/trellis-linear.ts +2 -2
- package/src/model/turboquant-codec.ts +83 -0
- package/src/model/universal/dense.ts +14 -13
- package/src/preflight.ts +31 -15
- package/src/prompt-cache.ts +42 -44
- package/src/sampler.ts +36 -25
- package/src/serve/chat-stage.ts +2 -2
- package/src/serve/generation-gateway.ts +39 -30
- package/src/serve/model-host.ts +3 -4
- package/src/serve/request-plan.ts +4 -3
- package/src/serve/request-prep.ts +2 -4
- package/src/serve/text-completion-stage.ts +2 -2
- package/src/server.ts +64 -108
- package/src/spec/assistant-rows.ts +92 -0
- package/src/spec/assistant-source.ts +28 -18
- package/src/spec/cache-identity.ts +22 -0
- package/src/spec/deepspec-rows.ts +29 -0
- package/src/spec/deepspec-source.ts +12 -10
- package/src/spec/dflash-source.ts +28 -4
- package/src/spec/drafter.ts +153 -142
- package/src/spec/dspark/deepspec-module.ts +95 -120
- package/src/spec/dspark/module-dflash.ts +126 -170
- package/src/spec/glm52-mtp-graph.ts +29 -0
- package/src/spec/glm52-mtp-rows.ts +132 -0
- package/src/spec/glm52-mtp-source.ts +53 -136
- package/src/spec/glm52-mtp-state.ts +33 -0
- package/src/spec/ngram-source.ts +92 -20
- package/src/spec/projected-draft-rows.ts +73 -0
- package/src/spec/qwen-mtp-module.ts +171 -0
- package/src/spec/qwen-mtp-rows.ts +219 -0
- package/src/spec/qwen-mtp-source.ts +81 -226
- package/src/spec/qwen-mtp-state.ts +30 -0
- package/src/spec/serve-loop.ts +70 -68
- package/src/spec/source.ts +80 -15
- package/src/spec/two-model-rows.ts +61 -0
- package/src/spec/two-model.ts +43 -101
- package/src/ssd-cache.ts +33 -21
- package/src/ssd-durability.ts +11 -7
- package/src/web/app.js +57 -57
- package/src/backends/mlx/memory-guard.ts +0 -40
- package/src/spec/prefix-state.ts +0 -45
- /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
|
-
|
|
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(), ...
|
|
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
|
-
|
|
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);
|
|
4
|
-
//
|
|
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
|
-
//
|
|
21
|
-
//
|
|
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
|
|
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
|
|
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
|
+
}
|