[PoC]: Support encode only models by Workflow Defined Engine

[noooop]

PTAL #8453

Briefly introduce

What new models need to be supported

These models are all from issues and are also very famous:

• xlm_roberta
• bge-m3
• bge-reranker-v2-m3
• bert
• bge v1.5 family
• Snowflake Arctic Embed (Family)
• gte-Qwen2
• This list is still growing

These models is roughly divided into three categories：

• Encode only models. (Bidirectional Transformers, causal=False), Often fine-tuned as retriever and reranker etc.
• Decode only models. (masked multi-head attention, causal=True). There are two interesting uses:
  • Output last hidden states as a feature extractor
  • Decode only retriever （I don't know of a better name），E.g. e5-mistral-7b （The only Embed model currently supported by vllm)
  • Whether it has been fine-tuned or not, there is almost no difference in the code.
• Enable bidirectional. LLM2Vec propose a simple unsupervised approach that can transform any decoder-only LLM into a strong text encoder.

What new features these new models have

What the above three categories have in common is that there is only the prefill stage. In order to make the terminology more precise, prefill only is used below.

You can think of prefill only as encode only fancy writing.

New features:

1. attention
   • Prefill only models requires simpler attention implementations, no need to consider kvcache, no decoding phase
   • We need to support enable_bidirectional flag manually or read hf config automatically, enable bidirectional.
2. scheduler
   • Prefill only models requires simpler scheduler, no need to consider kvcache and preemption
   • Prefill only models, there is no correlation between tasks, so it is easy to implement async scheduling
3. executer
   • In order to support async scheduling, model_input_builder needs to be separated from the runner.
   • The main thread executes scheduling and all CPU processing, and the gpu thread only executes h2d, execution model, d2h
   • If async scheduling and async execution are implemented, data parallelism is also easy to implement. Data parallelism is more efficient for small models

How engine Architecture needs to support these features flexibly and efficiently.

If we directly add new functions to existing modules, these modules are becoming increasingly complex, and sometimes new features must be compromised for compatibility. ultimately leading to suboptimal results

The most flexible and efficient way to support the prefill only models is to implement different modules for models of different architectures and load the required modules on demand.

I call this architecture Workflow Defined Engine, or WDE for short.

I divided the Engine into the following modules.

• InputProcessor: The llm models inputs strings, the reranker inputs pairs, and the multimodal model input is more complex...
• OutputProcessor: The retriever(embedding) models output embeddings, reranker models and classification models output Scores...
• ModelInputBuilder: Building model inputs and attention metadata
• AttnBackend: Support different AttnBackend and enable bidirectional
• Tokenizer: There may be different tokenizers
• Executor: Sync\Async\TP\PP\DP\maybe more
• Worker & runner: Support different devices\maybe more
• EngineArgs: Different models, different config may accept different parameters
• maybe more

With wde, there is no need for one module to be compatible with all functions. You can use the dynamic loading feature of python to load different modules at the highest level, for different models and different needs.

• Modules can be configured through Workflow, plug and play
• Flexibly support plug-ins, and developers can load their own modules.
• Workflow is really the best place to hide dirty codes.
  Some models cannot use the common Workflow. When you don’t know where to put the dirty code, you can always create a new workflow and link the model architecture to the new workflow to avoid leaving dirty code everywhere for the sake of compatibility.

Let's start splitting this pr and try to merge it into main

• wde core & bert model [Core]: (Last/N) Support prefill only models by Workflow Defined Engine #8964
  • [Core]: (2/N) Support prefill only models by Workflow Defined Engine - Prefill only attention #9124
  • [Core]: (1/N) Support prefill only models by Workflow Defined Engine - Prefill only scheduler #9181

[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.

🚀

[DarkLight1337]

This is a significant change from our current architecture. We'll consider incorporating this when we refactor our core framework.

[noooop]

As vllm supports more and more models and functions, they require different attention, scheduler, executor, and input output processor. . These modules are becoming increasingly complex, and sometimes new features must be compromised for compatibility. ultimately leading to suboptimal results

With wde, there is no need for one module to be compatible with all functions.
You can always use the workflow to load new modules at the highest level to support new functions.

I hope you like this new architecture @DarkLight1337

[DarkLight1337]

It looks nice for sure, but there are many abstractions that are difficult to adopt immediately. If you want us to use this architecture, I suggest that you split this PR up into smaller chunks and gradually refactor over the code base rather than doing everything all at once.

[noooop]

I want to do experiments on the encode only model and integrate it with the existing code only at entrypoints, so that the impact of modifications is minimal.

[noooop]

Do you have any better suggestions? @DarkLight1337

[noooop]

This PR is used as a demonstration. I can modify and resubmit a PR.

[noooop]

Don’t be too anxious
I will improve this PR and support more models and functions.
Until you find a suitable opportunity to merge it in

[DarkLight1337]

Do you have any better suggestions? @DarkLight1337

Not really, just do as you have said.

[DarkLight1337]

You should list out what are the features in your PR and how they correspond to #8779. Otherwise, people would have to read through your whole PR to understand what is going on.

[noooop]

I feel that although this PR has some similarities with #8779,

their focus is different and there is no way to compare their features one-to-one.

[DarkLight1337]

I feel that although this PR has some similarities with #8779, their focus is different and there is no way to compare their features one-to-one.

I'd say your PR focuses on tackling the second goal in #8779 (the new architecture will be extensible and modular). You should explain in detail how this is being achieved (in particular, what type of abstractions you are using?). That way, we can consider those aspects when planning how to refactor the existing code.

[noooop]

Many things are simple to write code but very complicated to explain.

Can you let them look at the code?

[DarkLight1337]

Many things are simple to write code but very complicated to explain.

Can you let them look at the code?

The thing is, people don't want to look at 10k lines of code to understand what is going on. If you want them to use this code, it is your responsibility to explain it.

[noooop]

OK, I'll try

[DarkLight1337]

Yeah, that should be good enough to start with.

Some models cannot use the common Workflow. When you don’t know where to put the dirty code, you can always create a new workflow and link the model architecture to the new workflow to avoid leaving dirty code everywhere for the sake of compatibility.

Probably need to address this. e.g. which types of models are not supported under this workflow? If it's a large category then we'll have to find a solution to it.

[DarkLight1337]

Looks better. You should also think about how to adopt this incrementally.

[DarkLight1337]

Sorry was afk, you can post it whenever you like.

[noooop]

@DarkLight1337

I think the first part supporting bert is pretty good, but it's still 6,000 lines of code

PTAL #8964

[noooop]

6000 lines of code support bert, I feel like a clown

[noooop]

@DarkLight1337

#9166 has many similarities with Workflow Defined Engine.

Can you invite @WoosukKwon to participate in the discussion of this PR?

[DarkLight1337]

I suggest you comment directly on his PR. You can also join our Slack workspace (see README) and ping him.