Masterx's picture
Add files using upload-large-folder tool
6342793 verified
|
Raw
History Blame Contribute Delete
3.14 kB
metadata
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 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; 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*.onnxfully static (fixed self- + cross-KV). Fastest on DirectML (metadata winstt_static_kv).
  • decoder_model_merged*_dyn.onnxgrowing 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.

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.