Masterx's picture
Add files using upload-large-folder tool
6342793 verified
|
Raw
History Blame Contribute Delete
3.14 kB
---
license: apache-2.0
language:
- ar
- de
- el
- en
- es
- fr
- it
- ja
- ko
- nl
- pl
- pt
- vi
- zh
pipeline_tag: automatic-speech-recognition
tags:
- audio
- speech-recognition
- transcription
- onnx
- directml
base_model:
- CohereLabs/cohere-transcribe-03-2026
---
# Cohere Transcribe 03-2026 — DirectML-safe, fully-static ONNX
ONNX export of [CohereLabs/cohere-transcribe-03-2026](https://huggingface.co/CohereLabs/cohere-transcribe-03-2026)
restructured to run **correctly and fast on the ONNX Runtime DirectML EP** (and every other EP).
Weight bytes are identical to
[onnx-community/cohere-transcribe-03-2026-ONNX](https://huggingface.co/onnx-community/cohere-transcribe-03-2026-ONNX);
the encoder proto additionally embeds the cross-attention projection weights hoisted out of the
decoder (see below).
## Why the stock export can't run well on DirectML
- The fused decoder attention (`com.microsoft.MultiHeadAttention` cross-attn + `GroupQueryAttention`
self-attn) **crashes / mis-computes** on the DML EP (ORT ≤ 1.24, unfixed).
- Even after decomposing to plain ops, the DML EP **re-fuses its graph on every autoregressive
shape change** (~34 ms/token fixed), so naïve GPU decode loses to CPU.
## What changed (three passes, `tools/onnx/cohere_decompose_attention.py`)
1. **Attention decomposed** to Reshape/Transpose/MatMul/Softmax/… with an explicit causal mask; the
two cross-attn `If` nodes flattened away.
2. **Loop-invariant cross-KV hoisted into the encoder** — the cross-attention K/V (which depend only
on the audio) are computed once per utterance in `encoder_model*.onnx` (extra `cross_attn.*`
outputs) instead of on every decoded token.
3. **Fully static decode** — the self-KV cache is a fixed-length buffer updated by a masked write,
the cross-KV are padded to a fixed length with a `cross_bias` mask, and every per-step shape is
constant. The DML EP now compiles the decoder **once** and reuses it, so per-token cost is
**constant ~4 ms regardless of audio length**, entirely on the GPU.
Two decoders are shipped per precision:
- `decoder_model_merged*.onnx` — **fully static** (fixed self- + cross-KV). Fastest on DirectML
(metadata `winstt_static_kv`).
- `decoder_model_merged*_dyn.onnx` — **growing self-KV** (cross still fixed). Faster on the CPU EP,
which doesn't benefit from static shapes. WinSTT loads this one when the decoder is CPU-bound.
Both are numerically identical to the fused graph on CPU (autoregressive logits parity: bit-exact
for fp32/int8/q4; fp16-rounding-level for fp16/q4f16). One padded encoder feeds both.
Measured decode per-token (fp32, ORT 1.24, RTX 3080 Ti), 3.4–4.4× faster than the prior hybrid:
| | DirectML (static) | CPU (dynamic) |
|---|---|---|
| any length (5 s … 66 s) | **~4 ms/token** | ~12 ms/token |
Produced by [WinSTT](https://github.com/dahshury/WinSTT).
## Files
- `onnx/encoder_model[_fp16|_int8|_q4|_q4f16].onnx` (+ external data) — hoisted, cross-padded encoder.
- `onnx/decoder_model_merged[_fp16|_int8|_q4|_q4f16].onnx` (+ `_dyn`) — static / dynamic decoders.
- Tokenizer / configs — unchanged.