fmt=7: fp8-e4m3-b128 passthrough weight format — CPU + Metal + repack tool (#524)

c/Makefile

@@ -199,7 +199,7 @@ else
 PYTHON    ?= python3
 endif
 CUDA_OBJ  =
-TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_mirror$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_temp_env$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_router_nan$(EXE) tests/test_pipe_block$(EXE) tests/test_e8_kernel$(EXE) tests/test_rans$(EXE)
+TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_st_mirror$(EXE) tests/test_st_pread$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_topp$(EXE) tests/test_sample_nan$(EXE) tests/test_temp_env$(EXE) tests/test_tok_o200k$(EXE) tests/test_kv_alloc$(EXE) tests/test_int3$(EXE) tests/test_int3_load$(EXE) tests/test_fp8_passthrough$(EXE) tests/test_fp8_load$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE) tests/test_logit_nan$(EXE) tests/test_router_nan$(EXE) tests/test_pipe_block$(EXE) tests/test_e8_kernel$(EXE) tests/test_qt_addrow$(EXE) tests/test_rans$(EXE)
 ifneq (,$(LINUX))
 TEST_BINS += tests/test_uring$(EXE)
 endif
@@ -478,6 +478,15 @@ tests/test_int3$(EXE): tests/test_int3.c colibri.c st.h uring.h json.h tok.h tok
 tests/test_int3_load$(EXE): tests/test_int3_load.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
+tests/test_fp8_passthrough$(EXE): tests/test_fp8_passthrough.c quant.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
+tests/test_fp8_load$(EXE): tests/test_fp8_load.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
+tests/test_qt_addrow$(EXE): tests/test_qt_addrow.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
 tests/test_logit_nan$(EXE): tests/test_logit_nan.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 

c/backend_metal.h

@@ -27,9 +27,14 @@ void coli_metal_stats(size_t *tensor_count, size_t *tensor_bytes);
 int  coli_metal_mem_info(size_t *used_bytes, size_t *total_bytes);
 
 /*
- * y[S,O] = (x[S,I] @ W[O,I]^T) * scale[o].
- * fmt matches QT in glm.c: 0=f32, 1=int8, 2=int4(packed), 3=int2(packed).
- * The first successful call wraps W and its row scales in GPU-visible buffers;
+ * y[S,O] = (x[S,I] @ W[O,I]^T) * scale[o]. fmt=7 (fp8 passthrough -- see
+ * colibri.c) instead folds a per-128x128-block scale into the accumulation
+ * -- see the shader comment in backend_metal.mm.
+ * fmt matches QT in colibri.c: 0=f32, 1=int8, 2=int4(packed), 3=int2(packed),
+ * 7=fp8-e4m3 (one raw byte/element, same layout as fmt=1) + per-128x128-
+ * block scale (scale array is [ceil(O/128),ceil(I/128)] floats; no group-size
+ * parameter needed -- the block is a fixed 128x128, not caller-configurable).
+ * The first successful call wraps W and its scales in GPU-visible buffers;
  * later calls reuse them (weights are assumed stable at the same address).
  * Returns 1 on success, 0 if Metal is unavailable or fmt is invalid.
  */

c/backend_metal.mm

@@ -9,6 +9,16 @@
 
 // ---- shader: general quantized GEMV, one threadgroup per output element (o,si) ----
 // y[si,o] = (sum_i dequant(W[o,i]) * x[si,i]) * scale[o]. fmt: 0=f32 1=i8 2=i4 3=i2.
+// fmt=7 (native FP8-e4m3 passthrough -- see colibri.c): raw byte
+// layout identical to fmt=1 (one e4m3 byte per element, no packing), but the scale is
+// per-128x128 BLOCK of [O,I] -- buffer(1) holds [ceil(O/128),ceil(I/128)] floats indexed
+// scale[(o/128)*nblkI+i/128]. Folded into acc like a per-group scale would be, for the
+// same reason: it is not constant across one output row. e4m3->float is bit manipulation
+// in-kernel (sign/exp/mantissa, OCP E4M3-FN: exp==0 subnormal, exp==0xF&&mant==0x7 is the
+// only NaN code) -- BW-bound kernel, decode ALU is free, no LUT texture. Must byte-match
+// quant.h's E4M3_LUT/e4m3_decode (the CPU reference) including the NaN policy, which is
+// why the metal-test suite runs all 256 byte codes through this exact kernel path (not
+// just spot values) -- see run_fp8_lut().
 static const char *SHADER = R"METAL(
 #include <metal_stdlib>
 using namespace metal;
@@ -53,14 +63,34 @@ kernel void mm_gemv(device const uchar* w      [[buffer(0)]],   // raw weight by
     for (int i = slane; i < I; i += 32) {
       uchar b = wr[i>>2]; int v = (b >> (2*(i&3))) & 0x3; acc += float(v-2) * xr[i];
     }
+  } else if (fmt == 7) {                            // fp8 e4m3 passthrough: one raw byte per
+                                                     // element (like fmt=1), scale per 128x128
+                                                     // block folded into acc (like a grouped fmt).
+    int nblkI = (I + 127) / 128;
+    device const uchar* wr = w + (long)o * I;
+    device const float* scl = scale + (long)(o/128) * nblkI;
+    for (int i = slane; i < I; i += 32) {
+      uchar b = wr[i];
+      uint sign = b >> 7, exp = (b >> 3) & 0xF, mant = b & 0x7;
+      float wv;
+      if (exp == 0xF && mant == 0x7) {
+        wv = as_type<float>(0x7fc00000u);           // qNaN -- matches quant.h's e4m3_decode
+      } else {
+        float mag = (exp == 0) ? (float(mant) * 0.001953125f)                 // subnormal: mant*2^-9
+                                : (1.0f + float(mant)*0.125f) * exp2(float(int(exp) - 7));
+        wv = sign ? -mag : mag;
+      }
+      acc += wv * xr[i] * scl[i/128];
+    }
   } else {                                          // f32
     device const float* wr = (device const float*)(w) + (long)o * I;
     device const float4* w4 = (device const float4*)wr;
     for (int c = slane; c < I8; c += 32) acc += dot(w4[2*c],x4[2*c]) + dot(w4[2*c+1],x4[2*c+1]);
     for (int i = I8*8 + slane; i < I; i += 32) acc += wr[i] * xr[i];
   }
   acc = simd_sum(acc);
-  if (slane == 0) y[row] = acc * scale[o];
+  // fmt==7 already folds its per-block scale into acc above -- do not scale again.
+  if (slane == 0) y[row] = (fmt == 7) ? acc : acc * scale[o];
 }
 
 // Batched bindless expert GEMV: each row gr belongs to expert erow[gr], whose weight and
@@ -423,8 +453,19 @@ static size_t fmt_bytes(int fmt, int I, int O) {
   if (fmt == 1) return (size_t)O * I;
   if (fmt == 2) return (size_t)O * ((I+1)/2);
   if (fmt == 3) return (size_t)O * ((I+3)/4);
+  if (fmt == 7) return (size_t)O * I;           // fp8 e4m3: one raw byte/element, same as fmt=1
   return (size_t)O * I * sizeof(float);
 }
+// fp8 (fmt=7) scale-array size: one f32 per 128x128 BLOCK -> ceil(O/128)*ceil(I/128) (2D,
+// not per-row -- quant.h isn't included here, so the ceil-div is inlined rather than sharing
+// colibri.c's qt_scale_bytes/quant.h's fp8_nblk). f32 is this build's implemented scale
+// ENCODING for fmt=7 (see quant.h/colibri.c) -- this file has no reason to know that a
+// UE8M0 encoding exists at all: qt_resolve_fmt refuses it on the CPU read path before any
+// tensor in that encoding could ever reach this Metal-side sizing helper.
+static size_t fmt_scale_bytes(int fmt, int I, int O) {
+  if (fmt == 7) return (size_t)((O + 127) / 128) * (size_t)((I + 127) / 128) * sizeof(float);
+  return (size_t)O * sizeof(float);
+}
 
 // Wrap host memory zero-copy if page-aligned, else copy into a shared buffer.
 static id<MTLBuffer> wrap(const void *p, size_t n) {
@@ -590,14 +631,16 @@ static size_t fmt_bytes(int fmt, int I, int O) {
 extern "C" int coli_metal_matmul(ColiMetalTensor **tp, float *y, const float *x,
                                  const void *weights, const float *scales,
                                  int fmt, int S, int I, int O) {
-  if (!g_dev || fmt < 0 || fmt > 3) return 0;
+  /* fmt==7 (fp8 passthrough) is an explicit allow-list entry, not folded into the 0..3
+   * contiguous range check below: it is not adjacent to it. */
+  if (!g_dev || fmt < 0 || (fmt > 3 && fmt != 7)) return 0;
   @autoreleasepool {
     ColiMetalTensor *t = *tp;
     if (!t) {
       t = new ColiMetalTensor();
       t->fmt = fmt; t->I = I; t->O = O; t->wbytes = fmt_bytes(fmt, I, O);
       t->w = wrap(weights, t->wbytes);
-      t->s = wrap(scales, (size_t)O * sizeof(float));
+      t->s = wrap(scales, fmt_scale_bytes(fmt, I, O));
       *tp = t;
       g_tensor_count++; g_tensor_bytes += t->wbytes;
     }