vadik82/hayai-ocr-v2-onnx

Hayai OCR v2.1 — ONNX mirror (+ KV-cache graphs)

Mirror of the upstream JustANormalTinkerer/hayai-ocr-v2-onnx (weights, base ONNX graphs, export/inference scripts) with one addition: KV-cache decode graphs (hayai_decoder_kv_prefill / hayai_decoder_kv_step), which the upstream does not ship. The mirror is self-sufficient: the graphs can be rebuilt from it even if the upstream disappears.

What's added, and what it buys

The upstream decoder is a single full-pass graph: to generate every next token it re-runs attention over the whole sequence — all vision tokens plus everything generated so far. Total decode cost grows quadratically with sequence length, and every step gets slower than the previous one.

The added pair of graphs replaces that with a classic two-phase KV-cache decode:

• —hayai_decoder_kv_prefill runs once per crop over the vision tokens + BOS and emits the key/value cache (present_k / present_v);
• —hayai_decoder_kv_step generates each subsequent token from a single position, attending to the cached keys/values (past_k / past_v → updated present_k / present_v).

The result: decode cost per token becomes constant instead of growing with the generated length — O(1) per token, O(N) total instead of O(N²). On typical crop OCR (dozens of tokens of output) this cuts decoder compute several-fold, with the gap widening on longer generations. The cache itself is small ([layers, batch, positions, 2 kv-heads, 64]), so the trade is negligible memory against a large latency win.

Numerically the split is transparent: the weights are identical to the base export, prefill inputs/outputs follow the hayai_decoder.onnx convention, and prefill outputs match the full forward (the export script has built-in parity checks).

Measured impact

End-to-end numbers from a production batch pipeline (RTX 3060, int8 dynamic quant graphs, greedy decode, a 5-page wave of ~46 crops):

That is a ×2.4–3.0 end-to-end speedup of the OCR stage from the KV split alone, at identical output quality (same weights, greedy decode in all rows). The exact ratio grows with the length of the generated text, since the full-pass decoder re-reads the whole sequence on every token while the KV step always costs the same.

All graphs come in three precisions — fp32, _fp16, _dynamic_quant (int8 dynamic quant) — as self-contained single files; switch precision by suffix. To re-export, drop export_onnx_kv.py into the root of this folder (next to modeling_hayai.py) and run:

pip install torch transformers safetensors onnx onnxruntime
python export_onnx_kv.py                 # fp32 + int8 dynamic quant
python export_onnx_kv.py --quant none    # fp32 only

Note on fp32

The upstream fp32 graphs store their weights in external hayai_*.onnx.data files — a per-file download of the .onnx alone will not pick them up. The fp16 and int8 dynamic quant graphs (base and KV) are self-contained single files.

License and attribution

Weights and base graphs — JustANormalTinkerer/hayai-ocr-v2 (training details, datasets and benchmarks in the model card), Apache-2.0. The KV graphs are derived from the same weights, same license.