mlx-bun 0.0.11 → 0.0.12
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 +43 -13
- package/THIRD_PARTY_LICENSES.md +9 -0
- package/docs/design/colibri-glm52-port.md +1479 -0
- package/docs/design/pre-colibri-stabilization.md +329 -0
- package/docs/investigations/colibri-oracle-pin.md +1284 -0
- package/docs/investigations/indexshare-performance-spike.md +357 -0
- package/docs/planning/release-notes-v0.0.12.md +80 -0
- package/docs/reference/cli.md +28 -2
- package/docs/reference/distribution.md +25 -1
- package/docs/reference/embedding.md +5 -2
- package/docs/reference/features-matrix.md +4 -2
- package/docs/reference/library-api.md +41 -13
- package/docs/reference/models.md +49 -2
- package/docs/reference/server-api.md +90 -9
- package/docs/reference/server-config.md +21 -14
- package/docs/reference/training.md +10 -1
- package/package.json +1 -1
- package/src/anthropic.ts +54 -44
- package/src/chat-template.ts +155 -10
- package/src/cli.ts +127 -8
- package/src/config.ts +71 -26
- package/src/diffusion/diffusion-generate.ts +14 -8
- package/src/download.ts +103 -12
- package/src/eval/ifeval.ts +63 -120
- package/src/eval/runner.ts +14 -5
- package/src/eval/tasks/ifeval.ts +167 -31
- package/src/expert-io.ts +391 -0
- package/src/expert-residency.ts +1134 -0
- package/src/expert-usage.ts +415 -0
- package/src/generate.ts +114 -17
- package/src/hf-push.ts +1 -1
- package/src/index.ts +14 -2
- package/src/jobs/runner.ts +37 -18
- package/src/kv-store.ts +147 -8
- package/src/lora.ts +77 -14
- package/src/media-fetch.ts +198 -56
- package/src/mlx/array.ts +43 -0
- package/src/mlx/autograd.ts +51 -14
- package/src/mlx/checkpoint.ts +46 -13
- package/src/mlx/custom-vjp.ts +83 -35
- package/src/model/factory.ts +96 -5
- package/src/model/gemma4-base.ts +32 -0
- package/src/model/glm52-atlas.ts +511 -0
- package/src/model/glm52-cache.ts +706 -0
- package/src/model/glm52-config.ts +262 -0
- package/src/model/glm52-container.ts +305 -0
- package/src/model/glm52-coupling.ts +388 -0
- package/src/model/glm52-dsa.ts +602 -0
- package/src/model/glm52-expert-layout.ts +253 -0
- package/src/model/glm52-layout.ts +152 -0
- package/src/model/glm52-memory.ts +473 -0
- package/src/model/glm52-mla.ts +786 -0
- package/src/model/glm52-moe.ts +223 -0
- package/src/model/glm52-pilot.ts +537 -0
- package/src/model/glm52-reference.ts +445 -0
- package/src/model/glm52-residency.ts +557 -0
- package/src/model/glm52-resident-weights.ts +135 -0
- package/src/model/glm52-streamed-experts.ts +453 -0
- package/src/model/glm52-streamed-metal.ts +610 -0
- package/src/model/glm52-weights.ts +503 -0
- package/src/model/glm52.ts +1153 -0
- package/src/model/support.ts +8 -1
- package/src/native/expert_io.c +523 -0
- package/src/native-pack.ts +6 -4
- package/src/responses.ts +62 -49
- package/src/serve/batch-scheduler.ts +95 -16
- package/src/serve/generation-gateway.ts +41 -16
- package/src/server.ts +362 -81
- package/src/spec/glm52-mtp-source.ts +392 -0
- package/src/spec/serve-loop.ts +99 -18
- package/src/spec/source.ts +31 -6
- package/src/ssd-cache.ts +6 -2
- package/src/tool-call.ts +30 -2
- package/src/train/job.ts +0 -10
- package/src/train/lora-params.ts +17 -6
- package/src/train/trainer.ts +59 -26
- package/src/web/app.html +13 -2
- package/src/web/app.js +30 -3
- package/src/web/src/status.ts +51 -3
- package/src/web-tools.ts +54 -2
- package/src/weights.ts +82 -19
|
@@ -0,0 +1,415 @@
|
|
|
1
|
+
import {
|
|
2
|
+
closeSync,
|
|
3
|
+
existsSync,
|
|
4
|
+
fsyncSync,
|
|
5
|
+
openSync,
|
|
6
|
+
readFileSync,
|
|
7
|
+
renameSync,
|
|
8
|
+
rmSync,
|
|
9
|
+
writeSync,
|
|
10
|
+
} from "node:fs";
|
|
11
|
+
|
|
12
|
+
const UINT32_MAX = 0xffff_ffff;
|
|
13
|
+
|
|
14
|
+
export interface ExpertUsageRoute {
|
|
15
|
+
readonly indices: ArrayLike<number>;
|
|
16
|
+
}
|
|
17
|
+
|
|
18
|
+
export interface ExpertUsageLedgerOptions {
|
|
19
|
+
readonly path: string;
|
|
20
|
+
readonly layers: readonly number[];
|
|
21
|
+
readonly expertsPerLayer: number;
|
|
22
|
+
/** A damaged profile is derived state: report it and start fresh. */
|
|
23
|
+
readonly onWarning?: (message: string) => void;
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
export interface ExpertUsageEntry {
|
|
27
|
+
readonly count: number;
|
|
28
|
+
readonly heat: number;
|
|
29
|
+
readonly lastAccess: number;
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
export interface ExpertUsageSnapshot {
|
|
33
|
+
readonly clock: number;
|
|
34
|
+
readonly routeSelections: number;
|
|
35
|
+
readonly totalCount: number;
|
|
36
|
+
readonly nonzeroCounts: number;
|
|
37
|
+
readonly dirty: boolean;
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
export interface ExpertUsageCount {
|
|
41
|
+
readonly layer: number;
|
|
42
|
+
readonly expertId: number;
|
|
43
|
+
readonly count: number;
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
export interface ExpertAutoPinPlan {
|
|
47
|
+
readonly historySelections: number;
|
|
48
|
+
readonly confidence: number;
|
|
49
|
+
readonly requestedBudgetBytes: number;
|
|
50
|
+
readonly usableBudgetBytes: number;
|
|
51
|
+
readonly plannedBytes: number;
|
|
52
|
+
readonly pins: readonly ExpertUsageCount[];
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
export interface ExpertAutoPinInput {
|
|
56
|
+
readonly ledger: ExpertUsageLedger;
|
|
57
|
+
/** Bytes available to the persistent LRU + pin tiers after working banks. */
|
|
58
|
+
readonly residentTierBudgetBytes: number;
|
|
59
|
+
/** One non-evictable LRU slot per managed layer. */
|
|
60
|
+
readonly mandatoryResidentBytes: number;
|
|
61
|
+
readonly slotBytes: (layer: number) => number;
|
|
62
|
+
readonly minHistorySelections?: number;
|
|
63
|
+
readonly fullConfidenceSelections?: number;
|
|
64
|
+
readonly pinShare?: number;
|
|
65
|
+
readonly minPinBudgetBytes?: number;
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
export interface ExpertLfruCandidate {
|
|
69
|
+
readonly layer: number;
|
|
70
|
+
readonly coldSlot: number;
|
|
71
|
+
readonly coldExpertId: number;
|
|
72
|
+
readonly hotExpertId: number;
|
|
73
|
+
readonly coldScore: number;
|
|
74
|
+
readonly hotScore: number;
|
|
75
|
+
readonly gain: number;
|
|
76
|
+
}
|
|
77
|
+
|
|
78
|
+
export function expertLfruScore(
|
|
79
|
+
heat: number,
|
|
80
|
+
lastAccess: number,
|
|
81
|
+
clock: number,
|
|
82
|
+
): number {
|
|
83
|
+
const age = (clock - lastAccess) >>> 0;
|
|
84
|
+
return heat * 256 + (age < 255 ? 255 - age : 0);
|
|
85
|
+
}
|
|
86
|
+
|
|
87
|
+
export function selectExpertLfruCandidates(
|
|
88
|
+
candidates: readonly ExpertLfruCandidate[],
|
|
89
|
+
limit = 4,
|
|
90
|
+
): ExpertLfruCandidate[] {
|
|
91
|
+
positiveSafeInteger(limit, "live-repin limit");
|
|
92
|
+
return candidates.slice().sort((left, right) =>
|
|
93
|
+
right.gain - left.gain ||
|
|
94
|
+
left.layer - right.layer ||
|
|
95
|
+
left.hotExpertId - right.hotExpertId)
|
|
96
|
+
.slice(0, limit);
|
|
97
|
+
}
|
|
98
|
+
|
|
99
|
+
function positiveSafeInteger(value: number, name: string): void {
|
|
100
|
+
if (!Number.isSafeInteger(value) || value <= 0)
|
|
101
|
+
throw new RangeError(`${name} must be a positive safe integer`);
|
|
102
|
+
}
|
|
103
|
+
|
|
104
|
+
function incrementUint32(value: number, amount: number): number {
|
|
105
|
+
return Math.min(UINT32_MAX, value + amount);
|
|
106
|
+
}
|
|
107
|
+
|
|
108
|
+
/**
|
|
109
|
+
* Colibri-compatible long-term expert frequency plus session-local heat and
|
|
110
|
+
* recency. The on-disk format is one `layer expert count` row per nonzero
|
|
111
|
+
* counter, matching Colibri's `.coli_usage` reader.
|
|
112
|
+
*/
|
|
113
|
+
export class ExpertUsageLedger {
|
|
114
|
+
readonly path: string;
|
|
115
|
+
readonly layers: readonly number[];
|
|
116
|
+
readonly expertsPerLayer: number;
|
|
117
|
+
#counts = new Map<number, Uint32Array>();
|
|
118
|
+
#heat = new Map<number, Uint32Array>();
|
|
119
|
+
#last = new Map<number, Uint32Array>();
|
|
120
|
+
#clock = 0;
|
|
121
|
+
#routeSelections = 0;
|
|
122
|
+
#dirty = false;
|
|
123
|
+
|
|
124
|
+
private constructor(options: ExpertUsageLedgerOptions) {
|
|
125
|
+
if (!options.path) throw new Error("expert usage path must not be empty");
|
|
126
|
+
positiveSafeInteger(options.expertsPerLayer, "expertsPerLayer");
|
|
127
|
+
if (options.expertsPerLayer > UINT32_MAX)
|
|
128
|
+
throw new RangeError("expertsPerLayer exceeds uint32 range");
|
|
129
|
+
const layers = options.layers.slice();
|
|
130
|
+
if (layers.length === 0 || new Set(layers).size !== layers.length)
|
|
131
|
+
throw new Error("expert usage layers must be non-empty and unique");
|
|
132
|
+
for (const layer of layers)
|
|
133
|
+
if (!Number.isSafeInteger(layer) || layer < 0)
|
|
134
|
+
throw new RangeError(`invalid expert usage layer ${layer}`);
|
|
135
|
+
this.path = options.path;
|
|
136
|
+
this.layers = Object.freeze(layers);
|
|
137
|
+
this.expertsPerLayer = options.expertsPerLayer;
|
|
138
|
+
for (const layer of layers) {
|
|
139
|
+
this.#counts.set(layer, new Uint32Array(options.expertsPerLayer));
|
|
140
|
+
this.#heat.set(layer, new Uint32Array(options.expertsPerLayer));
|
|
141
|
+
this.#last.set(layer, new Uint32Array(options.expertsPerLayer));
|
|
142
|
+
}
|
|
143
|
+
}
|
|
144
|
+
|
|
145
|
+
static open(options: ExpertUsageLedgerOptions): ExpertUsageLedger {
|
|
146
|
+
const ledger = new ExpertUsageLedger(options);
|
|
147
|
+
if (!existsSync(options.path)) return ledger;
|
|
148
|
+
try {
|
|
149
|
+
ledger.#load(readFileSync(options.path, "utf8"));
|
|
150
|
+
} catch (error) {
|
|
151
|
+
const detail = error instanceof Error ? error.message : String(error);
|
|
152
|
+
options.onWarning?.(
|
|
153
|
+
`ignoring invalid expert usage profile ${options.path}: ${detail}`,
|
|
154
|
+
);
|
|
155
|
+
ledger.#clearCounts();
|
|
156
|
+
}
|
|
157
|
+
return ledger;
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
/** Count every selected route before the execution plan deduplicates it. */
|
|
161
|
+
recordRoutes(layer: number, routes: readonly ExpertUsageRoute[]): void {
|
|
162
|
+
const counts = this.#counts.get(layer);
|
|
163
|
+
const heat = this.#heat.get(layer);
|
|
164
|
+
const last = this.#last.get(layer);
|
|
165
|
+
if (!counts || !heat || !last)
|
|
166
|
+
throw new RangeError(`layer ${layer} is not tracked by expert usage`);
|
|
167
|
+
|
|
168
|
+
const selected: number[] = [];
|
|
169
|
+
for (let row = 0; row < routes.length; row++) {
|
|
170
|
+
for (let rank = 0; rank < routes[row]!.indices.length; rank++) {
|
|
171
|
+
const expertId = Number(routes[row]!.indices[rank]);
|
|
172
|
+
if (!Number.isSafeInteger(expertId) || expertId < 0 ||
|
|
173
|
+
expertId >= this.expertsPerLayer) {
|
|
174
|
+
throw new RangeError(
|
|
175
|
+
`route ${row} rank ${rank} has invalid expert ${expertId}`,
|
|
176
|
+
);
|
|
177
|
+
}
|
|
178
|
+
selected.push(expertId);
|
|
179
|
+
}
|
|
180
|
+
}
|
|
181
|
+
|
|
182
|
+
for (const expertId of selected) {
|
|
183
|
+
counts[expertId] = incrementUint32(counts[expertId]!, 1);
|
|
184
|
+
heat[expertId] = incrementUint32(heat[expertId]!, 1);
|
|
185
|
+
this.#clock = (this.#clock + 1) >>> 0;
|
|
186
|
+
last[expertId] = this.#clock;
|
|
187
|
+
this.#routeSelections++;
|
|
188
|
+
}
|
|
189
|
+
if (selected.length) this.#dirty = true;
|
|
190
|
+
}
|
|
191
|
+
|
|
192
|
+
entry(layer: number, expertId: number): ExpertUsageEntry {
|
|
193
|
+
this.#validateEntry(layer, expertId);
|
|
194
|
+
return {
|
|
195
|
+
count: this.#counts.get(layer)![expertId]!,
|
|
196
|
+
heat: this.#heat.get(layer)![expertId]!,
|
|
197
|
+
lastAccess: this.#last.get(layer)![expertId]!,
|
|
198
|
+
};
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
snapshot(): ExpertUsageSnapshot {
|
|
202
|
+
let nonzeroCounts = 0;
|
|
203
|
+
let totalCount = 0;
|
|
204
|
+
for (const counts of this.#counts.values())
|
|
205
|
+
for (const count of counts) {
|
|
206
|
+
totalCount += count;
|
|
207
|
+
if (count !== 0) nonzeroCounts++;
|
|
208
|
+
}
|
|
209
|
+
return {
|
|
210
|
+
clock: this.#clock,
|
|
211
|
+
routeSelections: this.#routeSelections,
|
|
212
|
+
totalCount,
|
|
213
|
+
nonzeroCounts,
|
|
214
|
+
dirty: this.#dirty,
|
|
215
|
+
};
|
|
216
|
+
}
|
|
217
|
+
|
|
218
|
+
rankedCounts(): ExpertUsageCount[] {
|
|
219
|
+
const ranked: ExpertUsageCount[] = [];
|
|
220
|
+
for (const layer of this.layers) {
|
|
221
|
+
const counts = this.#counts.get(layer)!;
|
|
222
|
+
for (let expertId = 0; expertId < counts.length; expertId++) {
|
|
223
|
+
const count = counts[expertId]!;
|
|
224
|
+
if (count !== 0) ranked.push({ layer, expertId, count });
|
|
225
|
+
}
|
|
226
|
+
}
|
|
227
|
+
ranked.sort((left, right) =>
|
|
228
|
+
right.count - left.count ||
|
|
229
|
+
left.layer - right.layer ||
|
|
230
|
+
left.expertId - right.expertId);
|
|
231
|
+
return ranked;
|
|
232
|
+
}
|
|
233
|
+
|
|
234
|
+
lfruCandidate(
|
|
235
|
+
layer: number,
|
|
236
|
+
pinnedExpertIds: readonly number[],
|
|
237
|
+
): ExpertLfruCandidate | null {
|
|
238
|
+
if (pinnedExpertIds.length === 0) return null;
|
|
239
|
+
for (const expertId of pinnedExpertIds) this.#validateEntry(layer, expertId);
|
|
240
|
+
const pinned = new Set(pinnedExpertIds);
|
|
241
|
+
let coldSlot = 0;
|
|
242
|
+
let coldScore = this.#lfruScore(layer, pinnedExpertIds[0]!);
|
|
243
|
+
for (let slot = 1; slot < pinnedExpertIds.length; slot++) {
|
|
244
|
+
const score = this.#lfruScore(layer, pinnedExpertIds[slot]!);
|
|
245
|
+
if (score < coldScore) {
|
|
246
|
+
coldSlot = slot;
|
|
247
|
+
coldScore = score;
|
|
248
|
+
}
|
|
249
|
+
}
|
|
250
|
+
let hotExpertId = -1;
|
|
251
|
+
let hotScore = 0;
|
|
252
|
+
for (let expertId = 0; expertId < this.expertsPerLayer; expertId++) {
|
|
253
|
+
if (pinned.has(expertId)) continue;
|
|
254
|
+
const score = this.#lfruScore(layer, expertId);
|
|
255
|
+
if (hotExpertId < 0 || score > hotScore) {
|
|
256
|
+
hotExpertId = expertId;
|
|
257
|
+
hotScore = score;
|
|
258
|
+
}
|
|
259
|
+
}
|
|
260
|
+
if (hotExpertId < 0 ||
|
|
261
|
+
hotScore <= coldScore + Math.floor(coldScore / 4) + 4 * 256) {
|
|
262
|
+
return null;
|
|
263
|
+
}
|
|
264
|
+
return {
|
|
265
|
+
layer,
|
|
266
|
+
coldSlot,
|
|
267
|
+
coldExpertId: pinnedExpertIds[coldSlot]!,
|
|
268
|
+
hotExpertId,
|
|
269
|
+
coldScore,
|
|
270
|
+
hotScore,
|
|
271
|
+
gain: Math.floor((hotScore - coldScore) / 256),
|
|
272
|
+
};
|
|
273
|
+
}
|
|
274
|
+
|
|
275
|
+
decayHeat(): void {
|
|
276
|
+
for (const heat of this.#heat.values())
|
|
277
|
+
for (let expertId = 0; expertId < heat.length; expertId++)
|
|
278
|
+
heat[expertId] = heat[expertId]! >>> 1;
|
|
279
|
+
}
|
|
280
|
+
|
|
281
|
+
/** Persist the current profile with same-directory temp + atomic rename. */
|
|
282
|
+
flush(): void {
|
|
283
|
+
if (!this.#dirty) return;
|
|
284
|
+
const rows: string[] = [];
|
|
285
|
+
for (const layer of this.layers) {
|
|
286
|
+
const counts = this.#counts.get(layer)!;
|
|
287
|
+
for (let expertId = 0; expertId < counts.length; expertId++) {
|
|
288
|
+
const count = counts[expertId]!;
|
|
289
|
+
if (count !== 0) rows.push(`${layer} ${expertId} ${count}\n`);
|
|
290
|
+
}
|
|
291
|
+
}
|
|
292
|
+
const temp =
|
|
293
|
+
`${this.path}.tmp-${process.pid}-${Math.trunc(Bun.nanoseconds())}`;
|
|
294
|
+
let fd: number | null = null;
|
|
295
|
+
try {
|
|
296
|
+
fd = openSync(temp, "w", 0o600);
|
|
297
|
+
const data = Buffer.from(rows.join(""));
|
|
298
|
+
let written = 0;
|
|
299
|
+
while (written < data.length) {
|
|
300
|
+
written += writeSync(
|
|
301
|
+
fd,
|
|
302
|
+
data,
|
|
303
|
+
written,
|
|
304
|
+
data.length - written,
|
|
305
|
+
written,
|
|
306
|
+
);
|
|
307
|
+
}
|
|
308
|
+
fsyncSync(fd);
|
|
309
|
+
closeSync(fd);
|
|
310
|
+
fd = null;
|
|
311
|
+
renameSync(temp, this.path);
|
|
312
|
+
this.#dirty = false;
|
|
313
|
+
} catch (error) {
|
|
314
|
+
if (fd !== null) closeSync(fd);
|
|
315
|
+
try { rmSync(temp, { force: true }); } catch {}
|
|
316
|
+
throw error;
|
|
317
|
+
}
|
|
318
|
+
}
|
|
319
|
+
|
|
320
|
+
#load(text: string): void {
|
|
321
|
+
for (const [index, raw] of text.split("\n").entries()) {
|
|
322
|
+
const line = raw.trim();
|
|
323
|
+
if (!line) continue;
|
|
324
|
+
const fields = line.split(/\s+/);
|
|
325
|
+
if (fields.length !== 3)
|
|
326
|
+
throw new Error(`line ${index + 1}: expected layer expert count`);
|
|
327
|
+
const [layer, expertId, count] = fields.map(Number);
|
|
328
|
+
if (!Number.isSafeInteger(layer) || !this.#counts.has(layer!))
|
|
329
|
+
throw new Error(`line ${index + 1}: invalid layer ${fields[0]}`);
|
|
330
|
+
if (!Number.isSafeInteger(expertId) || expertId! < 0 ||
|
|
331
|
+
expertId! >= this.expertsPerLayer) {
|
|
332
|
+
throw new Error(`line ${index + 1}: invalid expert ${fields[1]}`);
|
|
333
|
+
}
|
|
334
|
+
if (!Number.isSafeInteger(count) || count! < 0 || count! > UINT32_MAX)
|
|
335
|
+
throw new Error(`line ${index + 1}: invalid count ${fields[2]}`);
|
|
336
|
+
const counts = this.#counts.get(layer!)!;
|
|
337
|
+
counts[expertId!] = incrementUint32(counts[expertId!]!, count!);
|
|
338
|
+
}
|
|
339
|
+
}
|
|
340
|
+
|
|
341
|
+
#clearCounts(): void {
|
|
342
|
+
for (const counts of this.#counts.values()) counts.fill(0);
|
|
343
|
+
}
|
|
344
|
+
|
|
345
|
+
#validateEntry(layer: number, expertId: number): void {
|
|
346
|
+
if (!this.#counts.has(layer))
|
|
347
|
+
throw new RangeError(`layer ${layer} is not tracked by expert usage`);
|
|
348
|
+
if (!Number.isSafeInteger(expertId) || expertId < 0 ||
|
|
349
|
+
expertId >= this.expertsPerLayer) {
|
|
350
|
+
throw new RangeError(`invalid expert ${expertId}`);
|
|
351
|
+
}
|
|
352
|
+
}
|
|
353
|
+
|
|
354
|
+
#lfruScore(layer: number, expertId: number): number {
|
|
355
|
+
const heat = this.#heat.get(layer)![expertId]!;
|
|
356
|
+
const last = this.#last.get(layer)![expertId]!;
|
|
357
|
+
return expertLfruScore(heat, last, this.#clock);
|
|
358
|
+
}
|
|
359
|
+
}
|
|
360
|
+
|
|
361
|
+
/**
|
|
362
|
+
* Faithful, budget-safe form of Colibri's startup AUTOPIN policy. History must
|
|
363
|
+
* reach 5k selections; confidence ramps to one at 200k; pins receive at most
|
|
364
|
+
* half of the resident-tier budget and are not created below a 0.5 GB tier.
|
|
365
|
+
*/
|
|
366
|
+
export function planExpertAutoPins(input: ExpertAutoPinInput): ExpertAutoPinPlan {
|
|
367
|
+
const minHistory = input.minHistorySelections ?? 5_000;
|
|
368
|
+
const fullConfidence = input.fullConfidenceSelections ?? 200_000;
|
|
369
|
+
const pinShare = input.pinShare ?? 0.5;
|
|
370
|
+
const minPinBudget = input.minPinBudgetBytes ?? 500_000_000;
|
|
371
|
+
positiveSafeInteger(minHistory, "auto-pin minHistorySelections");
|
|
372
|
+
positiveSafeInteger(fullConfidence, "auto-pin fullConfidenceSelections");
|
|
373
|
+
if (!Number.isFinite(pinShare) || pinShare <= 0 || pinShare > 1)
|
|
374
|
+
throw new RangeError("auto-pin pinShare must be in (0, 1]");
|
|
375
|
+
if (!Number.isSafeInteger(minPinBudget) || minPinBudget < 0)
|
|
376
|
+
throw new RangeError("auto-pin minPinBudgetBytes must be non-negative");
|
|
377
|
+
if (!Number.isSafeInteger(input.residentTierBudgetBytes) ||
|
|
378
|
+
input.residentTierBudgetBytes < 0 ||
|
|
379
|
+
!Number.isSafeInteger(input.mandatoryResidentBytes) ||
|
|
380
|
+
input.mandatoryResidentBytes < 0) {
|
|
381
|
+
throw new RangeError("auto-pin resident budgets must be non-negative safe integers");
|
|
382
|
+
}
|
|
383
|
+
|
|
384
|
+
const historySelections = input.ledger.snapshot().totalCount;
|
|
385
|
+
const confidence = Math.min(historySelections / fullConfidence, 1);
|
|
386
|
+
const requestedBudgetBytes = Math.floor(
|
|
387
|
+
input.residentTierBudgetBytes * pinShare * confidence,
|
|
388
|
+
);
|
|
389
|
+
const usableBudgetBytes = Math.min(
|
|
390
|
+
requestedBudgetBytes,
|
|
391
|
+
Math.max(
|
|
392
|
+
0,
|
|
393
|
+
input.residentTierBudgetBytes - input.mandatoryResidentBytes,
|
|
394
|
+
),
|
|
395
|
+
);
|
|
396
|
+
const pins: ExpertUsageCount[] = [];
|
|
397
|
+
let plannedBytes = 0;
|
|
398
|
+
if (historySelections >= minHistory && requestedBudgetBytes >= minPinBudget) {
|
|
399
|
+
for (const candidate of input.ledger.rankedCounts()) {
|
|
400
|
+
const bytes = input.slotBytes(candidate.layer);
|
|
401
|
+
positiveSafeInteger(bytes, `auto-pin slotBytes for layer ${candidate.layer}`);
|
|
402
|
+
if (plannedBytes + bytes > usableBudgetBytes) continue;
|
|
403
|
+
pins.push(candidate);
|
|
404
|
+
plannedBytes += bytes;
|
|
405
|
+
}
|
|
406
|
+
}
|
|
407
|
+
return {
|
|
408
|
+
historySelections,
|
|
409
|
+
confidence,
|
|
410
|
+
requestedBudgetBytes,
|
|
411
|
+
usableBudgetBytes,
|
|
412
|
+
plannedBytes,
|
|
413
|
+
pins: Object.freeze(pins),
|
|
414
|
+
};
|
|
415
|
+
}
|
package/src/generate.ts
CHANGED
|
@@ -117,6 +117,15 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
|
|
|
117
117
|
grammar?: GrammarController;
|
|
118
118
|
}
|
|
119
119
|
|
|
120
|
+
export function shouldUseGrammarJump(
|
|
121
|
+
options: Pick<GenerateOptions, "grammar" | "logprobs" | "topLogprobs">,
|
|
122
|
+
): boolean {
|
|
123
|
+
return options.grammar !== undefined &&
|
|
124
|
+
flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
|
|
125
|
+
!options.logprobs &&
|
|
126
|
+
!(options.topLogprobs && options.topLogprobs > 0);
|
|
127
|
+
}
|
|
128
|
+
|
|
120
129
|
/** Per-token logprob info (only present when GenerateOptions requested it). */
|
|
121
130
|
export interface TokenLogprobs {
|
|
122
131
|
/** logprob of the emitted token (requires options.logprobs). */
|
|
@@ -282,6 +291,8 @@ export interface GenerateStats {
|
|
|
282
291
|
spec?: {
|
|
283
292
|
drafted: number; accepted: number; targetCalls: number;
|
|
284
293
|
draftedByPos?: number[]; acceptedByPos?: number[];
|
|
294
|
+
rejected?: number; rounds?: number; acceptanceLengths?: number[];
|
|
295
|
+
tokensPerForward?: number; forwardsSaved?: number;
|
|
285
296
|
};
|
|
286
297
|
}
|
|
287
298
|
|
|
@@ -338,6 +349,33 @@ export class Generation implements AsyncIterable<GeneratedToken> {
|
|
|
338
349
|
const WIRE_THRESHOLD = 0.75;
|
|
339
350
|
let wiredScopeDepth = 0;
|
|
340
351
|
let wiredOldLimit = 0;
|
|
352
|
+
|
|
353
|
+
type WiredModelMemory = {
|
|
354
|
+
readonly weightsBytes: number;
|
|
355
|
+
readonly expertRuntime?: {
|
|
356
|
+
readonly plan: { readonly plannedBytes: number };
|
|
357
|
+
flushUsage?: () => void;
|
|
358
|
+
finishUsage?: () => Promise<void>;
|
|
359
|
+
} | null;
|
|
360
|
+
};
|
|
361
|
+
|
|
362
|
+
/** Bytes whose MLX graph must stay resident while a model executes. */
|
|
363
|
+
export function wiredWorkingSetBytes(model: WiredModelMemory): number {
|
|
364
|
+
const planned = model.expertRuntime?.plan.plannedBytes;
|
|
365
|
+
return typeof planned === "number" && Number.isFinite(planned) && planned > 0
|
|
366
|
+
? Math.max(model.weightsBytes, planned)
|
|
367
|
+
: model.weightsBytes;
|
|
368
|
+
}
|
|
369
|
+
|
|
370
|
+
export function modelNeedsWiredLimit(
|
|
371
|
+
model: WiredModelMemory,
|
|
372
|
+
recommendedBytes = maxRecommendedWorkingSetSize(),
|
|
373
|
+
force = process.env.MLX_BUN_FORCE_WIRE === "1",
|
|
374
|
+
): boolean {
|
|
375
|
+
return force ||
|
|
376
|
+
wiredWorkingSetBytes(model) > WIRE_THRESHOLD * recommendedBytes;
|
|
377
|
+
}
|
|
378
|
+
|
|
341
379
|
function enterWiredScope(): void {
|
|
342
380
|
if (wiredScopeDepth++ === 0)
|
|
343
381
|
wiredOldLimit = setWiredLimit(maxRecommendedWorkingSetSize());
|
|
@@ -362,9 +400,10 @@ export function generate(
|
|
|
362
400
|
? generateDiffusionInner(model, promptTokens, options)
|
|
363
401
|
: generateInner(model, promptTokens, options);
|
|
364
402
|
if (options.adapters?.length) inner = adapterScoped(model, options.adapters, inner);
|
|
365
|
-
const
|
|
366
|
-
|
|
367
|
-
|
|
403
|
+
const memoryModel: WiredModelMemory = model;
|
|
404
|
+
if (memoryModel.expertRuntime?.flushUsage)
|
|
405
|
+
inner = usageScoped(memoryModel, inner);
|
|
406
|
+
const wire = modelNeedsWiredLimit(model);
|
|
368
407
|
return new Generation(wire ? wiredScoped(inner) : inner);
|
|
369
408
|
}
|
|
370
409
|
|
|
@@ -437,6 +476,68 @@ async function* wiredScoped(
|
|
|
437
476
|
}
|
|
438
477
|
}
|
|
439
478
|
|
|
479
|
+
/** Publish the streamed model's route ledger on every generator exit path. */
|
|
480
|
+
async function* usageScoped(
|
|
481
|
+
model: WiredModelMemory,
|
|
482
|
+
inner: AsyncGenerator<GeneratedToken, GenerateStats>,
|
|
483
|
+
): AsyncGenerator<GeneratedToken, GenerateStats> {
|
|
484
|
+
try {
|
|
485
|
+
return yield* inner;
|
|
486
|
+
} finally {
|
|
487
|
+
if (model.expertRuntime?.finishUsage)
|
|
488
|
+
await model.expertRuntime.finishUsage();
|
|
489
|
+
else
|
|
490
|
+
model.expertRuntime?.flushUsage?.();
|
|
491
|
+
}
|
|
492
|
+
}
|
|
493
|
+
|
|
494
|
+
/** Apply generate()'s scoped wiring policy to non-generator execution paths. */
|
|
495
|
+
export async function withModelWiredLimit<T>(
|
|
496
|
+
model: WiredModelMemory,
|
|
497
|
+
run: () => Promise<T>,
|
|
498
|
+
): Promise<T> {
|
|
499
|
+
if (!modelNeedsWiredLimit(model)) return run();
|
|
500
|
+
enterWiredScope();
|
|
501
|
+
try {
|
|
502
|
+
return await run();
|
|
503
|
+
} finally {
|
|
504
|
+
exitWiredScope();
|
|
505
|
+
}
|
|
506
|
+
}
|
|
507
|
+
|
|
508
|
+
/** Apply generate()'s usage safe-point to direct/non-generator paths. */
|
|
509
|
+
export async function withModelUsageFlush<T>(
|
|
510
|
+
model: WiredModelMemory,
|
|
511
|
+
run: () => Promise<T>,
|
|
512
|
+
): Promise<T> {
|
|
513
|
+
try {
|
|
514
|
+
return await run();
|
|
515
|
+
} finally {
|
|
516
|
+
if (model.expertRuntime?.finishUsage)
|
|
517
|
+
await model.expertRuntime.finishUsage();
|
|
518
|
+
else
|
|
519
|
+
model.expertRuntime?.flushUsage?.();
|
|
520
|
+
}
|
|
521
|
+
}
|
|
522
|
+
|
|
523
|
+
/** Forward through either the ordinary synchronous model surface or a
|
|
524
|
+
* streamed-expert model whose layer execution must cross async I/O. Keeping
|
|
525
|
+
* this choice inside the shared generator preserves the complete generation
|
|
526
|
+
* contract (sampling, grammar, logprobs, prompt cache, and cancellation) for
|
|
527
|
+
* both execution styles. */
|
|
528
|
+
async function forwardHiddenForGeneration(
|
|
529
|
+
model: RuntimeModel,
|
|
530
|
+
ids: MlxArray,
|
|
531
|
+
cache: Cache[],
|
|
532
|
+
): Promise<MlxArray> {
|
|
533
|
+
const asyncModel = model as RuntimeModel & {
|
|
534
|
+
forwardHiddenAsync?: (ids: MlxArray, cache: Cache[]) => Promise<MlxArray>;
|
|
535
|
+
};
|
|
536
|
+
return typeof asyncModel.forwardHiddenAsync === "function"
|
|
537
|
+
? await asyncModel.forwardHiddenAsync(ids, cache)
|
|
538
|
+
: model.forwardHidden(ids, cache);
|
|
539
|
+
}
|
|
540
|
+
|
|
440
541
|
async function* generateInner(
|
|
441
542
|
model: RuntimeModel,
|
|
442
543
|
promptTokens: number[],
|
|
@@ -488,15 +589,14 @@ async function* generateInner(
|
|
|
488
589
|
if (e) for (const a of [e.sel, e.topIdx, e.topVals]) a?.dispose();
|
|
489
590
|
};
|
|
490
591
|
/** Read extras back to JS (forces eval — they were async-dispatched with the
|
|
491
|
-
* token
|
|
592
|
+
* token) without appending casts behind the next decode step. */
|
|
492
593
|
const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
|
|
493
594
|
if (!e) return undefined;
|
|
494
595
|
const out: TokenLogprobs = {};
|
|
495
|
-
if (e.sel) out.logprob = e.sel.
|
|
596
|
+
if (e.sel) out.logprob = e.sel.toFloat32Host()[0]!;
|
|
496
597
|
if (e.topIdx && e.topVals) {
|
|
497
|
-
|
|
498
|
-
const
|
|
499
|
-
const vals = e.topVals.toFloat32();
|
|
598
|
+
const ids = e.topIdx.toIntTokens();
|
|
599
|
+
const vals = e.topVals.toFloat32Host();
|
|
500
600
|
out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! })).sort(
|
|
501
601
|
(a, b) => b.logprob - a.logprob,
|
|
502
602
|
);
|
|
@@ -616,7 +716,7 @@ async function* generateInner(
|
|
|
616
716
|
while (pos < snapAt) {
|
|
617
717
|
const chunk = promptTokens.slice(pos, Math.min(pos + prefillChunkSize, snapAt));
|
|
618
718
|
const ids = ops.fromInt32(chunk, [1, chunk.length]);
|
|
619
|
-
const h = model
|
|
719
|
+
const h = await forwardHiddenForGeneration(model, ids, cache);
|
|
620
720
|
ids.dispose();
|
|
621
721
|
h.dispose();
|
|
622
722
|
evalCacheState(cache);
|
|
@@ -650,7 +750,7 @@ async function* generateInner(
|
|
|
650
750
|
: prefillChunkSize;
|
|
651
751
|
const chunk = promptTokens.slice(pos, pos + n);
|
|
652
752
|
const ids = ops.fromInt32(chunk, [1, chunk.length]);
|
|
653
|
-
const h = model
|
|
753
|
+
const h = await forwardHiddenForGeneration(model, ids, cache);
|
|
654
754
|
ids.dispose();
|
|
655
755
|
h.dispose(); // logits never computed for non-final chunks
|
|
656
756
|
evalCacheState(cache);
|
|
@@ -677,7 +777,7 @@ async function* generateInner(
|
|
|
677
777
|
// cover exactly the prompt, as before.
|
|
678
778
|
const lastChunk = promptTokens.slice(pos);
|
|
679
779
|
const ids0 = ops.fromInt32(lastChunk, [1, lastChunk.length]);
|
|
680
|
-
h0 = model
|
|
780
|
+
h0 = await forwardHiddenForGeneration(model, ids0, cache);
|
|
681
781
|
ids0.dispose();
|
|
682
782
|
}
|
|
683
783
|
if (options.snapshotAt === undefined || options.snapshotAt >= promptTokens.length)
|
|
@@ -736,10 +836,7 @@ async function* generateInner(
|
|
|
736
836
|
// (see GrammarController.jumpForward for the contract + the fidelity
|
|
737
837
|
// note on why this is opt-in). Excluded when logprobs are requested
|
|
738
838
|
// (jumped tokens are never sampled, so they'd have no logprobs rows).
|
|
739
|
-
const grammarJump =
|
|
740
|
-
options.grammar !== undefined &&
|
|
741
|
-
flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
|
|
742
|
-
!options.logprobs;
|
|
839
|
+
const grammarJump = shouldUseGrammarJump(options);
|
|
743
840
|
while (!stop) {
|
|
744
841
|
const cur = pending!;
|
|
745
842
|
const curExtras = pendingExtras;
|
|
@@ -798,7 +895,7 @@ async function* generateInner(
|
|
|
798
895
|
}
|
|
799
896
|
const chunk = [grammarTok, ...jumpEmit];
|
|
800
897
|
const ids = ops.fromInt32(chunk, [1, chunk.length]);
|
|
801
|
-
const h = model
|
|
898
|
+
const h = await forwardHiddenForGeneration(model, ids, cache);
|
|
802
899
|
ids.dispose();
|
|
803
900
|
// Every chunk token's KV is in the cache regardless of what follows.
|
|
804
901
|
forwarded.push(...chunk);
|
|
@@ -839,7 +936,7 @@ async function* generateInner(
|
|
|
839
936
|
}
|
|
840
937
|
if (!logits) {
|
|
841
938
|
const ids = ops.reshape(cur, [1, 1]);
|
|
842
|
-
const h = model
|
|
939
|
+
const h = await forwardHiddenForGeneration(model, ids, cache);
|
|
843
940
|
ids.dispose();
|
|
844
941
|
logits = model.logitsFromHidden(h);
|
|
845
942
|
h.dispose();
|
package/src/hf-push.ts
CHANGED
|
@@ -420,7 +420,7 @@ async function uploadLfsFiles(
|
|
|
420
420
|
const putRes = await fetch(upload.href, {
|
|
421
421
|
method: "PUT",
|
|
422
422
|
headers: { ...(upload.header ?? {}) },
|
|
423
|
-
body: item.bytes,
|
|
423
|
+
body: new Blob([Uint8Array.from(item.bytes)]),
|
|
424
424
|
});
|
|
425
425
|
if (!putRes.ok) throw await httpError(`LFS upload of ${item.file.repoPath}`, putRes);
|
|
426
426
|
onProgress?.(item.file.repoPath, item.file.size, item.file.size);
|
package/src/index.ts
CHANGED
|
@@ -12,7 +12,13 @@ export { loadModelConfig } from "./config";
|
|
|
12
12
|
export { Weights } from "./weights";
|
|
13
13
|
export { Gemma4Model } from "./model/gemma4";
|
|
14
14
|
export { MiniCPM5Model } from "./model/minicpm5";
|
|
15
|
-
export {
|
|
15
|
+
export {
|
|
16
|
+
createModel,
|
|
17
|
+
openGlm52RuntimeModel,
|
|
18
|
+
openModel,
|
|
19
|
+
type Glm52RuntimeOpenOptions,
|
|
20
|
+
type RuntimeModel,
|
|
21
|
+
} from "./model/factory";
|
|
16
22
|
export { Qwen3Model } from "./model/qwen3";
|
|
17
23
|
export { generate } from "./generate";
|
|
18
24
|
export { loadTokenizer } from "./tokenizer";
|
|
@@ -21,4 +27,10 @@ export { ChatTemplate } from "./chat-template";
|
|
|
21
27
|
export { chooseAutoModel, COEXIST_FRACTION, DEFAULT_REPO_ID, fit, largestRecommendedRepoId, recommendedRepoId, skuMatrix, thisMachine } from "./fit";
|
|
22
28
|
export { downloadModel } from "./download";
|
|
23
29
|
export { Registry } from "./registry";
|
|
24
|
-
export {
|
|
30
|
+
export {
|
|
31
|
+
createServer,
|
|
32
|
+
loadContext,
|
|
33
|
+
type LoadContextOptions,
|
|
34
|
+
type ServerContext,
|
|
35
|
+
type ServerOptions,
|
|
36
|
+
} from "./server";
|