Inferencing stuff and tests

cryptopoly · cryptopoly · commit b39c0bf1564c · 2026-04-15T19:06:19.000+01:00
diff --git a/backend_service/ddtree.py b/backend_service/ddtree.py
@@ -409,7 +409,12 @@ def generate_ddtree_mlx(
             cycles += 1
         else:
             # ── DDTree path ──────────────────────────────────
-            draft_logits_np = np.array(draft_logits[0].tolist(), dtype=np.float32)
+            # Evaluate and convert MLX tensor directly to NumPy via the
+            # buffer protocol — avoids the very slow .tolist() intermediate
+            # Python list, which is a major bottleneck for large vocabularies
+            # (150K+ tokens for Qwen models).
+            mx.eval(draft_logits)
+            draft_logits_np = np.array(draft_logits[0], dtype=np.float32)
 
             node_token_ids, node_depths, parents, child_maps, visibility = \
                 build_ddtree_tree(draft_logits_np, effective_budget)
diff --git a/backend_service/mlx_worker.py b/backend_service/mlx_worker.py
@@ -835,7 +835,14 @@ def _generate_dflash(self, request: dict[str, Any]) -> dict[str, Any]:
 
         gen_tokens = [int(token_id) for token_id in summary.get("generated_token_ids", [])]
         text = self.tokenizer.decode(gen_tokens).strip() if gen_tokens else ""
-        text = _strip_thinking_tokens(text) if text else ""
+        # Respect thinkingMode: only strip raw reasoning patterns when thinking
+        # is enabled. XML <think> tags are always processed regardless.
+        thinking_mode = request.get("thinkingMode") or "off"
+        if text:
+            think_filter = ThinkingTokenFilter(detect_raw_reasoning=(thinking_mode != "off"))
+            result = think_filter.feed(text)
+            flushed = think_filter.flush()
+            text = f"{result.text}{flushed.text}".strip()
         if not text:
             text = "Generation completed without decoded text."
 
@@ -911,7 +918,14 @@ def _generate_ddtree(self, request: dict[str, Any]) -> dict[str, Any]:
         # Decode output tokens
         gen_tokens = result["generated_tokens"]
         text = self.tokenizer.decode(gen_tokens).strip()
-        text = _strip_thinking_tokens(text) if text else ""
+        # Respect thinkingMode: only strip raw reasoning patterns when thinking
+        # is enabled. XML <think> tags are always processed regardless.
+        thinking_mode = request.get("thinkingMode") or "off"
+        if text:
+            think_filter = ThinkingTokenFilter(detect_raw_reasoning=(thinking_mode != "off"))
+            filter_result = think_filter.feed(text)
+            flushed = think_filter.flush()
+            text = f"{filter_result.text}{flushed.text}".strip()
         if not text:
             text = "Generation completed without decoded text."
 
@@ -1055,7 +1069,12 @@ def _generate_standard(self, request: dict[str, Any]) -> dict[str, Any]:
             )
 
         raw_text = "".join(text_parts).strip()
-        text = _strip_thinking_tokens(raw_text)
+        # Respect thinkingMode: only strip raw reasoning when thinking is on.
+        thinking_mode = request.get("thinkingMode") or "off"
+        think_filter = ThinkingTokenFilter(detect_raw_reasoning=(thinking_mode != "off"))
+        filter_result = think_filter.feed(raw_text)
+        flushed = think_filter.flush()
+        text = f"{filter_result.text}{flushed.text}".strip()
         if transcript_fallback:
             text, transcript_trimmed = _trim_transcript_continuation(text)
             if transcript_trimmed:
@@ -1085,8 +1104,42 @@ def stream_generate(self, request: dict[str, Any]) -> None:
             raise RuntimeError("No MLX model is loaded.")
 
         speculative_stream_fallback_note = None
-        # DFLASH doesn't support token-level streaming natively, so emit
-        # the full result as a single chunk in the streaming protocol.
+        # DFLASH/DDTree don't support token-level streaming natively, so
+        # emit the full result as a single chunk in the streaming protocol.
+        # Prefer DDTree (tree-based) when tree_budget > 0, else linear DFlash.
+        if self.speculative_decoding and self.tree_budget > 0 and self._ddtree_draft is not None:
+            try:
+                result = self._generate_ddtree(request)
+                if result.get("text"):
+                    _emit({"ok": True, "chunk": {"text": result["text"]}})
+                _emit({
+                    "ok": True,
+                    "done": True,
+                    "result": {
+                        "finishReason": result.get("finishReason", "stop"),
+                        "promptTokens": result.get("promptTokens", 0),
+                        "completionTokens": result.get("completionTokens", 0),
+                        "totalTokens": result.get("totalTokens", 0),
+                        "tokS": result.get("tokS", 0.0),
+                        "promptTokS": result.get("promptTokS", 0.0),
+                        "peakMemoryGb": result.get("peakMemoryGb", 0.0),
+                        "runtimeNote": result.get("runtimeNote"),
+                        "dflashAcceptanceRate": result.get("dflashAcceptanceRate"),
+                        "cacheStrategy": result.get("cacheStrategy"),
+                        "cacheBits": result.get("cacheBits"),
+                        "fp16Layers": result.get("fp16Layers"),
+                        "speculativeDecoding": result.get("speculativeDecoding"),
+                        "treeBudget": result.get("treeBudget"),
+                    },
+                })
+                return
+            except Exception as exc:
+                speculative_stream_fallback_note = (
+                    f"DDTree stream path failed ({exc}). "
+                    "Falling back to linear DFLASH."
+                )
+                # Fall through to linear DFLASH below
+
         if self.speculative_decoding and self._dflash_generator is not None:
             try:
                 result = self._generate_dflash(request)
diff --git a/scripts/build-llama-turbo.sh b/scripts/build-llama-turbo.sh
@@ -58,7 +58,11 @@ else
 fi
 
 # Platform-specific CMake flags
-CMAKE_FLAGS=(-DCMAKE_BUILD_TYPE=Release)
+# -DBUILD_SHARED_LIBS=OFF produces a self-contained static binary that
+# doesn't depend on .dylib/.so files at a build-time rpath. Without this
+# the installed binary crashes on dyld load because the shared libraries
+# are left behind in the build directory.
+CMAKE_FLAGS=(-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF)
 case "$(uname -s)" in
   Darwin)
     CMAKE_FLAGS+=(-DGGML_METAL=ON)
diff --git a/scripts/stage-runtime.mjs b/scripts/stage-runtime.mjs
@@ -69,6 +69,7 @@ function main() {
   }
 
   const chaosEngineBundle = stageVendoredChaosEngine(pythonInfo.executable);
+  const bundledOptionalPackages = stageOptionalRuntimePackages(pythonInfo.executable);
   validateBundledProjectImports(pythonInfo.executable);
   const llamaWarnings = stageLlamaBinaries();
   maybeSignEmbeddedRuntime();
@@ -88,6 +89,7 @@ function main() {
     llamaCli: fs.existsSync(path.join(binDest, binaryName("llama-cli"))) ? `bin/${binaryName("llama-cli")}` : null,
     pythonVersion: pythonInfo.versionTag,
     bundledCacheStrategies: chaosEngineBundle ? ["chaosengine"] : [],
+    bundledOptionalPackages: bundledOptionalPackages,
     warnings: llamaWarnings,
   };
 
@@ -169,6 +171,7 @@ function validateBundledPythonPackages(pythonBinary) {
     "requirements = {",
     "  'desktop': [('fastapi', 'fastapi'), ('huggingface_hub', 'huggingface_hub'), ('psutil', 'psutil'), ('uvicorn', 'uvicorn')],",
     "  'images': [('accelerate', 'accelerate'), ('diffusers', 'diffusers'), ('huggingface_hub', 'huggingface_hub'), ('PIL', 'pillow'), ('safetensors', 'safetensors'), ('torch', 'torch')],",
+    "  'inference': [('dflash_mlx', 'dflash-mlx'), ('turboquant', 'turboquant'), ('turboquant_mlx', 'turboquant-mlx-full')],",
     "}",
     "missing = {",
     "  group: [label for module, label in modules if importlib.util.find_spec(module) is None]",
@@ -182,14 +185,28 @@ function validateBundledPythonPackages(pythonBinary) {
     encoding: "utf8",
   }).trim();
   const missing = JSON.parse(payload);
-  const flatMissing = Object.entries(missing)
+
+  // Inference packages are optional — warn but never block the build.
+  const optionalGroups = new Set(["inference"]);
+  const requiredMissing = Object.entries(missing)
+    .filter(([group]) => !optionalGroups.has(group))
+    .flatMap(([group, values]) => values.length ? values.map((value) => `${group}:${value}`) : []);
+  const optionalMissing = Object.entries(missing)
+    .filter(([group]) => optionalGroups.has(group))
     .flatMap(([group, values]) => values.length ? values.map((value) => `${group}:${value}`) : []);
-  if (flatMissing.length === 0) {
+
+  if (optionalMissing.length) {
+    console.warn(
+      `[stage-runtime] info: optional inference packages not in build venv (${optionalMissing.join(", ")}). ` +
+      `DFlash/TurboQuant will require user install via the Setup page.`,
+    );
+  }
+  if (requiredMissing.length === 0) {
     return;
   }
 
   const message =
-    `Embedded Python is missing required runtime packages (${flatMissing.join(", ")}). ` +
+    `Embedded Python is missing required runtime packages (${requiredMissing.join(", ")}). ` +
     `Install them into the build venv with: ${pythonBinary} -m pip install -e ".[desktop,images]"`;
   if (strict) {
     throw new Error(message);
@@ -278,6 +295,82 @@ function resolveChaosEngineVendor() {
   };
 }
 
+function stageOptionalRuntimePackages(pythonBinary) {
+  // Pre-install optional runtime packages into the staged site-packages
+  // so that DFlash, TurboQuant, and RotorQuant work out of the box for
+  // new users without requiring manual pip installs via the Setup page.
+  //
+  // Each entry: [pip package name, import name used for verification]
+  const optionalPackages = [
+    ["dflash-mlx", "dflash_mlx"],
+    ["turboquant", "turboquant"],
+    ["turboquant-mlx-full", "turboquant_mlx"],
+  ];
+
+  const installed = [];
+  const skipped = [];
+
+  for (const [pipName, importName] of optionalPackages) {
+    // Check if already available in the build venv
+    const checkScript = `import importlib.util; exit(0 if importlib.util.find_spec("${importName}") else 1)`;
+    let available = false;
+    try {
+      execFileSync(pythonBinary, ["-c", checkScript], {
+        cwd: workspaceRoot,
+        stdio: "ignore",
+      });
+      available = true;
+    } catch {
+      available = false;
+    }
+
+    if (!available) {
+      const message = `Optional package "${pipName}" not found in build venv — skipping bundle`;
+      if (strict) {
+        console.warn(`[stage-runtime] warning: ${message}. Install with: ${pythonBinary} -m pip install ${pipName}`);
+      }
+      skipped.push(pipName);
+      continue;
+    }
+
+    try {
+      console.log(`[stage-runtime] bundling optional package: ${pipName}`);
+      execFileSync(
+        pythonBinary,
+        [
+          "-m", "pip", "install",
+          "--disable-pip-version-check",
+          "--no-deps",
+          "--no-compile",
+          "--upgrade",
+          "--target", sitePackagesDest,
+          pipName,
+        ],
+        {
+          cwd: workspaceRoot,
+          stdio: "inherit",
+        },
+      );
+      installed.push(pipName);
+    } catch (err) {
+      const message = `Failed to bundle optional package "${pipName}": ${err.message}`;
+      if (strict) {
+        throw new Error(message);
+      }
+      console.warn(`[stage-runtime] warning: ${message}`);
+      skipped.push(pipName);
+    }
+  }
+
+  if (installed.length) {
+    console.log(`[stage-runtime] bundled optional packages: ${installed.join(", ")}`);
+  }
+  if (skipped.length) {
+    console.log(`[stage-runtime] skipped optional packages (not in build venv): ${skipped.join(", ")}`);
+  }
+  return installed;
+}
+
 function stageLlamaBinaries() {
   const warnings = [];
   const sourceDir = process.env.CHAOSENGINE_LLAMA_BIN_DIR || defaultLlamaBinDir();
diff --git a/tests/inference-batch-dflash-comparison.json b/tests/inference-batch-dflash-comparison.json
@@ -0,0 +1,116 @@
+[
+  {
+    "_label": "Qwen3.5-9B MLX — native, NO DFlash (baseline)",
+    "modelRef": "mlx-community/Qwen3.5-9B-4bit",
+    "modelName": "Qwen3.5-9B-MLX-4bit",
+    "path": "/Users/dan/AI_Models/mlx-community/Qwen3.5-9B-MLX-4bit",
+    "backend": "mlx",
+    "cacheStrategy": "native",
+    "cacheBits": 0,
+    "fp16Layers": 0,
+    "contextTokens": 8192,
+    "maxTokens": 512,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": false,
+    "treeBudget": 0,
+    "thinkingMode": "off",
+    "prompt": "Explain the difference between TCP and UDP. Cover reliability, ordering, speed, and typical use cases."
+  },
+  {
+    "_label": "Qwen3.5-9B MLX — native, DFlash ON (tree=16)",
+    "modelRef": "mlx-community/Qwen3.5-9B-4bit",
+    "modelName": "Qwen3.5-9B-MLX-4bit",
+    "path": "/Users/dan/AI_Models/mlx-community/Qwen3.5-9B-MLX-4bit",
+    "backend": "mlx",
+    "cacheStrategy": "native",
+    "cacheBits": 0,
+    "fp16Layers": 0,
+    "contextTokens": 8192,
+    "maxTokens": 512,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": true,
+    "treeBudget": 16,
+    "thinkingMode": "off",
+    "prompt": "Explain the difference between TCP and UDP. Cover reliability, ordering, speed, and typical use cases."
+  },
+  {
+    "_label": "Qwen3.5-35B-A3B MLX — native, NO DFlash (baseline)",
+    "modelRef": "mlx-community/Qwen3.5-35B-A3B-4bit",
+    "modelName": "Qwen3.5-35B-A3B-4bit",
+    "path": "/Users/dan/AI_Models/mlx-community/Qwen3.5-35B-A3B-4bit",
+    "backend": "mlx",
+    "cacheStrategy": "native",
+    "cacheBits": 0,
+    "fp16Layers": 0,
+    "contextTokens": 8192,
+    "maxTokens": 512,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": false,
+    "treeBudget": 0,
+    "thinkingMode": "off",
+    "prompt": "Explain the difference between TCP and UDP. Cover reliability, ordering, speed, and typical use cases."
+  },
+  {
+    "_label": "Qwen3.5-35B-A3B MLX — native, DFlash ON (tree=16)",
+    "modelRef": "mlx-community/Qwen3.5-35B-A3B-4bit",
+    "modelName": "Qwen3.5-35B-A3B-4bit",
+    "path": "/Users/dan/AI_Models/mlx-community/Qwen3.5-35B-A3B-4bit",
+    "backend": "mlx",
+    "cacheStrategy": "native",
+    "cacheBits": 0,
+    "fp16Layers": 0,
+    "contextTokens": 8192,
+    "maxTokens": 512,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": true,
+    "treeBudget": 16,
+    "thinkingMode": "off",
+    "prompt": "Explain the difference between TCP and UDP. Cover reliability, ordering, speed, and typical use cases."
+  },
+  {
+    "_label": "Qwen3.5-9B GGUF — TurboQuant 4-bit",
+    "modelRef": "lmstudio-community/Qwen3.5-9B-GGUF",
+    "modelName": "Qwen3.5-9B-GGUF",
+    "path": "/Users/dan/.cache/huggingface/hub/models--lmstudio-community--Qwen3.5-9B-GGUF",
+    "backend": "gguf",
+    "cacheStrategy": "turboquant",
+    "cacheBits": 4,
+    "fp16Layers": 4,
+    "contextTokens": 4096,
+    "maxTokens": 256,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": false,
+    "treeBudget": 0,
+    "thinkingMode": "off",
+    "prompt": "What is the capital of France? Answer in one sentence."
+  },
+  {
+    "_label": "Qwen3.5-9B GGUF — RotorQuant 4-bit",
+    "modelRef": "lmstudio-community/Qwen3.5-9B-GGUF",
+    "modelName": "Qwen3.5-9B-GGUF",
+    "path": "/Users/dan/.cache/huggingface/hub/models--lmstudio-community--Qwen3.5-9B-GGUF",
+    "backend": "gguf",
+    "cacheStrategy": "rotorquant",
+    "cacheBits": 4,
+    "fp16Layers": 4,
+    "contextTokens": 4096,
+    "maxTokens": 256,
+    "temperature": 0.7,
+    "fusedAttention": false,
+    "fitModelInMemory": true,
+    "speculativeDecoding": false,
+    "treeBudget": 0,
+    "thinkingMode": "off",
+    "prompt": "What is a neural network? Answer in two sentences."
+  }
+]
diff --git a/tests/test_inference.py b/tests/test_inference.py
