Universal Quantization Router: Formats, Bitstreams & Layouts

This document provides a comprehensive technical reference for the Universal Quantization Router in m4-prefill-engine. It details the binary layout, unpacker implementations, SRAM staging strategies, layout routing transforms, and empirical hardware findings across six supported quantization formats.

1. Architectural Motivation

Open-source LLM inference on Apple Silicon has historically suffered from format fragmentation:

• MLX exclusively optimizes its proprietary affine 4-bit layout (MLX_4BIT).
• llama.cpp focuses on GGUF quantization formats (Q4_0, Q4_K, Q8_0).
• ExLlamaV2/V3 utilizes variable bit-depth affine and hierarchical codebooks on CUDA GPUs.
• BitNet introduces 1.58-bit ternary representations (${-1, 0, +1}$).

The Universal Quantization Router decouples the physical weight format from the computational core. It decodes all six formats on-the-fly inside high-speed threadgroup SRAM (sh_B[32][64]), feeding a single unified 2D BlockMMA execution engine.

 ┌──────────────────────────────────────────────────────────────────────────────────┐
 │                       UNIVERSAL QUANTIZATION ROUTING MATRIX                      │
 └──────────────────────────────────────────────────────────────────────────────────┘
   Input Activation A [M, K]         Quantized Weight Storage B [N, K/Block]
 ┌───────────────────────────┐     ┌──────────────────────────────────────────────┐
 │ • Standard [M, K]         │     │ • Q4_0 (GGUF 32-elem block, 4.50 bpw)        │
 │ • Direct-Head [H, M, D]   │     │ • MLX 4-bit (64-elem affine, 5.00 bpw)       │
 │                           │     │ • Q4_K (GGUF 256-elem super-block, 4.50 bpw) │
 └─────────────┬─────────────┘     │ • Var-Rate Affine (3/4/5-bit, 5.00 bpw)      │
               │                   │ • EXL3 (Hierarchical codebook, 4.50 bpw)     │
               │                   │ • Ternary 1.58-bit (BitNet, 3.00 bpw)        │
               │                   └──────────────────────┬───────────────────────┘
               │                                          │
               ▼                                          ▼
 ┌──────────────────────────────────────────────────────────────────────────────────┐
 │                    ON-THE-FLY UNPACKING & SRAM STAGING                           │
 │     Vectorized Bit-Shift Dequantization -> Threadgroup Memory sh_B[32][64]       │
 └────────────────────────────────────────┬─────────────────────────────────────────┘
                                          │
                                          ▼
 ┌──────────────────────────────────────────────────────────────────────────────────┐
 │                 UNIFIED 2D BLOCKMMA HARDWARE EXECUTION CORE                      │
 │    simdgroup_matrix<half, 8, 8> fragments -> 16.8 TFLOPS Hardware Coprocessor    │
 └──────────────────────────────────────────────────────────────────────────────────┘

2. Format Specifications & Memory Layouts

1. Q4_0 (GGUF Standard Symmetric 4-Bit)

• Block Size: 32 elements.
• Weight Footprint: 18 bytes per block (16 bytes packed nibbles + 2 bytes FP16 scale).
• Effective Bit-Depth: 4.50 bits/weight.
• Dequantization Formula: \(w_i = (q_i - 8) \times d\)

struct block_q4_0 {
    half d;          // FP16 Scale Factor (2 bytes)
    uint8_t qs[16];  // 32 packed 4-bit nibbles (16 bytes)
};

// Fast Bitwise Vector Unpacker:
inline void unpack_q4_0_block(block_q4_0 blk, thread half4 vl[4], thread half4 vh[4]) {
    half d = blk.d;
    half4 hd = half4(d);
    half4 h_off = half4(-8.0h * d);

    uint w0 = read_u32_unaligned(blk.qs + 0);
    uint w1 = read_u32_unaligned(blk.qs + 4);
    uint w2 = read_u32_unaligned(blk.qs + 8);
    uint w3 = read_u32_unaligned(blk.qs + 12);

    vl[0] = fma(half4(as_type<uchar4>(w0 & 0x0F0F0F0Fu)), hd, h_off);
    vh[0] = fma(half4(as_type<uchar4>((w0 >> 4) & 0x0F0F0F0Fu)), hd, h_off);
    vl[1] = fma(half4(as_type<uchar4>(w1 & 0x0F0F0F0Fu)), hd, h_off);
    vh[1] = fma(half4(as_type<uchar4>((w1 >> 4) & 0x0F0F0F0Fu)), hd, h_off);
    vl[2] = fma(half4(as_type<uchar4>(w2 & 0x0F0F0F0Fu)), hd, h_off);
    vh[2] = fma(half4(as_type<uchar4>((w2 >> 4) & 0x0F0F0F0Fu)), hd, h_off);
    vl[3] = fma(half4(as_type<uchar4>(w3 & 0x0F0F0F0Fu)), hd, h_off);
    vh[3] = fma(half4(as_type<uchar4>((w3 >> 4) & 0x0F0F0F0Fu)), hd, h_off);
}

2. MLX 4-Bit (Apple MLX Affine 4-Bit)

• Block Size: 32 elements.
• Weight Footprint: 20 bytes per block (16 bytes packed nibbles + 2 bytes FP16 scale + 2 bytes FP16 bias).
• Effective Bit-Depth: 5.00 bits/weight.
• Dequantization Formula: \(w_i = q_i \times d + \text{bias}\)

struct block_mlx_4bit {
    half d;          // FP16 Scale Factor (2 bytes)
    half bias;       // FP16 Bias Factor  (2 bytes)
    uint8_t qs[16];  // 32 packed 4-bit nibbles (16 bytes)
};

3. Q4_K (GGUF Super-Block Quantization)

• Block Size: 256 elements (composed of eight 32-element sub-blocks).
• Weight Footprint: 144 bytes per super-block (128 bytes nibbles + 4 bytes FP16 scales + 12 bytes 6-bit sub-block scales).
• Effective Bit-Depth: 4.50 bits/weight.
• Dequantization Formula: \(w_i = (q_i - 8) \times d_{\text{global}} \times s_{\text{sub}}\)

struct block_q4_K {
    half d;            // Global super-block scale (2 bytes)
    half dmin;         // Global super-block minimum scale (2 bytes)
    uint8_t scales[12];// Packed 6-bit sub-block scales/mins (12 bytes)
    uint8_t qs[128];   // 256 packed 4-bit nibbles (128 bytes)
};

4. Grouped Variable-Rate Affine (EXL2-Style)

• Block Size: 256 elements (8 sub-blocks of 32 elements).
• Weight Footprint: 160 bytes per super-block (mixed bit-depth streams: 3-bit, 4-bit, and 5-bit sub-blocks).
• Effective Bit-Depth: 5.00 bits/weight.
• Architecture: Grouped affine quantization with variable bit-depths and permutation modes.

struct block_var_rate_affine {
    half d;              // Super-block global scale (2 bytes)
    half bias;           // Super-block global bias (2 bytes)
    uint8_t scales[8];   // 8 x sub-block scales (8 bytes)
    uint8_t biases[8];   // 8 x sub-block biases (8 bytes)
    uint8_t modes[8];    // 8 x sub-block mode & permutation metadata (8 bytes)
    uint8_t _pad[12];    // Alignment padding to 160 bytes (12 bytes)
    uint8_t qs[128];     // Multi-bit packed quant payload stream (128 bytes)
};

5. EXL3 (Hierarchical Vector Codebook)

• Block Size: 256 elements.
• Weight Footprint: 144 bytes per super-block.
• Effective Bit-Depth: 4.50 bits/weight.
• Architecture: Hierarchical vector codebook centroids ($16 \times 4\text{D}$ codebook), 3-bit packed indices, sub-block residual corrections, and FP16 global scale/bias.

struct block_exl3 {
    half d;              // Super-block global scale (2 bytes)
    half bias;           // Super-block global bias (2 bytes)
    uint8_t scales[8];   // 8 x sub-block scales (8 bytes)
    uint8_t residuals[8];// 8 x sub-block residual scales (8 bytes)
    int8_t codebook[16]; // 16 x hierarchical vector codebook centroids (16 bytes)
    uint8_t _pad[12];    // Alignment padding to 144 bytes (12 bytes)
    uint8_t qs[96];      // 256 x 3-bit packed index streams (96 bytes = 8 x 12 bytes)
};

6. Ternary 1.58-Bit (BitNet / $1.58\text{ bpw}$)

• Block Size: 32 elements.
• Weight Footprint: 12 bytes per block (8 bytes packed ternary trits + 2 bytes FP16 scale + 2 bytes padding).
• Alphabet: ${-1, 0, +1}$.
• Effective Bit-Depth: 3.00 bits/weight (1.58 bits theoretical entropy).

struct block_ternary_1_58 {
    half d;              // Shared FP16 scale factor (2 bytes)
    half _pad;           // 4-byte struct alignment padding (2 bytes)
    uint32_t qs[2];      // 2 x 32-bit = 64 bits = 32 x 2-bit weights in {-1, 0, +1} (8 bytes)
};

3. The Ternary Reality on Apple Silicon (MMA vs. Vector ALU)

The Theoretical Hypothesis

In BitNet literature, ternary quantization is celebrated for replacing expensive floating-point multiplications with pure additions and subtractions: \(y = \sum_{w_i \in \{+1, -1, 0\}} x_i \cdot \text{sign}(w_i)\)

The Empirical Finding

To test this on physical Apple Silicon, the router implements two parallel execution paths:

1. TERNARY_1_58_VEC (True Add/Sub Bypass): Dequantizes trits and executes conditional branch-free vector additions/subtractions directly on Vector ALUs.
2. TERNARY_1_58_MMA (Bandwidth-Only FMA Fallback): Unpacks ternary trits into shared memory sh_B and feeds them directly into the 16.8 TFLOPS Hardware Matrix Coprocessor (simdgroup_multiply_accumulate).

Empirical Comparison on Apple M4 (LLaMA 8B Tier, $K=4096$)

Why MMA Wins on Apple Silicon

• Compute Asymmetry: Apple Silicon’s Vector ALUs offer ~7.4 TFLOPS peak, whereas the Matrix Coprocessor delivers 16.8 TFLOPS. Even with zero multiplications, Vector ALUs cannot overcome a $2.27\times$ raw compute deficit.
• Instruction Issue Bandwidth: Packing and masking trits for ALU additions requires multiple integer bit-manipulation instructions per element, saturating the GPU execution pipeline.
• The True Value of Ternary on Apple Silicon: The benefit of Ternary on M4 is memory bandwidth, not arithmetic bypass. At $3.0\text{ bpw}$, a complete 8B layer’s weights require only 6.00 MB (fitting entirely within the M4’s 24MB System-Level Cache), completely eliminating DRAM traffic.

4. Direct-Head Layout Routing ([M, K] -> [H, M, D])

Standard transformer engines project activations in flat row-major layout: \(X_{\text{out}} \in \mathbb{R}^{M \times (H \cdot D)}\) This requires a subsequent transpose kernel to partition activations into multi-head format: \(X_{\text{head}} \in \mathbb{R}^{H \times M \times D}\)

The Direct-Head Advantage

The Universal Quantization Router implements Direct-Head Routing:

• The GEMM kernel writes directly into head-strided memory: \(\text{Global Address} = h \times (M \times D) + m \times D + d\)
• Zero Dynamic Transpositions: Eliminates intermediate DRAM transpositions.
• Zero Dynamic Padding: Handles arbitrary sequence lengths ($M=33, 127, 129, 2047$) with zero boundary padding overhead.

5. Comprehensive Router Benchmark Matrix

Measured on Apple M4 (10-Core GPU, 16GB UMA) using bench_universal_router with mandatory 32MB SLC cache flushing and double-precision CPU verification.

8B Transformer Tier ($K=4096, N=4096$)