Masterx/ark-asr-3b-onnx

ARK-ASR-3B — INT8 ONNX

An ONNX Runtime export of Audio8/ARK-ASR-3B, which upstream publishes as safetensors only. Weights, architecture, and behaviour are unchanged — this repo just repackages the checkpoint so it can run on ONNX Runtime without PyTorch.

The file layout deliberately mirrors Audio8/ark-asr-0.6b-int8-onnx, so any runtime that already drives the 0.6B export works here by pointing at this repo instead.

Files

File	Purpose
audio_encoder_whisper_int8.onnx (+ .data)	Whisper encoder + the adapter's LayerNorm. audios [1,128,mel_seq] → encoded_audio_features [1,seq,1280]
audio_encoder_adapter_int8.onnx (+ .data)	MLP adapter. merged_audio_features [1,n,5120] → audio_embeddings [1,n,2048]
llm_kv_cpu_fp32_int8.onnx (+ .data)	Decoder + LM head over a static KV cache. Serves both prefill and single-token decode. Default — the faster tier.
llm_kv_cpu_fp32_int4.onnx (+ .data)	Same graph, weight-only int4 (MatMulNBits). 1.64 GB vs 3.09 GB, but ~50% slower on CPU — pick it for disk, not speed
embedding_fp16.onnx (+ embedding_fp16.data)	Token-embedding lookup (Gather), fp16 (halves 1.24 GB to 622 MB; the table is only ever read as f16)
runtime_manifest.json	Token ids, merge factor, encoder hidden size, default graph names
llm_kv_cpu_fp32_int8.json	Decoder geometry (layers, heads, max_total_len)

Precision: the audio encoder, adapter, and decoder are dynamically quantized to int8 (per-channel MatMul weights). The embedding table stays fp32.

Shapes

• 36 decoder layers, hidden 2048, 16 attention heads, 2 KV heads (GQA), head dim 128
• Static KV cache cache_{key,value}_i of shape [1, 512, 2, 128] — sequence-major (max_total_len is 512, down from 2048 in the 0.6B export. With a static cache the decoder attends over ALL max_total_len positions on every token however few are valid, so this is a throughput knob: measured ~20% faster end-to-end than the same model at 1024)
• The decoder returns key_delta_i / value_delta_i of shape [1, seq, 2, 128] — only the positions written by this call, which the host copies back into its cache at cache_position

Inference

Audio is mono 16 kHz. Mel features use the standard Whisper front end (128 bins, hop 160, n_fft 400) over the unpadded clip — the encoder's mel length is dynamic.

1. Encode → [1, seq, 1280]
2. Truncate seq to a multiple of 4 and reshape 4 frames into one row → [1, seq/4, 5120]
3. Adapter → [1, n, 2048]
4. Build the prompt <|user|><|begin_of_audio|> + N×<|audio|> + <|end_of_audio|> + Please transcribe this audio. + <|assistant|>, where N = max(((samples // 160) + 1) // 2 // 4, 1)
5. Look up the token embeddings and overwrite the <|audio|> slots with the adapter output
6. Prefill, then decode greedily, masking every <…> added-vocab token except <|im_end|>

That last step matters: the model's raw first-token argmax on ordinary speech is <|user|>. If it is left unmasked the decode stops immediately and returns an empty transcript.

Audio8/ark-asr-0.6b-int8-onnx ships an infer_ark_audio_onnx.py that implements exactly this loop and works against these files unchanged.

Verification

Checked against the source PyTorch checkpoint on real audio:

Check	Result
fp32 encoder graph vs PyTorch	max abs diff 3.6e-04
int8 encoder graph vs PyTorch	max abs diff 2.3 (quantization)
Prefill logits argmax	identical
Greedy transcript, ONNX vs PyTorch	identical
Prompt vs the shipped ArkasrProcessor	token-identical (47/47)

License

Apache-2.0, inherited from Audio8/ARK-ASR-3B. All credit for the model itself goes to Audio8 / AutoArk; this repo contributes only the export.

Revision history

• v2 — max_total_len 1024 → 512 (~20% faster end-to-end, measured on 3 s and 30 s clips); embedding stored as fp16 (1.24 GB → 622 MB); added a weight-only int4 LM tier.
• v1 — initial export (max_total_len 1024, int8 LM, fp32 embedding).