mlx-bun 0.0.13 → 0.3.0

package/src/fit.ts

@@ -10,6 +10,16 @@
 
 import { totalmem } from "node:os";
 import type { ModelConfig } from "./config";
+import type { MemoryPlan } from "./admission-plan";
+import {
+  kvBytesAt,
+  kvGeometry,
+  kvQuantBytesPerElement,
+  type KvSchemeOptions,
+} from "./kv-scheme";
+
+export { kvBytesAt, kvQuantBytesPerElement } from "./kv-scheme";
+export type FitKvScheme = KvSchemeOptions;
 
 /** Decode-efficiency vs theoretical bandwidth ceiling, measured on the
  *  reference machine (24.9 tok/s vs 30.3 ceiling @600 ctx). */
@@ -26,8 +36,6 @@ export const TRANSIENT_PER_TOKEN = 0.55e6;
  *  recommendedMaxWorkingSetSize is ~75% on consumer SKUs). */
 export const WIRED_FRACTION = 0.75;
 export const DEFAULT_CHUNK = 2048;
-/** KV cache element size (bf16; quantized KV lands in Phase 6). */
-const KV_BYTES = 2;
 
 export interface MachineSpec {
   name: string;
@@ -130,43 +138,9 @@ export function chooseAutoModel<T extends AutoPickCandidate>(
   return bySizeDesc.find(fitsCoexistBudget) ?? bySizeDesc.find(fitsFullBudget);
 }
 
-export interface FitReport {
-  fits: boolean;
-  contextTokens: number;
-  weightsBytes: number;
-  kvBytes: number;
-  transientBytes: number;
-  totalBytes: number;
-  usableBytes: number;
-  maxSafeContext: number;
-  predictedDecodeTps: number;
-}
-
-interface KvGeometry {
-  fullLayers: number;
-  slidingLayers: number;
-  fullBytesPerToken: number;
-  slidingBytesPerToken: number;
-  window: number;
-}
-
-function kvGeometry(config: ModelConfig): KvGeometry {
-  const t = config.text;
-  const slidingLayers = t.layerTypes.filter((l) => l === "sliding_attention").length;
-  const fullLayers = t.numHiddenLayers - slidingLayers;
-  const kEqVFactor = t.attentionKEqV ? 2 : 2; // k and v stored separately either way
-  return {
-    fullLayers,
-    slidingLayers,
-    fullBytesPerToken: fullLayers * kEqVFactor * t.numGlobalKeyValueHeads * t.globalHeadDim * KV_BYTES,
-    slidingBytesPerToken: slidingLayers * 2 * t.numKeyValueHeads * t.headDim * KV_BYTES,
-    window: t.slidingWindow,
-  };
-}
-
-export function kvBytesAt(config: ModelConfig, ctx: number): number {
-  const g = kvGeometry(config);
-  return g.fullBytesPerToken * ctx + g.slidingBytesPerToken * Math.min(ctx, g.window);
+export interface FitReport extends MemoryPlan {
+  readonly strategy: "generic-kv";
+  readonly predictedDecodeTps: number;
 }
 
 export function fit(
@@ -182,17 +156,20 @@ export function fit(
    *  replaces the machine-derived usable ceiling (ram × WIRED_FRACTION)
    *  outright — the budget IS the usable envelope. */
   usableBytes?: number,
+  /** Active KV-quant scheme; a quantized cache holds more context in the
+   *  same budget, so the solved ceiling (and admission) must bill it. */
+  kvScheme?: FitKvScheme,
 ): FitReport {
   const usable = usableBytes ?? machine.ramBytes * WIRED_FRACTION;
   const transient = Math.min(chunk, ctx) * TRANSIENT_PER_TOKEN;
-  const kv = kvBytesAt(config, ctx);
+  const kv = kvBytesAt(config, ctx, kvScheme);
   const total = weightsBytes + kv + transient;
 
   // solve max context: weights + kv(ctx) + transient ≤ usable.
   // Below the window both KV terms are linear in ctx; above it the
   // sliding term saturates and only full-attention layers keep growing.
-  const g = kvGeometry(config);
-  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN;
+  const g = kvGeometry(config, kvScheme);
+  const fixed = weightsBytes + chunk * TRANSIENT_PER_TOKEN + g.linearStateBytes;
   let maxCtx = 0;
   if (usable > fixed) {
     const budget = usable - fixed;
@@ -216,11 +193,14 @@ export function fit(
     (isMoe ? MOE_DECODE_EFFICIENCY : DECODE_EFFICIENCY);
 
   return {
+    schemaVersion: 1,
+    strategy: "generic-kv",
     fits: total <= usable,
     contextTokens: ctx,
     weightsBytes,
     kvBytes: kv,
     transientBytes: transient,
+    reserveBytes: 0,
     totalBytes: total,
     usableBytes: usable,
     maxSafeContext: maxCtx,
@@ -231,6 +211,7 @@ export function fit(
 /** The SKU matrix: which Apple Silicon configs run this model at `ctx`. */
 export function skuMatrix(
   config: ModelConfig, weightsBytes: number, ctx: number, expertsBytes = 0,
+  kvScheme?: FitKvScheme,
 ): { sku: string; ramGB: number; fits: boolean; maxContext: number; decodeTps: number }[] {
   const rows: ReturnType<typeof skuMatrix> = [];
   for (const sku of APPLE_SKUS) {
@@ -240,7 +221,7 @@ export function skuMatrix(
         ramBytes: ram * 2 ** 30,
         bandwidthGBs: sku.bandwidthGBs,
       };
-      const r = fit(config, weightsBytes, ctx, m, DEFAULT_CHUNK, expertsBytes);
+      const r = fit(config, weightsBytes, ctx, m, DEFAULT_CHUNK, expertsBytes, undefined, kvScheme);
       rows.push({
         sku: sku.chip,
         ramGB: ram,