perf(whisper): remove long-form decode overhead

whisper/mlx_whisper/decoding.py

@@ -332,67 +332,79 @@ def __init__(
         self.tokenizer = tokenizer
         self.sample_begin = sample_begin
         self.max_initial_timestamp_index = max_initial_timestamp_index
+        self._vocab = None
 
     def apply(self, logits: mx.array, tokens: mx.array) -> mx.array:
-        mask = np.zeros(logits.shape, np.float32)
+        if self._vocab is None:
+            self._vocab = mx.arange(logits.shape[-1])[None, :]
+        vocab = self._vocab
+        mask = mx.zeros(logits.shape, dtype=mx.bool_)
+
         # suppress <|notimestamps|> which is handled by without_timestamps
         if self.tokenizer.no_timestamps is not None:
-            mask[:, self.tokenizer.no_timestamps] = -np.inf
+            mask = mask | (vocab == self.tokenizer.no_timestamps)
 
         ## timestamps have to appear in pairs, except directly before EOT; mask logits accordingly
-        tokens = tokens.tolist()
-        for k in range(len(tokens)):
-            seq = tokens[k][self.sample_begin :]
-            last_was_timestamp = (
-                len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin
-            )
+        sampled = tokens.shape[1] - self.sample_begin
+        if sampled >= 1:
+            seq = tokens[:, self.sample_begin :]
+            last_was_timestamp = seq[:, -1] >= self.tokenizer.timestamp_begin
             penultimate_was_timestamp = (
-                len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin
+                mx.ones_like(last_was_timestamp)
+                if sampled < 2
+                else seq[:, -2] >= self.tokenizer.timestamp_begin
+            )
+            mask = mask | (
+                (last_was_timestamp & penultimate_was_timestamp)[:, None]
+                & (vocab >= self.tokenizer.timestamp_begin)
+            )
+            mask = mask | (
+                (last_was_timestamp & ~penultimate_was_timestamp)[:, None]
+                & (vocab < self.tokenizer.eot)
             )
 
-            if last_was_timestamp:
-                if penultimate_was_timestamp:  # has to be non-timestamp
-                    mask[k, self.tokenizer.timestamp_begin :] = -np.inf
-                else:  # cannot be normal text tokens
-                    mask[k, : self.tokenizer.eot] = -np.inf
+            # Preserve the legacy position-based lower bound for output parity.
+            # Token-value semantics are tracked separately in KT-642.
+            timestamp_positions = mx.where(
+                seq > self.tokenizer.timestamp_begin,
+                mx.arange(sampled)[None, :],
+                -1,
+            )
+            last_timestamp = timestamp_positions.max(axis=-1)
+            has_timestamp = last_timestamp >= 0
+            last_timestamp = last_timestamp + (
+                (last_timestamp == 0) | penultimate_was_timestamp
+            )
+            mask = mask | (
+                has_timestamp[:, None]
+                & (vocab >= self.tokenizer.timestamp_begin)
+                & (vocab < last_timestamp[:, None])
+            )
 
-            timestamps = [
-                i for i, v in enumerate(seq) if v > self.tokenizer.timestamp_begin
-            ]
-            if len(timestamps) > 0:
-                # timestamps shouldn't decrease; forbid timestamp tokens smaller than the last
-                # also force each segment to have a nonzero length, to prevent infinite looping
-                last_timestamp = timestamps[-1]
-                if not last_timestamp or penultimate_was_timestamp:
-                    last_timestamp += 1
-                mask[k, self.tokenizer.timestamp_begin : last_timestamp] = -np.inf
-
-        if len(tokens[0]) == self.sample_begin:
+        if tokens.shape[1] == self.sample_begin:
             # suppress generating non-timestamp tokens at the beginning
-            mask[:, : self.tokenizer.timestamp_begin] = -np.inf
+            mask = mask | (vocab < self.tokenizer.timestamp_begin)
 
             # apply the `max_initial_timestamp` option
             if self.max_initial_timestamp_index is not None:
                 last_allowed = (
                     self.tokenizer.timestamp_begin + self.max_initial_timestamp_index
                 )
-                mask[:, last_allowed + 1 :] = -np.inf
+                mask = mask | (vocab > last_allowed)
 
         # if sum of probability over timestamps is above any other token, sample timestamp
-        mask = mx.array(mask)
         logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
         timestamp_logprob = logprobs[:, self.tokenizer.timestamp_begin :].logsumexp(
             axis=-1, keepdims=True
         )
         max_text_token_logprob = logprobs[:, : self.tokenizer.timestamp_begin].max(
             axis=-1, keepdims=True
         )
-        mask[:, : self.tokenizer.timestamp_begin] = mx.where(
-            timestamp_logprob > max_text_token_logprob,
-            -mx.inf,
-            mask[:, : self.tokenizer.timestamp_begin],
+        mask = mask | (
+            (timestamp_logprob > max_text_token_logprob)
+            & (vocab < self.tokenizer.timestamp_begin)
         )
-        return logits + mask
+        return logits + mx.where(mask, -mx.inf, 0.0)
 
 
 class DecodingTask:

whisper/mlx_whisper/timing.py

@@ -7,7 +7,6 @@
 import mlx.core as mx
 import numba
 import numpy as np
-from scipy import signal
 
 from .audio import HOP_LENGTH, SAMPLE_RATE, TOKENS_PER_SECOND
 from .tokenizer import Tokenizer
@@ -20,28 +19,33 @@ def median_filter(x: np.ndarray, filter_width: int):
     """Apply a median filter of width `filter_width` along the last dimension of `x`"""
     pad_width = filter_width // 2
     if x.shape[-1] <= pad_width:
-        # F.pad requires the padding width to be smaller than the input dimension
+        # Reflect padding requires the padding width to be smaller than the input.
         return x
 
-    if (ndim := x.ndim) <= 2:
-        # `F.pad` does not support 1D or 2D inputs for reflect padding but supports 3D and 4D
-        x = x[None, None, :]
-
     assert (
         filter_width > 0 and filter_width % 2 == 1
     ), "`filter_width` should be an odd number"
-
-    x = np.pad(x, ((0, 0), (0, 0), (pad_width, pad_width)), mode="reflect")
-
-    # todo: more efficient version in mlx
-    result = signal.medfilt(x.astype(np.float32), kernel_size=(1, 1, filter_width))[
-        ..., pad_width:-pad_width
-    ]
-
-    if ndim <= 2:
-        result = result[0, 0]
-
-    return result
+    if filter_width == 1:
+        return x.astype(np.float32, copy=False)
+
+    values = mx.array(x.astype(np.float32, copy=False))
+    size = values.shape[-1]
+    left = mx.take(
+        values,
+        mx.array(list(range(pad_width, 0, -1))),
+        axis=-1,
+    )
+    right = mx.take(
+        values,
+        mx.array(list(range(size - 2, size - pad_width - 2, -1))),
+        axis=-1,
+    )
+    padded = mx.concatenate([left, values, right], axis=-1)
+    windows = mx.stack(
+        [padded[..., offset : offset + size] for offset in range(filter_width)],
+        axis=-1,
+    )
+    return np.array(mx.partition(windows, pad_width, axis=-1)[..., pad_width])
 
 
 @numba.jit(nopython=True)
@@ -69,7 +73,7 @@ def backtrace(trace: np.ndarray):
     return result[::-1, :].T
 
 
-@numba.jit(nopython=True, parallel=True)
+@numba.jit(nopython=True, parallel=True, cache=True)
 def dtw_cpu(x: np.ndarray):
     N, M = x.shape
     cost = np.ones((N + 1, M + 1), dtype=np.float32) * np.inf
@@ -145,7 +149,6 @@ def find_alignment(
     text_token_probs = mx.take_along_axis(
         token_probs, mx.array(text_tokens)[:, None], axis=1
     ).squeeze(1)
-    text_token_probs = np.array(text_token_probs)
 
     # heads * tokens * frames
     weights = mx.stack(
@@ -157,6 +160,8 @@ def find_alignment(
     mean = mx.mean(weights, axis=-2, keepdims=True)
     std = mx.var(weights, axis=-2, keepdims=True, ddof=0).sqrt()
     weights = (weights - mean) / std
+    mx.eval(text_token_probs, weights)
+    text_token_probs = np.array(text_token_probs)
     weights = median_filter(np.array(weights), medfilt_width)
 
     matrix = weights.mean(axis=0)

whisper/mlx_whisper/transcribe.py

@@ -209,6 +209,7 @@ def decode_with_fallback(segment: mx.array) -> DecodingResult:
             [temperature] if isinstance(temperature, (int, float)) else temperature
         )
         decode_result = None
+        audio_features = None
 
         for t in temperatures:
             kwargs = {**decode_options}
@@ -221,7 +222,11 @@ def decode_with_fallback(segment: mx.array) -> DecodingResult:
                 kwargs.pop("best_of", None)
 
             options = DecodingOptions(**kwargs, temperature=t)
-            decode_result = model.decode(segment, options)
+            decode_result = model.decode(
+                segment if audio_features is None else audio_features,
+                options,
+            )
+            audio_features = decode_result.audio_features
 
             needs_fallback = False
             if (