[ModelRunner V2] Support encoder-only attention (#49331)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Nick Hill committed
Jul 27, 2026 at 01:38 UTC
fdaa0d9e59238b6884f9515fa3245dea118edc66
8 files changed
+276
-4
tests/models/language/pooling/test_embedding.py
+54
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+import torch
6
+from transformers import AutoModel
7
8
from vllm.config import PoolerConfig
9
@@ -87,3 +89,55 @@ def test_models(
89
name_1="vllm",
90
tol=1e-2,
91
)
92
+
93
+
94
+@pytest.mark.parametrize(
95
+ "model",
96
+ [
97
+ "BAAI/bge-base-en-v1.5",
98
+ "intfloat/multilingual-e5-small",
99
+ ],
100
+)
101
+@torch.inference_mode()
102
+def test_encoder_only_model_runner_v2_attention(
103
+ hf_runner,
104
+ vllm_runner,
105
+ monkeypatch,
106
+ model: str,
107
+) -> None:
108
+ prompts = [
109
+ "short input",
110
+ "a longer input that exercises mixed sequence lengths",
111
+ ]
112
+
113
+ with hf_runner(model, dtype="float", auto_cls=AutoModel) as hf_model:
114
+ hf_outputs = []
115
+ for prompt in prompts:
116
+ inputs = hf_model.tokenizer(prompt, return_tensors="pt")
117
+ output = hf_model.model(**hf_model.wrap_device(inputs))
118
+ embedding = torch.nn.functional.normalize(
119
+ output.last_hidden_state[0, -1].float(), dim=0
120
+ )
121
+ hf_outputs.append(embedding.cpu().tolist())
122
+
123
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
124
+ with vllm_runner(
125
+ model,
126
+ runner="pooling",
127
+ dtype="float",
128
+ max_model_len=64,
129
+ max_num_seqs=2,
130
+ gpu_memory_utilization=0.25,
131
+ pooler_config=PoolerConfig(
132
+ task="embed", seq_pooling_type="LAST", use_activation=True
133
+ ),
134
+ ) as vllm_model:
135
+ vllm_outputs = vllm_model.embed(prompts)
136
+
137
+ check_embeddings_close(
138
+ embeddings_0_lst=hf_outputs,
139
+ embeddings_1_lst=vllm_outputs,
140
+ name_0="hf",
141
+ name_1="vllm",
142
+ tol=1e-2,
143
+ )
vllm/v1/worker/gpu/attn_utils.py
+12
@@ -46,6 +46,18 @@ class AttentionCGSupportInfo:
46
min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS
47
min_cg_attn_backend: str | None = None
48
49
+ def narrow(
50
+ self, support: AttentionCGSupport, backend: str | None
51
+ ) -> "AttentionCGSupportInfo":
52
+ """Return an info tightened by ``support`` if it is more restrictive.
53
+
54
+ Lets attention groups built outside ``init_attn_backend`` (e.g.
55
+ encoder-only layers) contribute to the runner's cudagraph decision.
56
+ """
57
+ if support.value < self.min_cg_support.value:
58
+ return AttentionCGSupportInfo(support, backend)
59
+ return self
60
+
61
62
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
63
kv_cache_spec: dict[str, KVCacheSpec] = {}
vllm/v1/worker/gpu/block_table.py
+1
-1
@@ -123,7 +123,7 @@ class BlockTables:
123
if self.num_kv_cache_groups == 1:
124
# Single group: write directly, skipping the per-write group lookup.
125
self.block_tables[0].apply_write()
126
- else:
126
+ elif self.num_kv_cache_groups > 1:
127
# Multiple groups: apply all block tables with one fused kernel.
128
assert self.fused_writer is not None
129
self.fused_writer.apply(
vllm/v1/worker/gpu/model_runner.py
+3
@@ -459,6 +459,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
459
self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
460
self.kv_cache_config, self.vllm_config, self.device
461
)
462
+ attn_cg_support = attn_cg_support.narrow(
463
+ *self.model_state.get_additional_cg_support()
464
+ )
465
self.block_tables = BlockTables(
466
block_sizes=block_sizes,
467
max_num_reqs=self.max_num_reqs,
vllm/v1/worker/gpu/model_states/__init__.py
+7
-1
@@ -4,7 +4,7 @@ import torch
4
import torch.nn as nn
5
6
from vllm.config import VllmConfig
7
-from vllm.model_executor.layers.attention import CrossAttention
7
+from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention
8
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
9
10
@@ -27,6 +27,12 @@ def init_model_state(
27
28
return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
29
30
+ # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache.
31
+ if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()):
32
+ from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState
33
+
34
+ return EncoderOnlyModelState(vllm_config, model, encoder_cache, device)
35
+
36
if vllm_config.model_config.is_hybrid:
37
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
38
vllm/v1/worker/gpu/model_states/encoder_only.py
new
+182
@@ -0,0 +1,182 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any, cast
4
+
5
+import torch
6
+import torch.nn as nn
7
+
8
+from vllm.config import VllmConfig, get_layers_from_vllm_config
9
+from vllm.config.compilation import CUDAGraphMode
10
+from vllm.model_executor.layers.attention import Attention
11
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
12
+from vllm.v1.attention.backend import (
13
+ AttentionCGSupport,
14
+ AttentionType,
15
+ CommonAttentionMetadata,
16
+)
17
+from vllm.v1.kv_cache_interface import (
18
+ AttentionSpec,
19
+ EncoderOnlyAttentionSpec,
20
+ KVCacheConfig,
21
+)
22
+from vllm.v1.worker.gpu.input_batch import InputBatch
23
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
24
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
25
+from vllm.v1.worker.utils import AttentionGroup
26
+
27
+
28
+class EncoderOnlyModelState(DefaultModelState):
29
+ """ModelState for encoder-only (BERT/RoBERTa) models.
30
+
31
+ Encoder attention needs no KV cache: it runs full, bidirectional
32
+ self-attention over each request's tokens in a single forward. Such layers
33
+ return no ``get_kv_cache_spec`` and therefore never enter
34
+ ``kv_cache_config.kv_cache_groups``, so the KV-backed attention path never
35
+ builds their metadata. We build their (non-causal) metadata here, keeping
36
+ the normal KV-backed path untouched.
37
+ """
38
+
39
+ def __init__(
40
+ self,
41
+ vllm_config: VllmConfig,
42
+ model: nn.Module,
43
+ encoder_cache: EncoderCache | None,
44
+ device: torch.device,
45
+ ):
46
+ super().__init__(vllm_config, model, encoder_cache, device)
47
+
48
+ cache_config = vllm_config.cache_config
49
+ if cache_config.cache_dtype == "auto":
50
+ kv_cache_dtype = self.dtype
51
+ else:
52
+ kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
53
+
54
+ # Build an attention group (and its non-causal metadata builder) for the
55
+ # encoder-only layers, grouped by backend + query/KV head config. Models
56
+ # are typically uniform, yielding a single group.
57
+ attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
58
+ groups: dict[tuple[tuple[str, str], int, int, int], AttentionGroup] = {}
59
+ for name, layer in attn_layers.items():
60
+ if layer.attn_type != AttentionType.ENCODER_ONLY:
61
+ continue
62
+ # No KV cache is bound for these layers; give them an (empty)
63
+ # device tensor so the attention forward context is well-formed.
64
+ layer.kv_cache = torch.empty(0, dtype=kv_cache_dtype, device=device)
65
+ backend = layer.get_attn_backend()
66
+ # Metadata builders require a uniform query-head count per group.
67
+ key = (
68
+ backend.full_cls_name(),
69
+ layer.num_heads,
70
+ layer.num_kv_heads,
71
+ layer.head_size,
72
+ )
73
+ group = groups.get(key)
74
+ if group is None:
75
+ spec = EncoderOnlyAttentionSpec(
76
+ block_size=cache_config.block_size,
77
+ num_kv_heads=layer.num_kv_heads,
78
+ head_size=layer.head_size,
79
+ dtype=kv_cache_dtype,
80
+ )
81
+ group = AttentionGroup(backend, [], spec, kv_cache_group_id=len(groups))
82
+ groups[key] = group
83
+ group.layer_names.append(name)
84
+
85
+ self.encoder_attn_groups = list(groups.values())
86
+ for group in self.encoder_attn_groups:
87
+ group.create_metadata_builders(vllm_config, device)
88
+
89
+ # Encoder attention reads neither the block table nor the slot mapping
90
+ # (full varlen self-attention over the batch's q/k/v), but the metadata
91
+ # builder expects both tensors to be present.
92
+ self._dummy_block_table = torch.zeros(
93
+ self.max_num_reqs, 1, dtype=torch.int32, device=device
94
+ )
95
+ self._dummy_slot_mapping = torch.zeros(
96
+ self.max_num_tokens, dtype=torch.int64, device=device
97
+ )
98
+
99
+ def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
100
+ # Encoder groups are built here rather than in init_attn_backend, so
101
+ # their cudagraph support must be surfaced to the runner separately.
102
+ support = AttentionCGSupport.ALWAYS
103
+ backend: str | None = None
104
+ for group in self.encoder_attn_groups:
105
+ builder = group.get_metadata_builder(0)
106
+ cg_support = builder.get_cudagraph_support(
107
+ self.vllm_config, cast(AttentionSpec, group.kv_cache_spec)
108
+ )
109
+ if cg_support.value < support.value:
110
+ support = cg_support
111
+ backend = group.backend.__name__
112
+ return support, backend
113
+
114
+ def prepare_attn(
115
+ self,
116
+ input_batch: InputBatch,
117
+ cudagraph_mode: CUDAGraphMode,
118
+ block_tables: tuple[torch.Tensor, ...],
119
+ slot_mappings: torch.Tensor,
120
+ attn_groups: list[list[AttentionGroup]],
121
+ kv_cache_config: KVCacheConfig,
122
+ for_capture: bool = False,
123
+ ) -> dict[str, Any]:
124
+ attn_metadata = super().prepare_attn(
125
+ input_batch,
126
+ cudagraph_mode,
127
+ block_tables,
128
+ slot_mappings,
129
+ attn_groups,
130
+ kv_cache_config,
131
+ for_capture,
132
+ )
133
+ attn_metadata.update(
134
+ self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture)
135
+ )
136
+ return attn_metadata
137
+
138
+ def _build_encoder_attn_metadata(
139
+ self,
140
+ input_batch: InputBatch,
141
+ cudagraph_mode: CUDAGraphMode,
142
+ for_capture: bool,
143
+ ) -> dict[str, Any]:
144
+ if cudagraph_mode == CUDAGraphMode.FULL:
145
+ num_reqs = input_batch.num_reqs_after_padding
146
+ num_tokens = input_batch.num_tokens_after_padding
147
+ else:
148
+ num_reqs = input_batch.num_reqs
149
+ num_tokens = input_batch.num_tokens
150
+ max_query_len = int(input_batch.num_scheduled_tokens.max())
151
+ if for_capture:
152
+ max_seq_len = self.max_model_len
153
+ else:
154
+ max_seq_len = int(input_batch.seq_lens_cpu_upper_bound[:num_reqs].max())
155
+
156
+ # The encoder builder forces ``causal=False`` regardless of this value.
157
+ common_attn_metadata = CommonAttentionMetadata(
158
+ query_start_loc=input_batch.query_start_loc,
159
+ query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
160
+ seq_lens=input_batch.seq_lens[:num_reqs],
161
+ num_reqs=num_reqs,
162
+ num_actual_tokens=num_tokens,
163
+ max_query_len=max_query_len,
164
+ max_seq_len=max_seq_len,
165
+ block_table_tensor=self._dummy_block_table[:num_reqs],
166
+ slot_mapping=self._dummy_slot_mapping[:num_tokens],
167
+ seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound[:num_reqs],
168
+ positions=input_batch.positions,
169
+ )
170
+
171
+ attn_metadata: dict[str, Any] = {}
172
+ for group in self.encoder_attn_groups:
173
+ builder = group.get_metadata_builder(0)
174
+ if for_capture:
175
+ metadata = builder.build_for_cudagraph_capture(common_attn_metadata)
176
+ else:
177
+ metadata = builder.build(
178
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
179
+ )
180
+ for name in group.layer_names:
181
+ attn_metadata[name] = metadata
182
+ return attn_metadata
vllm/v1/worker/gpu/model_states/interface.py
+10
@@ -9,6 +9,7 @@ import torch.nn as nn
9
from vllm.config import VllmConfig
10
from vllm.config.compilation import CUDAGraphMode
11
from vllm.tasks import GenerationTask
12
+from vllm.v1.attention.backend import AttentionCGSupport
13
from vllm.v1.core.sched.output import NewRequestData
14
from vllm.v1.kv_cache_interface import KVCacheConfig
15
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -95,6 +96,15 @@ class ModelState(ABC):
96
def apply_staged_writes(self) -> None:
97
return None
98
99
+ def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
100
+ """Cudagraph support of attention groups this ModelState builds outside
101
+ ``init_attn_backend`` (e.g. encoder-only layers).
102
+
103
+ Returns the minimum support level and its backend name. The default of
104
+ ``ALWAYS`` imposes no extra constraint on the runner's cudagraph mode.
105
+ """
106
+ return AttentionCGSupport.ALWAYS, None
107
+
108
def preprocess_state(
109
self,
110
input_batch: InputBatch,
vllm/v1/worker/gpu/warmup.py
+7
-2
@@ -217,9 +217,14 @@ def warmup_kernels(
217
model_runner.scheduler_config.max_num_seqs,
218
model_runner.scheduler_config.max_num_batched_tokens
219
// max(prompt_len, decode_query_len),
220
- # Reserve block 0 (null block) and ensure we have enough blocks.
221
- max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
220
)
221
+ if max_blocks_per_req > 0:
222
+ # Reserve block 0 (null block) and ensure we have enough blocks.
223
+ # Encoder-only models allocate no KV blocks, so this cap doesn't apply.
224
+ num_reqs = min(
225
+ num_reqs,
226
+ max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
227
+ )
228
229
req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
230