@mlx-node/lm 0.0.7 → 0.0.9

package/dist/index.d.ts

@@ -12,30 +12,25 @@
  * const result = await model.generate([{ role: 'user', content: 'Hello!' }]);
  * ```
  */
-export { Qwen3Tokenizer } from '@mlx-node/core';
 export { Qwen3Model } from './stream.js';
-export { Gemma4Model, Gemma4Model as Gemma4_Model } from './stream.js';
-export type { Gemma4Config } from '@mlx-node/core';
+export { Gemma4Model } from './stream.js';
 export { HarrierModel } from '@mlx-node/core';
-export type { HarrierConfig } from '@mlx-node/core';
-export { Qwen35Model, Qwen35Model as Qwen3_5Model } from './stream.js';
-export type { Qwen35Config, Qwen35GenerationConfig, Qwen35GenerationResult } from '@mlx-node/core';
+export { Qwen35Model } from './stream.js';
+export type { Qwen35Config, Qwen35ContextLimits } from '@mlx-node/core';
 export { Lfm2Model } from './stream.js';
-export type { Lfm2Config } from '@mlx-node/core';
 export { LFM2_CONFIGS, getLfm2Config } from './models/lfm2-configs.js';
-export { Qwen35MoeModel, Qwen35MoeModel as Qwen3_5MoeModel } from './stream.js';
-export type { Qwen35MoeConfig, Qwen35MoeGenerationConfig, Qwen35MoeGenerationResult } from '@mlx-node/core';
-export type { DType } from '@mlx-node/core';
-export type { SamplingConfig, BatchGenerationResult } from '@mlx-node/core';
+export { Qwen35MoeModel } from './stream.js';
+export type { Qwen35MoeConfig } from '@mlx-node/core';
+export { memoryStats } from '@mlx-node/core';
 export type { ChatConfig, ChatResult, ChatMessage, ToolCallResult, PerformanceMetrics } from '@mlx-node/core';
 export type { ChatStreamDelta, ChatStreamFinal, ChatStreamEvent } from './stream.js';
-export type { ChatStreamChunk, ChatStreamHandle } from '@mlx-node/core';
-export { _runChatStream } from './stream.js';
-export { ChatSession } from './chat-session.js';
-export type { ChatSessionOptions, SendOptions, SessionCapableModel } from './chat-session.js';
+export { _runChatStream, makeStreamingModel } from './stream.js';
+export type { NativeStreamingInstance, NativeStreamingMethod, StreamingInstance, StreamingModel } from './stream.js';
+export { ChatSession, ContextCapacityError, isContextCapacityError } from './chat-session.js';
+export type { ChatSessionOptions, SendOptions, SessionCapableModel, SessionContextLimits } from './chat-session.js';
 export { type Qwen3Config, QWEN3_CONFIGS, type GenerationResult, type GenerationConfig, getQwen3Config, } from './models/qwen3-configs.js';
-export { loadModel, loadSession, detectModelType, type ModelType } from './models/model-loader.js';
-export type { TrainableModel, LoadableModel, EmbeddingModel } from './interfaces.js';
+export { loadModel, loadSession, detectModelType, type LoadableModel, type TrainableModel, type ModelType, type LoadModelOptions, } from './models/model-loader.js';
+export { AGENT_PAGED_MODEL_TYPES, PagedConfigOverrideManager, QWEN35_PAGED_MODEL_TYPES, type PagedConfigOverrideManagerOptions, } from './models/paged-config-override.js';
 export { QWEN35_CONFIGS, getQwen35Config } from './models/qwen3_5-configs.js';
 export * from './tools/index.js';
 export { enableProfiling, disableProfiling } from './profiling.js';

package/dist/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;GAaG;AAGH,OAAO,EAAE,cAAc,EAAE,MAAM,gBAAgB,CAAC;AAChD,OAAO,EAAE,UAAU,EAAE,MAAM,aAAa,CAAC;AAGzC,OAAO,EAAE,WAAW,EAAE,WAAW,IAAI,YAAY,EAAE,MAAM,aAAa,CAAC;AACvE,YAAY,EAAE,YAAY,EAAE,MAAM,gBAAgB,CAAC;AAGnD,OAAO,EAAE,YAAY,EAAE,MAAM,gBAAgB,CAAC;AAC9C,YAAY,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AACpD,OAAO,EAAE,WAAW,EAAE,WAAW,IAAI,YAAY,EAAE,MAAM,aAAa,CAAC;AACvE,YAAY,EAAE,YAAY,EAAE,sBAAsB,EAAE,sBAAsB,EAAE,MAAM,gBAAgB,CAAC;AAGnG,OAAO,EAAE,SAAS,EAAE,MAAM,aAAa,CAAC;AACxC,YAAY,EAAE,UAAU,EAAE,MAAM,gBAAgB,CAAC;AACjD,OAAO,EAAE,YAAY,EAAE,aAAa,EAAE,MAAM,0BAA0B,CAAC;AAGvE,OAAO,EAAE,cAAc,EAAE,cAAc,IAAI,eAAe,EAAE,MAAM,aAAa,CAAC;AAChF,YAAY,EAAE,eAAe,EAAE,yBAAyB,EAAE,yBAAyB,EAAE,MAAM,gBAAgB,CAAC;AAK5G,YAAY,EAAE,KAAK,EAAE,MAAM,gBAAgB,CAAC;AAC5C,YAAY,EAAE,cAAc,EAAE,qBAAqB,EAAE,MAAM,gBAAgB,CAAC;AAG5E,YAAY,EAAE,UAAU,EAAE,UAAU,EAAE,WAAW,EAAE,cAAc,EAAE,kBAAkB,EAAE,MAAM,gBAAgB,CAAC;AAG9G,YAAY,EAAE,eAAe,EAAE,eAAe,EAAE,eAAe,EAAE,MAAM,aAAa,CAAC;AACrF,YAAY,EAAE,eAAe,EAAE,gBAAgB,EAAE,MAAM,gBAAgB,CAAC;AAMxE,OAAO,EAAE,cAAc,EAAE,MAAM,aAAa,CAAC;AAM7C,OAAO,EAAE,WAAW,EAAE,MAAM,mBAAmB,CAAC;AAChD,YAAY,EAAE,kBAAkB,EAAE,WAAW,EAAE,mBAAmB,EAAE,MAAM,mBAAmB,CAAC;AAG9F,OAAO,EACL,KAAK,WAAW,EAChB,aAAa,EACb,KAAK,gBAAgB,EACrB,KAAK,gBAAgB,EACrB,cAAc,GACf,MAAM,2BAA2B,CAAC;AAGnC,OAAO,EAAE,SAAS,EAAE,WAAW,EAAE,eAAe,EAAE,KAAK,SAAS,EAAE,MAAM,0BAA0B,CAAC;AAGnG,YAAY,EAAE,cAAc,EAAE,aAAa,EAAE,cAAc,EAAE,MAAM,iBAAiB,CAAC;AAErF,OAAO,EAAE,cAAc,EAAE,eAAe,EAAE,MAAM,6BAA6B,CAAC;AAG9E,cAAc,kBAAkB,CAAC;AAGjC,OAAO,EAAE,eAAe,EAAE,gBAAgB,EAAE,MAAM,gBAAgB,CAAC"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;GAaG;AAGH,OAAO,EAAE,UAAU,EAAE,MAAM,aAAa,CAAC;AAGzC,OAAO,EAAE,WAAW,EAAE,MAAM,aAAa,CAAC;AAG1C,OAAO,EAAE,YAAY,EAAE,MAAM,gBAAgB,CAAC;AAC9C,OAAO,EAAE,WAAW,EAAE,MAAM,aAAa,CAAC;AAC1C,YAAY,EAAE,YAAY,EAAE,mBAAmB,EAAE,MAAM,gBAAgB,CAAC;AAGxE,OAAO,EAAE,SAAS,EAAE,MAAM,aAAa,CAAC;AACxC,OAAO,EAAE,YAAY,EAAE,aAAa,EAAE,MAAM,0BAA0B,CAAC;AAGvE,OAAO,EAAE,cAAc,EAAE,MAAM,aAAa,CAAC;AAC7C,YAAY,EAAE,eAAe,EAAE,MAAM,gBAAgB,CAAC;AAsBtD,OAAO,EAAE,WAAW,EAAE,MAAM,gBAAgB,CAAC;AAG7C,YAAY,EAAE,UAAU,EAAE,UAAU,EAAE,WAAW,EAAE,cAAc,EAAE,kBAAkB,EAAE,MAAM,gBAAgB,CAAC;AAG9G,YAAY,EAAE,eAAe,EAAE,eAAe,EAAE,eAAe,EAAE,MAAM,aAAa,CAAC;AASrF,OAAO,EAAE,cAAc,EAAE,kBAAkB,EAAE,MAAM,aAAa,CAAC;AACjE,YAAY,EAAE,uBAAuB,EAAE,qBAAqB,EAAE,iBAAiB,EAAE,cAAc,EAAE,MAAM,aAAa,CAAC;AAMrH,OAAO,EAAE,WAAW,EAAE,oBAAoB,EAAE,sBAAsB,EAAE,MAAM,mBAAmB,CAAC;AAC9F,YAAY,EAAE,kBAAkB,EAAE,WAAW,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,MAAM,mBAAmB,CAAC;AAGpH,OAAO,EACL,KAAK,WAAW,EAChB,aAAa,EACb,KAAK,gBAAgB,EACrB,KAAK,gBAAgB,EACrB,cAAc,GACf,MAAM,2BAA2B,CAAC;AAGnC,OAAO,EACL,SAAS,EACT,WAAW,EACX,eAAe,EACf,KAAK,aAAa,EAClB,KAAK,cAAc,EACnB,KAAK,SAAS,EACd,KAAK,gBAAgB,GACtB,MAAM,0BAA0B,CAAC;AAElC,OAAO,EACL,uBAAuB,EACvB,0BAA0B,EAC1B,wBAAwB,EACxB,KAAK,iCAAiC,GACvC,MAAM,mCAAmC,CAAC;AAE3C,OAAO,EAAE,cAAc,EAAE,eAAe,EAAE,MAAM,6BAA6B,CAAC;AAG9E,cAAc,kBAAkB,CAAC;AAGjC,OAAO,EAAE,eAAe,EAAE,gBAAgB,EAAE,MAAM,gBAAgB,CAAC"}

package/dist/index.js

@@ -13,34 +13,58 @@
  * ```
  */
 // Model classes (for inference)
-export { Qwen3Tokenizer } from '@mlx-node/core';
 export { Qwen3Model } from './stream.js';
 // Gemma4 models
-export { Gemma4Model, Gemma4Model as Gemma4_Model } from './stream.js';
+export { Gemma4Model } from './stream.js';
 // Embedding models
 export { HarrierModel } from '@mlx-node/core';
-export { Qwen35Model, Qwen35Model as Qwen3_5Model } from './stream.js';
+export { Qwen35Model } from './stream.js';
 // LFM2 models
 export { Lfm2Model } from './stream.js';
 export { LFM2_CONFIGS, getLfm2Config } from './models/lfm2-configs.js';
 // MoE variant
-export { Qwen35MoeModel, Qwen35MoeModel as Qwen3_5MoeModel } from './stream.js';
+export { Qwen35MoeModel } from './stream.js';
+// Memory hygiene: most management is automatic — the decode loop
+// inside `@mlx-node/core` calls `mlx_clear_cache()` every 256 generated
+// tokens to prevent unbounded free-pool growth during long
+// generations, and `MLX_CACHE_LIMIT_GB` auto-tunes the Metal pool cap
+// at model load. Across-request drains are handled by the
+// `@mlx-node/server` idle sweeper (see `packages/server/src/idle-sweeper.ts`):
+// a single `clearCache()` fires after `idleClearCacheMs` of HTTP
+// inactivity once the in-flight request counter has returned to zero.
+// `memoryStats()` is re-exported as a read-only observability hook for
+// dashboards / debugging.
+//
+// `clearCache()` is DELIBERATELY not re-exported here: the native impl
+// routes through MLX's no-arg `synchronize()` which waits only on the
+// default stream, so calling it while a decode runs on a model's
+// custom stream risks racing live Metal command buffers. The only
+// safe caller today is `@mlx-node/server`'s idle sweeper (fires after
+// the in-flight request counter hits zero AND — for hot-load flows —
+// outside any `withSuspendedDrains()` bracket). Admin / cron code that
+// reaches for a manual drain should deep-import from `@mlx-node/core`
+// directly and read the `@internal` caveat there.
+export { memoryStats } from '@mlx-node/core';
 // Internal: exported for testing the callback-to-AsyncGenerator bridge
 // Not part of the public API — may change without notice.
 // `_runChatStream` is the generic adapter used by every model wrapper
 // (and the VLM package's QianfanOCR wrapper) to turn a callback-based
 // native stream into an `AsyncGenerator<ChatStreamEvent>`.
-export { _runChatStream } from './stream.js';
+// `makeStreamingModel` is the factory that builds each family's wrapper
+// subclass from its native class; the VLM package reuses it to build
+// `QianfanOCRModel`.
+export { _runChatStream, makeStreamingModel } from './stream.js';
 // Cross-model chat session wrapper (see chat-session.ts for design notes).
 // `SessionCapableModel` is the structural interface matched by every
 // generative model wrapper and used as the upper-bound for
 // `ChatSession<M>`; exported so the VLM wrapper can pin a compile-time
 // conformance assertion.
-export { ChatSession } from './chat-session.js';
+export { ChatSession, ContextCapacityError, isContextCapacityError } from './chat-session.js';
 // Model utilities (TypeScript-only)
 export { QWEN3_CONFIGS, getQwen3Config, } from './models/qwen3-configs.js';
 // Model loading
-export { loadModel, loadSession, detectModelType } from './models/model-loader.js';
+export { loadModel, loadSession, detectModelType, } from './models/model-loader.js';
+export { AGENT_PAGED_MODEL_TYPES, PagedConfigOverrideManager, QWEN35_PAGED_MODEL_TYPES, } from './models/paged-config-override.js';
 export { QWEN35_CONFIGS, getQwen35Config } from './models/qwen3_5-configs.js';
 // Tool calling utilities
 export * from './tools/index.js';

package/dist/interfaces.d.ts

@@ -1,18 +1,3 @@
-import type { Gemma4Model, HarrierModel, Lfm2Model, Qwen3Model, Qwen35Model, Qwen35MoeModel, QianfanOCRModel } from '@mlx-node/core';
-/**
- * Union of all model classes that can be used with training engines.
- * Uses the native (core) types so trainers can pass instances directly
- * to Rust engine factory methods without type conflicts.
- */
-export type TrainableModel = Qwen3Model | Qwen35Model | Qwen35MoeModel;
-/**
- * Union of all embedding model classes.
- * These models encode text into fixed-size embedding vectors.
- */
-export type EmbeddingModel = HarrierModel;
-/**
- * Union of all model classes that loadModel can return.
- * Includes trainable models, inference-only models, and embedding models.
- */
-export type LoadableModel = TrainableModel | QianfanOCRModel | EmbeddingModel | Gemma4Model | Lfm2Model;
+/** Compatibility type barrel; model unions are derived from the loader registry. */
+export type { LoadableModel, TrainableModel } from './models/model-loader.js';
 //# sourceMappingURL=interfaces.d.ts.map

package/dist/interfaces.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"interfaces.d.ts","sourceRoot":"","sources":["../src/interfaces.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACV,WAAW,EACX,YAAY,EACZ,SAAS,EACT,UAAU,EACV,WAAW,EACX,cAAc,EACd,eAAe,EAChB,MAAM,gBAAgB,CAAC;AAExB;;;;GAIG;AACH,MAAM,MAAM,cAAc,GAAG,UAAU,GAAG,WAAW,GAAG,cAAc,CAAC;AAEvE;;;GAGG;AACH,MAAM,MAAM,cAAc,GAAG,YAAY,CAAC;AAE1C;;;GAGG;AACH,MAAM,MAAM,aAAa,GAAG,cAAc,GAAG,eAAe,GAAG,cAAc,GAAG,WAAW,GAAG,SAAS,CAAC"}
+{"version":3,"file":"interfaces.d.ts","sourceRoot":"","sources":["../src/interfaces.ts"],"names":[],"mappings":"AAAA,oFAAoF;AACpF,YAAY,EAAE,aAAa,EAAE,cAAc,EAAE,MAAM,0BAA0B,CAAC"}

package/dist/models/lfm2-configs.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"lfm2-configs.d.ts","sourceRoot":"","sources":["../../src/models/lfm2-configs.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,gBAAgB,CAAC;AAEjD,eAAO,MAAM,YAAY,EAAE;IAAE,CAAC,GAAG,EAAE,MAAM,GAAG,UAAU,CAAA;CAwCrD,CAAC;AAEF,wBAAgB,aAAa,CAAC,IAAI,EAAE,MAAM,GAAG,UAAU,CAMtD"}
+{"version":3,"file":"lfm2-configs.d.ts","sourceRoot":"","sources":["../../src/models/lfm2-configs.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,gBAAgB,CAAC;AAEjD,eAAO,MAAM,YAAY,EAAE;IAAE,CAAC,GAAG,EAAE,MAAM,GAAG,UAAU,CAAA;CAmGrD,CAAC;AAEF,wBAAgB,aAAa,CAAC,IAAI,EAAE,MAAM,GAAG,UAAU,CAMtD"}

package/dist/models/lfm2-configs.js

@@ -1,4 +1,61 @@
 export const LFM2_CONFIGS = {
+    'lfm2.5-8b-a1b': {
+        vocabSize: 128000,
+        hiddenSize: 2048,
+        numHiddenLayers: 24,
+        numAttentionHeads: 32,
+        numKeyValueHeads: 8,
+        maxPositionEmbeddings: 128000,
+        normEps: 1e-5,
+        convBias: false,
+        convLCache: 3,
+        // block* fields are unused for MoE checkpoints (MoE dense layers use
+        // intermediateSize directly; MoE sparse layers use moeIntermediateSize).
+        // Required by the NAPI-derived Lfm2Config interface; set to 0 / defaults.
+        blockDim: 0,
+        blockFfDim: 0,
+        blockMultipleOf: 256,
+        blockFfnDimMultiplier: 1.0,
+        blockAutoAdjustFfDim: false,
+        ropeTheta: 5000000.0,
+        layerTypes: [
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+            'full_attention',
+            'conv',
+            'conv',
+        ],
+        tieEmbedding: true,
+        eosTokenId: 124900,
+        bosTokenId: 124894,
+        padTokenId: 124893,
+        intermediateSize: 7168,
+        moeIntermediateSize: 1792,
+        numExperts: 32,
+        numExpertsPerTok: 4,
+        numDenseLayers: 2,
+        normTopkProb: true,
+        useExpertBias: true,
+    },
     'lfm2.5-1.2b-thinking': {
         vocabSize: 65536,
         hiddenSize: 2048,
@@ -37,6 +94,8 @@ export const LFM2_CONFIGS = {
         eosTokenId: 7,
         bosTokenId: 1,
         padTokenId: 0,
+        normTopkProb: true,
+        useExpertBias: true,
     },
 };
 export function getLfm2Config(name) {

package/dist/models/model-loader.d.ts

@@ -3,16 +3,178 @@
  *
  * Handles loading pretrained weights from MLX format or converting from HuggingFace.
  */
+import { Gemma4Model as NativeGemma4Model, HarrierModel, Lfm2Model as NativeLfm2Model, QianfanOCRModel, Qwen3Model as NativeQwen3Model, Qwen35Model as NativeQwen35Model, Qwen35MoeModel as NativeQwen35MoeModel } from '@mlx-node/core';
 import { ChatSession, type SessionCapableModel } from '../chat-session.js';
-import type { LoadableModel } from '../interfaces.js';
-export type ModelType = 'qwen3' | 'qwen3_5' | 'qwen3_5_moe' | 'internvl_chat' | 'qianfan-ocr' | 'harrier' | 'gemma4' | 'lfm2';
+/** Optional settings for {@link loadModel} / {@link loadSession}. */
+export interface LoadModelOptions {
+    /**
+     * Gemma4 only: directory of an external draft checkpoint (config.json +
+     * model.safetensors) loaded alongside the target model for speculative
+     * decoding (forwarded as `Gemma4LoadOptions.draftModelPath`). Accepts
+     * either a DSpark draft or a Google gemma-4 assistant draft
+     * (`google/gemma-4-*-it-assistant`); the variant is auto-detected from
+     * the draft's config.json (`model_type` `gemma4_assistant` /
+     * `gemma4_unified_assistant` → assistant, `architectures` containing
+     * `Gemma4DSparkModel` → DSpark). When omitted, Gemma4 automatically loads
+     * an embedded draft from `<modelPath>/draft/` when present. Draft decoding
+     * runs on the flat KV-cache path, so the target checkpoint must not
+     * explicitly enable `use_block_paged_cache`. Setting this for any other
+     * model family is a hard error — no other loader accepts a draft model.
+     */
+    draftModelPath?: string;
+}
+interface NormalizedModelConfig {
+    readonly usesDefaultModelType: boolean;
+    readonly rawModelType: string | undefined;
+    readonly rawModelTypeLabel: string;
+    readonly architectures: ReadonlySet<string>;
+}
+interface ModelConfigMatchContext extends NormalizedModelConfig {
+    readonly modelType: string | undefined;
+}
+/**
+ * Ordered source of truth for every supported model family. Each entry owns
+ * its canonical `ModelType`, raw config aliases / architecture probes, loader,
+ * and `ChatSession` eligibility:
+ *
+ *   - `'trainable'` — GRPO/SFT-capable LM (Qwen3 family); chat-capable.
+ *   - `'loadable'`  — chat-capable LM with no trainer engine (Gemma4, LFM2).
+ *   - `'embedding'` — no chat surface (Harrier); rejected by `loadSession`.
+ *   - `'vlm'`       — VLM whose AsyncGenerator wrapper lives in
+ *                     `@mlx-node/vlm` (importing it here would create a
+ *                     circular package dependency), so `loadSession`
+ *                     rejects it and routes callers to `@mlx-node/vlm`.
+ *
+ * A base family is selected from an explicit alias or the single declarative
+ * nullish-model_type default, then architecture probes refine it in declaration
+ * order. Gemma's unified architecture is authoritative (matching the native
+ * loader); Harrier refines a Qwen3 base. Adding a family means adding one
+ * descriptor here, without a second normalization or dispatch branch.
+ */
+declare const MODEL_FAMILY_REGISTRY: readonly [{
+    readonly modelType: 'gemma4';
+    readonly kind: 'loadable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["gemma4", "gemma4_text", "gemma4_unified"];
+        readonly architectureProbe: ({ architectures }: ModelConfigMatchContext) => boolean;
+    };
+    readonly load: (modelPath: string, options?: LoadModelOptions) => Promise<import("../stream.js").StreamingInstance<typeof NativeGemma4Model, {
+        readonly recordModelPath: true;
+    }>>;
+    readonly nativeModelClass: typeof NativeGemma4Model;
+    readonly acceptsDraftModel: true;
+}, {
+    readonly modelType: 'harrier';
+    readonly kind: 'embedding';
+    readonly match: {
+        readonly rawModelTypes: readonly ["harrier"];
+        readonly architectureProbe: ({ modelType, architectures }: ModelConfigMatchContext) => boolean;
+    };
+    readonly load: (modelPath: string) => Promise<HarrierModel>;
+    readonly nativeModelClass: typeof HarrierModel;
+}, {
+    readonly modelType: 'qwen3';
+    readonly kind: 'trainable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["qwen3"];
+    };
+    readonly load: (modelPath: string) => Promise<import("../stream.js").StreamingInstance<typeof NativeQwen3Model, {
+        readonly recordModelPath: true;
+        readonly applyTemplate: false;
+    }>>;
+    readonly nativeModelClass: typeof NativeQwen3Model;
+    readonly defaultForNullishModelType: true;
+}, {
+    readonly modelType: 'qwen3_5';
+    readonly kind: 'trainable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["qwen3_5"];
+    };
+    readonly load: (modelPath: string) => Promise<import("../stream.js").StreamingInstance<typeof NativeQwen35Model, {
+        readonly recordModelPath: true;
+    }>>;
+    readonly nativeModelClass: typeof NativeQwen35Model;
+}, {
+    readonly modelType: 'qwen3_5_moe';
+    readonly kind: 'trainable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["qwen3_5_moe"];
+    };
+    readonly load: (modelPath: string) => Promise<import("../stream.js").StreamingInstance<typeof NativeQwen35MoeModel, {
+        readonly recordModelPath: true;
+    }>>;
+    readonly nativeModelClass: typeof NativeQwen35MoeModel;
+}, {
+    readonly modelType: 'lfm2';
+    readonly kind: 'loadable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["lfm2"];
+    };
+    readonly load: (modelPath: string) => Promise<import("../stream.js").StreamingInstance<typeof NativeLfm2Model, {
+        readonly recordModelPath: true;
+        readonly replayAssistantRawText: true;
+    }>>;
+    readonly nativeModelClass: typeof NativeLfm2Model;
+}, {
+    readonly modelType: 'lfm2_moe';
+    readonly kind: 'loadable';
+    readonly match: {
+        readonly rawModelTypes: readonly ["lfm2_moe"];
+    };
+    readonly load: (modelPath: string) => Promise<import("../stream.js").StreamingInstance<typeof NativeLfm2Model, {
+        readonly recordModelPath: true;
+        readonly replayAssistantRawText: true;
+    }>>;
+    readonly nativeModelClass: typeof NativeLfm2Model;
+}, {
+    readonly modelType: 'internvl_chat';
+    readonly kind: 'vlm';
+    readonly match: {
+        readonly rawModelTypes: readonly ["internvl_chat"];
+    };
+    readonly load: (modelPath: string) => Promise<QianfanOCRModel>;
+    readonly nativeModelClass: typeof QianfanOCRModel;
+}, {
+    readonly modelType: 'qianfan-ocr';
+    readonly kind: 'vlm';
+    readonly match: {
+        readonly rawModelTypes: readonly ["qianfan-ocr"];
+    };
+    readonly load: (modelPath: string) => Promise<QianfanOCRModel>;
+    readonly nativeModelClass: typeof QianfanOCRModel;
+}];
+export type ModelType = (typeof MODEL_FAMILY_REGISTRY)[number]['modelType'];
+type RegisteredModelFamily = (typeof MODEL_FAMILY_REGISTRY)[number];
+type RegisteredTrainableFamily = Extract<RegisteredModelFamily, {
+    readonly kind: 'trainable';
+}>;
+/**
+ * Union of the native `@mlx-node/core` model classes across every registered
+ * family — the public contract of {@link loadModel}. At runtime the chat
+ * families resolve to streaming-wrapper subclasses of these classes
+ * (AsyncGenerator `chatStream*` overrides), but the public type names the
+ * native classes so downstream code can pass instances directly to Rust
+ * engine factory methods without type conflicts.
+ */
+export type LoadableModel = InstanceType<RegisteredModelFamily['nativeModelClass']>;
+/**
+ * Union accepted by trainer APIs: registered wrapper results plus their native
+ * FFI instances. Both sides derive from the same trainable registry rows.
+ */
+export type TrainableModel = Awaited<ReturnType<RegisteredTrainableFamily['load']>> | InstanceType<RegisteredTrainableFamily['nativeModelClass']>;
 /**
  * Load a model from disk, auto-detecting architecture from config.json.
  *
  * Supports both language models (Qwen3, Qwen3.5) and vision-language models
  * (Qianfan-OCR / InternVL). Use `instanceof` to narrow the returned type.
+ *
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
+ * speculative decoding — gemma4 only; any other detected family rejects it.
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
+ * that embedded checkpoint is present.
  */
-export declare function loadModel(modelPath: string): Promise<LoadableModel>;
+export declare function loadModel(modelPath: string, options?: LoadModelOptions): Promise<LoadableModel>;
 /**
  * Load a model and wrap it in a {@link ChatSession} for multi-turn chat.
  *
@@ -28,7 +190,16 @@ export declare function loadModel(modelPath: string): Promise<LoadableModel>;
  *     package dependency), so callers who want a Qianfan-OCR session
  *     must import `QianfanOCRModel` from `@mlx-node/vlm` and construct
  *     `new ChatSession(model)` directly.
+ *
+ * `options.draftModelPath` attaches an external draft checkpoint (DSpark or
+ * Google gemma-4 assistant, auto-detected from the draft's config.json) for
+ * speculative decoding — gemma4 only; any other detected family rejects it.
+ * Without the option, Gemma4 loads `<modelPath>/draft/` automatically when
+ * that embedded checkpoint is present.
+ * The resulting session auto-enables the speculative path (the model
+ * reports `hasMtpWeights()`); pass `enableMtp: false` per call to opt out.
  */
-export declare function loadSession(modelPath: string): Promise<ChatSession<SessionCapableModel>>;
+export declare function loadSession(modelPath: string, options?: LoadModelOptions): Promise<ChatSession<SessionCapableModel>>;
 export declare function detectModelType(modelPath: string): Promise<ModelType>;
+export {};
 //# sourceMappingURL=model-loader.d.ts.map

package/dist/models/model-loader.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"model-loader.d.ts","sourceRoot":"","sources":["../../src/models/model-loader.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAOH,OAAO,EAAE,WAAW,EAAE,KAAK,mBAAmB,EAAE,MAAM,oBAAoB,CAAC;AAC3E,OAAO,KAAK,EAAE,aAAa,EAAkB,MAAM,kBAAkB,CAAC;AAGtE,MAAM,MAAM,SAAS,GACjB,OAAO,GACP,SAAS,GACT,aAAa,GACb,eAAe,GACf,aAAa,GACb,SAAS,GACT,QAAQ,GACR,MAAM,CAAC;AAaX;;;;;GAKG;AACH,wBAAsB,SAAS,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,aAAa,CAAC,CAoBzE;AAED;;;;;;;;;;;;;;;GAeG;AACH,wBAAsB,WAAW,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,WAAW,CAAC,mBAAmB,CAAC,CAAC,CAW9F;AAED,wBAAsB,eAAe,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,SAAS,CAAC,CAyB3E"}
+{"version":3,"file":"model-loader.d.ts","sourceRoot":"","sources":["../../src/models/model-loader.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAKH,OAAO,EACL,WAAW,IAAI,iBAAiB,EAChC,YAAY,EACZ,SAAS,IAAI,eAAe,EAC5B,eAAe,EACf,UAAU,IAAI,gBAAgB,EAC9B,WAAW,IAAI,iBAAiB,EAChC,cAAc,IAAI,oBAAoB,EACvC,MAAM,gBAAgB,CAAC;AAExB,OAAO,EAAE,WAAW,EAAE,KAAK,mBAAmB,EAAE,MAAM,oBAAoB,CAAC;AAG3E,qEAAqE;AACrE,MAAM,WAAW,gBAAgB;IAC/B;;;;;;;;;;;;;OAaG;IACH,cAAc,CAAC,EAAE,MAAM,CAAC;CACzB;AAID,UAAU,qBAAqB;IAC7B,QAAQ,CAAC,oBAAoB,EAAE,OAAO,CAAC;IACvC,QAAQ,CAAC,YAAY,EAAE,MAAM,GAAG,SAAS,CAAC;IAC1C,QAAQ,CAAC,iBAAiB,EAAE,MAAM,CAAC;IACnC,QAAQ,CAAC,aAAa,EAAE,WAAW,CAAC,MAAM,CAAC,CAAC;CAC7C;AAED,UAAU,uBAAwB,SAAQ,qBAAqB;IAC7D,QAAQ,CAAC,SAAS,EAAE,MAAM,GAAG,SAAS,CAAC;CACxC;AAgCD;;;;;;;;;;;;;;;;;;GAkBG;AACH,QAAA,MAAM,qBAAqB;wBAEZ,QAAQ;mBACb,UAAU;;;;;+BAKE,MAAM,YAAY,gBAAgB;;;;;;wBASzC,SAAS;mBACd,WAAW;;;;;+BAMC,MAAM;;;wBAIb,OAAO;mBACZ,WAAW;;iBACR,aAAa;;+BACJ,MAAM;;;;;;;wBAKb,SAAS;mBACd,WAAW;;iBACR,aAAa;;+BACJ,MAAM;;;;;wBAIb,aAAa;mBAClB,WAAW;;iBACR,aAAa;;+BACJ,MAAM;;;;;wBAIb,MAAM;mBACX,UAAU;;iBACP,aAAa;;+BACJ,MAAM;;;;;;wBAIb,UAAU;mBACf,UAAU;;iBACP,aAAa;;+BACJ,MAAM;;;;;;wBAIb,eAAe;mBACpB,KAAK;;iBACF,aAAa;;+BACJ,MAAM;;;wBAIb,aAAa;mBAClB,KAAK;;iBACF,aAAa;;+BACJ,MAAM;;EAGyB,CAAC;AAEtD,MAAM,MAAM,SAAS,GAAG,CAAC,OAAO,qBAAqB,CAAC,CAAC,MAAM,CAAC,CAAC,WAAW,CAAC,CAAC;AAE5E,KAAK,qBAAqB,GAAG,CAAC,OAAO,qBAAqB,CAAC,CAAC,MAAM,CAAC,CAAC;AACpE,KAAK,yBAAyB,GAAG,OAAO,CAAC,qBAAqB,EAAE;IAAE,QAAQ,CAAC,IAAI,EAAE,WAAW,CAAA;CAAE,CAAC,CAAC;AAEhG;;;;;;;GAOG;AACH,MAAM,MAAM,aAAa,GAAG,YAAY,CAAC,qBAAqB,CAAC,kBAAkB,CAAC,CAAC,CAAC;AAEpF;;;GAGG;AACH,MAAM,MAAM,cAAc,GACtB,OAAO,CAAC,UAAU,CAAC,yBAAyB,CAAC,MAAM,CAAC,CAAC,CAAC,GACtD,YAAY,CAAC,yBAAyB,CAAC,kBAAkB,CAAC,CAAC,CAAC;AAsIhE;;;;;;;;;;;GAWG;AACH,wBAAsB,SAAS,CAAC,SAAS,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,gBAAgB,GAAG,OAAO,CAAC,aAAa,CAAC,CAGrG;AAED;;;;;;;;;;;;;;;;;;;;;;;GAuBG;AACH,wBAAsB,WAAW,CAC/B,SAAS,EAAE,MAAM,EACjB,OAAO,CAAC,EAAE,gBAAgB,GACzB,OAAO,CAAC,WAAW,CAAC,mBAAmB,CAAC,CAAC,CAa3C;AAED,wBAAsB,eAAe,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,SAAS,CAAC,CAkB3E"}