Support embedding models in V1 with a dedicated model_runner

[maxdebayser]

This is an alternative to #16188 . In that other PR, I implemented embedding model support on the same model runner as the decoder models. This had the advantage that the code changes were fairly minimal. The other advantage in my opinion is that a single model runner implementation is less likely to become stale as new features and bug fixes only need to be applied to one code base. However, there were concerns about the performance implications and code complexity of a single implementation that tries to handle all cases.

In this PR I started by reverting all changes to the GPUModelRunner and created a GPUPoolingModelRunner basically by deleting everything that was related to sampling. In this state it was already passing the embedding model unit tests but there was still a lot of duplicated or unnecessary code.

Now I'm finished with the refactoring. Basically there is now a GPUBaseModelRunner that contains the common code and the GPUModelRunner and the GPUPoolingModelRunner implement the missing pieces.

There were a few issues where @22quinn spent some time thinking about:

• kv-cache management: For encode models this is unnecessary because the attention mask is not causal and therefore optimizations such as chunked prefill and prefix caching don't apply. However, there are encoder models with pooling based in the last hidden state where these optimizations are applicable. One example is the intfloat/e5-mistral-7b-instruct that we use in the unit tests.
• handling of m-rope, sliding window, multi-modal...: these thing are mostly orthogonal to pooling or sampling and went into the abstract base class
• input batch management: when chunked prefill is disabled, in pooling models each request only stays in the batch for one execute_model call. However, with chunked prefill execute_model is called several times and the same logic that is used in the sampling models applies.
• cascade attention: this only applies to decoding.

cc: @mgoin, @WoosukKwon , @DarkLight1337

[maxdebayser]

There are some merge conflicts because of the KV cache group PR that was reverted, but it seems that it will be added again once the maintainer fixes the bugs that have been found after it was merged in main. So I'm going to wait a little bit before trying to solve the conflicts.

[maxdebayser]

The kv cache group PR was redone so I've fixed the merge conflicts.

[22quinn]

Any benefit using torch.Tensor instead of list[int]? Same for prompt_lens

[maxdebayser]

For prompt_lens it helps that the array is already in tensor format to do things like torch.cumsum(prompt_lens, dim=0). For prompt_token_ids there is no strong reason, but it allows us to reuse the same _make_prompt_token_ids_tensor() function as in the non-pooling case.

[22quinn]

Same q here - why convert to tensor?

[maxdebayser]

See answer above.

[22quinn]

is it possible for pooler_output to be None here?

[maxdebayser]

Yes, during chunked prefill.

[facebook-github-bot]

@22quinn has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maxdebayser.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @maxdebayser.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork