Skip to content

SMOE Quantisation

S-MoE seals all routed expert weights into custom, hyper-efficient quantisation schemes—SMOE-Q2 (2-bit) and SMOE-Q4 (4-bit, the default since the Q4 pipeline proved that 2-bit uniform quantisation mathematically collapses fine-grained experts). This is not a compromise; it is an elegant solution to a fundamental physics constraint. Here is the mathematical rebellion that makes it lossless enough for true frontier-quality output, and the Metal kernel that executes it directly in register space.


Why Quantise: Conquering the NVMe Physics Constraint

Quantisation in S-MoE is not a trade-off for quality. It is a necessary mastery over bandwidth physics.

An Apple Silicon Mac's SSD delivers a magnificent ~7.4 GB/s. A single Qwen3-235B expert (three 4096 × 1536 matrices) at full bfloat16 precision weighs ~38 MB. With 8 experts active per layer and 94 MoE layers, a naive implementation would easily crush any NVMe drive.

Aggressive, intelligent quantisation is the weapon that makes streaming math possible (Qwen3-235B expert blobs, including scales):

PrecisionExpert blob sizeExperts/sec at 7.4 GB/s
bfloat16~38 MB~195
SMOE-Q4 (4-bit, default)~10 MB~740
SMOE-Q2 (2-bit)~5 MB~1,480

The Scheme: L2-MSE Optimal Symmetric Quantisation

SMOE deploys a ruthless grid-search optimiser to discover the mathematically perfect scale for each group of 64 weights.

Quantisation Levels

  • SMOE-Q2: 4 precise symmetric levels — [-1.0, -0.333, +0.333, +1.0]
  • SMOE-Q4: 16 symmetric levels across the normalized range

L2-MSE Scale Optimiser

Instead of blindly taking absmax as the scale (a naive approach that destroys sparsity), shatter_moe.py probes 64 candidate scales between 0.1 × absmax and 1.0 × absmax for every group of 64 weights. It selects the exact scale that minimises the L2 Mean Squared Error between the original and reconstructed weights.

python
# For a group of 64 bfloat16 weights:
absmax = max(abs(group))

# Grid search 64 candidate scales
candidate_scales = linspace(0.1, 1.0, 64) * absmax
best_scale = argmin(mse(group, decode(encode(group, s))) for s in candidate_scales)

normalized = group / best_scale   # ∈ [-1, +1]

# SMOE-Q2 encoding
code_q2 = clamp(round(normalized * 1.5 + 1.5), 0, 3)   # ∈ {0, 1, 2, 3}

# SMOE-Q4 encoding
code_q4 = clamp(round(normalized * 7.5 + 7.5), 0, 15)  # ∈ {0 .. 15}

Why the grid search matters: If a group is mostly zeroes with one loud outlier, absmax would enforce a massive scale, crushing the near-zero weights into noise. Our MSE optimiser deliberately chooses a smaller scale that clips the outlier—beautifully preserving the silence of the zero-weights. This natively defends MoE weight sparsity.

Bit Packing

  • SMOE-Q2: 4 codes tightly packed per byte, LSB-first (little-endian bit order)
  • SMOE-Q4: 2 codes seamlessly packed per byte, LSB-first

Choosing the Depth

The Sculptor defaults to --bits 4. Pass --bits 2 only for architectures whose experts tolerate it — on fine-grained frontier experts, Q2's four levels are too coarse and coherence collapses. Q4 is the production depth for Qwen3-235B.


Decoding: Metal Kernel (Fused Dequant-Multiply)

The Metal kernel liberates each weight directly in register space during the matrix-vector multiply. The weight is decoded and instantly multiplied by the input activation. A bloated float32 weight matrix is never materialized.

metal
// SMOE-Q2 decode (in-register, fused with multiply)
uint8_t code   = (packed[pack_idx] >> bit_shift) & 0x3;
float   scale  = (float)scales[group_idx];       // float16 → float32
float   weight = ((float)code - 1.5) / 1.5 * scale;
acc           += weight * input[col];            // elegant fused multiply-accumulate

// SMOE-Q4 decode
uint8_t code   = (packed[pack_idx] >> bit_shift) & 0xF;
float   weight = ((float)code - 7.5) / 7.5 * scale;
acc           += weight * input[col];

The kernel dynamically adapts based on the vault header bits field. A single, unified binary supports both Q2 and Q4 flawlessly without recompilation — the shaders themselves are JIT-compiled at startup from kernels.metal.


The Smoke Test

bash
make test-quant

This generates a random 128 × 64 tensor, aggressively quantises and dequantises it through the Q2 path (the same module powers Q4), and reports its triumph:

  RMSE=0.50xxx  absmax=1.2xxxx  SNR=5.x dB
  ✓  SMOE-Q2 smoke-test passed

Expected SNR Resonance

~6 dB on random Gaussian weights is standard for 2-bit quantisation. Real pre-trained MoE expert distributions are beautifully tight and yield a significantly higher SNR — and Q4's sixteen levels raise it further still.

MIT License.