[Frontend] Optimize beam search performance by limiting concurrency

[heheda12345]

Purpose

Current beam search will redo the prefill of all tokens for the generation of each token if the prefix cache cannot hold all tokens. This PR introduces a concurency_limit so that people can set it to a small enough value to achieve prefix cache hit

This PR also fixes a small bug in benchmark_throughput of beam search

Test Plan

python3 benchmark_throughput.py --model 'google/gemma-3-12b-it' --input-len 10000 --output-len 1000 --num-prompts 60 --n 2

Test Result

On h100, with concurrency 10 (about 10 min):

Throughput: 0.07 requests/s, 806.72 total tokens/s, 73.35 output tokens/s
Total num prompt tokens:  599907
Total num output tokens:  60000

Without concurrency optimization, I wait for 1 hour and the benchmark script doesn't finish.

---

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.

[gemini-code-assist]

Code Review

This pull request introduces a concurrency_limit parameter to the beam search functionality to optimize performance, especially when dealing with long prompts that may not fit into the prefix cache. The changes involve batching requests within the beam search loop. My review focuses on the correctness and maintainability of these changes. I've identified a consistent typo in the new API parameter concurency_limit which should be concurrency_limit. I've also found a case of variable shadowing that could impact code readability and maintainability. Both issues are of high severity and should be addressed.

[gemini-code-assist]

There's a typo in the new parameter concurency_limit. It should be concurrency_limit. This typo appears consistently across the PR in function arguments, variable names, and documentation. Please correct it everywhere for clarity and correctness, including in tests/conftest.py and tests/samplers/test_beam_search.py.

Suggested change

	concurency_limit: Optional[int] = None,
	concurrency_limit: Optional[int] = None,

[gemini-code-assist]

The loop variable i on this line shadows the i from the outer loop at line 721 (for i in range(0, len(prompts), concurency_limit):). This is generally bad practice as it can lead to confusion and potential bugs. Please use a more descriptive and non-conflicting name for the inner loop variable, for example beam_idx.

Suggested change

	for i in range(start, end):
	current_beam = all_beams[i]
	result = output[i]
	for beam_idx in range(start, end):
	current_beam = all_beams[beam_idx]
	result = output[beam_idx]

[DarkLight1337]

Can you update from main brance to fix the CI failure in entrypoints tests? Meanwhile the samplers test looks related to this PR