@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/turboquant/signs.ts

@@ -0,0 +1,130 @@
+// GENERATED by packages/inference/scripts/gen-turboquant-signs.ts; do not edit.
+//
+// TurboQuant's rotation signs per head dimension, from the construction the
+// vendored vllm-metal reference defines (`02d723a:lab/repro/vllm-metal-turboquant/
+// turboquant_reference.py`, `fwht()`): randint(0, 2, shape=(d,), key=key(42)),
+// mapped to 1 - 2·bit. The rotation is these signs then MLX's hadamard_transform.
+
+export const TURBOQUANT_HEAD_DIMS = [64, 96, 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,
+  ],
+  96: [
+  1, -1, -1, -1, -1, -1, 1, 1, 1, 1,
+  1, 1, 1, -1, 1, -1, 1, 1, 1, -1,
+  1, 1, 1, -1, -1, 1, 1, -1, -1, 1,
+  -1, -1, -1, -1, 1, 1, -1, 1, -1, 1,
+  1, -1, 1, 1, 1, 1, -1, 1, -1, 1,
+  1, -1, 1, 1, -1, 1, -1, -1, 1, 1,
+  -1, 1, -1, 1, -1, -1, -1, -1, -1, 1,
+  -1, -1, -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,
+  ],
+};

package/src/kernels/turboquant/tables.ts

@@ -0,0 +1,124 @@
+// GENERATED by `02d723a:scripts/regen.ts` turboquant — DO NOT EDIT.
+// Regenerate (from a 02d723a checkout):  bun scripts/regen.ts turboquant
+//
+// Source: `02d723a: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.
+//
+// - 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.
+//
+// The rotation's sign tables are signs.ts (scripts/gen-turboquant-signs.ts).
+
+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/layers/affine-verify-linear.ts

@@ -0,0 +1,75 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { affineRow, affineRowEligible } from "../kernels/quantization/affine-row";
+import { affineSplitMma, affineSplitMmaEligible, affineSplitMmaPlan, type AffineSplitMmaPlan } from "../kernels/quantization/affine-split-mma";
+import { affine3Rows, affine3RowsEligible } from "../kernels/quantization/affine3-rows";
+import { affine3Mma, affine3MmaEligible } from "../kernels/quantization/affine3-mma";
+import { affine4Rows, affine4RowsEligible } from "../kernels/quantization/affine4-rows";
+import { affine4Mma, affine4MmaEligible } from "../kernels/quantization/affine4-mma";
+import type { QuantizedLinear } from "./quantized-linear";
+
+// Decode and verify projections borrow an affine carrier. Each class runs its
+// fixed kernel and applies the request's external adapter delta afterward.
+
+
+function check(name: string, lin: QuantizedLinear, eligible: boolean): void {
+  if (!lin.biases || lin.bias || !eligible) throw new Error(`${name}: unsupported projection (width, group 64, affine, no additive bias)`);
+}
+
+/** 2..8 activation rows sharing each decoded weight. */
+export class Affine4RowsLinear {
+  constructor(readonly linear: QuantizedLinear) {
+    check("Affine4RowsLinear", linear, affine4RowsEligible(2, linear.w, linear.scales, linear.spec));
+  }
+  forward(x: MlxArray): MlxArray { return this.linear.delta.forward(x, affine4Rows(x, this.linear.w, this.linear.scales, this.linear.biases!)); }
+}
+
+/** One activation row (3- or 4-bit, the stored width): QuantizedLinear's bits (MLX's one-row quantized matmul),
+ *  each threadgroup computing its activation chunks once (M1 Max). */
+export class AffineRowLinear {
+  readonly #bits: 3 | 4;
+  constructor(readonly linear: QuantizedLinear) {
+    check("AffineRowLinear", linear, affineRowEligible(1, linear.w, linear.scales, linear.spec));
+    this.#bits = linear.spec.bits as 3 | 4;
+  }
+  forward(x: MlxArray): MlxArray { return this.linear.delta.forward(x, affineRow(x, this.linear.w, this.linear.scales, this.linear.biases!, this.#bits)); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit. */
+export class Affine4MmaLinear {
+  constructor(readonly linear: QuantizedLinear) {
+    check("Affine4MmaLinear", linear, affine4MmaEligible(1, linear.w, linear.scales, linear.spec));
+  }
+  forward(x: MlxArray): MlxArray { return this.linear.delta.forward(x, affine4Mma(x, this.linear.w, this.linear.scales, this.linear.biases!)); }
+}
+
+/** 3-bit: 2..8 activation rows sharing each decoded weight. */
+export class Affine3RowsLinear {
+  constructor(readonly linear: QuantizedLinear) {
+    check("Affine3RowsLinear", linear, affine3RowsEligible(2, linear.w, linear.scales, linear.spec));
+  }
+  forward(x: MlxArray): MlxArray { return this.linear.delta.forward(x, affine3Rows(x, this.linear.w, this.linear.scales, this.linear.biases!)); }
+}
+
+/** 3-bit: 1..8 activation rows on the simdgroup matrix unit. */
+export class Affine3MmaLinear {
+  constructor(readonly linear: QuantizedLinear) {
+    check("Affine3MmaLinear", linear, affine3MmaEligible(1, linear.w, linear.scales, linear.spec));
+  }
+  forward(x: MlxArray): MlxArray { return this.linear.delta.forward(x, affine3Mma(x, this.linear.w, this.linear.scales, this.linear.biases!)); }
+}
+
+/** Any number of activation rows, in blocks of 8 (3- or 4-bit, the stored width), on the simdgroup matrix unit,
+ *  output tiles sharing each activation fragment and K split across the threadgroup (M1 Max); the launch follows
+ *  the output width. */
+export class AffineSplitMmaLinear {
+  readonly #bits: 3 | 4;
+  readonly #plan: AffineSplitMmaPlan;
+  constructor(readonly linear: QuantizedLinear) {
+    check("AffineSplitMmaLinear", linear, affineSplitMmaEligible(1, linear.w, linear.scales, linear.spec));
+    this.#bits = linear.spec.bits as 3 | 4;
+    this.#plan = affineSplitMmaPlan(linear.outFeatures);
+  }
+  forward(x: MlxArray): MlxArray {
+    return this.linear.delta.forward(x, affineSplitMma(x, this.linear.w, this.linear.scales, this.linear.biases!, this.#bits, this.#plan));
+  }
+}

package/src/layers/ane-prefill-split.ts

@@ -0,0 +1,283 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { AneProgram, AneStreamedMatmul, mlpSliceProgram, streamedMatmulProgram, type AneProgramSpec } from "../kernels/ane/linear";
+import { trellisFillHalf, trellisFillHalfK3Interleaved, writeHalf } from "../kernels/ane/fill";
+import type { TrellisGeometry } from "../kernels/trellis/geometry";
+import { expandTrellis } from "../kernels/trellis/index";
+import { QuantizedLinear, QuantizedLinearPlain } from "./quantized-linear";
+import { TrellisLinear, trellisRows, trellisStoredRows } from "./trellis-linear";
+
+// Prompt-piece channel splits between the Apple Neural Engine and the GPU. The
+// ANE computes the first n output channels from streamed fp16 weights the GPU
+// writes into its program's buffers (zero-copy IOSurface wraps, made once per
+// set of buffers), the GPU computes the rest at the same time, and the results
+// combine. Each layer holds one program, built for the row count the layer is
+// constructed with, and serves inputs of up to that many rows: it writes the
+// rows it is given at the front of the program's input and reads as many back.
+// The rows behind them hold whatever the last evaluation left; no row it returns
+// reads them (each output row of the program depends only on its own input row).
+// dispose releases the program. Layers that never run at once may share, named
+// at construction (`share`): a layer whose program is identical to the other's
+// holds that same loaded program, otherwise its program loads on the other's
+// buffers (each at least as large); either lives until every layer holding it is
+// disposed. Loaded programs are a machine-wide resource (AneProgram), so a
+// graph's layers of one shape share one. Not bit-identical to the GPU path (fp16
+// weights and activations on the ANE); gate by KL.
+//
+// A layer's program depends only on its widths, the fraction, its row count and
+// its matmuls' reduction slices, never on the weights it streams, so the program
+// each split loads is named here as a function of those (`aneMlpSplitProgram`,
+// `aneAffineSplitProgram`): the install step compiles exactly what `build`
+// loads.
+
+/** Channels the ANE takes of R rows at `fraction` (64-aligned, at least 64 left on each side). */
+export function aneChannels(R: number, fraction: number): number {
+  return Math.min(R - 64, Math.max(64, Math.round((fraction * R) / 64) * 64));
+}
+
+/** The program either whole-MLP split of hidden width D and intermediate width R
+ *  loads at `fraction` for pieces of up to `rows` rows, its matmuls' reductions
+ *  in `slices` slices (mlpSliceProgram's default when absent). */
+export function aneMlpSplitProgram(D: number, R: number, fraction: number, rows: number, slices?: number): AneProgramSpec {
+  return mlpSliceProgram(D, aneChannels(R, fraction), rows, slices);
+}
+
+/** The program an affine projection split of K inputs and N outputs loads at
+ *  `fraction` for pieces of up to `rows` rows, its reduction in `slices` slices
+ *  (streamedMatmulProgram's default when absent). */
+export function aneAffineSplitProgram(K: number, N: number, fraction: number, rows: number, slices?: number): AneProgramSpec {
+  return streamedMatmulProgram(K, aneChannels(N, fraction), rows, slices);
+}
+
+type Fill = (codes: MlxArray, scales: MlxArray, g: TrellisGeometry, rows: number, dst: MlxArray, offset: number) => MlxArray;
+type Swiglu = (g: MlxArray, u: MlxArray) => MlxArray;
+type Rest = { gate: TrellisLinear; up: TrellisLinear; down: TrellisLinear };
+
+/** Dense axis-1 expansion queued with the ANE's GPU share, without an
+ * intermediate evaluation that would serialize gate, up and down. */
+export class AneTrellisExpandReduce {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    const lead = x.shape.slice(0, -1), M = lead.reduce((a, b) => a * b, 1);
+    using input = ops.reshape(x, [M, this.linear.inFeatures]);
+    using stored = expandTrellis(this.linear.codes, this.linear.scales, this.linear.geometry, x.dtype, this.linear.variant);
+    using transposed = ops.transposeAxes(stored, [1, 0]);
+    using product = ops.matmul(input, transposed);
+    return this.linear.delta.forward(x, ops.reshape(product, [...lead, this.linear.outFeatures]));
+  }
+}
+/** Dense axis-0 expansion with the same submission boundary. */
+export class AneTrellisExpandScatter {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    const lead = x.shape.slice(0, -1), M = lead.reduce((a, b) => a * b, 1);
+    using input = ops.reshape(x, [M, this.linear.inFeatures]);
+    using stored = expandTrellis(this.linear.codes, this.linear.scales, this.linear.geometry, x.dtype, this.linear.variant);
+    using product = ops.matmul(input, stored);
+    return this.linear.delta.forward(x, ops.reshape(product, [...lead, this.linear.outFeatures]));
+  }
+}
+type GpuRest = { gate: { forward(x: MlxArray): MlxArray }; up: { forward(x: MlxArray): MlxArray }; down: { forward(x: MlxArray): MlxArray } };
+/** The GPU projections attached to the split's stored row slices. */
+export type AneGpuMlpRows = (rest: Rest) => GpuRest;
+const expandedGpuRows: AneGpuMlpRows = rest => ({ gate: new AneTrellisExpandReduce(rest.gate),
+  up: new AneTrellisExpandReduce(rest.up), down: new AneTrellisExpandScatter(rest.down) });
+
+// The public builder also supports standalone small inputs. Its compatibility
+// projections preserve their original row-dependent kernels; written graph
+// bindings use the fixed expanded builders below.
+const legacyGpuRows: AneGpuMlpRows = rest => ({
+  gate: { forward: x => rest.gate.forward(x, true, false) },
+  up: { forward: x => rest.up.forward(x, true, false) },
+  down: { forward: x => rest.down.forward(x, true, false) },
+});
+
+/** Whole Trellis MLP: the ANE computes the first n intermediate channels end
+ *  to end (gate and up rows, SwiGLU, and the down projection's partial sum over
+ *  those channels); the GPU computes channels [n, R) through its attached projections; the
+ *  two [M, D] partials add. Measured on M4 Pro: one fused slice program at
+ *  10.2 TFLOPS (D 5120, n 9600, S 2048). `hidden` holds M <= `rows` rows.
+ *  The ANE starts as soon as the GPU has filled its buffers; the GPU's share is
+ *  then queued whole (its weight expansions unbounded, one submission) while
+ *  the host waits for the ANE. The sum is returned unevaluated: it reads the
+ *  ANE output, which the program evaluates before its buffers' next run
+ *  (`AneProgram.deferOutput`), unless the caller has by then. */
+function mlpSplit(program: AneProgram, rows: number, n: number, gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear,
+  rest: GpuRest, fillDown: Fill, hidden: MlxArray, swiglu: Swiglu): MlxArray {
+  const D = gate.geometry.inFeatures;
+  const lead = hidden.shape.slice(0, -1), M = lead.reduce((a, b) => a * b, 1);
+  if (M > rows) throw new Error(`ANE MLP split: ${M} rows exceed its program's ${rows}`);
+  {
+    // Inputs in symbol order: wd, wg, wu, x (borrowed wraps).
+    const wd = program.inputArray(0), wg = program.inputArray(1), wu = program.inputArray(2), x = program.inputArray(3);
+    using h2 = ops.reshape(hidden, [M * D]);
+    using d0 = fillDown(down.codes, down.scales, down.geometry, n, wd, 0);
+    using d1 = trellisFillHalf(gate.codes, gate.scales, gate.geometry, n, wg, 0);
+    using d2 = trellisFillHalf(up.codes, up.scales, up.geometry, n, wu, 0);
+    using d3 = writeHalf(x, h2, 0);
+    ops.evalAll([d0, d1, d2, d3]);
+  }
+  program.evalAsync();
+  let gpuOut: MlxArray;
+  try {
+    using g = rest.gate.forward(hidden);
+    using u = rest.up.forward(hidden);
+    using mid = swiglu(g, u);
+    gpuOut = rest.down.forward(mid);
+    ops.asyncEvalAll([gpuOut]);
+  } finally { program.wait(); }
+  try {
+    using used = program.output([M, D]);                    // zero-copy partial, the piece's rows
+    using b = used.astype(hidden.dtype);
+    using g2 = ops.reshape(gpuOut, [M, D]);
+    const sum = ops.add(b, g2);
+    program.deferOutput(sum);                               // evaluated (free once the result is) and freed by the program
+    return ops.reshape(sum, [...lead, D]);
+  } finally { gpuOut.dispose(); }
+}
+
+function checkMlp(name: string, gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear, interleavedDown: boolean): void {
+  const g = gate.geometry, d = down.geometry, rowMajor = (x: TrellisGeometry) => !x.blockInterleave && x.L === 12 && x.T === 256;
+  if (g.axis !== 1 || up.geometry.axis !== 1 || !rowMajor(g) || !rowMajor(up.geometry) || d.axis !== 0 ||
+      d.rows !== g.rows || d.cols !== g.inFeatures || interleavedDown !== (d.blockInterleave === 2) ||
+      (interleavedDown ? d.k !== 3 : !rowMajor(d)))
+    throw new Error(`${name}: unsupported MLP geometry`);
+}
+function disposeLinear(lin: TrellisLinear): void { lin.codes.dispose(); lin.scales.dispose(); }
+
+/** Either whole-MLP split; both build the same slice program (only the GPU
+ *  fill differs), so either can share the other's program. */
+export type AneMlpSplit = AneTrellisMlpSplit | AneTrellisMlpSplitK3i;
+
+/** Whole-MLP split for a down projection with row-major codes. */
+export class AneTrellisMlpSplit implements Disposable {
+  private constructor(readonly program: AneProgram, readonly rows: number, readonly channels: number, readonly gate: TrellisLinear,
+    readonly up: TrellisLinear, readonly down: TrellisLinear, readonly rest: Rest, readonly gpu: GpuRest) {}
+
+  /** Loads `aneMlpSplitProgram(D, R, fraction, rows, slices)` for pieces of up
+   *  to `rows` rows now, from the compile cache the install step filled (a
+   *  program not compiled on this Mac throws `AneBridgeError`). `share`: a split
+   *  that never runs at the same time as this one; this split holds its program
+   *  when the two are identical, otherwise loads one on its buffers (each at
+   *  least as large). */
+  static build(gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear, fraction: number, rows: number,
+    share?: AneMlpSplit, slices?: number, gpuRows: AneGpuMlpRows = legacyGpuRows): AneTrellisMlpSplit {
+    checkMlp("AneTrellisMlpSplit", gate, up, down, false);
+    const R = gate.geometry.rows, n = aneChannels(R, fraction);
+    const program = AneProgram.create(aneMlpSplitProgram(gate.geometry.inFeatures, R, fraction, rows, slices), share?.program);
+    const rest = { gate: trellisRows(gate, n, R), up: trellisRows(up, n, R), down: trellisStoredRows(down, n, R) };
+    return new AneTrellisMlpSplit(program, rows, n, gate, up, down, rest, gpuRows(rest));
+  }
+  /** hidden [.., M, D], M <= rows → MLP(hidden) [.., M, D], unevaluated (mlpSplit). */
+  forward(hidden: MlxArray, swiglu: Swiglu): MlxArray {
+    return mlpSplit(this.program, this.rows, this.channels, this.gate, this.up, this.down, this.gpu, trellisFillHalf, hidden, swiglu);
+  }
+  [Symbol.dispose](): void { this.dispose(); }
+  /** Releases this layer's program and frees its GPU rows; a shared program or buffers outlive it while another layer holds them. */
+  dispose(): void { this.program.dispose(); for (const lin of Object.values(this.rest)) disposeLinear(lin); }
+}
+
+/** Whole-MLP split for a 3-bit block-interleaved down projection. */
+export class AneTrellisMlpSplitK3i implements Disposable {
+  private constructor(readonly program: AneProgram, readonly rows: number, readonly channels: number, readonly gate: TrellisLinear,
+    readonly up: TrellisLinear, readonly down: TrellisLinear, readonly rest: Rest, readonly gpu: GpuRest) {}
+
+  /** As AneTrellisMlpSplit.build. */
+  static build(gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear, fraction: number, rows: number,
+    share?: AneMlpSplit, slices?: number, gpuRows: AneGpuMlpRows = legacyGpuRows): AneTrellisMlpSplitK3i {
+    checkMlp("AneTrellisMlpSplitK3i", gate, up, down, true);
+    const R = gate.geometry.rows, n = aneChannels(R, fraction);
+    const program = AneProgram.create(aneMlpSplitProgram(gate.geometry.inFeatures, R, fraction, rows, slices), share?.program);
+    const rest = { gate: trellisRows(gate, n, R), up: trellisRows(up, n, R), down: trellisStoredRows(down, n, R) };
+    return new AneTrellisMlpSplitK3i(program, rows, n, gate, up, down, rest, gpuRows(rest));
+  }
+  /** hidden [.., M, D], M <= rows → MLP(hidden) [.., M, D], unevaluated (mlpSplit). */
+  forward(hidden: MlxArray, swiglu: Swiglu): MlxArray {
+    return mlpSplit(this.program, this.rows, this.channels, this.gate, this.up, this.down, this.gpu, trellisFillHalfK3Interleaved,
+      hidden, swiglu);
+  }
+  [Symbol.dispose](): void { this.dispose(); }
+  /** As AneTrellisMlpSplit.dispose. */
+  dispose(): void { this.program.dispose(); for (const lin of Object.values(this.rest)) disposeLinear(lin); }
+}
+
+/** Fixed dense GPU share for written row-major ANE graph bindings. */
+export class AneTrellisExpandedMlpSplit {
+  static build(gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear, fraction: number, rows: number,
+    share?: AneMlpSplit, slices?: number): AneTrellisMlpSplit {
+    return AneTrellisMlpSplit.build(gate, up, down, fraction, rows, share, slices, expandedGpuRows);
+  }
+}
+/** Fixed dense GPU share for written block-interleaved ANE graph bindings. */
+export class AneTrellisExpandedMlpSplitK3i {
+  static build(gate: TrellisLinear, up: TrellisLinear, down: TrellisLinear, fraction: number, rows: number,
+    share?: AneMlpSplit, slices?: number): AneTrellisMlpSplitK3i {
+    return AneTrellisMlpSplitK3i.build(gate, up, down, fraction, rows, share, slices, expandedGpuRows);
+  }
+}
+
+/** The layer that computes an affine split's GPU rows from the projection's
+ *  remaining rows (default: QuantizedLinearPlain). */
+export type AneGpuRows = (rest: QuantizedLinear) => { forward(x: MlxArray): MlxArray };
+
+/** Affine projection split: the first n output rows of an MLX affine
+ *  projection run on the ANE from fp16 rows the GPU dequantizes straight into
+ *  the streamed program's weight buffer; the GPU computes the remaining rows;
+ *  outputs concatenate in row order. */
+export class AneAffineSplit implements Disposable {
+  private constructor(readonly program: AneStreamedMatmul, readonly head: { w: MlxArray; scales: MlxArray; biases: MlxArray },
+    readonly rest: QuantizedLinear, readonly gpu: { forward(x: MlxArray): MlxArray }, readonly spec: ops.QuantSpec,
+    readonly linear: QuantizedLinear) {}
+
+  /** Loads `aneAffineSplitProgram(K, N, fraction, rows, slices)` for pieces of
+   *  up to `rows` rows now, as AneTrellisMlpSplit.build loads its program.
+   *  `share`: as AneTrellisMlpSplit.build. `gpuRows`: the layer over the GPU's
+   *  rows. */
+  static build(lin: QuantizedLinear, fraction: number, rows: number, share?: AneAffineSplit, slices?: number,
+    gpuRows?: AneGpuRows): AneAffineSplit {
+    const N = lin.outFeatures, K = lin.inFeatures;
+    if (!lin.biases || lin.bias || lin.w.ndim !== 2 || lin.spec.mode !== "affine") throw new Error("AneAffineSplit: unsupported projection");
+    const n = aneChannels(N, fraction);
+    const program = AneStreamedMatmul.create(K, n, rows, share?.program, slices);
+    const rowsOf = (a: MlxArray, start: number, stop: number) => a.slice([start, 0], [stop, a.shape[1]!]);
+    const head = { w: rowsOf(lin.w, 0, n), scales: rowsOf(lin.scales, 0, n), biases: rowsOf(lin.biases, 0, n) };
+    const rest = new QuantizedLinear(rowsOf(lin.w, n, N), rowsOf(lin.scales, n, N), rowsOf(lin.biases, n, N), lin.spec);
+    return new AneAffineSplit(program, head, rest, gpuRows ? gpuRows(rest) : new QuantizedLinearPlain(rest), lin.spec, lin);
+  }
+
+  /** x [.., M, K], M <= rows → [.., M, N], unevaluated as the MLP split's
+   *  (mlpSplit): the ANE starts once its buffers are written, the GPU's rows
+   *  queued after. */
+  forward(x: MlxArray): MlxArray {
+    const ane = this.program, K = ane.cin, n = ane.cout, rows = ane.seq;
+    const lead = x.shape.slice(0, -1), M = lead.reduce((a, b) => a * b, 1);
+    if (M > rows) throw new Error(`AneAffineSplit: ${M} rows exceed its program's ${rows}`);
+    {
+      using dense = ops.dequantize(this.head.w, this.head.scales, this.head.biases, this.spec);
+      using flat = ops.reshape(dense, [n * K]);
+      using x2 = ops.reshape(x, [M * K]);
+      using d0 = writeHalf(ane.weights(), flat, 0);
+      using d1 = writeHalf(ane.input(), x2, 0);
+      ops.evalAll([d0, d1]);
+    }
+    ane.evalAsync();
+    let gpu: MlxArray;
+    try { gpu = this.gpu.forward(x); ops.asyncEvalAll([gpu]); }
+    finally { ane.wait(); }
+    try {
+      using used = ane.output(M);                             // the piece's rows of [rows, n] fp16
+      using b = used.astype(x.dtype);
+      using g2 = ops.reshape(gpu, [M, this.rest.outFeatures]);
+      const both = ops.concatAxis([b, g2], -1);
+      ane.program.deferOutput(both);                          // evaluated (free once the result is) and freed by the program
+      return this.linear.delta.forward(x, ops.reshape(both, [...lead, n + this.rest.outFeatures]));
+    } finally { gpu.dispose(); }
+  }
+
+  [Symbol.dispose](): void { this.dispose(); }
+  /** Releases this layer's program and frees its weight rows; a shared program or buffers outlive it while another layer holds them. */
+  dispose(): void {
+    this.program.dispose();
+    for (const a of [this.head.w, this.head.scales, this.head.biases, this.rest.w, this.rest.scales, this.rest.biases!]) a.dispose();
+  }
+}