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, 03:30 UTC
[CI/Build][CPU] Fix CPU CI and remove all CPU V0 files (#20560)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Jul 8, 2025 at 13:13 UTC
7721ef1786c49ec3738db5e61821183ee969d2a2
9 files changed
+785
-839
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+12
-12
@@ -48,10 +48,16 @@ function cpu_tests() {
48
# Run basic model test
49
docker exec cpu-test-"$NUMA_NODE" bash -c "
50
set -e
51
- pytest -v -s tests/kernels/attention/test_cache.py -m cpu_model
52
- pytest -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
53
- pytest -v -s tests/models/language/generation -m cpu_model
54
- VLLM_CPU_SGL_KERNEL=1 pytest -v -s tests/models/language/generation -m cpu_model
51
+ # Note: disable until supports V1
52
+ # pytest -v -s tests/kernels/attention/test_cache.py -m cpu_model
53
+ # pytest -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
54
+
55
+ # Note: disable Bart until supports V1
56
+ pytest -v -s tests/models/language/generation -m cpu_model \
57
+ --ignore=tests/models/language/generation/test_bart.py
58
+ VLLM_CPU_SGL_KERNEL=1 pytest -v -s tests/models/language/generation -m cpu_model \
59
+ --ignore=tests/models/language/generation/test_bart.py
60
+
61
pytest -v -s tests/models/language/pooling -m cpu_model
62
pytest -v -s tests/models/multimodal/generation \
63
--ignore=tests/models/multimodal/generation/test_mllama.py \
@@ -62,21 +68,15 @@ function cpu_tests() {
68
docker exec cpu-test-"$NUMA_NODE" bash -c "
69
set -e
70
pytest -s -v \
65
- tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_static_setup \
66
- tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_dynamic_per_token"
71
+ tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs[False-10-32-neuralmagic/Llama-3.2-1B-quantized.w8a8]"
72
73
+ # Note: disable it until supports V1
74
# Run AWQ test
75
# docker exec cpu-test-"$NUMA_NODE" bash -c "
76
# set -e
77
# VLLM_USE_V1=0 pytest -s -v \
78
# tests/quantization/test_ipex_quant.py"
79
74
- # Run chunked-prefill and prefix-cache test
75
- docker exec cpu-test-"$NUMA_NODE" bash -c "
76
- set -e
77
- pytest -s -v -k cpu_model \
78
- tests/basic_correctness/test_chunked_prefill.py"
79
-
80
# online serving
81
docker exec cpu-test-"$NUMA_NODE" bash -c "
82
set -e
tests/basic_correctness/test_chunked_prefill.py
-58
@@ -294,61 +294,3 @@ def test_with_prefix_caching(
294
name_0="w/o prefix caching",
295
name_1="with prefix caching",
296
)
297
-
298
-
299
-@pytest.mark.parametrize("model", ["facebook/opt-125m"])
300
-@pytest.mark.parametrize("dtype", ["bfloat16", "half"])
301
-@pytest.mark.parametrize("max_tokens", [32])
302
-@pytest.mark.parametrize("chunked_prefill_token_size", [1, 4, 16])
303
-@pytest.mark.parametrize("enforce_eager", [False])
304
-@pytest.mark.parametrize("attention_backend", ["TORCH_SDPA"])
305
-@pytest.mark.cpu_model
306
-@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
307
-def test_models_cpu(
308
- hf_runner: HfRunner,
309
- vllm_runner: VllmRunner,
310
- example_prompts,
311
- model: str,
312
- dtype: str,
313
- max_tokens: int,
314
- chunked_prefill_token_size: int,
315
- enforce_eager: bool,
316
- attention_backend: str,
317
- monkeypatch: pytest.MonkeyPatch,
318
-) -> None:
319
- test_models(
320
- hf_runner,
321
- vllm_runner,
322
- example_prompts,
323
- model,
324
- dtype,
325
- max_tokens,
326
- chunked_prefill_token_size,
327
- enforce_eager,
328
- 1,
329
- attention_backend,
330
- monkeypatch,
331
- )
332
-
333
-
334
-@pytest.mark.parametrize("max_tokens", [16])
335
-@pytest.mark.parametrize("enforce_eager", [False])
336
-@pytest.mark.parametrize("chunk_size", [30, 32])
337
-@pytest.mark.parametrize("dtype", ["bfloat16", "half"])
338
-@pytest.mark.cpu_model
339
-@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
340
-def test_with_prefix_caching_cpu(
341
- vllm_runner: VllmRunner,
342
- max_tokens: int,
343
- enforce_eager: bool,
344
- chunk_size: int,
345
- dtype: str,
346
-) -> None:
347
- test_with_prefix_caching(
348
- vllm_runner,
349
- max_tokens,
350
- enforce_eager,
351
- chunk_size,
352
- 1,
353
- dtype,
354
- )
tests/models/language/generation/test_common.py
+6
-2
@@ -39,7 +39,7 @@ AITER_MODEL_LIST = [
39
[
40
pytest.param(
41
"bigscience/bloom-560m", # bloom - testing alibi slopes
42
- marks=[pytest.mark.core_model, pytest.mark.cpu_model],
42
+ marks=[pytest.mark.core_model],
43
),
44
pytest.param(
45
"openai-community/gpt2", # gpt2
@@ -87,7 +87,11 @@ AITER_MODEL_LIST = [
87
pytest.param("bigcode/starcoder2-3b"), # starcoder2
88
pytest.param(
89
"TitanML/tiny-mixtral", # mixtral
90
- marks=[pytest.mark.core_model, pytest.mark.cpu_model],
90
+ marks=[pytest.mark.core_model],
91
+ ),
92
+ pytest.param(
93
+ "Qwen/Qwen1.5-MoE-A2.7B-Chat",
94
+ marks=[pytest.mark.cpu_model],
95
)
96
])
97
@pytest.mark.parametrize("max_tokens", [32])
tests/models/language/pooling/test_embedding.py
+11
-12
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import os
3
from typing import Optional
4
5
import pytest
@@ -29,8 +28,10 @@ def v1(run_with_both_engines):
28
# [Decoder-only]
29
pytest.param("BAAI/bge-multilingual-gemma2",
30
marks=[pytest.mark.core_model]),
32
- pytest.param("intfloat/e5-mistral-7b-instruct",
33
- marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
31
+ pytest.param(
32
+ "intfloat/e5-mistral-7b-instruct",
33
+ # CPU v1 doesn't support sliding window
34
+ marks=[pytest.mark.core_model]),
35
# the qwen models interfere with each other (see PR
36
# https://github.com/vllm-project/vllm/pull/18720).
37
# To avoid this problem, for now we skip v0 since it will be
@@ -38,11 +39,13 @@ 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]
41
- pytest.param("BAAI/bge-base-en-v1.5",
42
- marks=[
43
- pytest.mark.core_model, pytest.mark.cpu_model,
44
- pytest.mark.skip_v1
45
- ]),
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",
@@ -61,10 +64,6 @@ def test_models(
64
model,
65
monkeypatch,
66
) -> None:
64
- if model == "intfloat/e5-mistral-7b-instruct" and current_platform.is_cpu(
65
- ) and os.environ.get("VLLM_USE_V1", "0") == "1":
66
- pytest.skip("CPU V1 doesn't support sliding window")
67
-
67
if model == "BAAI/bge-multilingual-gemma2" and current_platform.is_rocm():
68
# ROCm Triton FA does not currently support sliding window attention
69
# switch to use ROCm CK FA backend
tests/models/language/pooling/test_reward.py
+5
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
+
5
import pytest
6
import torch
7
import torch.nn.functional as F
@@ -84,6 +86,9 @@ def test_prm_models(
86
dtype: str,
87
monkeypatch,
88
) -> None:
89
+ if current_platform.is_cpu() and os.environ.get("VLLM_USE_V1", "0") == "0":
90
+ pytest.skip("CPU only supports V1")
91
+
92
if current_platform.is_rocm():
93
# ROCm Triton FA does not currently support sliding window attention
94
# switch to use ROCm CK FA backend
tests/quantization/test_compressed_tensors.py
+2
-1
@@ -45,7 +45,8 @@ def use_v0_only(monkeypatch):
45
"""
46
This module relies on V0 internals, so set VLLM_USE_V1=0.
47
"""
48
- monkeypatch.setenv('VLLM_USE_V1', '0')
48
+ if not current_platform.is_cpu():
49
+ monkeypatch.setenv('VLLM_USE_V1', '0')
50
51
52
@pytest.mark.parametrize(
vllm/attention/backends/torch_sdpa.py
deleted
-546
@@ -1,546 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-""" Attention layer with torch scaled_dot_product_attention
4
- and PagedAttention."""
5
-from dataclasses import dataclass
6
-from typing import Any, Dict, List, Optional
7
-
8
-import torch
9
-from torch.nn.functional import scaled_dot_product_attention
10
-
11
-# yapf conflicts with isort for this block
12
-# yapf: disable
13
-from vllm.attention.backends.abstract import (AttentionImpl, AttentionLayer,
14
- AttentionMetadata, AttentionType,
15
- is_quantized_kv_cache)
16
-# yapf: enable
17
-from vllm.attention.ops.ipex_attn import PagedAttention, _use_ipex
18
-from vllm.attention.ops.paged_attn import PagedAttentionMetadata
19
-from vllm.logger import init_logger
20
-
21
-logger = init_logger(__name__)
22
-
23
-
24
-@dataclass
25
-class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
26
- """Metadata for TorchSDPABackend.
27
- """
28
- # Currently, input sequences can only contain all prompts
29
- # or all decoding. True if all sequences are prompts.
30
- chunked_prefill: bool
31
- seq_lens: Optional[List[int]] = None # For non-chunked prefill
32
-
33
- # For chunked prefill only
34
- max_query_len: Optional[int] = None
35
- max_kv_len: Optional[int] = None
36
- prefill_query_start_loc: Optional[torch.Tensor] = None
37
- kv_start_loc: Optional[torch.Tensor] = None
38
- prefill_block_tables: Optional[torch.Tensor] = None
39
-
40
- # For V1 logits index only
41
- query_start_loc: Optional[torch.Tensor] = None
42
-
43
- # Begin encoder attn & enc/dec cross-attn fields...
44
- # Encoder sequence lengths representation
45
- encoder_seq_lens: Optional[List[int]] = None
46
- encoder_seq_lens_tensor: Optional[torch.Tensor] = None
47
-
48
- # Maximum sequence length among encoder sequences
49
- max_encoder_seq_len: Optional[int] = None
50
-
51
- # Number of tokens input to encoder
52
- num_encoder_tokens: Optional[int] = None
53
-
54
- # Cross-attention memory-mapping data structures: slot mapping
55
- # and block tables
56
- cross_slot_mapping: Optional[torch.Tensor] = None
57
- cross_block_tables: Optional[torch.Tensor] = None
58
-
59
- def __post_init__(self):
60
- # Set during the execution of the first attention op.
61
- # It is a list because it is needed to set per prompt
62
- # when alibi slopes is used. It is because of the limitation
63
- # from xformer API.
64
- # will not appear in the __repr__ and __init__
65
- self.attn_bias: Optional[List[torch.Tensor]] = None
66
- self.encoder_attn_bias: Optional[List[torch.Tensor]] = None
67
- self.cross_attn_bias: Optional[List[torch.Tensor]] = None
68
-
69
- @property
70
- def is_all_encoder_attn_metadata_set(self):
71
- '''
72
- All attention metadata required for encoder attention is set.
73
- '''
74
- return ((self.encoder_seq_lens is not None)
75
- and (self.encoder_seq_lens_tensor is not None)
76
- and (self.max_encoder_seq_len is not None))
77
-
78
- @property
79
- def is_all_cross_attn_metadata_set(self):
80
- '''
81
- All attention metadata required for enc/dec cross-attention is set.
82
-
83
- Superset of encoder attention required metadata.
84
- '''
85
- return (self.is_all_encoder_attn_metadata_set
86
- and (self.cross_slot_mapping is not None)
87
- and (self.cross_block_tables is not None))
88
-
89
- @property
90
- def prefill_metadata(self) -> Optional["TorchSDPAMetadata"]:
91
- if self.num_prefill_tokens == 0:
92
- return None
93
- return self
94
-
95
- @property
96
- def decode_metadata(self) -> Optional["TorchSDPAMetadata"]:
97
- if self.num_decode_tokens == 0:
98
- return None
99
- return self
100
-
101
- def get_seq_lens(
102
- self,
103
- attn_type: str,
104
- ):
105
- '''
106
- Extract appropriate sequence lengths from attention metadata
107
- according to attention type.
108
-
109
- Arguments:
110
-
111
- * attn_metadata: Attention metadata structure associated with attention
112
- * attn_type: encoder attention, decoder self-attention,
113
- encoder/decoder cross-attention
114
-
115
- Returns:
116
- * Appropriate sequence lengths tensor for query
117
- * Appropriate sequence lengths tensor for key & value
118
- '''
119
-
120
- if (attn_type == AttentionType.DECODER
121
- or attn_type == AttentionType.ENCODER_ONLY):
122
- seq_lens_q = self.seq_lens
123
- seq_lens_kv = self.seq_lens
124
- elif attn_type == AttentionType.ENCODER:
125
- seq_lens_q = self.encoder_seq_lens
126
- seq_lens_kv = self.encoder_seq_lens
127
- elif attn_type == AttentionType.ENCODER_DECODER:
128
- seq_lens_q = self.seq_lens
129
- seq_lens_kv = self.encoder_seq_lens
130
- else:
131
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
132
- return seq_lens_q, seq_lens_kv
133
-
134
- def get_attn_bias(
135
- self,
136
- attn_type: str,
137
- ) -> Optional[List[torch.Tensor]]:
138
- '''
139
- Extract appropriate attention bias from attention metadata
140
- according to attention type.
141
-
142
- Arguments:
143
-
144
- * attn_metadata: Attention metadata structure associated with attention
145
- * attn_type: encoder attention, decoder self-attention,
146
- encoder/decoder cross-attention
147
-
148
- Returns:
149
- * Appropriate attention bias value given the attention type
150
- '''
151
-
152
- if (attn_type == AttentionType.DECODER
153
- or attn_type == AttentionType.ENCODER_ONLY):
154
- return self.attn_bias
155
- elif attn_type == AttentionType.ENCODER:
156
- return self.encoder_attn_bias
157
- elif attn_type == AttentionType.ENCODER_DECODER:
158
- return self.cross_attn_bias
159
- else:
160
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
161
-
162
- def set_attn_bias(
163
- self,
164
- attn_bias: List[torch.Tensor],
165
- attn_type: str,
166
- ) -> None:
167
- '''
168
- Update appropriate attention bias field of attention metadata,
169
- according to attention type.
170
-
171
- Arguments:
172
-
173
- * attn_metadata: Attention metadata structure associated with attention
174
- * attn_bias: The desired attention bias value
175
- * attn_type: encoder attention, decoder self-attention,
176
- encoder/decoder cross-attention
177
- '''
178
-
179
- if (attn_type == AttentionType.DECODER
180
- or attn_type == AttentionType.ENCODER_ONLY):
181
- self.attn_bias = attn_bias
182
- elif attn_type == AttentionType.ENCODER:
183
- self.encoder_attn_bias = attn_bias
184
- elif attn_type == AttentionType.ENCODER_DECODER:
185
- self.cross_attn_bias = attn_bias
186
- else:
187
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
188
-
189
- def get_seq_len_block_table_args(
190
- self,
191
- attn_type: str,
192
- ) -> tuple:
193
- '''
194
- The particular choice of sequence-length- and block-table-related
195
- attributes which should be extracted from attn_metadata is dependent
196
- on the type of attention operation.
197
-
198
- Decoder attn -> select entirely decoder self-attention-related fields
199
- Encoder/decoder cross-attn -> select encoder sequence lengths &
200
- cross-attn block-tables fields
201
- Encoder attn -> select encoder sequence lengths fields & no block tables
202
-
203
- Arguments:
204
-
205
- * attn_metadata: Attention metadata structure associated with attention
206
- * is_prompt: True if prefill, False otherwise
207
- * attn_type: encoder attention, decoder self-attention,
208
- encoder/decoder cross-attention
209
-
210
- Returns:
211
-
212
- * Appropriate sequence-lengths tensor
213
- * Appropriate max sequence-length scalar
214
- * Appropriate block tables (or None)
215
- '''
216
-
217
- if (attn_type == AttentionType.DECODER
218
- or attn_type == AttentionType.ENCODER_ONLY):
219
- # Decoder self-attention
220
- # Choose max_seq_len based on whether we are in prompt_run
221
- return (self.seq_lens_tensor, self.max_decode_seq_len,
222
- self.block_tables)
223
- elif attn_type == AttentionType.ENCODER_DECODER:
224
- # Enc/dec cross-attention KVs match encoder sequence length;
225
- # cross-attention utilizes special "cross" block tables
226
- return (self.encoder_seq_lens_tensor, self.max_encoder_seq_len,
227
- self.cross_block_tables)
228
- elif attn_type == AttentionType.ENCODER:
229
- # No block tables associated with encoder attention
230
- return (self.encoder_seq_lens_tensor, self.max_encoder_seq_len,
231
- None)
232
- else:
233
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
234
-
235
-
236
-class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
237
-
238
- def __init__(
239
- self,
240
- num_heads: int,
241
- head_size: int,
242
- scale: float,
243
- num_kv_heads: int,
244
- alibi_slopes: Optional[List[float]],
245
- sliding_window: Optional[int],
246
- kv_cache_dtype: str,
247
- blocksparse_params: Optional[Dict[str, Any]] = None,
248
- logits_soft_cap: Optional[float] = None,
249
- attn_type: str = AttentionType.DECODER,
250
- kv_sharing_target_layer_name: Optional[str] = None,
251
- use_irope: bool = False,
252
- ) -> None:
253
- if kv_sharing_target_layer_name is not None:
254
- raise NotImplementedError("KV sharing is not supported in V0.")
255
- if blocksparse_params is not None:
256
- raise ValueError(
257
- "Torch SPDA does not support block-sparse attention.")
258
- if logits_soft_cap is not None:
259
- logger.warning_once("Torch SPDA does not support logits soft cap. "
260
- "Outputs may be slightly off.")
261
- if use_irope:
262
- logger.warning_once(
263
- "Using irope in Torch SPDA is not supported yet, it will fall"
264
- " back to global attention for long context.")
265
- self.num_heads = num_heads
266
- self.head_size = head_size
267
- self.scale = float(scale)
268
- self.num_kv_heads = num_kv_heads
269
- if alibi_slopes is not None:
270
- alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
271
- self.alibi_slopes = alibi_slopes
272
- self.sliding_window = sliding_window
273
- self.kv_cache_dtype = kv_cache_dtype
274
-
275
- self.num_queries_per_kv = self.num_heads // self.num_kv_heads
276
- self.need_mask = (self.alibi_slopes is not None
277
- or self.sliding_window is not None)
278
-
279
- supported_head_sizes = PagedAttention.get_supported_head_sizes()
280
- if head_size not in supported_head_sizes:
281
- raise ValueError(
282
- f"Head size {head_size} is not supported by PagedAttention. "
283
- f"Supported head sizes are: {supported_head_sizes}.")
284
-
285
- if is_quantized_kv_cache(kv_cache_dtype) and not _use_ipex:
286
- raise NotImplementedError(
287
- "Torch SDPA backend FP8 KV cache requires "
288
- "intel_extension_for_pytorch support.")
289
- self.attn_type = attn_type
290
-
291
- def forward(
292
- self,
293
- layer: AttentionLayer,
294
- query: torch.Tensor,
295
- key: torch.Tensor,
296
- value: torch.Tensor,
297
- kv_cache: torch.Tensor,
298
- attn_metadata: TorchSDPAMetadata, # type: ignore
299
- output: Optional[torch.Tensor] = None,
300
- output_scale: Optional[torch.Tensor] = None,
301
- ) -> torch.Tensor:
302
- """Forward pass with torch SDPA and PagedAttention.
303
-
304
- Args:
305
- query: shape = [num_tokens, num_heads * head_size]
306
- key: shape = [num_tokens, num_kv_heads * head_size]
307
- value: shape = [num_tokens, num_kv_heads * head_size]
308
- kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
309
- NOTE: kv_cache will be an empty tensor with shape [0]
310
- for profiling run.
311
- attn_metadata: Metadata for attention.
312
- Returns:
313
- shape = [num_tokens, num_heads * head_size]
314
- """
315
- if output_scale is not None:
316
- raise NotImplementedError(
317
- "fused output quantization is not yet supported"
318
- " for TorchSDPABackendImpl")
319
-
320
- # For warming-up
321
- if attn_metadata is None:
322
- return query
323
-
324
- attn_type = self.attn_type
325
- if (attn_type == AttentionType.ENCODER
326
- and (not attn_metadata.is_all_encoder_attn_metadata_set)):
327
- raise AttributeError("Encoder attention requires setting "
328
- "encoder metadata attributes.")
329
- elif (attn_type == AttentionType.ENCODER_DECODER
330
- and (not attn_metadata.is_all_cross_attn_metadata_set)):
331
- raise AttributeError("Encoder/decoder cross-attention "
332
- "requires setting cross-attention "
333
- "metadata attributes.")
334
-
335
- # Reshape the query, key, and value tensors.
336
- query = query.view(-1, self.num_heads, self.head_size)
337
- if key is not None:
338
- assert value is not None
339
- key = key.view(-1, self.num_kv_heads, self.head_size)
340
- value = value.view(-1, self.num_kv_heads, self.head_size)
341
- else:
342
- assert value is None
343
-
344
- if (attn_type != AttentionType.ENCODER and kv_cache.numel() > 0):
345
- # KV-cache during decoder-self- or
346
- # encoder-decoder-cross-attention, but not
347
- # during encoder attention.
348
- #
349
- # Even if there are no new key/value pairs to cache,
350
- # we still need to break out key_cache and value_cache
351
- # i.e. for later use by paged attention
352
- key_cache, value_cache = PagedAttention.split_kv_cache(
353
- kv_cache, self.num_kv_heads, self.head_size)
354
-
355
- if (key is not None) and (value is not None):
356
- if attn_type == AttentionType.ENCODER_DECODER:
357
- # Update cross-attention KV cache (prefill-only)
358
- # During cross-attention decode, key & value will be None,
359
- # preventing this IF-statement branch from running
360
- updated_slot_mapping = attn_metadata.cross_slot_mapping
361
- else:
362
- # Update self-attention KV cache (prefill/decode)
363
- updated_slot_mapping = attn_metadata.slot_mapping
364
-
365
- PagedAttention.write_to_paged_cache(
366
- key, value, key_cache, value_cache, updated_slot_mapping,
367
- self.kv_cache_dtype, layer._k_scale, layer._v_scale)
368
-
369
- if attn_type != AttentionType.ENCODER:
370
- # Decoder self-attention supports chunked prefill.
371
- # Encoder/decoder cross-attention requires no chunked
372
- # prefill (100% prefill or 100% decode tokens, no mix)
373
- num_prefill_tokens = attn_metadata.num_prefill_tokens
374
- num_decode_tokens = attn_metadata.num_decode_tokens
375
- else:
376
- # Encoder attention - chunked prefill is not applicable;
377
- # derive token-count from query shape & and treat them
378
- # as 100% prefill tokens
379
- assert attn_metadata.num_encoder_tokens is not None
380
- num_prefill_tokens = attn_metadata.num_encoder_tokens
381
- num_decode_tokens = 0
382
-
383
- if attn_type == AttentionType.DECODER:
384
- # Only enforce this shape-constraint for decoder
385
- # self-attention
386
- assert key.shape[0] == num_prefill_tokens + num_decode_tokens
387
- assert value.shape[0] == num_prefill_tokens + num_decode_tokens
388
-
389
- output = torch.empty_like(query)
390
- if prefill_meta := attn_metadata.prefill_metadata:
391
- if not prefill_meta.prefill_metadata.chunked_prefill: # type: ignore
392
- assert attn_metadata.seq_lens is not None
393
- self._run_sdpa_forward(output,
394
- query,
395
- key,
396
- value,
397
- prefill_meta,
398
- attn_type=attn_type)
399
- else:
400
- # prefix-enabled attention
401
- assert not self.need_mask
402
- import intel_extension_for_pytorch.llm.modules as ipex_modules
403
- output = torch.empty_like(query)
404
- ipex_modules.PagedAttention.flash_attn_varlen_func(
405
- output[:prefill_meta.num_prefill_tokens, :, :],
406
- query[:prefill_meta.num_prefill_tokens, :, :],
407
- key_cache,
408
- value_cache,
409
- prefill_meta.prefill_query_start_loc,
410
- prefill_meta.kv_start_loc,
411
- prefill_meta.max_query_len,
412
- prefill_meta.max_kv_len,
413
- self.scale,
414
- True,
415
- prefill_meta.prefill_block_tables,
416
- self.alibi_slopes,
417
- )
418
-
419
- if decode_meta := attn_metadata.decode_metadata:
420
- assert attn_type != AttentionType.ENCODER_ONLY, (
421
- "Encoder-only models should not have decode metadata.")
422
- # Decoding run.
423
- (
424
- seq_lens_arg,
425
- max_seq_len_arg,
426
- block_tables_arg,
427
- ) = decode_meta.get_seq_len_block_table_args(attn_type)
428
-
429
- PagedAttention.forward_decode(
430
- output[attn_metadata.num_prefill_tokens:, :, :],
431
- query[attn_metadata.num_prefill_tokens:, :, :],
432
- key_cache,
433
- value_cache,
434
- block_tables_arg,
435
- seq_lens_arg,
436
- max_seq_len_arg,
437
- self.kv_cache_dtype,
438
- self.num_kv_heads,
439
- self.scale,
440
- self.alibi_slopes,
441
- layer._k_scale,
442
- layer._v_scale,
443
- )
444
-
445
- # Reshape the output tensor.
446
- return output.view(-1, self.num_heads * self.head_size)
447
-
448
- def _run_sdpa_forward(
449
- self,
450
- output: torch.Tensor,
451
- query: torch.Tensor,
452
- key: torch.Tensor,
453
- value: torch.Tensor,
454
- attn_metadata: TorchSDPAMetadata,
455
- attn_type: str = AttentionType.DECODER,
456
- ) -> None:
457
- if self.num_kv_heads != self.num_heads:
458
- key = key.repeat_interleave(self.num_queries_per_kv, dim=1)
459
- value = value.repeat_interleave(self.num_queries_per_kv, dim=1)
460
-
461
- attn_masks = attn_metadata.get_attn_bias(attn_type)
462
- if attn_masks is None:
463
- if self.alibi_slopes is not None:
464
- attn_masks = _make_alibi_bias(
465
- self.alibi_slopes, query.dtype,
466
- attn_metadata.seq_lens) # type: ignore
467
- elif self.sliding_window is not None:
468
- assert attn_metadata.seq_lens is not None
469
- attn_masks = _make_sliding_window_bias(
470
- attn_metadata.seq_lens, self.sliding_window,
471
- query.dtype) # type: ignore
472
- else:
473
- seq_lens, _ = attn_metadata.get_seq_lens(attn_type)
474
- attn_masks = [None] * len(seq_lens)
475
- attn_metadata.set_attn_bias(attn_masks, attn_type)
476
-
477
- query = query.movedim(0, query.dim() - 2)
478
- key = key.movedim(0, key.dim() - 2)
479
- value = value.movedim(0, value.dim() - 2)
480
-
481
- causal_attn = (attn_type == AttentionType.DECODER)
482
-
483
- seq_lens_q, seq_lens_kv = attn_metadata.get_seq_lens(attn_type)
484
- start_q, start_kv = 0, 0
485
- for seq_len_q, seq_len_kv, mask in zip(seq_lens_q, seq_lens_kv,
486
- attn_masks):
487
- end_q = start_q + seq_len_q
488
- end_kv = start_kv + seq_len_kv
489
- sub_out = scaled_dot_product_attention(
490
- query[None, :, start_q:end_q, :],
491
- key[None, :, start_kv:end_kv, :],
492
- value[None, :, start_kv:end_kv, :],
493
- attn_mask=mask,
494
- dropout_p=0.0,
495
- is_causal=causal_attn and mask is None,
496
- scale=self.scale).squeeze(0).movedim(query.dim() - 2, 0)
497
- output[start_q:end_q, :, :] = sub_out
498
- start_q, start_kv = end_q, end_kv
499
-
500
-
501
-def _make_alibi_bias(
502
- alibi_slopes: torch.Tensor,
503
- dtype: torch.dtype,
504
- seq_lens: List[int],
505
-) -> List[torch.Tensor]:
506
- attn_biases: List[torch.Tensor] = []
507
- for seq_len in seq_lens:
508
- bias = torch.arange(seq_len, dtype=dtype)
509
- # NOTE(zhuohan): HF uses
510
- # `bias = bias[None, :].repeat(seq_len, 1)`
511
- # here. We find that both biases give the same results, but
512
- # the bias below more accurately follows the original ALiBi
513
- # paper.
514
- bias = bias[None, :] - bias[:, None]
515
-
516
- num_heads = alibi_slopes.shape[0]
517
- bias = bias[None, :].repeat((num_heads, 1, 1))
518
- bias.mul_(alibi_slopes[:, None, None]).unsqueeze_(0)
519
- inf_mask = torch.empty(
520
- (1, seq_len, seq_len),
521
- dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1)
522
- attn_biases.append((bias + inf_mask).to(dtype))
523
-
524
- return attn_biases
525
-
526
-
527
-def _make_sliding_window_bias(
528
- seq_lens: List[int],
529
- window_size: Optional[int],
530
- dtype: torch.dtype,
531
-) -> List[torch.Tensor]:
532
- attn_biases: List[torch.Tensor] = []
533
- for seq_len in seq_lens:
534
- tensor = torch.full(
535
- (1, seq_len, seq_len),
536
- dtype=dtype,
537
- fill_value=1,
538
- )
539
- shift = 0
540
- mask = torch.tril(tensor, diagonal=shift).to(dtype) # type: ignore
541
- if window_size is not None:
542
- mask = torch.triu(mask, diagonal=shift - window_size + 1)
543
- mask = torch.log(mask)
544
- attn_biases.append(mask.to(dtype))
545
-
546
- return attn_biases
vllm/attention/ops/ipex_attn.py
deleted
-195
@@ -1,195 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import List, Optional, Tuple
5
-
6
-try:
7
- import intel_extension_for_pytorch.llm.modules as ipex_modules
8
- _use_ipex = True
9
-# AttributeError is to handle a bug in ipex https://github.com/intel/intel-extension-for-pytorch/pull/813
10
-except (ImportError, AttributeError):
11
- _use_ipex = False
12
-
13
-import torch
14
-
15
-from vllm import _custom_ops as ops
16
-
17
-
18
-class _PagedAttention:
19
-
20
- @staticmethod
21
- def get_supported_head_sizes() -> List[int]:
22
- return [32, 64, 80, 96, 112, 128, 192, 256]
23
-
24
- @staticmethod
25
- def get_kv_cache_shape(
26
- num_blocks: int,
27
- block_size: int,
28
- num_kv_heads: int,
29
- head_size: int,
30
- *args,
31
- ) -> Tuple[int, ...]:
32
- return 2, num_blocks, block_size * num_kv_heads * head_size
33
-
34
- @staticmethod
35
- def split_kv_cache(
36
- kv_cache: torch.Tensor,
37
- num_kv_heads: int,
38
- head_size: int,
39
- *args,
40
- ) -> Tuple[torch.Tensor, torch.Tensor]:
41
- x = 16 // kv_cache.element_size()
42
- num_blocks = kv_cache.shape[1]
43
-
44
- key_cache = kv_cache[0]
45
- key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x,
46
- -1, x)
47
- value_cache = kv_cache[1]
48
- value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
49
- return key_cache, value_cache
50
-
51
- @staticmethod
52
- def write_to_paged_cache(
53
- key: torch.Tensor,
54
- value: torch.Tensor,
55
- key_cache: torch.Tensor,
56
- value_cache: torch.Tensor,
57
- slot_mapping: torch.Tensor,
58
- kv_cache_dtype: str,
59
- k_scale: torch.Tensor,
60
- v_scale: torch.Tensor,
61
- *args,
62
- ) -> None:
63
- ops.reshape_and_cache(
64
- key,
65
- value,
66
- key_cache,
67
- value_cache,
68
- slot_mapping.flatten(),
69
- kv_cache_dtype,
70
- k_scale,
71
- v_scale,
72
- )
73
-
74
- @staticmethod
75
- def forward_decode(
76
- output: torch.Tensor,
77
- query: torch.Tensor,
78
- key_cache: torch.Tensor,
79
- value_cache: torch.Tensor,
80
- block_tables: torch.Tensor,
81
- context_lens: torch.Tensor,
82
- max_context_len: int,
83
- kv_cache_dtype: str,
84
- num_kv_heads: int,
85
- scale: float,
86
- alibi_slopes: Optional[torch.Tensor],
87
- k_scale: torch.Tensor,
88
- v_scale: torch.Tensor,
89
- *args,
90
- ) -> None:
91
- tp_rank: int = 0
92
- blocksparse_local_blocks: int = 0
93
- blocksparse_vert_stride: int = 0
94
- blocksparse_block_size: int = 64
95
- blocksparse_head_sliding_step: int = 0
96
- block_size = value_cache.shape[3]
97
-
98
- ops.paged_attention_v1(
99
- output,
100
- query,
101
- key_cache,
102
- value_cache,
103
- num_kv_heads,
104
- scale,
105
- block_tables,
106
- context_lens,
107
- block_size,
108
- max_context_len,
109
- alibi_slopes,
110
- kv_cache_dtype,
111
- k_scale,
112
- v_scale,
113
- tp_rank,
114
- blocksparse_local_blocks,
115
- blocksparse_vert_stride,
116
- blocksparse_block_size,
117
- blocksparse_head_sliding_step,
118
- )
119
-
120
- @staticmethod
121
- def copy_blocks(
122
- kv_caches: List[torch.Tensor],
123
- src_to_dists: torch.Tensor,
124
- *args,
125
- ) -> None:
126
- key_caches = [kv_cache[0] for kv_cache in kv_caches]
127
- value_caches = [kv_cache[1] for kv_cache in kv_caches]
128
- ops.copy_blocks(key_caches, value_caches, src_to_dists)
129
-
130
-
131
-class _IPEXPagedAttention(_PagedAttention):
132
-
133
- @staticmethod
134
- def split_kv_cache(
135
- kv_cache: torch.Tensor,
136
- num_kv_heads: int,
137
- head_size: int,
138
- *args,
139
- ) -> Tuple[torch.Tensor, torch.Tensor]:
140
- num_blocks = kv_cache.shape[1]
141
-
142
- key_cache = kv_cache[0]
143
- key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size)
144
- value_cache = kv_cache[1]
145
- value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size)
146
- return key_cache, value_cache
147
-
148
- @staticmethod
149
- def write_to_paged_cache(
150
- key: torch.Tensor,
151
- value: torch.Tensor,
152
- key_cache: torch.Tensor,
153
- value_cache: torch.Tensor,
154
- slot_mapping: torch.Tensor,
155
- kv_cache_dtype: str,
156
- k_scale: torch.Tensor,
157
- v_scale: torch.Tensor,
158
- *args,
159
- ) -> None:
160
- ipex_modules.PagedAttention.reshape_and_cache(
161
- key, value, key_cache, value_cache,
162
- slot_mapping.flatten().int())
163
-
164
- @staticmethod
165
- def forward_decode(
166
- output: torch.Tensor,
167
- query: torch.Tensor,
168
- key_cache: torch.Tensor,
169
- value_cache: torch.Tensor,
170
- block_tables: torch.Tensor,
171
- context_lens: torch.Tensor,
172
- max_context_len: int,
173
- kv_cache_dtype: str,
174
- num_kv_heads: int,
175
- scale: float,
176
- alibi_slopes: Optional[torch.Tensor],
177
- k_scale: torch.Tensor,
178
- v_scale: torch.Tensor,
179
- *args,
180
- ) -> None:
181
- block_size = value_cache.shape[2]
182
- head_mapping = torch.arange(
183
- 0,
184
- num_kv_heads,
185
- device="cpu",
186
- dtype=torch.int32,
187
- ).view(num_kv_heads,
188
- 1).repeat_interleave(query.size(1) // num_kv_heads).flatten()
189
- ipex_modules.PagedAttention.single_query_cached_kv_attention(
190
- output, query.contiguous(), key_cache, value_cache, head_mapping,
191
- scale, block_tables, context_lens, block_size, max_context_len,
192
- alibi_slopes)
193
-
194
-
195
-PagedAttention = _IPEXPagedAttention if _use_ipex else _PagedAttention
vllm/v1/attention/backends/cpu_attn.py
+749
-13
@@ -1,14 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import Any, Optional
5
+
6
import numpy as np
7
import torch
8
+from torch.nn.functional import scaled_dot_product_attention
9
6
-from vllm.attention.backends.abstract import (AttentionBackend,
7
- AttentionMetadata)
8
-from vllm.attention.backends.torch_sdpa import (TorchSDPABackendImpl,
9
- TorchSDPAMetadata)
10
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
+ AttentionLayer,
12
+ AttentionMetadata, AttentionType,
13
+ is_quantized_kv_cache)
14
from vllm.attention.backends.utils import CommonAttentionState
11
-from vllm.attention.ops.ipex_attn import PagedAttention
15
+from vllm.logger import init_logger
16
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
17
CommonAttentionMetadata)
18
from vllm.v1.core.sched.output import SchedulerOutput
@@ -17,18 +21,28 @@ from vllm.v1.worker.block_table import BlockTable
21
from vllm.v1.worker.cpu_model_runner import CPUModelRunner
22
from vllm.v1.worker.gpu_input_batch import InputBatch
23
24
+try:
25
+ import intel_extension_for_pytorch.llm.modules as ipex_modules
26
+ _use_ipex = True
27
+# AttributeError is to handle a bug in ipex
28
+# https://github.com/intel/intel-extension-for-pytorch/pull/813
29
+except (ImportError, AttributeError):
30
+ _use_ipex = False
31
+
32
+from vllm import _custom_ops as ops
33
+
34
+logger = init_logger(__name__)
35
+
36
37
class TorchSDPABackend(AttentionBackend):
38
accept_output_buffer: bool = False
39
24
- @classmethod
25
- def get_supported_head_sizes(cls) -> list[int]:
26
- return PagedAttention.get_supported_head_sizes()
27
-
40
@classmethod
41
def validate_head_size(cls, head_size: int) -> None:
30
- supported_head_sizes = cls.get_supported_head_sizes()
31
- if head_size not in supported_head_sizes:
42
+ attn_impl = _get_paged_attn_impl()
43
+ is_valid, supported_head_sizes = attn_impl.validate_head_size(
44
+ head_size)
45
+ if not is_valid:
46
attn_type = cls.__name__.removesuffix("Backend")
47
raise ValueError(
48
f"Head size {head_size} is not supported by {attn_type}. "
@@ -63,14 +77,239 @@ class TorchSDPABackend(AttentionBackend):
77
num_kv_heads: int,
78
head_size: int,
79
) -> tuple[int, ...]:
66
- return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
67
- num_kv_heads, head_size)
80
+ return _get_paged_attn_impl().get_kv_cache_shape(
81
+ num_blocks, block_size, num_kv_heads, head_size)
82
83
@staticmethod
84
def use_cascade_attention(*args, **kwargs) -> bool:
85
return False
86
87
88
+@dataclass
89
+class TorchSDPAMetadata(AttentionMetadata):
90
+ """Metadata for PagedAttention."""
91
+ # (batch_size,). The length of sequences (entire tokens seen so far) per
92
+ # sequence.
93
+ seq_lens_tensor: Optional[torch.Tensor]
94
+ # Maximum sequence length in the batch. 0 if it is prefill-only batch.
95
+ max_decode_seq_len: int
96
+ # (batch_size, max_blocks_per_seq).
97
+ # Block addresses per sequence. (Seq id -> list of physical block)
98
+ # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
99
+ # in the kv cache. Each block can contain up to block_size tokens.
100
+ # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
101
+ # captured.
102
+ block_tables: Optional[torch.Tensor]
103
+ """Metadata for TorchSDPABackend.
104
+ """
105
+ # Currently, input sequences can only contain all prompts
106
+ # or all decoding. True if all sequences are prompts.
107
+ chunked_prefill: bool
108
+ seq_lens: Optional[list[int]] = None # For non-chunked prefill
109
+
110
+ # For chunked prefill only
111
+ max_query_len: Optional[int] = None
112
+ max_kv_len: Optional[int] = None
113
+ prefill_query_start_loc: Optional[torch.Tensor] = None
114
+ kv_start_loc: Optional[torch.Tensor] = None
115
+ prefill_block_tables: Optional[torch.Tensor] = None
116
+
117
+ # For V1 logits index only
118
+ query_start_loc: Optional[torch.Tensor] = None
119
+
120
+ # Begin encoder attn & enc/dec cross-attn fields...
121
+ # Encoder sequence lengths representation
122
+ encoder_seq_lens: Optional[list[int]] = None
123
+ encoder_seq_lens_tensor: Optional[torch.Tensor] = None
124
+
125
+ # Maximum sequence length among encoder sequences
126
+ max_encoder_seq_len: Optional[int] = None
127
+
128
+ # Number of tokens input to encoder
129
+ num_encoder_tokens: Optional[int] = None
130
+
131
+ # Cross-attention memory-mapping data structures: slot mapping
132
+ # and block tables
133
+ cross_slot_mapping: Optional[torch.Tensor] = None
134
+ cross_block_tables: Optional[torch.Tensor] = None
135
+
136
+ def __post_init__(self):
137
+ # Set during the execution of the first attention op.
138
+ # It is a list because it is needed to set per prompt
139
+ # when alibi slopes is used. It is because of the limitation
140
+ # from xformer API.
141
+ # will not appear in the __repr__ and __init__
142
+ self.attn_bias: Optional[list[torch.Tensor]] = None
143
+ self.encoder_attn_bias: Optional[list[torch.Tensor]] = None
144
+ self.cross_attn_bias: Optional[list[torch.Tensor]] = None
145
+
146
+ @property
147
+ def is_all_encoder_attn_metadata_set(self):
148
+ '''
149
+ All attention metadata required for encoder attention is set.
150
+ '''
151
+ return ((self.encoder_seq_lens is not None)
152
+ and (self.encoder_seq_lens_tensor is not None)
153
+ and (self.max_encoder_seq_len is not None))
154
+
155
+ @property
156
+ def is_all_cross_attn_metadata_set(self):
157
+ '''
158
+ All attention metadata required for enc/dec cross-attention is set.
159
+
160
+ Superset of encoder attention required metadata.
161
+ '''
162
+ return (self.is_all_encoder_attn_metadata_set
163
+ and (self.cross_slot_mapping is not None)
164
+ and (self.cross_block_tables is not None))
165
+
166
+ @property
167
+ def prefill_metadata(self) -> Optional["TorchSDPAMetadata"]:
168
+ if self.num_prefill_tokens == 0:
169
+ return None
170
+ return self
171
+
172
+ @property
173
+ def decode_metadata(self) -> Optional["TorchSDPAMetadata"]:
174
+ if self.num_decode_tokens == 0:
175
+ return None
176
+ return self
177
+
178
+ def get_seq_lens(
179
+ self,
180
+ attn_type: str,
181
+ ):
182
+ '''
183
+ Extract appropriate sequence lengths from attention metadata
184
+ according to attention type.
185
+
186
+ Arguments:
187
+
188
+ * attn_metadata: Attention metadata structure associated with attention
189
+ * attn_type: encoder attention, decoder self-attention,
190
+ encoder/decoder cross-attention
191
+
192
+ Returns:
193
+ * Appropriate sequence lengths tensor for query
194
+ * Appropriate sequence lengths tensor for key & value
195
+ '''
196
+
197
+ if (attn_type == AttentionType.DECODER
198
+ or attn_type == AttentionType.ENCODER_ONLY):
199
+ seq_lens_q = self.seq_lens
200
+ seq_lens_kv = self.seq_lens
201
+ elif attn_type == AttentionType.ENCODER:
202
+ seq_lens_q = self.encoder_seq_lens
203
+ seq_lens_kv = self.encoder_seq_lens
204
+ elif attn_type == AttentionType.ENCODER_DECODER:
205
+ seq_lens_q = self.seq_lens
206
+ seq_lens_kv = self.encoder_seq_lens
207
+ else:
208
+ raise AttributeError(f"Invalid attention type {str(attn_type)}")
209
+ return seq_lens_q, seq_lens_kv
210
+
211
+ def get_attn_bias(
212
+ self,
213
+ attn_type: str,
214
+ ) -> Optional[list[torch.Tensor]]:
215
+ '''
216
+ Extract appropriate attention bias from attention metadata
217
+ according to attention type.
218
+
219
+ Arguments:
220
+
221
+ * attn_metadata: Attention metadata structure associated with attention
222
+ * attn_type: encoder attention, decoder self-attention,
223
+ encoder/decoder cross-attention
224
+
225
+ Returns:
226
+ * Appropriate attention bias value given the attention type
227
+ '''
228
+
229
+ if (attn_type == AttentionType.DECODER
230
+ or attn_type == AttentionType.ENCODER_ONLY):
231
+ return self.attn_bias
232
+ elif attn_type == AttentionType.ENCODER:
233
+ return self.encoder_attn_bias
234
+ elif attn_type == AttentionType.ENCODER_DECODER:
235
+ return self.cross_attn_bias
236
+ else:
237
+ raise AttributeError(f"Invalid attention type {str(attn_type)}")
238
+
239
+ def set_attn_bias(
240
+ self,
241
+ attn_bias: list[torch.Tensor],
242
+ attn_type: str,
243
+ ) -> None:
244
+ '''
245
+ Update appropriate attention bias field of attention metadata,
246
+ according to attention type.
247
+
248
+ Arguments:
249
+
250
+ * attn_metadata: Attention metadata structure associated with attention
251
+ * attn_bias: The desired attention bias value
252
+ * attn_type: encoder attention, decoder self-attention,
253
+ encoder/decoder cross-attention
254
+ '''
255
+
256
+ if (attn_type == AttentionType.DECODER
257
+ or attn_type == AttentionType.ENCODER_ONLY):
258
+ self.attn_bias = attn_bias
259
+ elif attn_type == AttentionType.ENCODER:
260
+ self.encoder_attn_bias = attn_bias
261
+ elif attn_type == AttentionType.ENCODER_DECODER:
262
+ self.cross_attn_bias = attn_bias
263
+ else:
264
+ raise AttributeError(f"Invalid attention type {str(attn_type)}")
265
+
266
+ def get_seq_len_block_table_args(
267
+ self,
268
+ attn_type: str,
269
+ ) -> tuple:
270
+ '''
271
+ The particular choice of sequence-length- and block-table-related
272
+ attributes which should be extracted from attn_metadata is dependent
273
+ on the type of attention operation.
274
+
275
+ Decoder attn -> select entirely decoder self-attention-related fields
276
+ Encoder/decoder cross-attn -> select encoder sequence lengths &
277
+ cross-attn block-tables fields
278
+ Encoder attn -> select encoder sequence lengths fields & no block tables
279
+
280
+ Arguments:
281
+
282
+ * attn_metadata: Attention metadata structure associated with attention
283
+ * is_prompt: True if prefill, False otherwise
284
+ * attn_type: encoder attention, decoder self-attention,
285
+ encoder/decoder cross-attention
286
+
287
+ Returns:
288
+
289
+ * Appropriate sequence-lengths tensor
290
+ * Appropriate max sequence-length scalar
291
+ * Appropriate block tables (or None)
292
+ '''
293
+
294
+ if (attn_type == AttentionType.DECODER
295
+ or attn_type == AttentionType.ENCODER_ONLY):
296
+ # Decoder self-attention
297
+ # Choose max_seq_len based on whether we are in prompt_run
298
+ return (self.seq_lens_tensor, self.max_decode_seq_len,
299
+ self.block_tables)
300
+ elif attn_type == AttentionType.ENCODER_DECODER:
301
+ # Enc/dec cross-attention KVs match encoder sequence length;
302
+ # cross-attention utilizes special "cross" block tables
303
+ return (self.encoder_seq_lens_tensor, self.max_encoder_seq_len,
304
+ self.cross_block_tables)
305
+ elif attn_type == AttentionType.ENCODER:
306
+ # No block tables associated with encoder attention
307
+ return (self.encoder_seq_lens_tensor, self.max_encoder_seq_len,
308
+ None)
309
+ else:
310
+ raise AttributeError(f"Invalid attention type {str(attn_type)}")
311
+
312
+
313
class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
314
315
def __init__(self, runner: CPUModelRunner, kv_cache_spec: AttentionSpec,
@@ -182,3 +421,500 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
421
)
422
423
return attn_metadata
424
+
425
+
426
+class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
427
+
428
+ def __init__(
429
+ self,
430
+ num_heads: int,
431
+ head_size: int,
432
+ scale: float,
433
+ num_kv_heads: int,
434
+ alibi_slopes: Optional[list[float]],
435
+ sliding_window: Optional[int],
436
+ kv_cache_dtype: str,
437
+ blocksparse_params: Optional[dict[str, Any]] = None,
438
+ logits_soft_cap: Optional[float] = None,
439
+ attn_type: str = AttentionType.DECODER,
440
+ kv_sharing_target_layer_name: Optional[str] = None,
441
+ use_irope: bool = False,
442
+ ) -> None:
443
+ if kv_sharing_target_layer_name is not None:
444
+ raise NotImplementedError("KV sharing is not supported in V0.")
445
+ if blocksparse_params is not None:
446
+ raise ValueError(
447
+ "Torch SPDA does not support block-sparse attention.")
448
+ if logits_soft_cap is not None:
449
+ logger.warning_once("Torch SPDA does not support logits soft cap. "
450
+ "Outputs may be slightly off.")
451
+ if use_irope:
452
+ logger.warning_once(
453
+ "Using irope in Torch SPDA is not supported yet, it will fall"
454
+ " back to global attention for long context.")
455
+ self.paged_attn_impl = _get_paged_attn_impl()
456
+ self.num_heads = num_heads
457
+ self.head_size = head_size
458
+ self.scale = float(scale)
459
+ self.num_kv_heads = num_kv_heads
460
+ if alibi_slopes is not None:
461
+ alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
462
+ self.alibi_slopes = alibi_slopes
463
+ self.sliding_window = sliding_window
464
+ self.kv_cache_dtype = kv_cache_dtype
465
+
466
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
467
+ self.need_mask = (self.alibi_slopes is not None
468
+ or self.sliding_window is not None)
469
+
470
+ if is_quantized_kv_cache(kv_cache_dtype) and not _use_ipex:
471
+ raise NotImplementedError(
472
+ "Torch SDPA backend FP8 KV cache requires "
473
+ "intel_extension_for_pytorch support.")
474
+ self.attn_type = attn_type
475
+
476
+ def forward(
477
+ self,
478
+ layer: AttentionLayer,
479
+ query: torch.Tensor,
480
+ key: torch.Tensor,
481
+ value: torch.Tensor,
482
+ kv_cache: torch.Tensor,
483
+ attn_metadata: TorchSDPAMetadata, # type: ignore
484
+ output: Optional[torch.Tensor] = None,
485
+ output_scale: Optional[torch.Tensor] = None,
486
+ ) -> torch.Tensor:
487
+ """Forward pass with torch SDPA and PagedAttention.
488
+
489
+ Args:
490
+ query: shape = [num_tokens, num_heads * head_size]
491
+ key: shape = [num_tokens, num_kv_heads * head_size]
492
+ value: shape = [num_tokens, num_kv_heads * head_size]
493
+ kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
494
+ NOTE: kv_cache will be an empty tensor with shape [0]
495
+ for profiling run.
496
+ attn_metadata: Metadata for attention.
497
+ Returns:
498
+ shape = [num_tokens, num_heads * head_size]
499
+ """
500
+ if output_scale is not None:
501
+ raise NotImplementedError(
502
+ "fused output quantization is not yet supported"
503
+ " for TorchSDPABackendImpl")
504
+
505
+ # For warming-up
506
+ if attn_metadata is None:
507
+ return query
508
+
509
+ attn_type = self.attn_type
510
+ if (attn_type == AttentionType.ENCODER
511
+ and (not attn_metadata.is_all_encoder_attn_metadata_set)):
512
+ raise AttributeError("Encoder attention requires setting "
513
+ "encoder metadata attributes.")
514
+ elif (attn_type == AttentionType.ENCODER_DECODER
515
+ and (not attn_metadata.is_all_cross_attn_metadata_set)):
516
+ raise AttributeError("Encoder/decoder cross-attention "
517
+ "requires setting cross-attention "
518
+ "metadata attributes.")
519
+
520
+ # Reshape the query, key, and value tensors.
521
+ query = query.view(-1, self.num_heads, self.head_size)
522
+ if key is not None:
523
+ assert value is not None
524
+ key = key.view(-1, self.num_kv_heads, self.head_size)
525
+ value = value.view(-1, self.num_kv_heads, self.head_size)
526
+ else:
527
+ assert value is None
528
+
529
+ if (attn_type != AttentionType.ENCODER and kv_cache.numel() > 0):
530
+ # KV-cache during decoder-self- or
531
+ # encoder-decoder-cross-attention, but not
532
+ # during encoder attention.
533
+ #
534
+ # Even if there are no new key/value pairs to cache,
535
+ # we still need to break out key_cache and value_cache
536
+ # i.e. for later use by paged attention
537
+ key_cache, value_cache = self.paged_attn_impl.split_kv_cache(
538
+ kv_cache, self.num_kv_heads, self.head_size)
539
+
540
+ if (key is not None) and (value is not None):
541
+ if attn_type == AttentionType.ENCODER_DECODER:
542
+ # Update cross-attention KV cache (prefill-only)
543
+ # During cross-attention decode, key & value will be None,
544
+ # preventing this IF-statement branch from running
545
+ updated_slot_mapping = attn_metadata.cross_slot_mapping
546
+ else:
547
+ # Update self-attention KV cache (prefill/decode)
548
+ updated_slot_mapping = attn_metadata.slot_mapping
549
+
550
+ self.paged_attn_impl.write_to_paged_cache(
551
+ key, value, key_cache, value_cache, updated_slot_mapping,
552
+ self.kv_cache_dtype, layer._k_scale, layer._v_scale)
553
+
554
+ if attn_type != AttentionType.ENCODER:
555
+ # Decoder self-attention supports chunked prefill.
556
+ # Encoder/decoder cross-attention requires no chunked
557
+ # prefill (100% prefill or 100% decode tokens, no mix)
558
+ num_prefill_tokens = attn_metadata.num_prefill_tokens
559
+ num_decode_tokens = attn_metadata.num_decode_tokens
560
+ else:
561
+ # Encoder attention - chunked prefill is not applicable;
562
+ # derive token-count from query shape & and treat them
563
+ # as 100% prefill tokens
564
+ assert attn_metadata.num_encoder_tokens is not None
565
+ num_prefill_tokens = attn_metadata.num_encoder_tokens
566
+ num_decode_tokens = 0
567
+
568
+ if attn_type == AttentionType.DECODER:
569
+ # Only enforce this shape-constraint for decoder
570
+ # self-attention
571
+ assert key.shape[0] == num_prefill_tokens + num_decode_tokens
572
+ assert value.shape[0] == num_prefill_tokens + num_decode_tokens
573
+
574
+ output = torch.empty_like(query)
575
+ if prefill_meta := attn_metadata.prefill_metadata:
576
+ if not prefill_meta.prefill_metadata.chunked_prefill: # type: ignore
577
+ assert attn_metadata.seq_lens is not None
578
+ self._run_sdpa_forward(output,
579
+ query,
580
+ key,
581
+ value,
582
+ prefill_meta,
583
+ attn_type=attn_type)
584
+ else:
585
+ # prefix-enabled attention
586
+ assert not self.need_mask
587
+ import intel_extension_for_pytorch.llm.modules as ipex_modules
588
+ output = torch.empty_like(query)
589
+ ipex_modules.PagedAttention.flash_attn_varlen_func(
590
+ output[:prefill_meta.num_prefill_tokens, :, :],
591
+ query[:prefill_meta.num_prefill_tokens, :, :],
592
+ key_cache,
593
+ value_cache,
594
+ prefill_meta.prefill_query_start_loc,
595
+ prefill_meta.kv_start_loc,
596
+ prefill_meta.max_query_len,
597
+ prefill_meta.max_kv_len,
598
+ self.scale,
599
+ True,
600
+ prefill_meta.prefill_block_tables,
601
+ self.alibi_slopes,
602
+ )
603
+
604
+ if decode_meta := attn_metadata.decode_metadata:
605
+ assert attn_type != AttentionType.ENCODER_ONLY, (
606
+ "Encoder-only models should not have decode metadata.")
607
+ # Decoding run.
608
+ (
609
+ seq_lens_arg,
610
+ max_seq_len_arg,
611
+ block_tables_arg,
612
+ ) = decode_meta.get_seq_len_block_table_args(attn_type)
613
+
614
+ self.paged_attn_impl.forward_decode(
615
+ output[attn_metadata.num_prefill_tokens:, :, :],
616
+ query[attn_metadata.num_prefill_tokens:, :, :],
617
+ key_cache,
618
+ value_cache,
619
+ block_tables_arg,
620
+ seq_lens_arg,
621
+ max_seq_len_arg,
622
+ self.kv_cache_dtype,
623
+ self.num_kv_heads,
624
+ self.scale,
625
+ self.alibi_slopes,
626
+ layer._k_scale,
627
+ layer._v_scale,
628
+ )
629
+
630
+ # Reshape the output tensor.
631
+ return output.view(-1, self.num_heads * self.head_size)
632
+
633
+ def _run_sdpa_forward(
634
+ self,
635
+ output: torch.Tensor,
636
+ query: torch.Tensor,
637
+ key: torch.Tensor,
638
+ value: torch.Tensor,
639
+ attn_metadata: TorchSDPAMetadata,
640
+ attn_type: str = AttentionType.DECODER,
641
+ ) -> None:
642
+ if self.num_kv_heads != self.num_heads:
643
+ key = key.repeat_interleave(self.num_queries_per_kv, dim=1)
644
+ value = value.repeat_interleave(self.num_queries_per_kv, dim=1)
645
+
646
+ attn_masks = attn_metadata.get_attn_bias(attn_type)
647
+ if attn_masks is None:
648
+ if self.alibi_slopes is not None:
649
+ attn_masks = _make_alibi_bias(
650
+ self.alibi_slopes, query.dtype,
651
+ attn_metadata.seq_lens) # type: ignore
652
+ elif self.sliding_window is not None:
653
+ assert attn_metadata.seq_lens is not None
654
+ attn_masks = _make_sliding_window_bias(
655
+ attn_metadata.seq_lens, self.sliding_window,
656
+ query.dtype) # type: ignore
657
+ else:
658
+ seq_lens, _ = attn_metadata.get_seq_lens(attn_type)
659
+ attn_masks = [None] * len(seq_lens)
660
+ attn_metadata.set_attn_bias(attn_masks, attn_type)
661
+
662
+ query = query.movedim(0, query.dim() - 2)
663
+ key = key.movedim(0, key.dim() - 2)
664
+ value = value.movedim(0, value.dim() - 2)
665
+
666
+ causal_attn = (attn_type == AttentionType.DECODER)
667
+
668
+ seq_lens_q, seq_lens_kv = attn_metadata.get_seq_lens(attn_type)
669
+ start_q, start_kv = 0, 0
670
+ for seq_len_q, seq_len_kv, mask in zip(seq_lens_q, seq_lens_kv,
671
+ attn_masks):
672
+ end_q = start_q + seq_len_q
673
+ end_kv = start_kv + seq_len_kv
674
+ sub_out = scaled_dot_product_attention(
675
+ query[None, :, start_q:end_q, :],
676
+ key[None, :, start_kv:end_kv, :],
677
+ value[None, :, start_kv:end_kv, :],
678
+ attn_mask=mask,
679
+ dropout_p=0.0,
680
+ is_causal=causal_attn and mask is None,
681
+ scale=self.scale).squeeze(0).movedim(query.dim() - 2, 0)
682
+ output[start_q:end_q, :, :] = sub_out
683
+ start_q, start_kv = end_q, end_kv
684
+
685
+
686
+def _make_alibi_bias(
687
+ alibi_slopes: torch.Tensor,
688
+ dtype: torch.dtype,
689
+ seq_lens: list[int],
690
+) -> list[torch.Tensor]:
691
+ attn_biases: list[torch.Tensor] = []
692
+ for seq_len in seq_lens:
693
+ bias = torch.arange(seq_len, dtype=dtype)
694
+ # NOTE(zhuohan): HF uses
695
+ # `bias = bias[None, :].repeat(seq_len, 1)`
696
+ # here. We find that both biases give the same results, but
697
+ # the bias below more accurately follows the original ALiBi
698
+ # paper.
699
+ bias = bias[None, :] - bias[:, None]
700
+
701
+ num_heads = alibi_slopes.shape[0]
702
+ bias = bias[None, :].repeat((num_heads, 1, 1))
703
+ bias.mul_(alibi_slopes[:, None, None]).unsqueeze_(0)
704
+ inf_mask = torch.empty(
705
+ (1, seq_len, seq_len),
706
+ dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1)
707
+ attn_biases.append((bias + inf_mask).to(dtype))
708
+
709
+ return attn_biases
710
+
711
+
712
+def _make_sliding_window_bias(
713
+ seq_lens: list[int],
714
+ window_size: Optional[int],
715
+ dtype: torch.dtype,
716
+) -> list[torch.Tensor]:
717
+ attn_biases: list[torch.Tensor] = []
718
+ for seq_len in seq_lens:
719
+ tensor = torch.full(
720
+ (1, seq_len, seq_len),
721
+ dtype=dtype,
722
+ fill_value=1,
723
+ )
724
+ shift = 0
725
+ mask = torch.tril(tensor, diagonal=shift).to(dtype) # type: ignore
726
+ if window_size is not None:
727
+ mask = torch.triu(mask, diagonal=shift - window_size + 1)
728
+ mask = torch.log(mask)
729
+ attn_biases.append(mask.to(dtype))
730
+
731
+ return attn_biases
732
+
733
+
734
+class _PagedAttention:
735
+
736
+ @staticmethod
737
+ def validate_head_size(head_size: int) -> tuple[bool, list[int]]:
738
+ SUPPORT_HS = [32, 64, 80, 96, 112, 128, 192, 256]
739
+ return head_size in SUPPORT_HS, SUPPORT_HS
740
+
741
+ @staticmethod
742
+ def get_kv_cache_shape(
743
+ num_blocks: int,
744
+ block_size: int,
745
+ num_kv_heads: int,
746
+ head_size: int,
747
+ *args,
748
+ ) -> tuple[int, ...]:
749
+ return 2, num_blocks, block_size * num_kv_heads * head_size
750
+
751
+ @staticmethod
752
+ def split_kv_cache(
753
+ kv_cache: torch.Tensor,
754
+ num_kv_heads: int,
755
+ head_size: int,
756
+ *args,
757
+ ) -> tuple[torch.Tensor, torch.Tensor]:
758
+ x = 16 // kv_cache.element_size()
759
+ num_blocks = kv_cache.shape[1]
760
+
761
+ key_cache = kv_cache[0]
762
+ key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x,
763
+ -1, x)
764
+ value_cache = kv_cache[1]
765
+ value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
766
+ return key_cache, value_cache
767
+
768
+ @staticmethod
769
+ def write_to_paged_cache(
770
+ key: torch.Tensor,
771
+ value: torch.Tensor,
772
+ key_cache: torch.Tensor,
773
+ value_cache: torch.Tensor,
774
+ slot_mapping: torch.Tensor,
775
+ kv_cache_dtype: str,
776
+ k_scale: torch.Tensor,
777
+ v_scale: torch.Tensor,
778
+ *args,
779
+ ) -> None:
780
+ ops.reshape_and_cache(
781
+ key,
782
+ value,
783
+ key_cache,
784
+ value_cache,
785
+ slot_mapping.flatten(),
786
+ kv_cache_dtype,
787
+ k_scale,
788
+ v_scale,
789
+ )
790
+
791
+ @staticmethod
792
+ def forward_decode(
793
+ output: torch.Tensor,
794
+ query: torch.Tensor,
795
+ key_cache: torch.Tensor,
796
+ value_cache: torch.Tensor,
797
+ block_tables: torch.Tensor,
798
+ context_lens: torch.Tensor,
799
+ max_context_len: int,
800
+ kv_cache_dtype: str,
801
+ num_kv_heads: int,
802
+ scale: float,
803
+ alibi_slopes: Optional[torch.Tensor],
804
+ k_scale: torch.Tensor,
805
+ v_scale: torch.Tensor,
806
+ *args,
807
+ ) -> None:
808
+ tp_rank: int = 0
809
+ blocksparse_local_blocks: int = 0
810
+ blocksparse_vert_stride: int = 0
811
+ blocksparse_block_size: int = 64
812
+ blocksparse_head_sliding_step: int = 0
813
+ block_size = value_cache.shape[3]
814
+
815
+ ops.paged_attention_v1(
816
+ output,
817
+ query,
818
+ key_cache,
819
+ value_cache,
820
+ num_kv_heads,
821
+ scale,
822
+ block_tables,
823
+ context_lens,
824
+ block_size,
825
+ max_context_len,
826
+ alibi_slopes,
827
+ kv_cache_dtype,
828
+ k_scale,
829
+ v_scale,
830
+ tp_rank,
831
+ blocksparse_local_blocks,
832
+ blocksparse_vert_stride,
833
+ blocksparse_block_size,
834
+ blocksparse_head_sliding_step,
835
+ )
836
+
837
+ @staticmethod
838
+ def copy_blocks(
839
+ kv_caches: list[torch.Tensor],
840
+ src_to_dists: torch.Tensor,
841
+ *args,
842
+ ) -> None:
843
+ key_caches = [kv_cache[0] for kv_cache in kv_caches]
844
+ value_caches = [kv_cache[1] for kv_cache in kv_caches]
845
+ ops.copy_blocks(key_caches, value_caches, src_to_dists)
846
+
847
+
848
+class _IPEXPagedAttention(_PagedAttention):
849
+
850
+ @staticmethod
851
+ def validate_head_size(head_size: int) -> tuple[bool, list[int]]:
852
+ return True, []
853
+
854
+ @staticmethod
855
+ def split_kv_cache(
856
+ kv_cache: torch.Tensor,
857
+ num_kv_heads: int,
858
+ head_size: int,
859
+ *args,
860
+ ) -> tuple[torch.Tensor, torch.Tensor]:
861
+ num_blocks = kv_cache.shape[1]
862
+
863
+ key_cache = kv_cache[0]
864
+ key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size)
865
+ value_cache = kv_cache[1]
866
+ value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size)
867
+ return key_cache, value_cache
868
+
869
+ @staticmethod
870
+ def write_to_paged_cache(
871
+ key: torch.Tensor,
872
+ value: torch.Tensor,
873
+ key_cache: torch.Tensor,
874
+ value_cache: torch.Tensor,
875
+ slot_mapping: torch.Tensor,
876
+ kv_cache_dtype: str,
877
+ k_scale: torch.Tensor,
878
+ v_scale: torch.Tensor,
879
+ *args,
880
+ ) -> None:
881
+ ipex_modules.PagedAttention.reshape_and_cache(
882
+ key, value, key_cache, value_cache,
883
+ slot_mapping.flatten().int())
884
+
885
+ @staticmethod
886
+ def forward_decode(
887
+ output: torch.Tensor,
888
+ query: torch.Tensor,
889
+ key_cache: torch.Tensor,
890
+ value_cache: torch.Tensor,
891
+ block_tables: torch.Tensor,
892
+ context_lens: torch.Tensor,
893
+ max_context_len: int,
894
+ kv_cache_dtype: str,
895
+ num_kv_heads: int,
896
+ scale: float,
897
+ alibi_slopes: Optional[torch.Tensor],
898
+ k_scale: torch.Tensor,
899
+ v_scale: torch.Tensor,
900
+ *args,
901
+ ) -> None:
902
+ block_size = value_cache.shape[2]
903
+ head_mapping = torch.arange(
904
+ 0,
905
+ num_kv_heads,
906
+ device="cpu",
907
+ dtype=torch.int32,
908
+ ).view(num_kv_heads,
909
+ 1).repeat_interleave(query.size(1) // num_kv_heads).flatten()
910
+ ipex_modules.PagedAttention.single_query_cached_kv_attention(
911
+ output, query.contiguous(), key_cache, value_cache, head_mapping,
912
+ scale, block_tables, context_lens, block_size, max_context_len,
913
+ alibi_slopes)
914
+
915
+
916
+def _get_paged_attn_impl():
917
+ if _use_ipex:
918
+ return _IPEXPagedAttention
919
+ else:
920
+ return _PagedAttention