[Core] Cross-attention KV caching and memory-management (towards eventual encoder/decoder model support)

[afeldman-nm]

This PR is a step towards encoder/decoder model support. This PR (1) allows a SequenceGroup to be associated with 0 or 1 encoder sequences, and (2) causes an encoder/decoder model to leverage a separate "cross-attention KV cache" when performing decoder cross-attention.

To that end, "cross-attention block tables" are added to the block manager (v1 and v2), in order to enable separate memory-mapping and memory-paging for cross-attention KVs.

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. (This PR) Support cross-attention KV cache & memory management (allocate/swap/free) in block manager
2. Invoke cross-attention operation via the Attention wrapper
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

In order to support cross-attention KV cache & memory management, this PR:

• Adds an optional "encoder sequence" member to SequenceGroup. A SequenceGroup may be associated with 0 or 1 encoder sequences.
• For each SequenceGroup that has an encoder sequence, block manager stores a "cross-attention block table", identified by the SequenceGroup's request id. The cross-attention block table maps the logical blocks of cross-attention KV cache to physical memory blocks
• Block manager allocate/swap/free methods detect when a SequenceGroup has an encoder sequence, and respond by incorporating the cross-attention block table into allocate/swap/free operations
• Adds encoder/decoder scenario unit tests for block manager v1/v2, under the tests/core (block manager v1) and tests/core/block (block manager v2) directories

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 scheme described above, requires that each SequenceGroup instance has a globally unique request_id, which we believe to be the case.

Note 3: 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!

[afeldman-nm]

FYI to reviewer - my PR is failing the buildkite/ci/pr/amd-distributed-tests test, with what appears to be a HuggingFace issue:

=========================== short test summary info ============================
FAILED distributed/test_chunked_prefill_distributed.py::test_models[16-5-half-meta-llama/Llama-2-7b-hf] - OSError: You are trying to access a gated repo.
Make sure to have access to it at https://huggingface.co/meta-llama/Llama-2-7b-hf.
401 Client Error. (Request ID: Root=1-66454838-0c0ff2125401d04d55e2ccb8;5098fddc-b480-4fae-9875-53c8cb4529bf)
Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-2-7b-hf/resolve/main/config.json.
Access to model meta-llama/Llama-2-7b-hf is restricted. You must be authenticated to access it.

This looks like a HuggingFace issue, i.e. not something I can fix. Is it possible to move forward with the PR review process in spite of this test failure?

[sroy745]

Thanks for the pr. LGTM

[joerunde]

Leaving my gray check of approval as well, thanks for taking the time to go through all the feedback!

[maxdebayser]

Thanks for addressing all review comments. LGTM

[njhill]

Thanks @afeldman-nm for the great work and awesome PR description/explanation.

And thanks for all the great reviews from others.

I just have some minor style suggestions

[njhill]

@afeldman-nm it looks like there's still a formatting update needed for yapf.

[njhill]

@afeldman-nm it looks like there's still a formatting update needed for yapf.