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, 23:45 UTC
Support encoder-only models without KV-Cache (#21270)
Signed-off-by: Max de Bayser <maxdebayser@gmail.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Russell Bryant <rbryant@redhat.com>
Maximilien de Bayser committed
Jul 26, 2025 at 10:09 UTC
1cd6eaba54a2461bcdca3958e5befd9c40907657
17 files changed
+352
-99
examples/offline_inference/prithvi_geospatial_mae.py
+1
-1
@@ -3,12 +3,12 @@
3
import argparse
4
import datetime
5
import os
6
-import re
6
from typing import Union
7
8
import albumentations
9
import numpy as np
10
import rasterio
11
+import regex as re
12
import torch
13
from einops import rearrange
14
from terratorch.datamodules import Sen1Floods11NonGeoDataModule
tests/conftest.py
+11
-2
@@ -1062,8 +1062,17 @@ class VllmRunner:
1062
return [req_output.outputs.score for req_output in req_outputs]
1063
1064
def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
1065
- executor = self.llm.llm_engine.model_executor
1066
- return executor.apply_model(func)
1065
+ if hasattr(self.llm.llm_engine, "model_executor"):
1066
+ # This works either in V0 or in V1 with
1067
+ # VLLM_ENABLE_V1_MULTIPROCESSING=0
1068
+ executor = self.llm.llm_engine.model_executor
1069
+ return executor.apply_model(func)
1070
+
1071
+ # This works in V1 with VLLM_ALLOW_INSECURE_SERIALIZATION=1
1072
+ def _apply_model(self):
1073
+ return func(self.get_model())
1074
+
1075
+ return self.llm.llm_engine.collective_rpc(_apply_model)
1076
1077
def __enter__(self):
1078
return self
tests/model_executor/test_model_load_with_params.py
+9
-3
@@ -22,10 +22,12 @@ REVISION_ROBERTA = os.environ.get("REVISION", "main")
22
23
@pytest.mark.skipif(current_platform.is_rocm(),
24
reason="Xformers backend is not supported on ROCm.")
25
-def test_model_loading_with_params(vllm_runner):
25
+def test_model_loading_with_params(vllm_runner, monkeypatch):
26
"""
27
Test parameter weight loading with tp>1.
28
"""
29
+ # to use apply_model
30
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
31
with vllm_runner(model_name=MODEL_NAME,
32
revision=REVISION,
33
dtype="float16",
@@ -61,10 +63,12 @@ def test_model_loading_with_params(vllm_runner):
63
64
@pytest.mark.skipif(current_platform.is_rocm(),
65
reason="Xformers backend is not supported on ROCm.")
64
-def test_roberta_model_loading_with_params(vllm_runner):
66
+def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
67
"""
68
Test parameter weight loading with tp>1.
69
"""
70
+ # to use apply_model
71
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
72
with vllm_runner(model_name=MODEL_NAME_ROBERTA,
73
revision=REVISION_ROBERTA,
74
dtype="float16",
@@ -101,10 +105,12 @@ def test_roberta_model_loading_with_params(vllm_runner):
105
106
@pytest.mark.skipif(current_platform.is_rocm(),
107
reason="Xformers backend is not supported on ROCm.")
104
-def test_facebook_roberta_model_loading_with_params(vllm_runner):
108
+def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch):
109
"""
110
Test loading roberta-base model with no lm_head.
111
"""
112
+ # to use apply_model
113
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
114
model_name = "FacebookAI/roberta-base"
115
with vllm_runner(model_name=model_name,
116
dtype="float16",
tests/models/language/pooling/test_embedding.py
+3
-11
@@ -39,17 +39,9 @@ def v1(run_with_both_engines):
39
pytest.param("ssmits/Qwen2-7B-Instruct-embed-base",
40
marks=[pytest.mark.skip_v0, pytest.mark.cpu_model]),
41
# [Encoder-only]
42
- pytest.param(
43
- "BAAI/bge-base-en-v1.5",
44
- marks=[
45
- # CPU only supports V1
46
- pytest.mark.core_model,
47
- pytest.mark.skip_v1
48
- ]),
49
- pytest.param("sentence-transformers/all-MiniLM-L12-v2",
50
- marks=[pytest.mark.skip_v1]),
51
- pytest.param("intfloat/multilingual-e5-small",
52
- marks=[pytest.mark.skip_v1]),
42
+ pytest.param("BAAI/bge-base-en-v1.5", marks=[pytest.mark.core_model]),
43
+ pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
44
+ pytest.param("intfloat/multilingual-e5-small"),
45
pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
46
marks=[pytest.mark.skip_v1]),
47
# [Cross-Encoder]
tests/models/language/pooling/test_jina.py
+8
@@ -23,6 +23,14 @@ RERANK_MODELS = [
23
]
24
25
26
+@pytest.fixture(autouse=True)
27
+def v1(run_with_both_engines):
28
+ # Simple autouse wrapper to run both engines for each test
29
+ # This can be promoted up to conftest.py to run for every
30
+ # test in a package
31
+ pass
32
+
33
+
34
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
35
def test_embed_models_mteb(hf_runner, vllm_runner,
36
model_info: EmbedModelInfo) -> None:
tests/v1/attention/utils.py
+1
@@ -93,6 +93,7 @@ def create_common_attn_metadata(
93
max_query_len=max_query_len,
94
block_table_tensor=block_table_tensor,
95
slot_mapping=slot_mapping,
96
+ causal=True,
97
)
98
99
tests/v1/test_oracle.py
-1
@@ -13,7 +13,6 @@ UNSUPPORTED_MODELS_V1 = [
13
"openai/whisper-large-v3", # transcription
14
"facebook/bart-large-cnn", # encoder decoder
15
"state-spaces/mamba-130m-hf", # mamba1
16
- "BAAI/bge-m3", # embedding
16
]
17
18
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
tests/v1/test_utils.py
+1
-2
@@ -1,9 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import re
5
-
4
import pytest
5
+import regex as re
6
import requests
7
import torch
8
vllm/engine/arg_utils.py
+2
-1
@@ -1649,7 +1649,8 @@ class EngineArgs:
1649
1650
if (self.max_num_seqs is None
1651
and usage_context in default_max_num_seqs):
1652
- self.max_num_seqs = default_max_num_seqs[usage_context]
1652
+ self.max_num_seqs = min(default_max_num_seqs[usage_context],
1653
+ self.max_num_batched_tokens or sys.maxsize)
1654
1655
logger.debug("Setting max_num_seqs to %d for %s usage context.",
1656
self.max_num_seqs, use_context_value)
vllm/model_executor/models/bert.py
+7
-11
@@ -12,7 +12,6 @@ from vllm.attention import Attention, AttentionType
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, PoolerConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
-from vllm.forward_context import get_forward_context
15
from vllm.model_executor.layers.activation import get_act_fn
16
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
QKVParallelLinear,
@@ -60,7 +59,6 @@ class BertEmbedding(nn.Module):
59
def forward(
60
self,
61
input_ids: torch.Tensor,
63
- seq_lens: torch.Tensor,
62
position_ids: torch.Tensor,
63
token_type_ids: Optional[torch.Tensor] = None,
64
) -> torch.Tensor:
@@ -119,7 +117,6 @@ class BertPooler(Pooler):
117
return pooled_output
118
119
122
-@support_torch_compile
120
class BertEncoder(nn.Module):
121
122
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
@@ -337,6 +334,7 @@ class BertOutput(nn.Module):
334
return hidden_states
335
336
337
+@support_torch_compile
338
class BertModel(nn.Module, SupportsQuant):
339
340
is_pooling_model = True
@@ -368,13 +366,9 @@ class BertModel(nn.Module, SupportsQuant):
366
if inputs_embeds is not None:
367
hidden_states = inputs_embeds
368
else:
371
- attn_metadata = get_forward_context().attn_metadata
372
- assert hasattr(attn_metadata, "seq_lens_tensor")
373
- hidden_states = self.embeddings(
374
- input_ids=input_ids,
375
- seq_lens=attn_metadata.seq_lens_tensor,
376
- position_ids=position_ids,
377
- token_type_ids=token_type_ids)
369
+ hidden_states = self.embeddings(input_ids=input_ids,
370
+ position_ids=position_ids,
371
+ token_type_ids=token_type_ids)
372
return self.encoder(hidden_states)
373
374
def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
@@ -447,7 +441,7 @@ class BertPoolingModel(BertModel):
441
return loaded_params
442
443
450
-class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
444
+class BertEmbeddingModel(nn.Module, SupportsQuant):
445
"""A model that uses Bert to provide embedding functionalities.
446
447
This class encapsulates the BertModel and provides an interface for
@@ -474,11 +468,13 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
468
self,
469
input_ids: Optional[torch.Tensor],
470
positions: torch.Tensor,
471
+ token_type_ids: Optional[torch.Tensor] = None,
472
intermediate_tensors: Optional[IntermediateTensors] = None,
473
inputs_embeds: Optional[torch.Tensor] = None,
474
) -> torch.Tensor:
475
return self.model(input_ids=input_ids,
476
position_ids=positions,
477
+ token_type_ids=token_type_ids,
478
inputs_embeds=inputs_embeds,
479
intermediate_tensors=intermediate_tensors)
480
vllm/model_executor/models/roberta.py
+64
-21
@@ -9,6 +9,7 @@ from torch import nn
9
from transformers import RobertaConfig
10
11
from vllm.config import VllmConfig
12
+from vllm.forward_context import get_forward_context
13
from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,
14
DispatchPooler, Pooler)
15
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -51,33 +52,12 @@ class RobertaEmbedding(nn.Module):
52
def forward(
53
self,
54
input_ids: torch.Tensor,
54
- seq_lens: torch.Tensor,
55
position_ids: torch.Tensor,
56
token_type_ids: Optional[torch.Tensor] = None,
57
) -> torch.Tensor:
58
input_shape = input_ids.size()
59
inputs_embeds = self.word_embeddings(input_ids)
60
61
- # Replace position ids because in RoBERTa models
62
- # they have to start at padding_idx + 1 and ignore
63
- # existing padding tokens
64
- # References:
65
- # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133
66
- # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669
67
- seq_lens_list = seq_lens.tolist()
68
- new_pos_list = []
69
- for positions, tokens in zip(position_ids.split(seq_lens_list),
70
- input_ids.split(seq_lens_list)):
71
- # Verify assumption that incoming position are
72
- # always a sequence from 0 to N.
73
- expected_pos = torch.arange(positions.size()[0],
74
- dtype=torch.long,
75
- device=inputs_embeds.device)
76
- assert torch.equal(positions, expected_pos)
77
- new_pos_list.append(
78
- create_position_ids_from_input_ids(tokens, self.padding_idx))
79
- position_ids = torch.cat(new_pos_list)
80
-
61
# Position embeddings.
62
position_embeddings = self.position_embeddings(position_ids)
63
if token_type_ids is None:
@@ -119,6 +99,32 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
99
_pooler: An instance of Pooler used for pooling operations.
100
"""
101
102
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
103
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
104
+ self.padding_idx = vllm_config.model_config.hf_config.pad_token_id
105
+
106
+ def forward(
107
+ self,
108
+ input_ids: Optional[torch.Tensor],
109
+ positions: torch.Tensor,
110
+ token_type_ids: Optional[torch.Tensor] = None,
111
+ intermediate_tensors: Optional[IntermediateTensors] = None,
112
+ inputs_embeds: Optional[torch.Tensor] = None,
113
+ ) -> torch.Tensor:
114
+
115
+ # Fix Roberta positions here outside of the CUDA graph.
116
+ # Because we need the to extract the sequences from
117
+ # input_ids the control flow is data dependent.
118
+ replace_roberta_positions(input_ids=input_ids,
119
+ position_ids=positions,
120
+ padding_idx=self.padding_idx)
121
+
122
+ return self.model(input_ids=input_ids,
123
+ position_ids=positions,
124
+ token_type_ids=token_type_ids,
125
+ inputs_embeds=inputs_embeds,
126
+ intermediate_tensors=intermediate_tensors)
127
+
128
def _build_model(self,
129
vllm_config: VllmConfig,
130
prefix: str = "") -> Union[BertModel, BertWithRope]:
@@ -175,6 +181,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
181
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
182
super().__init__()
183
config = vllm_config.model_config.hf_config
184
+ self.padding_idx = vllm_config.model_config.hf_config.pad_token_id
185
186
self.num_labels = config.num_labels
187
self.roberta = BertModel(vllm_config=vllm_config,
@@ -216,6 +223,9 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
223
inputs_embeds: Optional[torch.Tensor] = None,
224
token_type_ids: Optional[torch.Tensor] = None,
225
) -> torch.Tensor:
226
+ replace_roberta_positions(input_ids=input_ids,
227
+ position_ids=positions,
228
+ padding_idx=self.padding_idx)
229
return self.roberta(input_ids=input_ids,
230
position_ids=positions,
231
inputs_embeds=inputs_embeds,
@@ -245,3 +255,36 @@ def create_position_ids_from_input_ids(input_ids,
255
past_key_values_length) * mask
256
257
return incremental_indices.long() + padding_idx
258
+
259
+
260
+def replace_roberta_positions(input_ids: torch.Tensor,
261
+ position_ids: torch.Tensor,
262
+ padding_idx: int) -> None:
263
+
264
+ seq_lens: Optional[torch.Tensor] = None
265
+ attn_metadata = get_forward_context().attn_metadata
266
+ if attn_metadata is not None: # can be None during warmup
267
+ if isinstance(attn_metadata, dict):
268
+ attn_metadata = next(iter(attn_metadata.values()))
269
+ # TODO: remove "seq_lens_tensor" after V0 is removed
270
+ seq_lens = getattr(attn_metadata, "seq_lens_tensor",
271
+ getattr(attn_metadata, "seq_lens", None))
272
+
273
+ if seq_lens is not None:
274
+ assert isinstance(seq_lens, torch.Tensor)
275
+
276
+ # Replace position ids because in RoBERTa models
277
+ # they have to start at padding_idx + 1 and ignore
278
+ # existing padding tokens
279
+ # References:
280
+ # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133
281
+ # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669
282
+ token_list = torch.split(input_ids[:torch.sum(seq_lens)],
283
+ seq_lens.tolist())
284
+
285
+ offset = 0
286
+ for tokens in token_list:
287
+ length = tokens.shape[0]
288
+ position_ids[offset:offset+length] = \
289
+ create_position_ids_from_input_ids(tokens, padding_idx)
290
+ offset = offset + length
vllm/v1/attention/backends/flash_attn.py
+84
-7
@@ -130,6 +130,8 @@ class FlashAttentionMetadata:
130
prefix_scheduler_metadata: Optional[torch.Tensor] = None
131
max_num_splits: int = 0
132
133
+ causal: bool = True
134
+
135
136
def _get_sliding_window_configs(
137
vllm_config: VllmConfig) -> set[Optional[tuple[int, int]]]:
@@ -213,6 +215,7 @@ class FlashAttentionMetadataBuilder(
215
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
216
block_table_tensor = common_attn_metadata.block_table_tensor
217
slot_mapping = common_attn_metadata.slot_mapping
218
+ causal = common_attn_metadata.causal
219
220
# the overhead of the aot schedule is not worth it for spec-decode
221
aot_schedule = self.aot_schedule and not fast_build
@@ -288,7 +291,7 @@ class FlashAttentionMetadataBuilder(
291
max_query_len=max_query_len,
292
seqlens=seq_lens,
293
max_seq_len=max_seq_len,
291
- causal=True)
294
+ causal=causal)
295
296
if self.use_full_cuda_graph:
297
assert scheduler_metadata is not None
@@ -326,7 +329,7 @@ class FlashAttentionMetadataBuilder(
329
suffix_kv_lens=suffix_kv_lens,
330
prefix_scheduler_metadata=prefix_scheduler_metadata,
331
max_num_splits=max_num_splits,
329
- )
332
+ causal=causal)
333
return attn_metadata
334
335
def can_run_in_cudagraph(
@@ -375,11 +378,14 @@ class FlashAttentionImpl(AttentionImpl):
378
379
FlashAttentionBackend.validate_head_size(head_size)
380
378
- if attn_type != AttentionType.DECODER:
379
- raise NotImplementedError("Encoder self-attention and "
380
- "encoder/decoder cross-attention "
381
- "are not implemented for "
381
+ if attn_type not in [
382
+ AttentionType.DECODER, AttentionType.ENCODER_ONLY
383
+ ]:
384
+ raise NotImplementedError("Encoder/decoder cross-attention "
385
+ "is not implemented for "
386
"FlashAttentionImpl")
387
+
388
+ self.attn_type = attn_type
389
self.vllm_flash_attn_version = get_flash_attn_version()
390
if is_quantized_kv_cache(self.kv_cache_dtype) \
391
and not flash_attn_supports_fp8():
@@ -422,6 +428,8 @@ class FlashAttentionImpl(AttentionImpl):
428
# Profiling run.
429
return output
430
431
+ attn_type = self.attn_type
432
+
433
# IMPORTANT!
434
# NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
435
# eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
@@ -432,6 +440,18 @@ class FlashAttentionImpl(AttentionImpl):
440
# performance to make sure it does not introduce any overhead.
441
442
num_actual_tokens = attn_metadata.num_actual_tokens
443
+
444
+ # Handle encoder attention differently - no KV cache needed
445
+ if attn_type in (AttentionType.ENCODER_ONLY, ):
446
+ # For encoder attention,
447
+ # we use direct Q, K, V tensors without caching
448
+ return self._forward_encoder_attention(query[:num_actual_tokens],
449
+ key[:num_actual_tokens],
450
+ value[:num_actual_tokens],
451
+ output[:num_actual_tokens],
452
+ attn_metadata, layer)
453
+
454
+ # For decoder and cross-attention, use KV cache as before
455
key_cache, value_cache = kv_cache.unbind(0)
456
457
if self.kv_sharing_target_layer_name is None:
@@ -483,7 +503,7 @@ class FlashAttentionImpl(AttentionImpl):
503
seqused_k=seqused_k,
504
max_seqlen_k=max_seqlen_k,
505
softmax_scale=self.scale,
486
- causal=True,
506
+ causal=attn_metadata.causal,
507
alibi_slopes=self.alibi_slopes,
508
window_size=self.sliding_window,
509
block_table=block_table,
@@ -524,6 +544,63 @@ class FlashAttentionImpl(AttentionImpl):
544
)
545
return output
546
547
+ def _forward_encoder_attention(
548
+ self,
549
+ query: torch.Tensor,
550
+ key: torch.Tensor,
551
+ value: torch.Tensor,
552
+ output: torch.Tensor,
553
+ attn_metadata: FlashAttentionMetadata,
554
+ layer: torch.nn.Module,
555
+ ) -> torch.Tensor:
556
+ """Forward pass for encoder attention without KV cache.
557
+
558
+ Args:
559
+ query: shape = [num_encoder_tokens, num_heads, head_size]
560
+ key: shape = [num_encoder_tokens, num_kv_heads, head_size]
561
+ value: shape = [num_encoder_tokens, num_kv_heads, head_size]
562
+ output: shape = [num_encoder_tokens, num_heads, head_size]
563
+ attn_metadata: Encoder attention metadata
564
+ layer: The attention layer
565
+ """
566
+ # For encoder attention, process FP8 quantization if needed
567
+ if self.kv_cache_dtype.startswith("fp8"):
568
+ raise NotImplementedError(
569
+ "quantization is not supported for encoder attention")
570
+
571
+ # Use encoder-specific metadata for sequence information
572
+ cu_seqlens_q = attn_metadata.query_start_loc
573
+ cu_seqlens_k = attn_metadata.query_start_loc
574
+ max_seqlen_q = attn_metadata.max_query_len
575
+ max_seqlen_k = attn_metadata.max_query_len
576
+
577
+ descale_shape = (
578
+ cu_seqlens_q.shape[0] - 1, # type: ignore[union-attr]
579
+ self.num_kv_heads)
580
+
581
+ # Call flash attention directly on Q, K, V tensors
582
+ flash_attn_varlen_func(
583
+ q=query,
584
+ k=key,
585
+ v=value,
586
+ out=output,
587
+ cu_seqlens_q=cu_seqlens_q,
588
+ cu_seqlens_k=cu_seqlens_k,
589
+ max_seqlen_q=max_seqlen_q,
590
+ max_seqlen_k=max_seqlen_k,
591
+ softmax_scale=self.scale,
592
+ causal=False, # Encoder attention is bidirectional
593
+ alibi_slopes=self.alibi_slopes,
594
+ window_size=self.sliding_window,
595
+ softcap=self.logits_soft_cap,
596
+ fa_version=self.vllm_flash_attn_version,
597
+ q_descale=layer._q_scale.expand(descale_shape),
598
+ k_descale=layer._k_scale.expand(descale_shape),
599
+ v_descale=layer._v_scale.expand(descale_shape),
600
+ )
601
+
602
+ return output
603
+
604
605
def use_cascade_attention(
606
common_prefix_len: int,
vllm/v1/attention/backends/utils.py
+3
@@ -59,6 +59,8 @@ class CommonAttentionMetadata:
59
block_table_tensor: torch.Tensor
60
slot_mapping: torch.Tensor
61
62
+ causal: bool = True
63
+
64
65
M = TypeVar("M")
66
@@ -395,6 +397,7 @@ def make_local_attention_virtual_batches(
397
max_query_len=seqlens_q_local.max(),
398
block_table_tensor=block_table_local,
399
slot_mapping=common_attn_metadata.slot_mapping,
400
+ causal=True,
401
)
402
403
vllm/v1/engine/core.py
+6
@@ -111,6 +111,12 @@ class EngineCore:
111
"compatibility may not be maintained.",
112
vllm_config.scheduler_config.scheduler_cls)
113
114
+ if len(kv_cache_config.kv_cache_groups) == 0:
115
+ # Encoder models without KV cache don't support
116
+ # chunked prefill. But do SSM models?
117
+ logger.info("Disabling chunked prefill for model without KVCache")
118
+ vllm_config.scheduler_config.chunked_prefill_enabled = False
119
+
120
self.scheduler: SchedulerInterface = Scheduler(
121
vllm_config=vllm_config,
122
kv_cache_config=kv_cache_config,
vllm/v1/spec_decode/eagle.py
+1
@@ -330,6 +330,7 @@ class EagleProposer:
330
max_query_len=new_query_len_per_req.max().item(),
331
block_table_tensor=common_attn_metadata.block_table_tensor,
332
slot_mapping=common_attn_metadata.slot_mapping[token_indices],
333
+ causal=True,
334
)
335
336
return spec_common_attn_metadata, token_indices
vllm/v1/worker/cpu_model_runner.py
+4
@@ -4,6 +4,7 @@ from contextlib import contextmanager
4
from typing import Any
5
6
import torch
7
+import torch.nn as nn
8
9
from vllm.config import VllmConfig
10
from vllm.logger import init_logger
@@ -59,6 +60,9 @@ class CPUModelRunner(GPUModelRunner):
60
self.scheduler_config,
61
self.lora_config, self.device)
62
63
+ def get_model(self) -> nn.Module:
64
+ return self.model
65
+
66
def warming_up_model(self) -> None:
67
logger.info("Warming up model for the compilation...")
68
# Only generate graph for the generic shape
vllm/v1/worker/gpu_model_runner.py
+147
-39
@@ -126,6 +126,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
126
127
self.is_multimodal_model = model_config.is_multimodal_model
128
self.is_pooling_model = model_config.pooler_config is not None
129
+ self.is_encoder_only_model = False
130
self.model_supports_multimodal_raw_input = (
131
model_config.model_supports_multimodal_raw_input)
132
self.max_model_len = model_config.max_model_len
@@ -735,6 +736,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
736
spec_decode_common_attn_metadata = None
737
738
attn_metadata: dict[str, Any] = {}
739
+
740
+ # Prepare encoder attention metadata separately
741
+ # (encoder layers are not in KV cache groups)
742
+ if self.is_encoder_only_model:
743
+ common_attn_metadata, encoder_attn_metadata = \
744
+ self._build_encoder_only_attn_metadata(
745
+ scheduler_output)
746
+
747
+ # Add encoder attention metadata for all encoder layers
748
+ attention_layers = get_layers_from_vllm_config(
749
+ self.vllm_config, Attention)
750
+ for layer_name, attn_module in attention_layers.items():
751
+ if attn_module.attn_type == AttentionType.ENCODER_ONLY:
752
+ attn_metadata[layer_name] = encoder_attn_metadata
753
+
754
# Prepare the attention metadata for each KV cache group and make layers
755
# in the same group share the same metadata.
756
for kv_cache_group_id, kv_cache_group_spec in enumerate(
@@ -760,6 +776,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
776
max_query_len=max_num_scheduled_tokens,
777
block_table_tensor=blk_table_tensor,
778
slot_mapping=slot_mapping,
779
+ causal=True,
780
)
781
782
if self.speculative_config and \
@@ -2102,7 +2119,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2119
block_table_tensor=self.input_batch.block_table[
2120
kv_cache_group_id].get_device_tensor()[:num_reqs],
2121
slot_mapping=self.input_batch.
2105
- block_table[kv_cache_group_id].slot_mapping[:num_tokens])
2122
+ block_table[kv_cache_group_id].slot_mapping[:num_tokens],
2123
+ causal=True)
2124
2125
attn_metadata_i = self.attn_metadata_builders[
2126
kv_cache_group_id].build_for_cudagraph_capture(
@@ -2466,6 +2484,49 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2484
logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
2485
elapsed_time, cuda_graph_size / (1 << 30))
2486
2487
+ def _initialize_single_attn_backend(
2488
+ self, kv_cache_spec: KVCacheSpec
2489
+ ) -> tuple[AttentionBackend, AttentionMetadataBuilder]:
2490
+ if isinstance(kv_cache_spec, AttentionSpec):
2491
+ attn_backend_i = get_attn_backend(
2492
+ kv_cache_spec.head_size,
2493
+ self.dtype,
2494
+ kv_cache_spec.dtype,
2495
+ kv_cache_spec.block_size,
2496
+ self.model_config.is_attention_free,
2497
+ use_mla=kv_cache_spec.use_mla,
2498
+ )
2499
+ if attn_backend_i is None:
2500
+ error_msg = (f"Error with get_attn_backend: "
2501
+ f"{kv_cache_spec.head_size=}, "
2502
+ f"{self.dtype=}, {kv_cache_spec.dtype=}, "
2503
+ f"{kv_cache_spec.block_size=}, "
2504
+ f"{self.model_config.is_attention_free=}, "
2505
+ f"{kv_cache_spec.use_mla=}")
2506
+ logger.error(error_msg)
2507
+ raise NotImplementedError(
2508
+ "Non-Attention backend is not supported by V1 "
2509
+ "GPUModelRunner.")
2510
+ elif isinstance(kv_cache_spec, MambaSpec):
2511
+ attn_backend_i = Mamba2AttentionBackend
2512
+ else:
2513
+ raise ValueError(
2514
+ f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2515
+
2516
+ attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2517
+ kv_cache_spec,
2518
+ self.vllm_config,
2519
+ self.device,
2520
+ )
2521
+
2522
+ if (self.full_cuda_graph
2523
+ and not attn_metadata_builder_i.full_cudagraph_supported):
2524
+ raise ValueError(
2525
+ f"Full CUDAGraph not supported for "
2526
+ f"{attn_backend_i.__name__}. Turn off CompilationConfig."
2527
+ f"full_cuda_graph or use a different attention backend.")
2528
+ return attn_backend_i, attn_metadata_builder_i
2529
+
2530
def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
2531
"""
2532
Initialize the attention backends and attention metadata builders.
@@ -2476,48 +2537,45 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2537
for i, kv_cache_group_spec in enumerate(
2538
kv_cache_config.kv_cache_groups):
2539
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2479
- if isinstance(kv_cache_spec, AttentionSpec):
2480
- attn_backend_i = get_attn_backend(
2481
- kv_cache_spec.head_size,
2482
- self.dtype,
2483
- kv_cache_spec.dtype,
2484
- kv_cache_spec.block_size,
2485
- self.model_config.is_attention_free,
2486
- use_mla=kv_cache_spec.use_mla,
2487
- )
2488
- if attn_backend_i is None:
2489
- error_msg = (f"Error with get_attn_backend: "
2490
- f"{kv_cache_spec.head_size=}, "
2491
- f"{self.dtype=}, {kv_cache_spec.dtype=}, "
2492
- f"{kv_cache_spec.block_size=}, "
2493
- f"{self.model_config.is_attention_free=}, "
2494
- f"{kv_cache_spec.use_mla=}")
2495
- logger.error(error_msg)
2496
- raise NotImplementedError(
2497
- "Non-Attention backend is not supported by V1 "
2498
- "GPUModelRunner.")
2499
- elif isinstance(kv_cache_spec, MambaSpec):
2500
- attn_backend_i = Mamba2AttentionBackend
2501
- else:
2502
- raise ValueError(
2503
- f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2504
-
2505
- attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2506
- kv_cache_spec,
2507
- self.vllm_config,
2508
- self.device,
2509
- )
2510
-
2511
- if (self.full_cuda_graph
2512
- and not attn_metadata_builder_i.full_cudagraph_supported):
2513
- raise ValueError(
2514
- f"Full CUDAGraph not supported for "
2515
- f"{attn_backend_i.__name__}. Turn off CompilationConfig."
2516
- f"full_cuda_graph or use a different attention backend.")
2540
2541
+ attn_backend_i, attn_metadata_builder_i = \
2542
+ self._initialize_single_attn_backend(kv_cache_spec)
2543
self.attn_backends.append(attn_backend_i)
2544
self.attn_metadata_builders.append(attn_metadata_builder_i)
2545
2546
+ if len(self.attn_backends) > 0:
2547
+ return
2548
+
2549
+ # Check if model is encoder-only
2550
+ block_size = self.vllm_config.cache_config.block_size
2551
+ use_mla = self.vllm_config.model_config.use_mla
2552
+ attn_specs = list[AttentionSpec]()
2553
+ attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
2554
+ for attn_module in attn_layers.values():
2555
+
2556
+ if attn_module.attn_type == AttentionType.ENCODER_ONLY:
2557
+ assert attn_module.sliding_window is None, "Sliding "
2558
+ "window attention is not supported for encoder-only models"
2559
+
2560
+ attn_specs.append(
2561
+ FullAttentionSpec(block_size=block_size,
2562
+ num_kv_heads=attn_module.num_kv_heads,
2563
+ head_size=attn_module.head_size,
2564
+ dtype=self.kv_cache_dtype,
2565
+ use_mla=use_mla))
2566
+ else:
2567
+ raise ValueError("Expected only encoder-only layers")
2568
+
2569
+ if len(attn_specs) > 0:
2570
+ assert len(attn_specs) == len(attn_layers), \
2571
+ "All or none of the layers are expected to be encoder-only"
2572
+
2573
+ attn_backend, attn_metadata_builder = \
2574
+ self._initialize_single_attn_backend(attn_specs[0])
2575
+ self.attn_backends.append(attn_backend)
2576
+ self.attn_metadata_builders.append(attn_metadata_builder)
2577
+ self.is_encoder_only_model = True
2578
+
2579
def may_reinitialize_input_batch(self,
2580
kv_cache_config: KVCacheConfig) -> None:
2581
"""
@@ -2833,3 +2891,53 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2891
page_size_padded=page_size_padded)
2892
2893
return kv_cache_spec
2894
+
2895
+ def _build_encoder_only_attn_metadata(
2896
+ self, scheduler_output: "SchedulerOutput") -> \
2897
+ tuple[CommonAttentionMetadata, Any]:
2898
+ """Prepare encoder attention metadata for encoder-only models.
2899
+
2900
+ Args:
2901
+ scheduler_output: Scheduler output
2902
+
2903
+ Returns:
2904
+ dict[str, Any]: Encoder attention metadata
2905
+ """
2906
+ num_reqs = self.input_batch.num_reqs
2907
+ total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2908
+
2909
+ # Get the number of scheduled tokens for each request.
2910
+ req_ids = self.input_batch.req_ids
2911
+ tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
2912
+ max_num_scheduled_tokens = max(tokens)
2913
+
2914
+ # Use the first attention metadata builder
2915
+ # to create encoder attention metadata
2916
+ builder = self.attn_metadata_builders[0]
2917
+
2918
+ dummy_block_table = torch.zeros((num_reqs, 1),
2919
+ dtype=torch.int32,
2920
+ device=self.device)
2921
+ dummy_slot_mapping = torch.zeros((total_num_scheduled_tokens, ),
2922
+ dtype=torch.int32,
2923
+ device=self.device)
2924
+
2925
+ common_metadata = CommonAttentionMetadata(
2926
+ query_start_loc=self.query_start_loc[:num_reqs + 1],
2927
+ query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
2928
+ seq_lens=self.seq_lens[:num_reqs],
2929
+ seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
2930
+ num_computed_tokens_cpu=self.input_batch.
2931
+ num_computed_tokens_cpu_tensor[:num_reqs],
2932
+ num_reqs=num_reqs,
2933
+ num_actual_tokens=total_num_scheduled_tokens,
2934
+ max_query_len=max_num_scheduled_tokens,
2935
+ block_table_tensor=dummy_block_table,
2936
+ slot_mapping=dummy_slot_mapping,
2937
+ causal=False,
2938
+ )
2939
+
2940
+ return common_metadata, builder.build(
2941
+ common_prefix_len=0, # No cascade for encoder
2942
+ common_attn_metadata=common_metadata,
2943
+ )