@mlx-node/lm 0.0.12 → 0.0.15

package/src/models/paged-config-override.ts

@@ -0,0 +1,387 @@
+/**
+ * Process-local model-directory overrides for forcing block-paged KV caches.
+ *
+ * Native model loaders accept a directory path and read paging policy from
+ * `config.json`. This manager creates an isolated temporary clone containing a
+ * patched config plus symlinks to the source files, so callers can opt models
+ * into paging without mutating downloaded checkpoints. Loader-known Qwen3.5
+ * MTP sidecars are preserved explicitly; unrelated directories (including a
+ * Gemma4 `draft/`) remain hidden unless a caller explicitly preserves the
+ * Gemma draft for flat speculative decoding.
+ */
+
+import { mkdir, mkdtemp, readFile, readdir, rm, stat, symlink, writeFile } from 'node:fs/promises';
+import { tmpdir } from 'node:os';
+import { dirname, extname, isAbsolute, join, relative, resolve, sep } from 'node:path';
+
+import { ggufArchitecture, prepareMuseGlimmerGguf } from '@mlx-node/core';
+
+import { CHAT_FAMILY_IDS } from '../family-data.js';
+
+/** Families historically forced paged by `mlx launch claude`. */
+export const QWEN35_PAGED_MODEL_TYPES = ['qwen3_5', 'qwen3_5_moe'] as const;
+
+const QWEN35_CACHE_FLOOR_MODEL_TYPES = new Set<string>(QWEN35_PAGED_MODEL_TYPES);
+const DEFAULT_QWEN35_PAGED_CACHE_MB = 16_384;
+const DEFAULT_QWEN35_PAGED_CACHE_INITIAL_MB = 2_048;
+const QWEN35_MTP_DRAFTER_DIR = 'mtp-drafter';
+const QWEN35_DENSE_NESTED_MTP_SIDECAR = 'mtp/weights.safetensors';
+
+export interface PagedConfigOverrideManagerOptions {
+  /** Model types to force onto the paged path. Defaults to all agent chat families. */
+  modelTypes?: readonly string[];
+  /** Temporary-directory prefix. Primarily useful for diagnostics/tests. */
+  tempDirPrefix?: string;
+  /**
+   * Pass Gemma4 checkpoints with an embedded `draft/` through unchanged.
+   * This keeps native draft auto-discovery, which currently requires flat KV
+   * caches. Defaults to false so a paged override remains a paged contract.
+   */
+  preserveEmbeddedGemmaDraft?: boolean;
+}
+
+/**
+ * Owns one isolated set of temporary paged-config overrides.
+ *
+ * A manager is intentionally single-lifecycle: repeated resolution of one
+ * source returns the same override, and `cleanup()` permanently disposes the
+ * manager. Separate managers never share directories, so one launch cannot
+ * remove another launch's live override.
+ */
+export class PagedConfigOverrideManager {
+  private readonly modelTypes: ReadonlySet<string>;
+  private readonly tempDirPrefix: string;
+  private readonly preserveEmbeddedGemmaDraft: boolean;
+  private readonly overrides = new Map<string, Promise<string>>();
+  private readonly activeResolves = new Set<Promise<string>>();
+  private rootPromise: Promise<string> | undefined;
+  private cleanupPromise: Promise<void> | undefined;
+  private disposed = false;
+
+  constructor(options: PagedConfigOverrideManagerOptions = {}) {
+    this.modelTypes = new Set(options.modelTypes ?? CHAT_FAMILY_IDS);
+    this.tempDirPrefix = options.tempDirPrefix ?? 'mlx-paged-overrides-';
+    this.preserveEmbeddedGemmaDraft = options.preserveEmbeddedGemmaDraft ?? false;
+  }
+
+  /**
+   * Resolve `modelPath` to a paged-aware clone when its model type is managed.
+   * A caller-supplied canonical family takes precedence over the raw config
+   * type (for example, `gemma4` for a `gemma4_unified` checkpoint).
+   * Unmanaged, unreadable, or malformed checkpoints pass through unchanged.
+   *
+   * `persistPagedCache` is a tri-state cold-tier directive. `undefined` leaves
+   * the field untouched (families with no cold-tier opt-in). A boolean is
+   * AUTHORITATIVE: it writes `persist_paged_cache: <value>` into the cloned
+   * config, overriding whatever the source config.json carries (either alias),
+   * and forces a clone whenever the source's value disagrees so the directive
+   * actually reaches the loader. Callers gate the boolean to the families whose
+   * paged cold restore is sound — the allowlist in
+   * `packages/agent/src/cold-tier.ts`, mirrored by `COLD_RESTORE_FAMILIES` in
+   * `crates/mlx-core/src/cold_tier.rs`. It is a SET, not one family, and it is
+   * deliberately not repeated here: `MlxModelHost` calls this for every
+   * allowlisted family, and a second copy of the list would drift.
+   */
+  async resolve(modelPath: string, canonicalModelType?: string, persistPagedCache?: boolean): Promise<string> {
+    if (this.disposed) {
+      throw new Error('PagedConfigOverrideManager: resolve() called after cleanup()');
+    }
+
+    const operation = this.resolveInternal(modelPath, canonicalModelType, persistPagedCache);
+    this.activeResolves.add(operation);
+    try {
+      return await operation;
+    } finally {
+      this.activeResolves.delete(operation);
+    }
+  }
+
+  private async resolveInternal(
+    modelPath: string,
+    canonicalModelType?: string,
+    persistPagedCache?: boolean,
+  ): Promise<string> {
+    const sourcePath = isAbsolute(modelPath) ? modelPath : resolve(modelPath);
+    if (extname(sourcePath).toLowerCase() === '.gguf') {
+      let modelType = canonicalModelType;
+      if (modelType === undefined) {
+        try {
+          if (ggufArchitecture(sourcePath) === 'muse-glimmer') modelType = 'muse_glimmer';
+        } catch {
+          return modelPath;
+        }
+      }
+      if (modelType !== 'muse_glimmer' || !this.modelTypes.has(modelType)) return modelPath;
+      // A symlink to the GGUF would canonicalize back to the original config.
+      // Prepare once, then overlay the packed cache so paging/persistence
+      // directives reach the loader and the DFlash sidecar stays attached.
+      const prepared = await prepareMuseGlimmerGguf(sourcePath);
+      return this.resolveInternal(prepared, modelType, persistPagedCache);
+    }
+    let config: Record<string, unknown>;
+    try {
+      config = JSON.parse(await readFile(join(sourcePath, 'config.json'), 'utf-8')) as Record<string, unknown>;
+    } catch {
+      return modelPath;
+    }
+
+    const rawModelType = typeof config.model_type === 'string' ? config.model_type : null;
+    const modelType = canonicalModelType ?? rawModelType;
+    if (modelType === null || !this.modelTypes.has(modelType)) {
+      return modelPath;
+    }
+
+    // A boolean persist directive is authoritative: force a clone whenever the
+    // source config's EFFECTIVE value disagrees, so the directive actually reaches
+    // the loader (both `--no-persist-cache` off AND default-on). Mirror the native
+    // parser's snake-first precedence — `persist_paged_cache` wins when present,
+    // else the camelCase alias — so a config with snake=false AND camel=true does
+    // not read as `true` here (an OR) while native reads snake=false, silently
+    // dropping an authoritative `true` request.
+    const sourcePersist =
+      typeof config.persist_paged_cache === 'boolean' ? config.persist_paged_cache : config.persistPagedCache === true;
+    const persistOverrideNeeded = persistPagedCache !== undefined && persistPagedCache !== sourcePersist;
+
+    // Gemma4's DSpark / assistant speculative executor currently owns flat KV
+    // caches. Preserve native draft discovery only for explicit callers; the
+    // default paged clone intentionally omits subdirectories, hiding `draft/`
+    // while keeping the downloaded checkpoint unchanged.
+    const hasEmbeddedGemmaDraft = modelType === 'gemma4' && (await isDirectory(join(sourcePath, 'draft')));
+    if (this.preserveEmbeddedGemmaDraft && hasEmbeddedGemmaDraft) {
+      return modelPath;
+    }
+
+    const cacheFloorMb = QWEN35_CACHE_FLOOR_MODEL_TYPES.has(modelType) ? resolveQwen35CacheFloorMb() : undefined;
+    const pagedEnabled = config.use_block_paged_cache === true;
+    const configuredMemoryMb = positiveNumber(config.paged_cache_memory_mb);
+    const memorySatisfied = cacheFloorMb === undefined || (configuredMemoryMb ?? 0) >= cacheFloorMb;
+    // The clone writes `paged_cache_initial_memory_mb = min(initialMb, maxMb)`.
+    // A source whose field already equals that value needs no clone; an absent
+    // or different field does, or the new start-small default would silently
+    // skip every already-paged Qwen3.5 checkpoint (they lack the field).
+    const initialMb = QWEN35_CACHE_FLOOR_MODEL_TYPES.has(modelType) ? resolveQwen35InitialMb() : undefined;
+    const configuredInitialMb = positiveNumber(config.paged_cache_initial_memory_mb);
+    const initialSatisfied =
+      initialMb === undefined ||
+      configuredInitialMb === Math.min(initialMb, Math.max(configuredMemoryMb ?? 0, cacheFloorMb ?? 0));
+    // Even an already-paged Gemma config must be cloned when `draft/` exists:
+    // returning the source would expose the draft to native auto-discovery and
+    // trigger the flat-speculation/paged-cache conflict. An authoritative
+    // persist directive that disagrees with the source likewise blocks the
+    // pass-through so the resolved value reaches the loader.
+    if (pagedEnabled && memorySatisfied && initialSatisfied && !hasEmbeddedGemmaDraft && !persistOverrideNeeded) {
+      return modelPath;
+    }
+
+    // Memoize per (source, resolved family, persist directive): the same
+    // checkpoint resolved with a different persist tri-state must yield a distinct
+    // clone, not the first one cached under the bare path. `cleanup()` iterates
+    // every value, so multiple entries per path are all still disposed.
+    const persistKey = persistPagedCache === undefined ? 'u' : persistPagedCache ? 't' : 'f';
+    const cacheKey = `${sourcePath}\0${modelType}\0${persistKey}`;
+    const existing = this.overrides.get(cacheKey);
+    if (existing !== undefined) return existing;
+
+    const pending = this.createOverride(sourcePath, config, cacheFloorMb, initialMb, modelType, persistPagedCache);
+    this.overrides.set(cacheKey, pending);
+    try {
+      return await pending;
+    } catch (error) {
+      this.overrides.delete(cacheKey);
+      throw error;
+    }
+  }
+
+  /** Remove this manager's temporary root without affecting other managers. */
+  cleanup(): Promise<void> {
+    if (this.cleanupPromise !== undefined) return this.cleanupPromise;
+    this.disposed = true;
+    this.cleanupPromise = this.performCleanup();
+    return this.cleanupPromise;
+  }
+
+  private async performCleanup(): Promise<void> {
+    await Promise.allSettled(this.activeResolves);
+    this.activeResolves.clear();
+    await Promise.allSettled(this.overrides.values());
+    this.overrides.clear();
+    if (this.rootPromise === undefined) return;
+
+    const root = await this.rootPromise.catch(() => undefined);
+    if (root !== undefined) {
+      await rm(root, { recursive: true, force: true }).catch(() => undefined);
+    }
+  }
+
+  private async createOverride(
+    sourcePath: string,
+    sourceConfig: Record<string, unknown>,
+    cacheFloorMb: number | undefined,
+    initialMb: number | undefined,
+    modelType: string,
+    persistPagedCache: boolean | undefined,
+  ): Promise<string> {
+    const root = await this.getRoot();
+    const overrideDir = await mkdtemp(join(root, 'model-'));
+
+    const config: Record<string, unknown> = {
+      ...sourceConfig,
+      use_block_paged_cache: true,
+    };
+    if (persistPagedCache !== undefined) {
+      // Authoritative: the loader reads snake_case, so that spelling is the one
+      // that decides persistence. Reconcile a stray camelCase alias spread from
+      // the source config so it can never contradict the authoritative value.
+      config.persist_paged_cache = persistPagedCache;
+      if ('persistPagedCache' in config) config.persistPagedCache = persistPagedCache;
+    }
+    if (cacheFloorMb !== undefined) {
+      config.paged_cache_memory_mb = Math.max(positiveNumber(sourceConfig.paged_cache_memory_mb) ?? 0, cacheFloorMb);
+    }
+    if (initialMb !== undefined) {
+      // Authoritative initial (grow-on-demand) budget, clamped to the max just
+      // written above so it can never exceed the pool ceiling. `initialMb` is
+      // always >= 1 and the max is floored by `cacheFloorMb` (>= 1) for the
+      // same families, so this never writes 0 — which the native loader rejects.
+      const maxMb = positiveNumber(config.paged_cache_memory_mb) ?? 0;
+      config.paged_cache_initial_memory_mb = Math.min(initialMb, maxMb);
+    }
+    await writeFile(join(overrideDir, 'config.json'), JSON.stringify(config, null, 2), 'utf-8');
+
+    const sourceEntries = await readdir(sourcePath);
+    for (const name of sourceEntries) {
+      if (name === 'config.json') continue;
+      const source = join(sourcePath, name);
+      const destination = join(overrideDir, name);
+
+      let isFile: boolean;
+      try {
+        isFile = (await stat(source)).isFile();
+      } catch {
+        continue;
+      }
+      if (!isFile) continue;
+
+      try {
+        await symlink(source, destination);
+      } catch (error) {
+        if ((error as NodeJS.ErrnoException).code !== 'EEXIST') throw error;
+      }
+    }
+
+    if (QWEN35_CACHE_FLOOR_MODEL_TYPES.has(modelType)) {
+      await preserveQwen35MtpSidecars(sourcePath, overrideDir, sourceConfig, modelType);
+    }
+
+    return overrideDir;
+  }
+
+  private getRoot(): Promise<string> {
+    this.rootPromise ??= mkdtemp(join(tmpdir(), this.tempDirPrefix));
+    return this.rootPromise;
+  }
+}
+
+/**
+ * Keep only the nested paths that the native Qwen3.5 loaders inspect.
+ *
+ * Both dense and MoE loaders accept an mlx-vlm `mtp-drafter/` directory.
+ * Dense additionally accepts a config-declared relative sidecar and the
+ * conventional `mtp/weights.safetensors` fallback. Top-level sidecar files
+ * are already handled by the normal source-file loop above.
+ */
+async function preserveQwen35MtpSidecars(
+  sourcePath: string,
+  overrideDir: string,
+  sourceConfig: Record<string, unknown>,
+  modelType: string,
+): Promise<void> {
+  await symlinkDirectoryIfPresent(join(sourcePath, QWEN35_MTP_DRAFTER_DIR), join(overrideDir, QWEN35_MTP_DRAFTER_DIR));
+
+  // The MoE loader supports the split drafter directory, but deliberately has
+  // no `mtp.safetensors`-style sidecar discovery path.
+  if (modelType !== 'qwen3_5') return;
+
+  const relativeSidecars = new Set<string>([QWEN35_DENSE_NESTED_MTP_SIDECAR]);
+  const configuredSidecar = configuredQwen35MtpFile(sourceConfig);
+  if (configuredSidecar !== undefined && !isAbsolute(configuredSidecar)) {
+    relativeSidecars.add(configuredSidecar);
+  }
+
+  for (const sidecar of relativeSidecars) {
+    await symlinkContainedFileIfPresent(sourcePath, overrideDir, sidecar);
+  }
+}
+
+function configuredQwen35MtpFile(config: Record<string, unknown>): string | undefined {
+  const extra = config.mlx_lm_extra_tensors;
+  if (extra === null || typeof extra !== 'object' || Array.isArray(extra)) return undefined;
+  const mtpFile = (extra as Record<string, unknown>).mtp_file;
+  return typeof mtpFile === 'string' && mtpFile.trim() !== '' ? mtpFile : undefined;
+}
+
+async function symlinkContainedFileIfPresent(
+  sourceRoot: string,
+  destinationRoot: string,
+  relativePath: string,
+): Promise<void> {
+  const source = resolve(sourceRoot, relativePath);
+  const normalizedRelative = relative(sourceRoot, source);
+  if (
+    normalizedRelative === '' ||
+    normalizedRelative === '..' ||
+    normalizedRelative.startsWith(`..${sep}`) ||
+    isAbsolute(normalizedRelative)
+  ) {
+    return;
+  }
+
+  try {
+    if (!(await stat(source)).isFile()) return;
+  } catch {
+    return;
+  }
+
+  const destination = join(destinationRoot, normalizedRelative);
+  await mkdir(dirname(destination), { recursive: true });
+  await symlinkIfMissing(source, destination);
+}
+
+async function symlinkDirectoryIfPresent(source: string, destination: string): Promise<void> {
+  if (!(await isDirectory(source))) return;
+  await symlinkIfMissing(source, destination);
+}
+
+async function symlinkIfMissing(source: string, destination: string): Promise<void> {
+  try {
+    await symlink(source, destination);
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code !== 'EEXIST') throw error;
+  }
+}
+
+async function isDirectory(path: string): Promise<boolean> {
+  try {
+    return (await stat(path)).isDirectory();
+  } catch {
+    return false;
+  }
+}
+
+function positiveNumber(value: unknown): number | undefined {
+  return typeof value === 'number' && Number.isFinite(value) && value > 0 ? value : undefined;
+}
+
+function resolveQwen35CacheFloorMb(): number {
+  const raw = process.env.MLX_PAGED_CACHE_MEMORY_MB;
+  if (raw == null || raw === '') return DEFAULT_QWEN35_PAGED_CACHE_MB;
+  const parsed = Number.parseInt(raw, 10);
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_QWEN35_PAGED_CACHE_MB;
+}
+
+function resolveQwen35InitialMb(): number {
+  const raw = process.env.MLX_PAGED_CACHE_INITIAL_MB;
+  if (raw == null || raw === '') return DEFAULT_QWEN35_PAGED_CACHE_INITIAL_MB;
+  const parsed = Number.parseInt(raw, 10);
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_QWEN35_PAGED_CACHE_INITIAL_MB;
+}

package/src/models/qwen3-configs.ts

@@ -0,0 +1,113 @@
+/**
+ * Qwen3 Model Configurations and Type Definitions
+ *
+ * This module provides:
+ * - Default configurations for common Qwen3 model sizes
+ * - Type re-exports from Rust with enhanced documentation
+ * - Helper functions for config management
+ */
+
+import type {
+  Qwen3Config as RustQwen3Config,
+  GenerationConfig as RustGenerationConfig,
+  GenerationResult as RustGenerationResult,
+} from '@mlx-node/core';
+
+/**
+ * Configuration for Qwen3 models
+ *
+ * All fields are required when creating a model directly.
+ * Use QWEN3_CONFIGS for pre-configured model sizes.
+ */
+export type Qwen3Config = RustQwen3Config;
+
+/**
+ * Configuration for text generation
+ *
+ * Controls sampling behavior, temperature, and stopping criteria.
+ */
+export type GenerationConfig = RustGenerationConfig;
+
+/**
+ * Result from text generation with detailed metadata
+ *
+ * Includes generated tokens, log probabilities, finish reason, and token count.
+ */
+export type GenerationResult = RustGenerationResult;
+
+/**
+ * Default configurations for common Qwen3 models
+ *
+ * Includes optimized hyperparameters for:
+ * - qwen3-0.6b: Smallest model (1024 hidden size, 28 layers)
+ * - qwen3-1.8b: Medium model (1536 hidden size, 28 layers)
+ * - qwen3-7b: Large model (3072 hidden size, 32 layers)
+ */
+export const QWEN3_CONFIGS: { [key: string]: Qwen3Config } = {
+  'qwen3-0.6b': {
+    vocabSize: 151936,
+    hiddenSize: 1024,
+    numLayers: 28,
+    numHeads: 16,
+    numKvHeads: 8, // GQA with 2:1 ratio
+    headDim: 64, // hiddenSize / numHeads = 1024 / 16 = 64
+    intermediateSize: 3072,
+    rmsNormEps: 1e-6,
+    ropeTheta: 1000000.0,
+    maxPositionEmbeddings: 40960,
+    useQkNorm: true, // Qwen3 always uses QK normalization (core feature)
+    tieWordEmbeddings: true,
+    padTokenId: 151643,
+    eosTokenId: 151645,
+    bosTokenId: 151643,
+  },
+  'qwen3-1.8b': {
+    vocabSize: 151936,
+    hiddenSize: 1536,
+    numLayers: 28,
+    numHeads: 12,
+    numKvHeads: 2, // GQA with 6:1 ratio
+    headDim: 128, // hiddenSize / numHeads = 1536 / 12 = 128
+    intermediateSize: 8960,
+    rmsNormEps: 1e-6,
+    ropeTheta: 1000000.0,
+    maxPositionEmbeddings: 131072,
+    useQkNorm: true, // Qwen3 always uses QK normalization (core feature)
+    tieWordEmbeddings: false,
+    padTokenId: 151643,
+    eosTokenId: 151645,
+    bosTokenId: 151643,
+  },
+  'qwen3-7b': {
+    vocabSize: 151936,
+    hiddenSize: 3072,
+    numLayers: 32,
+    numHeads: 24,
+    numKvHeads: 4, // GQA with 6:1 ratio
+    headDim: 128, // hiddenSize / numHeads = 3072 / 24 = 128
+    intermediateSize: 18944,
+    rmsNormEps: 1e-6,
+    ropeTheta: 1000000.0,
+    maxPositionEmbeddings: 131072,
+    useQkNorm: true, // Qwen3 always uses QK normalization (core feature)
+    tieWordEmbeddings: false,
+    padTokenId: 151643,
+    eosTokenId: 151645,
+    bosTokenId: 151643,
+  },
+};
+
+/**
+ * Get a Qwen3 configuration by name
+ *
+ * @param name - Model name (e.g., "qwen3-0.6b", "qwen3-1.8b", "qwen3-7b")
+ * @returns Model configuration
+ * @throws Error if model name is not recognized
+ */
+export function getQwen3Config(name: string): Qwen3Config {
+  const config = QWEN3_CONFIGS[name];
+  if (!config) {
+    throw new Error(`Unknown model configuration: ${name}. Available models: ${Object.keys(QWEN3_CONFIGS).join(', ')}`);
+  }
+  return config;
+}

package/src/models/qwen3_5-configs.ts

@@ -0,0 +1,60 @@
+/**
+ * Qwen3.5 Model Configurations and Type Definitions
+ *
+ * Supports both dense and MoE variants. MoE fields are optional -
+ * when `numExperts` is undefined, the model uses dense MLP layers.
+ */
+
+import type { Qwen35Config as RustQwen35Config } from '@mlx-node/core';
+
+export type Qwen35Config = RustQwen35Config;
+
+/**
+ * Default configurations for common Qwen3.5 models
+ */
+export const QWEN35_CONFIGS: { [key: string]: Qwen35Config } = {
+  'qwen3.5-0.6b': {
+    vocabSize: 151936,
+    hiddenSize: 1024,
+    numLayers: 28,
+    numHeads: 16,
+    numKvHeads: 8,
+    intermediateSize: 3072,
+    rmsNormEps: 1e-6,
+    headDim: 64,
+    tieWordEmbeddings: true,
+    attentionBias: false,
+    maxPositionEmbeddings: 131072,
+    padTokenId: 151643,
+    eosTokenId: 151645,
+    bosTokenId: 151643,
+    linearNumValueHeads: 64,
+    linearNumKeyHeads: 16,
+    linearKeyHeadDim: 192,
+    linearValueHeadDim: 128,
+    linearConvKernelDim: 4,
+    fullAttentionInterval: 4,
+    partialRotaryFactor: 0.25,
+    ropeTheta: 100000.0,
+    // W1 (MTP): no MTP head in the stock 0.6B dense checkpoint.
+    // Real values are populated from `config.json` at load time; this
+    // preset is only used by tests that build a `Qwen35Config` from
+    // scratch.
+    nMtpLayers: 0,
+  },
+};
+
+/**
+ * Get a Qwen3.5 configuration by name
+ *
+ * @param name - Model name (e.g., "qwen3.5-0.6b")
+ * @returns Model configuration
+ * @throws Error if model name is not recognized
+ */
+export function getQwen35Config(name: string): Qwen35Config {
+  const config = QWEN35_CONFIGS[name];
+  if (!config) {
+    throw new Error(`Unknown Qwen3.5 config: ${name}. Available: ${Object.keys(QWEN35_CONFIGS).join(', ')}`);
+  }
+  return config;
+}

package/src/profiling.ts

@@ -0,0 +1,69 @@
+import { writeFile } from 'node:fs/promises';
+
+import { setProfilingEnabled, isProfilingEnabled, getProfilingData, resetProfilingData } from '@mlx-node/core';
+
+const ENV_VAR = 'MLX_PROFILE_DECODE';
+const envVarSet = !!process.env[ENV_VAR];
+let exitHandlerRegistered = false;
+
+// Auto-enable if env var set
+if (envVarSet) {
+  setProfilingEnabled(true);
+  registerExitHandler();
+}
+
+/**
+ * Enable profiling programmatically.
+ *
+ * When enabled, all subsequent model generate/chat calls will record
+ * timing, memory, and throughput data. Call `disableProfiling()` to
+ * stop recording and write the report.
+ *
+ * If `MLX_PROFILE_DECODE` env var is set, this is a no-op (env var
+ * takes precedence).
+ */
+export function enableProfiling(): void {
+  if (envVarSet) {
+    console.warn(`Warning: env var ${ENV_VAR} is set, ignoring explicit profiling API calls`);
+    return;
+  }
+  setProfilingEnabled(true);
+  resetProfilingData();
+  registerExitHandler();
+}
+
+/**
+ * Disable profiling and write the collected data to a JSON file.
+ *
+ * Returns the path to the written file, or empty string if no data
+ * was collected. If `MLX_PROFILE_DECODE` env var is set, this is a
+ * no-op (env var controls the lifecycle).
+ */
+export async function disableProfiling(): Promise<string> {
+  if (envVarSet) {
+    console.warn(`Warning: env var ${ENV_VAR} is set, ignoring explicit profiling API calls`);
+    return '';
+  }
+  setProfilingEnabled(false);
+  return writeProfilingReport();
+}
+
+async function writeProfilingReport(): Promise<string> {
+  const data = getProfilingData();
+  if (data.generations.length === 0) return '';
+  const path = `mlx-profile-${Date.now()}.json`;
+  await writeFile(path, JSON.stringify(data, null, 2));
+  console.info(`Profiling report written to ${path}`);
+  return path;
+}
+
+function registerExitHandler(): void {
+  if (exitHandlerRegistered) return;
+  exitHandlerRegistered = true;
+  process.on('beforeExit', async () => {
+    if (isProfilingEnabled()) {
+      setProfilingEnabled(false);
+      await writeProfilingReport();
+    }
+  });
+}