mlx-bun 0.0.11 → 0.0.13

package/src/expert-usage.ts

@@ -0,0 +1,415 @@
+import {
+  closeSync,
+  existsSync,
+  fsyncSync,
+  openSync,
+  readFileSync,
+  renameSync,
+  rmSync,
+  writeSync,
+} from "node:fs";
+
+const UINT32_MAX = 0xffff_ffff;
+
+export interface ExpertUsageRoute {
+  readonly indices: ArrayLike<number>;
+}
+
+export interface ExpertUsageLedgerOptions {
+  readonly path: string;
+  readonly layers: readonly number[];
+  readonly expertsPerLayer: number;
+  /** A damaged profile is derived state: report it and start fresh. */
+  readonly onWarning?: (message: string) => void;
+}
+
+export interface ExpertUsageEntry {
+  readonly count: number;
+  readonly heat: number;
+  readonly lastAccess: number;
+}
+
+export interface ExpertUsageSnapshot {
+  readonly clock: number;
+  readonly routeSelections: number;
+  readonly totalCount: number;
+  readonly nonzeroCounts: number;
+  readonly dirty: boolean;
+}
+
+export interface ExpertUsageCount {
+  readonly layer: number;
+  readonly expertId: number;
+  readonly count: number;
+}
+
+export interface ExpertAutoPinPlan {
+  readonly historySelections: number;
+  readonly confidence: number;
+  readonly requestedBudgetBytes: number;
+  readonly usableBudgetBytes: number;
+  readonly plannedBytes: number;
+  readonly pins: readonly ExpertUsageCount[];
+}
+
+export interface ExpertAutoPinInput {
+  readonly ledger: ExpertUsageLedger;
+  /** Bytes available to the persistent LRU + pin tiers after working banks. */
+  readonly residentTierBudgetBytes: number;
+  /** One non-evictable LRU slot per managed layer. */
+  readonly mandatoryResidentBytes: number;
+  readonly slotBytes: (layer: number) => number;
+  readonly minHistorySelections?: number;
+  readonly fullConfidenceSelections?: number;
+  readonly pinShare?: number;
+  readonly minPinBudgetBytes?: number;
+}
+
+export interface ExpertLfruCandidate {
+  readonly layer: number;
+  readonly coldSlot: number;
+  readonly coldExpertId: number;
+  readonly hotExpertId: number;
+  readonly coldScore: number;
+  readonly hotScore: number;
+  readonly gain: number;
+}
+
+export function expertLfruScore(
+  heat: number,
+  lastAccess: number,
+  clock: number,
+): number {
+  const age = (clock - lastAccess) >>> 0;
+  return heat * 256 + (age < 255 ? 255 - age : 0);
+}
+
+export function selectExpertLfruCandidates(
+  candidates: readonly ExpertLfruCandidate[],
+  limit = 4,
+): ExpertLfruCandidate[] {
+  positiveSafeInteger(limit, "live-repin limit");
+  return candidates.slice().sort((left, right) =>
+    right.gain - left.gain ||
+    left.layer - right.layer ||
+    left.hotExpertId - right.hotExpertId)
+    .slice(0, limit);
+}
+
+function positiveSafeInteger(value: number, name: string): void {
+  if (!Number.isSafeInteger(value) || value <= 0)
+    throw new RangeError(`${name} must be a positive safe integer`);
+}
+
+function incrementUint32(value: number, amount: number): number {
+  return Math.min(UINT32_MAX, value + amount);
+}
+
+/**
+ * Colibri-compatible long-term expert frequency plus session-local heat and
+ * recency. The on-disk format is one `layer expert count` row per nonzero
+ * counter, matching Colibri's `.coli_usage` reader.
+ */
+export class ExpertUsageLedger {
+  readonly path: string;
+  readonly layers: readonly number[];
+  readonly expertsPerLayer: number;
+  #counts = new Map<number, Uint32Array>();
+  #heat = new Map<number, Uint32Array>();
+  #last = new Map<number, Uint32Array>();
+  #clock = 0;
+  #routeSelections = 0;
+  #dirty = false;
+
+  private constructor(options: ExpertUsageLedgerOptions) {
+    if (!options.path) throw new Error("expert usage path must not be empty");
+    positiveSafeInteger(options.expertsPerLayer, "expertsPerLayer");
+    if (options.expertsPerLayer > UINT32_MAX)
+      throw new RangeError("expertsPerLayer exceeds uint32 range");
+    const layers = options.layers.slice();
+    if (layers.length === 0 || new Set(layers).size !== layers.length)
+      throw new Error("expert usage layers must be non-empty and unique");
+    for (const layer of layers)
+      if (!Number.isSafeInteger(layer) || layer < 0)
+        throw new RangeError(`invalid expert usage layer ${layer}`);
+    this.path = options.path;
+    this.layers = Object.freeze(layers);
+    this.expertsPerLayer = options.expertsPerLayer;
+    for (const layer of layers) {
+      this.#counts.set(layer, new Uint32Array(options.expertsPerLayer));
+      this.#heat.set(layer, new Uint32Array(options.expertsPerLayer));
+      this.#last.set(layer, new Uint32Array(options.expertsPerLayer));
+    }
+  }
+
+  static open(options: ExpertUsageLedgerOptions): ExpertUsageLedger {
+    const ledger = new ExpertUsageLedger(options);
+    if (!existsSync(options.path)) return ledger;
+    try {
+      ledger.#load(readFileSync(options.path, "utf8"));
+    } catch (error) {
+      const detail = error instanceof Error ? error.message : String(error);
+      options.onWarning?.(
+        `ignoring invalid expert usage profile ${options.path}: ${detail}`,
+      );
+      ledger.#clearCounts();
+    }
+    return ledger;
+  }
+
+  /** Count every selected route before the execution plan deduplicates it. */
+  recordRoutes(layer: number, routes: readonly ExpertUsageRoute[]): void {
+    const counts = this.#counts.get(layer);
+    const heat = this.#heat.get(layer);
+    const last = this.#last.get(layer);
+    if (!counts || !heat || !last)
+      throw new RangeError(`layer ${layer} is not tracked by expert usage`);
+
+    const selected: number[] = [];
+    for (let row = 0; row < routes.length; row++) {
+      for (let rank = 0; rank < routes[row]!.indices.length; rank++) {
+        const expertId = Number(routes[row]!.indices[rank]);
+        if (!Number.isSafeInteger(expertId) || expertId < 0 ||
+            expertId >= this.expertsPerLayer) {
+          throw new RangeError(
+            `route ${row} rank ${rank} has invalid expert ${expertId}`,
+          );
+        }
+        selected.push(expertId);
+      }
+    }
+
+    for (const expertId of selected) {
+      counts[expertId] = incrementUint32(counts[expertId]!, 1);
+      heat[expertId] = incrementUint32(heat[expertId]!, 1);
+      this.#clock = (this.#clock + 1) >>> 0;
+      last[expertId] = this.#clock;
+      this.#routeSelections++;
+    }
+    if (selected.length) this.#dirty = true;
+  }
+
+  entry(layer: number, expertId: number): ExpertUsageEntry {
+    this.#validateEntry(layer, expertId);
+    return {
+      count: this.#counts.get(layer)![expertId]!,
+      heat: this.#heat.get(layer)![expertId]!,
+      lastAccess: this.#last.get(layer)![expertId]!,
+    };
+  }
+
+  snapshot(): ExpertUsageSnapshot {
+    let nonzeroCounts = 0;
+    let totalCount = 0;
+    for (const counts of this.#counts.values())
+      for (const count of counts) {
+        totalCount += count;
+        if (count !== 0) nonzeroCounts++;
+      }
+    return {
+      clock: this.#clock,
+      routeSelections: this.#routeSelections,
+      totalCount,
+      nonzeroCounts,
+      dirty: this.#dirty,
+    };
+  }
+
+  rankedCounts(): ExpertUsageCount[] {
+    const ranked: ExpertUsageCount[] = [];
+    for (const layer of this.layers) {
+      const counts = this.#counts.get(layer)!;
+      for (let expertId = 0; expertId < counts.length; expertId++) {
+        const count = counts[expertId]!;
+        if (count !== 0) ranked.push({ layer, expertId, count });
+      }
+    }
+    ranked.sort((left, right) =>
+      right.count - left.count ||
+      left.layer - right.layer ||
+      left.expertId - right.expertId);
+    return ranked;
+  }
+
+  lfruCandidate(
+    layer: number,
+    pinnedExpertIds: readonly number[],
+  ): ExpertLfruCandidate | null {
+    if (pinnedExpertIds.length === 0) return null;
+    for (const expertId of pinnedExpertIds) this.#validateEntry(layer, expertId);
+    const pinned = new Set(pinnedExpertIds);
+    let coldSlot = 0;
+    let coldScore = this.#lfruScore(layer, pinnedExpertIds[0]!);
+    for (let slot = 1; slot < pinnedExpertIds.length; slot++) {
+      const score = this.#lfruScore(layer, pinnedExpertIds[slot]!);
+      if (score < coldScore) {
+        coldSlot = slot;
+        coldScore = score;
+      }
+    }
+    let hotExpertId = -1;
+    let hotScore = 0;
+    for (let expertId = 0; expertId < this.expertsPerLayer; expertId++) {
+      if (pinned.has(expertId)) continue;
+      const score = this.#lfruScore(layer, expertId);
+      if (hotExpertId < 0 || score > hotScore) {
+        hotExpertId = expertId;
+        hotScore = score;
+      }
+    }
+    if (hotExpertId < 0 ||
+        hotScore <= coldScore + Math.floor(coldScore / 4) + 4 * 256) {
+      return null;
+    }
+    return {
+      layer,
+      coldSlot,
+      coldExpertId: pinnedExpertIds[coldSlot]!,
+      hotExpertId,
+      coldScore,
+      hotScore,
+      gain: Math.floor((hotScore - coldScore) / 256),
+    };
+  }
+
+  decayHeat(): void {
+    for (const heat of this.#heat.values())
+      for (let expertId = 0; expertId < heat.length; expertId++)
+        heat[expertId] = heat[expertId]! >>> 1;
+  }
+
+  /** Persist the current profile with same-directory temp + atomic rename. */
+  flush(): void {
+    if (!this.#dirty) return;
+    const rows: string[] = [];
+    for (const layer of this.layers) {
+      const counts = this.#counts.get(layer)!;
+      for (let expertId = 0; expertId < counts.length; expertId++) {
+        const count = counts[expertId]!;
+        if (count !== 0) rows.push(`${layer} ${expertId} ${count}\n`);
+      }
+    }
+    const temp =
+      `${this.path}.tmp-${process.pid}-${Math.trunc(Bun.nanoseconds())}`;
+    let fd: number | null = null;
+    try {
+      fd = openSync(temp, "w", 0o600);
+      const data = Buffer.from(rows.join(""));
+      let written = 0;
+      while (written < data.length) {
+        written += writeSync(
+          fd,
+          data,
+          written,
+          data.length - written,
+          written,
+        );
+      }
+      fsyncSync(fd);
+      closeSync(fd);
+      fd = null;
+      renameSync(temp, this.path);
+      this.#dirty = false;
+    } catch (error) {
+      if (fd !== null) closeSync(fd);
+      try { rmSync(temp, { force: true }); } catch {}
+      throw error;
+    }
+  }
+
+  #load(text: string): void {
+    for (const [index, raw] of text.split("\n").entries()) {
+      const line = raw.trim();
+      if (!line) continue;
+      const fields = line.split(/\s+/);
+      if (fields.length !== 3)
+        throw new Error(`line ${index + 1}: expected layer expert count`);
+      const [layer, expertId, count] = fields.map(Number);
+      if (!Number.isSafeInteger(layer) || !this.#counts.has(layer!))
+        throw new Error(`line ${index + 1}: invalid layer ${fields[0]}`);
+      if (!Number.isSafeInteger(expertId) || expertId! < 0 ||
+          expertId! >= this.expertsPerLayer) {
+        throw new Error(`line ${index + 1}: invalid expert ${fields[1]}`);
+      }
+      if (!Number.isSafeInteger(count) || count! < 0 || count! > UINT32_MAX)
+        throw new Error(`line ${index + 1}: invalid count ${fields[2]}`);
+      const counts = this.#counts.get(layer!)!;
+      counts[expertId!] = incrementUint32(counts[expertId!]!, count!);
+    }
+  }
+
+  #clearCounts(): void {
+    for (const counts of this.#counts.values()) counts.fill(0);
+  }
+
+  #validateEntry(layer: number, expertId: number): void {
+    if (!this.#counts.has(layer))
+      throw new RangeError(`layer ${layer} is not tracked by expert usage`);
+    if (!Number.isSafeInteger(expertId) || expertId < 0 ||
+        expertId >= this.expertsPerLayer) {
+      throw new RangeError(`invalid expert ${expertId}`);
+    }
+  }
+
+  #lfruScore(layer: number, expertId: number): number {
+    const heat = this.#heat.get(layer)![expertId]!;
+    const last = this.#last.get(layer)![expertId]!;
+    return expertLfruScore(heat, last, this.#clock);
+  }
+}
+
+/**
+ * Faithful, budget-safe form of Colibri's startup AUTOPIN policy. History must
+ * reach 5k selections; confidence ramps to one at 200k; pins receive at most
+ * half of the resident-tier budget and are not created below a 0.5 GB tier.
+ */
+export function planExpertAutoPins(input: ExpertAutoPinInput): ExpertAutoPinPlan {
+  const minHistory = input.minHistorySelections ?? 5_000;
+  const fullConfidence = input.fullConfidenceSelections ?? 200_000;
+  const pinShare = input.pinShare ?? 0.5;
+  const minPinBudget = input.minPinBudgetBytes ?? 500_000_000;
+  positiveSafeInteger(minHistory, "auto-pin minHistorySelections");
+  positiveSafeInteger(fullConfidence, "auto-pin fullConfidenceSelections");
+  if (!Number.isFinite(pinShare) || pinShare <= 0 || pinShare > 1)
+    throw new RangeError("auto-pin pinShare must be in (0, 1]");
+  if (!Number.isSafeInteger(minPinBudget) || minPinBudget < 0)
+    throw new RangeError("auto-pin minPinBudgetBytes must be non-negative");
+  if (!Number.isSafeInteger(input.residentTierBudgetBytes) ||
+      input.residentTierBudgetBytes < 0 ||
+      !Number.isSafeInteger(input.mandatoryResidentBytes) ||
+      input.mandatoryResidentBytes < 0) {
+    throw new RangeError("auto-pin resident budgets must be non-negative safe integers");
+  }
+
+  const historySelections = input.ledger.snapshot().totalCount;
+  const confidence = Math.min(historySelections / fullConfidence, 1);
+  const requestedBudgetBytes = Math.floor(
+    input.residentTierBudgetBytes * pinShare * confidence,
+  );
+  const usableBudgetBytes = Math.min(
+    requestedBudgetBytes,
+    Math.max(
+      0,
+      input.residentTierBudgetBytes - input.mandatoryResidentBytes,
+    ),
+  );
+  const pins: ExpertUsageCount[] = [];
+  let plannedBytes = 0;
+  if (historySelections >= minHistory && requestedBudgetBytes >= minPinBudget) {
+    for (const candidate of input.ledger.rankedCounts()) {
+      const bytes = input.slotBytes(candidate.layer);
+      positiveSafeInteger(bytes, `auto-pin slotBytes for layer ${candidate.layer}`);
+      if (plannedBytes + bytes > usableBudgetBytes) continue;
+      pins.push(candidate);
+      plannedBytes += bytes;
+    }
+  }
+  return {
+    historySelections,
+    confidence,
+    requestedBudgetBytes,
+    usableBudgetBytes,
+    plannedBytes,
+    pins: Object.freeze(pins),
+  };
+}

package/src/generate.ts

@@ -117,6 +117,15 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
   grammar?: GrammarController;
 }
 
+export function shouldUseGrammarJump(
+  options: Pick<GenerateOptions, "grammar" | "logprobs" | "topLogprobs">,
+): boolean {
+  return options.grammar !== undefined &&
+    flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
+    !options.logprobs &&
+    !(options.topLogprobs && options.topLogprobs > 0);
+}
+
 /** Per-token logprob info (only present when GenerateOptions requested it). */
 export interface TokenLogprobs {
   /** logprob of the emitted token (requires options.logprobs). */
@@ -282,6 +291,8 @@ export interface GenerateStats {
   spec?: {
     drafted: number; accepted: number; targetCalls: number;
     draftedByPos?: number[]; acceptedByPos?: number[];
+    rejected?: number; rounds?: number; acceptanceLengths?: number[];
+    tokensPerForward?: number; forwardsSaved?: number;
   };
 }
 
@@ -338,6 +349,33 @@ export class Generation implements AsyncIterable<GeneratedToken> {
 const WIRE_THRESHOLD = 0.75;
 let wiredScopeDepth = 0;
 let wiredOldLimit = 0;
+
+type WiredModelMemory = {
+  readonly weightsBytes: number;
+  readonly expertRuntime?: {
+    readonly plan: { readonly plannedBytes: number };
+    flushUsage?: () => void;
+    finishUsage?: () => Promise<void>;
+  } | null;
+};
+
+/** Bytes whose MLX graph must stay resident while a model executes. */
+export function wiredWorkingSetBytes(model: WiredModelMemory): number {
+  const planned = model.expertRuntime?.plan.plannedBytes;
+  return typeof planned === "number" && Number.isFinite(planned) && planned > 0
+    ? Math.max(model.weightsBytes, planned)
+    : model.weightsBytes;
+}
+
+export function modelNeedsWiredLimit(
+  model: WiredModelMemory,
+  recommendedBytes = maxRecommendedWorkingSetSize(),
+  force = process.env.MLX_BUN_FORCE_WIRE === "1",
+): boolean {
+  return force ||
+    wiredWorkingSetBytes(model) > WIRE_THRESHOLD * recommendedBytes;
+}
+
 function enterWiredScope(): void {
   if (wiredScopeDepth++ === 0)
     wiredOldLimit = setWiredLimit(maxRecommendedWorkingSetSize());
@@ -362,9 +400,10 @@ export function generate(
       ? generateDiffusionInner(model, promptTokens, options)
       : generateInner(model, promptTokens, options);
   if (options.adapters?.length) inner = adapterScoped(model, options.adapters, inner);
-  const wire =
-    process.env.MLX_BUN_FORCE_WIRE === "1" ||
-    model.weightsBytes > WIRE_THRESHOLD * maxRecommendedWorkingSetSize();
+  const memoryModel: WiredModelMemory = model;
+  if (memoryModel.expertRuntime?.flushUsage)
+    inner = usageScoped(memoryModel, inner);
+  const wire = modelNeedsWiredLimit(model);
   return new Generation(wire ? wiredScoped(inner) : inner);
 }
 
@@ -437,6 +476,68 @@ async function* wiredScoped(
   }
 }
 
+/** Publish the streamed model's route ledger on every generator exit path. */
+async function* usageScoped(
+  model: WiredModelMemory,
+  inner: AsyncGenerator<GeneratedToken, GenerateStats>,
+): AsyncGenerator<GeneratedToken, GenerateStats> {
+  try {
+    return yield* inner;
+  } finally {
+    if (model.expertRuntime?.finishUsage)
+      await model.expertRuntime.finishUsage();
+    else
+      model.expertRuntime?.flushUsage?.();
+  }
+}
+
+/** Apply generate()'s scoped wiring policy to non-generator execution paths. */
+export async function withModelWiredLimit<T>(
+  model: WiredModelMemory,
+  run: () => Promise<T>,
+): Promise<T> {
+  if (!modelNeedsWiredLimit(model)) return run();
+  enterWiredScope();
+  try {
+    return await run();
+  } finally {
+    exitWiredScope();
+  }
+}
+
+/** Apply generate()'s usage safe-point to direct/non-generator paths. */
+export async function withModelUsageFlush<T>(
+  model: WiredModelMemory,
+  run: () => Promise<T>,
+): Promise<T> {
+  try {
+    return await run();
+  } finally {
+    if (model.expertRuntime?.finishUsage)
+      await model.expertRuntime.finishUsage();
+    else
+      model.expertRuntime?.flushUsage?.();
+  }
+}
+
+/** Forward through either the ordinary synchronous model surface or a
+ * streamed-expert model whose layer execution must cross async I/O. Keeping
+ * this choice inside the shared generator preserves the complete generation
+ * contract (sampling, grammar, logprobs, prompt cache, and cancellation) for
+ * both execution styles. */
+async function forwardHiddenForGeneration(
+  model: RuntimeModel,
+  ids: MlxArray,
+  cache: Cache[],
+): Promise<MlxArray> {
+  const asyncModel = model as RuntimeModel & {
+    forwardHiddenAsync?: (ids: MlxArray, cache: Cache[]) => Promise<MlxArray>;
+  };
+  return typeof asyncModel.forwardHiddenAsync === "function"
+    ? await asyncModel.forwardHiddenAsync(ids, cache)
+    : model.forwardHidden(ids, cache);
+}
+
 async function* generateInner(
   model: RuntimeModel,
   promptTokens: number[],
@@ -488,15 +589,14 @@ async function* generateInner(
     if (e) for (const a of [e.sel, e.topIdx, e.topVals]) a?.dispose();
   };
   /** Read extras back to JS (forces eval — they were async-dispatched with the
-   *  token, so this normally just copies) and dispose the device arrays. */
+   *  token) without appending casts behind the next decode step. */
   const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
     if (!e) return undefined;
     const out: TokenLogprobs = {};
-    if (e.sel) out.logprob = e.sel.toFloat32()[0]!;
+    if (e.sel) out.logprob = e.sel.toFloat32Host()[0]!;
     if (e.topIdx && e.topVals) {
-      // uint32 ids read via the f32 cast: exact for ids < 2^24 (vocab ≪ 16M)
-      const ids = e.topIdx.toFloat32();
-      const vals = e.topVals.toFloat32();
+      const ids = e.topIdx.toIntTokens();
+      const vals = e.topVals.toFloat32Host();
       out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! })).sort(
         (a, b) => b.logprob - a.logprob,
       );
@@ -616,7 +716,7 @@ async function* generateInner(
         while (pos < snapAt) {
           const chunk = promptTokens.slice(pos, Math.min(pos + prefillChunkSize, snapAt));
           const ids = ops.fromInt32(chunk, [1, chunk.length]);
-          const h = model.forwardHidden(ids, cache);
+          const h = await forwardHiddenForGeneration(model, ids, cache);
           ids.dispose();
           h.dispose();
           evalCacheState(cache);
@@ -650,7 +750,7 @@ async function* generateInner(
           : prefillChunkSize;
         const chunk = promptTokens.slice(pos, pos + n);
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const h = model.forwardHidden(ids, cache);
+        const h = await forwardHiddenForGeneration(model, ids, cache);
         ids.dispose();
         h.dispose(); // logits never computed for non-final chunks
         evalCacheState(cache);
@@ -677,7 +777,7 @@ async function* generateInner(
       // cover exactly the prompt, as before.
       const lastChunk = promptTokens.slice(pos);
       const ids0 = ops.fromInt32(lastChunk, [1, lastChunk.length]);
-      h0 = model.forwardHidden(ids0, cache);
+      h0 = await forwardHiddenForGeneration(model, ids0, cache);
       ids0.dispose();
     }
     if (options.snapshotAt === undefined || options.snapshotAt >= promptTokens.length)
@@ -736,10 +836,7 @@ async function* generateInner(
     // (see GrammarController.jumpForward for the contract + the fidelity
     // note on why this is opt-in). Excluded when logprobs are requested
     // (jumped tokens are never sampled, so they'd have no logprobs rows).
-    const grammarJump =
-      options.grammar !== undefined &&
-      flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
-      !options.logprobs;
+    const grammarJump = shouldUseGrammarJump(options);
     while (!stop) {
       const cur = pending!;
       const curExtras = pendingExtras;
@@ -798,7 +895,7 @@ async function* generateInner(
         }
         const chunk = [grammarTok, ...jumpEmit];
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const h = model.forwardHidden(ids, cache);
+        const h = await forwardHiddenForGeneration(model, ids, cache);
         ids.dispose();
         // Every chunk token's KV is in the cache regardless of what follows.
         forwarded.push(...chunk);
@@ -839,7 +936,7 @@ async function* generateInner(
         }
         if (!logits) {
           const ids = ops.reshape(cur, [1, 1]);
-          const h = model.forwardHidden(ids, cache);
+          const h = await forwardHiddenForGeneration(model, ids, cache);
           ids.dispose();
           logits = model.logitsFromHidden(h);
           h.dispose();

package/src/hf-push.ts

@@ -420,7 +420,7 @@ async function uploadLfsFiles(
     const putRes = await fetch(upload.href, {
       method: "PUT",
       headers: { ...(upload.header ?? {}) },
-      body: item.bytes,
+      body: new Blob([Uint8Array.from(item.bytes)]),
     });
     if (!putRes.ok) throw await httpError(`LFS upload of ${item.file.repoPath}`, putRes);
     onProgress?.(item.file.repoPath, item.file.size, item.file.size);

package/src/index.ts

@@ -12,7 +12,13 @@ export { loadModelConfig } from "./config";
 export { Weights } from "./weights";
 export { Gemma4Model } from "./model/gemma4";
 export { MiniCPM5Model } from "./model/minicpm5";
-export { createModel, type RuntimeModel } from "./model/factory";
+export {
+  createModel,
+  openGlm52RuntimeModel,
+  openModel,
+  type Glm52RuntimeOpenOptions,
+  type RuntimeModel,
+} from "./model/factory";
 export { Qwen3Model } from "./model/qwen3";
 export { generate } from "./generate";
 export { loadTokenizer } from "./tokenizer";
@@ -21,4 +27,10 @@ export { ChatTemplate } from "./chat-template";
 export { chooseAutoModel, COEXIST_FRACTION, DEFAULT_REPO_ID, fit, largestRecommendedRepoId, recommendedRepoId, skuMatrix, thisMachine } from "./fit";
 export { downloadModel } from "./download";
 export { Registry } from "./registry";
-export { createServer, loadContext } from "./server";
+export {
+  createServer,
+  loadContext,
+  type LoadContextOptions,
+  type ServerContext,
+  type ServerOptions,
+} from "./server";