mlx-bun 0.0.13 → 0.3.0

package/src/model/trellis-wide-prefill.ts

@@ -0,0 +1,94 @@
+// Packed Trellis prefill with MLX 0.32.2 GemvWide arithmetic.
+// Ported from ml-explore/mlx v0.32.2, mlx/backend/metal/kernels/gemv.h
+// (MIT). Decode weights to bf16 inside the native dot/unroll/shuffle order;
+// share each decoded row across every input vector in this narrow regime.
+import { MetalKernel } from "../mlx/metal-kernel";
+import { type MlxArray } from "../mlx/array";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
+import { type TrellisGeometry } from "./trellis-linear";
+
+/** Matches the bundled runtime's native reduction on M3 and newer GPUs. */
+export function nativeTrellisWidePrefill(m: number): boolean {
+  return m >= 5 && m <= 15 &&
+    Number(/^applegpu_[a-z](\d{2})/.exec(deviceArchitecture())?.[1] ?? 0) >= 15;
+}
+
+export function wideTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype): boolean {
+  return dtype === Dtype.bfloat16 && nativeTrellisWidePrefill(m) &&
+    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 5120 && g.outFeatures === 17408 && !g.blockInterleave;
+}
+
+let kernel: MetalKernel | undefined;
+function wideKernel() {
+  return kernel ??= new MetalKernel({ name: 'mlx_bun_trellis_wide_prefill',
+    inputNames: ['x', 'codes', 'scales'], outputNames: ['out'], source: String.raw`
+    constexpr uint unroll=8, rows_per_simd=32/KL, simdgroups=KL/8;
+    const uint lane=thread_index_in_simdgroup, sg=simdgroup_index_in_threadgroup;
+    const uint kl=lane%KL, row=threadgroup_position_in_grid.y*4+sg*rows_per_simd+lane/KL;
+    const uint vec0=threadgroup_position_in_grid.x*SM;
+    constexpr uint wpb=BT*BITS/32, row_words=KC*BITS/32;
+    const float scale=float(scales[row]);
+    constexpr uint n_v4=KC/4, n_main=n_v4-n_v4%(KL*unroll);
+    float result[SM]={0};
+    for(uint base=0;base<n_main;base+=KL*unroll) {
+      float4 wf[unroll];
+      #pragma clang loop unroll(full)
+      for(uint i=0;i<unroll;i++) {
+        #pragma clang loop unroll(full)
+        for(uint c=0;c<4;c++) {
+          const uint k=(base+i*KL+kl)*4+c, block=k/BT, t=k%BT;
+          const device uint* bw=codes+(ulong)row*row_words+block*wpb;
+          const uint pos=(BT-1-t)*BITS, wi=pos>>5, off=pos&31u;
+          uint win=bw[wi]>>off;
+          if(off+L>32)win|=bw[wi+1==wpb?0:wi+1]<<(32-off);
+          win&=(1u<<L)-1;
+          const uint z=win*34038481u+76625530u;
+          const uint p=(z&0x00FF00FFu)+((z>>8)&0x00FF00FFu);
+          const int code=int((p&0xFFFFu)+(p>>16))-510;
+          const float reciprocal=1.0f/147.800537109375f;
+          const float approx=float(code)*reciprocal;
+          const float residual=metal::fma(-approx,147.800537109375f,float(code));
+          const float value=metal::fma(residual,reciprocal,approx);
+          wf[i][c]=float(T(value*scale));
+        }
+      }
+      #pragma clang loop unroll(full)
+      for(uint v=0;v<SM;v++) {
+        const uint m=min(vec0+v,uint(M-1));
+        float acc=0;
+        #pragma clang loop unroll(full)
+        for(uint i=0;i<unroll;i++) {
+          const ulong offset=(ulong)m*KC+(base+i*KL+kl)*4;
+          const float4 xv=float4(float(x[offset]),float(x[offset+1]),float(x[offset+2]),float(x[offset+3]));
+          acc+=metal::dot(wf[i],xv);
+        }
+        result[v]+=acc;
+      }
+    }
+    #pragma clang loop unroll(full)
+    for(uint v=0;v<SM;v++) {
+      #pragma clang loop unroll(full)
+      for(ushort off=KL/2;off>=1;off>>=1) result[v]+=simd_shuffle_down(result[v],off);
+      if(kl==0 && vec0+v<M) out[(ulong)(vec0+v)*N+row]=T(result[v]);
+    }
+  ` });
+}
+/** Requires the caller to prove aligned, row-contiguous x, as Qwen does
+ * with its RMSNorm output. Preserve native fallback when that proof is absent. */
+export function wideTrellisPrefill(x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const [m, k] = x.shape;
+  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 15 ||
+      k !== 5120 || g.inFeatures !== k || g.axis !== 1 || g.T !== 256 || g.L !== 12 ||
+      ![2, 3, 4].includes(g.k) || g.blockInterleave || g.rows < 4 || g.rows % 4 ||
+      g.outFeatures !== g.rows || g.cols !== k || codes.dtype !== Dtype.uint32 ||
+      codes.ndim !== 2 || codes.shape[0] !== g.rows || codes.shape[1] !== k * g.k / 32 || scales.size !== g.rows)
+    throw new Error("unsupported wide trellis prefill geometry, dtype or layout");
+  const kl = m <= 5 ? 32 : 16;
+  return wideKernel().apply([x, codes, scales], {
+    outputs: [{ shape: [m, g.outFeatures], dtype: x.dtype }],
+    grid: [32, g.outFeatures / 4 * (kl / 8), 1], threadGroup: [32, kl / 8, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M: m, SM: m, KL: kl, KC: k, N: g.outFeatures, BT: g.T, L: g.L, BITS: g.k },
+  })[0]!;
+}

package/src/model/universal/dense.ts

@@ -3,7 +3,7 @@
 // like qwen3.ts, with every arch delta selected by the UniversalArgs
 // descriptor (archs.ts). Monolith path only: no compiled
 // decode, no kv-quant — slow, never broken. The bar is L1 bit-exactness
-// vs mlx-lm on this machine's GPU (tests/universal-parity.test.ts).
+// vs mlx-lm on this machine's GPU (tests/parity/universal-parity.test.ts).
 //
 // Porting discipline: each branch transcribes its mlx-lm source op-for-op
 // (q/k-norm before vs after reshape, rope order, scale points, softcap,

package/src/model/universal/modules.ts

@@ -11,6 +11,7 @@ import type { Weights } from "../../weights";
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import { QuantizedEmbedding, QuantizedLinear, RMSNorm } from "../gemma4-base";
+import { TrellisLinear } from "../trellis-linear";
 
 /** Records every tensor a universal load consumes; `finish` diffs against
  *  the shard index so a descriptor mistake is a LOAD error (unconsumed /
@@ -66,7 +67,7 @@ export class DenseLinear {
 }
 
 /** Either linear flavor behind one call surface. */
-export type AnyLinear = QuantizedLinear | DenseLinear;
+export type AnyLinear = QuantizedLinear | DenseLinear | TrellisLinear;
 
 /** Quantized when `.scales` exists (MLX-quantized checkpoints), else dense
  *  bf16/f16/f32 (Phase 1.5). Both flavors carry the optional ADDITIVE
@@ -74,6 +75,12 @@ export type AnyLinear = QuantizedLinear | DenseLinear;
 export function loadLinear(
   weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
 ): AnyLinear {
+  if (TrellisLinear.isTrellis(config, path)) {
+    const lin = TrellisLinear.load(weights, path, config);
+    audit.use(`${path}.weight`);
+    audit.use(`${path}.scales`);
+    return lin;
+  }
   if (weights.has(`${path}.scales`)) {
     const lin = QuantizedLinear.load(weights, path, config);
     audit.use(`${path}.weight`);

package/src/native/frame_extract.swift

@@ -0,0 +1,140 @@
+// AVFoundation frame-extraction sidecar probe (PLAN 14w) — the video
+// equivalent of the audio path's afconvert doctrine: macOS-native codecs,
+// zero vendored dependencies.
+//
+//   frame-extract <video> <outDir> [fps=2] [maxFrames=768]
+//     → writes frame-%04d.png (RGB) + prints "N <count> W <w> H <h>"
+//   frame-extract --make-fixture <out.mov>
+//     → writes the deterministic 6-frame 320x240 gradient test clip
+//
+// Sampling mirrors mlx-vlm's do_sample_frames: timestamps k/fps for
+// k = 0..min(ceil(duration*fps), maxFrames)-1, precise tolerance.
+
+import AVFoundation
+import Foundation
+import ImageIO
+import UniformTypeIdentifiers
+
+func die(_ msg: String) -> Never {
+    FileHandle.standardError.write((msg + "\n").data(using: .utf8)!)
+    exit(1)
+}
+
+func makeFixture(_ path: String) {
+    let url = URL(fileURLWithPath: path)
+    try? FileManager.default.removeItem(at: url)
+    guard let writer = try? AVAssetWriter(outputURL: url, fileType: .mov) else {
+        die("AVAssetWriter init failed")
+    }
+    let W = 320, H = 240, FRAMES = 6
+    let settings: [String: Any] = [
+        AVVideoCodecKey: AVVideoCodecType.h264,
+        AVVideoWidthKey: W,
+        AVVideoHeightKey: H,
+    ]
+    let input = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
+    let adaptor = AVAssetWriterInputPixelBufferAdaptor(
+        assetWriterInput: input,
+        sourcePixelBufferAttributes: [
+            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32ARGB,
+            kCVPixelBufferWidthKey as String: W,
+            kCVPixelBufferHeightKey as String: H,
+        ])
+    writer.add(input)
+    writer.startWriting()
+    writer.startSession(atSourceTime: .zero)
+    for f in 0..<FRAMES {
+        while !input.isReadyForMoreMediaData { Thread.sleep(forTimeInterval: 0.01) }
+        var pb: CVPixelBuffer?
+        CVPixelBufferCreate(nil, W, H, kCVPixelFormatType_32ARGB, nil, &pb)
+        guard let pixelBuffer = pb else { die("CVPixelBufferCreate failed") }
+        CVPixelBufferLockBaseAddress(pixelBuffer, [])
+        let base = CVPixelBufferGetBaseAddress(pixelBuffer)!
+        let bpr = CVPixelBufferGetBytesPerRow(pixelBuffer)
+        let p = base.assumingMemoryBound(to: UInt8.self)
+        for y in 0..<H {
+            for x in 0..<W {
+                let o = y * bpr + x * 4
+                p[o] = 255
+                p[o + 1] = UInt8((x * 255 / W + f * 20) % 256)
+                p[o + 2] = UInt8((y * 255 / H + f * 10) % 256)
+                p[o + 3] = UInt8(((x + y) * 255 / (W + H) + f * 30) % 256)
+            }
+        }
+        CVPixelBufferUnlockBaseAddress(pixelBuffer, [])
+        adaptor.append(pixelBuffer, withPresentationTime: CMTime(value: CMTimeValue(f), timescale: 2))
+    }
+    input.markAsFinished()
+    let sem = DispatchSemaphore(value: 0)
+    writer.finishWriting { sem.signal() }
+    sem.wait()
+    guard writer.status == .completed else { die("finishWriting status \(writer.status.rawValue)") }
+    print("fixture \(path)")
+}
+
+let args = CommandLine.arguments
+if args.count >= 3 && args[1] == "--make-fixture" {
+    makeFixture(args[2])
+    exit(0)
+}
+guard args.count >= 3 else { die("usage: frame-extract <video> <outDir> [fps] [maxFrames]") }
+let videoPath = args[1]
+let outDir = args[2]
+let fps = args.count > 3 ? Double(args[3]) ?? 2.0 : 2.0
+let maxFrames = args.count > 4 ? Int(args[4]) ?? 768 : 768
+
+let asset = AVURLAsset(url: URL(fileURLWithPath: videoPath))
+let sem = DispatchSemaphore(value: 0)
+var duration = CMTime.zero
+var loadError: Error?
+Task {
+    do { duration = try await asset.load(.duration) } catch { loadError = error }
+    sem.signal()
+}
+sem.wait()
+if let e = loadError { die("asset load failed: \(e.localizedDescription)") }
+// Fragmented/malformed containers can load with an indefinite duration —
+// CMTimeGetSeconds yields NaN and Int(NaN) is a Swift runtime TRAP. Fall
+// back to the frame cap; the extraction loop below already stops at the
+// first undecodable timestamp.
+let seconds = CMTimeGetSeconds(duration)
+let count = seconds.isFinite
+    ? max(1, min(Int((seconds * fps).rounded(.up)), maxFrames))
+    : maxFrames
+
+let gen = AVAssetImageGenerator(asset: asset)
+gen.appliesPreferredTrackTransform = true
+gen.requestedTimeToleranceBefore = .zero
+gen.requestedTimeToleranceAfter = .zero
+// Bound DECODED memory at the source: the language pipeline's video pixel
+// budget always downscales frames far below 1024px anyway (max_pixels
+// 32·32·768 across t·h·w), so capping the longest edge here loses nothing
+// it would have kept while keeping a 4K clip from materializing ~24 MB per
+// frame (768 frames ≈ 17.8 GiB) before preprocessing.
+gen.maximumSize = CGSize(width: 1024, height: 1024)
+
+try? FileManager.default.createDirectory(
+    atPath: outDir, withIntermediateDirectories: true)
+var written = 0
+var lastW = 0, lastH = 0
+for k in 0..<count {
+    let t = CMTime(seconds: Double(k) / fps, preferredTimescale: 600)
+    let img: CGImage
+    do {
+        img = try gen.copyCGImage(at: t, actualTime: nil)
+    } catch {
+        // Past-end sampling on short clips: stop at the last decodable frame.
+        break
+    }
+    lastW = img.width
+    lastH = img.height
+    let out = URL(fileURLWithPath: String(format: "%@/frame-%04d.png", outDir, k))
+    guard let dest = CGImageDestinationCreateWithURL(
+        out as CFURL, UTType.png.identifier as CFString, 1, nil)
+    else { die("png destination failed") }
+    CGImageDestinationAddImage(dest, img, nil)
+    guard CGImageDestinationFinalize(dest) else { die("png write failed") }
+    written += 1
+}
+guard written > 0 else { die("no frames decodable from \(videoPath)") }
+print("N \(written) W \(lastW) H \(lastH)")

package/src/native-pack.ts

@@ -8,25 +8,32 @@
 //
 // Resolution order for the libs (mirrored in src/mlx/ffi.ts):
 //   MLX_BUN_LIBMLXC env > beside the executable (sidecar dist/ layout,
-//   docs/reference/embedding.md) > the native-pack cache dir > homebrew.
+//   docs/reference/distribution.md) > the native-pack cache dir > homebrew.
 // Dev trees and embedders therefore never download anything.
 
 import { existsSync, mkdirSync, renameSync, rmSync } from "node:fs";
 import { homedir, tmpdir } from "node:os";
 import { dirname, join } from "node:path";
 import { downloadOne } from "./download";
+import { runtimeValue } from "./runtime-config";
 
-export const NATIVE_PACK_VERSION = "0.2.0";
+export const NATIVE_PACK_VERSION = "0.4.0";
+export const MLX_CORE_VERSION = "0.32.2";
 export const NATIVE_PACK_FILES = [
   "libmlxc.dylib", "libmlx.dylib", "libjaccl.dylib", "mlx.metallib",
-  "libmlx_bun_expert_io.dylib",
+  "libmlx_bun_expert_io.dylib", "mlx-bun-frame-extract",
 ] as const;
 
+// RELEASE-TIME BAKE: scripts/build-native-pack.sh prints the sha256/size for
+// the freshly built tarball — publish `native-v0.4.0` with both assets and
+// paste the values here BEFORE tagging the package release (the download
+// 404s until the GitHub release exists; dev trees resolve via homebrew/env
+// and never hit this path).
 const SHA256: Record<string, string> = {
-  arm64: "9bd3795c5ea8f52b18413501f2d68c32c264a20751302300a08dc04cd67df97c",
+  arm64: "f2b0857658aa0826aee14b94a6b83d4c86c08f2c07d716991de6afd038f5ca5d",
 };
 const SIZE: Record<string, number> = {
-  arm64: 52_307_647,
+  arm64: 41931583,
 };
 
 export function nativePackName(arch = process.arch): string {
@@ -37,7 +44,7 @@ export function nativePackName(arch = process.arch): string {
  *  the repo is private the download needs a token (GITHUB_TOKEN or
  *  `gh auth token`); once public it works anonymously. */
 export function nativePackUrl(arch = process.arch): string {
-  return process.env.MLX_BUN_NATIVE_PACK_URL ??
+  return runtimeValue("MLX_BUN_NATIVE_PACK_URL") ??
     `https://github.com/joshuarossi/mlx-bun/releases/download/native-v${NATIVE_PACK_VERSION}/${nativePackName(arch)}`;
 }
 
@@ -48,7 +55,8 @@ export function nativePackDir(arch = process.arch): string {
 /** Where the native runtime currently resolves from, or null if absent
  *  everywhere (= a fresh machine that needs the pack). */
 export function nativeRuntimeDir(): string | null {
-  if (process.env.MLX_BUN_LIBMLXC) return dirname(process.env.MLX_BUN_LIBMLXC);
+  const explicit = runtimeValue("MLX_BUN_LIBMLXC");
+  if (explicit) return dirname(explicit);
   const candidates = [
     dirname(process.execPath),
     nativePackDir(),

package/src/pi-provider.ts

@@ -17,37 +17,11 @@ import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
  *  derive it from the registry's `find()` return type. */
 export type PiModel = NonNullable<ReturnType<ModelRegistry["find"]>>;
 
-// Provider/model constants — kept aligned with src/harness-pi.ts so every
-// path (web embed, terminal embed, external pi) registers the same provider.
-export const PI_PROVIDER_ID = "mlx-bun";
-export const PI_LOCAL_MODEL_ID = "local";
-export const PI_API_KEY = "sk-mlx-bun-local";
-export const PI_API = "openai-completions" as const;
-
-export const DEFAULT_CONTEXT_WINDOW = 32_768;
-export const DEFAULT_MAX_TOKENS = 8_192;
-
-export interface PiProviderOptions {
-  /** Advertised context window. Default: DEFAULT_CONTEXT_WINDOW. */
-  contextWindow?: number;
-  /** Advertised max output tokens. Default: DEFAULT_MAX_TOKENS. */
-  maxTokens?: number;
-  /** Human-readable model label (e.g. the real repo id). Default: "mlx-bun (local)". */
-  name?: string;
-  /** The loaded model has a switchable reasoning channel (server's
-   *  `ctx.template.supportsThinking` — true for any model whose chat template
-   *  gates on enable_thinking: Qwen3.5, MiniCPM5, …). When true Pi engages
-   *  reasoning and sends the on/off switch the server maps to enable_thinking.
-   *  Default false. */
-  reasoning?: boolean;
-  /** The loaded model can accept images (server's `ctx.vision != null`).
-   *  This MUST match the real model: pi-ai serializes user content as OpenAI
-   *  multimodal content-parts (`[{type:"text",…}]`) for image-capable models
-   *  and as a plain string otherwise. A text-only chat template renders nothing
-   *  for a parts array, so declaring a non-vision model as image-capable makes
-   *  the user's turn vanish ("I don't see any message"). Default false. */
-  vision?: boolean;
-}
+import { PI_PROVIDER_ID, PI_API_KEY, PI_API, PI_LOCAL_MODEL_ID, piModelDefinition,
+  type PiModelOptions } from "./contracts/pi-provider";
+export { PI_PROVIDER_ID, PI_API_KEY, PI_API, PI_LOCAL_MODEL_ID,
+  DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from "./contracts/pi-provider";
+export interface PiProviderOptions extends PiModelOptions {}
 
 export interface PiProviderWiring {
   authStorage: AuthStorage;
@@ -70,39 +44,7 @@ export function buildPiProvider(baseUrl: string, opts: PiProviderOptions = {}):
     baseUrl,
     apiKey: PI_API_KEY,
     api: PI_API,
-    models: [
-      {
-        id: PI_LOCAL_MODEL_ID,
-        name: opts.name ?? "mlx-bun (local)",
-        api: PI_API,
-        reasoning: opts.reasoning ?? false,
-        // Compat flags that make pi emit the wire format our local server's
-        // chat templates expect (docs/custom-provider.md "Model Definition
-        // Reference"):
-        //  - supportsDeveloperRole:false → pi sends the system prompt as the
-        //    `system` role, not `developer` (our templates only know
-        //    system/user/assistant/tool).
-        //  - thinkingFormat:"qwen-chat-template" (reasoning models only) → pi
-        //    sends chat_template_kwargs.enable_thinking, the boolean our
-        //    Qwen3.5/MiniCPM5 templates gate on. pi computes the boolean from
-        //    the selected thinking level (off → false), so no thinkingLevelMap.
-        compat: {
-          supportsDeveloperRole: false,
-          ...(opts.reasoning ? { thinkingFormat: "qwen-chat-template" as const } : {}),
-        },
-        // input MUST match the loaded model. pi-ai serializes user content as a
-        // plain string for text-only models and as OpenAI content-parts
-        // (`[{type:"text",…}]`) once a model declares image input — and a
-        // text-only chat template renders NOTHING for a parts array, dropping
-        // the user's turn ("I don't see any message"). So declare image input
-        // only for an actual vision model; the web `ready` frame's `vision`
-        // flag is the same signal that drives the UI's image affordance.
-        input: opts.vision ? ["text", "image"] : ["text"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: opts.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
-        maxTokens: opts.maxTokens ?? DEFAULT_MAX_TOKENS,
-      },
-    ],
+    models: [piModelDefinition(opts)],
   });
 
   const model = modelRegistry.find(PI_PROVIDER_ID, PI_LOCAL_MODEL_ID);