[whisper] Add Flash Attention and batched decoding for up to 10x speedup

[ilyasmukiev]

Summary

Added Flash Attention and batched segment decoding to mlx-whisper, achieving 9.5x speedup on Apple Silicon.

Changes

1. Flash Attention (whisper.py)

• Replaced manual QKV attention with mx.fast.scaled_dot_product_attention
• Conditional path: uses flash attention by default, falls back to standard attention when word_timestamps=True (needs QK weights for DTW alignment)
• Proper mask handling for autoregressive decoding with KV cache

2. Batched decoding (transcribe.py)

• New batch_size parameter in transcribe() (default=1, fully backward-compatible)
• Pre-slices audio into fixed 30s chunks, stacks into batch tensor (N, 3000, n_mels), decodes simultaneously
• Per-segment temperature fallback for quality control
• batch_size=1 produces identical output to current code

Zero new dependencies. No breaking changes.

Benchmarks (M2 8GB, whisper-small, 5 min Russian audio)

Mode	Time	Realtime Factor	Speedup
Sequential (batch_size=1)	9.4s	4.8x RT	1x
Batched (batch_size=12)	6.6s	44.8x RT	9.5x

For a 15-hour video: ~20 minutes instead of ~3 hours.

Code

Full implementation with benchmarks: https://github.com/ilyasmukiev/mlx-whisper-pr

• Branch flash-attention-batch: minimal changes (Flash Attention + batch_size parameter only)
• Branch full-batching-vad-diarize: adds optional VAD (Silero) and speaker diarization

Standalone package: https://github.com/ilyasmukiev/mlx-whisper-fast

Notes

• Could not create a PR directly because gh repo fork returns HTTP 502 (repo too large?)
• Happy to submit a proper PR once the fork works
• The batched path uses fixed-stride chunking (no dynamic seeking), which is a deliberate trade-off for parallelism — same approach as WhisperX and lightning-whisper-mlx
• Related: Discussion Unexpected Processing Times for Short vs. Long Audio Files with MLX-Whisper #1275 where batching was acknowledged as possible but not implemented

[ilyasmukiev]

Update: Honest Benchmark Results (M2 8GB)

After thorough testing on a 1h14m video (M2 MacBook Air 8GB), the batching approach does not provide speedup on 8GB machines:

Config	Time	RTF	vs Sequential
whisper-small sequential	3.3 min	23.6x	baseline (fastest)
whisper-small batch=12	4.6 min	16.3x	0.7x (slower)
whisper-turbo sequential	6.0 min	12.4x	—
whisper-turbo batch=4	8.9 min	8.3x	0.67x (slower)

Root cause: On 8GB unified memory, the overhead from mx.stack + enlarged KV cache + memory bandwidth pressure exceeds the parallelism benefit. The autoregressive decoder is the bottleneck, and it doesn't parallelize well with batching.

Flash Attention (mx.fast.scaled_dot_product_attention) also showed no improvement (~0.96x) because Whisper's sequences are short (1500 encoder, ~50 decoder) — flash attention benefits long sequences.

Interesting finding: whisper-small is 2x faster than whisper-turbo (the default) on M2 8GB, with comparable quality. This alone cuts processing time from ~26 min to ~16 min for a 5h video.

Batching may still help on machines with 16GB+ RAM and large-v3 models where the encoder (32 layers) is the true bottleneck. I don't have hardware to test this.

The VAD (silence skipping) and speaker diarization modules remain useful as standalone features.

[m13v]

the whisper-small > turbo finding matches what we saw building on-device transcription for a macOS app. on 8GB machines the smaller models are consistently more responsive, and the decoder is always the bottleneck for short audio segments. one thing that helped us more than batching was aggressive VAD pre-filtering to skip silence before feeding chunks to the model - cuts total processing time by ~40% on typical meeting audio where half the recording is dead air. the flash attention result isn't surprising either, whisper's sequence lengths are just too short for it to matter.

[m13v]

fwiw here's how we handle the VAD + transcription pipeline in our macOS app - the TranscriptionService handles the chunking and silence detection before passing to WhisperKit: https://github.com/m13v/fazm/blob/main/Desktop/Sources/TranscriptionService.swift