@mlx-node/lm 0.0.7 → 0.0.8

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.
@@ -5,44 +5,215 @@
5
5
  */
6
6
  import { readFile } from 'node:fs/promises';
7
7
  import { join } from 'node:path';
8
- import { HarrierModel, QianfanOCRModel } from '@mlx-node/core';
8
+ import { Gemma4Model as NativeGemma4Model, HarrierModel, Lfm2Model as NativeLfm2Model, QianfanOCRModel, Qwen3Model as NativeQwen3Model, Qwen35Model as NativeQwen35Model, Qwen35MoeModel as NativeQwen35MoeModel, } from '@mlx-node/core';
9
9
  import { ChatSession } from '../chat-session.js';
10
10
  import { Gemma4Model, Lfm2Model, Qwen3Model, Qwen35Model, Qwen35MoeModel } from '../stream.js';
11
- const SUPPORTED_MODEL_TYPES = new Set([
12
- 'qwen3',
13
- 'qwen3_5',
14
- 'qwen3_5_moe',
15
- 'internvl_chat',
16
- 'qianfan-ocr',
17
- 'harrier',
18
- 'gemma4',
19
- 'lfm2',
20
- ]);
11
+ /**
12
+ * Ordered source of truth for every supported model family. Each entry owns
13
+ * its canonical `ModelType`, raw config aliases / architecture probes, loader,
14
+ * and `ChatSession` eligibility:
15
+ *
16
+ * - `'trainable'` — GRPO/SFT-capable LM (Qwen3 family); chat-capable.
17
+ * - `'loadable'` — chat-capable LM with no trainer engine (Gemma4, LFM2).
18
+ * - `'embedding'` — no chat surface (Harrier); rejected by `loadSession`.
19
+ * - `'vlm'` — VLM whose AsyncGenerator wrapper lives in
20
+ * `@mlx-node/vlm` (importing it here would create a
21
+ * circular package dependency), so `loadSession`
22
+ * rejects it and routes callers to `@mlx-node/vlm`.
23
+ *
24
+ * A base family is selected from an explicit alias or the single declarative
25
+ * nullish-model_type default, then architecture probes refine it in declaration
26
+ * order. Gemma's unified architecture is authoritative (matching the native
27
+ * loader); Harrier refines a Qwen3 base. Adding a family means adding one
28
+ * descriptor here, without a second normalization or dispatch branch.
29
+ */
30
+ const MODEL_FAMILY_REGISTRY = [
31
+ {
32
+ modelType: 'gemma4',
33
+ kind: 'loadable',
34
+ match: {
35
+ rawModelTypes: ['gemma4', 'gemma4_text', 'gemma4_unified'],
36
+ architectureProbe: ({ architectures }) => architectures.has('Gemma4UnifiedForConditionalGeneration'),
37
+ },
38
+ load: (modelPath, options) => Gemma4Model.load(modelPath, options?.draftModelPath === undefined ? null : { draftModelPath: options.draftModelPath }),
39
+ nativeModelClass: NativeGemma4Model,
40
+ acceptsDraftModel: true,
41
+ },
42
+ {
43
+ modelType: 'harrier',
44
+ kind: 'embedding',
45
+ match: {
46
+ rawModelTypes: ['harrier'],
47
+ architectureProbe: ({ modelType, architectures }) => modelType === 'qwen3' && architectures.has('Qwen3Model') && !architectures.has('Qwen3ForCausalLM'),
48
+ },
49
+ load: (modelPath) => HarrierModel.load(modelPath),
50
+ nativeModelClass: HarrierModel,
51
+ },
52
+ {
53
+ modelType: 'qwen3',
54
+ kind: 'trainable',
55
+ match: { rawModelTypes: ['qwen3'] },
56
+ load: (modelPath) => Qwen3Model.load(modelPath),
57
+ nativeModelClass: NativeQwen3Model,
58
+ defaultForNullishModelType: true,
59
+ },
60
+ {
61
+ modelType: 'qwen3_5',
62
+ kind: 'trainable',
63
+ match: { rawModelTypes: ['qwen3_5'] },
64
+ load: (modelPath) => Qwen35Model.load(modelPath),
65
+ nativeModelClass: NativeQwen35Model,
66
+ },
67
+ {
68
+ modelType: 'qwen3_5_moe',
69
+ kind: 'trainable',
70
+ match: { rawModelTypes: ['qwen3_5_moe'] },
71
+ load: (modelPath) => Qwen35MoeModel.load(modelPath),
72
+ nativeModelClass: NativeQwen35MoeModel,
73
+ },
74
+ {
75
+ modelType: 'lfm2',
76
+ kind: 'loadable',
77
+ match: { rawModelTypes: ['lfm2'] },
78
+ load: (modelPath) => Lfm2Model.load(modelPath),
79
+ nativeModelClass: NativeLfm2Model,
80
+ },
81
+ {
82
+ modelType: 'lfm2_moe',
83
+ kind: 'loadable',
84
+ match: { rawModelTypes: ['lfm2_moe'] },
85
+ load: (modelPath) => Lfm2Model.load(modelPath),
86
+ nativeModelClass: NativeLfm2Model,
87
+ },
88
+ {
89
+ modelType: 'internvl_chat',
90
+ kind: 'vlm',
91
+ match: { rawModelTypes: ['internvl_chat'] },
92
+ load: (modelPath) => QianfanOCRModel.load(modelPath),
93
+ nativeModelClass: QianfanOCRModel,
94
+ },
95
+ {
96
+ modelType: 'qianfan-ocr',
97
+ kind: 'vlm',
98
+ match: { rawModelTypes: ['qianfan-ocr'] },
99
+ load: (modelPath) => QianfanOCRModel.load(modelPath),
100
+ nativeModelClass: QianfanOCRModel,
101
+ },
102
+ ];
103
+ function buildModelFamilyIndex(registry) {
104
+ const byModelType = new Map();
105
+ const byRawModelType = new Map();
106
+ let defaultForNullishModelType;
107
+ for (const family of registry) {
108
+ const previousFamily = byModelType.get(family.modelType);
109
+ if (previousFamily !== undefined) {
110
+ throw new Error(`Duplicate canonical model type "${family.modelType}" in model family registry`);
111
+ }
112
+ byModelType.set(family.modelType, family);
113
+ for (const rawModelType of family.match.rawModelTypes) {
114
+ const previous = byRawModelType.get(rawModelType);
115
+ if (previous !== undefined) {
116
+ throw new Error(`Duplicate model_type alias "${rawModelType}" for "${previous.modelType}" and "${family.modelType}"`);
117
+ }
118
+ byRawModelType.set(rawModelType, family);
119
+ }
120
+ if (family.defaultForNullishModelType === true) {
121
+ if (defaultForNullishModelType !== undefined) {
122
+ throw new Error(`Duplicate nullish-model_type defaults for "${defaultForNullishModelType.modelType}" and "${family.modelType}"`);
123
+ }
124
+ defaultForNullishModelType = family;
125
+ }
126
+ }
127
+ if (defaultForNullishModelType === undefined) {
128
+ throw new Error('Model family registry must declare exactly one nullish-model_type default');
129
+ }
130
+ return { byModelType, byRawModelType, defaultForNullishModelType };
131
+ }
132
+ const MODEL_FAMILY_INDEX = buildModelFamilyIndex(MODEL_FAMILY_REGISTRY);
133
+ function findFamily(modelType) {
134
+ const family = MODEL_FAMILY_INDEX.byModelType.get(modelType);
135
+ if (family === undefined) {
136
+ throw new Error(`Internal error: missing model family descriptor for "${modelType}"`);
137
+ }
138
+ return family;
139
+ }
140
+ function matchesArchitectureProbe(family, config) {
141
+ return family.match.architectureProbe?.(config) === true;
142
+ }
143
+ class MalformedModelConfigError extends Error {
144
+ constructor(modelPath, reason) {
145
+ super(`Malformed config.json in ${modelPath}: ${reason}`);
146
+ this.name = 'MalformedModelConfigError';
147
+ }
148
+ }
149
+ /**
150
+ * Fail-closed validation: a config.json whose root is not a plain object,
151
+ * or whose `architectures` is neither an array nor a string, is rejected
152
+ * instead of coerced (coercion would fall through to the qwen3
153
+ * nullish-model_type default and silently misroute the checkpoint).
154
+ * Blessed lenient shapes stay accepted: `{}` root (qwen3 default),
155
+ * missing/`null` `architectures` (empty set), bare-string `architectures`
156
+ * (single-element set), and non-string array entries (filtered out).
157
+ */
158
+ function normalizeConfig(modelPath, config) {
159
+ if (typeof config !== 'object' || config === null || Array.isArray(config)) {
160
+ throw new MalformedModelConfigError(modelPath, 'root must be a JSON object');
161
+ }
162
+ const object = config;
163
+ const hasModelType = Object.hasOwn(object, 'model_type');
164
+ const rawModelTypeValue = hasModelType ? object.model_type : undefined;
165
+ const usesDefaultModelType = !hasModelType || rawModelTypeValue === null;
166
+ const rawModelType = typeof rawModelTypeValue === 'string' ? rawModelTypeValue : undefined;
167
+ const rawModelTypeLabel = hasModelType ? String(rawModelTypeValue) : '<missing>';
168
+ const rawArchitectures = 'architectures' in object ? object.architectures : undefined;
169
+ if (rawArchitectures !== undefined &&
170
+ rawArchitectures !== null &&
171
+ !Array.isArray(rawArchitectures) &&
172
+ typeof rawArchitectures !== 'string') {
173
+ throw new MalformedModelConfigError(modelPath, '"architectures" must be an array or a string');
174
+ }
175
+ const architectures = Array.isArray(rawArchitectures)
176
+ ? rawArchitectures.filter((architecture) => typeof architecture === 'string')
177
+ : typeof rawArchitectures === 'string'
178
+ ? [rawArchitectures]
179
+ : [];
180
+ return { usesDefaultModelType, rawModelType, rawModelTypeLabel, architectures: new Set(architectures) };
181
+ }
182
+ class UnsupportedModelTypeError extends Error {
183
+ constructor(modelPath, rawModelTypeLabel) {
184
+ super(`Unsupported model_type "${rawModelTypeLabel}" in ${modelPath}/config.json`);
185
+ this.name = 'UnsupportedModelTypeError';
186
+ }
187
+ }
188
+ /**
189
+ * Dispatch a load through the registry, validating gemma4-only options.
190
+ * `draftModelPath` reaches ONLY the gemma4 row; every other family rejects
191
+ * it loudly instead of silently ignoring a caller's speculative-decode
192
+ * intent.
193
+ */
194
+ function dispatchLoad(modelType, modelPath, options) {
195
+ const family = findFamily(modelType);
196
+ if (options?.draftModelPath !== undefined && family.acceptsDraftModel !== true) {
197
+ throw new Error(`draftModelPath (speculative-decoding draft) is only supported by gemma4 models; ` +
198
+ `${modelPath} has model_type "${modelType}"`);
199
+ }
200
+ return family.load(modelPath, options);
201
+ }
21
202
  /**
22
203
  * Load a model from disk, auto-detecting architecture from config.json.
23
204
  *
24
205
  * Supports both language models (Qwen3, Qwen3.5) and vision-language models
25
206
  * (Qianfan-OCR / InternVL). Use `instanceof` to narrow the returned type.
207
+ *
208
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
209
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
210
+ * speculative decoding — gemma4 only; any other detected family rejects it.
211
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
212
+ * that embedded checkpoint is present.
26
213
  */
27
- export async function loadModel(modelPath) {
214
+ export async function loadModel(modelPath, options) {
28
215
  const modelType = await detectModelType(modelPath);
29
- switch (modelType) {
30
- case 'qwen3_5_moe':
31
- return Qwen35MoeModel.load(modelPath);
32
- case 'qwen3_5':
33
- return Qwen35Model.load(modelPath);
34
- case 'qwen3':
35
- return Qwen3Model.load(modelPath);
36
- case 'harrier':
37
- return HarrierModel.load(modelPath);
38
- case 'internvl_chat':
39
- case 'qianfan-ocr':
40
- return QianfanOCRModel.load(modelPath);
41
- case 'gemma4':
42
- return Gemma4Model.load(modelPath);
43
- case 'lfm2':
44
- return Lfm2Model.load(modelPath);
45
- }
216
+ return dispatchLoad(modelType, modelPath, options);
46
217
  }
47
218
  /**
48
219
  * Load a model and wrap it in a {@link ChatSession} for multi-turn chat.
@@ -59,38 +230,44 @@ export async function loadModel(modelPath) {
59
230
  * package dependency), so callers who want a Qianfan-OCR session
60
231
  * must import `QianfanOCRModel` from `@mlx-node/vlm` and construct
61
232
  * `new ChatSession(model)` directly.
233
+ *
234
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
235
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
236
+ * speculative decoding — gemma4 only; any other detected family rejects it.
237
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
238
+ * that embedded checkpoint is present.
239
+ * The resulting session auto-enables the speculative path (the model
240
+ * reports `hasMtpWeights()`); pass `enableMtp: false` per call to opt out.
62
241
  */
63
- export async function loadSession(modelPath) {
64
- const m = await loadModel(modelPath);
65
- if (m instanceof HarrierModel) {
242
+ export async function loadSession(modelPath, options) {
243
+ const modelType = await detectModelType(modelPath);
244
+ const kind = findFamily(modelType).kind;
245
+ if (kind === 'embedding') {
66
246
  throw new Error('loadSession: embedding models (Harrier) cannot be wrapped in a ChatSession');
67
247
  }
68
- if (m instanceof QianfanOCRModel) {
248
+ if (kind === 'vlm') {
69
249
  throw new Error('loadSession: Qianfan-OCR / InternVL session support lives in @mlx-node/vlm. Import QianfanOCRModel from @mlx-node/vlm and construct ChatSession(model) directly.');
70
250
  }
251
+ const m = await dispatchLoad(modelType, modelPath, options);
71
252
  return new ChatSession(m);
72
253
  }
73
254
  export async function detectModelType(modelPath) {
74
255
  try {
75
256
  const raw = await readFile(join(modelPath, 'config.json'), 'utf-8');
76
- const config = JSON.parse(raw);
77
- const rawModelType = config.model_type ?? 'qwen3';
78
- // Normalize model_type: gemma4_text → gemma4
79
- let modelType = (rawModelType === 'gemma4_text' ? 'gemma4' : rawModelType);
80
- // Detect embedding models: Qwen3 backbone with base architecture (no ForCausalLM)
81
- if (modelType === 'qwen3') {
82
- const architectures = config.architectures ?? [];
83
- if (architectures.includes('Qwen3Model') && !architectures.includes('Qwen3ForCausalLM')) {
84
- modelType = 'harrier';
85
- }
86
- }
87
- if (!SUPPORTED_MODEL_TYPES.has(modelType)) {
88
- throw new Error(`Unsupported model_type "${modelType}" in ${modelPath}/config.json`);
89
- }
90
- return modelType;
257
+ const config = normalizeConfig(modelPath, JSON.parse(raw));
258
+ const baseFamily = config.usesDefaultModelType
259
+ ? MODEL_FAMILY_INDEX.defaultForNullishModelType
260
+ : config.rawModelType === undefined
261
+ ? undefined
262
+ : MODEL_FAMILY_INDEX.byRawModelType.get(config.rawModelType);
263
+ const matchContext = { ...config, modelType: baseFamily?.modelType };
264
+ const family = MODEL_FAMILY_REGISTRY.find((candidate) => matchesArchitectureProbe(candidate, matchContext)) ?? baseFamily;
265
+ if (family === undefined)
266
+ throw new UnsupportedModelTypeError(modelPath, config.rawModelTypeLabel);
267
+ return family.modelType;
91
268
  }
92
269
  catch (e) {
93
- if (e instanceof Error && e.message.startsWith('Unsupported model_type'))
270
+ if (e instanceof UnsupportedModelTypeError || e instanceof MalformedModelConfigError)
94
271
  throw e;
95
272
  throw new Error(`Cannot detect model type: config.json not found in ${modelPath}`);
96
273
  }
@@ -0,0 +1,60 @@
1
+ /**
2
+ * Process-local model-directory overrides for forcing block-paged KV caches.
3
+ *
4
+ * Native model loaders accept a directory path and read paging policy from
5
+ * `config.json`. This manager creates an isolated temporary clone containing a
6
+ * patched config plus symlinks to the source files, so callers can opt models
7
+ * into paging without mutating downloaded checkpoints. Loader-known Qwen3.5
8
+ * MTP sidecars are preserved explicitly; unrelated directories (including a
9
+ * Gemma4 `draft/`) remain hidden unless a caller explicitly preserves the
10
+ * Gemma draft for flat speculative decoding.
11
+ */
12
+ /** Every chat-capable family currently discovered by `@mlx-node/agent`. */
13
+ export declare const AGENT_PAGED_MODEL_TYPES: readonly ['qwen3', 'qwen3_5', 'qwen3_5_moe', 'gemma4', 'lfm2', 'lfm2_moe'];
14
+ /** Families historically forced paged by `mlx launch claude`. */
15
+ export declare const QWEN35_PAGED_MODEL_TYPES: readonly ['qwen3_5', 'qwen3_5_moe'];
16
+ export interface PagedConfigOverrideManagerOptions {
17
+ /** Model types to force onto the paged path. Defaults to all agent chat families. */
18
+ modelTypes?: readonly string[];
19
+ /** Temporary-directory prefix. Primarily useful for diagnostics/tests. */
20
+ tempDirPrefix?: string;
21
+ /**
22
+ * Pass Gemma4 checkpoints with an embedded `draft/` through unchanged.
23
+ * This keeps native draft auto-discovery, which currently requires flat KV
24
+ * caches. Defaults to false so a paged override remains a paged contract.
25
+ */
26
+ preserveEmbeddedGemmaDraft?: boolean;
27
+ }
28
+ /**
29
+ * Owns one isolated set of temporary paged-config overrides.
30
+ *
31
+ * A manager is intentionally single-lifecycle: repeated resolution of one
32
+ * source returns the same override, and `cleanup()` permanently disposes the
33
+ * manager. Separate managers never share directories, so one launch cannot
34
+ * remove another launch's live override.
35
+ */
36
+ export declare class PagedConfigOverrideManager {
37
+ private readonly modelTypes;
38
+ private readonly tempDirPrefix;
39
+ private readonly preserveEmbeddedGemmaDraft;
40
+ private readonly overrides;
41
+ private readonly activeResolves;
42
+ private rootPromise;
43
+ private cleanupPromise;
44
+ private disposed;
45
+ constructor(options?: PagedConfigOverrideManagerOptions);
46
+ /**
47
+ * Resolve `modelPath` to a paged-aware clone when its model type is managed.
48
+ * A caller-supplied canonical family takes precedence over the raw config
49
+ * type (for example, `gemma4` for a `gemma4_unified` checkpoint).
50
+ * Unmanaged, unreadable, or malformed checkpoints pass through unchanged.
51
+ */
52
+ resolve(modelPath: string, canonicalModelType?: string): Promise<string>;
53
+ private resolveInternal;
54
+ /** Remove this manager's temporary root without affecting other managers. */
55
+ cleanup(): Promise<void>;
56
+ private performCleanup;
57
+ private createOverride;
58
+ private getRoot;
59
+ }
60
+ //# sourceMappingURL=paged-config-override.d.ts.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"paged-config-override.d.ts","sourceRoot":"","sources":["../../src/models/paged-config-override.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;GAUG;AAMH,2EAA2E;AAC3E,eAAO,MAAM,uBAAuB,YAAI,OAAO,EAAE,SAAS,EAAE,aAAa,EAAE,QAAQ,EAAE,MAAM,EAAE,UAAU,CAAU,CAAC;AAElH,iEAAiE;AACjE,eAAO,MAAM,wBAAwB,YAAI,SAAS,EAAE,aAAa,CAAU,CAAC;AAO5E,MAAM,WAAW,iCAAiC;IAChD,qFAAqF;IACrF,UAAU,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;IAC/B,0EAA0E;IAC1E,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB;;;;OAIG;IACH,0BAA0B,CAAC,EAAE,OAAO,CAAC;CACtC;AAED;;;;;;;GAOG;AACH,qBAAa,0BAA0B;IACrC,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAsB;IACjD,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAS;IACvC,OAAO,CAAC,QAAQ,CAAC,0BAA0B,CAAU;IACrD,OAAO,CAAC,QAAQ,CAAC,SAAS,CAAsC;IAChE,OAAO,CAAC,QAAQ,CAAC,cAAc,CAA8B;IAC7D,OAAO,CAAC,WAAW,CAA8B;IACjD,OAAO,CAAC,cAAc,CAA4B;IAClD,OAAO,CAAC,QAAQ,CAAS;IAEzB,YAAY,OAAO,GAAE,iCAAsC,EAI1D;IAED;;;;;OAKG;IACG,OAAO,CAAC,SAAS,EAAE,MAAM,EAAE,kBAAkB,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC,CAY7E;YAEa,eAAe;IAgD7B,6EAA6E;IAC7E,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC,CAKvB;YAEa,cAAc;YAad,cAAc;IA8C5B,OAAO,CAAC,OAAO;CAIhB"}
@@ -0,0 +1,254 @@
1
+ /**
2
+ * Process-local model-directory overrides for forcing block-paged KV caches.
3
+ *
4
+ * Native model loaders accept a directory path and read paging policy from
5
+ * `config.json`. This manager creates an isolated temporary clone containing a
6
+ * patched config plus symlinks to the source files, so callers can opt models
7
+ * into paging without mutating downloaded checkpoints. Loader-known Qwen3.5
8
+ * MTP sidecars are preserved explicitly; unrelated directories (including a
9
+ * Gemma4 `draft/`) remain hidden unless a caller explicitly preserves the
10
+ * Gemma draft for flat speculative decoding.
11
+ */
12
+ import { mkdir, mkdtemp, readFile, readdir, rm, stat, symlink, writeFile } from 'node:fs/promises';
13
+ import { tmpdir } from 'node:os';
14
+ import { dirname, isAbsolute, join, relative, resolve, sep } from 'node:path';
15
+ /** Every chat-capable family currently discovered by `@mlx-node/agent`. */
16
+ export const AGENT_PAGED_MODEL_TYPES = ['qwen3', 'qwen3_5', 'qwen3_5_moe', 'gemma4', 'lfm2', 'lfm2_moe'];
17
+ /** Families historically forced paged by `mlx launch claude`. */
18
+ export const QWEN35_PAGED_MODEL_TYPES = ['qwen3_5', 'qwen3_5_moe'];
19
+ const QWEN35_CACHE_FLOOR_MODEL_TYPES = new Set(QWEN35_PAGED_MODEL_TYPES);
20
+ const DEFAULT_QWEN35_PAGED_CACHE_MB = 16_384;
21
+ const QWEN35_MTP_DRAFTER_DIR = 'mtp-drafter';
22
+ const QWEN35_DENSE_NESTED_MTP_SIDECAR = 'mtp/weights.safetensors';
23
+ /**
24
+ * Owns one isolated set of temporary paged-config overrides.
25
+ *
26
+ * A manager is intentionally single-lifecycle: repeated resolution of one
27
+ * source returns the same override, and `cleanup()` permanently disposes the
28
+ * manager. Separate managers never share directories, so one launch cannot
29
+ * remove another launch's live override.
30
+ */
31
+ export class PagedConfigOverrideManager {
32
+ modelTypes;
33
+ tempDirPrefix;
34
+ preserveEmbeddedGemmaDraft;
35
+ overrides = new Map();
36
+ activeResolves = new Set();
37
+ rootPromise;
38
+ cleanupPromise;
39
+ disposed = false;
40
+ constructor(options = {}) {
41
+ this.modelTypes = new Set(options.modelTypes ?? AGENT_PAGED_MODEL_TYPES);
42
+ this.tempDirPrefix = options.tempDirPrefix ?? 'mlx-paged-overrides-';
43
+ this.preserveEmbeddedGemmaDraft = options.preserveEmbeddedGemmaDraft ?? false;
44
+ }
45
+ /**
46
+ * Resolve `modelPath` to a paged-aware clone when its model type is managed.
47
+ * A caller-supplied canonical family takes precedence over the raw config
48
+ * type (for example, `gemma4` for a `gemma4_unified` checkpoint).
49
+ * Unmanaged, unreadable, or malformed checkpoints pass through unchanged.
50
+ */
51
+ async resolve(modelPath, canonicalModelType) {
52
+ if (this.disposed) {
53
+ throw new Error('PagedConfigOverrideManager: resolve() called after cleanup()');
54
+ }
55
+ const operation = this.resolveInternal(modelPath, canonicalModelType);
56
+ this.activeResolves.add(operation);
57
+ try {
58
+ return await operation;
59
+ }
60
+ finally {
61
+ this.activeResolves.delete(operation);
62
+ }
63
+ }
64
+ async resolveInternal(modelPath, canonicalModelType) {
65
+ const sourcePath = isAbsolute(modelPath) ? modelPath : resolve(modelPath);
66
+ let config;
67
+ try {
68
+ config = JSON.parse(await readFile(join(sourcePath, 'config.json'), 'utf-8'));
69
+ }
70
+ catch {
71
+ return modelPath;
72
+ }
73
+ const rawModelType = typeof config.model_type === 'string' ? config.model_type : null;
74
+ const modelType = canonicalModelType ?? rawModelType;
75
+ if (modelType === null || !this.modelTypes.has(modelType)) {
76
+ return modelPath;
77
+ }
78
+ // Gemma4's DSpark / assistant speculative executor currently owns flat KV
79
+ // caches. Preserve native draft discovery only for explicit callers; the
80
+ // default paged clone intentionally omits subdirectories, hiding `draft/`
81
+ // while keeping the downloaded checkpoint unchanged.
82
+ const hasEmbeddedGemmaDraft = modelType === 'gemma4' && (await isDirectory(join(sourcePath, 'draft')));
83
+ if (this.preserveEmbeddedGemmaDraft && hasEmbeddedGemmaDraft) {
84
+ return modelPath;
85
+ }
86
+ const cacheFloorMb = QWEN35_CACHE_FLOOR_MODEL_TYPES.has(modelType) ? resolveQwen35CacheFloorMb() : undefined;
87
+ const pagedEnabled = config.use_block_paged_cache === true;
88
+ const configuredMemoryMb = positiveNumber(config.paged_cache_memory_mb);
89
+ const memorySatisfied = cacheFloorMb === undefined || (configuredMemoryMb ?? 0) >= cacheFloorMb;
90
+ // Even an already-paged Gemma config must be cloned when `draft/` exists:
91
+ // returning the source would expose the draft to native auto-discovery and
92
+ // trigger the flat-speculation/paged-cache conflict.
93
+ if (pagedEnabled && memorySatisfied && !hasEmbeddedGemmaDraft) {
94
+ return modelPath;
95
+ }
96
+ const existing = this.overrides.get(sourcePath);
97
+ if (existing !== undefined)
98
+ return existing;
99
+ const pending = this.createOverride(sourcePath, config, cacheFloorMb, modelType);
100
+ this.overrides.set(sourcePath, pending);
101
+ try {
102
+ return await pending;
103
+ }
104
+ catch (error) {
105
+ this.overrides.delete(sourcePath);
106
+ throw error;
107
+ }
108
+ }
109
+ /** Remove this manager's temporary root without affecting other managers. */
110
+ cleanup() {
111
+ if (this.cleanupPromise !== undefined)
112
+ return this.cleanupPromise;
113
+ this.disposed = true;
114
+ this.cleanupPromise = this.performCleanup();
115
+ return this.cleanupPromise;
116
+ }
117
+ async performCleanup() {
118
+ await Promise.allSettled(this.activeResolves);
119
+ this.activeResolves.clear();
120
+ await Promise.allSettled(this.overrides.values());
121
+ this.overrides.clear();
122
+ if (this.rootPromise === undefined)
123
+ return;
124
+ const root = await this.rootPromise.catch(() => undefined);
125
+ if (root !== undefined) {
126
+ await rm(root, { recursive: true, force: true }).catch(() => undefined);
127
+ }
128
+ }
129
+ async createOverride(sourcePath, sourceConfig, cacheFloorMb, modelType) {
130
+ const root = await this.getRoot();
131
+ const overrideDir = await mkdtemp(join(root, 'model-'));
132
+ const config = {
133
+ ...sourceConfig,
134
+ use_block_paged_cache: true,
135
+ };
136
+ if (cacheFloorMb !== undefined) {
137
+ config.paged_cache_memory_mb = Math.max(positiveNumber(sourceConfig.paged_cache_memory_mb) ?? 0, cacheFloorMb);
138
+ }
139
+ await writeFile(join(overrideDir, 'config.json'), JSON.stringify(config, null, 2), 'utf-8');
140
+ const sourceEntries = await readdir(sourcePath);
141
+ for (const name of sourceEntries) {
142
+ if (name === 'config.json')
143
+ continue;
144
+ const source = join(sourcePath, name);
145
+ const destination = join(overrideDir, name);
146
+ let isFile;
147
+ try {
148
+ isFile = (await stat(source)).isFile();
149
+ }
150
+ catch {
151
+ continue;
152
+ }
153
+ if (!isFile)
154
+ continue;
155
+ try {
156
+ await symlink(source, destination);
157
+ }
158
+ catch (error) {
159
+ if (error.code !== 'EEXIST')
160
+ throw error;
161
+ }
162
+ }
163
+ if (QWEN35_CACHE_FLOOR_MODEL_TYPES.has(modelType)) {
164
+ await preserveQwen35MtpSidecars(sourcePath, overrideDir, sourceConfig, modelType);
165
+ }
166
+ return overrideDir;
167
+ }
168
+ getRoot() {
169
+ this.rootPromise ??= mkdtemp(join(tmpdir(), this.tempDirPrefix));
170
+ return this.rootPromise;
171
+ }
172
+ }
173
+ /**
174
+ * Keep only the nested paths that the native Qwen3.5 loaders inspect.
175
+ *
176
+ * Both dense and MoE loaders accept an mlx-vlm `mtp-drafter/` directory.
177
+ * Dense additionally accepts a config-declared relative sidecar and the
178
+ * conventional `mtp/weights.safetensors` fallback. Top-level sidecar files
179
+ * are already handled by the normal source-file loop above.
180
+ */
181
+ async function preserveQwen35MtpSidecars(sourcePath, overrideDir, sourceConfig, modelType) {
182
+ await symlinkDirectoryIfPresent(join(sourcePath, QWEN35_MTP_DRAFTER_DIR), join(overrideDir, QWEN35_MTP_DRAFTER_DIR));
183
+ // The MoE loader supports the split drafter directory, but deliberately has
184
+ // no `mtp.safetensors`-style sidecar discovery path.
185
+ if (modelType !== 'qwen3_5')
186
+ return;
187
+ const relativeSidecars = new Set([QWEN35_DENSE_NESTED_MTP_SIDECAR]);
188
+ const configuredSidecar = configuredQwen35MtpFile(sourceConfig);
189
+ if (configuredSidecar !== undefined && !isAbsolute(configuredSidecar)) {
190
+ relativeSidecars.add(configuredSidecar);
191
+ }
192
+ for (const sidecar of relativeSidecars) {
193
+ await symlinkContainedFileIfPresent(sourcePath, overrideDir, sidecar);
194
+ }
195
+ }
196
+ function configuredQwen35MtpFile(config) {
197
+ const extra = config.mlx_lm_extra_tensors;
198
+ if (extra === null || typeof extra !== 'object' || Array.isArray(extra))
199
+ return undefined;
200
+ const mtpFile = extra.mtp_file;
201
+ return typeof mtpFile === 'string' && mtpFile.trim() !== '' ? mtpFile : undefined;
202
+ }
203
+ async function symlinkContainedFileIfPresent(sourceRoot, destinationRoot, relativePath) {
204
+ const source = resolve(sourceRoot, relativePath);
205
+ const normalizedRelative = relative(sourceRoot, source);
206
+ if (normalizedRelative === '' ||
207
+ normalizedRelative === '..' ||
208
+ normalizedRelative.startsWith(`..${sep}`) ||
209
+ isAbsolute(normalizedRelative)) {
210
+ return;
211
+ }
212
+ try {
213
+ if (!(await stat(source)).isFile())
214
+ return;
215
+ }
216
+ catch {
217
+ return;
218
+ }
219
+ const destination = join(destinationRoot, normalizedRelative);
220
+ await mkdir(dirname(destination), { recursive: true });
221
+ await symlinkIfMissing(source, destination);
222
+ }
223
+ async function symlinkDirectoryIfPresent(source, destination) {
224
+ if (!(await isDirectory(source)))
225
+ return;
226
+ await symlinkIfMissing(source, destination);
227
+ }
228
+ async function symlinkIfMissing(source, destination) {
229
+ try {
230
+ await symlink(source, destination);
231
+ }
232
+ catch (error) {
233
+ if (error.code !== 'EEXIST')
234
+ throw error;
235
+ }
236
+ }
237
+ async function isDirectory(path) {
238
+ try {
239
+ return (await stat(path)).isDirectory();
240
+ }
241
+ catch {
242
+ return false;
243
+ }
244
+ }
245
+ function positiveNumber(value) {
246
+ return typeof value === 'number' && Number.isFinite(value) && value > 0 ? value : undefined;
247
+ }
248
+ function resolveQwen35CacheFloorMb() {
249
+ const raw = process.env.MLX_PAGED_CACHE_MEMORY_MB;
250
+ if (raw == null || raw === '')
251
+ return DEFAULT_QWEN35_PAGED_CACHE_MB;
252
+ const parsed = Number.parseInt(raw, 10);
253
+ return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_QWEN35_PAGED_CACHE_MB;
254
+ }
@@ -1 +1 @@
1
- {"version":3,"file":"qwen3_5-configs.d.ts","sourceRoot":"","sources":["../../src/models/qwen3_5-configs.ts"],"names":[],"mappings":"AAAA;;;;;GAKG;AAEH,OAAO,KAAK,EACV,YAAY,IAAI,gBAAgB,EAChC,sBAAsB,IAAI,0BAA0B,EACpD,sBAAsB,IAAI,0BAA0B,EACrD,MAAM,gBAAgB,CAAC;AAExB,MAAM,MAAM,YAAY,GAAG,gBAAgB,CAAC;AAC5C,MAAM,MAAM,sBAAsB,GAAG,0BAA0B,CAAC;AAChE,MAAM,MAAM,sBAAsB,GAAG,0BAA0B,CAAC;AAEhE;;GAEG;AACH,eAAO,MAAM,cAAc,EAAE;IAAE,CAAC,GAAG,EAAE,MAAM,GAAG,YAAY,CAAA;CAyBzD,CAAC;AAEF;;;;;;GAMG;AACH,wBAAgB,eAAe,CAAC,IAAI,EAAE,MAAM,GAAG,YAAY,CAM1D"}
1
+ {"version":3,"file":"qwen3_5-configs.d.ts","sourceRoot":"","sources":["../../src/models/qwen3_5-configs.ts"],"names":[],"mappings":"AAAA;;;;;GAKG;AAEH,OAAO,KAAK,EACV,YAAY,IAAI,gBAAgB,EAChC,sBAAsB,IAAI,0BAA0B,EACpD,sBAAsB,IAAI,0BAA0B,EACrD,MAAM,gBAAgB,CAAC;AAExB,MAAM,MAAM,YAAY,GAAG,gBAAgB,CAAC;AAC5C,MAAM,MAAM,sBAAsB,GAAG,0BAA0B,CAAC;AAChE,MAAM,MAAM,sBAAsB,GAAG,0BAA0B,CAAC;AAEhE;;GAEG;AACH,eAAO,MAAM,cAAc,EAAE;IAAE,CAAC,GAAG,EAAE,MAAM,GAAG,YAAY,CAAA;CA8BzD,CAAC;AAEF;;;;;;GAMG;AACH,wBAAgB,eAAe,CAAC,IAAI,EAAE,MAAM,GAAG,YAAY,CAM1D"}
@@ -31,6 +31,11 @@ export const QWEN35_CONFIGS = {
31
31
  fullAttentionInterval: 4,
32
32
  partialRotaryFactor: 0.25,
33
33
  ropeTheta: 100000.0,
34
+ // W1 (MTP): no MTP head in the stock 0.6B dense checkpoint.
35
+ // Real values are populated from `config.json` at load time; this
36
+ // preset is only used by tests that build a `Qwen35Config` from
37
+ // scratch.
38
+ nMtpLayers: 0,
34
39
  },
35
40
  };
36
41
  /**