⚡️ Speed up method BeamSearchDecoder.update by 80%

[codeflash-ai]

📄 80% (0.80x) speedup for BeamSearchDecoder.update in whisper/decoding.py

⏱️ Runtime : 214 milliseconds → 119 milliseconds (best of 56 runs)

📝 Explanation and details

Thank you for providing a detailed line profiler report. Most of the CPU time is spent in these "hot spots".

• logprobs[idx].topk(self.beam_size + 1) (and the nested for loop),
• new_logprob = (sum_logprobs[idx] + logprob).item()
• sequence = tuple(prefix + [token.item()])

Additional issues.

• There is heavy repeated Python looping (instead of batched tensor operations).
• Inefficient construction of sequences (list concatenation and then tuple conversion).
• Dictionary sorting and key construction is slow.

To speed up this beam search decoder, we need to vectorize as much of the candidate computation as possible, avoid repeated .item() and tolist() calls in Python, and batch candidate tracking.

Below is a rewritten version that.

• Vectorizes the candidate logprob and topk computations,
• Avoids Python loops for sequence reconstruction as much as possible,
• Avoids list/tuple concatenations where possible,
• Processes all beams for all audios in batches, not in innermost Python loops,
• Keeps finished sequences in the same data structure, but more efficiently.

Key changes:

• Use a single batch log_softmax and topk for all beams (removes one for-loop entirely).
• Batch build the new tokens by "expanding" and appending new candidates with tensor operations.
• Use efficient slicing and flattening for candidate management.
• Dict-based operations (scoring, finished-completed check) done only per audio, not per candidate.
• Eliminated redundant .tolist() and .item() (now done only when absolutely necessary, e.g., for Python dict key creation).
• Reduced the conversion of tensor-to-list-to-tuple in the inner loops.

The largest gain is from removing the innermost for-loops and vectorizing the probability+token expansion for all beams in all audios.

This should give a significant improvement.

---

If you want even more, you can use NamedTuple or hashable array types for sequences and more vectorized pruning of finished sequences, but that will require a more thorough refactoring of the API and outside code assumptions.

Let me know if you want deeper C++/Cython-level or further functional refactors!

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 33 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests Details

from typing import Optional, Tuple

# imports
import pytest  # used for our unit tests
import torch
# function to test
import torch.nn.functional as F
from torch import Tensor
from whisper.decoding import BeamSearchDecoder


# Dummy Inference class for testing
class DummyInference:
    def __init__(self):
        self.kv_rearranged = []

    def rearrange_kv_cache(self, source_indices):
        # Just keep a record of calls for test verification
        self.kv_rearranged.append(list(source_indices))

# Dummy TokenDecoder base class for compatibility
class TokenDecoder:
    pass
from whisper.decoding import BeamSearchDecoder

# unit tests

# ========== BASIC TEST CASES ==========

def test_basic_single_audio_single_beam():
    """
    Test with a single audio, beam_size=1, simple logits.
    Should select highest logit, and not complete unless EOT is picked.
    """
    beam_size = 1
    eot = 2
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    # tokens: 1 audio, 1 beam, current token is [0]
    tokens = torch.tensor([[0]])
    # logits: shape [1, vocab_size]
    logits = torch.tensor([[0.1, 0.2, 0.3]])  # EOT has highest logit
    sum_logprobs = torch.tensor([0.0])
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

def test_basic_multi_audio_multi_beam():
    """
    Test with two audio samples, beam_size=2, simple logits.
    Should keep top beams for each audio.
    """
    beam_size = 2
    eot = 3
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.tensor([[0], [1], [0], [1]])  # 2 audio, 2 beams each
    logits = torch.tensor([
        [0.1, 0.5, 0.2, 0.3],  # audio 1, beam 1
        [0.4, 0.2, 0.1, 0.3],  # audio 1, beam 2
        [0.2, 0.1, 0.6, 0.3],  # audio 2, beam 1
        [0.3, 0.2, 0.2, 0.4],  # audio 2, beam 2
    ])
    sum_logprobs = torch.zeros(4)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

def test_basic_patience():
    """
    Test that patience parameter increases max_candidates.
    """
    beam_size = 2
    patience = 2.0
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)

def test_basic_eot_not_in_logits():
    """
    Test that if EOT is not in topk, no finished sequences are added.
    """
    beam_size = 2
    eot = 4
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.tensor([[0], [1]])
    logits = torch.tensor([
        [0.1, 0.2, 0.3, 0.4, 0.0],  # EOT is last, lowest logit
        [0.2, 0.3, 0.4, 0.1, 0.0]
    ])
    sum_logprobs = torch.zeros(2)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

# ========== EDGE TEST CASES ==========

def test_edge_invalid_beam_size():
    """
    Test that invalid beam_size (0 or negative) raises AssertionError.
    """
    with pytest.raises(AssertionError):
        BeamSearchDecoder(0, 1, DummyInference())
    with pytest.raises(AssertionError):
        BeamSearchDecoder(-1, 1, DummyInference())


def test_edge_tokens_shape_mismatch():
    """
    Test that tokens.shape[0] not divisible by beam_size raises ValueError.
    """
    beam_size = 3
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.tensor([[0], [1], [2], [3]])  # 4 not divisible by 3
    logits = torch.randn(4, 5)
    sum_logprobs = torch.zeros(4)
    with pytest.raises(ValueError):
        decoder.update(tokens, logits, sum_logprobs)


def test_edge_max_candidates_limit():
    """
    Test that finished_sequences never exceeds max_candidates.
    """
    beam_size = 1
    eot = 2
    patience = 1.0
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)
    tokens = torch.tensor([[0]])
    logits = torch.tensor([[0.1, 0.2, 10.0]])  # EOT is highest
    sum_logprobs = torch.zeros(1)
    # Call update multiple times to try to overfill finished_sequences
    for _ in range(3):
        decoder.update(tokens, logits, sum_logprobs)

def test_edge_nonzero_sum_logprobs():
    """
    Test that nonzero sum_logprobs are handled correctly.
    """
    beam_size = 1
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.tensor([[0]])
    logits = torch.tensor([[0.0, 1.0]])
    sum_logprobs = torch.tensor([5.0])
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

def test_edge_eot_in_middle_of_vocab():
    """
    Test that EOT token works regardless of its index in vocab.
    """
    beam_size = 2
    eot = 1  # not last index
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.tensor([[0], [1]])
    logits = torch.tensor([
        [0.1, 10.0, 0.2],  # EOT is index 1
        [0.2, 10.0, 0.3]
    ])
    sum_logprobs = torch.zeros(2)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

# ========== LARGE SCALE TEST CASES ==========

def test_large_scale_beam_and_vocab():
    """
    Test with large beam_size and vocab size, but under 1000 elements.
    """
    beam_size = 50
    n_audio = 5
    vocab_size = 100
    eot = 99
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.zeros((n_audio * beam_size, 3), dtype=torch.long)
    logits = torch.rand((n_audio * beam_size, vocab_size))
    # Make EOT have highest logit for a few beams to test finishing
    for i in range(0, n_audio * beam_size, 10):
        logits[i, eot] = 100.0
    sum_logprobs = torch.zeros(n_audio * beam_size)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)
    # After a few more updates, at least some finished_sequences should be nonempty
    for _ in range(2):
        next_tokens, completed = decoder.update(next_tokens, logits, sum_logprobs)

def test_large_scale_multiple_updates():
    """
    Test repeated updates with large beam size and patience.
    """
    beam_size = 20
    n_audio = 3
    vocab_size = 50
    patience = 2.0
    eot = 49
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)
    tokens = torch.zeros((n_audio * beam_size, 4), dtype=torch.long)
    logits = torch.rand((n_audio * beam_size, vocab_size))
    # Force EOT to be top logit for some beams each update
    for i in range(0, n_audio * beam_size, 15):
        logits[i, eot] = 100.0
    sum_logprobs = torch.zeros(n_audio * beam_size)
    for _ in range(5):
        next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)
        tokens = next_tokens
    # Should not exceed max_candidates in any finished_sequences
    for fs in decoder.finished_sequences:
        pass

def test_large_scale_memory_limits():
    """
    Test that the function works with the largest allowed tensor sizes under 100MB.
    """
    # Each float32 is 4 bytes, so 100_000_000 / 4 = 25_000_000 elements max
    # We'll use much less: tokens (500, 10), logits (500, 100), sum_logprobs (500)
    beam_size = 50
    n_audio = 10
    vocab_size = 100
    eot = 99
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.zeros((n_audio * beam_size, 10), dtype=torch.long)
    logits = torch.rand((n_audio * beam_size, vocab_size))
    sum_logprobs = torch.zeros(n_audio * beam_size)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)

def test_large_scale_eot_completion():
    """
    Test that completed is True when all finished_sequences are full, even at scale.
    """
    beam_size = 10
    n_audio = 5
    patience = 1.0
    vocab_size = 20
    eot = 19
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)
    tokens = torch.zeros((n_audio * beam_size, 2), dtype=torch.long)
    logits = torch.zeros((n_audio * beam_size, vocab_size))
    # Set EOT to highest for all beams
    logits[:, eot] = 100.0
    sum_logprobs = torch.zeros(n_audio * beam_size)
    next_tokens, completed = decoder.update(tokens, logits, sum_logprobs)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import Optional, Tuple

# imports
import pytest  # used for our unit tests
import torch
# function to test
import torch.nn.functional as F
from torch import Tensor
from whisper.decoding import BeamSearchDecoder


# Dummy Inference class for testing rearrange_kv_cache
class DummyInference:
    def __init__(self):
        self.calls = []

    def rearrange_kv_cache(self, source_indices):
        # Just record the call for test verification
        self.calls.append(list(source_indices))

# Dummy TokenDecoder base class for compatibility
class TokenDecoder:
    pass
from whisper.decoding import BeamSearchDecoder

# unit tests

# -------------------- BASIC TEST CASES --------------------

def test_basic_single_beam_single_audio():
    # Test with beam_size=1, n_audio=1, 2 tokens in vocab
    beam_size = 1
    eot = 2
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)

    tokens = torch.tensor([[0]], dtype=torch.long)  # shape (1, 1)
    logits = torch.tensor([[2.0, 1.0, 0.0]])        # shape (1, 3)
    sum_logprobs = torch.tensor([0.0])

    # Expect topk beam_size+1 = 2: token 0 (score 2.0), token 1 (score 1.0)
    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())

def test_basic_multi_beam_single_audio():
    # Test with beam_size=2, n_audio=1, 3 tokens in vocab
    beam_size = 2
    eot = 2
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)

    tokens = torch.tensor([[0], [1]], dtype=torch.long)  # shape (2, 1)
    logits = torch.tensor([[0.0, 2.0, 1.0], [1.0, 0.0, 2.0]])  # shape (2, 3)
    sum_logprobs = torch.tensor([0.0, 0.0])

    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())
    # Both sequences should not end with eot
    for seq in new_tokens:
        pass

def test_basic_multi_audio():
    # Test with beam_size=1, n_audio=2, 2 tokens in vocab
    beam_size = 1
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)

    tokens = torch.tensor([[0], [1]], dtype=torch.long)  # shape (2, 1)
    logits = torch.tensor([[1.0, 2.0], [2.0, 1.0]])      # shape (2, 2)
    sum_logprobs = torch.tensor([0.0, 0.0])

    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())

def test_basic_eot_completion():
    # Test that completed is True when enough eot sequences are found
    beam_size = 1
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)

    tokens = torch.tensor([[0]], dtype=torch.long)  # shape (1, 1)
    logits = torch.tensor([[0.0, 10.0]])            # eot is most likely
    sum_logprobs = torch.tensor([0.0])

    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())
    # At least one finished sequence should end with eot
    found_eot = any(seq[-1] == eot for seq in decoder.finished_sequences[0])

# -------------------- EDGE CASES --------------------

def test_edge_invalid_beam_size():
    # Test that beam_size=0 raises assertion
    with pytest.raises(AssertionError):
        BeamSearchDecoder(0, 1, DummyInference())


def test_edge_tokens_shape_mismatch():
    # Test that tokens.shape[0] % beam_size != 0 raises ValueError
    decoder = BeamSearchDecoder(2, 1, DummyInference())
    tokens = torch.tensor([[0], [1], [2]], dtype=torch.long)  # shape (3, 1)
    logits = torch.tensor([[1.0, 2.0], [2.0, 1.0], [1.0, 2.0]])
    sum_logprobs = torch.tensor([0.0, 0.0, 0.0])
    with pytest.raises(ValueError):
        decoder.update(tokens, logits, sum_logprobs.clone())

def test_edge_empty_logits():
    # Test with logits of size 0
    decoder = BeamSearchDecoder(1, 1, DummyInference())
    tokens = torch.tensor([[0]], dtype=torch.long)
    logits = torch.empty((1, 0))
    sum_logprobs = torch.tensor([0.0])
    # Should raise an error because topk will fail
    with pytest.raises(RuntimeError):
        decoder.update(tokens, logits, sum_logprobs.clone())


def test_edge_max_candidates_patience():
    # Test that patience > 1 increases max_candidates
    beam_size = 2
    patience = 2.0
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)

def test_edge_finished_sequences_limit():
    # Test that finished_sequences does not exceed max_candidates
    beam_size = 1
    patience = 1.0
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference, patience=patience)
    tokens = torch.tensor([[0]], dtype=torch.long)
    logits = torch.tensor([[0.0, 10.0]])
    sum_logprobs = torch.tensor([0.0])
    # First update, should finish one sequence
    decoder.update(tokens, logits, sum_logprobs.clone())
    # Second update, should not add more if max_candidates reached
    decoder.update(tokens, logits, sum_logprobs.clone())

def test_edge_noncontiguous_input():
    # Test that non-contiguous tensors work
    beam_size = 2
    eot = 1
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.arange(4, dtype=torch.long).reshape(2, 2).t()  # shape (2,2), non-contiguous
    logits = torch.ones((2, 3))
    sum_logprobs = torch.zeros(2)
    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())

# -------------------- LARGE SCALE TEST CASES --------------------

def test_large_beam_and_vocab():
    # Test with large beam_size and vocab, but within 1000 elements
    beam_size = 20
    n_audio = 5
    vocab_size = 50
    eot = 49
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.zeros((n_audio * beam_size, 5), dtype=torch.long)
    logits = torch.randn((n_audio * beam_size, vocab_size))
    sum_logprobs = torch.zeros(n_audio * beam_size)
    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())


def test_large_multiple_updates_until_complete():
    # Simulate several updates until completion
    beam_size = 5
    n_audio = 2
    vocab_size = 10
    eot = 9
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.zeros((n_audio * beam_size, 2), dtype=torch.long)
    sum_logprobs = torch.zeros(n_audio * beam_size)
    for step in range(10):
        # Make eot more likely with each step
        logits = torch.randn((n_audio * beam_size, vocab_size))
        if step >= 8:
            logits[:, eot] += 10.0  # Make eot very likely
        new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())
        tokens = new_tokens
        sum_logprobs = torch.zeros_like(sum_logprobs)  # Reset for simplicity
        if completed:
            break

def test_large_memory_limit():
    # Test with tensors close to 100MB (but not exceeding)
    beam_size = 100
    n_audio = 5
    vocab_size = 200
    eot = 199
    inference = DummyInference()
    decoder = BeamSearchDecoder(beam_size, eot, inference)
    tokens = torch.zeros((n_audio * beam_size, 10), dtype=torch.long)
    logits = torch.randn((n_audio * beam_size, vocab_size))
    sum_logprobs = torch.zeros(n_audio * beam_size)
    new_tokens, completed = decoder.update(tokens, logits, sum_logprobs.clone())
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-BeamSearchDecoder.update-mayrurbj and push.