chore: import upstream snapshot with attribution

2026-07-13 12:34:46 +08:00
commit f4e68ed970
84 changed files with 14896 additions and 0 deletions
@@ -0,0 +1,87 @@
+// ============================================================
+// Per-token quantization: FP16 → INT8 + float32 scale
+// Target: Apple M5, Metal 4
+//
+// Each threadgroup handles one row (one token).
+// Threads cooperate to find absmax via simdgroup reduce,
+// then quantize in parallel.
+//
+// Host dispatch:
+//   threadgroup = (min(256, ceil(K/32)*32), 1, 1)
+//   grid = (M, 1, 1)
+// ============================================================
+
+#include <metal_stdlib>
+using namespace metal;
+
+kernel void
+quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
+                   device int8_t *A [[buffer(1)]],     // [M, K] INT8 output
+                   device float *scale [[buffer(2)]],  // [M] float32 scale
+                   constant uint &M [[buffer(3)]],
+                   constant uint &K [[buffer(4)]],
+                   uint gid [[threadgroup_position_in_grid]], // row index
+                   uint lid [[thread_index_in_threadgroup]],
+                   uint tg_size [[threads_per_threadgroup]]) {
+  if (gid >= M) {
+    return;
+  }
+
+  const device half *row_in = X + gid * K;
+  device int8_t *row_out = A + gid * K;
+
+  // Step 1: Find local absmax
+  float local_max = 0.0f;
+  for (uint i = lid; i < K; i += tg_size) {
+    float v = abs(float(row_in[i]));
+    local_max = max(local_max, v);
+  }
+
+  // Step 2: Simdgroup reduce max
+  float sg_max = simd_max(local_max);
+
+  // Step 3: Threadgroup reduce across simdgroups via shared memory
+  threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
+  threadgroup float shared_scale;
+  threadgroup float shared_inv_scale;
+  uint sg_id = lid / 32;
+  uint sg_lid = lid % 32;
+  if (sg_lid == 0) {
+    sg_maxes[sg_id] = sg_max;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Final reduce (first simdgroup only)
+  if (sg_id == 0) {
+    float row_max = 0.0f;
+    uint num_sgs = (tg_size + 31) / 32;
+    if (sg_lid < num_sgs) {
+      row_max = sg_maxes[sg_lid];
+    }
+    row_max = simd_max(row_max);
+
+    // Compute and broadcast scale
+    float s = row_max / 255.0f;
+    if (s == 0.0f) {
+      s = 1.0f;
+    }
+
+    if (sg_lid == 0) {
+      shared_scale = s;
+      shared_inv_scale = 1.0f / s;
+      // Store scale to output
+      scale[gid] = s;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Step 4: All threads read broadcasted scale
+  float inv_s = shared_inv_scale;
+
+  // Step 5: Quantize
+  for (uint i = lid; i < K; i += tg_size) {
+    float v = float(row_in[i]) * inv_s;
+    v = clamp(round(v), -128.0f, 127.0f);
+    row_out[i] = int8_t(v);
+  }
+}