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 15, 08:30 UTC
Enable hybrid attention models for Transformers backend (#18494)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
May 23, 2025 at 04:12 UTC
4b0da7b60e35bbeb2cbd656c54966ee2ceb4c4bd
4 files changed
+105
-29
docs/source/contributing/model/basic.md
+1
-1
@@ -117,7 +117,7 @@ For models with interleaving sliding windows (e.g. `google/gemma-2-2b-it` and `m
117
118
To support a model with interleaving sliding windows, we need to take care of the following details:
119
120
-- Make sure [this line](https://github.com/vllm-project/vllm/blob/996357e4808ca5eab97d4c97c7d25b3073f46aab/vllm/config.py#L308) evaluates `has_interleaved_attention` to `True` for this model, and set `self.hf_text_config.interleaved_sliding_window` to the format of interleaving sliding windows the model can understand. Then, `self.hf_text_config.sliding_window` will be deleted, and the model will be treated as a full-attention model.
120
+- Make sure the model's `config.json` contains `sliding_window_pattern`. vLLM then sets `self.hf_text_config.interleaved_sliding_window` to the value of `self.hf_text_config.sliding_window` and deletes `sliding_window` from `self.hf_text_config`. The model will then be treated as a full-attention model.
121
- In the modeling code, parse the correct sliding window value for every layer, and pass it to the attention layer's `per_layer_sliding_window` argument. For reference, check [this line](https://github.com/vllm-project/vllm/blob/996357e4808ca5eab97d4c97c7d25b3073f46aab/vllm/model_executor/models/llama.py#L171).
122
123
With these two steps, interleave sliding windows should work with the model.
tests/models/test_transformers.py
+42
-14
@@ -1,37 +1,50 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""Test the functionality of the Transformers backend."""
3
+from typing import Any, Optional, Union
4
+
5
import pytest
6
7
from vllm.platforms import current_platform
8
9
from ..conftest import HfRunner, VllmRunner
10
+from ..core.block.e2e.test_correctness_sliding_window import prep_prompts
11
from ..utils import multi_gpu_test
12
from .utils import check_logprobs_close
13
14
15
def check_implementation(
13
- hf_runner: type[HfRunner],
14
- vllm_runner: type[VllmRunner],
16
+ runner_ref: type[Union[HfRunner, VllmRunner]],
17
+ runner_test: type[VllmRunner],
18
example_prompts: list[str],
19
model: str,
20
+ kwargs_ref: Optional[dict[str, Any]] = None,
21
+ kwargs_test: Optional[dict[str, Any]] = None,
22
**kwargs,
23
):
24
+ if kwargs_ref is None:
25
+ kwargs_ref = {}
26
+ if kwargs_test is None:
27
+ kwargs_test = {}
28
+
29
max_tokens = 32
30
num_logprobs = 5
31
22
- with vllm_runner(model, **kwargs) as vllm_model:
23
- vllm_outputs = vllm_model.generate_greedy_logprobs(
24
- example_prompts, max_tokens, num_logprobs)
32
+ args = (example_prompts, max_tokens, num_logprobs)
33
+
34
+ with runner_test(model, **kwargs_test, **kwargs) as model_test:
35
+ outputs_test = model_test.generate_greedy_logprobs(*args)
36
26
- with hf_runner(model) as hf_model:
27
- hf_outputs = hf_model.generate_greedy_logprobs_limit(
28
- example_prompts, max_tokens, num_logprobs)
37
+ with runner_ref(model, **kwargs_ref) as model_ref:
38
+ if isinstance(model_ref, VllmRunner):
39
+ outputs_ref = model_ref.generate_greedy_logprobs(*args)
40
+ else:
41
+ outputs_ref = model_ref.generate_greedy_logprobs_limit(*args)
42
43
check_logprobs_close(
31
- outputs_0_lst=hf_outputs,
32
- outputs_1_lst=vllm_outputs,
33
- name_0="hf",
34
- name_1="vllm",
44
+ outputs_0_lst=outputs_ref,
45
+ outputs_1_lst=outputs_test,
46
+ name_0="ref",
47
+ name_1="test",
48
)
49
50
@@ -58,6 +71,18 @@ def test_models(
71
model_impl=model_impl)
72
73
74
+def test_hybrid_attention(vllm_runner: type[VllmRunner]) -> None:
75
+ prompts, _, _ = prep_prompts(4, (800, 801))
76
+ kwargs_ref = {"max_model_len": 8192, "enforce_eager": True}
77
+ kwargs_test = {"model_impl": "transformers", **kwargs_ref}
78
+ check_implementation(vllm_runner,
79
+ vllm_runner,
80
+ prompts,
81
+ model="hmellor/tiny-random-Gemma2ForCausalLM",
82
+ kwargs_ref=kwargs_ref,
83
+ kwargs_test=kwargs_test)
84
+
85
+
86
@multi_gpu_test(num_gpus=2)
87
def test_distributed(
88
hf_runner: type[HfRunner],
@@ -65,8 +90,11 @@ def test_distributed(
90
example_prompts,
91
):
92
kwargs = {"model_impl": "transformers", "tensor_parallel_size": 2}
68
- check_implementation(hf_runner, vllm_runner, example_prompts,
69
- "meta-llama/Llama-3.2-1B-Instruct", **kwargs)
93
+ check_implementation(hf_runner,
94
+ vllm_runner,
95
+ example_prompts,
96
+ "meta-llama/Llama-3.2-1B-Instruct",
97
+ kwargs_test=kwargs)
98
99
100
@pytest.mark.skipif(
vllm/config.py
+11
-8
@@ -533,13 +533,17 @@ class ModelConfig:
533
self.model, hf_token=self.hf_token, revision=self.revision)
534
self.dtype = _get_and_verify_dtype(self.hf_config, self.dtype)
535
536
- interleaved_attn_models = ["gemma2", "gemma3_text", "cohere2"]
536
+ # Workaround for Gemma 2 which uses interleaved sliding window
537
+ # attention, but it's not specified in its config. TODO: remove this
538
+ # when Gemma 2 is fixed in Transformers.
539
+ if self.hf_text_config.model_type == "gemma2":
540
+ self.hf_text_config.sliding_window_pattern = 2
541
+
542
sliding_window = getattr(self.hf_text_config, "sliding_window", None)
538
- has_interleaved_attention = (sliding_window is not None) and (
539
- isinstance(sliding_window, list) or
540
- (self.hf_text_config.model_type in interleaved_attn_models))
543
+ sliding_window_pattern = getattr(self.hf_text_config,
544
+ "sliding_window_pattern", None)
545
542
- if (not self.disable_sliding_window and has_interleaved_attention):
546
+ if not (self.disable_sliding_window or sliding_window_pattern is None):
547
if (backend :=
548
envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER"):
549
sliding_window_len_min = get_min_sliding_window(
@@ -1037,8 +1041,7 @@ class ModelConfig:
1041
if self.use_async_output_proc:
1042
self.use_async_output_proc = False
1043
1040
- def get_hf_config_sliding_window(
1041
- self) -> Union[Optional[int], list[Optional[int]]]:
1044
+ def get_hf_config_sliding_window(self) -> Optional[int]:
1045
"""Get the sliding window size, or None if disabled."""
1046
1047
# Some models, like Qwen2 and Qwen1.5, use `use_sliding_window` in
@@ -1049,7 +1052,7 @@ class ModelConfig:
1052
return None
1053
return getattr(self.hf_text_config, "sliding_window", None)
1054
1052
- def get_sliding_window(self) -> Optional[Union[int, list[Optional[int]]]]:
1055
+ def get_sliding_window(self) -> Optional[int]:
1056
"""Get the sliding window size, or None if disabled.
1057
"""
1058
# If user disables sliding window, return None.
vllm/model_executor/models/transformers.py
+51
-6
@@ -16,6 +16,7 @@
16
"""Wrapper around `transformers` models"""
17
import re
18
from collections.abc import Iterable
19
+from contextlib import nullcontext
20
from typing import Literal, Optional, Union
21
22
import torch
@@ -110,6 +111,33 @@ def replace_linear_class(
111
)
112
113
114
+class ConfigOverride:
115
+ """Context manager to temporarily override config attributes."""
116
+
117
+ def __init__(self, config: PretrainedConfig, **kwargs):
118
+ self.config = config
119
+ self.kwargs = kwargs
120
+ self.kwargs_original = {}
121
+ self.kwargs_delete = set()
122
+
123
+ def __enter__(self):
124
+ """Override config attributes."""
125
+ for key, value in self.kwargs.items():
126
+ if not hasattr(self.config, key):
127
+ self.kwargs_delete.add(key)
128
+ self.kwargs_original[key] = getattr(self.config, key, None)
129
+ setattr(self.config, key, value)
130
+ return self.config
131
+
132
+ def __exit__(self, exc_type, exc_value, traceback):
133
+ """Restore original config attributes."""
134
+ for key, value in self.kwargs_original.items():
135
+ if key in self.kwargs_delete:
136
+ delattr(self.config, key)
137
+ else:
138
+ setattr(self.config, key, value)
139
+
140
+
141
class TransformersModel(nn.Module):
142
143
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -135,8 +163,17 @@ class TransformersModel(nn.Module):
163
self.pp_rank = self.pp_group.rank_in_group
164
self.tp_size = get_tensor_model_parallel_world_size()
165
166
+ # vLLM handles interleaved sliding window attention by creating a new
167
+ # interleaved_sliding_window attribute and deleting the sliding_window
168
+ # attribute. This breaks the constructors in Transformers so we
169
+ # temporarily add the attribute back to construct the model.
170
+ config_override = nullcontext()
171
+ if hasattr(config, "interleaved_sliding_window"):
172
+ config_override = ConfigOverride(
173
+ config, sliding_window=config.interleaved_sliding_window)
174
+
175
# Use meta device to delay allocating GPU tensors
139
- with torch.device("meta"):
176
+ with torch.device("meta"), config_override:
177
# FIXME(Isotr0py): We need to refactor this part in the future to
178
# avoid registering an extra model layer, otherwise we will need a
179
# weights mapper to rename weights.
@@ -262,9 +299,17 @@ class TransformersModel(nn.Module):
299
num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
300
start, end = get_pp_indices(self.config.num_hidden_layers,
301
self.pp_rank, self.pp_size)
265
- return {
266
- i:
267
- Attention(
302
+
303
+ attention_instances = {}
304
+ for i in range(start, end):
305
+ # Handle interleaved sliding window attention
306
+ sliding_window = None
307
+ if (hasattr(self.config, "interleaved_sliding_window")
308
+ and hasattr(self.config, "sliding_window_pattern")
309
+ and ((i + 1) % self.config.sliding_window_pattern > 0)):
310
+ sliding_window = self.config.interleaved_sliding_window
311
+
312
+ attention_instances[i] = Attention(
313
num_heads=num_heads,
314
head_size=head_size,
315
# NOTE: We use Llama scale as default, if it's set by
@@ -273,9 +318,9 @@ class TransformersModel(nn.Module):
318
num_kv_heads=num_kv_heads,
319
cache_config=self.cache_config,
320
quant_config=self.quant_config,
321
+ per_layer_sliding_window=sliding_window,
322
prefix=f"{i}.attn")
277
- for i in range(start, end)
278
- }
323
+ return attention_instances
324
325
def init_buffers(self, module: nn.Module):
326
"""