A three-line modification restored Whisper's cross-attention to PyTorch's fused kernel during beam search, achieving 4.9x faster attention execution. This fix was merged upstream in openai/whisper#2812.
The headline metrics (4.9x faster kernel, 19% faster end-to-end transcription — the pull request states it as a 16% lower real-time factor, which is the same 1.19x) are the measurements in our pull request, taken on an RTX 3090 with FlashAttention. Independent testing on RTX 4060 without FlashAttention showed 5.4x attention speedup and approximately 13% end-to-end gains.
Silent performance fallbacks
PyTorch's scaled_dot_product_attention selects from multiple kernels — FlashAttention, memory-efficient variants, cuDNN paths, or unfused math versions — based on dtype, head dimensions, and tensor shapes. When tensors don't meet fast-kernel requirements, the system silently downgrades without warnings, maintaining correctness while sacrificing performance. Whisper contained such a fallback in its beam search implementation.
What was happening
Whisper's decoder efficiently computes encoder cross-attention keys and values once per 30-second audio window, reusing them across every beam. Only queries get duplicated per beam. At beam_size=5:
q [5, 20, 1, 64] ← one token, duplicated across 5 beams
k, v [1, 20, 1500, 64] ← 1500 encoder frames, computed once and shared
PyTorch broadcasts this automatically, but fused kernels require matching batch sizes. Consequently, every cross-attention call across all layers and decode steps fell back to the slow, unfused path.
Why it stayed hidden
- Correct output — all tests pass despite performance degradation
- No warnings — backend diagnostics must be explicitly enabled
- Greedy decoding unaffected — at
beam_size=1, shapes already match - Profiler obscurity — users see
aten::scaled_dot_product_attentionunless specifically searching foraten::_scaled_dot_product_attention_math
The fallback arrived with PR #2359 in September 2024. While that change accelerated greedy decoding, beam-search users on GPUs paid the unfused kernel penalty for approximately two years in the reference implementation of a widely deployed speech model.
How we found it
Neither the fallback nor its optimal fix emerges from code review or quick benchmarking. Discovery required systematically comparing candidates against measured costs using Artemis:
- Profile baseline end-to-end Whisper large-v3 to identify actual inference bottlenecks
- Execute Artemis Discovery with specified goals, benchmark, and target hardware
- Iterate through optimisation candidates with impact scoring
- Converge on a final solution through three Discovery runs producing 17 experiments and 51 candidates
Artemis identified the cross-attention fallback and generated a three-line patch, with every rejected candidate accompanied by clear reasoning.
The fix
if SDPA_AVAILABLE and MultiHeadAttention.use_sdpa:
+ if k.shape[0] == 1 and q.shape[0] != 1:
+ # Cross-attention K/V have batch 1 and broadcast against the
+ # beam-expanded query; the fused SDPA kernels reject the batch
+ # mismatch and fall back to math, so expand K/V to a stride-0
+ # view (no copy) to keep them on the fast path.
+ k = k.expand(q.shape[0], *k.shape[1:])
+ v = v.expand(q.shape[0], *v.shape[1:])
a = scaled_dot_product_attention(
q, k, v, is_causal=mask is not None and n_ctx > 1
)
.expand() repeats a tensor along size-1 dimensions without copying, returning a stride-zero view where all beams read identical memory. The tensor now reports matching batch dimensions, satisfying fused kernel preconditions. The narrow guard fires only when k's batch is 1 and q's differs — the specific beam-search configuration. Self-attention, greedy decoding, and matching batch sizes bypass the check with minimal overhead: one shape comparison.
Proving nothing regressed
Transcript verification. 200 LibriSpeech test-clean samples were processed through both implementations, comparing word error rates sample-by-sample. The merged PR reported 2.070% on its own 200-sample subset; before/after equality remained consistent regardless of normalisation differences.
Extended validation. Artemis ASR Bench evaluated nine scenarios over HTTP — clean audio, long-form content, noise, phone-quality speech, and overlapping speakers — stressing various decoder components. All nine scenarios passed exact-match validity requirements.
Above beam_size=1, both paths differed only at the final representable fp16 digit, reflecting accumulated rounding from different addition orders rather than different computations. The implementations were isolated as git worktrees one commit apart, verified to contain only these three lines of difference.
The numbers
Cross-attention calls improved from 802.5 µs to 147.8 µs — approximately 5.4x faster. End-to-end gains ranged from roughly 10 percent on short clean audio to 25 percent on overlapping speech, where demanding audio concentrates more time in the decode loop.
Gains require expanded decode batches: beam_size > 1 or best_of > 1. At beam_size=1, the guard doesn't execute, tensors remain unchanged, and overhead reduces to a single shape comparison. Above 1, the eliminated penalty grows with beam size — approximately 5x at beam size 2, roughly 7x at beam size 8.
Configuration: Whisper large-v3, fp16, beam_size=5, temperature=0.0, English language, NVIDIA RTX 4060 Laptop GPU 8GB, driver 555.97, CUDA 12.5, PyTorch 2.5.1+cu121, Python 3.12.4, Windows 11. Operator benchmark: n_head=20, head_dim=64, ctx_kv=1500, mean of 300 timed calls following 20 untimed warmup calls.
Run it on your own hardware
These results aren't hardware-specific. Original RTX 3090 measurements with FlashAttention showed cross-attention improving from 424 µs to 87 µs. This post's RTX 4060 results without FlashAttention applied the fix to the memory-efficient kernel instead — platform differences, not card limitations. Any fused backend gains once batch dimensions align.
Measure your own hardware:
from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
model.transcribe(audio, beam_size=5, fp16=True)
for e in prof.key_averages():
if "_scaled_dot_product_attention_math" in e.key:
print(f"{e.key}: {e.device_time_total/1000:.1f} ms across {e.count} calls")
Run with beam_size=5. Any printed output represents time the fused kernel could have handled. The fix lives on Whisper's main branch (not yet released), so install from source:
pip install git+https://github.com/openai/whisper.git
Profile afterward — the math-kernel entry should disappear. This pattern applies to any encoder-decoder caching cross-attention keys/values and expanding queries for beam search. The optimisation isn't Whisper-specific; similar architectures may carry identical performance overhead.
Merged upstream in openai/whisper#2812.
