mlx-bun 0.5.0 → 1.0.0

package/src/mlx/turboquant-tables.ts

@@ -1,237 +0,0 @@
-// GENERATED by scripts/regen.ts turboquant — DO NOT EDIT.
-// Regenerate:  bun scripts/regen.ts turboquant
-//
-// Source: lab/repro/vllm-metal-turboquant/turboquant_reference.py (vendored
-// vllm-metal oracle, see PROVENANCE.md in that directory), executed under
-// the oracle venv (mlx 0.31.2) via gen_goldens.py.
-//
-// - SIGN_VECTORS[headDim]: the fwht() sign table (±1), derived from
-//   mx.random.randint(0, 2, shape=(headDim,), key=mx.random.key(42)) —
-//   deterministic across machines and spill/restore by construction.
-// - LLOYD_MAX[bits]: Lloyd-Max (centroids, boundaries) for value
-//   quantization at the given bit width. bits=3 is the vendored reference's
-//   hardcoded fast-path table (CENTROIDS_3BIT/BOUNDARIES_3BIT); other widths
-//   are the reference's own iterative Lloyd-Max fit on a unit-normal source
-//   (mx.random.normal(key=mx.random.key(0)), 500 iterations), computed once
-//   here and hardcoded — no runtime iteration in the TS port.
-
-export const TURBOQUANT_HEAD_DIMS = [64, 128, 256, 512] as const;
-export type TurboQuantHeadDim = (typeof TURBOQUANT_HEAD_DIMS)[number];
-
-export const SIGN_VECTORS: Record<TurboQuantHeadDim, number[]> = {
-  64: [
-  1, -1, -1, -1, 1, 1, 1, -1, 1, 1,
-  -1, 1, -1, 1, -1, 1, -1, 1, 1, -1,
-  1, -1, -1, 1, -1, -1, -1, 1, -1, -1,
-  1, -1, 1, -1, 1, 1, 1, 1, 1, 1,
-  1, 1, 1, -1, 1, 1, -1, 1, 1, -1,
-  -1, -1, -1, 1, -1, -1, 1, -1, 1, 1,
-  1, -1, -1, 1,
-  ],
-  128: [
-  -1, 1, -1, -1, 1, 1, 1, 1, -1, -1,
-  -1, -1, -1, 1, -1, -1, -1, 1, 1, -1,
-  -1, -1, 1, 1, 1, -1, 1, -1, 1, -1,
-  -1, 1, 1, 1, 1, 1, -1, 1, -1, 1,
-  1, -1, 1, 1, 1, 1, 1, 1, 1, 1,
-  -1, -1, -1, 1, 1, -1, 1, -1, 1, -1,
-  1, -1, -1, 1, -1, 1, -1, 1, -1, 1,
-  -1, -1, -1, -1, -1, 1, -1, -1, 1, 1,
-  -1, -1, 1, 1, 1, 1, 1, -1, 1, -1,
-  1, -1, 1, 1, -1, 1, 1, -1, -1, -1,
-  -1, 1, -1, -1, -1, 1, 1, -1, -1, 1,
-  -1, 1, 1, 1, 1, 1, 1, -1, -1, 1,
-  1, 1, -1, 1, -1, 1, 1, 1,
-  ],
-  256: [
-  1, -1, 1, 1, -1, 1, 1, 1, 1, 1,
-  1, 1, -1, -1, 1, 1, 1, -1, -1, 1,
-  1, -1, -1, -1, 1, -1, 1, 1, 1, -1,
-  -1, 1, -1, 1, 1, 1, -1, 1, 1, 1,
-  1, -1, 1, -1, 1, 1, -1, 1, -1, 1,
-  -1, -1, 1, 1, -1, -1, 1, 1, 1, 1,
-  -1, 1, 1, 1, -1, -1, -1, -1, 1, -1,
-  -1, 1, 1, 1, -1, 1, -1, 1, 1, 1,
-  -1, 1, -1, -1, 1, -1, -1, 1, -1, 1,
-  1, -1, 1, -1, 1, -1, -1, 1, 1, -1,
-  1, 1, -1, -1, 1, -1, -1, -1, -1, 1,
-  -1, 1, 1, 1, -1, -1, -1, 1, 1, -1,
-  -1, -1, 1, 1, -1, 1, -1, -1, -1, -1,
-  1, -1, -1, 1, 1, -1, 1, -1, 1, 1,
-  -1, -1, -1, -1, -1, -1, -1, -1, 1, -1,
-  1, 1, 1, -1, 1, -1, -1, -1, 1, 1,
-  -1, -1, 1, 1, -1, -1, -1, 1, -1, 1,
-  -1, -1, -1, 1, -1, 1, -1, -1, -1, 1,
-  1, -1, 1, 1, 1, -1, 1, 1, -1, -1,
-  1, -1, -1, -1, 1, -1, 1, -1, -1, 1,
-  -1, 1, 1, -1, -1, 1, -1, -1, 1, 1,
-  1, 1, -1, -1, 1, 1, -1, -1, -1, 1,
-  -1, 1, -1, 1, -1, -1, -1, 1, -1, 1,
-  -1, 1, -1, -1, 1, 1, 1, 1, -1, -1,
-  -1, -1, 1, -1, 1, 1, 1, -1, 1, -1,
-  -1, -1, -1, -1, -1, -1,
-  ],
-  512: [
-  -1, 1, -1, -1, -1, -1, 1, -1, -1, -1,
-  1, -1, 1, 1, -1, 1, 1, -1, -1, 1,
-  -1, -1, 1, 1, 1, 1, 1, 1, -1, -1,
-  -1, 1, -1, 1, -1, -1, 1, 1, 1, 1,
-  -1, -1, 1, 1, 1, -1, -1, -1, -1, -1,
-  1, -1, -1, -1, -1, -1, 1, -1, 1, 1,
-  1, -1, -1, 1, -1, -1, 1, 1, 1, -1,
-  -1, 1, 1, 1, -1, 1, 1, -1, -1, -1,
-  -1, 1, 1, -1, -1, -1, 1, 1, 1, 1,
-  -1, 1, -1, 1, -1, -1, -1, 1, 1, 1,
-  1, -1, -1, 1, 1, -1, 1, 1, -1, 1,
-  -1, -1, 1, -1, 1, -1, 1, 1, -1, -1,
-  1, 1, -1, -1, -1, 1, -1, 1, -1, -1,
-  -1, -1, 1, 1, 1, -1, 1, 1, -1, -1,
-  -1, -1, -1, 1, -1, -1, -1, 1, 1, 1,
-  -1, 1, -1, -1, 1, -1, -1, 1, 1, -1,
-  -1, -1, 1, 1, 1, 1, -1, 1, 1, -1,
-  1, 1, -1, 1, 1, 1, -1, 1, -1, -1,
-  -1, 1, -1, -1, -1, 1, 1, -1, 1, -1,
-  -1, -1, 1, 1, 1, 1, -1, -1, 1, 1,
-  1, -1, -1, -1, -1, 1, -1, -1, 1, -1,
-  1, -1, 1, -1, 1, 1, 1, 1, -1, 1,
-  -1, 1, -1, 1, 1, -1, 1, -1, -1, -1,
-  1, -1, 1, -1, 1, -1, 1, 1, -1, -1,
-  1, -1, 1, -1, -1, -1, 1, -1, -1, -1,
-  -1, 1, -1, -1, -1, 1, 1, 1, -1, 1,
-  1, 1, 1, -1, 1, -1, -1, -1, -1, -1,
-  1, 1, -1, -1, -1, 1, 1, 1, 1, -1,
-  -1, 1, 1, 1, -1, 1, 1, 1, 1, 1,
-  1, -1, -1, 1, 1, 1, -1, 1, -1, 1,
-  -1, 1, 1, 1, 1, 1, -1, -1, 1, 1,
-  -1, 1, 1, -1, 1, -1, -1, 1, 1, -1,
-  -1, 1, -1, 1, 1, 1, -1, -1, -1, -1,
-  1, 1, 1, -1, 1, -1, 1, 1, -1, 1,
-  1, 1, 1, 1, 1, -1, 1, 1, -1, -1,
-  1, 1, -1, -1, -1, -1, -1, 1, 1, 1,
-  1, -1, 1, 1, -1, -1, -1, 1, -1, 1,
-  1, 1, 1, 1, 1, -1, -1, 1, -1, -1,
-  -1, 1, -1, 1, 1, -1, 1, 1, -1, 1,
-  -1, 1, 1, 1, -1, 1, -1, 1, 1, 1,
-  1, -1, 1, 1, 1, 1, 1, -1, 1, -1,
-  1, -1, 1, 1, 1, -1, 1, -1, 1, 1,
-  1, 1, -1, -1, -1, -1, 1, -1, -1, 1,
-  1, -1, -1, -1, -1, 1, 1, -1, 1, -1,
-  1, -1, -1, -1, 1, 1, 1, -1, -1, -1,
-  -1, -1, 1, -1, -1, -1, -1, -1, 1, -1,
-  1, 1, -1, 1, 1, -1, -1, -1, 1, 1,
-  1, 1, 1, -1, 1, 1, -1, 1, -1, 1,
-  -1, 1, -1, -1, 1, -1, -1, 1, -1, 1,
-  -1, -1, -1, 1, 1, -1, -1, -1, 1, -1,
-  1, -1, -1, 1, -1, 1, 1, -1, 1, -1,
-  1, 1,
-  ],
-};
-
-export const TURBOQUANT_VALUE_BITS = [2, 3, 4, 5, 8] as const;
-export type TurboQuantValueBits = (typeof TURBOQUANT_VALUE_BITS)[number];
-
-export interface LloydMaxTable {
-  centroids: number[];
-  boundaries: number[];
-}
-
-export const LLOYD_MAX: Record<TurboQuantValueBits, LloydMaxTable> = {
-  2: {
-    centroids: [
-  -1.5144749879837036, -0.45822522044181824, 0.44863229990005493, 1.512152075767517,
-    ],
-    boundaries: [
-  -0.9863501191139221, -0.004796460270881653, 0.9803922176361084,
-    ],
-  },
-  3: {
-    centroids: [
-  -2.151949882507324, -1.3439099788665771, -0.7560099959373474, -0.24508999288082123, 0.24508999288082123, 0.7560099959373474, 1.3439099788665771, 2.151949882507324,
-    ],
-    boundaries: [
-  -1.7479300498962402, -1.0499600172042847, -0.5005499720573425, 0, 0.5005499720573425, 1.0499600172042847, 1.7479300498962402,
-    ],
-  },
-  4: {
-    centroids: [
-  -2.765662670135498, -2.0995230674743652, -1.6401931047439575, -1.27662992477417, -0.9602491855621338, -0.673130452632904, -0.4038185179233551, -0.14272712171077728, 0.11716290563344955, 0.3790881633758545,
-  0.6475776433944702, 0.9337777495384216, 1.2523075342178345, 1.6217625141143799, 2.0757482051849365, 2.7207279205322266,
-    ],
-    boundaries: [
-  -2.4325928688049316, -1.8698580265045166, -1.458411455154419, -1.1184395551681519, -0.8166898488998413, -0.5384745001792908, -0.2732728123664856, -0.012782108038663864, 0.24812553822994232, 0.5133329033851624,
-  0.7906776666641235, 1.0930426120758057, 1.437035083770752, 1.8487553596496582, 2.398238182067871,
-    ],
-  },
-  5: {
-    centroids: [
-  -3.3482043743133545, -2.7744688987731934, -2.3917815685272217, -2.0885653495788574, -1.8382762670516968, -1.6273508071899414, -1.436827301979065, -1.2599420547485352, -1.0958681106567383, -0.9411738514900208,
-  -0.7932286262512207, -0.6509044766426086, -0.5113211274147034, -0.37372007966041565, -0.23732160031795502, -0.10278966277837753, 0.031794287264347076, 0.16419854760169983, 0.3001006245613098, 0.4395425021648407,
-  0.5808256268501282, 0.7275802493095398, 0.8738965392112732, 1.029397964477539, 1.1942607164382935, 1.3728704452514648, 1.5676456689834595, 1.778881311416626, 2.0225136280059814, 2.3095157146453857,
-  2.666351556777954, 3.2329483032226562,
-    ],
-    boundaries: [
-  -3.0613365173339844, -2.583125114440918, -2.24017333984375, -1.9634208679199219, -1.7328135967254639, -1.5320889949798584, -1.3483846187591553, -1.1779050827026367, -1.0185209512710571, -0.8672012090682983,
-  -0.7220665216445923, -0.581112802028656, -0.4425206184387207, -0.30552083253860474, -0.17005562782287598, -0.03549768775701523, 0.09799641370773315, 0.23214958608150482, 0.36982154846191406, 0.5101840496063232,
-  0.654202938079834, 0.8007383942604065, 0.9516472816467285, 1.1118292808532715, 1.2835655212402344, 1.4702579975128174, 1.6732635498046875, 1.9006974697113037, 2.1660146713256836, 2.48793363571167,
-  2.9496498107910156,
-    ],
-  },
-  8: {
-    centroids: [
-  -4.505678653717041, -4.454618453979492, -4.436285972595215, -4.401589393615723, -4.363239765167236, -4.332197189331055, -4.2975006103515625, -4.26280403137207, -4.228107929229736, -4.193411827087402,
-  -4.15871524810791, -4.124018669128418, -4.094437599182129, -4.045758247375488, -4.002621173858643, -4.002171516418457, -3.942392110824585, -3.916801691055298, -3.8863255977630615, -3.8477535247802734,
-  -3.7988643646240234, -3.762794256210327, -3.7445151805877686, -3.706960439682007, -3.67462158203125, -3.6314854621887207, -3.5950119495391846, -3.580256938934326, -3.539952516555786, -3.5038421154022217,
-  -3.4671568870544434, -3.4281671047210693, -3.398128032684326, -3.363133192062378, -3.3274667263031006, -3.2919743061065674, -3.2577857971191406, -3.2200968265533447, -3.192981719970703, -3.157238006591797,
-  -3.111802577972412, -3.0821340084075928, -3.050640344619751, -3.0148110389709473, -2.980158805847168, -2.946451425552368, -2.9111905097961426, -2.8747212886810303, -2.8412270545959473, -2.8061022758483887,
-  -2.7716867923736572, -2.7376010417938232, -2.7029218673706055, -2.6634461879730225, -2.630377769470215, -2.5996179580688477, -2.560898780822754, -2.525275230407715, -2.4947545528411865, -2.461374044418335,
-  -2.424009323120117, -2.3877711296081543, -2.3505144119262695, -2.3163986206054688, -2.282374858856201, -2.243652820587158, -2.2097959518432617, -2.1770639419555664, -2.14471697807312, -2.1113693714141846,
-  -2.078673839569092, -2.044494390487671, -2.00639271736145, -1.969869613647461, -1.9300988912582397, -1.8931699991226196, -1.8589746952056885, -1.8212106227874756, -1.7848811149597168, -1.7465486526489258,
-  -1.7103406190872192, -1.6743196249008179, -1.6393754482269287, -1.6040517091751099, -1.5708216428756714, -1.5365480184555054, -1.5035663843154907, -1.469175100326538, -1.4350656270980835, -1.4022008180618286,
-  -1.3688998222351074, -1.3355891704559326, -1.3018449544906616, -1.2698367834091187, -1.235532283782959, -1.20119047164917, -1.1664694547653198, -1.1305049657821655, -1.0959333181381226, -1.0614327192306519,
-  -1.0274933576583862, -0.9940271377563477, -0.9609262347221375, -0.9273853898048401, -0.8929082751274109, -0.8593462705612183, -0.8261474967002869, -0.7937442660331726, -0.7588282823562622, -0.7230017185211182,
-  -0.6871297359466553, -0.6517745852470398, -0.6162516474723816, -0.5811514854431152, -0.5462668538093567, -0.5111281275749207, -0.4750610589981079, -0.4394458830356598, -0.404982328414917, -0.36981964111328125,
-  -0.334137886762619, -0.29919183254241943, -0.26465246081352234, -0.23168867826461792, -0.20013625919818878, -0.1672147512435913, -0.13413792848587036, -0.1003282368183136, -0.06571532040834427, -0.03132394701242447,
-  0.0043428270146250725, 0.04037246108055115, 0.07451506704092026, 0.1076153963804245, 0.14043505489826202, 0.17426857352256775, 0.20812825858592987, 0.2430444061756134, 0.27797263860702515, 0.31494686007499695,
-  0.35088568925857544, 0.38743987679481506, 0.4222983717918396, 0.4568648338317871, 0.49128350615501404, 0.5256633758544922, 0.5586227178573608, 0.5910427570343018, 0.6249552369117737, 0.6594945788383484,
-  0.6942835450172424, 0.7278474569320679, 0.7624561786651611, 0.7963849902153015, 0.8309887647628784, 0.8644483685493469, 0.8981063961982727, 0.9320245385169983, 0.96662437915802, 1.0006996393203735,
-  1.0343793630599976, 1.0680633783340454, 1.101860523223877, 1.135563611984253, 1.1707100868225098, 1.2038687467575073, 1.2388564348220825, 1.2744282484054565, 1.3102103471755981, 1.3441187143325806,
-  1.3769688606262207, 1.4091657400131226, 1.4424594640731812, 1.477047085762024, 1.5127589702606201, 1.5482075214385986, 1.58390474319458, 1.6193201541900635, 1.655400037765503, 1.6924738883972168,
-  1.7273176908493042, 1.7618744373321533, 1.7985631227493286, 1.8349602222442627, 1.871187686920166, 1.9062069654464722, 1.9427249431610107, 1.9807021617889404, 2.0164709091186523, 2.049800395965576,
-  2.083641529083252, 2.1192877292633057, 2.1556057929992676, 2.189804792404175, 2.2249133586883545, 2.257265329360962, 2.2943572998046875, 2.3308892250061035, 2.3646092414855957, 2.3981246948242188,
-  2.4336130619049072, 2.4672765731811523, 2.5033960342407227, 2.535766124725342, 2.5712199211120605, 2.606968879699707, 2.6403279304504395, 2.6730761528015137, 2.7094709873199463, 2.745285749435425,
-  2.7794594764709473, 2.8145666122436523, 2.851815938949585, 2.885251045227051, 2.919672727584839, 2.9549803733825684, 2.987560510635376, 3.021491289138794, 3.061659097671509, 3.102267026901245,
-  3.1455001831054688, 3.1724164485931396, 3.2052037715911865, 3.2377841472625732, 3.2674646377563477, 3.3014495372772217, 3.3345274925231934, 3.3695874214172363, 3.4068915843963623, 3.442040205001831,
-  3.4636645317077637, 3.5024828910827637, 3.5508108139038086, 3.5749337673187256, 3.604726791381836, 3.6401684284210205, 3.692674160003662, 3.7173542976379395, 3.7488019466400146, 3.7978620529174805,
-  3.8167099952697754, 3.856578826904297, 3.890836000442505, 3.9217264652252197, 3.974006175994873, 4.009003639221191, 4.029621601104736, 4.079561233520508, 4.099013805389404, 4.1337103843688965,
-  4.172705173492432, 4.198393821716309, 4.237799167633057, 4.272495746612549, 4.307192325592041, 4.341888427734375,
-    ],
-    boundaries: [
-  -4.4801483154296875, -4.4454522132873535, -4.418937683105469, -4.382414817810059, -4.347718238830566, -4.314848899841309, -4.280152320861816, -4.245455741882324, -4.210760116577148, -4.176063537597656,
-  -4.141366958618164, -4.109228134155273, -4.070097923278809, -4.0241899490356445, -4.002396583557129, -3.9722819328308105, -3.9295969009399414, -3.9015636444091797, -3.867039680480957, -3.8233089447021484,
-  -3.780829429626465, -3.753654718399048, -3.7257378101348877, -3.690791130065918, -3.6530535221099854, -3.613248825073242, -3.587634563446045, -3.5601048469543457, -3.521897315979004, -3.485499382019043,
-  -3.447661876678467, -3.413147449493408, -3.3806304931640625, -3.3452999591827393, -3.309720516204834, -3.2748799324035645, -3.238941192626953, -3.2065391540527344, -3.17510986328125, -3.1345202922821045,
-  -3.096968173980713, -3.066387176513672, -3.0327258110046387, -2.9974849224090576, -2.9633049964904785, -2.928821086883545, -2.892955780029297, -2.857974052429199, -2.823664665222168, -2.7888946533203125,
-  -2.7546439170837402, -2.720261573791504, -2.6831841468811035, -2.646912097930908, -2.6149978637695312, -2.580258369445801, -2.5430870056152344, -2.5100150108337402, -2.4780642986297607, -2.4426918029785156,
-  -2.4058902263641357, -2.369142770767212, -2.333456516265869, -2.299386739730835, -2.2630138397216797, -2.22672438621521, -2.193429946899414, -2.160890579223633, -2.1280431747436523, -2.0950217247009277,
-  -2.061583995819092, -2.0254435539245605, -1.9881311655044556, -1.9499843120574951, -1.9116344451904297, -1.8760724067687988, -1.840092658996582, -1.8030458688735962, -1.7657148838043213, -1.7284445762634277,
-  -1.6923301219940186, -1.6568474769592285, -1.621713638305664, -1.5874366760253906, -1.5536848306655884, -1.520057201385498, -1.4863708019256592, -1.452120304107666, -1.418633222579956, -1.3855502605438232,
-  -1.35224449634552, -1.3187170028686523, -1.2858408689498901, -1.2526845932006836, -1.2183613777160645, -1.1838300228118896, -1.1484872102737427, -1.113219141960144, -1.0786830186843872, -1.044463038444519,
-  -1.0107603073120117, -0.9774767160415649, -0.9441558122634888, -0.9101468324661255, -0.8761272430419922, -0.8427468538284302, -0.8099458813667297, -0.776286244392395, -0.7409150004386902, -0.7050657272338867,
-  -0.6694521903991699, -0.6340131163597107, -0.5987015962600708, -0.5637091398239136, -0.5286974906921387, -0.4930945932865143, -0.45725345611572266, -0.4222140908241272, -0.3874009847640991, -0.35197877883911133,
-  -0.3166648745536804, -0.2819221615791321, -0.24817056953907013, -0.21591246128082275, -0.18367549777030945, -0.15067633986473083, -0.11723308265209198, -0.08302177488803864, -0.04851963371038437, -0.013490559533238411,
-  0.022357644513249397, 0.0574437640607357, 0.09106522798538208, 0.12402522563934326, 0.15735182166099548, 0.1911984086036682, 0.22558632493019104, 0.26050853729248047, 0.29645973443984985, 0.332916259765625,
-  0.36916279792785645, 0.4048691391944885, 0.43958160281181335, 0.47407418489456177, 0.5084734559059143, 0.5421430468559265, 0.5748327374458313, 0.6079989671707153, 0.642224907875061, 0.6768890619277954,
-  0.7110655307769775, 0.7451518177986145, 0.7794206142425537, 0.8136868476867676, 0.8477185964584351, 0.8812773823738098, 0.9150654673576355, 0.9493244886398315, 0.9836620092391968, 1.0175395011901855,
-  1.0512213706970215, 1.0849618911743164, 1.118712067604065, 1.1531368494033813, 1.1872894763946533, 1.221362590789795, 1.2566423416137695, 1.2923192977905273, 1.3271645307540894, 1.3605437278747559,
-  1.3930673599243164, 1.4258126020431519, 1.4597532749176025, 1.4949030876159668, 1.5304832458496094, 1.5660561323165894, 1.6016124486923218, 1.6373600959777832, 1.6739369630813599, 1.7098958492279053,
-  1.744596004486084, 1.7802188396453857, 1.8167617321014404, 1.8530739545822144, 1.8886973857879639, 1.9244658946990967, 1.9617135524749756, 1.9985865354537964, 2.0331356525421143, 2.066720962524414,
-  2.1014647483825684, 2.137446880340576, 2.1727051734924316, 2.2073590755462646, 2.241089344024658, 2.275811195373535, 2.3126232624053955, 2.3477492332458496, 2.3813669681549072, 2.4158687591552734,
-  2.4504446983337402, 2.4853363037109375, 2.5195810794830322, 2.553493022918701, 2.589094400405884, 2.6236484050750732, 2.6567020416259766, 2.6912736892700195, 2.7273783683776855, 2.7623724937438965,
-  2.7970130443573, 2.833191394805908, 2.8685336112976074, 2.9024620056152344, 2.937326431274414, 2.9712705612182617, 3.004525899887085, 3.0415751934051514, 3.081963062286377, 3.1238837242126465,
-  3.1589584350585938, 3.188810110092163, 3.22149395942688, 3.25262451171875, 3.284457206726074, 3.317988395690918, 3.352057456970215, 3.3882393836975098, 3.4244658946990967, 3.452852249145508,
-  3.4830737113952637, 3.526646852493286, 3.5628724098205566, 3.5898303985595703, 3.6224474906921387, 3.666421413421631, 3.705014228820801, 3.7330780029296875, 3.773332118988037, 3.807286024093628,
-  3.836644411087036, 3.8737072944641113, 3.9062812328338623, 3.947866439819336, 3.9915049076080322, 4.019312858581543, 4.054591178894043, 4.089287757873535, 4.11636209487915, 4.153207778930664,
-  4.185549736022949, 4.218096733093262, 4.255147457122803, 4.289844036102295, 4.324540138244629,
-    ],
-  },
-};

package/src/mmap.ts

@@ -1,90 +0,0 @@
-// libc mmap via bun:ffi.
-//
-// Bun.mmap panics (SIGTRAP, Bun 1.3.3) on files larger than 4 GB because
-// JSC ArrayBuffers cap at 2^32 bytes — and weight shards routinely exceed
-// that. We don't need a whole-file ArrayBuffer anyway: tensors go to mlx
-// as raw pointers (base + offset), and only small ranges (headers) are
-// ever viewed from JS.
-
-import { dlopen, FFIType, toArrayBuffer } from "bun:ffi";
-import { closeSync, fstatSync, openSync } from "node:fs";
-
-const libc = dlopen("/usr/lib/libSystem.B.dylib", {
-  mmap: {
-    args: [FFIType.u64, FFIType.u64, FFIType.i32, FFIType.i32, FFIType.i32, FFIType.i64],
-    returns: FFIType.u64,
-  },
-  munmap: { args: [FFIType.u64, FFIType.u64], returns: FFIType.i32 },
-  madvise: { args: [FFIType.u64, FFIType.u64, FFIType.i32], returns: FFIType.i32 },
-}).symbols;
-
-const PROT_READ = 0x1;
-const PROT_WRITE = 0x2;
-const MAP_SHARED = 0x1;
-const MAP_PRIVATE = 0x2;
-const MAP_FAILED = 0xffffffffffffffffn;
-
-// madvise advice (macOS <sys/mman.h>). DONTNEED drops clean file-backed
-// pages (they re-fault from the file); FREE_REUSABLE is the Apple-specific
-// lazy-reclaim variant.
-export const MADV_DONTNEED = 4;
-export const MADV_FREE = 5;
-export const MADV_FREE_REUSABLE = 7;
-
-export class MmapFile {
-  readonly path: string;
-  readonly base: bigint;
-  readonly size: number;
-  #mapped = true;
-
-  private constructor(path: string, base: bigint, size: number) {
-    this.path = path;
-    this.base = base;
-    this.size = size;
-  }
-
-  /** Open read-only (MAP_SHARED) or copy-on-write (MAP_PRIVATE +
-   *  PROT_WRITE — writes go to private pages, never the file; used for
-   *  KV-cache reloads where mlx may donate buffers). */
-  static open(path: string, mode: "ro" | "cow" = "ro"): MmapFile {
-    const fd = openSync(path, "r");
-    try {
-      const size = fstatSync(fd).size;
-      if (size === 0) throw new Error(`${path}: empty file`);
-      const prot = mode === "cow" ? PROT_READ | PROT_WRITE : PROT_READ;
-      const flags = mode === "cow" ? MAP_PRIVATE : MAP_SHARED;
-      const base = libc.mmap(0n, BigInt(size), prot, flags, fd, 0n);
-      if (base === MAP_FAILED) throw new Error(`${path}: mmap failed`);
-      return new MmapFile(path, base, size);
-    } finally {
-      closeSync(fd); // mapping survives close
-    }
-  }
-
-  /** Raw pointer to byte `offset` (for FFI consumers like mlx). */
-  pointer(offset: number): number {
-    if (offset < 0 || offset > this.size) throw new RangeError(`offset ${offset} out of range`);
-    return Number(this.base + BigInt(offset));
-  }
-
-  /** JS view of a range. Only for small ranges (headers, spot checks);
-   *  the view aliases the mapping — do not retain past munmap. */
-  view(offset: number, length: number): Uint8Array {
-    if (offset + length > this.size) throw new RangeError("view out of range");
-    if (length >= 2 ** 32) throw new RangeError("JS views cap at 4 GB");
-    return new Uint8Array(toArrayBuffer(this.pointer(offset) as never, 0, length));
-  }
-
-  /** madvise a page-aligned sub-range (offset + length should be page
-   *  multiples). e.g. advise(off, len, MADV_DONTNEED) to release clean
-   *  expert pages back to the OS; they re-fault from the file on next read. */
-  advise(offset: number, length: number, advice: number): number {
-    return Number(libc.madvise(this.base + BigInt(offset), BigInt(length), advice));
-  }
-
-  unmap(): void {
-    if (!this.#mapped) return;
-    this.#mapped = false;
-    libc.munmap(this.base, BigInt(this.size));
-  }
-}

package/src/model/artifact-identity.ts

@@ -1,73 +0,0 @@
-import { createHash, randomUUID } from "node:crypto";
-import { mkdir, readFile, realpath, rename, rm, stat, writeFile } from "node:fs/promises";
-import { homedir } from "node:os";
-import { join } from "node:path";
-
-/** Identity calculation owns bytes; this interface owns memo persistence. */
-export interface ArtifactIdentityStore {
-  get(key: string): Promise<string | undefined>;
-  put(key: string, digest: string): void;
-}
-
-export class FileArtifactIdentityStore implements ArtifactIdentityStore {
-  readonly #memory = new Map<string, string>();
-  readonly #pending = new Set<Promise<void>>();
-  constructor(readonly directory: string) {}
-
-  async get(key: string): Promise<string | undefined> {
-    const memory = this.#memory.get(key);
-    if (memory) return memory;
-    try {
-      const digest = await readFile(join(this.directory, key), "utf8");
-      if (!/^[0-9a-f]{64}$/.test(digest)) return undefined;
-      this.#memory.set(key, digest);
-      return digest;
-    } catch { return undefined; }
-  }
-
-  put(key: string, digest: string): void {
-    this.#memory.set(key, digest);
-    const temporary = join(this.directory, `${key}.${process.pid}.${randomUUID()}.tmp`);
-    const write = async () => {
-      try {
-        await mkdir(this.directory, { recursive: true });
-        await writeFile(temporary, digest);
-        await rename(temporary, join(this.directory, key));
-      } catch { await rm(temporary, { force: true }).catch(() => {}); }
-    };
-    const pending = write().finally(() => this.#pending.delete(pending));
-    this.#pending.add(pending);
-  }
-
-  async flush(): Promise<void> { await Promise.all(this.#pending); }
-}
-
-const defaultStore = new FileArtifactIdentityStore(join(homedir(), ".cache", "mlx-bun", "artifact-identities"));
-export interface ArtifactIdentityFile { name: string; path: string; }
-
-async function revision(files: readonly ArtifactIdentityFile[]): Promise<string> {
-  return JSON.stringify(await Promise.all(files.map(async file => {
-    const path = await realpath(file.path), info = await stat(path, { bigint: true });
-    return [file.name, path, info.dev, info.ino, info.size, info.mtimeNs, info.ctimeNs].map(String);
-  })));
-}
-
-/** Preserve the existing seed/name/byte digest, reusing it while file revisions
- * are unchanged. ctime catches same-size edits even when mtime is restored.
- * A memo miss computes the full identity; persistence never delays its caller. */
-export async function artifactIdentity(seed: string, input: readonly ArtifactIdentityFile[],
-  store: ArtifactIdentityStore = defaultStore): Promise<string> {
-  const files = [...input].sort((a, b) => a.name < b.name ? -1 : a.name > b.name ? 1 : 0);
-  const before = await revision(files);
-  const key = createHash("sha256").update(JSON.stringify([seed, before])).digest("hex");
-  const cached = await store.get(key);
-  if (cached) return cached;
-  const digest = createHash("sha256").update(seed);
-  for (const file of files) {
-    digest.update(file.name);
-    for await (const chunk of Bun.file(file.path).stream()) digest.update(chunk);
-  }
-  const value = digest.digest("hex");
-  if (before === await revision(files)) store.put(key, value);
-  return value;
-}

package/src/model/batched-kv.ts

@@ -1,235 +0,0 @@
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-import { materializeCopy } from "../mlx/materialize";
-import { KVCache, isPlainKvCache, type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
-import { FullPrefillPadding } from "./full-prefill-padding";
-
-/** Full-attention storage with independently advancing row positions.
- * Speculative rollback changes coverage, without moving the retained KV.
- * Subsequent appends overwrite each row's rejected suffix in place when MLX
- * can donate the buffer. Queue membership is outside this layout. */
-export class BatchedKVCache implements BatchableCache, PaddedPrefillCache {
-  keys: MlxArray | null = null;
-  values: MlxArray | null = null;
-  rowOffsets: number[] = [];
-  leftPad: number[] = [];
-  #rope: MlxArray | undefined;
-  #beforeRound: number[] | undefined;
-  readonly #padding = new FullPrefillPadding();
-
-  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
-  preparePrefill(padding: PrefillPadding): void {
-    this.#padding.prepare(this, padding); this.#positionsChanged();
-  }
-  finalizePrefill(): void {
-    const next = this.#padding.finalize(this.state(), this);
-    if (!next) return;
-    this.keys?.dispose(); this.values?.dispose();
-    this.keys = next[0] ?? null; this.values = next[1] ?? null;
-    this.#positionsChanged();
-  }
-
-  captureDonorRows(): import("./gemma4-base").KvDonorRows {
-    const B = this.rowOffsets.length, width = this.offset;
-    const keys = this.keys!.slice([0,0,0,0], [B,this.keys!.shape[1]!,width,this.keys!.shape[3]!]);
-    let values: MlxArray;
-    try { values = this.values!.slice([0,0,0,0], [B,this.values!.shape[1]!,width,this.values!.shape[3]!]); }
-    catch (error) { keys.dispose(); throw error; }
-    return { keys, values, offsets: [...this.rowOffsets], starts: [...this.leftPad], ends: this.#ends() };
-  }
-
-  signature(): string { return "kv:batched-positions"; }
-  get batchSize(): number | null { return this.rowOffsets.length || null; }
-  get offset(): number { return Math.max(0, ...this.#ends()); }
-  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
-  get ropeOffsetArr(): MlxArray | undefined {
-    // A common unpadded position uses the existing scalar RoPE operation.
-    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
-    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
-  }
-  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
-  makeEmptyBatch(): BatchedKVCache { return new BatchedKVCache(); }
-  bytesPerToken(): number {
-    if (!this.keys || !this.values) return 0;
-    return (this.keys.nbytes + this.values.nbytes) / (this.rowOffsets.length * this.keys.shape[2]!);
-  }
-  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
-  state(): MlxArray[] { return this.keys && this.values ? [this.keys, this.values] : []; }
-  isTrimmable(): boolean { return true; }
-  trim(count: number): void {
-    this.rowOffsets = this.rowOffsets.map(offset => offset - count);
-    this.#positionsChanged();
-  }
-  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
-  specRoundCommit(): void { this.#beforeRound = undefined; }
-  specRoundRollback(keep: number | readonly number[]): void {
-    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
-    this.#beforeRound = undefined;
-    this.#positionsChanged();
-  }
-
-  makeMask(N: number, window: number | null): Mask {
-    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
-    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
-      return { mode: N === 1 ? "" : "causal", arr: null };
-    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
-    using relative = ops.arange(0, N, 1, Dtype.int32);
-    using queries = ops.reshape(relative, [1, 1, N, 1]);
-    using position = ops.add(starts, queries);
-    using indices = ops.arange(0, S, 1, Dtype.int32);
-    using keys = ops.reshape(indices, [1, 1, 1, S]);
-    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
-    using causal = ops.lessEqual(keys, position);
-    using valid = ops.greaterEqual(keys, pads);
-    const mask = ops.logicalAnd(causal, valid);
-    if (window === null) return { mode: "array", arr: mask };
-    using windowSize = ops.fromInt32([window], []);
-    using limit = ops.add(keys, windowSize);
-    using within = ops.less(position, limit);
-    try { return { mode: "array", arr: ops.logicalAnd(mask, within) }; }
-    finally { mask.dispose(); }
-  }
-
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
-    const B = this.rowOffsets.length, N = k.shape[2]!, ends = this.#ends();
-    const needed = Math.max(...ends) + N;
-    const grow = (array: MlxArray) => {
-      const [rows, heads, capacity, width] = array.shape as [number, number, number, number];
-      if (capacity >= needed) return array;
-      using padding = ops.zeros([rows, heads, Math.ceil(needed / 256) * 256 - capacity, width], array.dtype);
-      const grown = ops.concatAxis([array, padding], 2);
-      array.dispose();
-      return grown;
-    };
-    const allocate = (input: MlxArray) => ops.zeros(
-      [B, input.shape[1]!, Math.ceil(needed / 256) * 256, input.shape[3]!], input.dtype);
-    this.keys = this.keys ? grow(this.keys) : allocate(k);
-    this.values = this.values ? grow(this.values) : allocate(v);
-    const common = ends.every(end => end === ends[0]);
-    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
-      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
-    const append = (array: MlxArray, values: MlxArray) => common
-      ? ops.sliceUpdate(array, values, [0, 0, ends[0]!, 0], [B, array.shape[1]!, ends[0]! + N, array.shape[3]!])
-      : ops.putAlongAxis(array, positions!, values, 2);
-    const keys = append(this.keys, k);
-    let values: MlxArray;
-    try { values = append(this.values, v); }
-    catch (error) { keys.dispose(); throw error; }
-    this.keys.dispose(); this.values.dispose();
-    this.keys = keys; this.values = values;
-    this.rowOffsets = this.rowOffsets.map(offset => offset + N);
-    this.#positionsChanged();
-    return [keys.slice([0, 0, 0, 0], [B, keys.shape[1]!, needed, keys.shape[3]!]),
-      values.slice([0, 0, 0, 0], [B, values.shape[1]!, needed, values.shape[3]!])];
-  }
-
-  mergeRows(rows: readonly Cache[]): void {
-    const only = rows.length === 1 ? rows[0] : undefined;
-    if (only && (only instanceof BatchedKVCache || isPlainKvCache(only))) {
-      // Preserve allocated capacity when a request starts alone. Immutable
-      // aliases transfer ownership without copying its prefill storage.
-      const keys = only.keys ? ops.copyOf(only.keys) : null;
-      let values: MlxArray | null;
-      try { values = only.values ? ops.copyOf(only.values) : null; }
-      catch (error) { keys?.dispose(); throw error; }
-      const offsets = only instanceof BatchedKVCache ? [...only.rowOffsets] : [only.offset];
-      const pads = only instanceof BatchedKVCache ? [...only.leftPad] : [0];
-      this.dispose(); this.keys = keys; this.values = values;
-      this.rowOffsets = offsets; this.leftPad = pads;
-      return;
-    }
-    const chunks: Array<{ keys: MlxArray | null; values: MlxArray | null; offsets: number[]; pads: number[]; width: number }> = [];
-    const held: MlxArray[] = [];
-    try {
-      for (const cache of rows) {
-        let keys: MlxArray | null, values: MlxArray | null, offsets: number[], pads: number[];
-        if (cache instanceof BatchedKVCache && cache.offset === 0) {
-          keys = values = null; offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
-        } else if (cache instanceof BatchedKVCache) {
-          const B = cache.rowOffsets.length, width = cache.offset;
-          keys = cache.keys!.slice([0, 0, 0, 0], [B, cache.keys!.shape[1]!, width, cache.keys!.shape[3]!]);
-          held.push(keys);
-          values = cache.values!.slice([0, 0, 0, 0], [B, cache.values!.shape[1]!, width, cache.values!.shape[3]!]);
-          held.push(values);
-          offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
-        } else if (isPlainKvCache(cache)) {
-          if (cache.offset === 0) keys = values = null;
-          else { [keys, values] = cache.temporalView(); held.push(keys, values); }
-          offsets = [cache.offset]; pads = [0];
-        } else throw new Error(`full-attention batch cannot merge ${cache.signature()}`);
-        chunks.push({ keys, values, offsets, pads, width: keys?.shape[2] ?? 0 });
-      }
-      const width = Math.max(0, ...chunks.map(chunk => chunk.width));
-      const padded = (array: MlxArray, count: number) => {
-        if (!count) return array;
-        using padding = ops.zeros([array.shape[0]!, array.shape[1]!, count, array.shape[3]!], array.dtype);
-        const result = ops.concatAxis([padding, array], 2); held.push(result); return result;
-      };
-      const merge = (field: "keys" | "values") => {
-        if (!width) return null; // Geometry arrives with the first append.
-        const prototype = chunks.find(chunk => chunk[field])![field]!;
-        return ops.concatAxis(chunks.map(chunk => {
-          const array = chunk[field];
-          if (array) return padded(array, width - chunk.width);
-          const empty = ops.zeros([chunk.offsets.length, prototype.shape[1]!, width, prototype.shape[3]!], prototype.dtype);
-          held.push(empty); return empty;
-        }), 0);
-      };
-      const keys = merge("keys");
-      let values: MlxArray | null;
-      try { values = merge("values"); }
-      catch (error) { keys?.dispose(); throw error; }
-      this.dispose();
-      this.keys = keys; this.values = values;
-      this.rowOffsets = chunks.flatMap(chunk => chunk.offsets);
-      this.leftPad = chunks.flatMap(chunk => chunk.pads.map(pad => pad + width - chunk.width));
-    } finally { for (const array of held) array.dispose(); }
-  }
-
-  filterRows(rows: readonly number[]): void {
-    if (!rows.length) { this.dispose(); return; }
-    const pads = rows.map(row => this.leftPad[row]!);
-    const sharedPadding = Math.min(...pads);
-    using indices = this.keys ? ops.fromInt32([...rows], [rows.length]) : null;
-    const select = (array: MlxArray | null): MlxArray | null => {
-      if (!array) return null;
-      const selected = ops.takeAxis(array, indices!, 0);
-      if (!sharedPadding) return selected;
-      // mlx-lm BatchKVCache.filter removes padding common to all survivors.
-      // Keeping it changes attention reduction geometry after retirement.
-      try { return selected.slice([0, 0, sharedPadding, 0], [...selected.shape]); }
-      finally { selected.dispose(); }
-    };
-    const keys = select(this.keys);
-    let values: MlxArray | null;
-    try { values = select(this.values); }
-    catch (error) { keys?.dispose(); throw error; }
-    this.keys?.dispose(); this.values?.dispose();
-    this.keys = keys; this.values = values;
-    this.rowOffsets = rows.map(row => this.rowOffsets[row]!);
-    this.leftPad = pads.map(pad => pad - sharedPadding);
-    this.#padding.filter(rows);
-    this.#positionsChanged();
-  }
-
-  extractRow(row: number): KVCache {
-    const start = this.leftPad[row]!, count = this.rowOffsets[row]!;
-    if (count === 0) return new KVCache();
-    const copy = (array: MlxArray) => {
-      using view = array.slice([row, 0, start, 0], [row + 1, array.shape[1]!, start + count, array.shape[3]!]);
-      return materializeCopy(view);
-    };
-    const keys = copy(this.keys!);
-    let values: MlxArray;
-    try { values = copy(this.values!); }
-    catch (error) { keys.dispose(); throw error; }
-    const cache = new KVCache(); cache.restoreState(keys, values, count); return cache;
-  }
-
-  dispose(): void {
-    this.keys?.dispose(); this.values?.dispose(); this.#positionsChanged();
-    this.keys = this.values = null;
-    this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
-  }
-}