Timing

Active contributors: Jong Wook Kim, ryanheise, taylorchu

Purpose

Compute word-level timestamps from a Whisper decode by aligning predicted text tokens to audio frames using DTW (Dynamic Time Warping) on the cross-attention weights of the model's "alignment heads." Provide both a CPU implementation (Numba JIT) and a GPU implementation (Triton kernel), with the GPU path falling back to CPU on launch failure.

Directory layout

whisper/
├── timing.py        # alignment + DTW driver + CPU kernels
└── triton_ops.py    # CUDA DTW + median-filter kernels

Key abstractions

Symbol	File	Description
WordTiming	whisper/timing.py	Dataclass: word, tokens, start, end, probability.
median_filter(x, filter_width)	whisper/timing.py	Median filter along the last axis. CUDA path tries triton_ops.median_filter_cuda; CPU falls back to unfold().sort().
dtw_cpu(x)	whisper/timing.py (numba)	CPU DTW + backtrace, JIT'd with @numba.jit(nopython=True, parallel=True).
dtw_cuda(x, BLOCK_SIZE=1024)	whisper/timing.py	GPU DTW that calls into triton_ops.dtw_kernel, then runs the backtrace on CPU.
dtw(x)	whisper/timing.py	Dispatches to GPU then CPU fallback.
backtrace(trace)	whisper/timing.py (numba)	Reconstructs the alignment path from a DTW trace matrix.
find_alignment(model, tokenizer, text_tokens, mel, num_frames, *, medfilt_width=7, qk_scale=1.0)	whisper/timing.py	Forward pass through the model that captures cross-attention QKs from the alignment heads, runs DTW, returns a list of WordTiming.
merge_punctuations(alignment, prepended, appended)	whisper/timing.py	Merges punctuation tokens into adjacent words.
add_word_timestamps(*, segments, model, tokenizer, mel, num_frames, ...)	whisper/timing.py	Top-level entry point used by transcribe(). Annotates each segment with words and adjusts segment-level start/end.
dtw_kernel	whisper/triton_ops.py	Triton JIT kernel implementing skewed-diagonal DTW.
median_kernel(filter_width)	whisper/triton_ops.py	Triton JIT kernel that implements bubble-sort median (specialized per width).
median_filter_cuda(x, filter_width)	whisper/triton_ops.py	Driver for median_kernel, called from whisper/timing.py:median_filter.

How it works

sequenceDiagram
    autonumber
    participant T as transcribe()
    participant A as add_word_timestamps
    participant F as find_alignment
    participant M as Whisper model
    participant D as dtw / median_filter
    T->>A: segments, mel, num_frames
    A->>F: text_tokens, mel
    F->>M: install hooks on cross_attn<br/>forward(mel, sot+text+eot)
    M-->>F: logits + cross-attn QKs<br/>per layer
    F->>F: stack QKs at alignment_heads<br/>(layer, head) indices
    F->>F: scale + softmax + standardize
    F->>D: median_filter(weights, 7)
    F->>D: dtw(-mean over heads)
    D-->>F: text_indices, time_indices
    F->>F: split tokens to words<br/>compute start/end per word
    F-->>A: List[WordTiming]
    A->>A: merge_punctuations
    A->>A: trim long words at boundaries
    A->>A: prefer segment-level start/end<br/>when first/last word is too long
    A-->>T: segments[i]['words'] populated

Capturing cross-attention

find_alignment registers forward hooks on each block.cross_attn of the decoder. The hook stashes the QK tensor (outs[-1][0], the second return value of MultiHeadAttention.forward's tuple) at the layer index. Because the SDPA fast path returns qk = None, find_alignment runs inside with disable_sdpa(): to force the manual attention path that returns real QK weights.

It then constructs an input prefix that includes <|notimestamps|> so the model does not emit timestamp tokens during the alignment forward (those would interfere with token-frame DTW), runs a single model(mel.unsqueeze(0), tokens.unsqueeze(0)) pass, and removes the hooks.

Scoring the QK matrix

The collected QKs are indexed by model.alignment_heads (the sparse boolean buffer set by set_alignment_heads):

weights = torch.stack([QKs[L][H] for L, H in model.alignment_heads.indices().T])
weights = weights[:, :, : num_frames // 2]   # skip silence padding past content
weights = (weights * qk_scale).softmax(dim=-1)
std, mean = torch.std_mean(weights, dim=-2, keepdim=True, unbiased=False)
weights = (weights - mean) / std             # standardize per token
weights = median_filter(weights, medfilt_width)
matrix = weights.mean(axis=0)                # average across alignment heads
matrix = matrix[len(tokenizer.sot_sequence) : -1]   # drop SOT prefix and EOT

The result is a (text_tokens, time_frames) log-likelihood-ish matrix. DTW is run on its negation so the path takes the most-aligned route.

DTW

DTW finds a monotone path through the matrix that minimizes total cost. The CPU implementation in dtw_cpu is a textbook two-loop DP with three predecessor moves (diagonal c0, up c1, left c2):

if c0 < c1 and c0 < c2:  c, t = c0, 0
elif c1 < c0 and c1 < c2: c, t = c1, 1
else:                     c, t = c2, 2
cost[i, j] = x[i-1, j-1] + c
trace[i, j] = t

backtrace follows the trace from the bottom-right corner back to the origin and returns (text_indices, time_indices).

The CUDA implementation rotates the cost matrix into a "skewed" layout so each anti-diagonal lives on one row and can be processed in a single Triton iteration (for k in range(1, N + M + 1)); BLOCK_SIZE-wide vectorized loads/stores update one anti-diagonal at a time. The backtrace itself still runs on CPU after the cost matrix is filled — the kernel just produces the trace matrix, not the path.

Median filter

median_filter(x, w) smooths each token's attention curve along time before averaging across heads. CPU does x.unfold(-1, w, 1).sort()[0][..., w // 2]. CUDA uses median_kernel(w) from whisper/triton_ops.py, which is a Triton JIT function whose source code is template-substituted with bubble-sort iterations specialized for the requested filter width — a width-7 median needs only six exchange passes.

From DTW path to word timings

Tokens are split into words with tokenizer.split_to_word_tokens(text_tokens + [eot]). Word boundaries in the token sequence map to specific points on the DTW path; the diff of text_indices tells which time index corresponds to each word boundary:

jumps = np.pad(np.diff(text_indices), (1, 0), constant_values=1).astype(bool)
jump_times = time_indices[jumps] / TOKENS_PER_SECOND
start_times = jump_times[word_boundaries[:-1]]
end_times = jump_times[word_boundaries[1:]]

A per-word probability is the mean of per-token probabilities over the word's token range.

Punctuation merging

merge_punctuations(alignment, prepend, append) walks the word list twice:

• First pass (right-to-left): if a word is itself punctuation that starts with a space and matches prepend_punctuations, it is pushed onto the next word.
• Second pass (left-to-right): if a word does not end with a space and the next word is in append_punctuations, it is appended to the previous word.

This produces output like "hello," and "(world)" where Whisper's tokenizer would otherwise emit punctuation as separate "words."

Long-word trimming

add_word_timestamps then trims unrealistically long words at sentence boundaries (where the audio padding likely got eaten by a final period) and at segment boundaries (where the first or last word might extend into silence). The thresholds use max_duration = 2 * median_word_duration, with median_word_duration clamped at 0.7 s.

It also optionally aligns the segment-level start / end with the word-level boundaries, picking whichever is more conservative depending on how badly the first/last words overshoot.

Integration points

• Imports from: whisper/audio.py (HOP_LENGTH, SAMPLE_RATE, TOKENS_PER_SECOND); whisper/tokenizer.py:Tokenizer; whisper/model.py:disable_sdpa; whisper/triton_ops.py (lazy import inside dtw_cuda and median_filter_cuda).
• Imported by: whisper/transcribe.py:add_word_timestamps; tests under tests/test_timing.py.

Entry points for modification

• Different alignment heads: change model.alignment_heads (call Whisper.set_alignment_heads(your_dump) after loading) — see _ALIGNMENT_HEADS in whisper/__init__.py.
• Different attention smoothing: change medfilt_width or replace the median filter with a different smoother in find_alignment.
• Different DTW step set: edit dtw_cpu and whisper/triton_ops.py:dtw_kernel together. Don't change one without the other.
• Word-splitting heuristics: see Tokenizer.split_to_word_tokens in Tokenizer.

See also: Decoding (the segments fed in here come from decode()) and Transcribe (the caller of add_word_timestamps).