Optimizations for beam_search by mgoin · Pull Request #29133 · vllm-project/vllm

mgoin · 2025-11-21T00:33:57Z

Purpose

Set detokenize=False in SamplingParams during the beam search since we manually detokenize at the end
Add skip_clone parameter to SamplingParams to skip the defensive clone in processor when trusted
Slightly tweak the code in append_logprobs_for_next_position to boost perf for default case (using FlatLogprobs was slower due to the amount of indexing needed)
Cache tokenizer related variables that trigger the tokenizer such as get_eos_token_id. These were tokenizing the eos text in every preprocess call for each request
Cache is_encoder_decoder as it is expensive to read through the model's config
Add a special case for kv_cache_manager::get_computed_blocks where we reuse the previous result when consecutive requests are the same, allowing us to skip find_longest_cache_hit for the rest of the requests

Benchmark command for beam_search with n=30 on H100:

vllm serve Qwen/Qwen3-8B --max-logprobs 100
python benchmarks/benchmark_beam_search_server.py \
    --base-url http://localhost:8000 \
    --beam-width 30 \
    --max-tokens 8 \
    --input-len 128 \
    --num-requests 10

Before:

Average latency: 282.96ms
Median latency: 282.68ms
Std deviation: 10.59ms
Min latency: 263.52ms
Max latency: 299.33ms

After

Average latency: 223.68ms
Median latency: 226.31ms
Std deviation: 15.06ms
Min latency: 200.60ms
Max latency: 243.34ms

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Signed-off-by: mgoin <mgoin64@gmail.com>

gemini-code-assist

Code Review

This pull request introduces several well-reasoned performance optimizations, primarily targeting beam search. The changes include adding various caching mechanisms, such as using cached_property, a new cache_if_not_none decorator, and memoizing results in KVCacheManager, as well as avoiding unnecessary work like cloning SamplingParams. Overall, these are solid improvements. However, I've identified a significant thread-safety issue in the new cache_if_not_none decorator that could lead to race conditions in a concurrent environment. I have provided a detailed comment with a suggested fix for this issue.

Signed-off-by: mgoin <mgoin64@gmail.com>

hmellor

LGTM, just one micro-optimisation nit

hmellor · 2025-11-21T12:09:27Z

    # into a dictionary twice is the same as doing it once.
-    topk_ranks = range(1, num_logprobs + 1)
-    ranks = itertools.chain((rank,), topk_ranks)
+    ranks = [rank] + list(range(1, num_logprobs + 1))


This is about 1.4x faster if we want to micro-optimise

Suggested change

ranks = [rank] + list(range(1, num_logprobs + 1))

ranks = [rank, *range(1, num_logprobs + 1)]

As the construction of ranks are fixed, I was thinking about simply get rid of this list, and just pass in the sampled rank here.

Then we could just do

# insert sampled token ... # insert top ranked tokens in order for rank, (token, logprob) in enumerate(zip(tokens[1:], logprobs[1:]))

heheda12345 · 2025-11-21T19:47:53Z

+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Simple benchmark script to test beam search performance via OpenAI API.


I think we have can benchmark beam search with vllm bench latency and vllm bench throughput

I didn't realize, will try to use it

heheda12345 · 2025-11-21T19:56:15Z

+        if (
+            request.block_hashes
+            and self._cache_hit_cache is not None
+            and max_cache_hit_length == self._cache_hit_cache.max_cache_hit_length


Does this constraint mean we need to recompute find_longest_cache_hit for each beam search iteration?

Yes, each beam search iteration submits n identical requests to the engine core and the core is unaware of this

heheda12345 · 2025-11-21T19:58:23Z

+            and self._cache_hit_cache is not None
+            and max_cache_hit_length == self._cache_hit_cache.max_cache_hit_length
+            and len(request.block_hashes) == self._cache_hit_cache.num_block_hashes
+            and request.block_hashes[-1] == self._cache_hit_cache.last_block_hash


does this mean the optimization works only when only the last partial block of the current request and the previous request is different and has no effect if different candidates have different prefix at the beginning? Do you know the cache hit rate of this optimization?

The intent of this optimization is just to deal with the case when the current request is identical to the previous request. My assumption was that because a given block_hash "includes" all the previous blocks, I could just look at the last one to check for equality. This is a huge boost to beam search n=30 because all requests perform the same search

Here are profiles showing that "Self duration" of schedule decreases from 6.060ms to 1.471ms

Before:

After:

heheda12345 · 2025-11-21T19:59:04Z

-            self.coordinator.find_longest_cache_hit(
-                request.block_hashes, max_cache_hit_length
+
+        # Optimization: Reuse previous result when consecutive requests are the same.


can you make this optimization optional to beam search? e.g., add a flag?

@mgoin @heheda12345 I thought about it before, and would love to see if it's possible. For find_longest_cache_hit, should we do binary search instead of linear search?

In this way, maybe we could consolidate the code without special handling for beam search?

heheda12345 · 2025-11-21T20:05:45Z

+            and max_cache_hit_length == self._cache_hit_cache.max_cache_hit_length
+            and len(request.block_hashes) == self._cache_hit_cache.num_block_hashes
+            and request.block_hashes[-1] == self._cache_hit_cache.last_block_hash
+        ):


We can reuse the computed_blocks of the previous request only when we can ensure these blocks are not evicted.

If the model only has full attention, I guess it's fine but want some proof on it

If the model has sliding window attention, I feel there may be some problem as the previous request may free the early blocks outside the sliding window and make these blocks evocable.

heheda12345 · 2025-11-21T20:10:28Z

We are using prefix caching to do beam search. Does it mean with block_size 16 and prompt length 30, we always need to recompute the prefill of 14 tokens?

njhill · 2025-11-22T04:57:31Z

        data that is expensive to copy. However, if not copied, the processor
        needs to support parallel decoding for multiple sequences
        See https://github.com/vllm-project/vllm/issues/3087
+


We can also cleanup the logitsprocessor stuff in this method since they aren't supported in SamplingParams anymore

njhill · 2025-11-22T05:05:53Z

            max_tokens=1,
            temperature=temperature,
+            detokenize=False,  # We detokenize the output after the search
+            skip_clone=True,  # Safe to reuse params without cloning


I think we could set this for all of the API server endpoints

Jialin · 2025-12-10T19:41:51Z

    # into a dictionary twice is the same as doing it once.
-    topk_ranks = range(1, num_logprobs + 1)
-    ranks = itertools.chain((rank,), topk_ranks)
+    ranks = [rank] + list(range(1, num_logprobs + 1))


As the construction of ranks are fixed, I was thinking about simply get rid of this list, and just pass in the sampled rank here.

Then we could just do

# insert sampled token ... # insert top ranked tokens in order for rank, (token, logprob) in enumerate(zip(tokens[1:], logprobs[1:]))

Jialin · 2025-12-10T19:44:02Z

+    skip_clone: bool = False
+    """Internal flag indicating that this SamplingParams instance is safe to
+    reuse without cloning. When True, clone() will return self without
+    performing a deep copy. This should only be set when the params object
+    is guaranteed to be dedicated to a single request and won't be modified
+    in ways that would affect other uses."""


IIUC, this should be helpful for n-gen and beam search, correct?

And just curious, what's the risk to make it default to true? Or even make it as a consistent behaviour?

Jialin · 2025-12-10T19:46:24Z

-            self.coordinator.find_longest_cache_hit(
-                request.block_hashes, max_cache_hit_length
+
+        # Optimization: Reuse previous result when consecutive requests are the same.


@mgoin @heheda12345 I thought about it before, and would love to see if it's possible. For find_longest_cache_hit, should we do binary search instead of linear search?

In this way, maybe we could consolidate the code without special handling for beam search?

Jialin · 2025-12-10T19:49:29Z

            logprobs=logprobs_num,
            max_tokens=1,
            temperature=temperature,
+            detokenize=False,  # We detokenize the output after the search


+1 on this. I'm wondering if we should directly remove tokenized string in Logprob. IMO,

detokenization should never live inside Engine.

clients are rarely used the data, if they do, it should be cleaner to do it on their side instead.

ducviet00 · 2025-12-23T16:47:46Z

Hi @mgoin, I’m currently using vLLM with beam search and was wondering if there are any updates on the status of this PR? I really appreciate all the work you're doing on this!!

github-actions · 2026-03-25T02:16:54Z

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

github-actions · 2026-04-24T02:18:13Z

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!

mgoin added 2 commits November 20, 2025 19:29

[Perf] Disable detokenization during beam_search

6c9d1ad

Signed-off-by: mgoin <mgoin64@gmail.com>

More optimizations for beam search

368becd

Signed-off-by: mgoin <mgoin64@gmail.com>

mgoin requested review from ApostaC, ProExpertProg, WoosukKwon, aarnphm, alexm-redhat, chaunceyjiang, heheda12345, hmellor, houseroad, njhill, robertgshaw2-redhat, tlrmchlsmth, yewentao256, youkaichao and ywang96 as code owners November 21, 2025 00:33

mergify Bot added frontend v1 labels Nov 21, 2025

Add benchmark script

3b0f07c

Signed-off-by: mgoin <mgoin64@gmail.com>

gemini-code-assist Bot reviewed Nov 21, 2025

View reviewed changes

Comment thread vllm/inputs/preprocess.py

mergify Bot added the performance Performance-related issues label Nov 21, 2025

Bench

ba40a10

Signed-off-by: mgoin <mgoin64@gmail.com>

hmellor approved these changes Nov 21, 2025

View reviewed changes

heheda12345 reviewed Nov 21, 2025

View reviewed changes

njhill reviewed Nov 22, 2025

View reviewed changes

Jialin reviewed Dec 10, 2025

View reviewed changes

AndreasKaratzas mentioned this pull request Dec 31, 2025

[CI Failure]: mi325_1: Samplers Test #29521

Closed

3 tasks

shireen-bean mentioned this pull request Jan 28, 2026

[Feature]: Optimized Beam Search Performance for TPUs vllm-project/tpu-inference#1551

Open

github-actions Bot added the stale Over 90 days of inactivity label Mar 25, 2026

github-actions Bot closed this Apr 24, 2026

This was referenced Jun 22, 2026

[Perf] Skip detokenization in online beam search GuyStone/vllm#1

Closed

[Perf] Skip detokenization in online beam search #46422

Merged

	ranks = [rank] + list(range(1, num_logprobs + 1))
	ranks = [rank, *range(1, num_logprobs + 1)]

Uh oh!

Uh oh!

Conversation

mgoin commented Nov 21, 2025 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

hmellor left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

heheda12345 commented Nov 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ducviet00 commented Dec 23, 2025

Uh oh!

github-actions Bot commented Mar 25, 2026

Uh oh!

github-actions Bot commented Apr 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

mgoin commented Nov 21, 2025 •

edited by github-actions Bot

Loading