feat(models): add MTP drafters for Gemma 4 QAT models

src/helpers/ipcHandlers.js

@@ -3724,7 +3724,10 @@ class IPCHandlers {
 
         const modelPath = require("path").join(modelManager.modelsDir, modelInfo.model.fileName);
 
-        await modelManager.serverManager.start(modelPath, modelManager.serverOptions(modelInfo));
+        await modelManager.serverManager.start(
+          modelPath,
+          await modelManager.serverStartOptions(modelInfo)
+        );
         modelManager.currentServerModelId = modelId;
 
         this.environmentManager.saveAllKeysToEnvFile().catch(() => {});

src/helpers/llamaServer.js

@@ -27,6 +27,10 @@ class LlamaServerManager {
     this.port = null;
     this.ready = false;
     this.modelPath = null;
+    // draftModelPath is the REQUESTED drafter (stable across identical requests, drives
+    // the start() restart check); activeDraftModelPath is the one that actually loaded.
+    this.draftModelPath = null;
+    this.activeDraftModelPath = null;
     this.startupPromise = null;
     this.healthCheckInterval = null;
     this.healthCheckFailures = 0;
@@ -107,7 +111,11 @@ class LlamaServerManager {
   async start(modelPath, options = {}) {
     if (this.startupPromise) return this.startupPromise;
 
-    if (this.ready && this.modelPath === modelPath) return;
+    // A change in drafter presence for the same model must still restart the
+    // server so the new speculative-decoding flags take effect.
+    const requestedDraftPath = options.draftModelPath || null;
+    if (this.ready && this.modelPath === modelPath && this.draftModelPath === requestedDraftPath)
+      return;
 
     if (this.process) {
       await this.stop();
@@ -128,6 +136,10 @@ class LlamaServerManager {
 
     this.port = await this.findAvailablePort();
     this.modelPath = modelPath;
+    // Store the REQUESTED drafter so start() compares against a stable value across
+    // identical requests; activeDraftModelPath tracks what actually loaded (see ctor).
+    this.draftModelPath = options.draftModelPath || null;
+    this.activeDraftModelPath = null;
 
     const baseArgs = [
       "--model",
@@ -145,17 +157,32 @@ class LlamaServerManager {
       "--jinja",
     ];
 
+    // Draft flags stay separate from baseArgs so the fallback ladder can retry without
+    // them when a stale (pre-b9763) binary rejects the MTP args at parse time.
+    const draftArgs = options.draftModelPath
+      ? [
+          "--model-draft",
+          options.draftModelPath,
+          "--spec-type",
+          "draft-mtp",
+          "--spec-draft-n-max",
+          "3",
+        ]
+      : [];
+
     if (process.platform === "darwin") {
-      const args = [...baseArgs, "--n-gpu-layers", String(options.gpuLayers ?? 99)];
+      // The metal binary is always the bundled pin, so it never rejects the draft flags.
+      const args = [...baseArgs, "--n-gpu-layers", String(options.gpuLayers ?? 99), ...draftArgs];
       await this._startWithBinary(
         binaryPaths.default,
         args,
         this._buildEnv(binaryPaths.default),
         STARTUP_TIMEOUT_MS
       );
       this.activeBackend = "metal";
+      this.activeDraftModelPath = this.draftModelPath;
     } else {
-      await this._startWithGpuFallback(binaryPaths, baseArgs, options);
+      await this._startWithGpuFallback(binaryPaths, baseArgs, options, draftArgs);
     }
 
     this.startHealthCheck();
@@ -164,41 +191,91 @@ class LlamaServerManager {
       port: this.port,
       model: path.basename(modelPath),
       backend: this.activeBackend,
+      mtp: this.activeDraftModelPath !== null,
     });
   }
 
-  async _startWithGpuFallback(binaryPaths, baseArgs, options) {
+  async _startWithGpuFallback(binaryPaths, baseArgs, options, draftArgs = []) {
     const gpuArgs = [...baseArgs, "--n-gpu-layers", String(options.gpuLayers ?? 99)];
     const cpuArgs = baseArgs;
+    const hasDraft = draftArgs.length > 0;
+
+    // Degrade ladder: GPU+MTP, then GPU alone (a live GPU beats speculation), then
+    // CPU+MTP (the bundled pin normally accepts the flags), then plain CPU. The
+    // no-draft rungs collapse into their twins when no drafter is declared, so a
+    // drafterless start keeps today's exact single vulkan->cpu fallback.
+    const rungs = [
+      {
+        backend: "vulkan",
+        name: "Vulkan",
+        binary: binaryPaths.vulkan,
+        args: [...gpuArgs, ...draftArgs],
+        mtp: hasDraft,
+        timeout: VULKAN_STARTUP_TIMEOUT_MS,
+        attemptMsg: "Attempting Vulkan backend startup",
+      },
+      {
+        backend: "vulkan",
+        name: "Vulkan",
+        binary: binaryPaths.vulkan,
+        args: gpuArgs,
+        mtp: false,
+        noDraft: true,
+        timeout: VULKAN_STARTUP_TIMEOUT_MS,
+        attemptMsg: "Attempting Vulkan backend startup",
+      },
+      {
+        backend: "cpu",
+        name: "CPU",
+        binary: binaryPaths.cpu,
+        args: [...cpuArgs, ...draftArgs],
+        mtp: hasDraft,
+        timeout: STARTUP_TIMEOUT_MS,
+        attemptMsg: "Starting with CPU backend",
+      },
+      {
+        backend: "cpu",
+        name: "CPU",
+        binary: binaryPaths.cpu,
+        args: cpuArgs,
+        mtp: false,
+        noDraft: true,
+        timeout: STARTUP_TIMEOUT_MS,
+        attemptMsg: "Starting with CPU backend",
+      },
+    ];
+
+    const ladder = rungs.filter((rung) => rung.binary && !(rung.noDraft && !hasDraft));
+    if (ladder.length === 0) throw new Error("No CPU llama-server binary available");
 
-    if (binaryPaths.vulkan) {
+    let lastError = null;
+    for (let i = 0; i < ladder.length; i++) {
+      const rung = ladder[i];
+      const next = ladder[i + 1];
       try {
-        debugLogger.debug("Attempting Vulkan backend startup");
+        debugLogger.debug(rung.attemptMsg);
         await this._startWithBinary(
-          binaryPaths.vulkan,
-          gpuArgs,
-          this._buildEnv(binaryPaths.vulkan),
-          VULKAN_STARTUP_TIMEOUT_MS
+          rung.binary,
+          rung.args,
+          this._buildEnv(rung.binary),
+          rung.timeout
         );
-        this.activeBackend = "vulkan";
+        this.activeBackend = rung.backend;
+        this.activeDraftModelPath = rung.mtp ? this.draftModelPath : null;
         return;
       } catch (err) {
-        debugLogger.warn("Vulkan backend failed, falling back to CPU", { error: err.message });
-        await this._killCurrentProcess();
-        this.port = await this.findAvailablePort();
+        lastError = err;
+        if (next) {
+          debugLogger.warn(`${rung.name} backend failed, falling back to ${next.name}`, {
+            error: err.message,
+          });
+          await this._killCurrentProcess();
+          this.port = await this.findAvailablePort();
+        }
       }
     }
 
-    if (!binaryPaths.cpu) throw new Error("No CPU llama-server binary available");
-
-    debugLogger.debug("Starting with CPU backend");
-    await this._startWithBinary(
-      binaryPaths.cpu,
-      cpuArgs,
-      this._buildEnv(binaryPaths.cpu),
-      STARTUP_TIMEOUT_MS
-    );
-    this.activeBackend = "cpu";
+    throw lastError || new Error("No CPU llama-server binary available");
   }
 
   _buildEnv(binaryPath) {
@@ -554,6 +631,8 @@ class LlamaServerManager {
     this.ready = false;
     this.port = null;
     this.modelPath = null;
+    this.draftModelPath = null;
+    this.activeDraftModelPath = null;
     this.activeBackend = null;
   }
 

src/helpers/modelManagerBridge.js

@@ -206,6 +206,13 @@ class ModelManager {
     };
   }
 
+  async serverStartOptions(modelInfo) {
+    const options = this.serverOptions(modelInfo);
+    const draftPath = await this.resolveDraftPath(modelInfo.model);
+    if (draftPath) options.draftModelPath = draftPath;
+    return options;
+  }
+
   async downloadModel(modelId, onProgress) {
     this.ensureInitialized();
     const modelInfo = this.findModelById(modelId);
@@ -236,7 +243,12 @@ class ModelManager {
     try {
       await this.ensureModelsDirExists();
 
-      const requiredBytes = model.sizeBytes || model.sizeMb * 1_000_000 || 0;
+      const hasDrafter = this.modelHasDrafter(model);
+
+      let requiredBytes = model.sizeBytes || model.sizeMb * 1_000_000 || 0;
+      if (requiredBytes > 0 && hasDrafter) {
+        requiredBytes += model.draftSizeBytes;
+      }
       if (requiredBytes > 0) {
         const spaceCheck = await checkDiskSpace(this.modelsDir, requiredBytes * 1.2);
         if (!spaceCheck.ok) {
@@ -251,20 +263,41 @@ class ModelManager {
 
       const downloadUrl = this.getDownloadUrl(provider, model);
 
+      // With a drafter, weight progress across both files by declared bytes and
+      // clamp it so the bar never regresses across the phase boundary. Without a
+      // drafter, keep today's single-file progress exactly.
+      const combinedTotal = hasDrafter ? (model.sizeBytes || 0) + model.draftSizeBytes : 0;
+      let lastCombined = 0;
+      const emitCombined = (rawCombined) => {
+        let combined = Math.min(rawCombined, combinedTotal);
+        if (combined < lastCombined) combined = lastCombined;
+        else lastCombined = combined;
+        const progress = combinedTotal > 0 ? (combined / combinedTotal) * 100 : 0;
+        this.downloadProgress.set(modelId, {
+          modelId,
+          progress,
+          downloadedSize: combined,
+          totalSize: combinedTotal,
+        });
+        if (onProgress) onProgress(progress, combined, combinedTotal);
+      };
+
       await sharedDownloadFile(downloadUrl, modelPath, {
         signal,
-        onProgress: (downloadedBytes, totalBytes) => {
-          const progress = totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0;
-          this.downloadProgress.set(modelId, {
-            modelId,
-            progress,
-            downloadedSize: downloadedBytes,
-            totalSize: totalBytes,
-          });
-          if (onProgress) {
-            onProgress(progress, downloadedBytes, totalBytes);
-          }
-        },
+        onProgress: hasDrafter
+          ? (downloadedBytes) => emitCombined(downloadedBytes)
+          : (downloadedBytes, totalBytes) => {
+              const progress = totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0;
+              this.downloadProgress.set(modelId, {
+                modelId,
+                progress,
+                downloadedSize: downloadedBytes,
+                totalSize: totalBytes,
+              });
+              if (onProgress) {
+                onProgress(progress, downloadedBytes, totalBytes);
+              }
+            },
       });
 
       const stats = await fsPromises.stat(modelPath);
@@ -277,6 +310,29 @@ class ModelManager {
         );
       }
 
+      // Drafter is opportunistic: a failure or cancel here leaves the main model
+      // fully usable, so never fail the download or delete the main file.
+      if (hasDrafter) {
+        const draftPath = path.join(this.modelsDir, model.draftFileName);
+        try {
+          await sharedDownloadFile(this.getDraftDownloadUrl(provider, model), draftPath, {
+            signal,
+            onProgress: (downloadedBytes) => emitCombined(stats.size + downloadedBytes),
+          });
+          const draftStats = await fsPromises.stat(draftPath);
+          if (draftStats.size < MIN_FILE_SIZE) {
+            await fsPromises.unlink(draftPath).catch(() => {});
+            debugLogger.warn("MTP drafter file too small, keeping model without it", { modelId });
+          }
+        } catch (draftError) {
+          await fsPromises.unlink(draftPath).catch(() => {});
+          debugLogger.warn("MTP drafter download failed, keeping model without it", {
+            modelId,
+            error: draftError.message,
+          });
+        }
+      }
+
       return modelPath;
     } catch (error) {
       if (error.isAbort) {
@@ -307,6 +363,24 @@ class ModelManager {
     return `${baseUrl}/${model.hfRepo}/resolve/main/${model.fileName}`;
   }
 
+  getDraftDownloadUrl(provider, model) {
+    const baseUrl = provider.baseUrl || "https://huggingface.co";
+    return `${baseUrl}/${model.draftHfRepo}/resolve/main/${model.draftFileName}`;
+  }
+
+  modelHasDrafter(model) {
+    return Boolean(model && model.draftHfRepo && model.draftFileName && model.draftSizeBytes);
+  }
+
+  // Opportunistic MTP drafter path: only when declared and the file passes the
+  // same >1MB validity gate as models. Returns null otherwise (start without MTP).
+  async resolveDraftPath(model) {
+    if (!this.modelHasDrafter(model)) return null;
+    const draftPath = path.join(this.modelsDir, model.draftFileName);
+    if (await this.checkModelValid(draftPath)) return draftPath;
+    return null;
+  }
+
   cancelDownload(modelId) {
     const entry = this.activeRequests.get(modelId);
     if (entry) {
@@ -330,6 +404,12 @@ class ModelManager {
     if (await this.checkFileExists(modelPath)) {
       await fsPromises.unlink(modelPath);
     }
+
+    // Remove the drafter too when present, best effort (ignore ENOENT).
+    if (modelInfo.model.draftFileName) {
+      const draftPath = path.join(this.modelsDir, modelInfo.model.draftFileName);
+      await fsPromises.unlink(draftPath).catch(() => {});
+    }
   }
 
   async deleteAllModels() {
@@ -409,7 +489,7 @@ class ModelManager {
         serverReady: this.serverManager.ready,
       });
 
-      await this.serverManager.start(modelPath, this.serverOptions(modelInfo));
+      await this.serverManager.start(modelPath, await this.serverStartOptions(modelInfo));
       this.currentServerModelId = modelId;
 
       debugLogger.logReasoning("INFERENCE_SERVER_STARTED", {
@@ -479,7 +559,7 @@ class ModelManager {
     if (!this.serverManager.isAvailable()) return false;
 
     try {
-      await this.serverManager.start(modelPath, this.serverOptions(modelInfo));
+      await this.serverManager.start(modelPath, await this.serverStartOptions(modelInfo));
       this.currentServerModelId = modelId;
       debugLogger.info("llama-server pre-warmed", { modelId });
       return true;

src/models/ModelRegistry.ts

@@ -15,6 +15,10 @@ export interface ModelDefinition {
   hfRepo: string;
   recommended?: boolean;
   supportsThinking?: boolean;
+  // Optional MTP speculative-decoding drafter downloaded alongside the main GGUF.
+  draftHfRepo?: string;
+  draftFileName?: string;
+  draftSizeBytes?: number;
 }
 
 export interface LocalProviderData {