mlx-bun 0.2.0 → 0.4.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/turboquant-codec.ts

@@ -0,0 +1,83 @@
+import type { MlxArray } from "../mlx/array";
+import * as tq from "../mlx/turboquant-ops";
+import type { KvCodec } from "../backends/mlx/kv-codec";
+
+/** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
+ *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
+ *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max
+ *  layout (docs/design/turboquant.md), asymmetric-affine for keys,
+ *  FWHT+Lloyd-Max for values. Field order is the kv-store.ts tensor-slot
+ *  contract — do not reorder without updating snapshotCache/loadKvCache. */
+export interface TurboQuantTensor {
+  kIdx: MlxArray;
+  kScales: MlxArray;
+  kZeros: MlxArray;
+  vPacked: MlxArray;
+  vScales: MlxArray;
+}
+
+export const disposeTurboQuant = (t: TurboQuantTensor): void => {
+  t.kIdx.dispose();
+  t.kScales.dispose();
+  t.kZeros.dispose();
+  t.vPacked.dispose();
+  t.vScales.dispose();
+};
+
+/** Numerical representation only. Storage owns capacity, positions, row
+ * membership and persistence; this codec retains the existing kernel choices. */
+export class TurboQuantCodec implements KvCodec<TurboQuantTensor> {
+  constructor(readonly kBits: number, readonly vBits: number, readonly fusedDecode: boolean) {}
+
+  /** Encode k/v [B,H,L,D] into a fresh TurboQuantTensor (packed at rest). */
+  encode(k: MlxArray, v: MlxArray): TurboQuantTensor {
+    const signed = this.kBits === 8;
+    const kEnc = tq.encodeKeys(k, this.kBits, signed);
+    const kIdxPacked = this.kBits < 8 ? tq.packBits(kEnc.indices, this.kBits) : kEnc.indices;
+    if (kIdxPacked !== kEnc.indices) kEnc.indices.dispose();
+    const vEnc = tq.encodeValues(v, this.vBits);
+    const vIdxPacked = this.vBits < 8 ? tq.packBits(vEnc.indices, this.vBits) : vEnc.indices;
+    if (vIdxPacked !== vEnc.indices) vEnc.indices.dispose();
+    return {
+      kIdx: kIdxPacked, kScales: kEnc.scales, kZeros: kEnc.zeros,
+      vPacked: vIdxPacked, vScales: vEnc.scales,
+    };
+  }
+
+  /** Decode a TurboQuantTensor's active [B,H,upTo,*] window back to
+   *  bf16 [B,H,upTo,headDim] (k, v). Caller owns the returned arrays.
+   *  deferV leaves V in the rotated domain (decodeValuesRotated) — the
+   *  caller un-rotates its attention output instead (tq.unrotateValues). */
+  decode(t: TurboQuantTensor, upTo: number, headDim: number, deferV = false): [MlxArray, MlxArray] {
+    const cut = (a: MlxArray): MlxArray => {
+      const [B, H, , D] = a.shape as [number, number, number, number];
+      return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
+    };
+    if (this.fusedDecode) {
+      const inputs = [cut(t.kIdx), cut(t.kScales), cut(t.kZeros), cut(t.vPacked), cut(t.vScales)] as const;
+      try {
+        const decoded = tq.tryDecodePackedKv(inputs, this.kBits, this.vBits, headDim, deferV);
+        if (decoded) return decoded;
+      } finally { for (const input of inputs) input.dispose(); }
+    }
+    const kIdxCut = cut(t.kIdx);
+    const kIdxUnpacked = this.kBits < 8 ? tq.unpackBits(kIdxCut, this.kBits, headDim) : kIdxCut;
+    if (kIdxUnpacked !== kIdxCut) kIdxCut.dispose();
+    const kScalesCut = cut(t.kScales);
+    const kZerosCut = cut(t.kZeros);
+    const k = tq.decodeKeys(kIdxUnpacked, kScalesCut, kZerosCut);
+    for (const a of [kIdxUnpacked, kScalesCut, kZerosCut]) a.dispose();
+
+    const vPackedCut = cut(t.vPacked);
+    const vIdxUnpacked = this.vBits < 8 ? tq.unpackBits(vPackedCut, this.vBits, headDim) : vPackedCut;
+    if (vIdxUnpacked !== vPackedCut) vPackedCut.dispose();
+    const vScalesCut = cut(t.vScales);
+    const v = deferV
+      ? tq.decodeValuesRotated(vIdxUnpacked, vScalesCut, this.vBits)
+      : tq.decodeValues(vIdxUnpacked, vScalesCut, this.vBits);
+    for (const a of [vIdxUnpacked, vScalesCut]) a.dispose();
+
+    return [k, v];
+  }
+
+}

package/src/model/universal/dense.ts

@@ -1,9 +1,11 @@
+import { captureKvAttention } from "../kv-attention-view";
 // UniversalDense — the Tier-0 config-driven dense llama-family module
 // (docs/design/generic-model-support.md §3.1). One module, shaped exactly
 // 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).
+// decode. Standard attention uses the shared plain/affine KV port; the
+// softcap branch retains plain KV. The bar is L1 bit-exactness
+// 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,
@@ -154,19 +156,18 @@ class UniversalAttention {
       ? this.rope.applyDynamic(k, offArr)
       : this.rope.apply(k, cache.offset));
 
-    const [keys, values] = cache.updateAndFetch(k, v);
-    k.dispose();
-    v.dispose();
-
     let attn: MlxArray;
-    if (a.attnLogitSoftcap !== null) {
-      attn = this.#softcapAttention(q, keys, values, mask, B, L);
-    } else {
-      attn = ops.sdpa(q, keys, values, a.attnScale, mask.mode, mask.arr);
-    }
-    keys.dispose();
-    values.dispose();
-    q.dispose();
+    try {
+      if (a.attnLogitSoftcap !== null) {
+        const [keys, values] = cache.updateAndFetch(k, v);
+        try { attn = this.#softcapAttention(q, keys, values, mask, B, L); }
+        finally { keys.dispose(); values.dispose(); }
+      } else {
+        const view = captureKvAttention(cache, k, v);
+        try { attn = view.attend(q, a.attnScale, mask); }
+        finally { view.dispose(); }
+      }
+    } finally { k.dispose(); v.dispose(); q.dispose(); }
 
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();

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-pack.ts

@@ -8,7 +8,7 @@
 //
 // 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";
@@ -17,22 +17,23 @@ import { dirname, join } from "node:path";
 import { downloadOne } from "./download";
 import { runtimeValue } from "./runtime-config";
 
-export const NATIVE_PACK_VERSION = "0.3.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", "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.3.0` with both assets and
+// 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: "b485fed858971ea6f072f349a0234c0e2d5030d3ef2ddfa942c8975eca2d3b2e",
+  arm64: "f2b0857658aa0826aee14b94a6b83d4c86c08f2c07d716991de6afd038f5ca5d",
 };
 const SIZE: Record<string, number> = {
-  arm64: 52299265,
+  arm64: 41931583,
 };
 
 export function nativePackName(arch = process.arch): string {

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);