mlx-bun 0.3.0 → 0.5.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/README.md +57 -32
- package/docs/README.md +58 -42
- package/docs/archive/investigations/benchmark-ledger-through-v0.4.0.md +3617 -0
- package/docs/archive/investigations/openjarvis-fit.md +198 -0
- package/docs/archive/investigations/qwen38-27b-campaign-log.md +1092 -0
- package/docs/archive/planning/release-notes-v0.4.0.md +45 -0
- package/docs/archive/planning/release-notes-v0.5.0.md +37 -0
- package/docs/design/batching.md +1535 -50
- package/docs/design/decode-speed-program.md +394 -1171
- package/docs/design/generic-model-support.md +195 -3
- package/docs/design/kv-cache.md +538 -225
- package/docs/design/speculative-decoding.md +552 -50
- package/docs/design/turboquant.md +140 -17
- package/docs/design/unified-engine-frontier-plan.md +259 -150
- package/docs/planning/PRODUCT_ROADMAP.md +21 -6
- package/docs/planning/public-readiness-handoff.md +244 -0
- package/docs/planning/public-readiness.md +149 -0
- package/docs/reference/benchmarks.md +3537 -1112
- package/docs/reference/cli.md +102 -2
- package/docs/reference/distribution.md +45 -13
- package/docs/reference/environment.md +159 -13
- package/docs/reference/library-api.md +18 -4
- package/docs/reference/memory.md +1 -1
- package/docs/reference/models.md +28 -7
- package/docs/reference/server-api.md +257 -33
- package/docs/reference/server-config.md +372 -151
- package/docs/reference/troubleshooting.md +13 -21
- package/package.json +2 -2
- package/src/audio/audiotoolbox.ts +112 -0
- package/src/audio/conformer.ts +8 -0
- package/src/audio/mic-capture.ts +98 -0
- package/src/audio/silero-vad.ts +379 -0
- package/src/audio/whisper-decode.ts +774 -0
- package/src/audio/whisper-fast.ts +732 -0
- package/src/audio/whisper-format.ts +73 -0
- package/src/audio/whisper-mel.ts +185 -0
- package/src/audio/whisper-timing.ts +274 -0
- package/src/audio/whisper-tokenizer.ts +239 -0
- package/src/audio/whisper-transcribe.ts +410 -0
- package/src/backends/mlx/assistant-target.ts +46 -0
- package/src/backends/mlx/autoregressive-draft-rows.ts +120 -0
- package/src/backends/mlx/autoregressive.ts +22 -1
- package/src/backends/mlx/batch-group.ts +464 -236
- 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 +107 -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/encoder-cache.ts +25 -0
- package/src/backends/mlx/fill-append.ts +21 -8
- package/src/backends/mlx/fill-group.ts +497 -0
- 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 +122 -19
- package/src/backends/mlx/kv-codec.ts +9 -0
- package/src/backends/mlx/kv-maintenance.ts +161 -0
- package/src/backends/mlx/kv-tensor-rows.ts +212 -0
- package/src/backends/mlx/mixed-iteration.ts +62 -0
- package/src/backends/mlx/model-prompt.ts +46 -32
- package/src/backends/mlx/model-serving.ts +5 -4
- package/src/backends/mlx/paged-cache-codec.ts +27 -0
- package/src/backends/mlx/prefill-cohort.ts +100 -0
- package/src/backends/mlx/prefill-policy.ts +26 -0
- package/src/backends/mlx/prefill-rows.ts +244 -0
- package/src/backends/mlx/prefill.ts +1 -1
- package/src/backends/mlx/projected-context-rows.ts +72 -0
- package/src/backends/mlx/prompt-input.ts +29 -0
- package/src/backends/mlx/qwen-prompt-input.ts +40 -0
- package/src/backends/mlx/request-state-policy.ts +54 -0
- package/src/backends/mlx/rollback.ts +45 -4
- package/src/backends/mlx/row-sampling.ts +49 -0
- package/src/backends/mlx/serial-executor.ts +51 -73
- package/src/backends/mlx/speculative-group.ts +438 -0
- package/src/backends/mlx/speculative-round.ts +161 -0
- package/src/backends/mlx/speculative.ts +24 -7
- 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 +173 -0
- package/src/cli.ts +333 -36
- package/src/contracts/execution.ts +21 -0
- package/src/contracts/object-cache.ts +8 -0
- package/src/contracts/pi-provider.ts +1 -1
- package/src/contracts/pi.ts +3 -1
- package/src/contracts/prefix-cache.ts +33 -0
- package/src/contracts/resources.ts +7 -0
- package/src/contracts/sampling.ts +23 -0
- package/src/contracts/scheduling.ts +27 -0
- package/src/engine/execution-plan.ts +38 -25
- package/src/engine/execution-tasks.ts +44 -0
- package/src/engine/namespaced-cache.ts +13 -0
- package/src/engine/resources.ts +12 -1
- package/src/engine/scheduler.ts +32 -7
- package/src/fill/fill-session.ts +14 -3
- package/src/fit.ts +29 -15
- package/src/generate.ts +105 -195
- package/src/grammar.ts +13 -0
- package/src/inference/draft-acceptance.ts +41 -0
- package/src/inference/draft-output.ts +49 -0
- package/src/inference/prefill.ts +6 -0
- package/src/inference/rollback.ts +2 -2
- package/src/kv-scheme.ts +92 -12
- package/src/kv-store.ts +183 -93
- package/src/lab/README.md +5 -1
- package/src/lab/paged-kv/paged-attention.ts +117 -0
- package/src/lab/paged-kv/paged-kv-rows.ts +111 -0
- package/src/lab/paged-kv/paged-kv.ts +126 -46
- package/src/md-modules.d.ts +5 -0
- package/src/mlx/array.ts +39 -1
- package/src/mlx/ffi.ts +19 -0
- package/src/mlx/normalized-argmax.ts +116 -0
- package/src/mlx/ops.ts +97 -10
- 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 +94 -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 +57 -0
- package/src/model/factory.ts +26 -0
- package/src/model/fingerprint.ts +12 -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 +183 -82
- package/src/model/gemma4.ts +116 -28
- 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 +70 -0
- package/src/model/minicpm5.ts +11 -7
- package/src/model/profile.ts +29 -4
- package/src/model/quantized-append-attention.ts +31 -0
- package/src/model/qwen3-delta-state.ts +74 -0
- package/src/model/qwen3-delta.ts +59 -36
- package/src/model/qwen3-mrope.ts +7 -9
- package/src/model/qwen3_5.ts +177 -75
- 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/support.ts +8 -0
- package/src/model/token-groups.ts +66 -0
- package/src/model/trellis-linear.ts +4 -4
- package/src/model/turboquant-codec.ts +89 -0
- package/src/model/universal/dense.ts +14 -13
- package/src/model/whisper.ts +435 -0
- package/src/native/mic_capture.swift +92 -0
- package/src/native-pack.ts +6 -5
- package/src/pi-web.ts +9 -1
- package/src/preflight.ts +31 -15
- package/src/prompt-cache.ts +306 -86
- package/src/runtime-config.ts +1 -5
- package/src/safetensors.ts +9 -2
- package/src/sampler.ts +37 -25
- package/src/serve/audio-routes.ts +323 -0
- package/src/serve/chat-request.ts +2 -0
- package/src/serve/chat-stage.ts +19 -14
- package/src/serve/completion-executor.ts +9 -1
- package/src/serve/discovery-routes.ts +17 -1
- package/src/serve/generated-token-history.ts +57 -0
- package/src/serve/generation-gateway.ts +88 -38
- package/src/serve/model-binding.ts +7 -1
- package/src/serve/model-host.ts +3 -4
- package/src/serve/preparation.ts +3 -3
- package/src/serve/prompt-response-trace.ts +14 -0
- package/src/serve/request-plan.ts +6 -3
- package/src/serve/request-prep.ts +40 -25
- package/src/serve/text-completion-stage.ts +2 -2
- package/src/serve/transcription-server.ts +64 -0
- package/src/serve/transcription-service.ts +502 -0
- package/src/server.ts +185 -261
- package/src/spec/assistant-rows.ts +93 -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 +54 -136
- package/src/spec/glm52-mtp-state.ts +33 -0
- package/src/spec/ngram-source.ts +115 -21
- 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 +230 -0
- package/src/spec/qwen-mtp-source.ts +84 -226
- package/src/spec/qwen-mtp-state.ts +30 -0
- package/src/spec/serve-loop.ts +73 -69
- package/src/spec/source.ts +93 -15
- package/src/spec/two-model-rows.ts +62 -0
- package/src/spec/two-model.ts +45 -103
- package/src/ssd-cache.ts +107 -36
- package/src/ssd-durability.ts +29 -27
- package/src/storage/host-buffer.ts +24 -0
- package/src/storage/kv-writer.ts +83 -0
- package/src/storage/kv-writer.worker.js +264 -0
- package/src/storage/retention-policy.ts +37 -0
- package/src/tiered-prompt-cache.ts +45 -0
- package/src/vision/embedder.ts +8 -0
- package/src/vision/prompt.ts +101 -65
- package/src/vision/qwen3vl-prompt.ts +101 -64
- package/src/vision/qwen3vl-tower.ts +11 -0
- package/src/vision/siglip.ts +8 -0
- package/src/web/app.html +12 -0
- package/src/web/app.js +217 -57
- package/src/web/src/chat.ts +4 -1
- package/src/web/src/composer.ts +3 -0
- package/src/web/src/voice.ts +141 -0
- package/src/web/sw.js +1 -1
- 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
|
@@ -35,7 +35,10 @@ export const FINGERPRINT = "418e9adc386ea67c";
|
|
|
35
35
|
export let generatedForwardUses = 0;
|
|
36
36
|
|
|
37
37
|
function disposeSharedKv(s: SharedKv): void {
|
|
38
|
-
if (s.kind === "
|
|
38
|
+
if (s.kind === "view") {
|
|
39
|
+
s.attention.dispose();
|
|
40
|
+
s.offsetArr?.dispose();
|
|
41
|
+
} else if (s.kind === "plain") {
|
|
39
42
|
s.keys.dispose();
|
|
40
43
|
s.values.dispose();
|
|
41
44
|
} else {
|
package/src/model/glm52-cache.ts
CHANGED
|
@@ -11,6 +11,8 @@
|
|
|
11
11
|
// row's logical position without reconstructing per-head K/V.
|
|
12
12
|
|
|
13
13
|
import { MlxArray } from "../mlx/array";
|
|
14
|
+
import { FullPrefillPadding } from "./full-prefill-padding";
|
|
15
|
+
import type { PrefillPadding } from "./gemma4-base";
|
|
14
16
|
import { Dtype } from "../mlx/ffi";
|
|
15
17
|
import * as ops from "../mlx/ops";
|
|
16
18
|
import {
|
|
@@ -210,6 +212,8 @@ export class Glm52Cache implements BatchableCache {
|
|
|
210
212
|
batchSize: number | null = null;
|
|
211
213
|
rowOffsets: number[] = [];
|
|
212
214
|
leftPad: number[] = [];
|
|
215
|
+
#beforeRound: number[] | undefined;
|
|
216
|
+
readonly #padding = new FullPrefillPadding(1);
|
|
213
217
|
|
|
214
218
|
constructor(geometry: MLACacheGeometry) {
|
|
215
219
|
positiveInteger("MLA kvLoraRank", geometry.kvLoraRank);
|
|
@@ -482,6 +486,56 @@ export class Glm52Cache implements BatchableCache {
|
|
|
482
486
|
this.rowOffsets = this.rowOffsets.map((value) => value - n);
|
|
483
487
|
}
|
|
484
488
|
|
|
489
|
+
preparePrefill(padding: PrefillPadding): void { this.#padding.prepare(this, padding); }
|
|
490
|
+
finalizePrefill(): void {
|
|
491
|
+
const next = this.#padding.finalize(this.state(), this);
|
|
492
|
+
if (!next) return;
|
|
493
|
+
this.latent?.dispose(); this.rope?.dispose();
|
|
494
|
+
this.latent = next[0] ?? null; this.rope = next[1] ?? null;
|
|
495
|
+
if (this.dsa && next[2]) this.dsa.restoreState(next[2], this.offset);
|
|
496
|
+
}
|
|
497
|
+
|
|
498
|
+
specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
|
|
499
|
+
specRoundCommit(): void { this.#beforeRound = undefined; }
|
|
500
|
+
|
|
501
|
+
/** Keep each accepted prefix in compressed form. Right alignment lets the
|
|
502
|
+
* next shared append reuse the existing MLA/DSA attention geometry. */
|
|
503
|
+
specRoundRollback(keep: number | readonly number[]): void {
|
|
504
|
+
const nextOffsets = this.#beforeRound!.map((before, row) => before +
|
|
505
|
+
(typeof keep === "number" ? keep : keep[row]!));
|
|
506
|
+
const rejected = this.rowOffsets.map((offset, row) => offset - nextOffsets[row]!);
|
|
507
|
+
if (rejected.every(count => count === rejected[0])) {
|
|
508
|
+
const count = this.batchSize!;
|
|
509
|
+
this.trim(rejected[0]!);
|
|
510
|
+
if (this.batchSize === null) {
|
|
511
|
+
const empty = this.makeEmptyBatch();
|
|
512
|
+
try { this.mergeRows(Array.from({ length: count }, () => empty)); }
|
|
513
|
+
finally { empty.dispose(); }
|
|
514
|
+
}
|
|
515
|
+
this.#beforeRound = undefined; return;
|
|
516
|
+
}
|
|
517
|
+
const width = this.offset;
|
|
518
|
+
const pads = this.leftPad.map((pad, row) => pad + rejected[row]!);
|
|
519
|
+
const removable = Math.min(...pads);
|
|
520
|
+
using indices = ops.fromInt32(rejected.flatMap(count =>
|
|
521
|
+
Array.from({ length: width }, (_, column) => (column - count + width) % width)),
|
|
522
|
+
[this.batchSize!, width, 1]);
|
|
523
|
+
const next: MlxArray[] = [];
|
|
524
|
+
try {
|
|
525
|
+
for (const array of this.state()) {
|
|
526
|
+
using shifted = ops.takeAlongAxis(array, indices, 1);
|
|
527
|
+
next.push(shifted.slice([0, removable, 0], [this.batchSize!, width, array.shape[2]!]));
|
|
528
|
+
}
|
|
529
|
+
} catch (error) { for (const array of next) array.dispose(); throw error; }
|
|
530
|
+
this.latent!.dispose(); this.rope!.dispose();
|
|
531
|
+
this.latent = next[0]!; this.rope = next[1]!;
|
|
532
|
+
if (this.dsa) this.dsa.restoreState(next[2]!, width - removable);
|
|
533
|
+
this.offset = width - removable;
|
|
534
|
+
this.rowOffsets = nextOffsets;
|
|
535
|
+
this.leftPad = pads.map(pad => pad - removable);
|
|
536
|
+
this.#beforeRound = undefined;
|
|
537
|
+
}
|
|
538
|
+
|
|
485
539
|
/** Exact logical bytes of valid f32 compressed state. */
|
|
486
540
|
get byteLength(): number {
|
|
487
541
|
if (this.batchSize === null) return 0;
|
|
@@ -529,8 +583,16 @@ export class Glm52Cache implements BatchableCache {
|
|
|
529
583
|
) {
|
|
530
584
|
throw new Error("MLA mergeRows cache geometry does not match");
|
|
531
585
|
}
|
|
532
|
-
if (generic.batchSize === null || generic.offset === 0)
|
|
533
|
-
|
|
586
|
+
if (generic.batchSize === null || generic.offset === 0) {
|
|
587
|
+
const count = generic.batchSize ?? 1;
|
|
588
|
+
for (let row = 0; row < count; row++) {
|
|
589
|
+
states.push({ latent: ops.zeros([1, 0, this.kvLoraRank], Dtype.float32),
|
|
590
|
+
rope: ops.zeros([1, 0, this.ropeHeadDim], Dtype.float32),
|
|
591
|
+
dsa: this.dsa ? ops.zeros([1, 0, this.dsa.headDim], Dtype.float32) : null });
|
|
592
|
+
offsets.push(0);
|
|
593
|
+
}
|
|
594
|
+
continue;
|
|
595
|
+
}
|
|
534
596
|
for (let row = 0; row < generic.batchSize; row++) {
|
|
535
597
|
states.push(generic.fetchRow(row));
|
|
536
598
|
offsets.push(generic.rowOffset(row));
|
|
@@ -577,7 +639,10 @@ export class Glm52Cache implements BatchableCache {
|
|
|
577
639
|
this.offset = width;
|
|
578
640
|
this.rowOffsets = [...offsets];
|
|
579
641
|
this.leftPad = offsets.map((value) => width - value);
|
|
580
|
-
if (this.dsa)
|
|
642
|
+
if (this.dsa) {
|
|
643
|
+
if (width) this.dsa.restoreState(dsa!, width);
|
|
644
|
+
else { this.dsa.data = dsa!; this.dsa.batchSize = offsets.length; }
|
|
645
|
+
}
|
|
581
646
|
} finally {
|
|
582
647
|
for (const state of states) {
|
|
583
648
|
state.latent.dispose();
|
|
@@ -591,6 +656,7 @@ export class Glm52Cache implements BatchableCache {
|
|
|
591
656
|
extractRow(row: number): Glm52Cache {
|
|
592
657
|
if (this.batchSize === null) throw new Error("MLA cache is empty");
|
|
593
658
|
validateRowIndex(row, this.batchSize, "MLA extract");
|
|
659
|
+
if (this.rowOffsets[row] === 0) return this.makeEmptyBatch();
|
|
594
660
|
const state = this.fetchRow(row);
|
|
595
661
|
const copy = (array: MlxArray): MlxArray => ops.mulScalar(array, 1);
|
|
596
662
|
const latent = copy(state.latent);
|
|
@@ -666,7 +732,15 @@ export class Glm52Cache implements BatchableCache {
|
|
|
666
732
|
this.offset = nextWidth;
|
|
667
733
|
this.rowOffsets = nextOffsets;
|
|
668
734
|
this.leftPad = selectedPad.map((value) => value - removablePad);
|
|
669
|
-
|
|
735
|
+
this.#padding.filter(keep);
|
|
736
|
+
if (this.dsa) {
|
|
737
|
+
if (nextWidth) this.dsa.restoreState(dsa!, nextWidth);
|
|
738
|
+
else {
|
|
739
|
+
// A retained cold row still owns batch membership, even though the
|
|
740
|
+
// strict persisted-state restore contract requires nonempty tokens.
|
|
741
|
+
this.dsa.dispose(); this.dsa.data = dsa!; this.dsa.batchSize = keep.length;
|
|
742
|
+
}
|
|
743
|
+
}
|
|
670
744
|
}
|
|
671
745
|
|
|
672
746
|
projectedBytes(tokens: number): number {
|
|
@@ -709,6 +783,8 @@ export class Glm52Cache implements BatchableCache {
|
|
|
709
783
|
this.batchSize = null;
|
|
710
784
|
this.rowOffsets = [];
|
|
711
785
|
this.leftPad = [];
|
|
786
|
+
this.#beforeRound = undefined;
|
|
787
|
+
this.#padding.clear();
|
|
712
788
|
}
|
|
713
789
|
}
|
|
714
790
|
|
package/src/model/glm52-mla.ts
CHANGED
|
@@ -484,7 +484,7 @@ export class Glm52Mla {
|
|
|
484
484
|
}
|
|
485
485
|
return this.attendAbsorbed(projection, state, selectedPositions);
|
|
486
486
|
}
|
|
487
|
-
return this.attendReconstructed(projection, state);
|
|
487
|
+
return this.attendReconstructed(projection, state, cache.leftPad);
|
|
488
488
|
} finally {
|
|
489
489
|
disposeState(state);
|
|
490
490
|
disposeProjection(projection);
|
|
@@ -495,6 +495,7 @@ export class Glm52Mla {
|
|
|
495
495
|
attendReconstructed(
|
|
496
496
|
projection: Pick<Glm52MlaProjection, "qNope" | "qRope">,
|
|
497
497
|
state: Pick<MLACompressedState, "latent" | "rope">,
|
|
498
|
+
leftPadding: readonly number[] = [],
|
|
498
499
|
): MlxArray {
|
|
499
500
|
const [batch, tokens] = projection.qNope.shape;
|
|
500
501
|
const keys = state.latent.shape[1]!;
|
|
@@ -570,14 +571,19 @@ export class Glm52Mla {
|
|
|
570
571
|
|
|
571
572
|
const prefixTokens = keys - tokens!;
|
|
572
573
|
let maskedScores = scores;
|
|
573
|
-
|
|
574
|
-
|
|
575
|
-
|
|
574
|
+
const padded = leftPadding.some(value => value > 0);
|
|
575
|
+
if (tokens! > 1 || padded) {
|
|
576
|
+
const maskRows = padded ? batch! : 1;
|
|
577
|
+
const maskValues = new Float32Array(maskRows * tokens! * keys);
|
|
578
|
+
for (let row = 0; row < maskRows; row++) for (let queryToken = 0; queryToken < tokens!; queryToken++) {
|
|
576
579
|
const lastVisible = prefixTokens + queryToken;
|
|
580
|
+
const start = (row * tokens! + queryToken) * keys;
|
|
581
|
+
for (let keyToken = 0; keyToken < (leftPadding[row] ?? 0); keyToken++)
|
|
582
|
+
maskValues[start + keyToken] = -Infinity;
|
|
577
583
|
for (let keyToken = lastVisible + 1; keyToken < keys; keyToken++)
|
|
578
|
-
maskValues[
|
|
584
|
+
maskValues[start + keyToken] = -Infinity;
|
|
579
585
|
}
|
|
580
|
-
const mask = MlxArray.fromFloat32(maskValues, [
|
|
586
|
+
const mask = MlxArray.fromFloat32(maskValues, [maskRows, 1, tokens!, keys]);
|
|
581
587
|
maskedScores = ops.add(scores, mask);
|
|
582
588
|
mask.dispose();
|
|
583
589
|
scores.dispose();
|
package/src/model/glm52.ts
CHANGED
|
@@ -144,8 +144,6 @@ class Glm52DsaIndexer {
|
|
|
144
144
|
forceDense = false,
|
|
145
145
|
): DsaStep {
|
|
146
146
|
const [batch, tokens, hidden] = input.shape;
|
|
147
|
-
if (batch! > 1 && tokens !== 1)
|
|
148
|
-
throw new Error("batched DSA supports ordinary single-token decode only");
|
|
149
147
|
if (hidden !== this.config.hiddenSize)
|
|
150
148
|
throw new Error(`DSA hidden width ${hidden} != ${this.config.hiddenSize}`);
|
|
151
149
|
|
|
@@ -228,7 +226,10 @@ class Glm52DsaIndexer {
|
|
|
228
226
|
}
|
|
229
227
|
const contextLength = allKeys.shape[1]!;
|
|
230
228
|
let selectedPositions: MlxArray | Glm52MlaBatchedSelection | null;
|
|
231
|
-
if (
|
|
229
|
+
if (forceDense || contextLength <= this.config.indexTopk) {
|
|
230
|
+
state.selectFullDense(this.layer, contextLength);
|
|
231
|
+
selectedPositions = null;
|
|
232
|
+
} else if (batch! > 1) {
|
|
232
233
|
const rowSelections: Array<MlxArray | null> = [];
|
|
233
234
|
for (let row = 0; row < batch!; row++) {
|
|
234
235
|
const rowContext = (cache.rowOffsets[row] ?? 0) + tokens!;
|
|
@@ -269,15 +270,6 @@ class Glm52DsaIndexer {
|
|
|
269
270
|
scores.dispose();
|
|
270
271
|
}
|
|
271
272
|
selectedPositions = { rows: rowSelections };
|
|
272
|
-
} else if (forceDense) {
|
|
273
|
-
state.selectFullDense(this.layer, contextLength);
|
|
274
|
-
selectedPositions = null;
|
|
275
|
-
} else if (contextLength <= this.config.indexTopk) {
|
|
276
|
-
state.selectFull(
|
|
277
|
-
this.layer,
|
|
278
|
-
new Float32Array(contextLength),
|
|
279
|
-
);
|
|
280
|
-
selectedPositions = null;
|
|
281
273
|
} else {
|
|
282
274
|
if (tokens !== 1) {
|
|
283
275
|
if (allKeys !== key) allKeys.dispose();
|
|
@@ -533,6 +525,7 @@ export class Glm52DecoderLayer {
|
|
|
533
525
|
this.inputNorm,
|
|
534
526
|
this.config.rmsNormEps,
|
|
535
527
|
);
|
|
528
|
+
const batch = input.shape[0]!;
|
|
536
529
|
const tokens = input.shape[1]!;
|
|
537
530
|
const denseBenchmarkPrefill =
|
|
538
531
|
this.denseBenchmarkPrefillTokens > 0 &&
|
|
@@ -552,10 +545,10 @@ export class Glm52DecoderLayer {
|
|
|
552
545
|
for (let token = 0; token < tokens; token++) {
|
|
553
546
|
const row = normalized.slice(
|
|
554
547
|
[0, token, 0],
|
|
555
|
-
[
|
|
548
|
+
[batch, token + 1, this.config.hiddenSize],
|
|
556
549
|
);
|
|
557
550
|
let dsaKey: MlxArray | null = null;
|
|
558
|
-
let positions: MlxArray | null = null;
|
|
551
|
+
let positions: MlxArray | Glm52MlaBatchedSelection | null = null;
|
|
559
552
|
try {
|
|
560
553
|
if (this.dsa) {
|
|
561
554
|
const step = this.dsa.projectAndSelect(
|
|
@@ -564,18 +557,12 @@ export class Glm52DecoderLayer {
|
|
|
564
557
|
dsaState,
|
|
565
558
|
);
|
|
566
559
|
dsaKey = step.key;
|
|
567
|
-
if (
|
|
568
|
-
step.selectedPositions !== null &&
|
|
569
|
-
!(step.selectedPositions instanceof MlxArray)
|
|
570
|
-
) {
|
|
571
|
-
throw new Error("serial DSA verification received batched selection");
|
|
572
|
-
}
|
|
573
560
|
positions = step.selectedPositions;
|
|
574
561
|
} else {
|
|
575
|
-
positions =
|
|
576
|
-
|
|
577
|
-
|
|
578
|
-
|
|
562
|
+
positions = batch > 1
|
|
563
|
+
? { rows: Array.from({ length: batch }, (_, row) =>
|
|
564
|
+
dsaState.selectSharedPositions(this.layer, (cache.rowOffsets[row] ?? 0) + 1, row)) }
|
|
565
|
+
: dsaState.selectSharedPositions(this.layer, cache.offset + 1);
|
|
579
566
|
}
|
|
580
567
|
rows.push(this.mla.forward(
|
|
581
568
|
row,
|
|
@@ -609,7 +596,9 @@ export class Glm52DecoderLayer {
|
|
|
609
596
|
dsaKey = step.key;
|
|
610
597
|
positions = step.selectedPositions;
|
|
611
598
|
} else if (dsaState) {
|
|
612
|
-
positions = cache.
|
|
599
|
+
positions = denseBenchmarkPrefill || cache.offset + tokens <= this.config.indexTopk
|
|
600
|
+
? null
|
|
601
|
+
: cache.batchSize !== null && cache.batchSize > 1
|
|
613
602
|
? {
|
|
614
603
|
rows: cache.rowOffsets.map((offset, row) =>
|
|
615
604
|
dsaState.selectSharedPositions(
|
|
@@ -871,18 +860,13 @@ export class Glm52Model {
|
|
|
871
860
|
tokenCount > GLM52_DSA_SPARSE_DECODE_BATCH_MAX &&
|
|
872
861
|
firstOffset + tokenCount > this.glmConfig.indexTopk
|
|
873
862
|
) {
|
|
874
|
-
|
|
875
|
-
|
|
876
|
-
|
|
877
|
-
);
|
|
878
|
-
}
|
|
879
|
-
// Colibri selects independently for every query row. G2 preserves that
|
|
880
|
-
// behavior with a deliberately serial correctness fallback once a
|
|
881
|
-
// multi-token prefill crosses index_topk; G6 may vectorize it.
|
|
863
|
+
// Bound sparse-selection state for long prefills while retaining all
|
|
864
|
+
// request rows in each forward. Short verify windows keep the residual
|
|
865
|
+
// and expert graph batched across both requests and query positions.
|
|
882
866
|
const rows: MlxArray[] = [];
|
|
883
867
|
try {
|
|
884
868
|
for (let token = 0; token < tokenCount; token++) {
|
|
885
|
-
const row = ids.slice([0, token], [
|
|
869
|
+
const row = ids.slice([0, token], [ids.shape[0]!, token + 1]);
|
|
886
870
|
try {
|
|
887
871
|
rows.push(this.forwardHidden(row, cache));
|
|
888
872
|
} finally {
|
|
@@ -922,15 +906,10 @@ export class Glm52Model {
|
|
|
922
906
|
tokenCount > GLM52_DSA_SPARSE_DECODE_BATCH_MAX &&
|
|
923
907
|
firstOffset + tokenCount > this.glmConfig.indexTopk
|
|
924
908
|
) {
|
|
925
|
-
if (ids.shape.length !== 2 || ids.shape[0] !== 1) {
|
|
926
|
-
throw new Error(
|
|
927
|
-
"G3 DSA sparse prefill requires serial [1,T] token ids",
|
|
928
|
-
);
|
|
929
|
-
}
|
|
930
909
|
const rows: MlxArray[] = [];
|
|
931
910
|
try {
|
|
932
911
|
for (let token = 0; token < tokenCount; token++) {
|
|
933
|
-
const row = ids.slice([0, token], [
|
|
912
|
+
const row = ids.slice([0, token], [ids.shape[0]!, token + 1]);
|
|
934
913
|
try {
|
|
935
914
|
rows.push(await this.forwardHiddenAsync(row, cache));
|
|
936
915
|
} finally {
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
import type { MlxArray } from "../mlx/array";
|
|
2
|
+
import * as ops from "../mlx/ops";
|
|
3
|
+
import type { Cache, KvAttentionView, KvDonorAttention } from "./gemma4-base";
|
|
4
|
+
import { quantizedSdpa, disposeTriple } from "./gemma4-base";
|
|
5
|
+
import { quantizedAppendAttention } from "./quantized-append-attention";
|
|
6
|
+
|
|
7
|
+
export function captureKvDonorAttention(cache: Cache): KvDonorAttention {
|
|
8
|
+
if (cache.captureDonorAttention) return cache.captureDonorAttention();
|
|
9
|
+
const view = cache.captureDonorRows!();
|
|
10
|
+
return {
|
|
11
|
+
width: view.keys.shape[2]!, dtype: view.keys.dtype,
|
|
12
|
+
offsets: view.offsets, starts: view.starts, ends: view.ends,
|
|
13
|
+
attend: (q, scale, mask) => ops.sdpa(q, view.keys, view.values, scale, mask.mode, mask.arr),
|
|
14
|
+
dispose() { view.keys.dispose(); view.values.dispose(); },
|
|
15
|
+
};
|
|
16
|
+
}
|
|
17
|
+
|
|
18
|
+
export function quantizedDonorAttention(keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
|
|
19
|
+
validity: Pick<KvDonorAttention, "offsets" | "starts" | "ends">,
|
|
20
|
+
groupSize: number, bits: number): KvDonorAttention {
|
|
21
|
+
return {
|
|
22
|
+
width: keys.packed.shape[2]!, dtype: keys.scales.dtype, ...validity,
|
|
23
|
+
attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
|
|
24
|
+
dispose() { disposeTriple(keys); disposeTriple(values); },
|
|
25
|
+
};
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
export function captureKvAttention(cache: Cache, k: MlxArray, v: MlxArray): KvAttentionView {
|
|
29
|
+
if (cache.attentionState) return cache.attentionState.appendAndFetch(k, v);
|
|
30
|
+
const quantized = cache.quantizedAttention;
|
|
31
|
+
if (quantized) {
|
|
32
|
+
const { groupSize, bits } = quantized;
|
|
33
|
+
const [keys, values] = quantized.updateAndFetchQuantized(k, v);
|
|
34
|
+
return {
|
|
35
|
+
attend: (q, scale, mask, independentPositions) => independentPositions && q.shape[2]! > 1
|
|
36
|
+
? quantizedAppendAttention(q, keys, values, scale, groupSize, bits)
|
|
37
|
+
: quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
|
|
38
|
+
dispose() { disposeTriple(keys); disposeTriple(values); },
|
|
39
|
+
};
|
|
40
|
+
}
|
|
41
|
+
const [keys, values] = cache.updateAndFetch(k, v);
|
|
42
|
+
return {
|
|
43
|
+
attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
|
|
44
|
+
dispose() { keys.dispose(); values.dispose(); },
|
|
45
|
+
};
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
/** Takes ownership of already aligned per-row donor views. Each storage codec
|
|
50
|
+
* keeps its own attention arithmetic while the caller supplies one query batch. */
|
|
51
|
+
export function combineKvDonorAttention(rows: readonly KvDonorAttention[]): KvDonorAttention {
|
|
52
|
+
return { width: rows[0]!.width, dtype: rows[0]!.dtype,
|
|
53
|
+
offsets: rows.flatMap(row => [...row.offsets]), starts: rows.flatMap(row => [...row.starts]),
|
|
54
|
+
ends: rows.flatMap(row => [...row.ends]),
|
|
55
|
+
attend(q, scale, mask) {
|
|
56
|
+
const outputs: MlxArray[] = [];
|
|
57
|
+
try {
|
|
58
|
+
for (const [index, row] of rows.entries()) {
|
|
59
|
+
using query = q.slice([index, 0, 0, 0], [index + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
|
|
60
|
+
const from = mask.arr?.shape.map(() => 0), to = mask.arr ? [...mask.arr.shape] : undefined;
|
|
61
|
+
if (from && to && to.length === 4 && to[0]! > 1) { from[0] = index; to[0] = index + 1; }
|
|
62
|
+
using arr = mask.arr ? mask.arr.slice(from!, to!) : null;
|
|
63
|
+
outputs.push(row.attend(query, scale, { mode: mask.mode, arr }));
|
|
64
|
+
}
|
|
65
|
+
return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
|
|
66
|
+
} finally { for (const output of outputs) output.dispose(); }
|
|
67
|
+
},
|
|
68
|
+
dispose() { for (const row of rows) row.dispose(); },
|
|
69
|
+
};
|
|
70
|
+
}
|
package/src/model/minicpm5.ts
CHANGED
|
@@ -14,7 +14,6 @@ import {
|
|
|
14
14
|
KVCache,
|
|
15
15
|
LoraState,
|
|
16
16
|
QuantizedKVCache,
|
|
17
|
-
TurboQuantKVCache,
|
|
18
17
|
QuantizedEmbedding,
|
|
19
18
|
QuantizedLinear,
|
|
20
19
|
quantizedSdpa,
|
|
@@ -143,23 +142,28 @@ export class LlamaAttention {
|
|
|
143
142
|
q = disposing(q, ropeStep(q));
|
|
144
143
|
k = disposing(k, ropeStep(k));
|
|
145
144
|
let attn: MlxArray;
|
|
146
|
-
if (cache
|
|
147
|
-
const
|
|
145
|
+
if (cache.attentionState) {
|
|
146
|
+
const view = cache.attentionState.appendAndFetch(k, v);
|
|
147
|
+
k.dispose(); v.dispose();
|
|
148
|
+
try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
|
|
149
|
+
} else if (cache.quantizedAttention) {
|
|
150
|
+
const quantized = cache.quantizedAttention;
|
|
151
|
+
const [keys, values] = quantized.updateAndFetchQuantized(k, v);
|
|
148
152
|
k.dispose();
|
|
149
153
|
v.dispose();
|
|
150
|
-
attn = quantizedSdpa(q, keys, values, this.scale, mask,
|
|
154
|
+
attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
|
|
151
155
|
disposeTriple(keys);
|
|
152
156
|
disposeTriple(values);
|
|
153
|
-
} else if (cache
|
|
157
|
+
} else if (cache.rotatedValueAttention) {
|
|
154
158
|
// Deferred-V: sdpa runs on rotated values, then one InvFWHT on the
|
|
155
159
|
// output (linear in V) — see TurboQuantKVCache.updateAndFetchDeferredV.
|
|
156
|
-
const [keys, values] = cache.updateAndFetchDeferredV(k, v);
|
|
160
|
+
const [keys, values] = cache.rotatedValueAttention.updateAndFetchDeferredV(k, v);
|
|
157
161
|
k.dispose();
|
|
158
162
|
v.dispose();
|
|
159
163
|
const rotated = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
|
|
160
164
|
keys.dispose();
|
|
161
165
|
values.dispose();
|
|
162
|
-
attn = tqUnrotateValues(rotated);
|
|
166
|
+
attn = cache.rotatedValueAttention.captureValueTransform?.()(rotated) ?? tqUnrotateValues(rotated);
|
|
163
167
|
rotated.dispose();
|
|
164
168
|
} else {
|
|
165
169
|
const [keys, values] = cache.updateAndFetch(k, v);
|
package/src/model/profile.ts
CHANGED
|
@@ -7,6 +7,7 @@ import {
|
|
|
7
7
|
isQwen35Config,
|
|
8
8
|
isQwen3Config,
|
|
9
9
|
isQwen3MoeConfig,
|
|
10
|
+
isWhisperConfig,
|
|
10
11
|
} from "./support";
|
|
11
12
|
import { GENERATED } from "./generated";
|
|
12
13
|
import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
|
|
@@ -16,6 +17,7 @@ export const ENGINE_CAPABILITIES = Object.freeze([
|
|
|
16
17
|
"colibri-container",
|
|
17
18
|
"diffusion",
|
|
18
19
|
"diffusion-gemma-graph",
|
|
20
|
+
"encoder-decoder",
|
|
19
21
|
"gemma4-graph",
|
|
20
22
|
"generated-graph",
|
|
21
23
|
"glm5.2-graph",
|
|
@@ -30,6 +32,7 @@ export const ENGINE_CAPABILITIES = Object.freeze([
|
|
|
30
32
|
"streamed-experts",
|
|
31
33
|
"universal-dense-graph",
|
|
32
34
|
"vision-sidecar",
|
|
35
|
+
"whisper-graph",
|
|
33
36
|
] as const);
|
|
34
37
|
|
|
35
38
|
export type EngineCapability = typeof ENGINE_CAPABILITIES[number];
|
|
@@ -37,6 +40,7 @@ export type FidelityTier = "l1" | "l2" | "l3";
|
|
|
37
40
|
|
|
38
41
|
export type FidelityTarget =
|
|
39
42
|
| Readonly<{ tier: "l1"; oracle: "mlx-lm"; claim: "bit-exact" }>
|
|
43
|
+
| Readonly<{ tier: "l1"; oracle: "mlx-whisper"; claim: "bit-exact" }>
|
|
40
44
|
| Readonly<{ tier: "l2"; oracle: "mlx-optiq"; claim: "bit-exact" }>
|
|
41
45
|
| Readonly<{ tier: "l3"; oracle: null; claim: "measured" }>;
|
|
42
46
|
|
|
@@ -49,8 +53,9 @@ export type ModelGraph =
|
|
|
49
53
|
| "qwen3-moe"
|
|
50
54
|
| "diffusion-gemma"
|
|
51
55
|
| "glm5.2"
|
|
52
|
-
| "universal-dense"
|
|
53
|
-
|
|
56
|
+
| "universal-dense"
|
|
57
|
+
| "whisper";
|
|
58
|
+
export type GenerationLoop = "autoregressive" | "diffusion" | "encoder-decoder";
|
|
54
59
|
export type ModelSpecialization = "artifact" | "dedicated" | "generated" | "generic";
|
|
55
60
|
|
|
56
61
|
export interface ModelExecutionComposition {
|
|
@@ -105,6 +110,11 @@ export interface ResolveModelProfileOptions {
|
|
|
105
110
|
const L1: FidelityTarget = Object.freeze({
|
|
106
111
|
tier: "l1", oracle: "mlx-lm", claim: "bit-exact",
|
|
107
112
|
});
|
|
113
|
+
/** L1 for speech: the capability's own reference implementation (mlx-whisper)
|
|
114
|
+
* plays the mlx-lm role — no mlx-lm arm exists for Whisper. */
|
|
115
|
+
const L1_WHISPER: FidelityTarget = Object.freeze({
|
|
116
|
+
tier: "l1", oracle: "mlx-whisper", claim: "bit-exact",
|
|
117
|
+
});
|
|
108
118
|
const L2: FidelityTarget = Object.freeze({
|
|
109
119
|
tier: "l2", oracle: "mlx-optiq", claim: "bit-exact",
|
|
110
120
|
});
|
|
@@ -231,6 +241,14 @@ const FAMILY_PROFILES = {
|
|
|
231
241
|
loader: "colibri", graph: "glm5.2", loop: "autoregressive", specialization: "dedicated",
|
|
232
242
|
},
|
|
233
243
|
}),
|
|
244
|
+
whisper: freezeProfile({
|
|
245
|
+
id: "whisper-dedicated",
|
|
246
|
+
fidelity: L1_WHISPER,
|
|
247
|
+
requiredCapabilities: ["safetensors", "encoder-decoder", "whisper-graph"],
|
|
248
|
+
execution: {
|
|
249
|
+
loader: "safetensors", graph: "whisper", loop: "encoder-decoder", specialization: "dedicated",
|
|
250
|
+
},
|
|
251
|
+
}),
|
|
234
252
|
universal: freezeProfile({
|
|
235
253
|
id: "universal-dense",
|
|
236
254
|
fidelity: L1,
|
|
@@ -263,6 +281,7 @@ function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
|
|
|
263
281
|
if (isQwen35Config(config)) return FAMILY_PROFILES.qwen35;
|
|
264
282
|
if (isQwen3MoeConfig(config)) return FAMILY_PROFILES.qwen3Moe;
|
|
265
283
|
if (isQwen3Config(config)) return FAMILY_PROFILES.qwen3;
|
|
284
|
+
if (isWhisperConfig(config)) return FAMILY_PROFILES.whisper;
|
|
266
285
|
if (config.modelType.startsWith("gemma4"))
|
|
267
286
|
return GENERATED.has(fingerprint) ? FAMILY_PROFILES.gemma4Generated : FAMILY_PROFILES.gemma4;
|
|
268
287
|
if (genericArgsFor(config)) return FAMILY_PROFILES.universal;
|
|
@@ -271,7 +290,7 @@ function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
|
|
|
271
290
|
throw new Error(
|
|
272
291
|
`unsupported model_type "${config.modelType}"` +
|
|
273
292
|
(arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
|
|
274
|
-
` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
|
|
293
|
+
` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, whisper, MiniCPM5;` +
|
|
275
294
|
` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
|
|
276
295
|
);
|
|
277
296
|
}
|
|
@@ -301,6 +320,7 @@ const GRAPH_METADATA: Readonly<Record<ModelGraph, GraphMetadata>> = Object.freez
|
|
|
301
320
|
accepts: isGlm52Config,
|
|
302
321
|
capabilities: ["glm5.2-graph", "streamed-experts"],
|
|
303
322
|
},
|
|
323
|
+
"whisper": { accepts: isWhisperConfig, capabilities: ["whisper-graph"] },
|
|
304
324
|
"universal-dense": {
|
|
305
325
|
accepts: (config) => {
|
|
306
326
|
try { return genericArgsFor(config) !== null; } catch { return false; }
|
|
@@ -318,6 +338,7 @@ function executionCapabilities(profile: ModelProfile): EngineCapability[] {
|
|
|
318
338
|
? "colibri-container"
|
|
319
339
|
: "safetensors"];
|
|
320
340
|
if (profile.execution.loop === "diffusion") required.push("diffusion");
|
|
341
|
+
else if (profile.execution.loop === "encoder-decoder") required.push("encoder-decoder");
|
|
321
342
|
else required.push("autoregressive");
|
|
322
343
|
required.push(...GRAPH_METADATA[profile.execution.graph].capabilities);
|
|
323
344
|
if (profile.execution.specialization === "generated") required.push("generated-graph");
|
|
@@ -337,6 +358,8 @@ function validateProfile(profile: ModelProfile): void {
|
|
|
337
358
|
const fidelityOk =
|
|
338
359
|
(profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-lm" &&
|
|
339
360
|
profile.fidelity.claim === "bit-exact") ||
|
|
361
|
+
(profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-whisper" &&
|
|
362
|
+
profile.fidelity.claim === "bit-exact" && profile.execution.graph === "whisper") ||
|
|
340
363
|
(profile.fidelity.tier === "l2" && profile.fidelity.oracle === "mlx-optiq" &&
|
|
341
364
|
profile.fidelity.claim === "bit-exact") ||
|
|
342
365
|
(profile.fidelity.tier === "l3" && profile.fidelity.oracle === null &&
|
|
@@ -345,7 +368,9 @@ function validateProfile(profile: ModelProfile): void {
|
|
|
345
368
|
const expectedLoader = profile.execution.graph === "glm5.2" ? "colibri" : "safetensors";
|
|
346
369
|
const expectedLoop = profile.execution.graph === "diffusion-gemma"
|
|
347
370
|
? "diffusion"
|
|
348
|
-
: "
|
|
371
|
+
: profile.execution.graph === "whisper"
|
|
372
|
+
? "encoder-decoder"
|
|
373
|
+
: "autoregressive";
|
|
349
374
|
if (profile.execution.loader !== expectedLoader || profile.execution.loop !== expectedLoop)
|
|
350
375
|
throw new Error(
|
|
351
376
|
`model profile ${profile.id} has an invalid execution composition for ` +
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
import type { MlxArray } from "../mlx/array";
|
|
2
|
+
import * as ops from "../mlx/ops";
|
|
3
|
+
import { quantizedSdpa } from "./gemma4-base";
|
|
4
|
+
|
|
5
|
+
/** Committed positions retain single-token attention arithmetic. Projection
|
|
6
|
+
* work can still share the span; each query sees its own causal KV prefix. */
|
|
7
|
+
export function quantizedAppendAttention(
|
|
8
|
+
q: MlxArray, keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
|
|
9
|
+
scale: number, groupSize: number, bits: number,
|
|
10
|
+
): MlxArray {
|
|
11
|
+
const [batch, heads, length, dim] = q.shape as [number, number, number, number];
|
|
12
|
+
const prefix = keys.packed.shape[2]! - length;
|
|
13
|
+
const outputs: MlxArray[] = [];
|
|
14
|
+
const trim = (tensor: ops.QuantizedTensor, end: number) => {
|
|
15
|
+
const slice = (a: MlxArray) => a.slice([0, 0, 0, 0],
|
|
16
|
+
[a.shape[0]!, a.shape[1]!, end, a.shape[3]!]);
|
|
17
|
+
return { packed: slice(tensor.packed), scales: slice(tensor.scales), biases: slice(tensor.biases) };
|
|
18
|
+
};
|
|
19
|
+
try {
|
|
20
|
+
for (let row = 0; row < length; row++) {
|
|
21
|
+
using queryView = q.slice([0, 0, row, 0], [batch, heads, row + 1, dim]);
|
|
22
|
+
using query = ops.contiguous(queryView);
|
|
23
|
+
const k = trim(keys, prefix + row + 1), v = trim(values, prefix + row + 1);
|
|
24
|
+
try { outputs.push(quantizedSdpa(query, k, v, scale, { mode: "", arr: null }, groupSize, bits)); }
|
|
25
|
+
finally {
|
|
26
|
+
for (const t of [k, v]) { t.packed.dispose(); t.scales.dispose(); t.biases.dispose(); }
|
|
27
|
+
}
|
|
28
|
+
}
|
|
29
|
+
return ops.concatAxis(outputs, 2);
|
|
30
|
+
} finally { for (const output of outputs) output.dispose(); }
|
|
31
|
+
}
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
import { MlxArray } from "../mlx/array";
|
|
2
|
+
import { Dtype } from "../mlx/ffi";
|
|
3
|
+
import { MetalKernel } from "../mlx/metal-kernel";
|
|
4
|
+
import * as ops from "../mlx/ops";
|
|
5
|
+
import { computeG } from "./qwen3-delta";
|
|
6
|
+
|
|
7
|
+
// The same recurrent arithmetic as gated_delta_step, with no query/output
|
|
8
|
+
// calculation. Each row replays only its retained prefix of the verify window.
|
|
9
|
+
const SOURCE = String.raw`
|
|
10
|
+
const int T = Tin[0];
|
|
11
|
+
auto n = thread_position_in_grid.z;
|
|
12
|
+
auto b_idx = n / Hv;
|
|
13
|
+
auto hv_idx = n % Hv;
|
|
14
|
+
auto hk_idx = hv_idx / (Hv / Hk);
|
|
15
|
+
constexpr int n_per_t = Dk / 32;
|
|
16
|
+
auto dk_idx = thread_position_in_threadgroup.x;
|
|
17
|
+
auto dv_idx = thread_position_in_grid.y;
|
|
18
|
+
auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
|
|
19
|
+
auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
|
|
20
|
+
auto g_ = g + b_idx * T * Hv;
|
|
21
|
+
auto beta_ = beta + b_idx * T * Hv;
|
|
22
|
+
auto i_state = state_in + (n * Dv + dv_idx) * Dk;
|
|
23
|
+
auto o_state = state_out + (n * Dv + dv_idx) * Dk;
|
|
24
|
+
float state[n_per_t];
|
|
25
|
+
for (int i = 0; i < n_per_t; ++i)
|
|
26
|
+
state[i] = static_cast<float>(i_state[n_per_t * dk_idx + i]);
|
|
27
|
+
const int keep = lengths[b_idx];
|
|
28
|
+
for (int t = 0; t < keep; ++t) {
|
|
29
|
+
float kv_mem = 0.0f;
|
|
30
|
+
for (int i = 0; i < n_per_t; ++i) {
|
|
31
|
+
auto s_idx = n_per_t * dk_idx + i;
|
|
32
|
+
state[i] = state[i] * g_[hv_idx];
|
|
33
|
+
kv_mem += state[i] * k_[s_idx];
|
|
34
|
+
}
|
|
35
|
+
kv_mem = simd_sum(kv_mem);
|
|
36
|
+
auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
|
|
37
|
+
for (int i = 0; i < n_per_t; ++i) {
|
|
38
|
+
auto s_idx = n_per_t * dk_idx + i;
|
|
39
|
+
state[i] = state[i] + k_[s_idx] * delta;
|
|
40
|
+
}
|
|
41
|
+
k_ += Hk * Dk;
|
|
42
|
+
v_ += Hv * Dv;
|
|
43
|
+
g_ += Hv;
|
|
44
|
+
beta_ += Hv;
|
|
45
|
+
}
|
|
46
|
+
for (int i = 0; i < n_per_t; ++i)
|
|
47
|
+
o_state[n_per_t * dk_idx + i] = state[i];
|
|
48
|
+
`;
|
|
49
|
+
|
|
50
|
+
let kernel: MetalKernel | undefined;
|
|
51
|
+
|
|
52
|
+
/** State-only prefix replay. Inputs keep their original [B,T,...] strides;
|
|
53
|
+
* lengths are request-owned accepted counts, never scheduling decisions. */
|
|
54
|
+
export function gatedDeltaState(
|
|
55
|
+
k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
|
|
56
|
+
aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
|
|
57
|
+
lengths: readonly number[],
|
|
58
|
+
): MlxArray {
|
|
59
|
+
const [B, T, Hk, Dk] = k.shape as [number, number, number, number];
|
|
60
|
+
const [, , Hv, Dv] = v.shape as [number, number, number, number];
|
|
61
|
+
using beta = ops.sigmoid(b);
|
|
62
|
+
using g = computeG(aLog, a, dtBias);
|
|
63
|
+
using zero = state ? null : ops.zeros([B, Hv, Dv, Dk], Dtype.float32);
|
|
64
|
+
using t = ops.fromInt32([T], [1]);
|
|
65
|
+
using kept = ops.fromInt32([...lengths], [B]);
|
|
66
|
+
kernel ??= new MetalKernel({ name: "gated_delta_prefix_state",
|
|
67
|
+
inputNames: ["k", "v", "g", "beta", "state_in", "Tin", "lengths"],
|
|
68
|
+
outputNames: ["state_out"], source: SOURCE, ensureRowContiguous: true });
|
|
69
|
+
return kernel.apply([k, v, g, beta, state ?? zero!, t, kept], {
|
|
70
|
+
outputs: [{ shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 }],
|
|
71
|
+
grid: [32, Dv, B * Hv], threadGroup: [32, 4, 1],
|
|
72
|
+
templateInts: { Dk, Dv, Hk, Hv },
|
|
73
|
+
})[0]!;
|
|
74
|
+
}
|