From 0137df0783f93f617a83dda1ec275ca38b2f143d Mon Sep 17 00:00:00 2001
From: unohee <unohee.official@gmail.com>
Date: Mon, 27 Jul 2026 08:48:56 +0900
Subject: [PATCH] perf(whisper): reuse encoded audio for word alignment

---
 whisper/mlx_whisper/timing.py     | 11 +++++++++--
 whisper/mlx_whisper/transcribe.py |  1 +
 whisper/mlx_whisper/whisper.py    |  6 ++++--
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/whisper/mlx_whisper/timing.py b/whisper/mlx_whisper/timing.py
index 07b81186c..750fc4ae7 100644
--- a/whisper/mlx_whisper/timing.py
+++ b/whisper/mlx_whisper/timing.py
@@ -2,7 +2,7 @@
 
 import itertools
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
 import mlx.core as mx
 import numba
@@ -116,6 +116,7 @@ def find_alignment(
     mel: mx.array,
     num_frames: int,
     *,
+    audio_features: Optional[mx.array] = None,
     medfilt_width: int = 7,
     qk_scale: float = 1.0,
 ) -> List[WordTiming]:
@@ -131,7 +132,13 @@ def find_alignment(
         ]
     )
 
-    logits, cross_qk = model.forward_with_cross_qk(mel[None, :], tokens[None, :])
+    if audio_features is not None and audio_features.ndim == 2:
+        audio_features = audio_features[None, :]
+    logits, cross_qk = model.forward_with_cross_qk(
+        mel[None, :],
+        tokens[None, :],
+        audio_features=audio_features,
+    )
     # consider only the logits associated with predicting text
     sampled_logits = logits[0][len(tokenizer.sot_sequence) : -2, : tokenizer.eot]
     token_probs = mx.softmax(sampled_logits, precise=True, axis=-1)
diff --git a/whisper/mlx_whisper/transcribe.py b/whisper/mlx_whisper/transcribe.py
index bced16a58..a86caecb9 100644
--- a/whisper/mlx_whisper/transcribe.py
+++ b/whisper/mlx_whisper/transcribe.py
@@ -418,6 +418,7 @@ def next_words_segment(segments: List[dict]) -> Optional[dict]:
                         tokenizer=tokenizer,
                         mel=mel_segment,
                         num_frames=segment_size,
+                        audio_features=result.audio_features,
                         prepend_punctuations=prepend_punctuations,
                         append_punctuations=append_punctuations,
                         last_speech_timestamp=last_speech_timestamp,
diff --git a/whisper/mlx_whisper/whisper.py b/whisper/mlx_whisper/whisper.py
index 5c85195cd..9f3b4d490 100644
--- a/whisper/mlx_whisper/whisper.py
+++ b/whisper/mlx_whisper/whisper.py
@@ -247,8 +247,10 @@ def embed_audio(self, mel):
     def logits(self, tokens, audio_features):
         return self.decoder(tokens, audio_features)[0]
 
-    def forward_with_cross_qk(self, mel, tokens):
-        logits, _, cross_qk = self.decoder(tokens, self.encoder(mel))
+    def forward_with_cross_qk(self, mel, tokens, audio_features=None):
+        if audio_features is None:
+            audio_features = self.encoder(mel)
+        logits, _, cross_qk = self.decoder(tokens, audio_features)
         return logits, cross_qk
 
     def __call__(self, mel, tokens):