[Frontend] merge beam search implementations

[LunrEclipse]

Merged implementation of AsyncEngine and MQLLMEngine's beam_search into EngineClient(Protocol)

Manually testing conducted to verify that requests are still ran in parallel and output is correct.

server side:

$ vllm serve meta-llama/Meta-Llama-3-8B

client side:

Completion

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="key123",
)

prompt = "Capital of France is"

try:
    completion = client.completions.create(
        model="meta-llama/Meta-Llama-3-8B",
        prompt=prompt,
        max_tokens=4,
        extra_body={'use_beam_search': True, 'best_of': 3}
    )
    print(completion.choices[0].text)
except Exception as e:
    print(e)

Chat

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="key123",
)

prompt = "Capital of France is"

try:
    completion = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3-8B",
        messages = [
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=10,
        extra_body={'use_beam_search': True, 'best_of': 3, 'temperature': 0}
    )
    print(completion)
    print(completion.choices[0].message.content)
except Exception as e:
    print(e)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[njhill]

Thanks @LunrEclipse, this looks good to me. I think along with this we should change EngineClient from a Protocol to an ABC. It doesn't make much sense to have a method impl in an Protocol.

[nFunctor]

Since I tried to implement stop logic elsewhere I'd like to know why are we dealing with eos like that instead of putting ignore_eos into sampling params? Strictly speaking this is not the goal of the PR so feel free to ignore this comment. Thanks

[njhill]

@nFunctor yes this PR is just consolidating the existing logic, let's address that in your follow-on one.

[youkaichao]

@njhill thanks for shepherding this pr!
@LunrEclipse please address the review from @njhill .

I'll be afk and will hand it over to @njhill for review.

[LunrEclipse]

@njhill Thank you for the review! I've gone ahead and pushed changes based on your feedback

[njhill]

Thanks @LunrEclipse.

Not related to this PR specifically, but couldn't the beam search impl still be kept behind the EngineClient.generate API? I.e. we just intercept the existing beam_search and associated params in SamplingParams ... so that the outward function remains the same?

[LunrEclipse]

@njhill Yeah, it's definitely doable if we add logic to the EngineClient.generate methods to check if beam_search is true and then yield different results there than checking inside the engine itself.

[russellb]

@njhill I hadn't got to it yet, but FYI, mypy complains about this line. It's passing a float where it expects a dict, at least according to the typing.

[njhill]

Thanks @russellb yes this looks wrong! Though not really due to this PR which just moved/consolidated the existing logic.

Probably we should keep a logprobs list in BeamSearchSequence in addition to tokens, and set this.

I think this new external beam search impl still needs a bit more work in general.

[russellb]

Yeah, I knew you had just moved the code. I just wanted to highlight it in case it was a super quick fix for you. Thanks for sharing your thoughts! I'll probably get to it at some point as I keep hacking through the type checking. It seems pretty valuable since it's found multiple bugs in my digging so far!

[youkaichao]

Thanks @LunrEclipse.

Not related to this PR specifically, but couldn't the beam search impl still be kept behind the EngineClient.generate API? I.e. we just intercept the existing beam_search and associated params in SamplingParams ... so that the outward function remains the same?

it is possible following the spirit of #9302 . we can just create another BeamSearchSequenceGroup class. @LunrEclipse @njhill

[njhill]

@youkaichao yes I think we should refactor things a bit.. and also move the impl into beam_search.py instead of protocol.py, etc.