mlx-bun 0.2.0 → 0.4.0

package/src/serve/completion-executor.ts

@@ -27,9 +27,10 @@ import {
   type RequestRejection,
 } from "./request-plan";
 import type { PromptResponseTrace } from "./prompt-response-trace";
+import { cleanupFailure } from "../engine/resources";
 
 export interface CompletionEngine {
-  place(shape: RequestShape): GenerationPlacement;
+  place(shape: RequestShape, options?: GenerateOptions): GenerationPlacement;
   run(
     promptIds: number[],
     options: GenerateOptions & { stopSequences?: string[] },
@@ -118,6 +119,8 @@ export interface CompletionUsage {
   completionTokens: number;
   totalTokens: number;
   speculation?: GenerateStats["spec"];
+  /** Token fast-forwarding telemetry (serial lane, MLX_BUN_FILL=strict). */
+  fill?: GenerateStats["fill"];
 }
 
 export interface CompletionSummary {
@@ -199,7 +202,7 @@ function finalLane(
   stats: GenerateStats,
 ): Lane {
   if (mechanism === "continuous") return "batched";
-  return stats.spec ? "serial+spec" : lane;
+  return stats.spec ? "serial+spec" : "serial";
 }
 
 function finishReason(
@@ -209,7 +212,7 @@ function finishReason(
 ): CompletionFinishReason {
   if (result.toolCalls.length) return "tool_calls";
   if (result.stopped) return "stop";
-  return stats.generatedTokens >= maxTokens ? "length" : "stop";
+  return stats.finishReason ?? (stats.generatedTokens >= maxTokens ? "length" : "stop");
 }
 
 function combineControl(
@@ -251,12 +254,14 @@ export class CompletionExecutor {
     try {
       if (usageProgress) control.onUsageProgress!(usageProgress);
       const closePlacement = control.trace?.begin("completion.placement");
-      const enginePlacement = this.engine.place(planned.shape);
+      const enginePlacement = this.engine.place(planned.shape, planned.options);
       closePlacement?.();
       if (enginePlacement.shape !== planned.shape)
         throw new Error("generation placement does not belong to this request shape");
       const mechanism = enginePlacement.mechanism;
-      const lane = initialLane(mechanism, planned.shape);
+      const lane = enginePlacement.execution
+        ? (mechanism === "continuous" ? "batched" : enginePlacement.execution.method === "speculative" ? "serial+spec" : "serial")
+        : initialLane(mechanism, planned.shape);
       recordLane(input.requestId, lane);
       const placement = { mechanism, lane, shape: planned.shape };
       input.onPlacement?.(placement);
@@ -314,6 +319,8 @@ export class CompletionExecutor {
         usageProgress.totalTokens = stats.promptTokens + stats.generatedTokens;
         if (stats.spec) usageProgress.speculation = stats.spec;
         else delete usageProgress.speculation;
+        if (stats.fill) usageProgress.fill = stats.fill;
+        else delete usageProgress.fill;
       }
       const summary = {
         content: result.content,
@@ -332,15 +339,16 @@ export class CompletionExecutor {
           completionTokens: stats.generatedTokens,
           totalTokens: stats.promptTokens + stats.generatedTokens,
           ...(stats.spec ? { speculation: stats.spec } : {}),
+          ...(stats.fill ? { fill: stats.fill } : {}),
         },
         logprobs: logprobs?.payload() ?? null,
       };
+      planned.dispose(); // retained preparation reservations outlive the native run
       closeCompletion?.();
       return summary;
     } catch (error) {
       closeCompletion?.();
-      planned.dispose();
-      throw error;
+      cleanupFailure(error, () => planned.dispose());
     }
   }
 }

package/src/serve/discovery-routes.ts

@@ -1,9 +1,8 @@
+import { modelServingBinding } from "../backends/mlx/model-serving";
 import pkgJson from "../../package.json" with { type: "json" };
-import type { ServerContext } from "../server";
+import type { ServingContext } from "./model-host";
 import type { GenerationGateway } from "./generation-gateway";
 import { fit } from "../fit";
-import { isEmbeddingModel } from "../embed";
-import { Glm52Model } from "../model/glm52";
 
 const pkgVersion = (pkgJson as { version: string }).version;
 
@@ -32,7 +31,7 @@ export interface DiscoveryRoutes {
 }
 
 export function createDiscoveryRoutes(
-  ctx: ServerContext,
+  ctx: ServingContext,
   gateway: Pick<GenerationGateway, "batchMode">,
   startedAt: number,
 ): DiscoveryRoutes {
@@ -46,7 +45,7 @@ export function createDiscoveryRoutes(
     async handle(url, request) {
       switch (matchDiscoveryRoute(request.method, url.pathname)) {
         case "library": {
-          if (!libraryCache || Date.now() - libraryCache.at > 30_000) {
+          if (url.searchParams.get("refresh") === "1" || !libraryCache || Date.now() - libraryCache.at > 30_000) {
             const { Registry, visionCapable, audioCapable } = await import("../registry");
             const { loadModelConfig } = await import("../config");
             const { supportTier } = await import("../model/support");
@@ -132,7 +131,7 @@ export function createDiscoveryRoutes(
             top_p: ctx.genDefaults.topP ?? null,
             top_k: ctx.genDefaults.topK ?? null,
           };
-          const isGlm52 = ctx.model instanceof Glm52Model;
+          const capabilities = modelServingBinding(ctx).discovery;
           const data: Array<Record<string, unknown>> = [{
             id: ctx.modelId,
             object: "model",
@@ -147,10 +146,10 @@ export function createDiscoveryRoutes(
             batch_mode: gateway.batchMode,
             tools: true,
             structured_output: true,
-            embeddings: isEmbeddingModel(ctx.model),
-            adapters: !isGlm52,
-            training: !isGlm52,
-            dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
+            embeddings: capabilities.embeddings,
+            adapters: capabilities.adapters,
+            training: capabilities.training,
+            dsa: capabilities.dsa,
             mtp: ctx.draft?.provider.id === "glm52-native-mtp",
             capabilities: {
               chat_completions: true,
@@ -161,11 +160,11 @@ export function createDiscoveryRoutes(
               tools: true,
               structured_output: true,
               logprobs: true,
-              embeddings: isEmbeddingModel(ctx.model),
+              embeddings: capabilities.embeddings,
               vision: !!(ctx.vision || ctx.loadVision),
               audio: !!(ctx.audio || ctx.loadAudio),
-              adapters: !isGlm52,
-              training: !isGlm52,
+              adapters: capabilities.adapters,
+              training: capabilities.training,
             },
             gen_defaults: genDefaults,
           }];