[Frontend] API support for beam search for MQLLMEngine

[LunrEclipse]

Added support for running a higher layer of beam search across openai API server with frontend multiprocessing.

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.

🚀

[youkaichao]

and also update

vllm/tests/entrypoints/openai/test_completion.py

Line 518 in 151ef4e

except BadRequestError as e:

[youkaichao]

this is great! left some comments.

[robertgshaw2-neuralmagic]

Is beam search supported with LoRA?

It seems like it does not work through the AsyncLLMEngine since we aren't passing the lora_request parameter there. So I think you could collapse these two cases since lora_request will always be none when passed to MQLLMEngine

[robertgshaw2-neuralmagic]

--disable-frontend-multiprocessing will no longer resolve this error. There also is no other case here so this could be an assert.

[robertgshaw2-neuralmagic]

On second thought, you don't actually need the assert

You can beam_search to the EngineClientProtocol

[LunrEclipse]

I don't think the base EngineClientProtocol has a beam_search function. I can add it in though.

[robertgshaw2-neuralmagic]

It does not. But the EngineClientProtocol defines the behavior of AsyncLLMEngine and MQLLMEngine. So now that both support it, you can expand EngineClientProtocol to include the beam_search api

[youkaichao]

I cannot find EngineClientProtocol . does it exist now? @robertgshaw2-neuralmagic

[LunrEclipse]

I think it's this?

[youkaichao]

okay, it should be EngineClient class. but MQLLMEngineClient does not inherit from EngineClient . we can make it a future step to absorb beam search implementation into the EngineClient .

[robertgshaw2-neuralmagic]

EngineClient is a protocol. MQLLMEngine should inherit from this. If it doesn’t, I’ll submit a PR to make it (since we support the full API). On train so AFK

Either way, we are about to collapse MQLLMEngine and AsyncLLMEngine once we have PP working, so the concept of an EngineClient will be removed once this is done

[youkaichao]

sounds good. I'll go ahead and merge this pr after it is ready. and after you make MQLLMEngine inherit from EngineClient, we can merge separate beam search implementation in one place.

[njhill]

Given that it's currently a Protocol, MQLLMEngine technically doesn't need to subclass it directly. But it would probably be good to anyhow and actually we could consider changing it to an ABC instead.

I agree with @robertgshaw2-neuralmagic that this method should just be added to EngineClient though and we should not need these type assertions.

Not directly related to this PR but I also think we should consider renaming it to something like AsyncEngineClient, and have a way to obtain an instance of an AsyncEngineClient which doesn't involve explicit construction. And that would replace explicit use of AsyncLLMEngine.

[youkaichao]

changes are welcome on this part!

[robertgshaw2-neuralmagic]

I would add some sort of kwarg to show why it is None here

[youkaichao]

Suggested change

	def create_sort_beams_key_function(tokenizer, length_penalty):
	def create_sort_beams_key_function(eos_token_id: int, length_penalty: float):

[youkaichao]

what's the difference between request_id and request_id_item ?

[LunrEclipse]

Ahh yeah it's supposed to be request_id_item, I was a bit careless when I was doing some refactoring. I think completion supports multiple prompts so each prompt has it's own request_id_item for a general request_id.

[youkaichao]

move them to vllm/sequence.py ?

[LunrEclipse]

There's a circular import error if I put these classes in vllm/sequence.py as BeamSearchSequence is needed in vllm/utils.py, but vllm/sequence.py indirectly imports from vllm/utils.py

[youkaichao]

logically, vllm/utils.py should not import vllm/sequence.py . we should change the code if this is the case.

[LunrEclipse]

should i also move the create_sort_beams_key_function to vllm/sequence.py? That will solve the issue

[youkaichao]

sure, go ahead!

[njhill]

I'm not sure whether these should go in sequence.py since that holds "internal" data-structures used within the scheduler etc. and if I understand correctly, BeamSearchSequence etc. are only used in the outer layer(s). Maybe better to have a dedicated file an the appropriate place in the tree for this?

[youkaichao]

Maybe better to have a dedicated file an the appropriate place in the tree for this?

makes sense. how about vllm/beam_search.py ?

[youkaichao]

code LGTM, and we can merge after all tests pass. thanks again for the great work @LunrEclipse !

[njhill]

I added some comments, might be good to address those first (unless you guys disagree with them).

[youkaichao]

DM with @njhill , it is good to merge now. we can also add followup pr if needed.

[nFunctor]

Hi, I'm currently testing this new beam search with the nightly build.

• The logs in the server are quite extensive, is this intended? They are not suppressed by --disable-log-requests
• If I understand correctly, the num_beams is now the n of completion request, and not the "best_of" that used to be in extra-body? Your example uses best_of of the extra body but when I read the code for BeamSearchParams it seems to rely on n.

Sorry if those are to be addressed anyway, I am trying to figure out how it works before we get a new vllm release. Thanks!

PS should have probably asked it in #9087 , sorry, still that's your PR as well

[youkaichao]

The logs in the server are quite extensive, is this intended? They are not suppressed by --disable-log-requests

every beam search request will join and leave the request queue again and again. so it is expected. using --disable-log-requests is recommended.

If I understand correctly, the num_beams is now the n of completion request, and not the "best_of" that used to be in extra-body?

I think so. @LunrEclipse can you confirm it?

[nFunctor]

@youkaichao well, what I meant is that --disable-log-requests does not help at all. I think what I see is logging.info(output) that is placed after output = [x[0] for x in output] in the code.

I am trying to make a PR for stop logic where I suppress this logging in passing.