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.
Files changed (231) hide show
  1. package/README.md +57 -32
  2. package/docs/README.md +58 -42
  3. package/docs/archive/investigations/benchmark-ledger-through-v0.4.0.md +3617 -0
  4. package/docs/archive/investigations/openjarvis-fit.md +198 -0
  5. package/docs/archive/investigations/qwen38-27b-campaign-log.md +1092 -0
  6. package/docs/archive/planning/release-notes-v0.4.0.md +45 -0
  7. package/docs/archive/planning/release-notes-v0.5.0.md +37 -0
  8. package/docs/design/batching.md +1535 -50
  9. package/docs/design/decode-speed-program.md +394 -1171
  10. package/docs/design/generic-model-support.md +195 -3
  11. package/docs/design/kv-cache.md +538 -225
  12. package/docs/design/speculative-decoding.md +552 -50
  13. package/docs/design/turboquant.md +140 -17
  14. package/docs/design/unified-engine-frontier-plan.md +259 -150
  15. package/docs/planning/PRODUCT_ROADMAP.md +21 -6
  16. package/docs/planning/public-readiness-handoff.md +244 -0
  17. package/docs/planning/public-readiness.md +149 -0
  18. package/docs/reference/benchmarks.md +3537 -1112
  19. package/docs/reference/cli.md +102 -2
  20. package/docs/reference/distribution.md +45 -13
  21. package/docs/reference/environment.md +159 -13
  22. package/docs/reference/library-api.md +18 -4
  23. package/docs/reference/memory.md +1 -1
  24. package/docs/reference/models.md +28 -7
  25. package/docs/reference/server-api.md +257 -33
  26. package/docs/reference/server-config.md +372 -151
  27. package/docs/reference/troubleshooting.md +13 -21
  28. package/package.json +2 -2
  29. package/src/audio/audiotoolbox.ts +112 -0
  30. package/src/audio/conformer.ts +8 -0
  31. package/src/audio/mic-capture.ts +98 -0
  32. package/src/audio/silero-vad.ts +379 -0
  33. package/src/audio/whisper-decode.ts +774 -0
  34. package/src/audio/whisper-fast.ts +732 -0
  35. package/src/audio/whisper-format.ts +73 -0
  36. package/src/audio/whisper-mel.ts +185 -0
  37. package/src/audio/whisper-timing.ts +274 -0
  38. package/src/audio/whisper-tokenizer.ts +239 -0
  39. package/src/audio/whisper-transcribe.ts +410 -0
  40. package/src/backends/mlx/assistant-target.ts +46 -0
  41. package/src/backends/mlx/autoregressive-draft-rows.ts +120 -0
  42. package/src/backends/mlx/autoregressive.ts +22 -1
  43. package/src/backends/mlx/batch-group.ts +464 -236
  44. package/src/backends/mlx/cache-attachment.ts +67 -0
  45. package/src/backends/mlx/cache-layout.ts +31 -0
  46. package/src/backends/mlx/checkpoint-state.ts +53 -0
  47. package/src/backends/mlx/continuation-persistence.ts +107 -0
  48. package/src/backends/mlx/continuation-request.ts +43 -0
  49. package/src/backends/mlx/continuation.ts +77 -0
  50. package/src/backends/mlx/draft-target.ts +7 -22
  51. package/src/backends/mlx/encoder-cache.ts +25 -0
  52. package/src/backends/mlx/fill-append.ts +21 -8
  53. package/src/backends/mlx/fill-group.ts +497 -0
  54. package/src/backends/mlx/full-kv-row-append.ts +32 -0
  55. package/src/backends/mlx/full-kv-row-donor.ts +54 -0
  56. package/src/backends/mlx/full-transitioning-kv-rows.ts +65 -0
  57. package/src/backends/mlx/gateway-binding.ts +122 -19
  58. package/src/backends/mlx/kv-codec.ts +9 -0
  59. package/src/backends/mlx/kv-maintenance.ts +161 -0
  60. package/src/backends/mlx/kv-tensor-rows.ts +212 -0
  61. package/src/backends/mlx/mixed-iteration.ts +62 -0
  62. package/src/backends/mlx/model-prompt.ts +46 -32
  63. package/src/backends/mlx/model-serving.ts +5 -4
  64. package/src/backends/mlx/paged-cache-codec.ts +27 -0
  65. package/src/backends/mlx/prefill-cohort.ts +100 -0
  66. package/src/backends/mlx/prefill-policy.ts +26 -0
  67. package/src/backends/mlx/prefill-rows.ts +244 -0
  68. package/src/backends/mlx/prefill.ts +1 -1
  69. package/src/backends/mlx/projected-context-rows.ts +72 -0
  70. package/src/backends/mlx/prompt-input.ts +29 -0
  71. package/src/backends/mlx/qwen-prompt-input.ts +40 -0
  72. package/src/backends/mlx/request-state-policy.ts +54 -0
  73. package/src/backends/mlx/rollback.ts +45 -4
  74. package/src/backends/mlx/row-sampling.ts +49 -0
  75. package/src/backends/mlx/serial-executor.ts +51 -73
  76. package/src/backends/mlx/speculative-group.ts +438 -0
  77. package/src/backends/mlx/speculative-round.ts +161 -0
  78. package/src/backends/mlx/speculative.ts +24 -7
  79. package/src/backends/mlx/state-rows.ts +104 -0
  80. package/src/backends/mlx/state-views.ts +21 -1
  81. package/src/backends/mlx/target-layout-capability.ts +35 -0
  82. package/src/backends/mlx/transitioning-kv-rows.ts +173 -0
  83. package/src/cli.ts +333 -36
  84. package/src/contracts/execution.ts +21 -0
  85. package/src/contracts/object-cache.ts +8 -0
  86. package/src/contracts/pi-provider.ts +1 -1
  87. package/src/contracts/pi.ts +3 -1
  88. package/src/contracts/prefix-cache.ts +33 -0
  89. package/src/contracts/resources.ts +7 -0
  90. package/src/contracts/sampling.ts +23 -0
  91. package/src/contracts/scheduling.ts +27 -0
  92. package/src/engine/execution-plan.ts +38 -25
  93. package/src/engine/execution-tasks.ts +44 -0
  94. package/src/engine/namespaced-cache.ts +13 -0
  95. package/src/engine/resources.ts +12 -1
  96. package/src/engine/scheduler.ts +32 -7
  97. package/src/fill/fill-session.ts +14 -3
  98. package/src/fit.ts +29 -15
  99. package/src/generate.ts +105 -195
  100. package/src/grammar.ts +13 -0
  101. package/src/inference/draft-acceptance.ts +41 -0
  102. package/src/inference/draft-output.ts +49 -0
  103. package/src/inference/prefill.ts +6 -0
  104. package/src/inference/rollback.ts +2 -2
  105. package/src/kv-scheme.ts +92 -12
  106. package/src/kv-store.ts +183 -93
  107. package/src/lab/README.md +5 -1
  108. package/src/lab/paged-kv/paged-attention.ts +117 -0
  109. package/src/lab/paged-kv/paged-kv-rows.ts +111 -0
  110. package/src/lab/paged-kv/paged-kv.ts +126 -46
  111. package/src/md-modules.d.ts +5 -0
  112. package/src/mlx/array.ts +39 -1
  113. package/src/mlx/ffi.ts +19 -0
  114. package/src/mlx/normalized-argmax.ts +116 -0
  115. package/src/mlx/ops.ts +97 -10
  116. package/src/model/artifact-identity.ts +73 -0
  117. package/src/model/batched-kv.ts +235 -0
  118. package/src/model/batched-mask.ts +12 -4
  119. package/src/model/batched-quant.ts +8 -2
  120. package/src/model/batched-quantized-kv.ts +144 -0
  121. package/src/model/batched-rotating-quant.ts +76 -15
  122. package/src/model/batched-rotating-state.ts +76 -0
  123. package/src/model/batched-rotating.ts +64 -22
  124. package/src/model/batched-row-storage.ts +52 -0
  125. package/src/model/batched-ssm.ts +52 -0
  126. package/src/model/batched-turboquant-kv.ts +94 -0
  127. package/src/model/compiled-decode.ts +7 -1
  128. package/src/model/decoded-kv-donor.ts +12 -0
  129. package/src/model/delayed-quantized-kv.ts +72 -0
  130. package/src/model/delayed-rotating-quantized-kv.ts +93 -0
  131. package/src/model/delayed-turboquant-kv.ts +57 -0
  132. package/src/model/factory.ts +26 -0
  133. package/src/model/fingerprint.ts +12 -0
  134. package/src/model/full-prefill-padding.ts +56 -0
  135. package/src/model/full-prefill-row.ts +105 -0
  136. package/src/model/gemma4-base.ts +183 -82
  137. package/src/model/gemma4.ts +116 -28
  138. package/src/model/generated/gemma4-12b.ts +4 -1
  139. package/src/model/generated/gemma4-26b.ts +4 -1
  140. package/src/model/generated/gemma4-e4b.ts +4 -1
  141. package/src/model/glm52-cache.ts +80 -4
  142. package/src/model/glm52-mla.ts +12 -6
  143. package/src/model/glm52.ts +19 -40
  144. package/src/model/kv-attention-view.ts +70 -0
  145. package/src/model/minicpm5.ts +11 -7
  146. package/src/model/profile.ts +29 -4
  147. package/src/model/quantized-append-attention.ts +31 -0
  148. package/src/model/qwen3-delta-state.ts +74 -0
  149. package/src/model/qwen3-delta.ts +59 -36
  150. package/src/model/qwen3-mrope.ts +7 -9
  151. package/src/model/qwen3_5.ts +177 -75
  152. package/src/model/rotating-kv-donor.ts +35 -0
  153. package/src/model/rotating-kv-layout.ts +262 -0
  154. package/src/model/rotating-row-transaction.ts +35 -0
  155. package/src/model/speculative-rotating-kv.ts +101 -0
  156. package/src/model/ssm-prefill-padding.ts +39 -0
  157. package/src/model/support.ts +8 -0
  158. package/src/model/token-groups.ts +66 -0
  159. package/src/model/trellis-linear.ts +4 -4
  160. package/src/model/turboquant-codec.ts +89 -0
  161. package/src/model/universal/dense.ts +14 -13
  162. package/src/model/whisper.ts +435 -0
  163. package/src/native/mic_capture.swift +92 -0
  164. package/src/native-pack.ts +6 -5
  165. package/src/pi-web.ts +9 -1
  166. package/src/preflight.ts +31 -15
  167. package/src/prompt-cache.ts +306 -86
  168. package/src/runtime-config.ts +1 -5
  169. package/src/safetensors.ts +9 -2
  170. package/src/sampler.ts +37 -25
  171. package/src/serve/audio-routes.ts +323 -0
  172. package/src/serve/chat-request.ts +2 -0
  173. package/src/serve/chat-stage.ts +19 -14
  174. package/src/serve/completion-executor.ts +9 -1
  175. package/src/serve/discovery-routes.ts +17 -1
  176. package/src/serve/generated-token-history.ts +57 -0
  177. package/src/serve/generation-gateway.ts +88 -38
  178. package/src/serve/model-binding.ts +7 -1
  179. package/src/serve/model-host.ts +3 -4
  180. package/src/serve/preparation.ts +3 -3
  181. package/src/serve/prompt-response-trace.ts +14 -0
  182. package/src/serve/request-plan.ts +6 -3
  183. package/src/serve/request-prep.ts +40 -25
  184. package/src/serve/text-completion-stage.ts +2 -2
  185. package/src/serve/transcription-server.ts +64 -0
  186. package/src/serve/transcription-service.ts +502 -0
  187. package/src/server.ts +185 -261
  188. package/src/spec/assistant-rows.ts +93 -0
  189. package/src/spec/assistant-source.ts +28 -18
  190. package/src/spec/cache-identity.ts +22 -0
  191. package/src/spec/deepspec-rows.ts +29 -0
  192. package/src/spec/deepspec-source.ts +12 -10
  193. package/src/spec/dflash-source.ts +28 -4
  194. package/src/spec/drafter.ts +153 -142
  195. package/src/spec/dspark/deepspec-module.ts +95 -120
  196. package/src/spec/dspark/module-dflash.ts +126 -170
  197. package/src/spec/glm52-mtp-graph.ts +29 -0
  198. package/src/spec/glm52-mtp-rows.ts +132 -0
  199. package/src/spec/glm52-mtp-source.ts +54 -136
  200. package/src/spec/glm52-mtp-state.ts +33 -0
  201. package/src/spec/ngram-source.ts +115 -21
  202. package/src/spec/projected-draft-rows.ts +73 -0
  203. package/src/spec/qwen-mtp-module.ts +171 -0
  204. package/src/spec/qwen-mtp-rows.ts +230 -0
  205. package/src/spec/qwen-mtp-source.ts +84 -226
  206. package/src/spec/qwen-mtp-state.ts +30 -0
  207. package/src/spec/serve-loop.ts +73 -69
  208. package/src/spec/source.ts +93 -15
  209. package/src/spec/two-model-rows.ts +62 -0
  210. package/src/spec/two-model.ts +45 -103
  211. package/src/ssd-cache.ts +107 -36
  212. package/src/ssd-durability.ts +29 -27
  213. package/src/storage/host-buffer.ts +24 -0
  214. package/src/storage/kv-writer.ts +83 -0
  215. package/src/storage/kv-writer.worker.js +264 -0
  216. package/src/storage/retention-policy.ts +37 -0
  217. package/src/tiered-prompt-cache.ts +45 -0
  218. package/src/vision/embedder.ts +8 -0
  219. package/src/vision/prompt.ts +101 -65
  220. package/src/vision/qwen3vl-prompt.ts +101 -64
  221. package/src/vision/qwen3vl-tower.ts +11 -0
  222. package/src/vision/siglip.ts +8 -0
  223. package/src/web/app.html +12 -0
  224. package/src/web/app.js +217 -57
  225. package/src/web/src/chat.ts +4 -1
  226. package/src/web/src/composer.ts +3 -0
  227. package/src/web/src/voice.ts +141 -0
  228. package/src/web/sw.js +1 -1
  229. package/src/backends/mlx/memory-guard.ts +0 -40
  230. package/src/spec/prefix-state.ts +0 -45
  231. /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 === "plain") {
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 {
@@ -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
- throw new Error("MLA mergeRows cannot merge an empty row");
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) this.dsa.restoreState(dsa!, width);
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
- if (this.dsa) this.dsa.restoreState(dsa!, nextWidth);
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
 
@@ -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
- if (tokens! > 1) {
574
- const maskValues = new Float32Array(tokens! * keys);
575
- for (let queryToken = 0; queryToken < tokens!; queryToken++) {
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[queryToken * keys + keyToken] = -Infinity;
584
+ maskValues[start + keyToken] = -Infinity;
579
585
  }
580
- const mask = MlxArray.fromFloat32(maskValues, [1, 1, tokens!, keys]);
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();
@@ -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 (batch! > 1) {
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
- [1, token + 1, this.config.hiddenSize],
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 = dsaState.selectSharedPositions(
576
- this.layer,
577
- cache.offset + 1,
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.batchSize !== null && cache.batchSize > 1
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
- if (ids.shape.length !== 2 || ids.shape[0] !== 1) {
875
- throw new Error(
876
- "G2 DSA sparse prefill requires serial [1,T] token ids",
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], [1, token + 1]);
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], [1, token + 1]);
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
+ }
@@ -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 instanceof QuantizedKVCache) {
147
- const [keys, values] = cache.updateAndFetchQuantized(k, v);
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, cache.groupSize, cache.bits);
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 instanceof TurboQuantKVCache) {
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);
@@ -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
- export type GenerationLoop = "autoregressive" | "diffusion";
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
- : "autoregressive";
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
+ }