# Which exl3 codebook is the better QUANTISER, and the answer to "should we just
# convert the weights to mul1 for speed?"  (#1861)
#
#   python3 codebook_quality.py
#
# NO.  mcg is the better codebook at every width this repo serves, and the gap
# WIDENS with K: 1.46x at 4 bpw, 2.24x at 5, 4.04x at 6 -- 6 bpw being the
# DFlash2 drafter's width.  mul1 wins only at K=2, which nothing here uses.
#
# WHAT IS MEASURED.  The per-step resolution FLOOR of the trellis:
# E_s E_w [ min over the 2^K reachable successors of (w - v)^2 ], w ~ N(0,1),
# each codebook first normalised to unit variance.  A Viterbi search over the
# sequence does better than picking the nearest value at each step, so this is
# NOT the achieved distortion -- it is a bound, computed identically for both,
# on the resolution the search gets to choose from.
#
# THE DIRECTION OF THE TRELLIS STEP DECIDES THE VERDICT, so it is read off the
# kernel rather than assumed: `win8_aligned<4>` emits out[7] = b & 0xffff and
# out[6] = bits[4:20], so weight t+1's window is weight t's shifted RIGHT by K
# with K new bits entering at the top.  Assuming a LEFT shift reverses the
# result -- it says mul1 wins at every K.  That wrong answer was produced once
# in this session before the direction was checked.
#
# THE MECHANISM IS NOT WHAT IT LOOKS LIKE.  mcg has 10746 distinct fp16 values
# and mul1 has 913, which invites a story about collisions -- and that story is
# WRONG: both codebooks reach essentially all 2^K distinct values per step
# (measured: 15.99 and 16.00 of 16 at K=4; 63.77 and 64.00 of 64 at K=6).
#
# What actually separates them is where those values SIT.  mul1's reachable set
# is narrower AND BADLY CENTRED, and the offset does not improve with K:
#
#          set std      set span     |set mean|
#   K=4    mcg 0.944    mcg 3.395    mcg  0.105
#          mul1 0.844   mul1 3.137   mul1 0.442
#   K=6    mcg 0.994    mcg 4.733    mcg  0.037     <- centres as K grows
#          mul1 0.865   mul1 3.911   mul1 0.434     <- does not
#
# mul1's value is an affine function of a BYTE SUM of idx * 0x83DCD12D, and the
# K new bits enter at the TOP of idx -- so they move the high bytes of the
# product and barely touch the low ones.  The reachable set is therefore a
# cluster whose centre is pinned by the state's fixed low-order bits, roughly
# 0.44 sigma off zero however many successors there are.  mcg's lop3 over the
# whole 32-bit MCG product mixes every byte, so its set both spreads and
# centres as K grows.  That is a structural property of the two constructions,
# not a tuning artefact, and it is presumably why exl3 makes mcg the default.
#
# THE TRADE, STATED PLAINLY.  Converting the 13 mcg checkpoints here to mul1
# would buy 3.2% of the int8 GEMM's runtime (receipt_bench_codebook_power.txt)
# and cost 1.5-4x the per-step quantiser resolution.  It is the wrong direction.
#
# Lloyd-Max is printed as a SCALE, not a target: it is the optimal SCALAR
# quantiser at 2^K levels, and a trellis quantiser beats it by spending state.

per-step reachable-set distortion, N(0,1) source, unit-variance codebooks
 K          mcg         mul1   mul1/mcg   Lloyd-Max 2^K
 2    4.271e-01    3.292e-01     0.771x       1.175e-01
 3    1.237e-01    1.845e-01     1.491x       3.455e-02
 4    6.526e-02    9.496e-02     1.455x       9.495e-03
 5    2.342e-02    5.254e-02     2.244x       2.498e-03
 6    8.408e-03    3.397e-02     4.040x       6.340e-04

why: the shape of the reachable set (in units of the source's sigma)
 K  book   set std   set span   |set mean|
 2   mcg    0.7295     1.8828       0.4950
 2  mul1    0.7822     2.0881       0.4681
 3   mcg    0.9275     2.9230       0.1732
 3  mul1    0.8300     2.6256       0.4472
 4   mcg    0.9445     3.3949       0.1048
 4  mul1    0.8438     3.1365       0.4416
 5   mcg    0.9804     4.1243       0.0533
 5  mul1    0.8612     3.5972       0.4351
 6   mcg    0.9937     4.7330       0.0375
 6  mul1    0.8645     3.9108       0.4335
