[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.