core/nn/whisper_decoder library
Whisper (openai) text decoder architecture.
Layout mirrors openai-whisper/whisper/model.py::TextDecoder:
x = token_embedding(tokens) + positional_embeddingoffset:offset+T
for block in N:
x = x + attn(attn_ln(x), mask=causal)
x = x + cross_attn(cross_attn_ln(x), xa)
x = x + mlp(mlp_ln(x))
x = ln(x)
logits = x @ token_embedding.weight.T
Whisper tiny.en: d_model=384, n_head=6, n_layer=4, n_ctx=448, vocab=51864.
Device-agnostic (CPU or GPU). Multi-head attention is expressed as
per-head Linear projections + Tensor.scaledDotProductAttention,
so every step stays on-device. Cross-attention K/V of the encoder
memory is cached per-block (as device tensors) via
WhisperDecoderBlock.primeCrossAttn, so each decode step pays
O(T_text) rather than O(T_audio).