@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 37e8182bfe
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[v1] Add Whisper model support (encoder-decoder) (#21088)

Signed-off-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: NickLucche <nlucches@redhat.com>

Russell Bryant committed Sep 10, 2025 at 16:53 UTC 37e8182bfe155f9decd1277471f1b1b238d50c70
31 files changed +429 -92
.buildkite/test-pipeline.yaml
+3 -3
@@ -321,7 +321,6 @@ steps:
321 - python3 offline_inference/vision_language_pooling.py --seed 0
322 - python3 offline_inference/vision_language_multi_image.py --seed 0
323 - VLLM_USE_V1=0 python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
324 - - python3 offline_inference/encoder_decoder.py
324 - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
325 - python3 offline_inference/basic/classify.py
326 - python3 offline_inference/basic/embed.py
@@ -644,7 +643,7 @@ steps:
643 - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
644 - pip freeze | grep -E 'torch'
645 - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
647 - - cd .. && pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
646 + - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
647
648 - label: Multi-Modal Models Test (Extended) 1
649 mirror_hardwares: [amdexperimental]
@@ -818,7 +817,8 @@ steps:
817 # Avoid importing model tests that cause CUDA reinitialization error
818 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
819 - pytest models/language -v -s -m 'distributed(num_gpus=2)'
821 - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)'
820 + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
821 + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
822 # test sequence parallel
823 - pytest -v -s distributed/test_sequence_parallel.py
824 # this test fails consistently.
examples/offline_inference/encoder_decoder.py
+2
@@ -5,6 +5,8 @@ Demonstrate prompting of text-to-text
5 encoder/decoder models, specifically BART and mBART.
6
7 This script is refactored to allow model selection via command-line arguments.
8 +
9 +NOTE: This example is not yet supported in V1.
10 """
11
12 import argparse
examples/offline_inference/encoder_decoder_multimodal.py
+3
@@ -5,6 +5,7 @@ This example shows how to use vLLM for running offline inference with
5 the explicit/implicit prompt format on enc-dec LMMs for text generation.
6 """
7
8 +import os
9 import time
10 from collections.abc import Sequence
11 from dataclasses import asdict
@@ -130,6 +131,8 @@ def run_mllama():
131
132
133 def run_whisper():
134 + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
135 +
136 engine_args = EngineArgs(
137 model="openai/whisper-large-v3-turbo",
138 max_model_len=448,
tests/encoder_decoder/test_e2e_correctness.py
+1
@@ -63,6 +63,7 @@ def clear_cache():
63 current_platform.is_cpu(),
64 reason="CPU backend is not currently supported with encoder/decoder models"
65 )
66 +@pytest.mark.skip(reason="bart not supported in V1")
67 def test_encoder_decoder_e2e(
68 hf_runner,
69 vllm_runner,
tests/entrypoints/openai/test_encoder_decoder.py
+1
@@ -30,6 +30,7 @@ async def client(server):
30
31 @pytest.mark.asyncio
32 @pytest.mark.parametrize("model_name", [MODEL_NAME])
33 +@pytest.mark.skip(reason="bart is not yet supported in V1")
34 async def test_single_completion(client: openai.AsyncOpenAI, model_name: str):
35 completion = await client.completions.create(model=model_name,
36 prompt="Hello, my name is",
tests/models/language/generation/test_bart.py
+2
@@ -178,6 +178,7 @@ def run_test(
178 @pytest.mark.parametrize("max_tokens", [64])
179 @pytest.mark.parametrize("num_logprobs", [5])
180 @pytest.mark.parametrize("decoder_prompt_type", list(DecoderPromptType))
181 +@pytest.mark.skip(reason="bart not supported in V1")
182 def test_models(hf_runner, vllm_runner, example_encoder_decoder_prompts, model,
183 dtype, max_tokens, num_logprobs, decoder_prompt_type) -> None:
184
@@ -201,6 +202,7 @@ def test_models(hf_runner, vllm_runner, example_encoder_decoder_prompts, model,
202 @pytest.mark.parametrize("max_tokens", [64])
203 @pytest.mark.parametrize("num_logprobs", [5])
204 @pytest.mark.parametrize("decoder_prompt_type", [DecoderPromptType.CUSTOM])
205 +@pytest.mark.skip(reason="bart not supported in V1")
206 def test_models_distributed(hf_runner, vllm_runner,
207 example_encoder_decoder_prompts,
208 distributed_executor_backend, model, dtype,
tests/models/multimodal/generation/test_whisper.py
+1 -2
@@ -122,8 +122,7 @@ def run_test(
122
123
124 @pytest.mark.core_model
125 -@pytest.mark.parametrize(
126 - "model", ["openai/whisper-small", "openai/whisper-large-v3-turbo"])
125 +@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
126 @create_new_process_for_each_test()
127 def test_models(vllm_runner, model) -> None:
128 run_test(
tests/models/multimodal/processing/test_tensor_schema.py
+1
@@ -31,6 +31,7 @@ from ...utils import dummy_hf_overrides
31
32 ARCH_TO_SKIP = {
33 "MolmoForCausalLM": "incompatible requirements",
34 + "Florence2ForConditionalGeneration": "not supported in V1",
35 }
36 ARCH_NEEDS_EXTRAS = [
37 "InternVLChatModel",
tests/models/test_initialization.py
+6
@@ -68,6 +68,12 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
68 # has cc==8.9 which hasn't supported FA3 yet. Remove this hack when
69 # L4 supports FA3.
70 m.setenv("VLLM_ATTENTION_BACKEND", "TRITON_ATTN_VLLM_V1")
71 + if model_arch == "Florence2ForConditionalGeneration":
72 + # An encoder-decoder model that's V0-only. Just skip it
73 + # since V0 is about to be removed.
74 + pytest.skip("Skipping Florence2ForConditionalGeneration")
75 + if model_arch == "WhisperForConditionalGeneration":
76 + m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
77 LLM(
78 model_info.default,
79 tokenizer=model_info.tokenizer,
tests/v1/test_oracle.py
-1
@@ -10,7 +10,6 @@ from vllm.engine.arg_utils import AsyncEngineArgs
10 from vllm.engine.async_llm_engine import AsyncLLMEngine
11
12 UNSUPPORTED_MODELS_V1 = [
13 - "openai/whisper-large-v3", # transcription
13 "facebook/bart-large-cnn", # encoder decoder
14 ]
15
vllm/attention/layers/cross_attention.py new
+160
@@ -0,0 +1,160 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +from copy import copy
5 +from typing import Optional
6 +
7 +import numpy as np
8 +import torch
9 +from transformers import CacheConfig
10 +
11 +from vllm import envs
12 +from vllm.attention.backends.abstract import (AttentionBackend,
13 + AttentionMetadata, AttentionType)
14 +from vllm.attention.layer import Attention
15 +from vllm.attention.selector import get_attn_backend
16 +from vllm.config import VllmConfig
17 +from vllm.logger import init_logger
18 +from vllm.multimodal import MULTIMODAL_REGISTRY
19 +from vllm.utils import cdiv
20 +from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
21 + subclass_attention_backend)
22 +from vllm.v1.kv_cache_interface import CrossAttentionSpec
23 +
24 +logger = init_logger(__name__)
25 +
26 +
27 +def _get_max_encoder_len(vllm_config: VllmConfig) -> int:
28 + return MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(
29 + vllm_config.model_config)
30 +
31 +
32 +def _get_cross_slot_mapping(encoder_seq_lens: np.ndarray,
33 + block_table_tensor: torch.Tensor,
34 + kv_cache_spec: CrossAttentionSpec,
35 + device: torch.device) -> torch.Tensor:
36 + """Get cross-attention slot mappings."""
37 +
38 + block_size = kv_cache_spec.block_size
39 + slot_mappings = []
40 +
41 + # Find indices with non-zero encoder sequence lengths
42 + # The majority of parallel requests will be running the
43 + # decoder, so this list should be relatively small.
44 + active_indices = np.nonzero(encoder_seq_lens)[0]
45 +
46 + for req_index in active_indices:
47 + encoder_seq_len = encoder_seq_lens[req_index].item()
48 +
49 + # Calculate the number of blocks needed for this request
50 + num_blocks_needed = cdiv(encoder_seq_len, block_size)
51 +
52 + # Get the block IDs for this request from the tensor
53 + req_block_ids = block_table_tensor[req_index]
54 +
55 + # Get only the blocks we need (first num_blocks_needed blocks)
56 + needed_block_ids = req_block_ids[:num_blocks_needed]
57 +
58 + # All needed blocks are allocated
59 + i_values = torch.arange(encoder_seq_len,
60 + dtype=torch.int64,
61 + device=device)
62 + block_indices = i_values // block_size
63 + block_offsets = i_values % block_size
64 + block_numbers = needed_block_ids[block_indices]
65 + slot_mapping = block_numbers * block_size + block_offsets
66 +
67 + slot_mappings.append(slot_mapping)
68 +
69 + if slot_mappings:
70 + return torch.cat(slot_mappings)
71 + else:
72 + return torch.empty(0, dtype=torch.int64, device=device)
73 +
74 +
75 +@functools.lru_cache
76 +def create_cross_attention_backend(
77 + underlying_attn_backend: AttentionBackend, ) -> type[AttentionBackend]:
78 + prefix = "CrossAttention_"
79 + underlying_builder = underlying_attn_backend.get_builder_cls()
80 +
81 + class CrossAttentionBuilder(underlying_builder): # type: ignore
82 +
83 + def build(self,
84 + common_prefix_len: int,
85 + common_attn_metadata: CommonAttentionMetadata,
86 + fast_build: bool = False) -> AttentionMetadata:
87 + new_metadata = copy(common_attn_metadata)
88 + new_metadata.causal = False
89 + max_encoder_len = _get_max_encoder_len(self.vllm_config)
90 + new_metadata.max_seq_len = max_encoder_len
91 +
92 + new_metadata.seq_lens = torch.full(
93 + (new_metadata.num_reqs, ),
94 + max_encoder_len,
95 + dtype=torch.int32,
96 + device=self.device,
97 + )
98 + new_metadata.seq_lens_cpu = torch.full(
99 + (new_metadata.num_reqs, ),
100 + max_encoder_len,
101 + dtype=torch.int32,
102 + device="cpu",
103 + )
104 + new_metadata.slot_mapping = _get_cross_slot_mapping(
105 + new_metadata.encoder_seq_lens, new_metadata.block_table_tensor,
106 + self.kv_cache_spec, self.device)
107 + return super().build(common_prefix_len, new_metadata, fast_build)
108 +
109 + attn_backend = subclass_attention_backend(
110 + name_prefix=prefix,
111 + attention_backend_cls=underlying_attn_backend,
112 + builder_cls=CrossAttentionBuilder)
113 +
114 + return attn_backend
115 +
116 +
117 +class CrossAttention(Attention):
118 + """
119 + Cross-attention for encoder-decoder models.
120 + Handles attention between decoder queries and encoder keys/values.
121 + """
122 +
123 + def __init__(self,
124 + num_heads: int,
125 + head_size: int,
126 + scale: float,
127 + cache_config: Optional[CacheConfig] = None,
128 + attn_type: Optional[str] = None,
129 + **kwargs):
130 + dtype = torch.get_default_dtype()
131 +
132 + if cache_config is not None:
133 + kv_cache_dtype = cache_config.cache_dtype
134 + block_size = cache_config.block_size
135 + else:
136 + kv_cache_dtype = "auto"
137 + block_size = 16
138 +
139 + if envs.VLLM_USE_V1:
140 + underlying_attn_backend = get_attn_backend(head_size, dtype,
141 + kv_cache_dtype,
142 + block_size)
143 +
144 + attn_backend = create_cross_attention_backend(
145 + underlying_attn_backend)
146 + else:
147 + # in v0 cross attention is handled inside the backends
148 + attn_backend = None
149 +
150 + if attn_type is not None:
151 + assert attn_type == AttentionType.ENCODER_DECODER, (
152 + "CrossAttention only supports AttentionType.ENCODER_DECODER")
153 +
154 + super().__init__(num_heads=num_heads,
155 + head_size=head_size,
156 + scale=scale,
157 + cache_config=cache_config,
158 + attn_backend=attn_backend,
159 + attn_type=AttentionType.ENCODER_DECODER,
160 + **kwargs)
vllm/config/__init__.py
+28 -9
@@ -8,6 +8,7 @@ import enum
8 import hashlib
9 import inspect
10 import json
11 +import os
12 import textwrap
13 import warnings
14 from collections.abc import Mapping
@@ -41,6 +42,7 @@ from vllm.config.scheduler import SchedulerConfig, SchedulerPolicy
42 from vllm.config.utils import ConfigType, config
43 from vllm.logger import init_logger
44 from vllm.model_executor.layers.quantization import QuantizationMethods
45 +from vllm.multimodal import MULTIMODAL_REGISTRY
46 from vllm.platforms import current_platform
47 from vllm.transformers_utils.config import (
48 ConfigFormat, get_config, get_hf_image_processor_config,
@@ -3509,16 +3511,33 @@ class VllmConfig:
3511
3512 disable_chunked_prefill_reasons: list[str] = []
3513
3512 - if self.model_config and self.model_config.pooler_config:
3513 - pooling_type = self.model_config.pooler_config.pooling_type
3514 - if pooling_type is None or pooling_type.lower() != "last":
3515 - disable_chunked_prefill_reasons.append(
3516 - "Only \"last\" pooling supports chunked "
3517 - "prefill and prefix caching; disabling both.")
3518 - elif not getattr(self.model_config.hf_config, "is_causal", True):
3514 + if self.model_config:
3515 + if self.model_config.pooler_config:
3516 + pooling_type = self.model_config.pooler_config.pooling_type
3517 + if pooling_type is None or pooling_type.lower() != "last":
3518 + disable_chunked_prefill_reasons.append(
3519 + "Only \"last\" pooling supports chunked "
3520 + "prefill and prefix caching; disabling both.")
3521 + elif self.model_config.is_encoder_decoder:
3522 + self.scheduler_config.max_num_encoder_input_tokens = \
3523 + MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(self.model_config)
3524 + logger.debug(
3525 + "Encoder-decoder model detected: setting "
3526 + "`max_num_encoder_input_tokens` to encoder length (%s)",
3527 + self.scheduler_config.max_num_encoder_input_tokens)
3528 + self.scheduler_config.disable_chunked_mm_input = True
3529 disable_chunked_prefill_reasons.append(
3520 - "Only models using causal attention supports chunked "
3521 - "prefill and prefix caching; disabling both.")
3530 + "Encoder-decoder models do not support chunked prefill nor"
3531 + " prefix caching; disabling both.")
3532 + if (self.model_config.architecture
3533 + == "WhisperForConditionalGeneration"
3534 + and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD")
3535 + != "spawn"):
3536 + logger.warning(
3537 + "Whisper is known to have issues with "
3538 + "forked workers. If startup is hanging, "
3539 + "try setting 'VLLM_WORKER_MULTIPROC_METHOD' "
3540 + "to 'spawn'.")
3541
3542 if disable_chunked_prefill_reasons:
3543 for reason in disable_chunked_prefill_reasons:
vllm/model_executor/models/voxtral.py
-1
@@ -600,7 +600,6 @@ class VoxtralEncoderModel(nn.Module):
600 self.whisper_encoder = WhisperEncoder(vllm_config=vllm_config,
601 prefix=maybe_prefix(
602 prefix, "whisper_encoder"),
603 - is_standalone_encoder=True,
603 init_in_fp32=True)
604 mel_filters = mel_filter_bank(
605 num_frequency_bins=1 + self.config.window_size // 2,
vllm/model_executor/models/whisper.py
+51 -23
@@ -15,6 +15,7 @@ from transformers.models.whisper.modeling_whisper import sinusoids
15
16 from vllm.attention import Attention, AttentionType
17 from vllm.attention.layer import MultiHeadAttention
18 +from vllm.attention.layers.cross_attention import CrossAttention
19 from vllm.config import (CacheConfig, ModelConfig, SpeechToTextConfig,
20 VllmConfig)
21 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -43,7 +44,7 @@ from vllm.transformers_utils.processor import cached_get_processor
44 from vllm.utils.tensor_schema import TensorSchema, TensorShape
45
46 from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,
46 - SupportsTranscription, SupportsV0Only)
47 + SupportsTranscription)
48 from .utils import (AutoWeightsLoader, WeightsMapper, cast_overflow_tensors,
49 make_layers)
50
@@ -124,6 +125,34 @@ class WhisperAudioInputs(TensorSchema):
125 TensorShape("b", "nmb", "t")]
126
127
128 +class WhisperEncoderAttention(MultiHeadAttention):
129 + """Multi-headed attention for Whisper encoder with 2D tensor support."""
130 +
131 + def forward(
132 + self,
133 + query: torch.Tensor,
134 + key: torch.Tensor,
135 + value: torch.Tensor,
136 + ) -> torch.Tensor:
137 + """
138 + Input shape: batch_size x seq_len x hidden_size
139 + or seq_len x hidden_size
140 + """
141 + is_2d = query.dim() == 2
142 + if is_2d:
143 + query = query.unsqueeze(0)
144 + key = key.unsqueeze(0)
145 + value = value.unsqueeze(0)
146 +
147 + # Call the parent forward method
148 + out = super().forward(query, key, value)
149 +
150 + if is_2d:
151 + out = out.squeeze(0)
152 +
153 + return out
154 +
155 +
156 class WhisperPositionalEmbedding(nn.Embedding):
157
158 def __init__(self, num_positions: int, embedding_dim: int):
@@ -144,7 +173,6 @@ class WhisperAttention(nn.Module):
173 cache_config: Optional[CacheConfig] = None,
174 quant_config: Optional[QuantizationConfig] = None,
175 prefix: str = "",
147 - standalone_encoder: bool = False,
176 ):
177 super().__init__()
178 self.embed_dim = embed_dim
@@ -180,14 +208,25 @@ class WhisperAttention(nn.Module):
208 quant_config=quant_config,
209 prefix=f"{prefix}.out_proj",
210 )
183 - if standalone_encoder:
184 - self.attn = MultiHeadAttention(
211 + if attn_type == AttentionType.ENCODER:
212 + self.attn = WhisperEncoderAttention(
213 self.num_heads,
214 self.head_dim,
215 self.scaling,
216 num_kv_heads=self.num_kv_heads,
217 )
190 - else:
218 + elif self.attn_type == AttentionType.ENCODER_DECODER:
219 + self.attn = CrossAttention(
220 + self.num_heads,
221 + self.head_dim,
222 + self.scaling,
223 + num_kv_heads=self.num_kv_heads,
224 + cache_config=cache_config,
225 + quant_config=quant_config,
226 + prefix=f"{prefix}.attn",
227 + attn_type=self.attn_type,
228 + )
229 + else: # AttentionType.DECODER (regular decoder self-attention)
230 self.attn = Attention(
231 self.num_heads,
232 self.head_dim,
@@ -332,11 +371,7 @@ class WhisperMLP(nn.Module):
371
372 class WhisperEncoderLayer(nn.Module):
373
335 - def __init__(self,
336 - *,
337 - vllm_config: VllmConfig,
338 - prefix: str = "",
339 - is_standalone_encoder: bool = False):
374 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
375 super().__init__()
376 config = vllm_config.model_config.hf_config
377 cache_config = vllm_config.cache_config
@@ -350,7 +385,6 @@ class WhisperEncoderLayer(nn.Module):
385 cache_config=cache_config,
386 quant_config=quant_config,
387 prefix=f"{prefix}.self_attn",
353 - standalone_encoder=is_standalone_encoder,
388 )
389 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
390 self.mlp = WhisperMLP(
@@ -446,12 +480,10 @@ class WhisperEncoder(nn.Module):
480 *,
481 vllm_config: VllmConfig,
482 prefix: str = "",
449 - is_standalone_encoder: bool = False,
483 init_in_fp32: bool = False):
484 super().__init__()
485 config = vllm_config.model_config.hf_config
486 embed_dim = config.d_model
454 - self.is_standalone_encoder = is_standalone_encoder
487 self.num_mel_bins = config.num_mel_bins
488 self.max_source_positions = config.max_source_positions
489 self.embed_scale = (math.sqrt(embed_dim)
@@ -469,9 +501,7 @@ class WhisperEncoder(nn.Module):
501 self.start_layer, self.end_layer, self.layers = make_layers(
502 config.encoder_layers,
503 lambda prefix: WhisperEncoderLayer(vllm_config=vllm_config,
472 - prefix=f"{prefix}.layers",
473 - is_standalone_encoder=
474 - is_standalone_encoder),
504 + prefix=f"{prefix}.layers"),
505 prefix=f"{prefix}.layers",
506 )
507 self.layer_norm = nn.LayerNorm(config.d_model)
@@ -752,7 +782,7 @@ class WhisperMultiModalProcessor(
782 info=WhisperProcessingInfo,
783 dummy_inputs=WhisperDummyInputsBuilder)
784 class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,
755 - SupportsMultiModal, SupportsV0Only):
785 + SupportsMultiModal):
786 packed_modules_mapping = {
787 "self_attn.qkv_proj": [
788 "self_attn.q_proj",
@@ -880,19 +910,17 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,
910
911 def get_multimodal_embeddings(self,
912 **kwargs: object) -> MultiModalEmbeddings:
883 - # TODO: This method does not obey the interface for SupportsMultiModal.
884 - # Refactor this once encoder/decoder support is implemented in V1.
913 + # Required as part of SupportsMultiModal interface.
914 audio_input = self._parse_and_validate_audio_input(**kwargs)
886 - return self.model.get_encoder_outputs(audio_input["input_features"])
915 + return [self.model.get_encoder_outputs(audio_input["input_features"])]
916
917 def get_input_embeddings(
918 self,
919 input_ids: torch.Tensor,
920 multimodal_embeddings: Optional[NestedTensors] = None,
921 ) -> torch.Tensor:
893 - # TODO: This method just returns the decoder sequence embeddings since
894 - # Whisper does not have encoder text tokens. Refactor this once
895 - # encoder/decoder support is implemented in V1.
922 + # This method just returns the decoder sequence embeddings since
923 + # Whisper does not have encoder text tokens.
924 return self.model.decoder.get_input_embeddings(input_ids)
925
926 def _parse_and_validate_audio_input(
vllm/transformers_utils/configs/mistral.py
+1
@@ -157,6 +157,7 @@ def _remap_mistral_audio_args(config: dict) -> dict:
157 encoder_attention_heads=encoder_args["n_heads"],
158 vocab_size=encoder_args["vocab_size"],
159 max_source_positions=encoder_args["max_source_positions"],
160 + is_encoder_decoder=False, # Override WhisperConfig default
161 )
162 }
163 if quant_config:
vllm/v1/attention/backends/cpu_attn.py
+2 -2
@@ -317,8 +317,8 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
317
318 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
319 vllm_config: VllmConfig, device: torch.device) -> None:
320 - self.kv_cache_spec = kv_cache_spec
321 - self.vllm_config = vllm_config
320 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
321 +
322 self.scheduler_config = vllm_config.scheduler_config
323
324 # For reorder
vllm/v1/attention/backends/flash_attn.py
+1 -2
@@ -177,12 +177,11 @@ class FlashAttentionMetadataBuilder(
177
178 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
179 vllm_config: VllmConfig, device: torch.device):
180 - self.vllm_config = vllm_config
180 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
181 self.model_config = vllm_config.model_config
182 self.parallel_config = vllm_config.parallel_config
183 self.cache_config = vllm_config.cache_config
184 self.compilation_config = vllm_config.compilation_config
185 - self.device = device
185
186 self.num_heads_q = self.model_config.get_num_attention_heads(
187 self.parallel_config)
vllm/v1/attention/backends/flashinfer.py
+1 -3
@@ -163,11 +163,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
163
164 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
165 vllm_config: VllmConfig, device: torch.device):
166 - self.device = device
167 - self.vllm_config = vllm_config
166 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
167 self.cache_config = vllm_config.cache_config
168 self.model_config = vllm_config.model_config
170 - self.kv_cache_spec = kv_cache_spec
169 self._workspace_buffer = None
170 self._prefill_wrapper = None # Wrapper for prefill/append
171 self._decode_wrapper = None # Wrapper for decode (general shape)
vllm/v1/attention/backends/flex_attention.py
+2 -1
@@ -516,10 +516,11 @@ class FlexAttentionMetadataBuilder(
516
517 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
518 vllm_config: VllmConfig, device: torch.device):
519 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
520 +
521 self.model_config = vllm_config.model_config
522 self.parallel_config = vllm_config.parallel_config
523 self.cache_config = vllm_config.cache_config
522 - self.device = device
524
525 self.num_heads_q = self.model_config.get_num_attention_heads(
526 self.parallel_config)
vllm/v1/attention/backends/linear_attn.py
+1 -1
@@ -39,8 +39,8 @@ class LinearAttentionMetadataBuilder(
39
40 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
41 vllm_config: VllmConfig, device: torch.device):
42 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
43 assert isinstance(kv_cache_spec, MambaSpec)
43 - self.kv_cache_spec = kv_cache_spec
44
45 def build(self,
46 common_prefix_len: int,
vllm/v1/attention/backends/mamba_attn.py
+3 -6
@@ -22,12 +22,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
22
23 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
24 vllm_config: VllmConfig, device: torch.device):
25 - assert isinstance(kv_cache_spec, MambaSpec)
26 - self.kv_cache_spec = kv_cache_spec
27 - self.device = device
28 - self.vllm_config = vllm_config
29 - self.layer_names = layer_names
25 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
26
27 + assert isinstance(kv_cache_spec, MambaSpec)
28 self.compilation_config = vllm_config.compilation_config
29 self.decode_cudagraph_max_bs = min(
30 self.vllm_config.scheduler_config.max_num_seqs,
@@ -52,4 +49,4 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
49
50 m.max_query_len = 1 # decode-only
51
55 - return self.build(0, m)
\ No newline at end of file
52 + return self.build(0, m)
vllm/v1/attention/backends/rocm_aiter_fa.py
+2 -3
@@ -236,11 +236,11 @@ class AiterFlashAttentionMetadataBuilder(
236
237 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
238 vllm_config: VllmConfig, device: torch.device):
239 - self.vllm_config = vllm_config
239 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
240 +
241 self.model_config = vllm_config.model_config
242 self.parallel_config = vllm_config.parallel_config
243 self.cache_config = vllm_config.cache_config
243 - self.device = device
244
245 self.num_heads_q = self.model_config.get_num_attention_heads(
246 self.parallel_config)
@@ -248,7 +248,6 @@ class AiterFlashAttentionMetadataBuilder(
248 self.parallel_config)
249 self.headdim = self.model_config.get_head_size()
250 self.block_size = kv_cache_spec.block_size
251 - self.kv_cache_spec = kv_cache_spec
251 # Sliding window size to be used with the AOT scheduler will be
252 # populated on first build() call.
253 self.aot_sliding_window: Optional[tuple[int, int]] = None
vllm/v1/attention/backends/short_conv_attn.py
+1 -1
@@ -45,8 +45,8 @@ class ShortConvAttentionMetadataBuilder(
45
46 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
47 vllm_config: VllmConfig, device: torch.device):
48 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
49 assert isinstance(kv_cache_spec, MambaSpec)
49 - self.kv_cache_spec = kv_cache_spec
50
51 def build(self,
52 common_prefix_len: int,
vllm/v1/attention/backends/tree_attn.py
+2 -1
@@ -165,7 +165,8 @@ class TreeAttentionMetadataBuilder(
165 vllm_config: VllmConfig,
166 device: torch.device,
167 ):
168 - self.kv_cache_spec = kv_cache_spec
168 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
169 +
170 self.block_size = kv_cache_spec.block_size
171
172 spec_config = vllm_config.speculative_config
vllm/v1/attention/backends/triton_attn.py
+2 -2
@@ -66,9 +66,9 @@ class TritonAttentionMetadataBuilder(
66
67 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
68 vllm_config: VllmConfig, device: torch.device):
69 - self.device = device
69 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
70 +
71 self.block_size = kv_cache_spec.block_size
71 - self.kv_cache_spec = kv_cache_spec
72
73 model_config = vllm_config.model_config
74 self.num_heads_q = model_config.get_num_attention_heads(
vllm/v1/attention/backends/utils.py
+6
@@ -72,6 +72,9 @@ class CommonAttentionMetadata:
72 logits_indices_padded: Optional[torch.Tensor] = None
73 num_logits_indices: Optional[int] = None
74
75 + # Needed by CrossAttentionBuilder
76 + encoder_seq_lens: Optional[np.ndarray] = None
77 +
78
79 @dataclass
80 class UbatchSlice:
@@ -193,6 +196,9 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
196 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
197 vllm_config: VllmConfig, device: torch.device):
198 self.kv_cache_spec = kv_cache_spec
199 + self.layer_names = layer_names
200 + self.vllm_config = vllm_config
201 + self.device = device
202
203 @abstractmethod
204 def build(self,
vllm/v1/attention/backends/xformers.py
+2 -1
@@ -206,8 +206,9 @@ class XFormersAttentionMetadataBuilder(
206 vllm_config: VllmConfig,
207 device: torch.device,
208 ):
209 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
210 +
211 assert XFORMERS_AVAILABLE
210 - self.kv_cache_spec = kv_cache_spec
212 self.block_size = kv_cache_spec.block_size
213 self._num_decodes = 0
214 self._num_decode_tokens = 0
vllm/v1/core/sched/scheduler.py
+21 -11
@@ -144,8 +144,8 @@ class Scheduler(SchedulerInterface):
144 )
145
146 # NOTE(woosuk): Here, "encoder" includes the vision encoder (and
147 - # projector if needed). Currently, we assume that the encoder also
148 - # has the Transformer architecture (e.g., ViT).
147 + # projector if needed) for MM models as well as encoder-decoder
148 + # transformers.
149 self.max_num_encoder_input_tokens = encoder_compute_budget
150 # NOTE: For the models without encoder (e.g., text-only models),
151 # the encoder cache will not be initialized because cache size is 0
@@ -775,15 +775,19 @@ class Scheduler(SchedulerInterface):
775 # in the decoder's KV cache.
776 continue
777
778 - # The same encoder input has already been scheduled in the current
779 - # step.
780 - if request.mm_hashes[i] in mm_hashes_to_schedule:
781 - continue
778 + if not self.is_encoder_decoder:
779 + # We are not using the encoder cache for encoder-decoder models,
780 + # yet.
781 + if request.mm_hashes[i] in mm_hashes_to_schedule:
782 + # The same encoder input has already been scheduled in the
783 + # current step.
784 + continue
785
783 - if self.encoder_cache_manager.check_and_update_cache(request, i):
784 - # The encoder input is already computed and cached from a
785 - # previous step.
786 - continue
786 + if self.encoder_cache_manager.check_and_update_cache(
787 + request, i):
788 + # The encoder input is already computed and cached from a
789 + # previous step.
790 + continue
791
792 # If no encoder input chunking is allowed, we do not want to
793 # partially schedule a multimodal item. If the scheduled range would
@@ -1047,7 +1051,13 @@ class Scheduler(SchedulerInterface):
1051 mm_positions = request.mm_positions[input_id]
1052 start_pos = mm_positions.offset
1053 num_tokens = mm_positions.length
1050 - if start_pos + num_tokens <= request.num_computed_tokens:
1054 + if self.is_encoder_decoder and request.num_computed_tokens > 0:
1055 + # With Whisper, as soon as we've generated a single token,
1056 + # we know we're done with the encoder input. Cross Attention
1057 + # KVs have been calculated and cached already.
1058 + self.encoder_cache_manager.free_encoder_input(
1059 + request, input_id)
1060 + elif start_pos + num_tokens <= request.num_computed_tokens:
1061 # The encoder output is already processed and stored
1062 # in the decoder's KV cache.
1063 self.encoder_cache_manager.free_encoder_input(
vllm/v1/engine/processor.py
-5
@@ -325,7 +325,6 @@ class Processor:
325 ) -> tuple[Optional[str], EngineCoreRequest]:
326
327 # TODO(woosuk): Support pooling models.
328 - # TODO(woosuk): Support encoder-decoder models.
328 self._validate_lora(lora_request)
329 self._validate_params(params, lora_request)
330 if trace_headers is not None:
@@ -384,10 +383,6 @@ class Processor:
383
384 encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
385
387 - # TODO: Impl encoder-decoder
388 - if encoder_inputs is not None:
389 - raise NotImplementedError
390 -
386 sampling_params = None
387 pooling_params = None
388 if isinstance(params, SamplingParams):
vllm/v1/worker/gpu_model_runner.py
+111 -13
@@ -61,12 +61,16 @@ from vllm.v1.attention.backends.utils import (
61 create_fast_prefill_custom_backend,
62 reorder_batch_to_split_decodes_and_prefills)
63 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
64 +# yapf conflicts with isort for this block
65 +# yapf: disable
66 from vllm.v1.kv_cache_interface import (AttentionSpec,
67 ChunkedLocalAttentionSpec,
68 + CrossAttentionSpec,
69 EncoderOnlyAttentionSpec,
70 FullAttentionSpec, KVCacheConfig,
71 KVCacheGroupSpec, KVCacheSpec,
72 MambaSpec, SlidingWindowSpec)
73 +# yapf: enable
74 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
75 DraftTokenIds, LogprobsLists, LogprobsTensors,
76 ModelRunnerOutput, SamplerOutput)
@@ -208,6 +212,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
212 self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
213 model_config)
214
215 + if self.model_config.is_encoder_decoder:
216 + # Maximum length of the encoder input, only for encoder-decoder
217 + # models.
218 + self.max_encoder_len = self.mm_registry.\
219 + get_encdec_max_encoder_len(model_config)
220 + else:
221 + self.max_encoder_len = 0
222 +
223 # Sampler
224 self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
225
@@ -265,7 +277,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
277 # the block_sizes in the kv cache config.
278 self.input_batch = InputBatch(
279 max_num_reqs=self.max_num_reqs,
268 - max_model_len=self.max_model_len,
280 + # We need to use the encoder length for encoder-decoer
281 + # because of KV cache for cross-attention.
282 + max_model_len=max(self.max_model_len, self.max_encoder_len),
283 max_num_batched_tokens=self.max_num_tokens,
284 device=self.device,
285 pin_memory=self.pin_memory,
@@ -798,6 +812,24 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
812 src=self.input_batch.prev_sampled_token_ids[
813 prev_common_req_indices_tensor, 0])
814
815 + def _get_encoder_seq_lens(
816 + self,
817 + scheduler_output: "SchedulerOutput",
818 + kv_cache_spec: KVCacheSpec,
819 + num_reqs: int,
820 + ) -> Optional[np.ndarray]:
821 + if not isinstance(kv_cache_spec, CrossAttentionSpec):
822 + return None
823 +
824 + # Build encoder_seq_lens array mapping request indices to
825 + # encoder lengths for inputs scheduled in this batch
826 + encoder_seq_lens = np.zeros(num_reqs, dtype=np.int32)
827 + for req_id in scheduler_output.scheduled_encoder_inputs:
828 + req_index = self.input_batch.req_id_to_index[req_id]
829 + encoder_seq_lens[req_index] = self.max_encoder_len
830 +
831 + return encoder_seq_lens
832 +
833 def _prepare_inputs(
834 self,
835 scheduler_output: "SchedulerOutput",
@@ -937,6 +969,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
969 # in the same group share the same metadata.
970 for kv_cache_group_id, kv_cache_group_spec in enumerate(
971 self.kv_cache_config.kv_cache_groups):
972 + encoder_seq_lens = self._get_encoder_seq_lens(
973 + scheduler_output, kv_cache_group_spec.kv_cache_spec, num_reqs)
974
975 if isinstance(kv_cache_group_spec.kv_cache_spec,
976 EncoderOnlyAttentionSpec):
@@ -981,6 +1015,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1015 logits_indices_padded=logits_indices_padded,
1016 num_logits_indices=logits_indices.size(0),
1017 causal=True,
1018 + encoder_seq_lens=encoder_seq_lens,
1019 )
1020
1021 if self.speculative_config and \
@@ -1253,10 +1288,24 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1288 self.kv_sharing_fast_prefill_logits_indices[:num_logits_padded])
1289 return logits_indices_padded
1290
1256 - def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
1291 + def _batch_mm_kwargs_from_scheduler(
1292 + self,
1293 + scheduler_output: "SchedulerOutput",
1294 + ) -> tuple[list[MultiModalKwargsItem], list[tuple[str, PlaceholderRange]]]:
1295 + """Batch multimodal kwargs from scheduled encoder inputs.
1296 +
1297 + Args:
1298 + scheduler_output: The scheduler output containing scheduled encoder
1299 + inputs.
1300 +
1301 + Returns:
1302 + A tuple of (mm_kwargs, req_ids_pos) where:
1303 + - mm_kwargs: List of multimodal kwargs items to be batched
1304 + - mm_hashes_pos: List of (mm_hash, position_info) tuples
1305 + """
1306 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
1307 if not scheduled_encoder_inputs:
1259 - return
1308 + return [], []
1309 # Batch the multi-modal inputs.
1310 mm_kwargs = list[MultiModalKwargsItem]()
1311 # list of tuple (mm_hash, position_info)
@@ -1270,6 +1319,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1319 mm_hashes_pos.append(
1320 (mm_hash, req_state.mm_positions[mm_input_id]))
1321
1322 + return mm_kwargs, mm_hashes_pos
1323 +
1324 + def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
1325 + # Batch the multi-modal inputs using the helper method.
1326 + mm_kwargs, mm_hashes_pos = self._batch_mm_kwargs_from_scheduler(
1327 + scheduler_output)
1328 +
1329 + if not mm_kwargs:
1330 + return
1331 +
1332 # Batch mm inputs as much as we can: if a request in the batch has
1333 # multiple modalities or a different modality than the previous one,
1334 # we process it separately to preserve item order.
@@ -1360,6 +1419,35 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1419 mm_embeds.append(mm_embeds_item)
1420 return mm_embeds
1421
1422 + def _extract_encoder_inputs(
1423 + self,
1424 + scheduler_output: "SchedulerOutput",
1425 + ) -> dict[str, torch.Tensor]:
1426 + """Extract encoder inputs for encoder-decoder models.
1427 +
1428 + This method extracts multimodal input features from scheduled encoder
1429 + inputs and formats them for the encoder-decoder model forward pass.
1430 + """
1431 + # Batch the multi-modal inputs using the helper method.
1432 + mm_kwargs, _ = self._batch_mm_kwargs_from_scheduler(scheduler_output)
1433 +
1434 + if not mm_kwargs:
1435 + return {}
1436 +
1437 + # Group MM kwargs by modality and extract features
1438 + encoder_features = {}
1439 + for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
1440 + mm_kwargs,
1441 + device=self.device,
1442 + pin_memory=self.pin_memory,
1443 + ):
1444 + # Add the grouped features to encoder_features dict
1445 + # This allows the model to receive them as kwargs (e.g.,
1446 + # input_features=...)
1447 + encoder_features.update(mm_kwargs_group)
1448 +
1449 + return encoder_features
1450 +
1451 def get_model(self) -> nn.Module:
1452 # get raw model out of the cudagraph wrapper.
1453 if isinstance(self.model, CUDAGraphWrapper):
@@ -1631,7 +1719,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1719
1720 # _prepare_inputs may reorder the batch, so we must gather multi
1721 # modal outputs after that to ensure the correct order
1634 - if self.supports_mm_inputs and get_pp_group().is_first_rank:
1722 + if (self.supports_mm_inputs and get_pp_group().is_first_rank
1723 + and not self.model_config.is_encoder_decoder):
1724 # Run the multimodal encoder if any.
1725 self._execute_mm_encoder(scheduler_output)
1726 mm_embeds = self._gather_mm_embeddings(scheduler_output)
@@ -1673,6 +1762,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1762 intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1763 num_input_tokens, intermediate_tensors, True)
1764
1765 + if (self.model_config.is_encoder_decoder
1766 + and scheduler_output.scheduled_encoder_inputs):
1767 + encoder_inputs = self._extract_encoder_inputs(scheduler_output)
1768 + model_kwargs.update(encoder_inputs)
1769 +
1770 return (
1771 num_scheduled_tokens,
1772 num_input_tokens,
@@ -2591,17 +2685,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2685
2686 with self.maybe_dummy_run_with_lora(self.lora_config,
2687 num_scheduled_tokens, remove_lora):
2594 - if self.supports_mm_inputs:
2688 + model_kwargs = self._init_model_kwargs(num_tokens)
2689 + if (self.supports_mm_inputs
2690 + and not self.model_config.is_encoder_decoder):
2691 input_ids = None
2692 inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
2693 model_kwargs = {
2598 - **self._init_model_kwargs(num_tokens),
2694 + **model_kwargs,
2695 **self._dummy_mm_kwargs(num_reqs),
2696 }
2697 else:
2698 input_ids = self.input_ids.gpu[:num_tokens]
2699 inputs_embeds = None
2604 - model_kwargs = self._init_model_kwargs(num_tokens)
2700
2701 if self.uses_mrope:
2702 positions = self.mrope_positions.gpu[:, :num_tokens]
@@ -2823,7 +2918,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2918 mm_budget = self.mm_budget
2919 assert mm_budget is not None
2920
2826 - # TODO: handle encoder-decoder models once we support them.
2921 if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
2922 # NOTE: Currently model is profiled with a single non-text
2923 # modality with the max possible input tokens even when
@@ -3170,7 +3264,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3264 "for more details.")
3265 self.input_batch = InputBatch(
3266 max_num_reqs=self.max_num_reqs,
3173 - max_model_len=self.max_model_len,
3267 + max_model_len=max(self.max_model_len, self.max_encoder_len),
3268 max_num_batched_tokens=self.max_num_tokens,
3269 device=self.device,
3270 pin_memory=self.pin_memory,
@@ -3443,7 +3537,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3537 attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
3538 for layer_name, attn_module in attn_layers.items():
3539 if attn_module.attn_type == AttentionType.ENCODER_ONLY:
3446 - attn_spec = EncoderOnlyAttentionSpec(
3540 + attn_spec: AttentionSpec = EncoderOnlyAttentionSpec(
3541 block_size=block_size,
3542 num_kv_heads=attn_module.num_kv_heads,
3543 head_size=attn_module.head_size,
@@ -3485,7 +3579,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3579 self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
3580 continue
3581
3488 - # TODO: Support other attention modules, e.g., cross-attention
3582 # TODO(lucas): move the attention specs into the model layers like
3583 # the attention backends
3584 if attn_module.attn_type == AttentionType.DECODER:
@@ -3513,12 +3606,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3606 head_size=attn_module.head_size,
3607 dtype=self.kv_cache_dtype,
3608 use_mla=use_mla)
3609 + elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
3610 + kv_cache_spec[layer_name] = CrossAttentionSpec(
3611 + block_size=block_size,
3612 + num_kv_heads=attn_module.num_kv_heads,
3613 + head_size=attn_module.head_size,
3614 + dtype=self.kv_cache_dtype,
3615 + use_mla=use_mla)
3616 elif attn_module.attn_type in (AttentionType.ENCODER,
3617 AttentionType.ENCODER_ONLY):
3618 # encoder-only attention does not need KV cache.
3619 continue
3520 - elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
3521 - raise NotImplementedError
3620 else:
3621 raise ValueError(
3622 f"Unknown attention type: {attn_module.attn_type}")
vllm/v1/worker/utils.py
+12 -1
@@ -12,6 +12,7 @@ from vllm.model_executor.models.interfaces import MultiModalEmbeddings
12 from vllm.model_executor.models.utils import extract_layer_index
13 from vllm.multimodal.cache import processor_only_cache_from_config
14 from vllm.multimodal.registry import MultiModalRegistry
15 +from vllm.platforms import current_platform
16 from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
17 from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
18 from vllm.v1.kv_cache_interface import KVCacheGroupSpec
@@ -269,7 +270,17 @@ def bind_kv_cache(
270 # One typical case is encoder-decoder model, e.g., bart.
271 # The cross attention and self attention in the same decoder layer
272 # has different layer_name but the same layer_index.
272 - raise NotImplementedError
273 +
274 + # TODO - analyze where runner_kv_caches is used and the right
275 + # way to ensure it properly reflects multiple attention layers
276 + # in the same decoder block.
277 + if current_platform.is_cuda():
278 + # We know that the GPU runner is not impacted by this
279 + # case. Some test code depends on runner_kv_caches, but
280 + # not in a way that's impacted by ignoring this.
281 + pass
282 + else:
283 + raise NotImplementedError
284 layer_name = layer_names[0]
285 runner_kv_caches.append(kv_caches[layer_name])
286