◧ ANE Knowledge Base
EN ES

Chapters · 06

Case Study: WhisperKit — Production ASR on the ANE (Argmax)

Part of the ANE Knowledge Base. Sources: argmaxinc/whisperkittools (cloned at references/whisperkittools) and its core dependency argmaxtools 0.1.23 (PyPI snapshot vendored at references/argmaxtools-0.1.23-pypi-snapshot; the GitHub repo is no longer public).

WhisperKit is Argmax's on-device speech recognition stack. whisperkittools is the Python side: it converts OpenAI Whisper checkpoints (Hugging Face) into ANE-optimized Core ML models consumed by the WhisperKit Swift runtime, applies compression, and benchmarks the results (published to hf.co/argmaxinc/whisperkit-coreml). It is the best public example of Apple's four principles applied at production scale by a third party — and of the paradigms that had to be invented beyond them: autoregressive KV-cached decoding, pluggable attention implementations, context prefill, programmatic ANE-dispatch verification, and mixed-bit compression (doc 07).

1. Architecture: one model becomes four .mlmodelc components#

A Whisper deployment is decomposed into independent Core ML models, each traced/converted separately (see scripts/generate_model.py::rearrange_quantized_variants):

Component Source Role
MelSpectrogram.mlmodelc whisperkit/audio_encoder.py::WhisperMelSpectrogram audio → log-mel features (torch.stft + mel filterbank as an exportable nn.Module)
AudioEncoder.mlmodelc whisperkit/audio_encoder.py::WhisperAudioEncoder mel features → encoder embeddings (runs once per 30 s window)
TextDecoder.mlmodelc whisperkit/text_decoder.py::WhisperTextDecoder one token per call, KV-cached autoregressive decode
TextDecoderContextPrefill.mlmodelc (optional) text_decoder.py::WhisperTextDecoderContextPrefill KV-cache lookup table for task/language prefixes

Why decompose? Each component has a different execution cadence (once per window vs once per token), different optimal SDPA implementation, and different compression sensitivity. Decomposition also lets the runtime schedule them independently.

2. The Apple principles, verbatim — via argmaxtools.nn#

argmaxtools is Argmax's generalized ANE transformer library (their equivalent of ane_transformers, extended for modern LLM-era needs). The lineage is explicit — nn.py cites Apple's repo directly above its LayerNorm.

KB principle Where it appears in argmaxtools/whisperkittools
P1 — BC1S layout, nn.Conv2d 1×1 Attention.__init__: q/k/v/o_proj = nn.Conv2d(embed_dim, ..., 1); FFN likewise (argmaxtools/nn.py). All I/O shapes are (batch, embed_dim, 1, seq_len)
P1 — channel-dim LayerNorm argmaxtools.nn.LayerNorm normalizes dim=1, optional clip_mag clamp — copied from Apple's reference (citation at nn.py:498-499)
P2 — chunking / split softmax _sdpa.SplitHeadsQ: per-head splits + per-head softmax dim=1 plus query-sequence chunking (chunk_size=256) — an extension of Apple's P2
P3 — minimal copies, einsum Same einsum pair (bchq,bkhc->bkhq / bkhq,bchk->bchq) in SplitHeadsQ; single key transpose ("transposeThenSplit is more efficient" pattern)
P4 — bandwidth-boundness Weight palettization (doc 07); single-token decode is inherently bandwidth-bound, hence compression is a latency optimization too
FP16 hygiene Masks use -1e4 (causal mask, decoder_key_padding_mask); lazy softmax denominators clamped (.clamp(min=1e-6))
Checkpoint compatibility hooks argmaxtools/utils.py::linear_to_conv2d_map_{attention,ffn,mlp} — generalized with name-alias tables (q_proj/query_proj/linear_q/…) so one hook adapts many upstream naming conventions

New micro-optimization not in Apple's work: the key projection bias is dropped entirely — k_proj = nn.Conv2d(..., bias=False) with the comment "key bias is redundant due to softmax invariance" (nn.py:88-89); the state-dict hook discards incoming k_proj.bias. (Adding a constant to every attention logit column doesn't change the softmax — softmax over keys is applied per query against all keys shifted equally... precisely: a bias on K contributes q·b, constant across keys for a given query, and softmax is shift-invariant.)

Also unlike Apple's port: argmaxtools.LayerNorm keeps the standard w*x + b order, so no bias/scale correction hook is needed.

3. New paradigm #1: SDPA as a pluggable, per-component strategy#

Apple's 2022 article prescribed one attention shape. Argmax's key insight is that the optimal SDPA implementation depends on the workload, so Attention takes a runtime-configurable sdpa_implementation (argmaxtools/_sdpa.py) with a common interface:

Implementation Technique Best for
Cat (default) No splits: 4D view to (B, h, c, x), one einsum bhcq,bhck->bhqk, one softmax dim=3. Views are free; minimal op count Short query sequences — used for TextDecoder (q_seq_len = 1)
SplitHeadsQ Apple-style per-head split + split softmax (dim=1), plus chunking the query sequence every 256 positions Long sequences — used for AudioEncoder (1500 tokens)
SplitKV Memory-efficient attention with lazy (online) softmax (Rabe & Staats 2021): split the key/value sequence into chunks, keep running max/sum, merge Very long KV sequences where the full attention matrix won't fit cache
SharedSplitKVCached Two-chunk lazy softmax over (key_cache, current_key) — attends to a batch-1 shared cache from a batch-N query Batched KV-cached decode (e.g., token-tree / speculative patterns)

The production defaults in scripts/generate_model.py are the punchline:

--audio-encoder-sdpa-implementation  default: SplitHeadsQ   # long seq → split softmax wins
--text-decoder-sdpa-implementation   default: Cat           # seq_len=1 → compact graph wins

Lesson for the KB: Apple's split-softmax recipe is not universally optimal — for single-token decoding, the chunking overhead outweighs cache-residency benefits. Choose per workload; keep the model interface stable so the choice is a one-line swap.

Attention also generalizes beyond Whisper: AttentionHeadType supports MHA / GQA / MQA (KV-head tiling via repeat_kv), plus optional RoPE (_positional_encoding.py) and RMSNorm — i.e., the same library is Llama-ready.

4. New paradigm #2: KV-cached autoregressive decoding on the ANE#

The ANE needs static shapes, so the classic "growing KV cache" must be re-engineered. WhisperTextDecoder (with AttentionType.KVCachedSelfAttention) shows the whole pattern:

Cache as explicit model I/O, in BC1S, fused across layers#

key_cache / value_cache inputs: (batch, embed_dim * n_layers, 1, max_seq_len)

All layers' caches travel as one tensor concatenated on the channel axis, split(d_model, dim=1) inside the graph (text_decoder.py:207). Outputs return only the single-token updates (key_cache_updates, value_cache_updates, shape (..., 1)); the Swift runtime writes them into its persistent buffer. One I/O pair instead of 2×n_layers tensors.

Fixed-length cache + masks instead of dynamic shapes#

  • The cache is always max_seq_len long (e.g., 448). Unused slots are disabled with decoder_key_padding_mask (-1e4 additive, the KB's FP16-safe convention).
  • kv_cache_update_mask is a one-hot vector marking the current token position; inside attention the cache update is a masked blend — pure elementwise ops, ANE-friendly:
# argmaxtools/nn.py::_finalize_kv
key_cache   = key_cache   * (1. - kv_cache_update_mask) + current_key   * kv_cache_update_mask
value_cache = value_cache * (1. - kv_cache_update_mask) + current_value * kv_cache_update_mask
  • cache_length (an int input) selects the position embedding for the current step: embed_positions(cache_length).

Cross-attention cache: compute once, delete the projections#

Whisper's decoder cross-attends to a fixed encoder output, so K/V are computed once per audio window. StatefulKVCachedEncoderDecoderCrossAttention literally does delattr(self, "k_proj")/"v_proj" — those projections belong to the encoder-side pass, and the decoder only consumes cached tensors (nn.py:452-463).

Stateful variant: Core ML MLState (iOS 18 / macOS 15)#

StatefulKVCachedAttention registers the caches as register_buffers; the conversion helper maps them to ct.StateType inputs (argmaxtools/test_utils.py::_create_coreml_model), and inference uses model.make_state() + predict(..., state=...). The cache then lives inside the Core ML model — no per-token cache I/O traffic at all. Deployment-target gated: states require iOS18/macOS15+.

5. New paradigm #3: context prefill as a lookup-table model#

Every Whisper transcription window starts with the same 3-token prefix pattern: <|startoftranscript|> <|language|> <|task|>. Instead of running the decoder 3 times, WhisperTextDecoderContextPrefill:

  1. Enumerates all valid (language × task) prefixes (~99 languages × 2 tasks),
  2. Precomputes the decoder KV caches for each prefix and stores them flattened in two nn.Embedding lookup tables,
  3. Exports as a tiny Core ML model: (task, language) → (key_cache_prefill, value_cache_prefill).

The subtle insight (documented in the code, text_decoder.py:342-345): the prefix caches should depend on the encoder output, but due to causal masking during training, decoder KV embeddings for the forced prefix tokens are uncorrelated with the audio — so they can be estimated once with a batch-mean over random encoder outputs. Runtime saves 3 of ~N decoder forward passes per window and starts decoding at cache_length=3.

6. Word timestamps via alignment-head attention weights#

For token-level timestamps, configure_for_token_timestamps() flips _return_w on the cross-attention of specific alignment heads (from the model's generation_config), harvests their attention weights with a register_forward_hook, and returns their mean as a 4th model output, alignment_heads_weights (text_decoder.py:143-170). The DTW alignment then happens in the Swift runtime. This mirrors Apple's own pattern of harvesting attention weights via hooks rather than plumbing them through return values.

7. Whisper-specific ANE adaptations#

  • Conv1d stem → Conv2d: Whisper's two Conv1d layers are executed as F.conv2d with the Conv1d weights unsqueezed (weight[:, :, None, :]) so the whole graph stays 4D/BC1S (audio_encoder.py::pre_transformer_proj).
  • Mel spectrogram as a model: torch.stft + Hann window + mel filterbank wrapped in an nn.Module with register_buffers → its own .mlmodelc; keeps DSP off the app's CPU code path.
  • Tied embeddings for logits: final projection reuses embed_tokens.weight via F.linear — no separate vocab projection matrix to store or palettize.

8. New paradigm #4: verification as CI, not as a manual Xcode step#

The KB's doc 05 verifies dispatch with the Xcode Performance tab. argmaxtools automates all of it in unittest mixins (argmaxtools/test_utils.py):

  • Correctness: PSNR between PyTorch and Core ML outputs must exceed 35 dB (compute_psnr, TEST_PSNR_THR).
  • Speedup: median latency on CPU_AND_NE vs CPU_ONLY (same .mlpackage, reloaded with a different compute unit) must beat a threshold; plus a FLOP counter for TFlop/s.
  • Compression must not regress: 1-bit palettized variant must retain ≥ 0.95× speed.
  • Programmatic compute plan (_print_compute_plan, coremltools ≥ 8.1): loads the compiled .mlmodelc with ct.models.compute_plan.MLComputePlan and logs, per op, the dispatch device, supported devices, and estimated cost share — summarized as ANE support coverage % and ANE dispatch %, dumped to a .mlcomputeplan.json. This is the Xcode performance report, scriptable in CI.
  • Multifunction models (ct.utils.MultiFunctionDescriptor, iOS 18): several input-shape variants exported as functions of one .mlpackage with shared weights — the modern answer to "tracing bakes in static shapes".
  • Model bisecting (ct.models.utils.bisect_model): split an oversized model into a chunked pipeline.
  • Reproducibility metadata: InferenceContextSpec/AppleSiliconContextMixin capture device (gpu_core_count, RAM, chip name), OS, code commit, and model commit for every benchmark; Core ML metadata fields (whisperkit_version, per-input descriptions) are stamped on every exported model (whisperkit/test_utils.py::set_metadata_for_whisper_decoder).
  • End-to-end evals: whisperkit-evaluate-model computes WER over librispeech/earnings22/Common Voice via HF datasets; results published to a public HF space. Quality is regression-tested, not eyeballed.

Notable pragmatic detail: the generation pipeline relaxes TEST_MIN_SPEEDUP_VS_CPU to 0.3 (generate_model.py:24) — for some components (a KV-cached decoder step is tiny), beating the CPU isn't the point; not blocking the pipeline is.

9. Transferable lessons (KB checklist additions)#

  1. Decompose the pipeline into per-cadence Core ML models (feature extractor / encoder / decoder / prefill LUT).
  2. Make SDPA pluggable and choose per component: split softmax for long sequences, compact Cat for single-token decode.
  3. KV cache on ANE = fixed max length + one-hot update mask + padding mask, cache fused across layers on the channel axis; or MLState on iOS18+.
  4. Cross-attention K/V: compute once per context, delete the projections from the decoder graph.
  5. Precompute anything enumerable into an embedding-LUT model (context prefill).
  6. Drop mathematically redundant parameters (key bias).
  7. Verify with PSNR thresholds + programmatic compute plans in CI, not manual Xcode inspection.
  8. Compress for bandwidth-bound regimes — see doc 07.

Next: 07 — Model Compression for the ANE.

Generated from the knowledge base markdown — every claim traces to a cited source.