[Kernel] Correctly invoke prefill & decode kernels for cross-attention (towards eventual encoder/decoder model support)

[afeldman-nm]

This PR is a step towards encoder/decoder model support. This PR modifies the xFormers backend* such that (1) the attention impl can implement cross-attention, and (2) the attention metadata data structure can represent the necessary metadata for invoking cross-attention.

* FlashAttention backend support for encoder/decoder models is left as future work

---

A quick overview of the plan for supporting encoder/decoder models in vLLM:

• Architectural assumptions:
  • The encoder/decoder model comprises one non-autoregressive encoder module and one autoregressive decoder module.
  • A single inference call to the model consumes an encoder prompt and a decoder prompt. The model output is the result of decoder inference against the decoder prompt, conditional on the encoder hidden states which result from applying the encoder to the encoder prompt. The encoder hidden states are not part of the overall model output
  • Thus, encoder inference is a prerequisite for decoder inference. The decoder consumes encoder hidden states via cross-attention, which is not present in decoder-only models.
  • It is assumed that these architectural details are handled inside the model definition; however, to support the inference process for such models, vLLM core must be changed to accommodate cross-attention
• Encoder/decoder inference process & cross-attention:

  • Prefill phase: (1) Non-autoregressive encoder inference yields encoder hidden states in a single pass; no KV caching occurs. (2) decoder prefill yields first-token-prediction & cached KVs. Within the decoder, cross-attention layers cache the KVs derived from encoder hidden states:

    • Key_{cross-attn, layer-n} = W_{K, cross-attn, layer-n} x (Encoder hidden states)

    • Value_{cross-attn, layer-n} = W_{V, cross-attn, layer-n} x (Encoder hidden states)

    • Note that all cross-attention layers consume the same encoder hidden states; however each cross-attention layers' keys and values differ because each layer has unique W_{K, cross-attn, layer-n} and W_{V, cross-attn, layer-n}. Therefore, the cross-attention KV cache must store KVs for each decoder layer, even though these KVs are all derived from a single set of encoder hidden states.

    • Note that self-attention layer behavior is unchanged compared to what it would be in a decoder-only model (cache KVs computed from the previous decoder layer outputs.)

  • Decode phase: during each iteration of the autoregressive decode process,

    • Each self-attention layer appends the last predicted token's KVs to the KV cache, and then utilizes cached KVs for next-token prediction (again, this is unchanged compared to a decoder-only model)
    • Each cross-attention layer has read-only access to cross-attention KVs, to use for next-token prediction. The cross-attention KV cache is never modified after prefill

---

To implement the above encoder/decoder inference process, the following functionality will be added to vLLM over the course of multiple PRs:

1. Support cross-attention KV cache & memory management (allocate/swap/free) in block manager
   • See PR: [Core] Cross-attention KV caching and memory-management (towards eventual encoder/decoder model support) #4837
2. (This PR) Invoke cross-attention operation via the Attention wrapper & Attention metadata data structure
3. Modify ModelRunner to construct input tensors & Attention metadata structures for cross-attention
4. Small changes to LLM engine & scheduler so that vLLM requests can include an encoder input prompt

Note 1: because this PR makes an incremental contribution (cross-attention KV-caching and memory management), this PR will not enable end-to-end encoder/decoder support (this will rely on later PRs.)

Note 2: the best effort is being made to ensure that encoder/decoder models are compatible with existing vLLM features. At this time, encoder/decoder models are unlikely to be compatible with the following vLLM features:

• Speculative decoding
• Chunked prefill
• Automatic prefix caching
• Sliding window
• Flash attention
• CUDA graph

INCREMENTAL FIX TOWARDS #187

BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE

---

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Model] for adding a new model or improving an existing model. Model name should appear in the title.
• [Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
• [Kernel] for changes affecting CUDA kernels or other compute kernels.
• [Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
• [Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• We adhere to Google Python style guide and Google C++ style guide.
• Pass all linter checks. Please use format.sh to format your code.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
• Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

• After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
• After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
• After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
• Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

[WoosukKwon]

@afeldman-nm Sorry for the delay. Let me take a look at the PR this afternoon.

[WoosukKwon]

@afeldman-nm Thanks for the PR.

Overall, I'm OK with this change. The code in xformers.py is a bit difficult to follow, but it's understandable given the inherent complexity.

The only thing I'd like to ask is to avoid using \ in the code as guided in https://google.github.io/styleguide/pyguide.html#32-line-length Please use parentheses if the line overflows. Also, please considering adding , to the last argument of the methods (particularly in the forward method of the attention backends) if that reduces the line changes.

[afeldman-nm]

@afeldman-nm Thanks for the PR.

Overall, I'm OK with this change. The code in xformers.py is a bit difficult to follow, but it's understandable given the inherent complexity.

The only thing I'd like to ask is to avoid using \ in the code as guided in https://google.github.io/styleguide/pyguide.html#32-line-length Please use parentheses if the line overflows. Also, please considering adding , to the last argument of the methods (particularly in the forward method of the attention backends) if that reduces the line changes.

@WoosukKwon Thanks, I have addressed your feedback (both the specific points of feedback as well as the more general styleguide suggestsions.)

I agree it is hard to make xformers comprehensible given the complexity, perhaps this can wait for a future PR.

Please let me know if you have any other changes to recommend, or if the fixes look satisfactory.

[afeldman-nm]

Thanks @sroy745 @maxdebayser @njhill @WoosukKwon @robertgshaw2-neuralmagic for all of your excellent help in landing this PR!