From fb4aba041a4eaba640910e9095ec85a0e3678e0b Mon Sep 17 00:00:00 2001 From: Shaw Date: Mon, 19 Oct 2026 21:57:35 -0700 Subject: [PATCH] tools/parakeet: fv5-b5x3t7-v1 code layout, Phonon-2 GGUF 331 -> 209 MB The five-value encoder weights were stored as 4-bit level codes (fv5-pair16-v1), so the GGUF was 330.6 MB, against 177 MB for Fermion's own container. fv5-b5x3t7-v1 stores three panel vectors' levels as one base-5 code per lane (< 125, so 7 bits) and packs 8 code vectors into 7: 2.35-2.41 bits per weight, 208.9 MB in all. Both layouts decode to the same int8 panels, so results are bit-identical. Every tier gets a decoder: scalar, NEON (dotprod and i8mm), and AVX2. The loader accepts both formats; the converter writes b5x3t7 by default and pair16 with --layout pair16. --verify checks both against the container. tools/bench-fv5-packing.cpp measures the decode cost per GEMM shape. Co-Authored-By: Claude Opus 5.5 --- tools/parakeet/CMakeLists.txt | 7 + tools/parakeet/README.md | 33 +++- tools/parakeet/convert_phonon2_to_gguf.py | 96 ++++++++-- tools/parakeet/pins.json | 14 +- tools/parakeet/src/parakeet.cpp | 20 +- tools/parakeet/src/pk-kernels-avx2.cpp | 57 +++++- tools/parakeet/src/pk-kernels-dotprod.cpp | 2 +- tools/parakeet/src/pk-kernels-i8mm.cpp | 2 +- tools/parakeet/src/pk-kernels-neon.inc | 68 +++++++ tools/parakeet/src/pk-kernels.cpp | 25 ++- tools/parakeet/src/pk-kernels.h | 18 ++ tools/parakeet/src/pk-model.h | 1 + tools/parakeet/src/pk-ops.cpp | 16 +- tools/parakeet/src/pk-ops.h | 8 +- .../parakeet/tests/test_parakeet_kernels.cpp | 60 +++++- tools/parakeet/tools/bench-fv5-packing.cpp | 175 ++++++++++++++++++ 16 files changed, 565 insertions(+), 37 deletions(-) create mode 100644 tools/parakeet/tools/bench-fv5-packing.cpp diff --git a/tools/parakeet/CMakeLists.txt b/tools/parakeet/CMakeLists.txt index d7a049fac..f68ab1ad2 100644 --- a/tools/parakeet/CMakeLists.txt +++ b/tools/parakeet/CMakeLists.txt @@ -95,6 +95,13 @@ if (NOT CMAKE_SYSTEM_NAME STREQUAL "iOS" AND NOT CMAKE_SYSTEM_NAME STREQUAL "And endif() endif() +if (NOT WIN32 AND NOT CMAKE_SYSTEM_NAME STREQUAL "iOS") + # weight-layout study (not installed): cost of denser fv5 code packings + add_executable(bench-fv5-packing tools/bench-fv5-packing.cpp) + target_link_libraries(bench-fv5-packing PRIVATE parakeet_lib) + target_include_directories(bench-fv5-packing PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) +endif() + if (LLAMA_BUILD_TESTS OR PARAKEET_BUILD_TESTS) add_subdirectory(tests) endif() diff --git a/tools/parakeet/README.md b/tools/parakeet/README.md index 8556c04f2..d1c0e5078 100644 --- a/tools/parakeet/README.md +++ b/tools/parakeet/README.md @@ -40,9 +40,33 @@ directory is new code. ## Kernels (`src/pk-kernels*.cpp`) -The weights stay in the GGUF in kernel order (`fv5-pair16-v1`): 4-bit level codes, two rows by 16 k per -16-byte block. The model is mmapped and used in place. Per 8-row block, one `vqtbl1q_s8` per 16 weights -decodes the codes into an int8 panel that stays in L1. The panel is then multiplied against all tokens: +The weights stay in the GGUF in kernel order and are used in place from the mmapped file. Per 8-row block, +the codes are decoded into an int8 panel that stays in L1, and the panel is then multiplied against all +tokens. Two code layouts decode to the same panel, so they give bit-identical results: + +| layout | bits per weight | Phonon-2 GGUF | decode per 16 weights | +|---|---|---|---| +| `fv5-b5x3t7-v1` (default) | 2.35–2.41 | 208.9 MB | 1/3 of: 2 `UMULL` + `UZP2` (x / 5), `TBL` (/ 25), 2 `MLS`, then 3 x (`ORR` + `TBL`); 1/8 of a 7-`SRI` bit gather | +| `fv5-pair16-v1` | 4 | 330.6 MB | `AND` or `USHR`, one `vqtbl1q_s8` | + +- **`fv5-pair16-v1`**: 4-bit level codes, two rows by 16 k per 16-byte block. +- **`fv5-b5x3t7-v1`**: the panel's 16-byte vectors are taken three at a time. Lane j of a code vector is the + base-5 number of the three levels, which is below 125, so 8 code vectors share 7 stored vectors (the eighth + code is spread over their top bits). `pk-kernels.h` specifies it. +- The five-value alphabet carries log2(5) = 2.32 bits, and the levels' entropy is 1.95 bits. Fermion's own + container gets 2.11 bits with a variable-length bit per non-zero weight, but that cannot be decoded at an + arbitrary block without prefix counts. + +What the denser layout costs is decode time. It is fixed per encoder pass, so it only shows on short windows. +These are `tools/bench-fv5-packing.cpp` results, decode plus tile over all 24 layers, on one thread: + +| N tokens (80 ms each) | 8 | 16 | 32 | 64 | 125 | 375 | +|---|---|---|---|---|---|---| +| TB336FU Cortex-A76, dotprod: pair16 → b5x3t7 (ms) | 101 → 121 | 176 → 191 | 324 → 330 | 596 → 613 | 1139 → 1142 | 3661 → 3602 | +| M4 Max, i8mm: pair16 → b5x3t7 (ms) | 25 → 33 | 45 → 52 | 85 → 91 | 167 → 171 | 321 → 326 | 957 → 960 | + +The converter writes `fv5-b5x3t7-v1` unless given `--layout pair16`. The loader accepts both. + | tier | instruction | runtime check | |---|---|---| @@ -97,6 +121,7 @@ a **pause-delimited segment**, from the segment start to the newest sample. - `convert_phonon2_to_gguf.py` reads the pinned `phonon-2.bps.tar.zst`, or an unpacked `model.fermion` and `config.json`, and writes the GGUF. - It checks every input against `pins.json` (HF revision `7e153e40…`, size and sha256). + - `--layout b5x3t7` (default) or `--layout pair16` picks the code layout. - `--verify` re-reads the GGUF and checks that all 264 five-value records and 11 int tables expand bit-identically to the container, that all 723 records are consumed, and that the mel filterbank equals librosa's. @@ -105,6 +130,8 @@ a **pause-delimited segment**, from the segment start to the newest sample. fp32 `transformers.ParakeetForTDT`. The fixtures hold the PCM, the normalised mel, the subsampling, layer-0 and encoder outputs, and the token ids and text. - `tests/test_parakeet_parity.cpp` compares the native pipeline with those fixtures stage by stage. +- `tools/bench-fv5-packing.cpp` times the two code layouts' decode and decode + tile per GEMM shape, and + checks that they decode to the same panel. - `tools/parakeet-bench.cpp` and `tools/bench_parakeet.py` measure RTFx (wall and CPU time), first-window partial latency, peak RSS and WER, plus stream WER and stream/batch agreement with `--stream`. - libelizainference integration lives in `tools/omnivoice/src/eliza-asr-parakeet.inc`. diff --git a/tools/parakeet/convert_phonon2_to_gguf.py b/tools/parakeet/convert_phonon2_to_gguf.py index 91cb562bb..bc4f54e4d 100644 --- a/tools/parakeet/convert_phonon2_to_gguf.py +++ b/tools/parakeet/convert_phonon2_to_gguf.py @@ -19,10 +19,14 @@ fp16: raw half floats. GGUF layout written here (all lossless with respect to the container): - * five-value linears -> `.fv5` (I8 bytes [O/2, I], 4-bit level - indices in the `fv5-pair16-v1` kernel order, see pack_pair16; level 0..4 - = {-hi,-lo,0,+lo,+hi}) and `.lohi` (F32 [O,2]). q/k/v are - concatenated into one `attn.qkv`. The file is mmapped and used in place. + * five-value linears (level 0..4 = {-hi,-lo,0,+lo,+hi}) -> the codes plus + `.lohi` (F32 [O,2]). q/k/v are concatenated into one `attn.qkv`. + The file is mmapped and used in place. `--layout` picks the code layout: + b5x3t7 (default): `.fv5t7`, I8 [O/8, fv5t7 block bytes], the + `fv5-b5x3t7-v1` order (see pack_b5x3t7), 2.35-2.41 bits per weight; + pair16: `.fv5`, I8 [O/2, I], 4-bit level indices in the + `fv5-pair16-v1` order (see pack_pair16), 4 bits per weight. + Both decode to the same int8 panels, so they give identical results. * int6 tables -> `.q` (I8, the exact integers) + `.scale` (F32). * fp16 records -> F32 (exact widening). * BatchNorm is folded into the depthwise conv (F32 weight + bias); this is @@ -42,6 +46,7 @@ convert_phonon2_to_gguf.py --archive phonon-2.bps.tar.zst --output phonon-2-fv5.gguf [--verify] convert_phonon2_to_gguf.py --container model.fermion --nemo-config config.json \ --notice NOTICE --output phonon-2-fv5.gguf [--verify] + add --layout pair16 for the 4-bit `fv5-pair16-v1` layout """ from __future__ import annotations @@ -300,6 +305,64 @@ def unpack_pair16(packed: np.ndarray) -> np.ndarray: return seq.reshape(2 * p, k).astype(np.uint8) +def fv5t7_block_bytes(k: int) -> int: + t = (k // 2 + 2) // 3 + return (t + 7) // 8 * 7 * 16 + + +def pack_b5x3t7(levels: np.ndarray) -> np.ndarray: + """Kernel layout `fv5-b5x3t7-v1`, [O/8, fv5t7_block_bytes(K)] bytes. + + Per block of 8 rows, take the K/2 16-byte panel vectors of the int8 GEMM + (vector v = 4c + p holds row 2p at k 8c..8c+7, then row 2p+1) as level + indices, three vectors at a time: code lane j = l(3t) + 5 l(3t+1) + + 25 l(3t+2) < 125 (vectors past K/2 count as level 2). Codes fit in 7 bits, + so 8 code vectors share 7 stored vectors: byte m of lane j = code m | + (bit m of code 7) << 7; a short last group is padded with code 62 (levels + 2, 2, 2). A NEON kernel recovers the digits with two multiply-highs and + two multiply-subtracts per 16 codes, then makes each panel vector with + the same 16-byte LUT lookup as pair16.""" + o, k = levels.shape + if o % 8 or k % 16: + die(f"five-value shape {levels.shape}: rows must be a multiple of 8 and width a multiple of 16") + nb, v = o // 8, k // 2 + t = (v + 2) // 3 + g = (t + 7) // 8 + pan = levels.astype(np.uint8).reshape(nb, 4, 2, k // 8, 8).transpose(0, 3, 1, 2, 4).reshape(nb, v, 16) + pan = np.concatenate([pan, np.full((nb, 3 * t - v, 16), 2, np.uint8)], axis=1).astype(np.uint16) + pan = pan.reshape(nb, t, 3, 16) + code = pan[:, :, 0] + 5 * pan[:, :, 1] + 25 * pan[:, :, 2] + code = np.concatenate([code, np.full((nb, 8 * g - t, 16), 62, np.uint16)], axis=1).reshape(nb, g, 8, 16) + out = np.empty((nb, g, 7, 16), np.uint8) + for m in range(7): + out[:, :, m] = code[:, :, m] | (((code[:, :, 7] >> m) & 1) << 7) + return out.reshape(nb, g * 112) + + +def unpack_b5x3t7(packed: np.ndarray, k: int) -> np.ndarray: + nb, nbytes = packed.shape + if nbytes != fv5t7_block_bytes(k): + die(f"unpack_b5x3t7: {nbytes} bytes per block, expected {fv5t7_block_bytes(k)} for K={k}") + v = k // 2 + t = (v + 2) // 3 + g = (t + 7) // 8 + b = packed.reshape(nb, g, 7, 16).astype(np.uint16) + code = np.empty((nb, g, 8, 16), np.uint16) + code[:, :, :7] = b & 0x7F + code[:, :, 7] = sum(((b[:, :, m] >> 7) << m) for m in range(7)) + code = code.reshape(nb, 8 * g, 16) + if code[:, :t].max() > 124 or (8 * g > t and np.any(code[:, t:] != 62)): + die("unpack_b5x3t7: code out of range") + pan = np.stack([code[:, :t] % 5, code[:, :t] // 5 % 5, code[:, :t] // 25], axis=2).reshape(nb, 3 * t, 16) + if np.any(pan[:, v:] != 2): + die("unpack_b5x3t7: padding is not level 2") + pan = pan[:, :v].reshape(nb, k // 8, 4, 2, 8).transpose(0, 2, 3, 1, 4) + return pan.reshape(nb * 8, k).astype(np.uint8) + + +LAYOUTS = {"b5x3t7": ("fv5t7", "fv5-b5x3t7-v1"), "pair16": ("fv5", "fv5-pair16-v1")} + + def fv5_expand(levels: np.ndarray, lohi: np.ndarray) -> np.ndarray: lo = lohi[:, 0:1].astype(np.float16) hi = lohi[:, 1:2].astype(np.float16) @@ -314,11 +377,13 @@ def container_fv5_expand(r: dict) -> np.ndarray: # --------------------------------------------------------------- writer class Builder: - def __init__(self, recs: dict): + def __init__(self, recs: dict, layout: str): self.recs = recs + self.layout = layout self.used: set[str] = set() self.tensors: list[tuple[str, np.ndarray]] = [] self.fv5_src: dict[str, list[str]] = {} + self.fv5_k: dict[str, int] = {} self.int_src: dict[str, str] = {} def take(self, name: str, kind: str): @@ -345,9 +410,13 @@ def fv5(self, out: str, names: list[str], o_expect: int, i_expect: int): levels = np.concatenate(lv, axis=0) if levels.shape[0] != o_expect: die(f"{out}: rows {levels.shape[0]} != {o_expect}") - self.add(out + ".fv5", pack_pair16(levels).view(np.int8)) + if self.layout == "b5x3t7": + self.add(out + ".fv5t7", pack_b5x3t7(levels).view(np.int8)) + else: + self.add(out + ".fv5", pack_pair16(levels).view(np.int8)) self.add(out + ".lohi", np.concatenate(lohi, axis=0).astype(np.float32)) self.fv5_src[out] = names + self.fv5_k[out] = i_expect def intq(self, out: str, name: str, shape2d: tuple[int, int] | None = None): r = self.take(name, "int") @@ -365,8 +434,8 @@ def f32(self, out: str, name: str, shape=None): self.add(out, v) -def build(recs: dict) -> Builder: - b = Builder(recs) +def build(recs: dict, layout: str) -> Builder: + b = Builder(recs, layout) d, ff, L = HP["d_model"], HP["ff_dim"], HP["n_layers"] C = HP["sub_channels"] # subsampling (dw_striding): conv0 (1->C, 3x3 s2) relu, [dw 3x3 s2, pw 1x1, relu] x2, linear @@ -450,7 +519,7 @@ def write_gguf(out: Path, b: Builder, labels: list[str], notice: str, pins: dict w.add_string(f"{NS}.source.{k}", v) w.add_string(f"{NS}.notice", notice) w.add_string(f"{NS}.converter.blob_sha1", git_blob_sha(Path(__file__))) - w.add_string(f"{NS}.weights.format", "fv5-pair16-v1") + w.add_string(f"{NS}.weights.format", LAYOUTS[b.layout][1]) for key in ("d_model", "n_layers", "n_heads", "ff_dim", "conv_kernel", "sub_channels", "sub_factor", "n_mels", "sample_rate", "n_fft", "win_length", "hop_length", "pos_max_len", "pred_hidden", "pred_layers", "joint_hidden", "vocab", "blank_id", "max_symbols"): @@ -479,8 +548,11 @@ def verify(out: Path, b: Builder, recs: dict): t = {x.name: np.array(x.data) for x in r.tensors} nfv = nint = 0 for name, srcs in b.fv5_src.items(): - lv = unpack_pair16(t[name + ".fv5"].view(np.uint8)) lohi = t[name + ".lohi"] + if b.layout == "b5x3t7": + lv = unpack_b5x3t7(t[name + ".fv5t7"].view(np.uint8), b.fv5_k[name]) + else: + lv = unpack_pair16(t[name + ".fv5"].view(np.uint8)) got = fv5_expand(lv, lohi) row = 0 for s in srcs: @@ -520,6 +592,8 @@ def main(): ap.add_argument("--notice", type=Path, help="NOTICE from the model repo (default: next to --archive)") ap.add_argument("--output", type=Path, required=True) ap.add_argument("--verify", action="store_true") + ap.add_argument("--layout", choices=sorted(LAYOUTS), default="b5x3t7", + help="five-value code layout (default b5x3t7, 2.35-2.41 bits per weight; pair16 is 4 bits)") a = ap.parse_args() pins = json.loads(PINS_PATH.read_text()) @@ -564,7 +638,7 @@ def main(): recs, index = read_container(container) if len(index) != 723: die(f"container has {len(index)} records, expected 723") - b = build(recs) + b = build(recs, a.layout) src = {"container_sha256": pins["archive_members"]["model.fermion"]["sha256"], "archive_sha256": pins["files"]["phonon-2.bps.tar.zst"]["sha256"]} write_gguf(a.output, b, labels, notice.decode("utf-8"), pins, src) diff --git a/tools/parakeet/pins.json b/tools/parakeet/pins.json index f10f433c1..4f2ab1b88 100644 --- a/tools/parakeet/pins.json +++ b/tools/parakeet/pins.json @@ -63,9 +63,15 @@ }, "gguf": { "file": "phonon-2-fv5.gguf", - "size": 330632256, - "sha256": "dc552eb786486167dde3ac80bbe49a6088c8ef39c03c257d5a090fc57b2c47b3", - "format": "fv5-pair16-v1", - "note": "convert_phonon2_to_gguf.py --archive output; byte-reproducible. Changes whenever the converter changes, because its git blob sha is in the metadata." + "size": 208932352, + "sha256": "c0f6ad42f0a5b481904416077dde0a1a78e5b8642295fdf2356ace822588c3df", + "format": "fv5-b5x3t7-v1", + "note": "convert_phonon2_to_gguf.py --archive output (default --layout b5x3t7); byte-reproducible. Changes whenever the converter changes, because its git blob sha is in the metadata.", + "pair16": { + "size": 330632256, + "sha256": "cc0ee571b01c549cef3a30bdd2626624c5e9059e1aaa7d02be56adab12e57fca", + "format": "fv5-pair16-v1", + "note": "--layout pair16 output of the same converter. Same weights and bit-identical results; 4 bits per weight." + } } } diff --git a/tools/parakeet/src/parakeet.cpp b/tools/parakeet/src/parakeet.cpp index d63a257c0..350d6cc65 100644 --- a/tools/parakeet/src/parakeet.cpp +++ b/tools/parakeet/src/parakeet.cpp @@ -209,16 +209,17 @@ struct loader { } void fv5(wmat & w, const std::string & name, int O, int K) { - const int8_t * codes = i8(name + ".fv5", { K, O / 2 }); - const float * lohi = f32(name + ".lohi", { 2, O }); - if (!err.empty()) { - return; - } if (O % 8 || K % 16) { err = name + ": rows must be a multiple of 8 and width a multiple of 16"; return; } - wmat_from_fv5(w, (const uint8_t *) codes, lohi, O, K); + const int8_t * codes = m->fv5t7 ? i8(name + ".fv5t7", { (int64_t) fv5t7_block_bytes(K), O / 8 }) + : i8(name + ".fv5", { K, O / 2 }); + const float * lohi = f32(name + ".lohi", { 2, O }); + if (!err.empty()) { + return; + } + wmat_from_fv5(w, (const uint8_t *) codes, lohi, O, K, m->fv5t7); m->fv5_grid_err = std::max(m->fv5_grid_err, w.grid_err); } void intw(wmat & w, const std::string & name, int O, int K, const float * bias, const int * perm = nullptr) { @@ -311,8 +312,11 @@ extern "C" pk_model * pk_model_load(const char * path, pk_params params, char * return nullptr; } const char * fmt = str_kv("parakeet.weights.format"); - if (!fmt || std::strcmp(fmt, "fv5-pair16-v1") != 0) { - set_err(err, err_len, "parakeet.weights.format is not fv5-pair16-v1 (re-run convert_phonon2_to_gguf.py)"); + if (fmt && std::strcmp(fmt, "fv5-b5x3t7-v1") == 0) { + m->fv5t7 = true; + } else if (!fmt || std::strcmp(fmt, "fv5-pair16-v1") != 0) { + set_err(err, err_len, + "parakeet.weights.format is neither fv5-b5x3t7-v1 nor fv5-pair16-v1 (re-run convert_phonon2_to_gguf.py)"); return nullptr; } if (const char * notice = str_kv("parakeet.notice")) { diff --git a/tools/parakeet/src/pk-kernels-avx2.cpp b/tools/parakeet/src/pk-kernels-avx2.cpp index c92e94b7d..2ad5f3c49 100644 --- a/tools/parakeet/src/pk-kernels-avx2.cpp +++ b/tools/parakeet/src/pk-kernels-avx2.cpp @@ -10,6 +10,8 @@ #if PK_HAVE_X86_EXT #include +#include + namespace pk { namespace { @@ -32,6 +34,59 @@ void decode_fv5_avx2(const uint8_t * codes, size_t pair_stride, const int8_t * l } } +// `fv5-b5x3t7-v1` with the NEON tier's arithmetic: q = ((x + 1) * 51) >> 8 +// on 16-bit lanes, digits by subtraction, then the pair16 PSHUFB lookup. +inline __m128i fv5t7_div5(__m128i x) { + const __m128i z = _mm_setzero_si128(), k51 = _mm_set1_epi16(51), one = _mm_set1_epi16(1); + const __m128i lo = _mm_srli_epi16(_mm_mullo_epi16(_mm_add_epi16(_mm_unpacklo_epi8(x, z), one), k51), 8); + const __m128i hi = _mm_srli_epi16(_mm_mullo_epi16(_mm_add_epi16(_mm_unpackhi_epi8(x, z), one), k51), 8); + return _mm_packus_epi16(lo, hi); +} + +inline __m128i fv5t7_times5(__m128i q) { + const __m128i q2 = _mm_add_epi8(q, q); + return _mm_add_epi8(_mm_add_epi8(q2, q2), q); +} + +void decode_fv5t7_avx2(const uint8_t * codes, const int8_t * luts, int K, int8_t * panel) { + const __m128i lut[4] = { _mm_loadu_si128((const __m128i *) luts), _mm_loadu_si128((const __m128i *) (luts + 16)), + _mm_loadu_si128((const __m128i *) (luts + 32)), _mm_loadu_si128((const __m128i *) (luts + 48)) }; + const __m128i row_hi = _mm_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8); + const __m128i m7 = _mm_set1_epi8(0x7F); + const size_t pbytes = (size_t) K * 8; + const int G = (fv5t7_codes(K) + 7) / 8; + alignas(16) int8_t stage[384]; + for (int g = 0; g < G; ++g) { + const uint8_t * s = codes + (size_t) g * 112; + __m128i x[8]; + __m128i c7 = _mm_setzero_si128(); + for (int m = 0; m < 7; ++m) { + const __m128i v = _mm_loadu_si128((const __m128i *) (s + m * 16)); + x[m] = _mm_and_si128(v, m7); + // bit 7 of every byte -> bit m (a 16-bit shift only brings in + // bits above m from the neighbouring byte, and those are masked) + c7 = _mm_or_si128(c7, _mm_and_si128(_mm_srli_epi16(v, 7 - m), _mm_set1_epi8((char) (1 << m)))); + } + x[7] = c7; + const size_t off = (size_t) g * 384; + const bool full = off + 384 <= pbytes; + int8_t * dst = full ? panel + off : stage; + for (int m = 0; m < 8; ++m) { + const __m128i q1 = fv5t7_div5(x[m]); + const __m128i q2 = fv5t7_div5(q1); + const __m128i d0 = _mm_sub_epi8(x[m], fv5t7_times5(q1)); + const __m128i d1 = _mm_sub_epi8(q1, fv5t7_times5(q2)); + const int v0 = 3 * m; // panel vector index within the group; row pair = v % 4 + _mm_storeu_si128((__m128i *) (dst + (size_t) v0 * 16), _mm_shuffle_epi8(lut[v0 % 4], _mm_or_si128(d0, row_hi))); + _mm_storeu_si128((__m128i *) (dst + (size_t) (v0 + 1) * 16), _mm_shuffle_epi8(lut[(v0 + 1) % 4], _mm_or_si128(d1, row_hi))); + _mm_storeu_si128((__m128i *) (dst + (size_t) (v0 + 2) * 16), _mm_shuffle_epi8(lut[(v0 + 2) % 4], _mm_or_si128(q2, row_hi))); + } + if (!full) { + std::memcpy(panel + off, stage, pbytes - off); + } + } +} + inline __m256i load_s8_as_s16(const int8_t * p) { return _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *) p)); } @@ -129,7 +184,7 @@ void matvec_avx2(const int8_t * w, int O, int K, const int8_t * x, int32_t * y) } // namespace const kernel_table & kernels_avx2() { - static const kernel_table t = { tier::avx2, "avx2", decode_fv5_avx2, tile_avx2, matvec_avx2 }; + static const kernel_table t = { tier::avx2, "avx2", decode_fv5_avx2, decode_fv5t7_avx2, tile_avx2, matvec_avx2 }; return t; } diff --git a/tools/parakeet/src/pk-kernels-dotprod.cpp b/tools/parakeet/src/pk-kernels-dotprod.cpp index ca631117d..cd3b6d29b 100644 --- a/tools/parakeet/src/pk-kernels-dotprod.cpp +++ b/tools/parakeet/src/pk-kernels-dotprod.cpp @@ -55,7 +55,7 @@ void tile_dotprod(const int8_t * panel, const int8_t * x, size_t x_stride, int K } // namespace const kernel_table & kernels_dotprod() { - static const kernel_table t = { tier::dotprod, "dotprod", decode_fv5_neon, tile_dotprod, matvec_dotprod }; + static const kernel_table t = { tier::dotprod, "dotprod", decode_fv5_neon, decode_fv5t7_neon, tile_dotprod, matvec_dotprod }; return t; } diff --git a/tools/parakeet/src/pk-kernels-i8mm.cpp b/tools/parakeet/src/pk-kernels-i8mm.cpp index 6ecff5dc8..f6a5a7942 100644 --- a/tools/parakeet/src/pk-kernels-i8mm.cpp +++ b/tools/parakeet/src/pk-kernels-i8mm.cpp @@ -51,7 +51,7 @@ void tile_i8mm(const int8_t * panel, const int8_t * x, size_t x_stride, int K, i } // namespace const kernel_table & kernels_i8mm() { - static const kernel_table t = { tier::i8mm, "i8mm", decode_fv5_neon, tile_i8mm, matvec_dotprod }; + static const kernel_table t = { tier::i8mm, "i8mm", decode_fv5_neon, decode_fv5t7_neon, tile_i8mm, matvec_dotprod }; return t; } diff --git a/tools/parakeet/src/pk-kernels-neon.inc b/tools/parakeet/src/pk-kernels-neon.inc index 2ecbf4703..e5ec1a89e 100644 --- a/tools/parakeet/src/pk-kernels-neon.inc +++ b/tools/parakeet/src/pk-kernels-neon.inc @@ -2,6 +2,8 @@ // and i8mm translation units (each is compiled with its own -march). #include +#include + namespace pk { namespace { @@ -29,6 +31,72 @@ inline void decode_fv5_neon(const uint8_t * codes, size_t pair_stride, const int } } +// `fv5-b5x3t7-v1`. A code x < 125 splits into base-5 digits with +// q1 = x / 5 = ((x + 1) * 51) >> 8 (exact on 0..124: two UMULLs and a UZP2 +// of the high bytes), d0 = x - 5 q1, q2 = q1 / 5 from a 32-byte TBL (q1 < 25) +// and d1 = q1 - 5 q2. Each digit becomes a panel vector with the pair16 LUT +// lookup. A group of 8 code vectors makes 24 panel vectors, a multiple of 4, +// so every group starts at row pair 0. +inline uint8x16_t fv5t7_div5(uint8x16_t x) { + const uint8x16_t x1 = vaddq_u8(x, vdupq_n_u8(1)); + const uint16x8_t lo = vmull_u8(vget_low_u8(x1), vdup_n_u8(51)); + const uint16x8_t hi = vmull_high_u8(x1, vdupq_n_u8(51)); + return vuzp2q_u8(vreinterpretq_u8_u16(lo), vreinterpretq_u8_u16(hi)); +} + +template +inline void fv5t7_emit(const int8x16_t * lut, uint8x16_t row_hi, const uint8x16x2_t & div5_tab, uint8x16_t x, + int8_t * dst) { + const uint8x16_t five = vdupq_n_u8(5); + const uint8x16_t q1 = fv5t7_div5(x); + const uint8x16_t q2 = vqtbl2q_u8(div5_tab, q1); + const uint8x16_t d0 = vmlsq_u8(x, q1, five); + const uint8x16_t d1 = vmlsq_u8(q1, q2, five); + vst1q_s8(dst, vqtbl1q_s8(lut[P0 % 4], vorrq_u8(d0, row_hi))); + vst1q_s8(dst + 16, vqtbl1q_s8(lut[(P0 + 1) % 4], vorrq_u8(d1, row_hi))); + vst1q_s8(dst + 32, vqtbl1q_s8(lut[(P0 + 2) % 4], vorrq_u8(q2, row_hi))); +} + +inline void decode_fv5t7_neon(const uint8_t * codes, const int8_t * luts, int K, int8_t * panel) { + const int8x16_t lut[4] = { vld1q_s8(luts), vld1q_s8(luts + 16), vld1q_s8(luts + 32), vld1q_s8(luts + 48) }; + static const uint8_t hi8[16] = { 0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8 }; + const uint8x16_t row_hi = vld1q_u8(hi8); + const uint8x16_t m7 = vdupq_n_u8(0x7F); + static const uint8_t div5[32] = { 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, + 3, 3, 3, 3, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6 }; + const uint8x16x2_t div5_tab = { { vld1q_u8(div5), vld1q_u8(div5 + 16) } }; + const size_t pbytes = (size_t) K * 8; + const int G = (fv5t7_codes(K) + 7) / 8; + alignas(16) int8_t stage[384]; + for (int g = 0; g < G; ++g) { + const uint8_t * s = codes + (size_t) g * 112; + const uint8x16_t v0 = vld1q_u8(s), v1 = vld1q_u8(s + 16), v2 = vld1q_u8(s + 32), v3 = vld1q_u8(s + 48); + const uint8x16_t v4 = vld1q_u8(s + 64), v5 = vld1q_u8(s + 80), v6 = vld1q_u8(s + 96); + // code 7 = the top bits of v0..v6 (bit m from v_m) + uint8x16_t c7 = vsriq_n_u8(v6, v5, 1); + c7 = vsriq_n_u8(c7, v4, 2); + c7 = vsriq_n_u8(c7, v3, 3); + c7 = vsriq_n_u8(c7, v2, 4); + c7 = vsriq_n_u8(c7, v1, 5); + c7 = vsriq_n_u8(c7, v0, 6); + c7 = vshrq_n_u8(c7, 1); + const size_t off = (size_t) g * 384; + const bool full = off + 384 <= pbytes; // the last group can run past K/2 vectors + int8_t * dst = full ? panel + off : stage; + fv5t7_emit<0>(lut, row_hi, div5_tab, vandq_u8(v0, m7), dst); + fv5t7_emit<3>(lut, row_hi, div5_tab, vandq_u8(v1, m7), dst + 48); + fv5t7_emit<2>(lut, row_hi, div5_tab, vandq_u8(v2, m7), dst + 96); + fv5t7_emit<1>(lut, row_hi, div5_tab, vandq_u8(v3, m7), dst + 144); + fv5t7_emit<0>(lut, row_hi, div5_tab, vandq_u8(v4, m7), dst + 192); + fv5t7_emit<3>(lut, row_hi, div5_tab, vandq_u8(v5, m7), dst + 240); + fv5t7_emit<2>(lut, row_hi, div5_tab, vandq_u8(v6, m7), dst + 288); + fv5t7_emit<1>(lut, row_hi, div5_tab, c7, dst + 336); + if (!full) { + std::memcpy(panel + off, stage, pbytes - off); + } + } +} + inline void matvec_dotprod(const int8_t * w, int O, int K, const int8_t * x, int32_t * y) { int o = 0; for (; o + 4 <= O; o += 4) { diff --git a/tools/parakeet/src/pk-kernels.cpp b/tools/parakeet/src/pk-kernels.cpp index 3f83a8ff0..a8345ccf2 100644 --- a/tools/parakeet/src/pk-kernels.cpp +++ b/tools/parakeet/src/pk-kernels.cpp @@ -44,6 +44,29 @@ static void decode_fv5_scalar(const uint8_t * codes, size_t pair_stride, const i } } +static void decode_fv5t7_scalar(const uint8_t * codes, const int8_t * luts, int K, int8_t * panel) { + const int V = K / 2, T = fv5t7_codes(K); + for (int t = 0; t < T; ++t) { + const uint8_t * grp = codes + (size_t) (t / 8) * 112; + const int m = t % 8; + for (int j = 0; j < 16; ++j) { + int x; + if (m < 7) { + x = grp[m * 16 + j] & 0x7F; + } else { + x = 0; + for (int i = 0; i < 7; ++i) x |= (grp[i * 16 + j] >> 7) << i; + } + for (int d = 0; d < 3; ++d, x /= 5) { + const int v = 3 * t + d; + if (v < V) { + panel[(size_t) v * 16 + j] = luts[(v % 4) * 16 + x % 5 + (j >= 8 ? 8 : 0)]; + } + } + } + } +} + static void tile_scalar(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out) { int32_t acc[8][8]; std::memset(acc, 0, sizeof(acc)); @@ -85,7 +108,7 @@ static void matvec_scalar(const int8_t * w, int O, int K, const int8_t * x, int3 } const kernel_table & kernels_scalar() { - static const kernel_table t = { tier::scalar, "scalar", decode_fv5_scalar, tile_scalar, matvec_scalar }; + static const kernel_table t = { tier::scalar, "scalar", decode_fv5_scalar, decode_fv5t7_scalar, tile_scalar, matvec_scalar }; return t; } diff --git a/tools/parakeet/src/pk-kernels.h b/tools/parakeet/src/pk-kernels.h index 43cc5fa3e..4b69af4ad 100644 --- a/tools/parakeet/src/pk-kernels.h +++ b/tools/parakeet/src/pk-kernels.h @@ -20,6 +20,17 @@ // table per row pair maps the nibble to the int8 level: // 0..4 -> row 2p {-H,-L,0,L,H}, 8..12 -> row 2p+1. // +// fv5 t7 `fv5-b5x3t7-v1` (2.35-2.41 bits per weight): per block of +// 8 rows, the K/2 panel vectors above are taken three at a +// time. Code vector t holds, in lane j, the base-5 number +// l(3t) + 5 l(3t+1) + 25 l(3t+2), where l(v) is the level +// 0..4 of lane j of panel vector v (missing vectors past K/2 +// count as level 2). Codes are < 128, so code vectors go in +// groups of 8 into 7 stored vectors: byte m (0..6) of lane j +// is code m | (bit m of code 7) << 7. A short last group is +// padded with code 62 (levels 2, 2, 2). Decoding uses the same +// 16-byte LUTs as pair16 (level + 8 for row 2p+1). +// // Every tier computes exact int32 dot products over the same int8 operands, // so all tiers are bit-identical (tests/test_parakeet_kernels.cpp). #pragma once @@ -37,6 +48,9 @@ struct kernel_table { // decode the fv5 codes of one 8-row block (4 row pairs starting at // `codes`, `pair_stride` bytes apart) into a weight panel. void (*decode_fv5)(const uint8_t * codes, size_t pair_stride, const int8_t * luts, int K, int8_t * panel); + // same panel from one 8-row block in `fv5-b5x3t7-v1` (fv5t7_block_bytes(K) + // bytes at `codes`). + void (*decode_fv5t7)(const uint8_t * codes, const int8_t * luts, int K, int8_t * panel); // int32 tile: 8 rows x 2*np tokens (np = 1..4 token pairs). `panel` is // the 8-row weight panel, `x` the first token pair, `x_stride` the bytes // between token pairs. out[t * 8 + r] (token-major, so the epilogue @@ -47,6 +61,10 @@ struct kernel_table { void (*matvec)(const int8_t * w, int O, int K, const int8_t * x, int32_t * y); }; +// `fv5-b5x3t7-v1` sizes: code vectors and stored bytes per 8-row block. +inline int fv5t7_codes(int K) { return (K / 2 + 2) / 3; } +inline size_t fv5t7_block_bytes(int K) { return (size_t) ((fv5t7_codes(K) + 7) / 8) * 7 * 16; } + // Best tier this CPU supports, or `want` when it is available. const kernel_table & kernels_for(int want /*0 = best*/); bool cpu_has_dotprod(); diff --git a/tools/parakeet/src/pk-model.h b/tools/parakeet/src/pk-model.h index 50444aa27..7b4fca7e2 100644 --- a/tools/parakeet/src/pk-model.h +++ b/tools/parakeet/src/pk-model.h @@ -94,6 +94,7 @@ struct pk_model { pk::workspace ws; double fv5_grid_err = 0.0; + bool fv5t7 = false; // weights in `fv5-b5x3t7-v1` (else `fv5-pair16-v1`) // relative-position keys W_pos . e(p) per layer for p in [-rc_P, rc_P], // row c <-> p = rc_P - c, fp16. They depend only on p, so they are diff --git a/tools/parakeet/src/pk-ops.cpp b/tools/parakeet/src/pk-ops.cpp index 748aa5802..339f3e4a3 100644 --- a/tools/parakeet/src/pk-ops.cpp +++ b/tools/parakeet/src/pk-ops.cpp @@ -263,10 +263,11 @@ void quantize_acts(pool & P, const float * X, int N, int K, size_t ldx, qact & o // ------------------------------------------------------------ weights -void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K) { +void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K, bool t7) { w.O = O; w.K = K; w.fv5 = codes; + w.fv5t7 = t7; w.lut.assign((size_t) (O / 2) * 16, 0); w.ws.assign(O, 0.0f); w.grid_err = 0.0; @@ -306,6 +307,15 @@ void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, i } } +void decode_fv5_block(const kernel_table & kt, const wmat & W, int rb, int8_t * panel) { + const int K = W.K; + if (W.fv5t7) { + kt.decode_fv5t7(W.fv5 + (size_t) rb * fv5t7_block_bytes(K), W.lut.data() + (size_t) rb * 64, K, panel); + } else { + kt.decode_fv5(W.fv5 + (size_t) rb * 4 * K, (size_t) K, W.lut.data() + (size_t) rb * 64, K, panel); + } +} + void wmat_from_int8(wmat & w, const int8_t * q, const float * scale, int O, int K, const int * perm) { w.O = O; w.K = K; @@ -343,7 +353,7 @@ static void gemm_f32_debug(pool & P, const kernel_table & kt, const wmat & W, co for (int rb = b; rb < e; ++rb) { const int8_t * panel; if (W.fv5) { - kt.decode_fv5(W.fv5 + (size_t) rb * 4 * K, (size_t) K, W.lut.data() + (size_t) rb * 64, K, scratch.data()); + decode_fv5_block(kt, W, rb, scratch.data()); panel = scratch.data(); } else { panel = W.panels.data() + (size_t) rb * 8 * K; @@ -403,7 +413,7 @@ void gemm_q8(pool & P, const kernel_table & kt, const wmat & W, const qact & X, for (int rb = rb0; rb < rb1; ++rb) { if (W.fv5) { int8_t * dst = scratch.data() + (size_t) (rb - rb0) * 8 * K; - kt.decode_fv5(W.fv5 + (size_t) rb * 4 * K, (size_t) K, W.lut.data() + (size_t) rb * 64, K, dst); + decode_fv5_block(kt, W, rb, dst); panels[rb - rb0] = dst; } else { panels[rb - rb0] = W.panels.data() + (size_t) rb * 8 * K; diff --git a/tools/parakeet/src/pk-ops.h b/tools/parakeet/src/pk-ops.h index f08b381b4..406a4db10 100644 --- a/tools/parakeet/src/pk-ops.h +++ b/tools/parakeet/src/pk-ops.h @@ -33,8 +33,10 @@ void quantize_acts(pool & P, const float * X, int N, int K, size_t ldx, qact & o struct wmat { int O = 0, K = 0; - // five-value: codes in `fv5-pair16-v1`, one 16-byte LUT per row pair + // five-value: codes in `fv5-pair16-v1` or, with fv5t7, `fv5-b5x3t7-v1`; + // one 16-byte LUT per row pair either way const uint8_t * fv5 = nullptr; + bool fv5t7 = false; std::vector lut; // int tables: weight panels prepared at load ([O/8][K/8][4][16]) std::vector panels; @@ -44,7 +46,9 @@ struct wmat { double grid_err = 0.0; }; -void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K); +void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K, bool t7 = false); +// decode the 8-row block `rb` of a five-value matrix into a weight panel +void decode_fv5_block(const kernel_table & kt, const wmat & W, int rb, int8_t * panel); // row-major int8 [O][K] with per-row scales; `perm` (optional) permutes the // input columns: new column j reads old column perm[j]. void wmat_from_int8(wmat & w, const int8_t * q, const float * scale, int O, int K, const int * perm = nullptr); diff --git a/tools/parakeet/tests/test_parakeet_kernels.cpp b/tools/parakeet/tests/test_parakeet_kernels.cpp index 8476b463a..67c8af8a3 100644 --- a/tools/parakeet/tests/test_parakeet_kernels.cpp +++ b/tools/parakeet/tests/test_parakeet_kernels.cpp @@ -1,7 +1,8 @@ // tools/parakeet/tests/test_parakeet_kernels.cpp - every kernel tier must give // bit-identical int32 results to the scalar reference (decode, tile, matvec, // and the whole gemm_q8 driver), and gemm_q8 must match a double-precision -// GEMM over the same int8 operands. +// GEMM over the same int8 operands. The `fv5-b5x3t7-v1` layout must decode to +// exactly the `fv5-pair16-v1` panels, on every tier. #include "pk-kernels.h" #include "pk-ops.h" #include "pk-pool.h" @@ -44,6 +45,36 @@ static void make_fv5(std::mt19937 & rng, int O, int K, std::vector & co } } +// Repack pair16 codes as `fv5-b5x3t7-v1` (pk-kernels.h), independently of +// the decoders: level of lane j in panel vector v of block rb, from pair16. +static std::vector pair16_to_t7(const std::vector & codes, int O, int K) { + auto level = [&](int rb, int v, int j) -> int { + if (v >= K / 2) return 2; + const int c = v / 4, p = rb * 4 + v % 4; + const int g = c / 2, half = c % 2; + const uint8_t b = codes[(size_t) p * K + 16 * g + j]; + return (half ? (b >> 4) : (b & 15)) - (j >= 8 ? 8 : 0); + }; + const int T = fv5t7_codes(K), G = (T + 7) / 8; + const size_t bb = fv5t7_block_bytes(K); + std::vector out((size_t) O / 8 * bb, 0); + for (int rb = 0; rb < O / 8; ++rb) { + for (int g = 0; g < G; ++g) { + for (int j = 0; j < 16; ++j) { + int code[8]; + for (int m = 0; m < 8; ++m) { + const int t = 8 * g + m; + code[m] = t < T ? level(rb, 3 * t, j) + 5 * level(rb, 3 * t + 1, j) + 25 * level(rb, 3 * t + 2, j) : 62; + } + for (int m = 0; m < 7; ++m) { + out[rb * bb + ((size_t) g * 7 + m) * 16 + j] = (uint8_t) (code[m] | (((code[7] >> m) & 1) << 7)); + } + } + } + } + return out; +} + int main() { std::mt19937 rng(1234); const kernel_table & ref = kernels_scalar(); @@ -59,13 +90,30 @@ int main() { for (auto * t : tiers) std::printf(", %s", t->name); std::printf("\n"); - for (int K : { 256, 1024, 4096 }) { + for (int K : { 256, 1024, 4096, 48, 80 }) { const int O = 64, N = 21; std::vector codes; std::vector lohi; make_fv5(rng, O, K, codes, lohi); wmat W; wmat_from_fv5(W, codes.data(), lohi.data(), O, K); + // fv5-b5x3t7-v1: every block decodes to the pair16 panel on every tier + const std::vector t7 = pair16_to_t7(codes, O, K); + wmat W7; + wmat_from_fv5(W7, t7.data(), lohi.data(), O, K, true); + { + std::vector p0((size_t) 8 * K), p1((size_t) 8 * K); + for (int rb = 0; rb < O / 8; ++rb) { + decode_fv5_block(ref, W, rb, p0.data()); + decode_fv5_block(ref, W7, rb, p1.data()); + CHECK(std::memcmp(p0.data(), p1.data(), p0.size()) == 0, "scalar t7 decode differs (K=%d rb=%d)", K, rb); + for (auto * t : tiers) { + std::memset(p1.data(), 0x5a, p1.size()); + decode_fv5_block(*t, W7, rb, p1.data()); + CHECK(std::memcmp(p0.data(), p1.data(), p0.size()) == 0, "%s t7 decode differs (K=%d rb=%d)", t->name, K, rb); + } + } + } CHECK(W.grid_err < 0.01, "grid error %g", W.grid_err); std::vector X((size_t) N * K); std::normal_distribution nd(0.0f, 1.0f); @@ -111,6 +159,14 @@ int main() { gemm_q8(P3, *t, W, qa, y1.data(), O, epi::store); CHECK(std::memcmp(y0.data(), y1.data(), y0.size() * 4) == 0, "%s gemm differs (K=%d)", t->name, K); } + for (const kernel_table * t : { &ref }) { + gemm_q8(P3, *t, W7, qa, y1.data(), O, epi::store); + CHECK(std::memcmp(y0.data(), y1.data(), y0.size() * 4) == 0, "%s t7 gemm differs (K=%d)", t->name, K); + } + for (auto * t : tiers) { + gemm_q8(P3, *t, W7, qa, y1.data(), O, epi::store); + CHECK(std::memcmp(y0.data(), y1.data(), y0.size() * 4) == 0, "%s t7 gemm differs (K=%d)", t->name, K); + } // and the driver result is the dequantised product double maxrel = 0; for (int t = 0; t < N; ++t) { diff --git a/tools/parakeet/tools/bench-fv5-packing.cpp b/tools/parakeet/tools/bench-fv5-packing.cpp new file mode 100644 index 000000000..1e3bcc271 --- /dev/null +++ b/tools/parakeet/tools/bench-fv5-packing.cpp @@ -0,0 +1,175 @@ +// tools/parakeet/tools/bench-fv5-packing.cpp - cost of the `fv5-b5x3t7-v1` +// code layout (2.35-2.41 bits per weight) against `fv5-pair16-v1` (4 bits). +// +// Both layouts decode into the same int8 weight panel (pk-kernels.h), so only +// the decode step differs. The bench packs random levels with the real +// Phonon-2 level distribution in both layouts, checks that the selected tier +// decodes them to the same panel, then times decode alone and decode + the +// tier's int8 tile over N tokens, on one thread, for the encoder's linear +// shapes. It is what the gemm_q8 driver does per 8-row block. +// +// usage: bench-fv5-packing [reps] (PK_KERNEL selects the tier) +#include "pk-kernels.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using clk = std::chrono::steady_clock; + +// level of lane j in panel vector v for an 8-row block of `lv` ([8][K], 0..4) +inline uint8_t panel_level(const uint8_t * lv, int K, int v, int j) { + const int c = v / 4, p = v % 4; + const int row = 2 * p + (j >= 8 ? 1 : 0); + return lv[(size_t) row * K + 8 * c + (j & 7)]; +} + +// pair16 codes for one block: [4 pairs][K bytes], pair stride K +void pack_pair16(const uint8_t * lv, int K, uint8_t * out) { + for (int p = 0; p < 4; ++p) { + for (int g = 0; g < K / 16; ++g) { + for (int j = 0; j < 16; ++j) { + const int r = 2 * p + (j >= 8 ? 1 : 0); + const int k = 16 * g + (j & 7); + const uint8_t a = (uint8_t) (lv[(size_t) r * K + k] + (j >= 8 ? 8 : 0)); + const uint8_t b = (uint8_t) (lv[(size_t) r * K + k + 8] + (j >= 8 ? 8 : 0)); + out[(size_t) p * K + g * 16 + j] = (uint8_t) (a | (b << 4)); + } + } + } +} + +int b5x3_vectors(int K) { return pk::fv5t7_codes(K); } +size_t b5x3_bytes(int K) { return (size_t) b5x3_vectors(K) * 16; } +int t7_groups(int K) { return (b5x3_vectors(K) + 7) / 8; } +size_t t7_bytes(int K) { return pk::fv5t7_block_bytes(K); } + +void pack_b5x3(const uint8_t * lv, int K, uint8_t * out) { + const int V = K / 2, T = b5x3_vectors(K); + for (int t = 0; t < T; ++t) { + for (int j = 0; j < 16; ++j) { + int code = 0, mul = 1; + for (int d = 0; d < 3; ++d) { + const int v = 3 * t + d; + code += (v < V ? panel_level(lv, K, v, j) : 2) * mul; + mul *= 5; + } + out[(size_t) t * 16 + j] = (uint8_t) code; + } + } +} + +void pack_t7(const uint8_t * b5, int K, uint8_t * out) { + const int T = b5x3_vectors(K), Gn = t7_groups(K); + for (int g = 0; g < Gn; ++g) { + for (int j = 0; j < 16; ++j) { + uint8_t c[8]; + for (int m = 0; m < 8; ++m) { + const int t = 8 * g + m; + c[m] = t < T ? b5[(size_t) t * 16 + j] : 62; // 62 = levels {2,2,2} + } + for (int m = 0; m < 7; ++m) { + out[((size_t) g * 7 + m) * 16 + j] = (uint8_t) (c[m] | (((c[7] >> m) & 1) << 7)); + } + } + } +} + +double now_s() { return std::chrono::duration(clk::now().time_since_epoch()).count(); } + +} // namespace + +int main(int argc, char ** argv) { + const int reps = argc > 1 ? std::atoi(argv[1]) : 3; + const pk::kernel_table & kt = pk::kernels_for(0); + std::printf("tile tier: %s\n", kt.name); + for (int x = 0; x < 125; ++x) { + if (((x + 1) * 51) >> 8 != x / 5) { std::printf("div5 identity fails at %d\n", x); return 1; } + } + // encoder linears per conformer block (O x K): ff1 up/down, qkv, out, pos, + // conv pw1/pw2, ff2 up/down + struct shape { int O, K, per_block; }; + const shape shapes[] = { { 4096, 1024, 2 }, { 1024, 4096, 2 }, { 3072, 1024, 1 }, { 1024, 1024, 3 }, { 2048, 1024, 1 } }; + const int Ns[] = { 8, 16, 32, 64, 125, 375 }; + // level distribution of the real Phonon-2 encoder (container statistics) + std::mt19937 rng(1234); + std::discrete_distribution dist({ 0.0909, 0.1526, 0.5133, 0.1525, 0.0908 }); + + std::printf("%-6s %-5s %-6s %10s %10s %10s | %s\n", "O", "K", "layout", "bytes/blk", "dec ns/blk", "dec GW/s", + "decode+tile ms per GEMM at N ="); + for (const shape & sh : shapes) { + const int O = sh.O, K = sh.K, nb = O / 8; + std::vector lv((size_t) O * K); + for (auto & v : lv) v = (uint8_t) dist(rng); + std::vector luts((size_t) nb * 64); + for (int i = 0; i < nb * 4; ++i) { + const int8_t H = (int8_t) (64 + rng() % 64), L = (int8_t) (1 + rng() % 40); + const int8_t row[5] = { (int8_t) -H, (int8_t) -L, 0, L, H }; + for (int r = 0; r < 2; ++r) for (int l = 0; l < 5; ++l) luts[(size_t) i * 16 + r * 8 + l] = row[l]; + } + // packed codes per layout + std::vector p16((size_t) nb * 4 * K), b5((size_t) nb * b5x3_bytes(K)), t7((size_t) nb * t7_bytes(K)); + for (int rb = 0; rb < nb; ++rb) { + pack_pair16(lv.data() + (size_t) rb * 8 * K, K, p16.data() + (size_t) rb * 4 * K); + pack_b5x3(lv.data() + (size_t) rb * 8 * K, K, b5.data() + (size_t) rb * b5x3_bytes(K)); + pack_t7(b5.data() + (size_t) rb * b5x3_bytes(K), K, t7.data() + (size_t) rb * t7_bytes(K)); + } + // correctness: both layouts give the same panel + std::vector ref((size_t) 8 * K), got((size_t) 8 * K); + for (int rb = 0; rb < nb; ++rb) { + kt.decode_fv5(p16.data() + (size_t) rb * 4 * K, (size_t) K, luts.data() + (size_t) rb * 64, K, ref.data()); + kt.decode_fv5t7(t7.data() + (size_t) rb * t7_bytes(K), luts.data() + (size_t) rb * 64, K, got.data()); + if (std::memcmp(ref.data(), got.data(), ref.size())) { std::printf("fv5t7 panel mismatch\n"); return 1; } + } + // activations: N tokens, pair layout of pk-kernels.h + const int Nmax = 375, npmax = (Nmax + 1) / 2; + std::vector x((size_t) npmax * K * 2); + for (auto & v : x) v = (int8_t) ((int) (rng() % 255) - 127); + const size_t ps = (size_t) K * 2; + + struct layout { const char * name; size_t bytes; int id; }; + const layout ls[] = { { "pair16", (size_t) 4 * K, 0 }, { "b5x3t7", t7_bytes(K), 1 } }; + auto decode = [&](int id, int rb, int8_t * dst) { + if (id == 0) { + kt.decode_fv5(p16.data() + (size_t) rb * 4 * K, (size_t) K, luts.data() + (size_t) rb * 64, K, dst); + } else { + kt.decode_fv5t7(t7.data() + (size_t) rb * t7_bytes(K), luts.data() + (size_t) rb * 64, K, dst); + } + }; + std::vector panel((size_t) 8 * K); + alignas(16) int32_t acc[64]; + for (const layout & L : ls) { + double best_dec = 1e30; + for (int r = 0; r < reps * 3; ++r) { + const double t0 = now_s(); + for (int rb = 0; rb < nb; ++rb) decode(L.id, rb, panel.data()); + best_dec = std::min(best_dec, now_s() - t0); + } + std::printf("%-6d %-5d %-6s %10zu %10.0f %10.2f |", O, K, L.name, L.bytes, 1e9 * best_dec / nb, + (double) O * K / best_dec / 1e9); + for (int N : Ns) { + const int np = (N + 1) / 2, ntt = (np + 3) / 4; + double best = 1e30; + for (int r = 0; r < reps; ++r) { + const double t0 = now_s(); + for (int rb = 0; rb < nb; ++rb) { + decode(L.id, rb, panel.data()); + for (int tt = 0; tt < ntt; ++tt) { + kt.tile(panel.data(), x.data() + (size_t) tt * 4 * ps, ps, K, std::min(4, np - tt * 4), acc); + } + } + best = std::min(best, now_s() - t0); + } + std::printf(" %d:%.3f", N, 1e3 * best); + } + std::printf("\n"); + } + } + return 0; +}