mlx-bun 0.0.9 → 0.0.11

package/src/mlx/turboquant-tables.ts

@@ -0,0 +1,237 @@
+// GENERATED by scripts/regen-turboquant-goldens.ts — DO NOT EDIT.
+// Regenerate:  bun scripts/regen-turboquant-goldens.ts
+//
+// 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/model/batched-mask.ts

@@ -17,7 +17,7 @@
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
-import { createCausalMask, type Cache, type Mask } from "./gemma4-base";
+import { createCausalMask, KVCache, type Cache, type Mask } from "./gemma4-base";
 
 /** Padding-aware key-validity mask for batched DECODE with left-padded rows.
  *
@@ -190,6 +190,67 @@ export function mergeKVRows(
   return { keys, values, leftPad, width };
 }
 
+/** mlx-lm `BatchKVCache.extend`: append ONE new right-justified row to an
+ *  existing batched [B,H,S,D] buffer WITHOUT re-extracting the running rows
+ *  (the O(B·S) slice storm mergeKVRows pays on a join). Pad whichever side
+ *  is narrower on the LEFT to the common width, then a single B-axis concat.
+ *  Existing rows' left-pads grow by the batch-side pad (they never shrink —
+ *  mlx-lm extend semantics; the decode mask covers pads either way).
+ *  Caller owns the result; inputs are not disposed. */
+export function extendKVRows(
+  keys: MlxArray, values: MlxArray, leftPad: number[],
+  row: { keys: MlxArray; values: MlxArray },
+): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
+  const S = keys.shape[2]!;
+  const Lr = row.keys.shape[2]!;
+  const width = Math.max(S, Lr);
+  const padTo = (a: MlxArray, pad: number): MlxArray => {
+    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
+    const [B, H, , D] = a.shape as [number, number, number, number];
+    const z = ops.zeros([B, H, pad, D], a.dtype);
+    const out = ops.concatAxis([z, a], 2);
+    z.dispose();
+    return out;
+  };
+  const bk = padTo(keys, width - S);
+  const bv = padTo(values, width - S);
+  const rk = padTo(row.keys, width - Lr);
+  const rv = padTo(row.values, width - Lr);
+  const outK = ops.concatAxis([bk, rk], 0);
+  const outV = ops.concatAxis([bv, rv], 0);
+  for (const a of [bk, bv, rk, rv]) a.dispose();
+  return {
+    keys: outK,
+    values: outV,
+    leftPad: [...leftPad.map((p) => p + (width - S)), width - Lr],
+    width,
+  };
+}
+
+/** mlx-lm `BatchKVCache.extract` (models/cache.py:1080): pull row `i` out of
+ *  a batched full-attention cache into a fresh SERIAL KVCache — the row slice
+ *  past its left padding, materialized as an OWNED contiguous copy (extraction
+ *  outlives the batch; the source buffers are about to be filter()ed or
+ *  disposed). `leftPad` is the scheduler's per-row padding for this row (full
+ *  layers share one offset; the scheduler owns the pad list). Bit-exact vs a
+ *  solo run by construction: merge/extend/filter are byte-preserving per row
+ *  (tests/batched-decode-parity, tests/batched-extract) and this is a pure
+ *  slice+copy of those bytes. Off the hot decode path. */
+export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCache {
+  if (!cache.keys || !cache.values) throw new Error("extractKVRow: empty cache");
+  const S = cache.offset; // valid columns (buffer may be STEP-padded past it)
+  const cut = (a: MlxArray): MlxArray => {
+    const [, H, , D] = a.shape as [number, number, number, number];
+    const view = a.slice([i, 0, leftPad, 0], [i + 1, H, S, D]);
+    const own = ops.contiguous(view);
+    view.dispose();
+    return own;
+  };
+  const out = new KVCache();
+  out.restoreState(cut(cache.keys), cut(cache.values), S - leftPad);
+  return out;
+}
+
 /** mlx-lm `filter`: keep only `keep` (sorted row indices) along the batch
  *  axis — eviction of finished sequences from a batched [B,H,S,D] buffer.
  *  Caller owns the result; inputs are not disposed. */

package/src/model/batched-quant.ts

@@ -0,0 +1,206 @@
+// Batched QUANTIZED KV for continuous batching (Phase 3.1 of
+// docs/design/unified-engine-frontier-plan.md) — the quantized twins of
+// batched-mask.ts's dynamic-B ops, over (packed, scales, biases) triples.
+//
+// Why this is safe surgery: mlx quantization packs along HEAD_DIM (the last
+// axis; packed last-dim = D/(32/bits), scales/biases last-dim = D/groupSize),
+// so the TOKEN axis (axis 2) stays token-granular in all three tensors —
+// left-padding, concat, and row-filter at token granularity never touch a
+// quantization group. Pad columns hold zeros; they are never attended (the
+// padding mask covers them) and never dequantized into anything a real row
+// reads. QuantizedKVCache.updateAndFetchQuantized already handles arbitrary
+// starting widths (it trims to `offset` and regrows in STEP chunks), so a
+// merged exact-width buffer needs no STEP alignment here.
+//
+// Composition contract (the L2 oracle binds per ROW): each joiner solo-
+// prefills with the SAME per-layer conversion the serial path runs
+// (maybeQuantizeKv at chunk boundaries — see BatchScheduler.#quantizeSolo),
+// so a row's quantized bytes are bit-exact vs serial `--kv-quant config` by
+// construction; this module only re-arranges those bytes across the batch
+// axis. Gates: tests/batched-kv-quant-parity.test.ts.
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { QuantizedKVCache, type Mask } from "./gemma4-base";
+import { buildBatchedDecodeMask } from "./batched-mask";
+
+export interface QuantRow {
+  keys: ops.QuantizedTensor;
+  values: ops.QuantizedTensor;
+}
+
+const disposeTriple = (t: ops.QuantizedTensor): void => {
+  t.packed.dispose();
+  t.scales.dispose();
+  t.biases.dispose();
+};
+
+/** Apply an MlxArray op to each member of a triple. Each tensor keeps its
+ *  OWN last-dim (packed vs scales/biases differ) — helpers must never share
+ *  one `D` across the triple (recon hazard, 2026-07-05). */
+const tripleMap = (
+  t: ops.QuantizedTensor, f: (a: MlxArray) => MlxArray,
+): ops.QuantizedTensor => ({ packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) });
+
+/** Left-pad one tensor of a triple by `pad` tokens (axis 2) with zeros of
+ *  its own dtype/last-dim. pad=0 returns a fresh full view (mirrors
+ *  batched-mask.ts leftPadTo so ownership is uniform). */
+const leftPadTo = (a: MlxArray, pad: number): MlxArray => {
+  if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]);
+  const [B, H, , D] = a.shape as [number, number, number, number];
+  const z = ops.zeros([B, H, pad, D], a.dtype);
+  const out = ops.concatAxis([z, a], 2);
+  z.dispose();
+  return out;
+};
+
+/** Quantized twin of mergeKVRows: stack N single-row quantized KV slices
+ *  into one left-padded [B,H,Smax,*] triple pair. Caller owns the result;
+ *  input rows are not disposed. */
+export function mergeQuantRows(rows: QuantRow[]): {
+  keys: ops.QuantizedTensor; values: ops.QuantizedTensor; leftPad: number[]; width: number;
+} {
+  const lens = rows.map((r) => r.keys.packed.shape[2]!);
+  const width = Math.max(...lens);
+  const leftPad = lens.map((l) => width - l);
+  const padded = rows.map((r, i) => ({
+    keys: tripleMap(r.keys, (a) => leftPadTo(a, leftPad[i]!)),
+    values: tripleMap(r.values, (a) => leftPadTo(a, leftPad[i]!)),
+  }));
+  const cat = (pick: (r: QuantRow) => ops.QuantizedTensor): ops.QuantizedTensor => ({
+    packed: ops.concatAxis(padded.map((r) => pick(r).packed), 0),
+    scales: ops.concatAxis(padded.map((r) => pick(r).scales), 0),
+    biases: ops.concatAxis(padded.map((r) => pick(r).biases), 0),
+  });
+  const keys = cat((r) => r.keys);
+  const values = cat((r) => r.values);
+  for (const r of padded) { disposeTriple(r.keys); disposeTriple(r.values); }
+  return { keys, values, leftPad, width };
+}
+
+/** Quantized twin of extendKVRows: append ONE right-justified quantized row
+ *  to an existing batched triple pair in one pad + one B-axis concat.
+ *  Existing pads grow, never shrink (mlx-lm extend semantics). Caller owns
+ *  the result; inputs are not disposed. */
+export function extendQuantRows(
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, leftPad: number[], row: QuantRow,
+): { keys: ops.QuantizedTensor; values: ops.QuantizedTensor; leftPad: number[]; width: number } {
+  const S = keys.packed.shape[2]!;
+  const Lr = row.keys.packed.shape[2]!;
+  const width = Math.max(S, Lr);
+  const bk = tripleMap(keys, (a) => leftPadTo(a, width - S));
+  const bv = tripleMap(values, (a) => leftPadTo(a, width - S));
+  const rk = tripleMap(row.keys, (a) => leftPadTo(a, width - Lr));
+  const rv = tripleMap(row.values, (a) => leftPadTo(a, width - Lr));
+  const cat = (b: ops.QuantizedTensor, r: ops.QuantizedTensor): ops.QuantizedTensor => ({
+    packed: ops.concatAxis([b.packed, r.packed], 0),
+    scales: ops.concatAxis([b.scales, r.scales], 0),
+    biases: ops.concatAxis([b.biases, r.biases], 0),
+  });
+  const outK = cat(bk, rk);
+  const outV = cat(bv, rv);
+  for (const t of [bk, bv, rk, rv]) disposeTriple(t);
+  return {
+    keys: outK, values: outV,
+    leftPad: [...leftPad.map((p) => p + (width - S)), width - Lr],
+    width,
+  };
+}
+
+/** Quantized twin of extractKVRow (mlx-lm BatchKVCache.extract,
+ *  models/cache.py:1080, over triples): pull row `i` past its left padding
+ *  into a fresh SERIAL QuantizedKVCache with OWNED contiguous copies.
+ *  Token-axis slicing is byte-safe because quantization packs along
+ *  HEAD_DIM (file header) — a row's quantized bytes come out exactly as
+ *  the serial oracle wrote them. */
+export function extractQuantRow(
+  cache: QuantizedKVCache, leftPad: number, i: number,
+): QuantizedKVCache {
+  if (!cache.keys || !cache.values) throw new Error("extractQuantRow: empty cache");
+  const S = cache.offset;
+  const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+    tripleMap(t, (a) => {
+      const [, H, , D] = a.shape as [number, number, number, number];
+      const view = a.slice([i, 0, leftPad, 0], [i + 1, H, S, D]);
+      const own = ops.contiguous(view);
+      view.dispose();
+      return own;
+    });
+  const out = new QuantizedKVCache(cache.groupSize, cache.bits);
+  out.restoreState(cut(cache.keys), cut(cache.values), S - leftPad);
+  return out;
+}
+
+/** Quantized twin of filterKVRows: keep `keep` rows along the batch axis.
+ *  Caller owns the result; inputs are not disposed. */
+export function filterQuantRows(
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, keep: number[],
+): { keys: ops.QuantizedTensor; values: ops.QuantizedTensor } {
+  const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
+  const take = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+    tripleMap(t, (a) => ops.takeAxis(a, idx, 0));
+  const k = take(keys);
+  const v = take(values);
+  idx.dispose();
+  return { keys: k, values: v };
+}
+
+/** Per-step cache wrapper for batched QUANTIZED decode with left-padded
+ *  rows — the quantized twin of BatchedDecodeMaskCache. Subclasses
+ *  QuantizedKVCache so attention's `instanceof` dispatch routes to
+ *  quantizedSdpa exactly as for a serial quantized cache; delegates KV
+ *  storage to the persistent batched inner and overrides the two things a
+ *  padded batch needs: the padding-aware decode mask and per-row RoPE
+ *  positions. Ephemeral — the scheduler builds a fresh wrapper each step
+ *  and calls releaseRopeArr() after (never dispose(), which would free the
+ *  inner). Rope is captured ONCE per step before updateAndFetch (attention
+ *  reads it pre-write; see minicpm5.ts LlamaAttention.forward), so an
+ *  eagerly-built [B] position array is exact. */
+export class BatchedQuantDecodeMaskCache extends QuantizedKVCache {
+  override readonly ropeOffsetArr: MlxArray;
+
+  constructor(
+    private readonly inner: QuantizedKVCache,
+    private readonly B: number,
+    private readonly leftPad: number[],
+  ) {
+    super(inner.groupSize, inner.bits);
+    this.offset = inner.offset;
+    const data = new Int32Array(B);
+    for (let b = 0; b < B; b++) data[b] = inner.offset - leftPad[b]!;
+    this.ropeOffsetArr = MlxArray.fromInt32(data, [B]);
+  }
+
+  override updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    const r = this.inner.updateAndFetchQuantized(k, v);
+    this.offset = this.inner.offset;
+    return r;
+  }
+
+  override makeMask(N: number, _windowSize: number | null): Mask {
+    // Same convention as BatchedDecodeMaskCache: built pre-write, S spans
+    // offset + N; quantizedSdpaUnfused expands it for the 5-D GQA scores.
+    const S = this.inner.offset + N;
+    return { mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, null) };
+  }
+
+  /** Free the per-step RoPE array WITHOUT touching the persistent inner —
+   *  name and signature must match the scheduler's structural
+   *  `(c as { releaseRopeArr?: () => void })` cast. */
+  releaseRopeArr(): void {
+    this.ropeOffsetArr.dispose();
+  }
+
+  override state(): MlxArray[] {
+    return this.inner.state();
+  }
+  override isTrimmable(): boolean {
+    return this.inner.isTrimmable();
+  }
+  override trim(n: number): void {
+    this.inner.trim(n);
+  }
+  override dispose(): void {
+    this.inner.dispose();
+  }
+}