Model

Active contributors: Jong Wook Kim

Purpose

Implement the Whisper Transformer encoder–decoder. The encoder turns a log-Mel spectrogram into audio features; the decoder autoregressively predicts text tokens that are conditioned both on its own past tokens (self-attention with a causal mask) and on the encoder output (cross-attention). This file also defines the KV-caching mechanism, the alignment-head bookkeeping for word timestamps, and the dtype-flexible LayerNorm/Linear/Conv1d overrides.

Directory layout

whisper/
└── model.py

Key abstractions

Symbol	File / location	Description
ModelDimensions	whisper/model.py	Frozen-shape dataclass: n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer, n_vocab, n_text_ctx, n_text_state, n_text_head, n_text_layer. Stored in every checkpoint.
LayerNorm	whisper/model.py	nn.LayerNorm that always computes in fp32 then casts back to the input dtype.
Linear	whisper/model.py	nn.Linear that casts weight and bias to the input dtype, so the module runs in fp16 when the input is fp16.
Conv1d	whisper/model.py	Same idea as Linear, applied to _conv_forward.
sinusoids(length, channels, max_timescale=10000)	whisper/model.py	Standard sinusoidal positional embedding helper used by the encoder.
MultiHeadAttention	whisper/model.py	Q/K/V/O projections plus qkv_attention. Optionally uses torch.nn.functional.scaled_dot_product_attention.
disable_sdpa()	whisper/model.py	Context manager that flips MultiHeadAttention.use_sdpa to False so the manual attention path runs.
ResidualAttentionBlock	whisper/model.py	LN → self-attn → (LN → cross-attn) → LN → MLP → residual. Used in both encoder and decoder.
AudioEncoder	whisper/model.py	Two strided convs (mels → state, stride 2) → sinusoidal pos. embedding → N self-attention blocks → LN.
TextDecoder	whisper/model.py	Token embedding + learned pos. embedding → N self+cross-attention blocks → LN → tied output projection.
Whisper	whisper/model.py	Wraps AudioEncoder and TextDecoder. Exposes forward, embed_audio, logits, set_alignment_heads, install_kv_cache_hooks, plus decode, detect_language, transcribe bound from sibling modules.

How it works

graph TD
    subgraph "Whisper.forward"
        mel[mel: B x n_mels x 3000] --> encoder[AudioEncoder]
        tokens[tokens: B x L] --> decoder[TextDecoder]
        encoder --> features[audio features]
        features --> decoder
        decoder --> logits[B x L x n_vocab]
    end

Dtype handling

The custom LayerNorm, Linear, and Conv1d make the model dtype-agnostic. When you call mel.half() and run model.encoder(mel), the encoder's weights are implicitly cast to fp16 inside each layer's forward. This keeps the nn.Module graph identical for fp32 and fp16; there is no separate model_fp16 codepath.

LayerNorm is special: it always computes in fp32 (super().forward(x.float())) before casting back, which avoids known fp16 instability in normalization layers.

Attention

MultiHeadAttention.qkv_attention has two implementations:

• When MultiHeadAttention.use_sdpa is true and PyTorch's scaled_dot_product_attention is available (added in v20240930), the optimized SDPA kernel is called and the QK matrix is not returned (qk = None).
• Otherwise the attention is computed manually with explicit softmax over q @ k^T. The QK tensor is detached and returned.

The manual path is required when consumers (e.g. word-timestamp DTW alignment in whisper/timing.py:find_alignment) need access to attention weights. Use the disable_sdpa() context manager to force it.

KV cache

install_kv_cache_hooks(cache=None) walks the decoder and registers a forward hook on every MultiHeadAttention.key and MultiHeadAttention.value linear projection. The hook closure inspects whether this is the first call (no entry yet for the module, or the output is the full encoder context) or a subsequent autoregressive step (smaller output that should be appended to the cached tensor):

def save_to_cache(module, _, output):
    if module not in cache or output.shape[1] > self.dims.n_text_ctx:
        cache[module] = output  # first call (or cross-attention)
    else:
        cache[module] = torch.cat([cache[module], output], dim=1).detach()
    return cache[module]

Returning the cached tensor is what makes it visible to the rest of the attention computation. The matching cleanup_caching removes the hooks. whisper/decoding.py:PyTorchInference is the main consumer.

Cross-attention versus self-attention

ResidualAttentionBlock(cross_attention=True) adds a second MultiHeadAttention and LayerNorm. In forward, when xa (the encoder output) is passed in, the cross-attention layer attends from the decoder hidden state to the encoder features. MultiHeadAttention.forward short-circuits the K/V projection on subsequent autoregressive steps (because the cache stores them after the first call):

if kv_cache is None or xa is None or self.key not in kv_cache:
    k = self.key(x if xa is None else xa)
    v = self.value(x if xa is None else xa)
else:
    k = kv_cache[self.key]
    v = kv_cache[self.value]

So cross-attention K/V are computed exactly once per audio segment.

Alignment heads

The decoder has n_text_layer × n_text_head cross-attention heads, but only a few of them carry consistent audio-text alignment information. set_alignment_heads(dump) decompresses a base85 + gzip-compressed boolean mask of shape (n_text_layer, n_text_head) and stores it as a sparse tensor buffer:

array = np.frombuffer(gzip.decompress(base64.b85decode(dump)), dtype=bool).copy()
mask = torch.from_numpy(array).reshape(self.dims.n_text_layer, self.dims.n_text_head)
self.register_buffer("alignment_heads", mask.to_sparse(), persistent=False)

The masks are stored per model in _ALIGNMENT_HEADS in whisper/__init__.py and applied in load_model(). If a custom checkpoint is loaded by path, no mask is set and the model's default — the last half of decoder layers — is used.

Default forward

Whisper.forward(mel, tokens) runs the encoder and then the decoder. It is mostly used by tests; production code calls embed_audio once per chunk and logits (or decoder directly) per autoregressive step so the encoder runs only once.

Whisper.is_multilingual checks n_vocab >= 51865. Whisper.num_languages returns n_vocab - 51765 - int(is_multilingual) — i.e. the count of language tokens currently present in the vocabulary, derived from the vocabulary size.

Integration points

• Imports from: whisper/decoding.py (decode, detect_language are bound onto Whisper); whisper/transcribe.py (transcribe is bound onto Whisper).
• Imported by: whisper/__init__.py (Whisper, ModelDimensions); whisper/decoding.py (disable_sdpa-aware code paths); whisper/timing.py (disable_sdpa).
• Side effects: registers buffers, installs forward hooks (caller is responsible for cleanup).

Entry points for modification

• New attention variants: subclass MultiHeadAttention and replace it in ResidualAttentionBlock.__init__.
• Different positional encodings: replace the register_buffer("positional_embedding", sinusoids(...)) call in AudioEncoder.__init__ and the nn.Parameter(torch.empty(n_ctx, n_state)) in TextDecoder.__init__.
• Changing alignment heads for a custom checkpoint: call set_alignment_heads(your_dump) after load_model(path).
• Quantization or LoRA: substitute Linear for a quantized variant; the dtype-flexible base means most of the framework will keep working.