Add ANE quantization probe and expand project scope to GPT-2

[m0at]

Summary

This PR adds a comprehensive quantization probe for Apple Neural Engine (quant_probe.m) and significantly expands the project README to document GPT-2 inference, M5 hardware investigation, and fused kernel benchmarking work.

Key Changes

New Files

• quant_probe.m (551 lines) — Systematic probe to determine if ANE executes int8/int4 quantized operations natively or only dequantizes to fp16:
  • Tests 5 quantization schemes: FP16 baseline, INT8 affine dequantize, UINT4 affine dequantize, UINT4 blockwise (block_size=32), and LUT4 palettized
  • Generates MIL programs at runtime for each quantization type with proper weight blob builders
  • Benchmarks each variant across representative transformer dimensions (768→768, 768→2048, 2048→768, 1024→1024, 2048→2048)
  • Interprets results: 2x speedup indicates native int8, 4x indicates native int4, same speed indicates dequant-only
  • Includes weight blob builders for fp16, int8, int4 (packed nibbles), scale/zero-point tensors, and LUT indices

Documentation Updates

• README.md — Completely restructured to reflect project evolution:
  • Reframed as "Running Transformers on Apple's Neural Engine" with GPT-2 as primary focus
  • Added GPT-2 inference section with two-phase architecture (ANE prefill + CPU KV-cache decode)
  • Documented M5 hardware investigation findings (weight swapping, weightsBuffer override, QoS effects, chaining API rejection)
  • Added fused kernel benchmark results and dispatch overhead analysis
  • Expanded file structure to include all new test files and investigation documents
  • Clarified that project extends upstream maderix/ANEtransformers work

Implementation Details

Quantization Probe Architecture:

• Reuses ANE private API boilerplate (class resolution, IOSurface creation, mach_timebase timing)
• MIL generators for each quantization scheme with proper tensor type annotations (int8, uint4, uint8)
• Uses constexpr_affine_dequantize for per-channel quantization and constexpr_blockwise_shift_scale for block-wise int4
• Uses constexpr_lut_to_dense for palettized weights with shared 16-entry lookup table
• Benchmarks with 10 warmup iterations + 100 timed iterations, reports ms/eval and TFLOPS

Weight Blob Format:

• Consistent 128-byte header structure across all blob types (magic bytes, size, offset)
• FP16: 2 bytes per weight; INT8: 1 byte per weight; INT4: packed nibbles (2 per byte)
• Scale/zero-point blobs for affine dequantization; blockwise scale as [oc, n_blocks, 1, 1]
• LUT blob as [1, 1, 16] shared across all output channels

The probe enables empirical determination of ANE quantization capabilities without relying on undocumented compiler behavior or reverse-engineering binary outputs.

https://claude.ai/code/session_01U5HLjsm4iUzL9iDaHbxeRB