mlx-bun 0.0.13 → 0.2.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 +31 -21
- package/bin/mlx-bun.mjs +1 -1
- package/docs/{design → archive}/hlg-sampling.md +3 -3
- package/docs/{investigations → archive/investigations}/hlg-sampling-investigation.md +1 -1
- package/docs/archive/investigations/pareto-specialized-runtime-findings.md +92 -0
- package/docs/{planning → archive/planning}/chat-ui-vision.md +1 -1
- package/docs/{planning → archive/planning}/release-notes-v0.0.9.md +1 -1
- package/docs/{planning → archive/planning}/website-readme-pass-plan.md +1 -1
- package/docs/design/colibri-glm52-port.md +1 -1
- package/docs/design/compat-cli-surface-design.md +1 -1
- package/docs/design/decode-speed-program.md +1 -1
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dspark-serving-program.md +1 -1
- package/docs/design/dspark-speculative-decoding.md +53 -3
- package/docs/design/grammar-spec-batching-integration.md +1 -1
- package/docs/design/mlx-lm-serving-execution-seam.md +81 -0
- package/docs/design/orpo-dynamic-lambda.md +1 -1
- package/docs/design/orpo-future-enhancements.md +1 -1
- package/docs/design/orpo-training.md +2 -2
- package/docs/design/parity-tier-dag.md +1 -1
- package/docs/design/repo-cleanup-plan.md +46 -4
- package/docs/design/segmented-backward-training.md +2 -2
- package/docs/design/ssd-kv-cold-tier.md +127 -0
- package/docs/design/turboquant-weights.md +467 -0
- package/docs/design/w4a16-compute-precision-spike.md +1 -1
- package/docs/design/web-chat-redesign.md +5 -5
- package/docs/planning/chunk-task-roadmap.md +1 -1
- package/docs/planning/curve-sampler-research-plan.md +4 -4
- package/docs/planning/memory-docs-and-dag-plan.md +1 -1
- package/docs/planning/release-notes-v0.2.0.md +49 -0
- package/docs/planning/video-series.md +421 -0
- package/docs/reference/cli.md +15 -4
- package/docs/reference/distribution.md +1 -0
- package/docs/reference/features-matrix.md +3 -2
- package/docs/reference/library-api.md +41 -1
- package/docs/reference/models.md +33 -0
- package/docs/reference/server-api.md +69 -14
- package/docs/reference/server-config.md +57 -33
- package/package.json +2 -2
- package/src/anthropic.ts +92 -3
- package/src/cli.ts +131 -36
- package/src/eval/kl.ts +27 -20
- package/src/eval/runner.ts +23 -4
- package/src/eval/tasks/bfcl.ts +2 -1
- package/src/eval/tasks/gsm8k.ts +2 -1
- package/src/eval/tasks/hashhop.ts +2 -1
- package/src/eval/tasks/humaneval.ts +5 -4
- package/src/eval/tasks/ifeval.ts +2 -1
- package/src/eval/tasks/mmlu.ts +2 -1
- package/src/expert-io.ts +3 -1
- package/src/expert-offload.ts +3 -1
- package/src/expert-trace.ts +3 -1
- package/src/fit.ts +24 -43
- package/src/flags.ts +7 -9
- package/src/generate.ts +107 -92
- package/src/grammar.ts +2 -1
- package/src/harness-pi.ts +1 -1
- package/src/index.ts +20 -0
- package/src/jobs/job-entry.ts +3 -2
- package/src/kv-scheme.ts +199 -0
- package/src/kv-store.ts +370 -205
- package/src/media-fetch.ts +19 -3
- package/src/memory/model.ts +2 -1
- package/src/memory/vault.ts +2 -1
- package/src/memory-plan.ts +18 -0
- package/src/mlx/ffi.ts +8 -11
- package/src/mlx/ops.ts +31 -6
- package/src/model/batched-mask.ts +1 -1
- package/src/model/batched-quant.ts +1 -1
- package/src/model/batched-rotating-quant.ts +43 -97
- package/src/model/batched-rotating-state.ts +73 -0
- package/src/model/batched-rotating.ts +56 -105
- package/src/model/batched-row-storage.ts +127 -0
- package/src/model/compiled-decode.ts +3 -2
- package/src/model/factory.ts +44 -49
- package/src/model/gemma4-base.ts +163 -10
- package/src/model/generated/gemma4-12b.ts +50 -146
- package/src/model/generated/gemma4-26b.ts +32 -92
- package/src/model/generated/gemma4-e4b.ts +26 -74
- package/src/model/glm52-cache.ts +19 -9
- package/src/model/glm52-memory.ts +28 -1
- package/src/model/minicpm5.ts +2 -1
- package/src/model/profile.ts +433 -0
- package/src/model/qwen3-delta.ts +98 -2
- package/src/model/qwen3-mrope.ts +252 -0
- package/src/model/qwen3_5.ts +217 -25
- package/src/native/frame_extract.swift +140 -0
- package/src/native-pack.ts +13 -6
- package/src/pi-web.ts +6 -5
- package/src/prompt-cache.ts +2 -2
- package/src/quantize/atomic-output.ts +38 -0
- package/src/quantize/calibration.ts +2 -1
- package/src/quantize/config-writer.ts +6 -0
- package/src/quantize/index.ts +28 -1
- package/src/quantize/job.ts +11 -0
- package/src/quantize/quantizer.ts +265 -100
- package/src/quantize/rotate.ts +603 -0
- package/src/quantize/safetensors-writer.ts +79 -1
- package/src/quantize/sensitivity.ts +41 -59
- package/src/quantize/weight-transform.ts +292 -0
- package/src/responses.ts +102 -2
- package/src/runtime-config.ts +67 -0
- package/src/sampler.ts +272 -0
- package/src/serve/admin-routes.ts +184 -0
- package/src/serve/aux-routes.ts +132 -0
- package/src/serve/batch-scheduler.ts +200 -87
- package/src/serve/completion-executor.ts +346 -0
- package/src/serve/completion-sink.ts +157 -0
- package/src/serve/discovery-routes.ts +205 -0
- package/src/serve/generation-gateway.ts +116 -109
- package/src/serve/kv-budget.ts +14 -0
- package/src/serve/lab-routes.ts +338 -0
- package/src/serve/model-admin-routes.ts +153 -0
- package/src/serve/prompt-response-trace.ts +147 -0
- package/src/serve/request-plan.ts +158 -0
- package/src/serve/static-routes.ts +101 -0
- package/src/server.ts +704 -1400
- package/src/spec/dflash-source.ts +2 -1
- package/src/spec/dspark/module-dflash.ts +1 -1
- package/src/spec/glm52-mtp-source.ts +8 -0
- package/src/spec/qwen-mtp-source.ts +431 -0
- package/src/spec/serve-loop.ts +36 -68
- package/src/ssd-cache.ts +17 -14
- package/src/ssd-durability.ts +195 -0
- package/src/train/flash-cce.ts +12 -11
- package/src/train/forward.ts +32 -1
- package/src/train/loss.ts +2 -1
- package/src/train/segmented.ts +4 -3
- package/src/train/trainer.ts +3 -2
- package/src/vision/preprocess.ts +1 -1
- package/src/vision/prompt.ts +47 -1
- package/src/vision/qwen3vl-preprocess.ts +409 -0
- package/src/vision/qwen3vl-prompt.ts +196 -0
- package/src/vision/qwen3vl-tower.ts +513 -0
- package/src/vision/video-frames.ts +148 -0
- package/src/web/app.js +10 -3
- package/src/web/src/chat.ts +12 -5
- package/src/weights.ts +81 -26
- /package/docs/{investigations → archive/investigations}/chunk-finetune-experiment.md +0 -0
- /package/docs/{investigations → archive/investigations}/colibri-oracle-pin.md +0 -0
- /package/docs/{investigations → archive/investigations}/curve-bisector-routing.md +0 -0
- /package/docs/{investigations → archive/investigations}/curve-runs/bisector-route.json +0 -0
- /package/docs/{investigations → archive/investigations}/curve-runs/bisector-trace.json +0 -0
- /package/docs/{investigations → archive/investigations}/curve-runs/curve-terrain.json +0 -0
- /package/docs/{investigations → archive/investigations}/curve-runs/curve-terrain.svg +0 -0
- /package/docs/{investigations → archive/investigations}/curve-terrain.html +0 -0
- /package/docs/{investigations → archive/investigations}/decode-roofline-lookagain.md +0 -0
- /package/docs/{investigations → archive/investigations}/dspark-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/expert-offload-single-user-moe.md +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-ab.html +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-explorer.html +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/cliff.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/cross.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/edges.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/k20-grids.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/k20-peaks.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/lw-curve.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-figs/terrain.svg +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-report.html +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-ab.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-blind-prompts.md +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-blind.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-cliff.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-cross.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-edge.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-eval1.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-frontier.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-grid.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-grid2.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-lw-coarse.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/hlg-map.json +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/run0-original.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runA-recommended.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAA-cliff-cpm5.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAB-grid2-e4b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAC-peak-e4b-partial.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAD-ab-samples.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAE-edge-e4b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAF-logit-calibration.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runAG-eval-config2.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runB-gain-variants.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runC-pivots.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runD-oetf.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runE-eotf.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runF-shaper-working.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runG-shaper-sweep.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runH-wsweep.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runI-asweep.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runJ-knobmap.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runK-frontier.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runL-widemap.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runM-frontier-emb.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runN-verify.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runO-frontier-seed5000.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runP-grid-AxSm.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runQ-grid-tgW.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runR-blind-transcript.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runS-cross-coherence.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runT-cross-cpm5.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runU-cross-12b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runV-lw-e4b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runW-lw-12b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runX-lw-cpm5.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runY-cliff-e4b.log +0 -0
- /package/docs/{investigations → archive/investigations}/hlg-runs/runZ-cliff-12b.log +0 -0
- /package/docs/{investigations → archive/investigations}/indexshare-performance-spike.md +0 -0
- /package/docs/{investigations → archive/investigations}/kernel-perf-review-2026-07.md +0 -0
- /package/docs/{investigations → archive/investigations}/lab-build-journal.md +0 -0
- /package/docs/{investigations → archive/investigations}/openwiki-evaluation.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-base-uf-experiment-and-directions.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-flash-cce-pin-leak.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-uf-testing-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/pi-builtin-investigation.md +0 -0
- /package/docs/{investigations → archive/investigations}/segmented-backward-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/starter-model-port-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/steel-flash-cce-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/trainer-validation-experiment.md +0 -0
- /package/docs/{planning → archive/planning}/chat-ui-vision-memory-first.md +0 -0
- /package/docs/{planning → archive/planning}/journal.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.10.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.11.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.12.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.13.md +0 -0
- /package/docs/{planning → archive/planning}/web-ui-pass-plan.md +0 -0
|
@@ -0,0 +1,513 @@
|
|
|
1
|
+
// Qwen3-VL vision tower (Qwen3.8's `vision_tower`, PLAN 14v) — faithful port
|
|
2
|
+
// of mlx-vlm qwen3_vl/vision.py at the qwen3_5 configuration: depth 27,
|
|
3
|
+
// hidden 1152, heads 16 (head_dim 72), patch 16, temporal 2, merge 2,
|
|
4
|
+
// intermediate 4304 with gelu_pytorch_tanh, learned 48×48 pos-embed grid
|
|
5
|
+
// bilinearly interpolated, 2D half/half rotary (dim 36) over full head_dim,
|
|
6
|
+
// pre-LN blocks with biased LayerNorms, merger LayerNorm→fc1→GELU(erf)→fc2
|
|
7
|
+
// to out_hidden 5120. deepstack is EMPTY for qwen3_5 and not ported.
|
|
8
|
+
//
|
|
9
|
+
// Weights: the artifact ships the tower as a bf16 sidecar
|
|
10
|
+
// (optiq/optiq_vision.safetensors, 333 tensors, MLX conv layout). The
|
|
11
|
+
// Conv3d patch embed has kernel == stride, so it is computed as one matmul
|
|
12
|
+
// against the weight flattened to the preprocessor's (c, t, py, px) row
|
|
13
|
+
// order — same values, GEMM accumulation (gate: tower output vs the mlx-vlm
|
|
14
|
+
// capture; e2e tokens are the decisive gate, the SigLIP precedent).
|
|
15
|
+
//
|
|
16
|
+
// Input rows arrive MERGE-BLOCK-MAJOR from qwen3vl-preprocess.ts; rope
|
|
17
|
+
// h/w positions and the pos-embed permutation are built in the same order.
|
|
18
|
+
|
|
19
|
+
import { ptr, read } from "bun:ffi";
|
|
20
|
+
import { existsSync, readFileSync } from "node:fs";
|
|
21
|
+
import { MlxArray, cpuStream } from "../mlx/array";
|
|
22
|
+
import { C, Dtype } from "../mlx/ffi";
|
|
23
|
+
import * as ops from "../mlx/ops";
|
|
24
|
+
import type { Qwen3VLPreprocessed } from "./qwen3vl-preprocess";
|
|
25
|
+
import {
|
|
26
|
+
QWEN3VL_MERGE_SIZE,
|
|
27
|
+
QWEN3VL_PATCH_SIZE,
|
|
28
|
+
QWEN3VL_TEMPORAL_PATCH_SIZE,
|
|
29
|
+
} from "./qwen3vl-preprocess";
|
|
30
|
+
|
|
31
|
+
const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
|
|
32
|
+
|
|
33
|
+
const DEPTH = 27;
|
|
34
|
+
const HIDDEN = 1152;
|
|
35
|
+
const HEADS = 16;
|
|
36
|
+
const HEAD_DIM = HIDDEN / HEADS; // 72
|
|
37
|
+
const ROT_DIM = HEAD_DIM / 2; // 36 (VisionRotaryEmbedding dim)
|
|
38
|
+
const INTERMEDIATE = 4304;
|
|
39
|
+
const OUT_HIDDEN = 5120;
|
|
40
|
+
const POS_GRID = 48; // sqrt(num_position_embeddings 2304)
|
|
41
|
+
const ROPE_THETA = 10000;
|
|
42
|
+
|
|
43
|
+
function disposing(old: MlxArray, next: MlxArray): MlxArray {
|
|
44
|
+
old.dispose();
|
|
45
|
+
return next;
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
interface Block {
|
|
49
|
+
norm1W: MlxArray; norm1B: MlxArray;
|
|
50
|
+
norm2W: MlxArray; norm2B: MlxArray;
|
|
51
|
+
qkvW: MlxArray; qkvB: MlxArray;
|
|
52
|
+
projW: MlxArray; projB: MlxArray;
|
|
53
|
+
fc1W: MlxArray; fc1B: MlxArray;
|
|
54
|
+
fc2W: MlxArray; fc2B: MlxArray;
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
export class Qwen3VLVisionTower {
|
|
58
|
+
#patchW: MlxArray; // [1536, 1152] — flattened conv, already transposed
|
|
59
|
+
#patchB: MlxArray;
|
|
60
|
+
#posEmbed: MlxArray; // [2304, 1152]
|
|
61
|
+
#blocks: Block[];
|
|
62
|
+
#mergerNormW: MlxArray; #mergerNormB: MlxArray;
|
|
63
|
+
#mergerFc1W: MlxArray; #mergerFc1B: MlxArray;
|
|
64
|
+
#mergerFc2W: MlxArray; #mergerFc2B: MlxArray;
|
|
65
|
+
|
|
66
|
+
private constructor(w: Map<string, MlxArray>) {
|
|
67
|
+
const t = (n: string): MlxArray => {
|
|
68
|
+
const a = w.get(`vision_tower.${n}`);
|
|
69
|
+
if (!a) throw new Error(`vision sidecar missing tensor vision_tower.${n}`);
|
|
70
|
+
return a;
|
|
71
|
+
};
|
|
72
|
+
// Conv3d weight [1152, 2, 16, 16, 3] (out, kT, kH, kW, in) — used with
|
|
73
|
+
// the REAL conv3d op (kernel == stride): a value-equivalent GEMM rounds
|
|
74
|
+
// differently in bf16 and the ulp seeds diverge through 27 blocks.
|
|
75
|
+
this.#patchW = t("patch_embed.proj.weight");
|
|
76
|
+
this.#patchB = t("patch_embed.proj.bias");
|
|
77
|
+
this.#posEmbed = t("pos_embed.weight");
|
|
78
|
+
this.#blocks = [];
|
|
79
|
+
for (let i = 0; i < DEPTH; i++) {
|
|
80
|
+
const p = `blocks.${i}`;
|
|
81
|
+
this.#blocks.push({
|
|
82
|
+
norm1W: t(`${p}.norm1.weight`), norm1B: t(`${p}.norm1.bias`),
|
|
83
|
+
norm2W: t(`${p}.norm2.weight`), norm2B: t(`${p}.norm2.bias`),
|
|
84
|
+
qkvW: t(`${p}.attn.qkv.weight`), qkvB: t(`${p}.attn.qkv.bias`),
|
|
85
|
+
projW: t(`${p}.attn.proj.weight`), projB: t(`${p}.attn.proj.bias`),
|
|
86
|
+
fc1W: t(`${p}.mlp.linear_fc1.weight`), fc1B: t(`${p}.mlp.linear_fc1.bias`),
|
|
87
|
+
fc2W: t(`${p}.mlp.linear_fc2.weight`), fc2B: t(`${p}.mlp.linear_fc2.bias`),
|
|
88
|
+
});
|
|
89
|
+
}
|
|
90
|
+
this.#mergerNormW = t("merger.norm.weight");
|
|
91
|
+
this.#mergerNormB = t("merger.norm.bias");
|
|
92
|
+
this.#mergerFc1W = t("merger.linear_fc1.weight");
|
|
93
|
+
this.#mergerFc1B = t("merger.linear_fc1.bias");
|
|
94
|
+
this.#mergerFc2W = t("merger.linear_fc2.weight");
|
|
95
|
+
this.#mergerFc2B = t("merger.linear_fc2.bias");
|
|
96
|
+
}
|
|
97
|
+
|
|
98
|
+
/** The 333 tensor names the tower models, by constructed name. */
|
|
99
|
+
static #tensorNames(): string[] {
|
|
100
|
+
const names: string[] = [
|
|
101
|
+
"vision_tower.patch_embed.proj.weight", "vision_tower.patch_embed.proj.bias",
|
|
102
|
+
"vision_tower.pos_embed.weight",
|
|
103
|
+
"vision_tower.merger.norm.weight", "vision_tower.merger.norm.bias",
|
|
104
|
+
"vision_tower.merger.linear_fc1.weight", "vision_tower.merger.linear_fc1.bias",
|
|
105
|
+
"vision_tower.merger.linear_fc2.weight", "vision_tower.merger.linear_fc2.bias",
|
|
106
|
+
];
|
|
107
|
+
for (let i = 0; i < DEPTH; i++) {
|
|
108
|
+
for (const nm of [
|
|
109
|
+
"norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias",
|
|
110
|
+
"attn.qkv.weight", "attn.qkv.bias", "attn.proj.weight", "attn.proj.bias",
|
|
111
|
+
"mlp.linear_fc1.weight", "mlp.linear_fc1.bias",
|
|
112
|
+
"mlp.linear_fc2.weight", "mlp.linear_fc2.bias",
|
|
113
|
+
]) names.push(`vision_tower.blocks.${i}.${nm}`);
|
|
114
|
+
}
|
|
115
|
+
return names;
|
|
116
|
+
}
|
|
117
|
+
|
|
118
|
+
/** Load one safetensors file's map and move `wanted` tensors into `out`
|
|
119
|
+
* (missing names are skipped; the map is always freed). */
|
|
120
|
+
static #takeFromFile(path: string, wanted: string[], out: Map<string, MlxArray>): void {
|
|
121
|
+
const arrMap = new BigUint64Array([C.mlx_map_string_to_array_new()]);
|
|
122
|
+
const metaMap = new BigUint64Array([C.mlx_map_string_to_string_new()]);
|
|
123
|
+
const arrMapPtr = ptr(arrMap);
|
|
124
|
+
const metaMapPtr = ptr(metaMap);
|
|
125
|
+
const status = C.mlx_load_safetensors(arrMapPtr, metaMapPtr, ptr(cstr(path)), cpuStream);
|
|
126
|
+
C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
|
|
127
|
+
const handle = read.u64(arrMapPtr, 0);
|
|
128
|
+
if (status !== 0) {
|
|
129
|
+
// The out-param map was allocated regardless — free it on the error
|
|
130
|
+
// path too (2026-08-18 review; one-shot leak, but a leak).
|
|
131
|
+
C.mlx_map_string_to_array_free(handle);
|
|
132
|
+
throw new Error(`failed to load ${path}`);
|
|
133
|
+
}
|
|
134
|
+
try {
|
|
135
|
+
for (const name of wanted) {
|
|
136
|
+
if (out.has(name)) continue;
|
|
137
|
+
const slot = new BigUint64Array([C.mlx_array_new()]);
|
|
138
|
+
const slotPtr = ptr(slot);
|
|
139
|
+
if (C.mlx_map_string_to_array_get(slotPtr, handle, ptr(cstr(name))) !== 0) {
|
|
140
|
+
C.mlx_array_free(read.u64(slotPtr, 0));
|
|
141
|
+
continue;
|
|
142
|
+
}
|
|
143
|
+
out.set(name, new MlxArray(read.u64(slotPtr, 0)));
|
|
144
|
+
}
|
|
145
|
+
} finally {
|
|
146
|
+
C.mlx_map_string_to_array_free(handle);
|
|
147
|
+
}
|
|
148
|
+
}
|
|
149
|
+
|
|
150
|
+
static load(modelDir: string): Qwen3VLVisionTower {
|
|
151
|
+
const names = Qwen3VLVisionTower.#tensorNames();
|
|
152
|
+
const weights = new Map<string, MlxArray>();
|
|
153
|
+
try {
|
|
154
|
+
const sidecar = `${modelDir}/optiq/optiq_vision.safetensors`;
|
|
155
|
+
if (existsSync(sidecar)) {
|
|
156
|
+
// OptiQ-convention sidecar (their artifacts + ours pre-2026-08-18).
|
|
157
|
+
Qwen3VLVisionTower.#takeFromFile(sidecar, names, weights);
|
|
158
|
+
} else {
|
|
159
|
+
// In-main fallback: the bf16 vision_tower.* tensors live in the main
|
|
160
|
+
// safetensors shards (the mlx-vlm convention — one copy serves both
|
|
161
|
+
// ecosystems; sidecar duplication retired 2026-08-18).
|
|
162
|
+
const idxPath = `${modelDir}/model.safetensors.index.json`;
|
|
163
|
+
let files: string[];
|
|
164
|
+
if (existsSync(idxPath)) {
|
|
165
|
+
const idx = JSON.parse(readFileSync(idxPath, "utf8")) as
|
|
166
|
+
{ weight_map: Record<string, string> };
|
|
167
|
+
files = [...new Set(
|
|
168
|
+
names.map((n) => idx.weight_map[n]).filter((f): f is string => !!f),
|
|
169
|
+
)];
|
|
170
|
+
if (files.length === 0)
|
|
171
|
+
throw new Error(`no vision_tower tensors in ${modelDir} (index has none)`);
|
|
172
|
+
} else {
|
|
173
|
+
files = ["model.safetensors"];
|
|
174
|
+
}
|
|
175
|
+
for (const f of files)
|
|
176
|
+
Qwen3VLVisionTower.#takeFromFile(`${modelDir}/${f}`, names, weights);
|
|
177
|
+
}
|
|
178
|
+
for (const name of names)
|
|
179
|
+
if (!weights.has(name))
|
|
180
|
+
throw new Error(`qwen vision weights missing tensor ${name} (sidecar or in-main)`);
|
|
181
|
+
} catch (e) {
|
|
182
|
+
for (const [, a] of weights) a.dispose();
|
|
183
|
+
throw e;
|
|
184
|
+
}
|
|
185
|
+
return new Qwen3VLVisionTower(weights);
|
|
186
|
+
}
|
|
187
|
+
|
|
188
|
+
/** One vision block: pre-LN attention (+residual), pre-LN MLP (+residual).
|
|
189
|
+
* Consumes `h`. */
|
|
190
|
+
#block(
|
|
191
|
+
h: MlxArray, blk: Block, cos: MlxArray, sin: MlxArray,
|
|
192
|
+
scale: number, n: number, segLen: number = n,
|
|
193
|
+
): MlxArray {
|
|
194
|
+
const xn = ops.layerNorm(h, blk.norm1W, blk.norm1B, 1e-6);
|
|
195
|
+
const qkv = this.#linear(xn, blk.qkvW, blk.qkvB); // [n, 3456]
|
|
196
|
+
xn.dispose();
|
|
197
|
+
const qkv4 = disposing(qkv, ops.reshape(qkv, [n, 3, HEADS, HEAD_DIM]));
|
|
198
|
+
const [qs, ks, vs] = ops.split(qkv4, [1, 2], 1) as [MlxArray, MlxArray, MlxArray];
|
|
199
|
+
qkv4.dispose();
|
|
200
|
+
let q = disposing(qs, ops.reshape(qs, [n, HEADS, HEAD_DIM]));
|
|
201
|
+
let k = disposing(ks, ops.reshape(ks, [n, HEADS, HEAD_DIM]));
|
|
202
|
+
const v = disposing(vs, ops.reshape(vs, [n, HEADS, HEAD_DIM]));
|
|
203
|
+
q = disposing(q, this.#applyRope(q, cos, sin));
|
|
204
|
+
k = disposing(k, this.#applyRope(k, cos, sin));
|
|
205
|
+
// [1, HEADS, n, HEAD_DIM] full (non-causal) attention — one image =
|
|
206
|
+
// one cu_seqlens segment, so no mask. The reference's ensure_fused_sdpa
|
|
207
|
+
// ZERO-PADS head_dim 72 → 80 (nearest fused-kernel width), runs fused
|
|
208
|
+
// sdpa there, and slices back; the padded kernel's accumulation is the
|
|
209
|
+
// parity-relevant arithmetic, so replicate it exactly.
|
|
210
|
+
const PAD_DIM = 80;
|
|
211
|
+
const padTo = (x: MlxArray): MlxArray => {
|
|
212
|
+
const z = ops.zeros([n, HEADS, PAD_DIM - HEAD_DIM], Dtype.bfloat16);
|
|
213
|
+
const p = ops.concatAxis([x, z], -1);
|
|
214
|
+
z.dispose();
|
|
215
|
+
return disposing(x, p);
|
|
216
|
+
};
|
|
217
|
+
const lift = (x: MlxArray): MlxArray => {
|
|
218
|
+
const p = padTo(x); // consumes x
|
|
219
|
+
const tr = ops.transposeAxes(p, [1, 0, 2]);
|
|
220
|
+
p.dispose();
|
|
221
|
+
return disposing(tr, ops.expandDims(tr, 0));
|
|
222
|
+
};
|
|
223
|
+
const qT = lift(q);
|
|
224
|
+
const kT = lift(k);
|
|
225
|
+
const vT = lift(v);
|
|
226
|
+
// cu_seqlens split: one sdpa per temporal segment (single segment for
|
|
227
|
+
// images), outputs concatenated in order — the reference's exact walk.
|
|
228
|
+
let attnP: MlxArray;
|
|
229
|
+
if (segLen === n) {
|
|
230
|
+
attnP = ops.sdpa(qT, kT, vT, scale, "", null);
|
|
231
|
+
} else {
|
|
232
|
+
const outs: MlxArray[] = [];
|
|
233
|
+
for (let s = 0; s < n; s += segLen) {
|
|
234
|
+
const qs = qT.slice([0, 0, s, 0], [1, HEADS, s + segLen, PAD_DIM]);
|
|
235
|
+
const ks = kT.slice([0, 0, s, 0], [1, HEADS, s + segLen, PAD_DIM]);
|
|
236
|
+
const vsS = vT.slice([0, 0, s, 0], [1, HEADS, s + segLen, PAD_DIM]);
|
|
237
|
+
outs.push(ops.sdpa(qs, ks, vsS, scale, "", null));
|
|
238
|
+
qs.dispose(); ks.dispose(); vsS.dispose();
|
|
239
|
+
}
|
|
240
|
+
attnP = ops.concatAxis(outs, 2);
|
|
241
|
+
for (const o of outs) o.dispose();
|
|
242
|
+
}
|
|
243
|
+
qT.dispose(); kT.dispose(); vT.dispose();
|
|
244
|
+
const attn = disposing(
|
|
245
|
+
attnP, attnP.slice([0, 0, 0, 0], [1, HEADS, n, HEAD_DIM]),
|
|
246
|
+
);
|
|
247
|
+
const attnT = disposing(attn, ops.transposeAxes(attn, [0, 2, 1, 3]));
|
|
248
|
+
const merged = disposing(attnT, ops.reshape(attnT, [n, HIDDEN]));
|
|
249
|
+
const proj = disposing(merged, this.#linear(merged, blk.projW, blk.projB));
|
|
250
|
+
h = disposing(h, disposing(proj, ops.add(h, proj)));
|
|
251
|
+
const mn = ops.layerNorm(h, blk.norm2W, blk.norm2B, 1e-6);
|
|
252
|
+
const f1 = disposing(mn, this.#linear(mn, blk.fc1W, blk.fc1B));
|
|
253
|
+
const act = disposing(f1, ops.geluApprox(f1)); // gelu_pytorch_tanh
|
|
254
|
+
const f2 = disposing(act, this.#linear(act, blk.fc2W, blk.fc2B));
|
|
255
|
+
return disposing(h, disposing(f2, ops.add(h, f2)));
|
|
256
|
+
}
|
|
257
|
+
|
|
258
|
+
/** nn.Linear on bf16 — mx.addmm(bias, x, W^T), the reference's exact bias
|
|
259
|
+
* path (a separate matmul+add rounds differently and seeded the ±1-ulp
|
|
260
|
+
* block divergence found in the stage bisect). */
|
|
261
|
+
#linear(x: MlxArray, w: MlxArray, b: MlxArray): MlxArray {
|
|
262
|
+
const wT = ops.transposeAxes(w, [1, 0]);
|
|
263
|
+
const y = ops.addmm(b, x, wT);
|
|
264
|
+
wT.dispose();
|
|
265
|
+
return y;
|
|
266
|
+
}
|
|
267
|
+
|
|
268
|
+
/** Bilinearly interpolated learned pos embeds, in merge-block-major row
|
|
269
|
+
* order — port of fast_pos_embed_interpolate + the merge permutation. */
|
|
270
|
+
#interpolatedPosEmbeds(gridH: number, gridW: number): MlxArray {
|
|
271
|
+
const n = gridH * gridW;
|
|
272
|
+
const idx = [new Int32Array(n), new Int32Array(n), new Int32Array(n), new Int32Array(n)];
|
|
273
|
+
const wts = [new Float32Array(n), new Float32Array(n), new Float32Array(n), new Float32Array(n)];
|
|
274
|
+
const lin = (count: number): Float32Array => {
|
|
275
|
+
// mx.linspace(0, POS_GRID-1, count) in f32.
|
|
276
|
+
const outArr = new Float32Array(count);
|
|
277
|
+
if (count === 1) { outArr[0] = 0; return outArr; }
|
|
278
|
+
const step = Math.fround((POS_GRID - 1) / (count - 1));
|
|
279
|
+
for (let i = 0; i < count; i++) outArr[i] = Math.fround(i * step);
|
|
280
|
+
return outArr;
|
|
281
|
+
};
|
|
282
|
+
const hIdx = lin(gridH);
|
|
283
|
+
const wIdx = lin(gridW);
|
|
284
|
+
for (let h = 0; h < gridH; h++) {
|
|
285
|
+
const hf = Math.trunc(hIdx[h]!);
|
|
286
|
+
const hc = Math.min(hf + 1, POS_GRID - 1);
|
|
287
|
+
const dh = Math.fround(hIdx[h]! - hf);
|
|
288
|
+
for (let x = 0; x < gridW; x++) {
|
|
289
|
+
const wf = Math.trunc(wIdx[x]!);
|
|
290
|
+
const wc = Math.min(wf + 1, POS_GRID - 1);
|
|
291
|
+
const dw = Math.fround(wIdx[x]! - wf);
|
|
292
|
+
const i = h * gridW + x;
|
|
293
|
+
idx[0]![i] = hf * POS_GRID + wf;
|
|
294
|
+
idx[1]![i] = hf * POS_GRID + wc;
|
|
295
|
+
idx[2]![i] = hc * POS_GRID + wf;
|
|
296
|
+
idx[3]![i] = hc * POS_GRID + wc;
|
|
297
|
+
wts[0]![i] = Math.fround((1 - dh) * (1 - dw));
|
|
298
|
+
wts[1]![i] = Math.fround((1 - dh) * dw);
|
|
299
|
+
wts[2]![i] = Math.fround(dh * (1 - dw));
|
|
300
|
+
wts[3]![i] = Math.fround(dh * dw);
|
|
301
|
+
}
|
|
302
|
+
}
|
|
303
|
+
// pos_embeds[k] = pos_embed[idx_k] * w_k (bf16, like the reference's
|
|
304
|
+
// weight tensor cast); summed pairwise in reference order.
|
|
305
|
+
let acc: MlxArray | null = null;
|
|
306
|
+
for (let k = 0; k < 4; k++) {
|
|
307
|
+
const ids = MlxArray.fromInt32(idx[k]!, [n]);
|
|
308
|
+
const rows = ops.takeAxis(this.#posEmbed, ids, 0); // [n, 1152] bf16
|
|
309
|
+
ids.dispose();
|
|
310
|
+
const wf = MlxArray.fromFloat32(wts[k]!, [n, 1]);
|
|
311
|
+
const wb = disposing(wf, wf.astype(Dtype.bfloat16));
|
|
312
|
+
const term = ops.mul(rows, wb);
|
|
313
|
+
rows.dispose();
|
|
314
|
+
wb.dispose();
|
|
315
|
+
acc = acc ? disposing(acc, disposing(term, ops.add(acc, term))) : term;
|
|
316
|
+
}
|
|
317
|
+
// Merge-block-major permutation (rows currently h-major patch order).
|
|
318
|
+
const ms = QWEN3VL_MERGE_SIZE;
|
|
319
|
+
const permIdx = new Int32Array(n);
|
|
320
|
+
let r = 0;
|
|
321
|
+
for (let h1 = 0; h1 < gridH / ms; h1++)
|
|
322
|
+
for (let w1 = 0; w1 < gridW / ms; w1++)
|
|
323
|
+
for (let mh = 0; mh < ms; mh++)
|
|
324
|
+
for (let mw = 0; mw < ms; mw++)
|
|
325
|
+
permIdx[r++] = (h1 * ms + mh) * gridW + (w1 * ms + mw);
|
|
326
|
+
const pIds = MlxArray.fromInt32(permIdx, [n]);
|
|
327
|
+
const out = ops.takeAxis(acc!, pIds, 0);
|
|
328
|
+
pIds.dispose();
|
|
329
|
+
acc!.dispose();
|
|
330
|
+
return out;
|
|
331
|
+
}
|
|
332
|
+
|
|
333
|
+
/** 2D rotary cos/sin [n, HEAD_DIM] f32 in merge-block-major order —
|
|
334
|
+
* port of rot_pos_emb + apply_rotary_pos_emb_vision's tile(·,2). */
|
|
335
|
+
#ropeCosSin(gridH: number, gridW: number): [MlxArray, MlxArray] {
|
|
336
|
+
const n = gridH * gridW;
|
|
337
|
+
const half = ROT_DIM / 2; // 18 inv freqs
|
|
338
|
+
// inv_freq on-device, the reference's exact graph — Metal powf is not
|
|
339
|
+
// correctly rounded, so a JS Math.pow emulation differs by 1 ulp on some
|
|
340
|
+
// entries (seen as 2.4e-7 in the freq gate before this).
|
|
341
|
+
const invFreq = (() => {
|
|
342
|
+
const ar = ops.arange(0, ROT_DIM, 2, Dtype.float32);
|
|
343
|
+
const ex = disposing(ar, ops.mulScalar(ar, 1 / ROT_DIM));
|
|
344
|
+
const theta = MlxArray.fromFloat32(Float32Array.from([ROPE_THETA]), [1]);
|
|
345
|
+
const p = ops.pow(theta, ex);
|
|
346
|
+
theta.dispose();
|
|
347
|
+
ex.dispose();
|
|
348
|
+
const one = MlxArray.fromFloat32(Float32Array.from([1]), [1]);
|
|
349
|
+
const inv = disposing(p, ops.div(one, p));
|
|
350
|
+
one.dispose();
|
|
351
|
+
const v = inv.toFloat32();
|
|
352
|
+
inv.dispose();
|
|
353
|
+
return v;
|
|
354
|
+
})();
|
|
355
|
+
const ms = QWEN3VL_MERGE_SIZE;
|
|
356
|
+
const freqs = new Float32Array(n * ROT_DIM);
|
|
357
|
+
let r = 0;
|
|
358
|
+
for (let h1 = 0; h1 < gridH / ms; h1++) {
|
|
359
|
+
for (let w1 = 0; w1 < gridW / ms; w1++) {
|
|
360
|
+
for (let mh = 0; mh < ms; mh++) {
|
|
361
|
+
for (let mw = 0; mw < ms; mw++) {
|
|
362
|
+
const hPos = h1 * ms + mh;
|
|
363
|
+
const wPos = w1 * ms + mw;
|
|
364
|
+
const base = r * ROT_DIM;
|
|
365
|
+
for (let i = 0; i < half; i++) {
|
|
366
|
+
freqs[base + i] = Math.fround(hPos * invFreq[i]!);
|
|
367
|
+
freqs[base + half + i] = Math.fround(wPos * invFreq[i]!);
|
|
368
|
+
}
|
|
369
|
+
r++;
|
|
370
|
+
}
|
|
371
|
+
}
|
|
372
|
+
}
|
|
373
|
+
}
|
|
374
|
+
const f = MlxArray.fromFloat32(freqs, [n, ROT_DIM]);
|
|
375
|
+
const cosH = ops.cos(f);
|
|
376
|
+
const sinH = ops.sin(f);
|
|
377
|
+
f.dispose();
|
|
378
|
+
// tile(·, (1, 2)) over the last dim → [n, HEAD_DIM].
|
|
379
|
+
const cos = disposing(cosH, ops.concatAxis([cosH, cosH], -1));
|
|
380
|
+
const sin = disposing(sinH, ops.concatAxis([sinH, sinH], -1));
|
|
381
|
+
return [cos, sin];
|
|
382
|
+
}
|
|
383
|
+
|
|
384
|
+
/** (t·cos) + (rotate_half(t)·sin), f32 math, cast back to bf16. */
|
|
385
|
+
#applyRope(x: MlxArray, cos: MlxArray, sin: MlxArray): MlxArray {
|
|
386
|
+
// x [n, HEADS, HEAD_DIM]; cos/sin [n, 1, HEAD_DIM] broadcast over heads.
|
|
387
|
+
const half = HEAD_DIM / 2;
|
|
388
|
+
const n = x.shape[0]!;
|
|
389
|
+
const x1 = x.slice([0, 0, 0], [n, HEADS, half]);
|
|
390
|
+
const x2 = x.slice([0, 0, half], [n, HEADS, HEAD_DIM]);
|
|
391
|
+
const negX2 = ops.mulScalar(x2, -1);
|
|
392
|
+
x2.dispose();
|
|
393
|
+
const rot = disposing(negX2, ops.concatAxis([negX2, x1], -1));
|
|
394
|
+
x1.dispose();
|
|
395
|
+
const a = ops.mul(x, cos);
|
|
396
|
+
const b = disposing(rot, ops.mul(rot, sin));
|
|
397
|
+
const sum = ops.add(a, b);
|
|
398
|
+
a.dispose();
|
|
399
|
+
b.dispose();
|
|
400
|
+
return disposing(sum, sum.astype(Dtype.bfloat16));
|
|
401
|
+
}
|
|
402
|
+
|
|
403
|
+
/** Stage-instrumented encode for parity bisection (tests/debug): returns
|
|
404
|
+
* named f32 copies of each pipeline stage alongside the final output. */
|
|
405
|
+
encodeStages(pp: Qwen3VLPreprocessed): Map<string, Float32Array> {
|
|
406
|
+
const stages = new Map<string, Float32Array>();
|
|
407
|
+
const out = this.encode(pp, (name, arr) => {
|
|
408
|
+
const f = arr.astype(Dtype.float32);
|
|
409
|
+
stages.set(name, f.toFloat32());
|
|
410
|
+
f.dispose();
|
|
411
|
+
});
|
|
412
|
+
const f = out.astype(Dtype.float32);
|
|
413
|
+
stages.set("merged", f.toFloat32());
|
|
414
|
+
f.dispose();
|
|
415
|
+
out.dispose();
|
|
416
|
+
return stages;
|
|
417
|
+
}
|
|
418
|
+
|
|
419
|
+
/** Parity-bisect helper: run ONLY the block chain (+ optional merger) on a
|
|
420
|
+
* caller-provided hidden grid — used to separate patch-embed seed ulp
|
|
421
|
+
* from block-math divergence in the stage gates. */
|
|
422
|
+
debugRunBlocks(
|
|
423
|
+
hF32: Float32Array, gridH: number, gridW: number, throughMerger = false,
|
|
424
|
+
): Float32Array {
|
|
425
|
+
const n = gridH * gridW;
|
|
426
|
+
const hf = MlxArray.fromFloat32(hF32, [n, HIDDEN]);
|
|
427
|
+
let h = disposing(hf, hf.astype(Dtype.bfloat16));
|
|
428
|
+
const [cosF, sinF] = this.#ropeCosSin(gridH, gridW);
|
|
429
|
+
const cos = disposing(cosF, ops.expandDims(cosF, 1));
|
|
430
|
+
const sin = disposing(sinF, ops.expandDims(sinF, 1));
|
|
431
|
+
const scale = Math.pow(HEAD_DIM, -0.5);
|
|
432
|
+
for (const blk of this.#blocks) h = this.#block(h, blk, cos, sin, scale, n);
|
|
433
|
+
cos.dispose();
|
|
434
|
+
sin.dispose();
|
|
435
|
+
if (throughMerger) {
|
|
436
|
+
const mnorm = disposing(h, ops.layerNorm(h, this.#mergerNormW, this.#mergerNormB, 1e-6));
|
|
437
|
+
const grouped = disposing(mnorm, ops.reshape(mnorm, [n / 4, HIDDEN * 4]));
|
|
438
|
+
const f1 = disposing(grouped, this.#linear(grouped, this.#mergerFc1W, this.#mergerFc1B));
|
|
439
|
+
const act = disposing(f1, ops.geluPrecise(f1));
|
|
440
|
+
h = disposing(act, this.#linear(act, this.#mergerFc2W, this.#mergerFc2B));
|
|
441
|
+
}
|
|
442
|
+
const f = disposing(h, h.astype(Dtype.float32));
|
|
443
|
+
const out = f.toFloat32();
|
|
444
|
+
f.dispose();
|
|
445
|
+
return out;
|
|
446
|
+
}
|
|
447
|
+
|
|
448
|
+
/** Encode one preprocessed image → [imageTokens, 5120] bf16. */
|
|
449
|
+
encode(
|
|
450
|
+
pp: Qwen3VLPreprocessed,
|
|
451
|
+
tap?: (name: string, arr: MlxArray) => void,
|
|
452
|
+
): MlxArray {
|
|
453
|
+
const [gridT, gridH, gridW] = pp.gridThw;
|
|
454
|
+
const n = pp.rows;
|
|
455
|
+
const pvF = MlxArray.fromFloat32(pp.pixelValues, [n, pp.cols]);
|
|
456
|
+
const pv = disposing(pvF, pvF.astype(Dtype.bfloat16));
|
|
457
|
+
// PatchEmbed: rows (c,t,py,px) → [N, tps, ps, ps, C] NDHWC → Conv3d with
|
|
458
|
+
// kernel == stride — the reference's exact reshape/moveaxis/conv graph.
|
|
459
|
+
const tps = QWEN3VL_TEMPORAL_PATCH_SIZE;
|
|
460
|
+
const ps = QWEN3VL_PATCH_SIZE;
|
|
461
|
+
const r5 = disposing(pv, ops.reshape(pv, [n, 3, tps, ps, ps]));
|
|
462
|
+
const tr = disposing(r5, ops.transposeAxes(r5, [0, 2, 3, 4, 1]));
|
|
463
|
+
const ndhwc = disposing(tr, ops.contiguous(tr));
|
|
464
|
+
const convOut = disposing(
|
|
465
|
+
ndhwc, ops.conv3d(ndhwc, this.#patchW, [tps, ps, ps]),
|
|
466
|
+
);
|
|
467
|
+
let h = disposing(convOut, ops.reshape(convOut, [n, HIDDEN]));
|
|
468
|
+
h = disposing(h, ops.add(h, this.#patchB));
|
|
469
|
+
tap?.("patch", h);
|
|
470
|
+
// Per-temporal-group pos/rope embeds are identical — tile the (h,w)
|
|
471
|
+
// block gridT times (the reference tiles coords/pos_embed per frame).
|
|
472
|
+
const tile = (a: MlxArray): MlxArray => {
|
|
473
|
+
if (gridT === 1) return a;
|
|
474
|
+
const parts = new Array(gridT).fill(a);
|
|
475
|
+
return disposing(a, ops.concatAxis(parts, 0));
|
|
476
|
+
};
|
|
477
|
+
const pos = tile(this.#interpolatedPosEmbeds(gridH, gridW));
|
|
478
|
+
tap?.("pos", pos);
|
|
479
|
+
h = disposing(h, disposing(pos, ops.add(h, pos)));
|
|
480
|
+
|
|
481
|
+
const [cosF0, sinF0] = this.#ropeCosSin(gridH, gridW);
|
|
482
|
+
const cosF = tile(cosF0);
|
|
483
|
+
const sinF = tile(sinF0);
|
|
484
|
+
const cos = disposing(cosF, ops.expandDims(cosF, 1)); // [n,1,72]
|
|
485
|
+
const sin = disposing(sinF, ops.expandDims(sinF, 1));
|
|
486
|
+
|
|
487
|
+
const scale = Math.pow(HEAD_DIM, -0.5);
|
|
488
|
+
// Attention never crosses temporal groups (cu_seqlens = h·w per frame
|
|
489
|
+
// pair): images run one segment, videos gridT segments.
|
|
490
|
+
const segLen = gridH * gridW;
|
|
491
|
+
let blkIdx = 0;
|
|
492
|
+
for (const blk of this.#blocks) {
|
|
493
|
+
h = this.#block(h, blk, cos, sin, scale, n, segLen);
|
|
494
|
+
tap?.(`blk${blkIdx}`, h);
|
|
495
|
+
blkIdx++;
|
|
496
|
+
}
|
|
497
|
+
tap?.("blocks", h);
|
|
498
|
+
cos.dispose();
|
|
499
|
+
sin.dispose();
|
|
500
|
+
|
|
501
|
+
// Merger: LayerNorm(1152) → reshape 4× → fc1 → GELU(erf) → fc2.
|
|
502
|
+
const mnorm = disposing(h, ops.layerNorm(h, this.#mergerNormW, this.#mergerNormB, 1e-6));
|
|
503
|
+
const grouped = disposing(mnorm, ops.reshape(mnorm, [n / 4, HIDDEN * 4]));
|
|
504
|
+
const f1 = disposing(grouped, this.#linear(grouped, this.#mergerFc1W, this.#mergerFc1B));
|
|
505
|
+
const act = disposing(f1, ops.geluPrecise(f1)); // reference nn.GELU()
|
|
506
|
+
return disposing(act, this.#linear(act, this.#mergerFc2W, this.#mergerFc2B));
|
|
507
|
+
}
|
|
508
|
+
|
|
509
|
+
dispose(): void {
|
|
510
|
+
// Sidecar arrays are process-pinned (house rule: mmap-backed weights are
|
|
511
|
+
// never handed a JS dtor); nothing to free eagerly.
|
|
512
|
+
}
|
|
513
|
+
}
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
// Video file → RGB frames via the AVFoundation sidecar (PLAN 14w) — the
|
|
2
|
+
// video twin of the audio path's afconvert doctrine: macOS-native codecs
|
|
3
|
+
// (H.264/HEVC/ProRes/…), zero vendored decoders. The sidecar
|
|
4
|
+
// (src/native/frame_extract.swift, shipped as `mlx-bun-frame-extract` beside
|
|
5
|
+
// the binary and in the native pack) samples frames at fps=2 with
|
|
6
|
+
// zero-tolerance timestamps — mirroring mlx-vlm's do_sample_frames — and
|
|
7
|
+
// writes PNGs that decode through the SAME fast-png path as image inputs.
|
|
8
|
+
//
|
|
9
|
+
// Resolution order (afconvert needs none of this because /usr/bin ships it;
|
|
10
|
+
// our extractor is the missing CLI face on Apple's decoder):
|
|
11
|
+
// MLX_BUN_FRAME_EXTRACT env → beside the executable (release bundle) →
|
|
12
|
+
// the native-pack cache → dist-native (dev tree) → compile-on-demand from
|
|
13
|
+
// src/native (from-source installs with Xcode CLT; cached in dist-native).
|
|
14
|
+
|
|
15
|
+
import { existsSync, mkdirSync } from "node:fs";
|
|
16
|
+
import { rm } from "node:fs/promises";
|
|
17
|
+
import { randomUUID } from "node:crypto";
|
|
18
|
+
import { tmpdir } from "node:os";
|
|
19
|
+
import { dirname, join } from "node:path";
|
|
20
|
+
import { nativePackDir } from "../native-pack";
|
|
21
|
+
import { decodeImage, type RGBImage } from "./preprocess";
|
|
22
|
+
import { runtimeValue } from "../runtime-config";
|
|
23
|
+
|
|
24
|
+
export const VIDEO_SAMPLE_FPS = 2;
|
|
25
|
+
export const VIDEO_MAX_FRAMES = 768; // Qwen3VLVideoProcessor max_frames
|
|
26
|
+
|
|
27
|
+
const BIN = "mlx-bun-frame-extract";
|
|
28
|
+
|
|
29
|
+
/** Memoize the PROMISE, not the result: a second video request arriving
|
|
30
|
+
* while the first compile is in flight must await the same compile, not
|
|
31
|
+
* observe a half-set cache and spuriously report the sidecar missing. */
|
|
32
|
+
let devCompile: Promise<string | null> | null = null;
|
|
33
|
+
|
|
34
|
+
/** Compile the sidecar from the repo tree (dev/from-source convenience —
|
|
35
|
+
* needs swiftc from the Xcode CLT). Null when the source or toolchain is
|
|
36
|
+
* unavailable. */
|
|
37
|
+
function compileFromSource(): Promise<string | null> {
|
|
38
|
+
devCompile ??= (async () => {
|
|
39
|
+
const src = join(import.meta.dir, "..", "native", "frame_extract.swift");
|
|
40
|
+
if (!existsSync(src)) return null;
|
|
41
|
+
const out = join(import.meta.dir, "..", "..", "dist-native", BIN);
|
|
42
|
+
try {
|
|
43
|
+
mkdirSync(dirname(out), { recursive: true });
|
|
44
|
+
const proc = Bun.spawn(["swiftc", "-O", src, "-o", out], {
|
|
45
|
+
stdout: "ignore", stderr: "pipe",
|
|
46
|
+
});
|
|
47
|
+
if ((await proc.exited) === 0 && existsSync(out)) return out;
|
|
48
|
+
} catch {
|
|
49
|
+
// no swiftc — resolver returns null and serving 400s with the reason
|
|
50
|
+
}
|
|
51
|
+
return null;
|
|
52
|
+
})();
|
|
53
|
+
return devCompile;
|
|
54
|
+
}
|
|
55
|
+
|
|
56
|
+
/** Locate (or dev-compile) the extractor; null ⟹ video input unavailable
|
|
57
|
+
* on this install (surfaced as a clear 400, never a crash). */
|
|
58
|
+
export async function resolveFrameExtract(): Promise<string | null> {
|
|
59
|
+
const explicit = runtimeValue("MLX_BUN_FRAME_EXTRACT");
|
|
60
|
+
if (explicit) return explicit;
|
|
61
|
+
const candidates = [
|
|
62
|
+
join(dirname(process.execPath), BIN),
|
|
63
|
+
join(nativePackDir(), BIN),
|
|
64
|
+
join(import.meta.dir, "..", "..", "dist-native", BIN),
|
|
65
|
+
];
|
|
66
|
+
for (const c of candidates) if (existsSync(c)) return c;
|
|
67
|
+
return compileFromSource();
|
|
68
|
+
}
|
|
69
|
+
|
|
70
|
+
/** Decode a video container to sampled RGB frames. Throws with the
|
|
71
|
+
* sidecar's diagnostic on failure (the serve layer surfaces it as a 400,
|
|
72
|
+
* the afconvert pattern — timeout kill included). */
|
|
73
|
+
export async function extractVideoFrames(
|
|
74
|
+
bytes: Uint8Array,
|
|
75
|
+
opts: { fps?: number; maxFrames?: number } = {},
|
|
76
|
+
): Promise<RGBImage[]> {
|
|
77
|
+
const bin = await resolveFrameExtract();
|
|
78
|
+
if (!bin)
|
|
79
|
+
throw new Error(
|
|
80
|
+
"video input needs the frame-extraction sidecar (mlx-bun-frame-extract) " +
|
|
81
|
+
"— ships beside the release binary and in the native pack; from a " +
|
|
82
|
+
"source tree it compiles on first use via the Xcode CLT (swiftc)",
|
|
83
|
+
);
|
|
84
|
+
const base = join(tmpdir(), `mlx-bun-video-${randomUUID()}`);
|
|
85
|
+
// AVFoundation infers the container from the path EXTENSION — sniff the
|
|
86
|
+
// ISO-BMFF `ftyp` box (mp4/mov share it; the `qt ` major brand is
|
|
87
|
+
// QuickTime) so the temp file opens. Unknown magic defaults to .mp4.
|
|
88
|
+
const isFtyp = bytes.length >= 12 &&
|
|
89
|
+
bytes[4] === 0x66 && bytes[5] === 0x74 && bytes[6] === 0x79 && bytes[7] === 0x70;
|
|
90
|
+
const brand = isFtyp
|
|
91
|
+
? String.fromCharCode(bytes[8]!, bytes[9]!, bytes[10]!, bytes[11]!)
|
|
92
|
+
: "";
|
|
93
|
+
const ext = brand.startsWith("qt") ? "mov" : "mp4";
|
|
94
|
+
const inPath = `${base}.${ext}`;
|
|
95
|
+
const outDir = `${base}.frames`;
|
|
96
|
+
try {
|
|
97
|
+
// Inside the try so a part-way write failure (ENOSPC on a large clip)
|
|
98
|
+
// still hits the finally's cleanup instead of stranding the temp file.
|
|
99
|
+
await Bun.write(inPath, bytes);
|
|
100
|
+
const proc = Bun.spawn(
|
|
101
|
+
[
|
|
102
|
+
bin, inPath, outDir,
|
|
103
|
+
String(opts.fps ?? VIDEO_SAMPLE_FPS),
|
|
104
|
+
String(opts.maxFrames ?? VIDEO_MAX_FRAMES),
|
|
105
|
+
],
|
|
106
|
+
{ stdout: "pipe", stderr: "pipe" },
|
|
107
|
+
);
|
|
108
|
+
// Decoding is fast (hardware) but long clips take real time; a hung
|
|
109
|
+
// decoder must not pin the request forever — kill → nonzero → 400.
|
|
110
|
+
const killer = setTimeout(() => proc.kill(), 120_000);
|
|
111
|
+
const exit = await proc.exited;
|
|
112
|
+
clearTimeout(killer);
|
|
113
|
+
if (exit !== 0) {
|
|
114
|
+
const err = (await new Response(proc.stderr).text()).trim();
|
|
115
|
+
throw new Error(
|
|
116
|
+
`video decode failed (frame-extract exit ${exit}): ` +
|
|
117
|
+
`${err.split("\n").at(-1) || "unknown error"} — supported containers: ` +
|
|
118
|
+
`whatever AVFoundation decodes on this macOS (mp4/mov H.264/HEVC/ProRes/…)`,
|
|
119
|
+
);
|
|
120
|
+
}
|
|
121
|
+
const head = (await new Response(proc.stdout).text()).trim();
|
|
122
|
+
const n = Number(/^N (\d+)/.exec(head)?.[1] ?? 0);
|
|
123
|
+
if (!n) throw new Error(`video decode produced no frames (${head})`);
|
|
124
|
+
// Aggregate decoded-RGB budget (belt over the sidecar's 1024px
|
|
125
|
+
// longest-edge cap): a valid low-bitrate clip must not balloon into
|
|
126
|
+
// gigabytes of frame buffers before preprocessing.
|
|
127
|
+
const w = Number(/W (\d+)/.exec(head)?.[1] ?? 0);
|
|
128
|
+
const hh = Number(/H (\d+)/.exec(head)?.[1] ?? 0);
|
|
129
|
+
const DECODED_BUDGET = 1536 * 1024 * 1024; // 1.5 GiB
|
|
130
|
+
if (n * w * hh * 3 > DECODED_BUDGET)
|
|
131
|
+
throw new Error(
|
|
132
|
+
`video decodes to ${(n * w * hh * 3 / 1024 / 1024).toFixed(0)} MB of ` +
|
|
133
|
+
`frames (${n} × ${w}x${hh}) — over the ${DECODED_BUDGET / 1024 / 1024} MB ` +
|
|
134
|
+
`budget; send a shorter or smaller clip`,
|
|
135
|
+
);
|
|
136
|
+
const frames: RGBImage[] = [];
|
|
137
|
+
for (let i = 0; i < n; i++) {
|
|
138
|
+
const f = join(outDir, `frame-${String(i).padStart(4, "0")}.png`);
|
|
139
|
+
frames.push(await decodeImage(new Uint8Array(await Bun.file(f).arrayBuffer())));
|
|
140
|
+
}
|
|
141
|
+
return frames;
|
|
142
|
+
} finally {
|
|
143
|
+
await Promise.allSettled([
|
|
144
|
+
rm(inPath, { force: true }),
|
|
145
|
+
rm(outDir, { recursive: true, force: true }),
|
|
146
|
+
]);
|
|
147
|
+
}
|
|
148
|
+
}
|