Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions tools/parakeet/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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()
33 changes: 30 additions & 3 deletions tools/parakeet/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
|---|---|---|
Expand Down Expand Up @@ -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.
Expand All @@ -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`.
Expand Down
96 changes: 85 additions & 11 deletions tools/parakeet/convert_phonon2_to_gguf.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,14 @@
fp16: raw half floats.

GGUF layout written here (all lossless with respect to the container):
* five-value linears -> `<name>.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 `<name>.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
`<name>.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): `<name>.fv5t7`, I8 [O/8, fv5t7 block bytes], the
`fv5-b5x3t7-v1` order (see pack_b5x3t7), 2.35-2.41 bits per weight;
pair16: `<name>.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 -> `<name>.q` (I8, the exact integers) + `<name>.scale` (F32).
* fp16 records -> F32 (exact widening).
* BatchNorm is folded into the depthwise conv (F32 weight + bias); this is
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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):
Expand All @@ -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")
Expand All @@ -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
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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())

Expand Down Expand Up @@ -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)
Expand Down
14 changes: 10 additions & 4 deletions tools/parakeet/pins.json
Original file line number Diff line number Diff line change
Expand Up @@ -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."
}
}
}
20 changes: 12 additions & 8 deletions tools/parakeet/src/parakeet.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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")) {
Expand Down
57 changes: 56 additions & 1 deletion tools/parakeet/src/pk-kernels-avx2.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
#if PK_HAVE_X86_EXT
#include <immintrin.h>

#include <cstring>

namespace pk {
namespace {

Expand All @@ -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));
}
Expand Down Expand Up @@ -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;
}

Expand Down
2 changes: 1 addition & 1 deletion tools/parakeet/src/pk-kernels-dotprod.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
Loading
Loading