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
[V0 deprecation] Remove V0 CPU/XPU/TPU backends (#20412)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Woosuk Kwon committed
Jul 6, 2025 at 08:48 UTC
e202dd2736bc575b11250b15311512d19d3225d5
20 files changed
+46
-5034
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+4
-4
@@ -66,10 +66,10 @@ function cpu_tests() {
66
tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_dynamic_per_token"
67
68
# Run AWQ test
69
- docker exec cpu-test-"$NUMA_NODE" bash -c "
70
- set -e
71
- VLLM_USE_V1=0 pytest -s -v \
72
- tests/quantization/test_ipex_quant.py"
69
+ # docker exec cpu-test-"$NUMA_NODE" bash -c "
70
+ # set -e
71
+ # VLLM_USE_V1=0 pytest -s -v \
72
+ # tests/quantization/test_ipex_quant.py"
73
74
# Run chunked-prefill and prefix-cache test
75
docker exec cpu-test-"$NUMA_NODE" bash -c "
.buildkite/scripts/hardware_ci/run-xpu-test.sh
-2
@@ -26,7 +26,5 @@ docker run \
26
--name "${container_name}" \
27
"${image_name}" \
28
sh -c '
29
- VLLM_USE_V1=0 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m
30
- VLLM_USE_V1=0 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m -tp 2
29
VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
30
'
examples/online_serving/chart-helm/values.yaml
+1
-1
@@ -8,7 +8,7 @@ image:
8
# -- Image tag
9
tag: "latest"
10
# -- Container launch command
11
- command: ["vllm", "serve", "/data/", "--served-model-name", "opt-125m", "--dtype", "float32", "--block-size", "16", "--host", "0.0.0.0", "--port", "8000"]
11
+ command: ["vllm", "serve", "/data/", "--served-model-name", "opt-125m", "--enforce-eager", "--dtype", "bfloat16", "--block-size", "16", "--host", "0.0.0.0", "--port", "8000"]
12
13
# -- Container port
14
containerPort: 8000
tests/kernels/attention/test_attention_selector.py
+10
-7
@@ -36,7 +36,8 @@ DEVICE_REGULAR_ATTN_BACKENDS = {
36
DEVICE_MLA_BLOCK_SIZES = {
37
"cuda": [16, 64], # CUDA supports both standard and extended block sizes
38
"hip": [16, 1], # HIP requires special handling for block_size=1
39
- "cpu": [16] # CPU uses fixed block size from test cases
39
+ # "cpu": [16] # CPU uses fixed block size from test cases
40
+ "cpu": [] # FIXME(woosuk): Temporarily disable CPU tests
41
}
42
43
@@ -81,14 +82,14 @@ def test_env(
82
m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
83
84
if device == "cpu":
85
+ if not use_v1:
86
+ pytest.skip("CPU backend only supports V1")
87
+
88
with patch("vllm.attention.selector.current_platform",
89
CpuPlatform()):
90
backend = get_attn_backend(16, torch.float16, torch.float16,
91
block_size, False)
88
- if use_v1:
89
- assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
90
- else:
91
- assert backend.get_name() == "TORCH_SDPA"
92
+ assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
93
94
elif device == "hip":
95
with patch("vllm.attention.selector.current_platform",
@@ -193,12 +194,14 @@ def test_fp32_fallback(
194
m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
195
196
if device == "cpu":
197
+ if not use_v1:
198
+ pytest.skip("CPU backend only supports V1")
199
+
200
with patch("vllm.attention.selector.current_platform",
201
CpuPlatform()):
202
backend = get_attn_backend(16, torch.float32, torch.float32,
203
16, False)
200
- assert (backend.get_name() == "TORCH_SDPA_VLLM_V1"
201
- if use_v1 else "TORCH_SDPA")
204
+ assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
205
206
elif device == "cuda":
207
with patch("vllm.attention.selector.current_platform",
vllm/attention/backends/cpu_mla.py
deleted
-307
@@ -1,307 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass
5
-from typing import Any, Dict, List, Optional, Tuple, Type
6
-
7
-import torch
8
-
9
-import vllm._custom_ops as ops
10
-from vllm._ipex_ops import ipex_ops
11
-from vllm.attention.backends.abstract import (AttentionBackend,
12
- AttentionMetadataBuilder,
13
- AttentionType,
14
- is_quantized_kv_cache)
15
-from vllm.attention.backends.mla.common import MLACommonImpl, MLACommonState
16
-from vllm.attention.backends.torch_sdpa import TorchSDPAMetadata
17
-from vllm.utils import make_tensor_with_pad
18
-from vllm.worker.cpu_model_runner import ModelInputForCPUBuilder
19
-
20
-
21
-class CPUMLABackend(AttentionBackend):
22
-
23
- @staticmethod
24
- def get_name() -> str:
25
- return "CPU_MLA"
26
-
27
- @staticmethod
28
- def get_metadata_cls() -> Type["CPUMLAMetadata"]:
29
- return CPUMLAMetadata
30
-
31
- @staticmethod
32
- def get_builder_cls() -> Type["CPUMLAMetadataBuilder"]:
33
- return CPUMLAMetadataBuilder
34
-
35
- @staticmethod
36
- def get_state_cls() -> Type["MLACommonState"]:
37
- return MLACommonState
38
-
39
- @staticmethod
40
- def get_impl_cls() -> Type["CPUMLAImpl"]:
41
- return CPUMLAImpl
42
-
43
- @staticmethod
44
- def get_kv_cache_shape(
45
- num_blocks: int,
46
- block_size: int,
47
- num_kv_heads: int, # assumed to be 1 for MLA
48
- head_size: int,
49
- ) -> Tuple[int, ...]:
50
- return (num_blocks, block_size, head_size)
51
-
52
- @staticmethod
53
- def swap_blocks(
54
- src_kv_cache: torch.Tensor,
55
- dst_kv_cache: torch.Tensor,
56
- src_to_dst: torch.Tensor,
57
- ) -> None:
58
- ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
59
-
60
- @staticmethod
61
- def copy_blocks(
62
- kv_caches: List[torch.Tensor],
63
- src_to_dists: torch.Tensor,
64
- ) -> None:
65
- ops.copy_blocks_mla(kv_caches, src_to_dists)
66
-
67
- @staticmethod
68
- def get_supported_head_sizes() -> List[int]:
69
- return [576]
70
-
71
-
72
-@dataclass
73
-class CPUMLAMetadata(TorchSDPAMetadata):
74
- # New for MLA
75
- # Input positions for rotrary embeddings since for MLA the rotary
76
- # position embeddings are applied inside the attention backend
77
- input_positions: torch.Tensor = None
78
-
79
- # required by MLACommonImpl
80
- is_profile_run: bool = False
81
-
82
-
83
-class CPUMLAMetadataBuilder(AttentionMetadataBuilder[CPUMLAMetadata]):
84
-
85
- def __init__(self, input_builder: ModelInputForCPUBuilder) -> None:
86
- self.chunked_prefill = input_builder.chunked_prefill
87
- self.input_builder = input_builder
88
- assert not self.chunked_prefill, \
89
- "chunked prefill is currently not supported"
90
-
91
- def prepare(self):
92
- self.input_data = self.input_builder.input_data
93
-
94
- def build(self, seq_lens, query_lens, cuda_graph_pad_size, batch_size):
95
- input_data = self.input_data
96
- prefill_seq_lens = seq_lens[0:input_data.num_prefills]
97
- prefill_query_lens = query_lens[0:input_data.num_prefills]
98
- slot_mapping = torch.tensor(input_data.slot_mapping,
99
- dtype=torch.long,
100
- device="cpu")
101
-
102
- # metadata for prefill
103
- if input_data.num_prefills > 0:
104
- query_lens_tensor = torch.tensor(prefill_query_lens,
105
- dtype=torch.int32,
106
- device="cpu")
107
- kv_lens_tensor = torch.tensor(prefill_seq_lens,
108
- dtype=torch.int32,
109
- device="cpu")
110
- query_start_loc = torch.zeros(input_data.num_prefills + 1,
111
- dtype=torch.int32,
112
- device="cpu")
113
- kv_start_loc = torch.zeros(input_data.num_prefills + 1,
114
- dtype=torch.int32,
115
- device="cpu")
116
- torch.cumsum(query_lens_tensor,
117
- dim=0,
118
- dtype=torch.int32,
119
- out=query_start_loc[1:])
120
- torch.cumsum(kv_lens_tensor,
121
- dim=0,
122
- dtype=torch.int32,
123
- out=kv_start_loc[1:])
124
- max_query_len = max(prefill_query_lens)
125
- max_kv_len = max(prefill_seq_lens)
126
-
127
- # for chunked-prefill
128
- if self.chunked_prefill:
129
- prefill_block_tables = make_tensor_with_pad(
130
- self.input_data.prefill_block_tables,
131
- pad=0,
132
- dtype=torch.int32,
133
- device="cpu",
134
- )
135
- else:
136
- prefill_block_tables = None
137
-
138
- else:
139
- query_start_loc = None
140
- kv_start_loc = None
141
- max_query_len = None
142
- max_kv_len = None
143
- prefill_block_tables = None
144
-
145
- # metadata for decode
146
- if input_data.num_decode_tokens != 0:
147
- seq_lens_tensor = torch.tensor(
148
- input_data.seq_lens[input_data.num_prefills:],
149
- dtype=torch.int32,
150
- device="cpu",
151
- )
152
- block_tables = make_tensor_with_pad(
153
- self.input_data.decode_block_tables,
154
- pad=0,
155
- dtype=torch.int32,
156
- device="cpu",
157
- )
158
- else:
159
- block_tables = torch.tensor([])
160
- seq_lens_tensor = torch.tensor(
161
- input_data.seq_lens[:input_data.num_prefills],
162
- dtype=torch.int32,
163
- device="cpu",
164
- )
165
-
166
- # For multi-modal models
167
- placeholder_index_maps = None
168
- if len(input_data.multi_modal_inputs_list) != 0:
169
- placeholder_index_maps = {
170
- modality: placeholder_map.index_map()
171
- for modality, placeholder_map in
172
- input_data.multi_modal_placeholder_maps.items()
173
- }
174
-
175
- return CPUMLAMetadata(
176
- chunked_prefill=self.chunked_prefill,
177
- seq_lens=prefill_seq_lens,
178
- seq_lens_tensor=seq_lens_tensor,
179
- max_query_len=max_query_len,
180
- max_kv_len=max_kv_len,
181
- prefill_query_start_loc=query_start_loc,
182
- kv_start_loc=kv_start_loc,
183
- max_decode_seq_len=input_data.max_decode_seq_len,
184
- num_prefills=input_data.num_prefills,
185
- num_prefill_tokens=input_data.num_prefill_tokens,
186
- num_decode_tokens=input_data.num_decode_tokens,
187
- block_tables=block_tables,
188
- prefill_block_tables=prefill_block_tables,
189
- slot_mapping=slot_mapping,
190
- multi_modal_placeholder_index_maps=placeholder_index_maps,
191
- enable_kv_scales_calculation=False,
192
- input_positions=torch.tensor([self.input_data.input_positions]))
193
-
194
-
195
-class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):
196
-
197
- def __init__(
198
- self,
199
- num_heads: int,
200
- head_size: int,
201
- scale: float,
202
- num_kv_heads: int,
203
- alibi_slopes: Optional[List[float]],
204
- sliding_window: Optional[int],
205
- kv_cache_dtype: str,
206
- blocksparse_params: Optional[Dict[str, Any]],
207
- logits_soft_cap: Optional[float],
208
- attn_type: str,
209
- kv_sharing_target_layer_name: Optional[str],
210
- # MLA Specific Arguments
211
- **mla_args) -> None:
212
- super().__init__(num_heads, head_size, scale, num_kv_heads,
213
- alibi_slopes, sliding_window, kv_cache_dtype,
214
- blocksparse_params, logits_soft_cap, attn_type,
215
- kv_sharing_target_layer_name, **mla_args)
216
-
217
- unsupported_features = [
218
- alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
219
- ]
220
- if any(unsupported_features):
221
- raise NotImplementedError(
222
- "CPUMLAImpl does not support one of the following: "
223
- "alibi_slopes, sliding_window, blocksparse_params, "
224
- "logits_soft_cap")
225
-
226
- if attn_type != AttentionType.DECODER:
227
- raise NotImplementedError("Encoder self-attention and "
228
- "encoder/decoder cross-attention "
229
- "are not implemented for "
230
- "CPUMLAImpl")
231
-
232
- # states is implemented.
233
- if is_quantized_kv_cache(self.kv_cache_dtype):
234
- raise NotImplementedError(
235
- "CPUMLAImpl with FP8 KV cache not yet supported")
236
-
237
- def _forward_prefill(
238
- self,
239
- q: torch.Tensor,
240
- kv_c_normed: torch.Tensor,
241
- k_pe: torch.Tensor,
242
- kv_c_and_k_pe_cache: torch.Tensor,
243
- attn_metadata: CPUMLAMetadata, # type: ignore[override]
244
- ) -> torch.Tensor:
245
-
246
- prefill_metadata = attn_metadata.prefill_metadata
247
- assert prefill_metadata is not None
248
-
249
- kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
250
- -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
251
- k_nope, v = kv_nope\
252
- .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
253
-
254
- k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
255
-
256
- # For MLA the v head dim is smaller than qk head dim so we pad out
257
- # v with 0s to match the qk head dim
258
- v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
259
- value=0)
260
-
261
- output = torch.empty_like(q)
262
- ipex_ops.varlen_attention(
263
- query=q,
264
- key=k,
265
- value=v_padded,
266
- out=output,
267
- seqlen_q=prefill_metadata.prefill_query_start_loc,
268
- seqlen_k=prefill_metadata.prefill_query_start_loc,
269
- max_seqlen_q=prefill_metadata.max_query_len,
270
- max_seqlen_k=prefill_metadata.max_query_len,
271
- pdropout=0.0,
272
- softmax_scale=self.scale,
273
- zero_tensors=False,
274
- is_causal=True,
275
- return_softmax=False,
276
- gen_=None,
277
- logits_soft_cap=0.0,
278
- window_size_left=-1,
279
- window_size_right=-1,
280
- alibi_slopes=None,
281
- )
282
-
283
- # remove padding
284
- output = output.view(-1, self.num_heads,
285
- q.shape[-1])[..., :v.shape[-1]]
286
- return output.reshape(-1, self.num_heads * v.shape[-1])
287
-
288
- def _forward_decode(
289
- self,
290
- q_nope: torch.Tensor,
291
- q_pe: torch.Tensor,
292
- kv_c_and_k_pe_cache: torch.Tensor,
293
- attn_metadata: CPUMLAMetadata, # type: ignore[override]
294
- ) -> torch.Tensor:
295
- assert kv_c_and_k_pe_cache.numel() > 0
296
-
297
- decode_meta = attn_metadata.decode_metadata
298
- assert decode_meta is not None
299
-
300
- q = torch.cat([q_nope, q_pe], dim=-1)
301
- o = q.new_empty(q.shape[0], self.num_heads, self.kv_lora_rank)
302
-
303
- # Run MQA
304
- ops.mla_decode_kvcache_cpu(o, q, kv_c_and_k_pe_cache, self.scale,
305
- decode_meta.block_tables,
306
- decode_meta.seq_lens_tensor)
307
- return self._v_up_proj(o)
vllm/attention/backends/ipex_attn.py
deleted
-403
@@ -1,403 +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, Tuple, Type
7
-
8
-import torch
9
-
10
-from vllm._ipex_ops import ipex_ops
11
-from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12
- AttentionLayer,
13
- AttentionMetadata, AttentionType,
14
- is_quantized_kv_cache)
15
-from vllm.attention.backends.utils import CommonAttentionState
16
-from vllm.attention.ops.paged_attn import (PagedAttention,
17
- PagedAttentionMetadata)
18
-from vllm.logger import init_logger
19
-
20
-logger = init_logger(__name__)
21
-
22
-_PARTITION_SIZE = 512
23
-
24
-
25
-class IpexAttnBackend(AttentionBackend):
26
-
27
- @staticmethod
28
- def get_name() -> str:
29
- return "IPEX"
30
-
31
- @staticmethod
32
- def get_impl_cls() -> Type["IpexAttnBackendImpl"]:
33
- return IpexAttnBackendImpl
34
-
35
- @staticmethod
36
- def get_metadata_cls() -> Type["IpexAttnMetadata"]:
37
- return IpexAttnMetadata
38
-
39
- @staticmethod
40
- def get_state_cls() -> Type["CommonAttentionState"]:
41
- return CommonAttentionState
42
-
43
- @staticmethod
44
- def get_kv_cache_shape(
45
- num_blocks: int,
46
- block_size: int,
47
- num_kv_heads: int,
48
- head_size: int,
49
- ) -> Tuple[int, ...]:
50
- return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
51
- num_kv_heads, head_size)
52
-
53
- @staticmethod
54
- def swap_blocks(
55
- src_kv_cache: torch.Tensor,
56
- dst_kv_cache: torch.Tensor,
57
- src_to_dst: torch.Tensor,
58
- ) -> None:
59
- from vllm._ipex_ops import ipex_ops as ops
60
- ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
61
-
62
- @staticmethod
63
- def copy_blocks(
64
- kv_caches: List[torch.Tensor],
65
- src_to_dists: torch.Tensor,
66
- ) -> None:
67
- from vllm._ipex_ops import ipex_ops as ops
68
- key_caches = [kv_cache[0] for kv_cache in kv_caches]
69
- value_caches = [kv_cache[1] for kv_cache in kv_caches]
70
- ops.copy_blocks(key_caches, value_caches, src_to_dists)
71
-
72
-
73
-@dataclass
74
-class IpexAttnMetadata(AttentionMetadata, PagedAttentionMetadata):
75
- """Metadata for IpexAttnBackend.
76
- """
77
- # Currently, input sequences can only contain all prompts
78
- # or all decoding. True if all sequences are prompts.
79
- is_prompt: bool
80
- slot_mapping: torch.Tensor
81
- seq_lens: Optional[List[int]]
82
- seqlen_q: Optional[torch.Tensor]
83
- max_seqlen: Optional[int]
84
-
85
- def __post_init__(self):
86
- # Set during the execution of the first attention op.
87
- # It is a list because it is needed to set per prompt
88
- # when alibi slopes is used. It is because of the limitation
89
- # from xformer API.
90
- # will not appear in the __repr__ and __init__
91
- self.attn_bias: Optional[List[torch.Tensor]] = None
92
-
93
- @property
94
- def prefill_metadata(self) -> Optional["IpexAttnMetadata"]:
95
- # Currently chunked prefill is not supported
96
- if self.num_decode_tokens == 0:
97
- assert self.num_prefills > 0
98
- return self
99
-
100
- return None
101
-
102
- @property
103
- def decode_metadata(self) -> Optional["IpexAttnMetadata"]:
104
- # Currently chunked prefill is not supported
105
- if self.num_prefills > 0:
106
- assert self.num_decode_tokens == 0
107
- return None
108
-
109
- return self
110
-
111
-
112
-class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):
113
-
114
- def __init__(
115
- self,
116
- num_heads: int,
117
- head_size: int,
118
- scale: float,
119
- num_kv_heads: int,
120
- alibi_slopes: Optional[List[float]],
121
- sliding_window: Optional[int],
122
- kv_cache_dtype: str,
123
- blocksparse_params: Optional[Dict[str, Any]] = None,
124
- logits_soft_cap: Optional[float] = None,
125
- attn_type: str = AttentionType.DECODER,
126
- kv_sharing_target_layer_name: Optional[str] = None,
127
- use_irope: bool = False,
128
- ) -> None:
129
- if kv_sharing_target_layer_name is not None:
130
- raise NotImplementedError("KV sharing is not supported in V0.")
131
- if use_irope:
132
- logger.warning_once(
133
- "Using irope in Ipex is not supported yet, it will fall"
134
- " back to global attention for long context.")
135
- if blocksparse_params is not None:
136
- raise ValueError(
137
- "IPEX backend does not support block-sparse attention.")
138
- self.num_heads = num_heads
139
- self.head_size = head_size
140
- self.scale = float(scale)
141
- self.num_kv_heads = num_kv_heads
142
- if alibi_slopes is not None:
143
- alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
144
- self.alibi_slopes = alibi_slopes
145
- self.sliding_window = sliding_window
146
- self.kv_cache_dtype = kv_cache_dtype
147
-
148
- self.num_queries_per_kv = self.num_heads // self.num_kv_heads
149
- self.need_mask = (self.sliding_window is not None)
150
- if logits_soft_cap is None:
151
- logits_soft_cap = -1
152
- self.logits_soft_cap = logits_soft_cap
153
-
154
- supported_head_sizes = PagedAttention.get_supported_head_sizes()
155
- if head_size not in supported_head_sizes:
156
- raise ValueError(
157
- f"Head size {head_size} is not supported by PagedAttention. "
158
- f"Supported head sizes are: {supported_head_sizes}.")
159
- if is_quantized_kv_cache(kv_cache_dtype):
160
- raise NotImplementedError(
161
- "IPEX backend does not support FP8 KV cache. "
162
- "Please use xFormers backend instead.")
163
- if attn_type != AttentionType.DECODER:
164
- raise NotImplementedError("Encoder self-attention and "
165
- "encoder/decoder cross-attention "
166
- "are not implemented for "
167
- "IpexAttnBackendImpl")
168
-
169
- def split_kv_cache(
170
- self,
171
- kv_cache: torch.Tensor,
172
- num_kv_heads: int,
173
- head_size: int,
174
- ) -> Tuple[torch.Tensor, torch.Tensor]:
175
- x = 1
176
- num_blocks = kv_cache.shape[1]
177
-
178
- key_cache = kv_cache[0]
179
- key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x,
180
- -1, x)
181
- value_cache = kv_cache[1]
182
- value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
183
- return key_cache, value_cache
184
-
185
- def forward(
186
- self,
187
- layer: AttentionLayer,
188
- query: torch.Tensor,
189
- key: torch.Tensor,
190
- value: torch.Tensor,
191
- kv_cache: torch.Tensor,
192
- attn_metadata: IpexAttnMetadata, # type: ignore
193
- output: Optional[torch.Tensor] = None,
194
- output_scale: Optional[torch.Tensor] = None,
195
- ) -> torch.Tensor:
196
- """Forward pass with IPEX varlen_attention and PagedAttention.
197
-
198
- Args:
199
- query: shape = [num_tokens, num_heads * head_size]
200
- key: shape = [num_tokens, num_kv_heads * head_size]
201
- value: shape = [num_tokens, num_kv_heads * head_size]
202
- kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
203
- NOTE: kv_cache will be an empty tensor with shape [0]
204
- for profiling run.
205
- attn_metadata: Metadata for attention.
206
- Returns:
207
- shape = [num_tokens, num_heads * head_size]
208
- """
209
- if output_scale is not None:
210
- raise NotImplementedError(
211
- "fused output quantization is not yet supported"
212
- " for IpexAttentionImpl")
213
-
214
- assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
215
- num_tokens, hidden_size = query.shape
216
- # Reshape the query, key, and value tensors.
217
- query = query.view(-1, self.num_heads, self.head_size)
218
- key = key.view(-1, self.num_kv_heads, self.head_size)
219
- value = value.view(-1, self.num_kv_heads, self.head_size)
220
-
221
- if kv_cache.numel() > 0:
222
- key_cache, value_cache = self.split_kv_cache(
223
- kv_cache, self.num_kv_heads, self.head_size)
224
- ipex_ops.reshape_and_cache(
225
- key,
226
- value,
227
- key_cache,
228
- value_cache,
229
- attn_metadata.slot_mapping.flatten(),
230
- self.kv_cache_dtype,
231
- layer._k_scale_float,
232
- layer._v_scale_float,
233
- )
234
-
235
- if attn_metadata.is_prompt:
236
- assert attn_metadata.seq_lens is not None
237
- if (kv_cache.numel() == 0
238
- or attn_metadata.block_tables.numel() == 0):
239
- if self.num_kv_heads != self.num_heads:
240
- key = key.repeat_interleave(self.num_queries_per_kv, dim=1)
241
- value = value.repeat_interleave(self.num_queries_per_kv,
242
- dim=1)
243
-
244
- if attn_metadata.attn_bias is None:
245
- if self.sliding_window is not None:
246
- att_masks = _make_sliding_window_bias(
247
- attn_metadata.seq_lens, self.sliding_window,
248
- query.dtype) # type: ignore
249
- else:
250
- att_masks = _make_sliding_window_bias(
251
- attn_metadata.seq_lens, None, dtype=query.dtype)
252
- attn_metadata.attn_bias = att_masks
253
-
254
- output = torch.empty(
255
- (num_tokens, self.num_heads, self.head_size),
256
- dtype=query.dtype,
257
- device=query.device)
258
- ipex_ops.varlen_attention(
259
- query,
260
- key,
261
- value,
262
- output,
263
- attn_metadata.seqlen_q,
264
- attn_metadata.seqlen_q,
265
- self.alibi_slopes,
266
- attn_metadata.max_seqlen,
267
- attn_metadata.max_seqlen,
268
- pdropout=0.0,
269
- softmax_scale=self.scale,
270
- zero_tensors=False,
271
- is_causal=True,
272
- return_softmax=False,
273
- gen_=None,
274
- window_size_left=-1,
275
- window_size_right=-1,
276
- logits_soft_cap=self.logits_soft_cap,
277
- )
278
- else:
279
- # prefix-enabled attention
280
- raise RuntimeError(
281
- "IPEX backend doesn't support prefix decoding.")
282
-
283
- else:
284
- # Decoding run.
285
- max_seq_len = attn_metadata.max_decode_seq_len
286
- output = torch.empty_like(query)
287
- block_size = value_cache.shape[3]
288
- num_seqs, num_heads, head_size = query.shape
289
- max_num_partitions = ((max_seq_len + _PARTITION_SIZE - 1) //
290
- _PARTITION_SIZE)
291
- # NOTE(woosuk): We use a simple heuristic to decide whether to use
292
- # PagedAttention V1 or V2. If the number of partitions is 1, we use
293
- # V1 to avoid the overhead of reduction. Also, if the number of
294
- # sequences or heads is large, we use V1 since there is enough work
295
- # to parallelize.
296
- # TODO(woosuk): Tune this heuristic.
297
- # For context len > 8192, use V2 kernel to avoid shared memory
298
- # shortage.
299
- use_v1 = (max_seq_len <= 8192 and
300
- (max_num_partitions == 1 or num_seqs * num_heads > 512))
301
- if use_v1:
302
- # Run PagedAttention V1.
303
- ipex_ops.paged_attention_v1(
304
- output,
305
- query,
306
- key_cache,
307
- value_cache,
308
- self.num_kv_heads,
309
- self.scale,
310
- attn_metadata.block_tables,
311
- attn_metadata.seq_lens_tensor,
312
- block_size,
313
- max_seq_len,
314
- self.alibi_slopes,
315
- self.kv_cache_dtype,
316
- layer._k_scale_float,
317
- layer._v_scale_float,
318
- )
319
- else:
320
- # Run PagedAttention V2.
321
- assert _PARTITION_SIZE % block_size == 0
322
- tmp_output = torch.empty(
323
- size=(num_seqs, num_heads, max_num_partitions, head_size),
324
- dtype=output.dtype,
325
- device=output.device,
326
- )
327
- exp_sums = torch.empty(
328
- size=(num_seqs, num_heads, max_num_partitions),
329
- dtype=torch.float32,
330
- device=output.device,
331
- )
332
- max_logits = torch.empty_like(exp_sums)
333
- ipex_ops.paged_attention_v2(
334
- output,
335
- exp_sums,
336
- max_logits,
337
- tmp_output,
338
- query,
339
- key_cache,
340
- value_cache,
341
- self.num_kv_heads,
342
- self.scale,
343
- attn_metadata.block_tables,
344
- attn_metadata.seq_lens_tensor,
345
- block_size,
346
- max_seq_len,
347
- self.alibi_slopes,
348
- self.kv_cache_dtype,
349
- layer._k_scale_float,
350
- layer._v_scale_float,
351
- )
352
-
353
- # Reshape the output tensor.
354
- return output.view(-1, self.num_heads * self.head_size)
355
-
356
-
357
-def _make_alibi_bias(
358
- alibi_slopes: torch.Tensor,
359
- dtype: torch.dtype,
360
- seq_lens: List[int],
361
-) -> List[torch.Tensor]:
362
- attn_biases = []
363
- for seq_len in seq_lens:
364
- bias = torch.arange(seq_len, dtype=dtype, device=alibi_slopes.device)
365
- # NOTE(zhuohan): HF uses
366
- # `bias = bias[None, :].repeat(seq_len, 1)`
367
- # here. We find that both biases give the same results, but
368
- # the bias below more accurately follows the original ALiBi
369
- # paper.
370
- bias = bias[None, :] - bias[:, None]
371
-
372
- num_heads = alibi_slopes.shape[0]
373
- bias = bias[None, :].repeat((num_heads, 1, 1))
374
- bias.mul_(alibi_slopes[:, None, None])
375
- inf_mask = torch.empty(
376
- (1, seq_len, seq_len),
377
- dtype=bias.dtype,
378
- device=alibi_slopes.device).fill_(-torch.inf).triu_(diagonal=1)
379
- attn_biases.append((bias + inf_mask).to(dtype))
380
-
381
- return attn_biases
382
-
383
-
384
-def _make_sliding_window_bias(
385
- seq_lens: List[int],
386
- window_size: Optional[int],
387
- dtype: torch.dtype,
388
-) -> List[torch.Tensor]:
389
- attn_biases = []
390
- for seq_len in seq_lens:
391
- tensor = torch.full(
392
- (1, seq_len, seq_len),
393
- dtype=dtype,
394
- fill_value=1,
395
- )
396
- shift = 0
397
- mask = torch.tril(tensor, diagonal=shift).to(dtype) # type: ignore
398
- if window_size is not None:
399
- mask = torch.triu(mask, diagonal=shift - window_size + 1)
400
- mask = torch.log(mask)
401
- attn_biases.append(mask.to(dtype))
402
-
403
- return attn_biases
vllm/attention/backends/pallas.py
deleted
-356
@@ -1,356 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass
5
-from typing import Any, Dict, List, Optional, Tuple, Type
6
-
7
-import torch
8
-import torch_xla.experimental.custom_kernel # Required to register custom ops.
9
-
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
15
-from vllm.logger import init_logger
16
-
17
-logger = init_logger(__name__)
18
-
19
-
20
-class PallasAttentionBackend(AttentionBackend):
21
-
22
- @staticmethod
23
- def get_name() -> str:
24
- return "PALLAS"
25
-
26
- @staticmethod
27
- def get_impl_cls() -> Type["PallasAttentionBackendImpl"]:
28
- return PallasAttentionBackendImpl
29
-
30
- @staticmethod
31
- def get_metadata_cls() -> Type["PallasMetadata"]:
32
- return PallasMetadata
33
-
34
- @staticmethod
35
- def get_state_cls() -> Type["CommonAttentionState"]:
36
- return CommonAttentionState
37
-
38
- @staticmethod
39
- def get_kv_cache_shape(
40
- num_blocks: int,
41
- block_size: int,
42
- num_kv_heads: int,
43
- head_size: int,
44
- ) -> Tuple[int, ...]:
45
- return (num_kv_heads, num_blocks, block_size, head_size)
46
-
47
- @staticmethod
48
- def swap_blocks(
49
- src_kv_cache: torch.Tensor,
50
- dst_kv_cache: torch.Tensor,
51
- src_to_dst: torch.Tensor,
52
- ) -> None:
53
- raise RuntimeError("swap_blocks is not used for the TPU backend.")
54
-
55
- @torch.compile(backend="openxla")
56
- @staticmethod
57
- def copy_blocks(
58
- kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
59
- src_to_dists: Tuple[torch.Tensor, torch.Tensor],
60
- ) -> None:
61
- src_indices, dst_indices = src_to_dists
62
- for k_cache, v_cache in kv_caches:
63
- torch.ops.xla.dynamo_set_buffer_donor_(k_cache, True)
64
- k_cache[:, dst_indices] = k_cache[:, src_indices]
65
- torch.ops.xla.dynamo_set_buffer_donor_(v_cache, True)
66
- v_cache[:, dst_indices] = v_cache[:, src_indices]
67
-
68
-
69
-@dataclass
70
-class PallasMetadata(AttentionMetadata):
71
-
72
- # Currently, input sequences can only contain all prefills
73
- # or all decoding.
74
- block_tables: Optional[torch.Tensor] = None
75
- context_lens: Optional[torch.Tensor] = None
76
- effective_query_lens: Optional[torch.Tensor] = None
77
-
78
- @property
79
- def prefill_metadata(self) -> Optional["PallasMetadata"]:
80
- if self.num_prefills == 0:
81
- return None
82
-
83
- assert self.num_decode_tokens == 0
84
- return self
85
-
86
- @property
87
- def decode_metadata(self) -> Optional["PallasMetadata"]:
88
- if self.num_decode_tokens == 0:
89
- return None
90
-
91
- assert self.num_prefills == 0
92
- assert self.num_prefill_tokens == 0
93
- assert self.block_tables is not None
94
- assert self.context_lens is not None
95
- return self
96
-
97
-
98
-class PallasAttentionBackendImpl(AttentionImpl):
99
-
100
- def __init__(
101
- self,
102
- num_heads: int,
103
- head_size: int,
104
- scale: float,
105
- num_kv_heads: int,
106
- alibi_slopes: Optional[List[float]],
107
- sliding_window: Optional[int],
108
- kv_cache_dtype: str,
109
- blocksparse_params: Optional[Dict[str, Any]] = None,
110
- logits_soft_cap: Optional[float] = None,
111
- attn_type: str = AttentionType.DECODER,
112
- kv_sharing_target_layer_name: Optional[str] = None,
113
- use_irope: bool = False,
114
- ) -> None:
115
- if kv_sharing_target_layer_name is not None:
116
- raise NotImplementedError("KV sharing is not supported in V0.")
117
- if use_irope:
118
- logger.warning_once(
119
- "Using irope in Pallas is not supported yet, it will fall back "
120
- "to global attention for long context.")
121
- self.num_heads = num_heads
122
- self.head_size = head_size
123
- self.scale = float(scale)
124
- self.num_kv_heads = num_kv_heads
125
-
126
- self.num_queries_per_kv = self.num_heads // self.num_kv_heads
127
- self.logits_soft_cap = logits_soft_cap
128
- if head_size % 128 != 0:
129
- raise NotImplementedError(
130
- f"Head size must be a multiple of 128, found {head_size}.")
131
- if alibi_slopes is not None:
132
- raise NotImplementedError("Alibi slopes is not supported.")
133
- if sliding_window is not None:
134
- raise NotImplementedError("Sliding window is not supported.")
135
- if is_quantized_kv_cache(kv_cache_dtype):
136
- raise NotImplementedError("FP8 KV cache dtype is not supported.")
137
- if blocksparse_params is not None:
138
- raise NotImplementedError("Blocksparse is not supported.")
139
-
140
- if torch_xla.tpu.version() < 4:
141
- raise NotImplementedError("TPU version must be 4 or higher.")
142
-
143
- self.megacore_mode = None
144
- tpu_env = torch_xla.tpu.get_tpu_env()
145
- tpu_type = (tpu_env.get("ACCELERATOR_TYPE", None)
146
- or tpu_env.get("TYPE", None)
147
- or tpu_env.get("TPU_ACCELERATOR_TYPE", None))
148
- assert tpu_type is not None
149
- tpu_type = tpu_type.lower()
150
-
151
- if (("lite" not in tpu_type) and ("v6" not in tpu_type)):
152
- if self.num_kv_heads % 2 == 0:
153
- self.megacore_mode = "kv_head"
154
- else:
155
- # NOTE(woosuk): If the batch size is not a multiple of 2, the
156
- # megacore mode will be None.
157
- self.megacore_mode = "batch"
158
-
159
- if attn_type != AttentionType.DECODER:
160
- raise NotImplementedError("Encoder self-attention and "
161
- "encoder/decoder cross-attention "
162
- "are not implemented for "
163
- "PallasAttentionBackendImpl")
164
-
165
- def forward(
166
- self,
167
- layer: AttentionLayer,
168
- query: torch.Tensor,
169
- key: torch.Tensor,
170
- value: torch.Tensor,
171
- kv_cache: Tuple[torch.Tensor, torch.Tensor],
172
- attn_metadata: PallasMetadata,
173
- output: Optional[torch.Tensor] = None,
174
- output_scale: Optional[torch.Tensor] = None,
175
- ) -> torch.Tensor:
176
- """Forward pass with Pallas attention.
177
-
178
- Args:
179
- query: shape = [batch_size, seq_len, num_heads * head_size]
180
- key: shape = [batch_size, seq_len, num_kv_heads * head_size]
181
- value: shape = [batch_size, seq_len, num_kv_heads * head_size]
182
- kv_cache[0] = [num_kv_heads, num_blocks, block_size, head_size]
183
- kv_cache[1] = [num_kv_heads, num_blocks, block_size, head_size]
184
- NOTE: kv_cache[0] and kv_cache[1] will be an empty tensor
185
- with shape [0] for profiling run.
186
- attn_metadata: Metadata for attention.
187
- Returns:
188
- shape = [batch_size, seq_len, num_heads * head_size]
189
- """
190
- if output_scale is not None:
191
- raise NotImplementedError(
192
- "fused output quantization is not yet supported"
193
- " for PallasAttentionImpl")
194
-
195
- assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
196
- batch_size, seq_len, hidden_size = query.shape
197
- query = query.view(batch_size, seq_len, self.num_heads, self.head_size)
198
- key = key.view(batch_size, seq_len, self.num_kv_heads, self.head_size)
199
- value = value.view(batch_size, seq_len, self.num_kv_heads,
200
- self.head_size)
201
-
202
- if kv_cache[0].numel() > 0:
203
- slot_mapping = attn_metadata.slot_mapping
204
- key_cache, value_cache = kv_cache
205
- write_to_kv_cache(key, value, key_cache, value_cache, slot_mapping)
206
-
207
- query = query * self.scale
208
- if attn_metadata.num_prefills > 0:
209
- if attn_metadata.block_tables is None:
210
- # Prefill without paged KV cache.
211
- assert seq_len % 16 == 0, (
212
- "Pallas FlashAttention kernel requires seq_len to be a "
213
- f"multiple of 16 but got {seq_len}")
214
-
215
- # Handle GQA/MQA.
216
- if self.num_kv_heads != self.num_heads:
217
- key = key.repeat_interleave(self.num_queries_per_kv,
218
- dim=-2)
219
- key = key.view(batch_size, seq_len, self.num_heads,
220
- self.head_size)
221
- value = value.repeat_interleave(self.num_queries_per_kv,
222
- dim=-2)
223
- value = value.view(batch_size, seq_len, self.num_heads,
224
- self.head_size)
225
- # FlashAttention kernel requires the input shape to be
226
- # [batch_size, num_heads, seq_len, d_model]
227
- # while the input is [batch_size, seq_len, num_heads, d_model].
228
- # Permute the input to match the required format.
229
- output = torch.ops.xla.flash_attention(
230
- query.permute(0, 2, 1, 3),
231
- key.permute(0, 2, 1, 3),
232
- value.permute(0, 2, 1, 3),
233
- True,
234
- )
235
- output = output.permute(0, 2, 1, 3)
236
- else:
237
- # Prefill with paged KV cache.
238
- # TODO(woosuk): Tune the below knobs.
239
- num_kv_pages_per_compute_block = 16
240
- num_queries_per_compute_block = 16
241
- assert seq_len % num_queries_per_compute_block == 0
242
- output = torch.ops.xla.multi_queries_paged_attention(
243
- query,
244
- key_cache,
245
- value_cache,
246
- attn_metadata.context_lens,
247
- attn_metadata.block_tables,
248
- attn_metadata.effective_query_lens,
249
- num_kv_pages_per_compute_block,
250
- num_queries_per_compute_block,
251
- use_kernel=True,
252
- attn_logits_soft_cap=self.logits_soft_cap,
253
- )
254
- else:
255
- # Decoding run.
256
- assert kv_cache[0].numel() > 0
257
- query = query.squeeze(dim=1)
258
- pages_per_compute_block = 16 # TODO(woosuk): Tune this value.
259
-
260
- assert attn_metadata.block_tables is not None
261
- assert attn_metadata.context_lens is not None
262
- # NOTE(woosuk): The PagedAttention Pallas kernel stores the entire
263
- # block table in SMEM. Therefore, if the block table is too large,
264
- # the kernel compilation will fail. To avoid this, we split the
265
- # batch dimension into smaller chunks and run the kernel multiple
266
- # times.
267
- MAX_SMEM_USAGE = 512 * 1024
268
- size_per_seq = 4 * attn_metadata.block_tables.shape[1]
269
- max_num_seq = MAX_SMEM_USAGE // size_per_seq
270
-
271
- if batch_size <= max_num_seq:
272
- output = paged_attention(
273
- query,
274
- key_cache,
275
- value_cache,
276
- attn_metadata.context_lens,
277
- attn_metadata.block_tables,
278
- pages_per_compute_block,
279
- self.megacore_mode,
280
- attn_logits_soft_cap=self.logits_soft_cap,
281
- )
282
- else:
283
- chunk_size = max_num_seq
284
- # Make sure the chunk size is a multiple of 2.
285
- chunk_size = chunk_size // 2 * 2
286
- num_chunks = (batch_size + chunk_size - 1) // chunk_size
287
-
288
- output = torch.empty_like(query)
289
- for chunk_idx in range(num_chunks):
290
- chunk_start = chunk_idx * chunk_size
291
- chunk_end = chunk_start + chunk_size
292
- # NOTE(woosuk): We skip this line because it causes Dynamo
293
- # compilation error. Instead, we rely on the slice operation
294
- # to handle the out-of-bound case.
295
- # chunk_end = min(chunk_end, batch_size)
296
- chunk_output = paged_attention(
297
- query[chunk_start:chunk_end],
298
- key_cache,
299
- value_cache,
300
- attn_metadata.context_lens[chunk_start:chunk_end],
301
- attn_metadata.block_tables[chunk_start:chunk_end],
302
- pages_per_compute_block,
303
- self.megacore_mode,
304
- attn_logits_soft_cap=self.logits_soft_cap,
305
- )
306
- output[chunk_start:chunk_end] = chunk_output
307
-
308
- # Reshape the output tensor.
309
- return output.reshape(batch_size, seq_len, hidden_size)
310
-
311
-
312
-def write_to_kv_cache(
313
- key: torch.Tensor,
314
- value: torch.Tensor,
315
- key_cache: torch.Tensor,
316
- value_cache: torch.Tensor,
317
- slot_mapping: torch.Tensor,
318
-) -> None:
319
- torch.ops.xla.dynamo_set_buffer_donor_(key_cache, True)
320
- torch.ops.xla.dynamo_set_buffer_donor_(value_cache, True)
321
-
322
- key = key.flatten(0, 2)
323
- value = value.flatten(0, 2)
324
- key_cache = key_cache.flatten(0, 2)
325
- value_cache = value_cache.flatten(0, 2)
326
- key_cache.index_copy_(0, slot_mapping, key)
327
- value_cache.index_copy_(0, slot_mapping, value)
328
-
329
-
330
-def paged_attention(
331
- query: torch.Tensor,
332
- key_cache: torch.Tensor,
333
- value_cache: torch.Tensor,
334
- context_lens: torch.Tensor,
335
- block_tables: torch.Tensor,
336
- pages_per_compute_block: int,
337
- megacore_mode: Optional[str],
338
- *,
339
- attn_logits_soft_cap: Optional[float],
340
-) -> torch.Tensor:
341
- batch_size = query.shape[0]
342
- if megacore_mode == "batch" and batch_size % 2 != 0:
343
- megacore_mode = None
344
- else:
345
- megacore_mode = megacore_mode
346
-
347
- return torch.ops.xla.paged_attention(
348
- query,
349
- key_cache,
350
- value_cache,
351
- context_lens,
352
- block_tables,
353
- pages_per_compute_block,
354
- megacore_mode=megacore_mode,
355
- attn_logits_soft_cap=attn_logits_soft_cap,
356
- )
vllm/attention/backends/torch_sdpa.py
+3
-164
@@ -3,78 +3,24 @@
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, Tuple, Type
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 (AttentionBackend, AttentionImpl,
14
- AttentionLayer,
15
- AttentionMetadata,
16
- AttentionMetadataBuilder,
17
- AttentionType,
13
+from vllm.attention.backends.abstract import (AttentionImpl, AttentionLayer,
14
+ AttentionMetadata, AttentionType,
15
is_quantized_kv_cache)
16
# yapf: enable
20
-from vllm.attention.backends.utils import CommonAttentionState
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
24
-from vllm.utils import make_tensor_with_pad
25
-from vllm.worker.cpu_model_runner import ModelInputForCPUBuilder
20
21
logger = init_logger(__name__)
22
23
30
-class TorchSDPABackend(AttentionBackend):
31
-
32
- @staticmethod
33
- def get_name() -> str:
34
- return "TORCH_SDPA"
35
-
36
- @staticmethod
37
- def get_impl_cls() -> Type["TorchSDPABackendImpl"]:
38
- return TorchSDPABackendImpl
39
-
40
- @staticmethod
41
- def get_metadata_cls() -> Type["AttentionMetadata"]:
42
- return TorchSDPAMetadata
43
-
44
- @staticmethod
45
- def get_state_cls() -> Type["CommonAttentionState"]:
46
- return CommonAttentionState
47
-
48
- @staticmethod
49
- def get_builder_cls() -> Type["TorchSDPAMetadataBuilder"]:
50
- return TorchSDPAMetadataBuilder
51
-
52
- @staticmethod
53
- def get_kv_cache_shape(
54
- num_blocks: int,
55
- block_size: int,
56
- num_kv_heads: int,
57
- head_size: int,
58
- ) -> Tuple[int, ...]:
59
- return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
60
- num_kv_heads, head_size)
61
-
62
- @staticmethod
63
- def swap_blocks(
64
- src_kv_cache: torch.Tensor,
65
- dst_kv_cache: torch.Tensor,
66
- src_to_dst: torch.Tensor,
67
- ) -> None:
68
- raise NotImplementedError("Swap is not supported in TorchSDPABackend.")
69
-
70
- @staticmethod
71
- def copy_blocks(
72
- kv_caches: List[torch.Tensor],
73
- src_to_dists: torch.Tensor,
74
- ) -> None:
75
- PagedAttention.copy_blocks(kv_caches, src_to_dists)
76
-
77
-
24
@dataclass
25
class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
26
"""Metadata for TorchSDPABackend.
@@ -287,113 +233,6 @@ class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
233
raise AttributeError(f"Invalid attention type {str(attn_type)}")
234
235
290
-class TorchSDPAMetadataBuilder(AttentionMetadataBuilder[TorchSDPAMetadata]):
291
-
292
- def __init__(self, input_builder: ModelInputForCPUBuilder) -> None:
293
- self.chunked_prefill = input_builder.chunked_prefill
294
- self.input_builder = input_builder
295
-
296
- def prepare(self):
297
- self.input_data = self.input_builder.input_data
298
-
299
- def build(self, seq_lens: List[int], query_lens: List[int],
300
- cuda_graph_pad_size: int, batch_size: int) -> TorchSDPAMetadata:
301
- input_data = self.input_data
302
- prefill_seq_lens = seq_lens[0:input_data.num_prefills]
303
- prefill_query_lens = query_lens[0:input_data.num_prefills]
304
- slot_mapping = torch.tensor(input_data.slot_mapping,
305
- dtype=torch.long,
306
- device="cpu")
307
-
308
- # For chunked-prefill
309
- if self.chunked_prefill and input_data.num_prefill_tokens != 0:
310
- prefill_block_tables = make_tensor_with_pad(
311
- self.input_data.prefill_block_tables,
312
- pad=0,
313
- dtype=torch.int32,
314
- device="cpu",
315
- )
316
- query_lens_tensor = torch.tensor(prefill_query_lens,
317
- dtype=torch.int32,
318
- device="cpu")
319
- kv_lens_tensor = torch.tensor(prefill_seq_lens,
320
- dtype=torch.int32,
321
- device="cpu")
322
- query_start_loc = torch.zeros(input_data.num_prefills + 1,
323
- dtype=torch.int32,
324
- device="cpu")
325
- kv_start_loc = torch.zeros(input_data.num_prefills + 1,
326
- dtype=torch.int32,
327
- device="cpu")
328
- torch.cumsum(query_lens_tensor,
329
- dim=0,
330
- dtype=torch.int32,
331
- out=query_start_loc[1:])
332
- torch.cumsum(kv_lens_tensor,
333
- dim=0,
334
- dtype=torch.int32,
335
- out=kv_start_loc[1:])
336
- max_query_len = max(prefill_query_lens)
337
- max_kv_len = max(prefill_seq_lens)
338
- else:
339
- prefill_block_tables = None
340
- query_start_loc = None
341
- kv_start_loc = None
342
- max_query_len = None
343
- max_kv_len = None
344
-
345
- # For paged attention
346
- if input_data.num_decode_tokens != 0:
347
- seq_lens_tensor = torch.tensor(
348
- input_data.seq_lens[input_data.num_prefills:],
349
- dtype=torch.int32,
350
- device="cpu",
351
- )
352
- block_tables = make_tensor_with_pad(
353
- self.input_data.decode_block_tables,
354
- pad=0,
355
- dtype=torch.int32,
356
- device="cpu",
357
- )
358
- else:
359
- block_tables = torch.tensor([])
360
- seq_lens_tensor = torch.tensor(
361
- input_data.seq_lens[:input_data.num_prefills],
362
- dtype=torch.int32,
363
- device="cpu",
364
- )
365
-
366
- # For multi-modal models
367
- placeholder_index_maps = None
368
- if len(input_data.multi_modal_inputs_list) != 0:
369
- placeholder_index_maps = {
370
- modality: placeholder_map.index_map()
371
- for modality, placeholder_map in
372
- input_data.multi_modal_placeholder_maps.items()
373
- }
374
-
375
- attn_metadata = TorchSDPAMetadata(
376
- chunked_prefill=self.chunked_prefill,
377
- seq_lens=prefill_seq_lens,
378
- seq_lens_tensor=seq_lens_tensor,
379
- max_query_len=max_query_len,
380
- max_kv_len=max_kv_len,
381
- prefill_query_start_loc=query_start_loc,
382
- kv_start_loc=kv_start_loc,
383
- max_decode_seq_len=input_data.max_decode_seq_len,
384
- num_prefills=input_data.num_prefills,
385
- num_prefill_tokens=input_data.num_prefill_tokens,
386
- num_decode_tokens=input_data.num_decode_tokens,
387
- block_tables=block_tables,
388
- prefill_block_tables=prefill_block_tables,
389
- slot_mapping=slot_mapping,
390
- multi_modal_placeholder_index_maps=placeholder_index_maps,
391
- enable_kv_scales_calculation=False,
392
- )
393
-
394
- return attn_metadata
395
-
396
-
236
class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
237
238
def __init__(
vllm/platforms/cpu.py
+6
-20
@@ -64,13 +64,11 @@ class CpuPlatform(Platform):
64
if selected_backend and selected_backend != _Backend.TORCH_SDPA:
65
logger.info("Cannot use %s backend on CPU.", selected_backend)
66
if use_mla:
67
- logger.info("Using CPU MLA backend.")
68
- return "vllm.attention.backends.cpu_mla.CPUMLABackend"
67
+ raise NotImplementedError("MLA is not supported on CPU.")
68
logger.info("Using Torch SDPA backend.")
70
- if use_v1:
71
- return "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
72
- else:
73
- return "vllm.attention.backends.torch_sdpa.TorchSDPABackend"
69
+ if not use_v1:
70
+ raise ValueError("CPU backend only supports V1.")
71
+ return "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
72
73
@classmethod
74
def get_device_total_memory(cls, device_id: int = 0) -> int:
@@ -147,26 +145,14 @@ class CpuPlatform(Platform):
145
parallel_config.distributed_executor_backend)
146
parallel_config.distributed_executor_backend = "mp"
147
if parallel_config.worker_cls == "auto":
150
- if vllm_config.speculative_config:
151
- parallel_config.worker_cls = \
152
- "vllm.spec_decode.spec_decode_worker.create_spec_worker"
153
- parallel_config.sd_worker_cls = \
154
- "vllm.worker.cpu_worker.CPUWorker"
155
- else:
156
- if envs.VLLM_USE_V1:
157
- parallel_config.worker_cls = \
158
- "vllm.v1.worker.cpu_worker.CPUWorker"
159
- else:
160
- parallel_config.worker_cls = \
161
- "vllm.worker.cpu_worker.CPUWorker"
148
+ parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker"
149
150
# Note: workaround for v1 gpu_model_runner
151
from vllm.config import CompilationLevel
152
vllm_config.compilation_config.cudagraph_capture_sizes = []
153
154
compilation_config = vllm_config.compilation_config
168
- if (envs.VLLM_USE_V1 and vllm_config.compilation_config.level
169
- == CompilationLevel.PIECEWISE):
155
+ if vllm_config.compilation_config.level == CompilationLevel.PIECEWISE:
156
157
# Note: vLLM V1 is using PIECEWISE level compilation, which will
158
# take time to compile kernels just-in-time with the inductor
vllm/platforms/tpu.py
+16
-35
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Optional, Union, cast
6
import torch
7
from tpu_info import device
8
9
-import vllm.envs as envs
9
from vllm.inputs import ProcessorInputs, PromptType
10
from vllm.logger import init_logger
11
from vllm.sampling_params import SamplingParams, SamplingType
@@ -50,12 +49,10 @@ class TpuPlatform(Platform):
49
and selected_backend != _Backend.PALLAS_VLLM_V1):
50
logger.info("Cannot use %s backend on TPU.", selected_backend)
51
53
- if use_v1:
54
- logger.info("Using Pallas V1 backend.")
55
- return "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
56
- else:
57
- logger.info("Using Pallas backend.")
58
- return "vllm.attention.backends.pallas.PallasAttentionBackend"
52
+ if not use_v1:
53
+ raise ValueError("TPU backend only supports V1.")
54
+ logger.info("Using Pallas V1 backend.")
55
+ return "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
56
57
@classmethod
58
def get_device_name(cls, device_id: int = 0) -> str:
@@ -68,7 +65,7 @@ class TpuPlatform(Platform):
65
66
@classmethod
67
def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
71
- return not envs.VLLM_USE_V1
68
+ return False
69
70
@classmethod
71
def get_punica_wrapper(cls) -> str:
@@ -117,31 +114,19 @@ class TpuPlatform(Platform):
114
"Using bfloat16 instead.", vllm_config.model_config.dtype)
115
vllm_config.model_config.dtype = torch.bfloat16
116
120
- if envs.VLLM_USE_V1:
121
- from vllm.v1.attention.backends.pallas import (
122
- PallasAttentionBackend)
123
- cache_config.block_size = PallasAttentionBackend.get_page_size(
124
- vllm_config) # type: ignore[assignment]
117
+ from vllm.v1.attention.backends.pallas import PallasAttentionBackend
118
+ cache_config.block_size = PallasAttentionBackend.get_page_size(
119
+ vllm_config) # type: ignore[assignment]
120
121
parallel_config = vllm_config.parallel_config
122
scheduler_config = vllm_config.scheduler_config
123
if parallel_config.worker_cls == "auto":
124
if scheduler_config.is_multi_step:
130
- if envs.VLLM_USE_V1:
131
- raise NotImplementedError(
132
- "Multi-step scheduling is not supported (and not "
133
- "needed) on vLLM V1. Please launch without "
134
- "--num-scheduler-steps.")
135
- else:
136
- parallel_config.worker_cls = \
137
- "vllm.worker.multi_step_tpu_worker.MultiStepTPUWorker"
138
- else:
139
- if envs.VLLM_USE_V1:
140
- parallel_config.worker_cls = \
141
- "vllm.v1.worker.tpu_worker.TPUWorker"
142
- else:
143
- parallel_config.worker_cls = \
144
- "vllm.worker.tpu_worker.TPUWorker"
125
+ raise NotImplementedError(
126
+ "Multi-step scheduling is not supported (and not "
127
+ "needed) on vLLM V1. Please launch without "
128
+ "--num-scheduler-steps.")
129
+ parallel_config.worker_cls = "vllm.v1.worker.tpu_worker.TPUWorker"
130
131
assert not vllm_config.speculative_config, (
132
"Speculative decoding is not yet supported for TPU backend")
@@ -189,13 +174,9 @@ class TpuPlatform(Platform):
174
processed_inputs: ProcessorInputs,
175
) -> None:
176
"""Raises if this request is unsupported on this platform"""
192
- if isinstance(params, SamplingParams):
193
- if params.guided_decoding is not None and not envs.VLLM_USE_V1:
194
- raise ValueError("Structured output is not supported on "
195
- f"{cls.device_name} V0.")
196
- if params.sampling_type == SamplingType.RANDOM_SEED:
197
- raise ValueError(
198
- "Torch XLA does not support per-request seed.")
177
+ if (isinstance(params, SamplingParams)
178
+ and params.sampling_type == SamplingType.RANDOM_SEED):
179
+ raise ValueError("Torch XLA does not support per-request seed.")
180
181
182
try:
vllm/platforms/xpu.py
+6
-15
@@ -39,12 +39,10 @@ class XPUPlatform(Platform):
39
if selected_backend != _Backend.IPEX:
40
logger.info("Cannot use %s backend on XPU.", selected_backend)
41
use_v1 = envs.VLLM_USE_V1
42
- if use_v1:
43
- logger.info("Using Flash Attention backend on V1 engine.")
44
- return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
45
- else:
46
- logger.info("Using IPEX attention backend.")
47
- return "vllm.attention.backends.ipex_attn.IpexAttnBackend"
42
+ if not use_v1:
43
+ raise ValueError("XPU backend only supports V1.")
44
+ logger.info("Using Flash Attention backend on V1 engine.")
45
+ return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
46
47
@classmethod
48
def get_device_capability(
@@ -77,10 +75,7 @@ class XPUPlatform(Platform):
75
cache_config = vllm_config.cache_config
76
# in V1(or with ipex chunked prefill) block_size is 64
77
if cache_config and cache_config.block_size is None:
80
- if envs.VLLM_USE_V1:
81
- cache_config.block_size = 64
82
- else:
83
- cache_config.block_size = 16
78
+ cache_config.block_size = 64
79
80
# Instances created using VllmConfig() typically have model_config as
81
# None by default. The modification involves adding a check to prevent
@@ -106,11 +101,7 @@ class XPUPlatform(Platform):
101
102
# check and update parallel config
103
parallel_config = vllm_config.parallel_config
109
- if envs.VLLM_USE_V1:
110
- parallel_config.worker_cls =\
111
- "vllm.v1.worker.xpu_worker.XPUWorker"
112
- else:
113
- parallel_config.worker_cls = "vllm.worker.xpu_worker.XPUWorker"
104
+ parallel_config.worker_cls = "vllm.v1.worker.xpu_worker.XPUWorker"
105
106
if parallel_config.distributed_executor_backend is None:
107
if parallel_config.world_size > 1:
vllm/worker/cpu_enc_dec_model_runner.py
deleted
-326
@@ -1,326 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
6
-
7
-import torch
8
-
9
-from vllm.attention import AttentionMetadata
10
-from vllm.forward_context import set_forward_context
11
-from vllm.model_executor import SamplingMetadata
12
-from vllm.model_executor.layers.sampler import SamplerOutput
13
-from vllm.multimodal import MultiModalKwargs
14
-from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
15
-from vllm.utils import make_tensor_with_pad
16
-from vllm.worker.cpu_model_runner import (CPUModelRunnerBase,
17
- ModelInputForCPUBuilder,
18
- ModelInputForCPUWithSamplingMetadata)
19
-from vllm.worker.model_runner_base import (
20
- _add_attn_metadata_broadcastable_dict,
21
- _add_sampling_metadata_broadcastable_dict)
22
-
23
-if TYPE_CHECKING:
24
- from vllm.attention.backends.abstract import AttentionBackend
25
-
26
-
27
-@dataclasses.dataclass(frozen=True)
28
-class EncoderDecoderModelInputForCPU(ModelInputForCPUWithSamplingMetadata):
29
- """
30
- Used by the EncoderDecoderModelRunner.
31
- """
32
- encoder_input_tokens: Optional[torch.Tensor] = None
33
- encoder_input_positions: Optional[torch.Tensor] = None
34
-
35
- def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
36
- tensor_dict = {
37
- "input_tokens": self.input_tokens,
38
- "input_positions": self.input_positions,
39
- "encoder_input_tokens": self.encoder_input_tokens,
40
- "encoder_input_positions": self.encoder_input_positions,
41
- "multi_modal_kwargs": self.multi_modal_kwargs,
42
- }
43
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
44
- _add_sampling_metadata_broadcastable_dict(tensor_dict,
45
- self.sampling_metadata)
46
- return tensor_dict
47
-
48
- @classmethod
49
- def from_broadcasted_tensor_dict(
50
- cls,
51
- tensor_dict: Dict[str, Any],
52
- attn_backend: Optional["AttentionBackend"] = None,
53
- ) -> "EncoderDecoderModelInputForCPU":
54
- return cast(
55
- EncoderDecoderModelInputForCPU,
56
- super().from_broadcasted_tensor_dict(tensor_dict, attn_backend))
57
-
58
-
59
-class CPUEncoderDecoderModelRunner(
60
- CPUModelRunnerBase[EncoderDecoderModelInputForCPU]):
61
- _model_input_cls: Type[EncoderDecoderModelInputForCPU] = (
62
- EncoderDecoderModelInputForCPU)
63
- _builder_cls: Type[ModelInputForCPUBuilder] = ModelInputForCPUBuilder
64
-
65
- def _list_to_int32_tensor(
66
- self,
67
- _list: List[int],
68
- ) -> torch.Tensor:
69
- return torch.tensor(_list, dtype=torch.int32, device=self.device)
70
-
71
- def _list_to_long_tensor(
72
- self,
73
- _list: List[int],
74
- ) -> torch.Tensor:
75
- return torch.tensor(_list, dtype=torch.long, device=self.device)
76
-
77
- def _empty_int32_tensor(self) -> torch.Tensor:
78
- return self._list_to_int32_tensor([])
79
-
80
- def _empty_long_tensor(self) -> torch.Tensor:
81
- return self._list_to_long_tensor([])
82
-
83
- def make_model_input_from_broadcasted_tensor_dict(
84
- self, tensor_dict: Dict[str,
85
- Any]) -> EncoderDecoderModelInputForCPU:
86
- return EncoderDecoderModelInputForCPU.from_broadcasted_tensor_dict(
87
- tensor_dict,
88
- attn_backend=self.attn_backend,
89
- )
90
-
91
- def prepare_model_input(
92
- self,
93
- seq_group_metadata_list: List[SequenceGroupMetadata],
94
- virtual_engine: int = 0,
95
- finished_requests_ids: Optional[List[str]] = None
96
- ) -> EncoderDecoderModelInputForCPU:
97
- model_input = self._prepare_model_input_tensors(
98
- seq_group_metadata_list, finished_requests_ids)
99
- (
100
- attn_metadata,
101
- encoder_input_tokens_tensor,
102
- encoder_input_positions_tensor,
103
- ) = self._prepare_encoder_model_input_tensors(seq_group_metadata_list,
104
- model_input)
105
- # Sampling metadata is only required for the final pp group
106
- generators = self.get_generators(finished_requests_ids)
107
- sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list,
108
- model_input.seq_lens,
109
- model_input.query_lens,
110
- self.device,
111
- pin_memory=False,
112
- generators=generators)
113
- return dataclasses.replace(
114
- model_input,
115
- sampling_metadata=sampling_metadata,
116
- attn_metadata=attn_metadata,
117
- encoder_input_tokens=encoder_input_tokens_tensor,
118
- encoder_input_positions=encoder_input_positions_tensor,
119
- virtual_engine=virtual_engine,
120
- )
121
-
122
- def _prepare_encoder_model_input_tensors(
123
- self,
124
- seq_group_metadata_list: List[SequenceGroupMetadata],
125
- model_input: EncoderDecoderModelInputForCPU,
126
- ) -> Tuple[AttentionMetadata, Optional[torch.Tensor],
127
- Optional[torch.Tensor]]:
128
- """Helper method to prepare the encoder- and cross-attn-related
129
- model inputs based on a given sequence group. These additional inputs
130
- are used to augment an already-computed `EncoderDecoderModelInput`
131
- data structure which already has decoder-related model inputs
132
- populated.
133
-
134
- Sets the following attn_metadata fields:
135
- * `num_encoder_tokens`
136
- * `encoder_seq_lens`
137
- * `encoder_seq_lens_tensor`
138
- * `max_encoder_seq_len`
139
- * `cross_slot_mapping`
140
- * `cross_block_tables`
141
-
142
- Constructs a new model inputs data structure, based on
143
- (1) the existing fields in the `model_inputs` argument,
144
- and (2) the following additional fields which are
145
- computed (or in the case of `attn_metadata`, updated)
146
- by this function:
147
- * attn_metadata
148
- * encoder_input_tokens
149
- * encoder_input_positions
150
-
151
- Arguments:
152
-
153
- * seq_group_metadata_list: list of sequence groups for which to
154
- compute inputs
155
- * model_inputs: model inputs data structure with decoder-oriented
156
- fields already computed.
157
-
158
- Return:
159
-
160
- * Updated model inputs data structure
161
- """
162
-
163
- if len(seq_group_metadata_list) == 0:
164
- return (model_input.attn_metadata, None, None)
165
-
166
- # Since we are not supporting chunked prefill either the entire
167
- # batch is prefill or it is decode
168
- is_prompt = seq_group_metadata_list[0].is_prompt
169
-
170
- # Build encoder inputs
171
- encoder_seq_lens: List[int] = []
172
- if is_prompt:
173
- # Prefill phase.
174
- cross_block_tables = self._empty_int32_tensor().view(
175
- len(seq_group_metadata_list), -1)
176
-
177
- # Extract input tokens/positions, cross-attention slot-mapping,
178
- # & seq len from each sequence group metadata
179
- (
180
- encoder_input_tokens,
181
- encoder_input_positions,
182
- cross_slot_mapping,
183
- ) = (
184
- [],
185
- [],
186
- [],
187
- )
188
- for seq_group_metadata in seq_group_metadata_list:
189
- # Build seq lens
190
- seq_len = seq_group_metadata.encoder_seq_data.get_len()
191
- token_ids = seq_group_metadata.encoder_seq_data.get_token_ids()
192
- encoder_seq_lens.append(seq_len)
193
-
194
- # Build slot mapping
195
- for i in range(0, seq_len):
196
- block_number = seq_group_metadata.cross_block_table[
197
- i // self.block_size]
198
- block_offset = i % self.block_size
199
- slot = block_number * self.block_size + block_offset
200
- cross_slot_mapping.append(slot)
201
-
202
- # Build encoder input tokens
203
- encoder_input_tokens.extend(token_ids)
204
- encoder_input_positions.extend(list(range(0, seq_len)))
205
-
206
- # Convert tokens/positions & cross-attention
207
- # slot-mapping to encoder input tensors
208
- encoder_input_tokens_tensor = self._list_to_long_tensor(
209
- encoder_input_tokens)
210
- encoder_input_positions_tensor = self._list_to_long_tensor(
211
- encoder_input_positions)
212
- cross_slot_mapping_tensor = self._list_to_long_tensor(
213
- cross_slot_mapping)
214
-
215
- else:
216
- # Decode phase.
217
- encoder_input_tokens_tensor = self._empty_long_tensor()
218
- encoder_input_positions_tensor = self._empty_long_tensor()
219
- cross_slot_mapping_tensor = self._empty_long_tensor()
220
- # Extract cross-attention block tables &
221
- # seq len from each sequence group metadata.
222
- # Cross-attention block tables are empty
223
- # during vLLM memory profiling.
224
- cross_block_tables = []
225
- for seq_group_metadata in seq_group_metadata_list:
226
- for _ in range(len(seq_group_metadata.seq_data)):
227
- encoder_seq_lens.append(
228
- seq_group_metadata.encoder_seq_data.get_len())
229
- cross_block_table = seq_group_metadata.cross_block_table
230
- cross_block_tables.append([] if (
231
- cross_block_table is None) else cross_block_table)
232
-
233
- max_len_of_block_table = max(
234
- len(block_table) for block_table in cross_block_tables)
235
-
236
- cross_block_tables = make_tensor_with_pad(
237
- cross_block_tables,
238
- max_len=max_len_of_block_table,
239
- pad=0,
240
- dtype=torch.int32,
241
- device=self.device,
242
- )
243
-
244
- # Compute encoder sequence lengths & encoder
245
- # sequence starting offset tensors
246
- max_encoder_seq_len = max(encoder_seq_lens, default=0)
247
- encoder_seq_lens_tensor = self._list_to_int32_tensor(encoder_seq_lens)
248
- encoder_seq_start_loc = torch.zeros(encoder_seq_lens_tensor.shape[0] +
249
- 1,
250
- dtype=torch.int32,
251
- device=self.device)
252
- torch.cumsum(encoder_seq_lens_tensor,
253
- dim=0,
254
- dtype=encoder_seq_start_loc.dtype,
255
- out=encoder_seq_start_loc[1:])
256
-
257
- # Update attention metadata with encoder-oriented attributes
258
- attn_metadata = model_input.attn_metadata
259
- assert attn_metadata is not None
260
- (
261
- attn_metadata.num_encoder_tokens,
262
- attn_metadata.encoder_seq_lens,
263
- attn_metadata.encoder_seq_lens_tensor,
264
- attn_metadata.max_encoder_seq_len,
265
- attn_metadata.cross_slot_mapping,
266
- attn_metadata.cross_block_tables,
267
- ) = (
268
- sum(encoder_seq_lens),
269
- encoder_seq_lens,
270
- encoder_seq_lens_tensor,
271
- max_encoder_seq_len,
272
- cross_slot_mapping_tensor,
273
- cross_block_tables,
274
- )
275
-
276
- return (attn_metadata, encoder_input_tokens_tensor,
277
- encoder_input_positions_tensor)
278
-
279
- @torch.no_grad()
280
- def execute_model(
281
- self,
282
- model_input: EncoderDecoderModelInputForCPU,
283
- kv_caches: List[torch.Tensor],
284
- intermediate_tensors: Optional[IntermediateTensors] = None,
285
- num_steps: int = 1,
286
- ) -> Optional[List[SamplerOutput]]:
287
- if num_steps > 1:
288
- raise ValueError(
289
- "CPU worker does not support multi-step execution.")
290
-
291
- model_executable = self.model
292
- execute_model_kwargs = {
293
- "input_ids":
294
- model_input.input_tokens,
295
- "positions":
296
- model_input.input_positions,
297
- "encoder_input_ids":
298
- model_input.encoder_input_tokens,
299
- "encoder_positions":
300
- model_input.encoder_input_positions,
301
- **MultiModalKwargs.as_kwargs(
302
- model_input.multi_modal_kwargs or {},
303
- device=self.device,
304
- ),
305
- "intermediate_tensors":
306
- intermediate_tensors,
307
- }
308
-
309
- with set_forward_context(model_input.attn_metadata, self.vllm_config,
310
- model_input.virtual_engine):
311
- hidden_states = model_executable(**execute_model_kwargs)
312
-
313
- # Compute the logits.
314
- logits = self.model.compute_logits(hidden_states,
315
- model_input.sampling_metadata)
316
-
317
- # Only perform sampling in the driver worker.
318
- if not self.is_driver_worker:
319
- return []
320
-
321
- # Sample the next token.
322
- output = self.sampler(
323
- logits=logits,
324
- sampling_metadata=model_input.sampling_metadata,
325
- )
326
- return [output]
vllm/worker/cpu_model_runner.py
deleted
-671
@@ -1,671 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-import weakref
6
-from collections import defaultdict
7
-from dataclasses import dataclass
8
-from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Set, Type,
9
- TypeVar, Union)
10
-
11
-import torch
12
-from torch import nn
13
-
14
-from vllm.attention import AttentionMetadata, get_attn_backend
15
-from vllm.config import VllmConfig
16
-from vllm.forward_context import set_forward_context
17
-from vllm.logger import init_logger
18
-from vllm.lora.layers import LoRAMapping
19
-from vllm.lora.request import LoRARequest
20
-from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
21
-from vllm.model_executor import SamplingMetadata
22
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
23
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
24
-from vllm.model_executor.model_loader import get_model
25
-from vllm.model_executor.models import supports_lora, supports_multimodal
26
-from vllm.multimodal import (BatchedTensorInputs, MultiModalKwargs,
27
- MultiModalPlaceholderMap)
28
-from vllm.sequence import (IntermediateTensors, SequenceData,
29
- SequenceGroupMetadata)
30
-from vllm.worker.model_runner_base import (
31
- ModelRunnerBase, ModelRunnerInputBase, ModelRunnerInputBuilderBase,
32
- _add_attn_metadata_broadcastable_dict,
33
- _add_sampling_metadata_broadcastable_dict,
34
- _init_attn_metadata_from_tensor_dict,
35
- _init_sampling_metadata_from_tensor_dict)
36
-
37
-if TYPE_CHECKING:
38
- from vllm.attention.backends.abstract import AttentionBackend
39
-
40
-logger = init_logger(__name__)
41
-
42
-TModelInputForCPU = TypeVar('TModelInputForCPU', bound="ModelInputForCPU")
43
-_PAD_SLOT_ID = -1
44
-
45
-
46
-@dataclass(frozen=True)
47
-class ModelInputForCPU(ModelRunnerInputBase):
48
- """
49
- Base class contains metadata needed for the base model forward pass on CPU
50
- """
51
- input_tokens: Optional[torch.Tensor] = None
52
- input_positions: Optional[torch.Tensor] = None
53
- token_type_ids: Optional[torch.Tensor] = None
54
- attn_metadata: Optional["AttentionMetadata"] = None
55
- multi_modal_kwargs: Optional[BatchedTensorInputs] = None
56
- virtual_engine: Optional[int] = None
57
- seq_lens: Optional[List[int]] = None
58
- query_lens: Optional[List[int]] = None
59
- lora_mapping: Optional["LoRAMapping"] = None
60
- lora_requests: Optional[Set[LoRARequest]] = None
61
-
62
- def as_broadcastable_tensor_dict(
63
- self) -> Dict[str, Union[int, torch.Tensor]]:
64
- tensor_dict = {
65
- "input_tokens": self.input_tokens,
66
- "input_positions": self.input_positions,
67
- "token_type_ids": self.token_type_ids,
68
- "multi_modal_kwargs": self.multi_modal_kwargs,
69
- "lora_requests": self.lora_requests,
70
- "lora_mapping": self.lora_mapping,
71
- }
72
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
73
-
74
- return tensor_dict
75
-
76
- @classmethod
77
- def from_broadcasted_tensor_dict(
78
- cls: Type[TModelInputForCPU],
79
- tensor_dict: Dict[str, Any],
80
- attn_backend: Optional["AttentionBackend"] = None
81
- ) -> TModelInputForCPU:
82
- if attn_backend is not None:
83
- tensor_dict = _init_attn_metadata_from_tensor_dict(
84
- attn_backend, tensor_dict)
85
- return cls(**tensor_dict)
86
-
87
-
88
-@dataclass(frozen=True)
89
-class ModelInputForCPUWithSamplingMetadata(ModelInputForCPU):
90
- """
91
- Used by the ModelRunner.
92
- """
93
- sampling_metadata: Optional["SamplingMetadata"] = None
94
- is_prompt: Optional[bool] = None
95
-
96
- def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
97
- tensor_dict = {
98
- "input_tokens": self.input_tokens,
99
- "input_positions": self.input_positions,
100
- "token_type_ids": self.token_type_ids,
101
- "multi_modal_kwargs": self.multi_modal_kwargs,
102
- }
103
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
104
- _add_sampling_metadata_broadcastable_dict(tensor_dict,
105
- self.sampling_metadata)
106
- return tensor_dict
107
-
108
- @classmethod
109
- def from_broadcasted_tensor_dict(
110
- cls,
111
- tensor_dict: Dict[str, Any],
112
- attn_backend: Optional["AttentionBackend"] = None,
113
- ) -> "ModelInputForCPUWithSamplingMetadata":
114
- tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)
115
- if attn_backend is not None:
116
- tensor_dict = _init_attn_metadata_from_tensor_dict(
117
- attn_backend, tensor_dict)
118
- return cls(**tensor_dict)
119
-
120
-
121
-class ModelInputForCPUBuilder(ModelRunnerInputBuilderBase[ModelInputForCPU]):
122
-
123
- class ModelInputData:
124
-
125
- def __init__(self, use_mrope: bool):
126
- self.use_mrope = use_mrope
127
- self.input_tokens: List[int] = []
128
- self.input_positions: List[int] = []
129
- self.token_type_ids: Optional[List[int]] = []
130
- self.seq_lens: List[int] = []
131
- self.query_lens: List[int] = []
132
- self.prefill_block_tables: List[List[int]] = []
133
- self.decode_block_tables: List[List[int]] = []
134
- self.max_decode_seq_len: int = 0
135
- self.num_prefills: int = 0
136
- self.num_prefill_tokens: int = 0
137
- self.num_decode_tokens: int = 0
138
- self.slot_mapping: List[int] = []
139
- self.multi_modal_inputs_list: List[MultiModalKwargs] = []
140
- self.multi_modal_placeholder_maps: Dict[
141
- str, MultiModalPlaceholderMap] = defaultdict(
142
- MultiModalPlaceholderMap)
143
- self.input_mrope_positions: List[List[int]] = [[]
144
- for _ in range(3)]
145
-
146
- def __init__(self,
147
- runner: "CPUModelRunner",
148
- finished_requests_ids: Optional[List[str]] = None) -> None:
149
- super().__init__()
150
- self.runner = runner
151
- self.chunked_prefill = (runner.scheduler_config.chunked_prefill_enabled
152
- or runner.cache_config.enable_prefix_caching)
153
- self.model_input_cls = self.runner._model_input_cls
154
- self.attn_backend = self.runner.attn_backend
155
- self.sliding_window = self.runner.sliding_window
156
- self.block_size = self.runner.block_size
157
- self.device = self.runner.device
158
- self.enable_lora = self.runner.lora_config is not None
159
- if self.runner.attn_backend is not None:
160
- # spec decode (e.g. Medusa) does not have atten backend
161
- attn_backend = self.runner.attn_backend
162
- self.att_metadata_builder = attn_backend.get_builder_cls()(self)
163
-
164
- def prepare(self,
165
- finished_requests_ids: Optional[List[str]] = None) -> None:
166
- self.seq_group_metadata_list: List[SequenceGroupMetadata] = []
167
- self.input_data = ModelInputForCPUBuilder.ModelInputData(
168
- self.runner.model_config.uses_mrope)
169
- self.att_metadata_builder.prepare()
170
-
171
- def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
172
- self.seq_group_metadata_list.append(seq_group_metadata)
173
-
174
- def set_seq_group_list(
175
- self, seq_group_metadata_list: List[SequenceGroupMetadata]):
176
- self.seq_group_metadata_list = seq_group_metadata_list
177
-
178
- def build(self) -> ModelInputForCPU:
179
- self._build_input_data()
180
-
181
- input_data = self.input_data
182
- input_tokens = torch.tensor(input_data.input_tokens,
183
- dtype=torch.long,
184
- device="cpu")
185
- input_positions = torch.tensor(
186
- input_data.input_positions
187
- if not any(input_data.input_mrope_positions) else
188
- input_data.input_mrope_positions,
189
- dtype=torch.long,
190
- device="cpu")
191
- token_type_ids = torch.tensor(input_data.token_type_ids,
192
- dtype=torch.long,
193
- device="cpu") \
194
- if input_data.token_type_ids else None
195
-
196
- # For multi-modal models
197
- multi_modal_kwargs = None
198
- if len(input_data.multi_modal_inputs_list) != 0:
199
- multi_modal_kwargs = MultiModalKwargs.batch(
200
- input_data.multi_modal_inputs_list)
201
-
202
- attn_metadata = self.att_metadata_builder.build(
203
- input_data.seq_lens, input_data.query_lens, -1, -1)
204
-
205
- is_prompt = (self.seq_group_metadata_list[0].is_prompt
206
- if self.seq_group_metadata_list else None)
207
- # LoRA data.
208
- lora_requests = set()
209
- lora_mapping = None
210
- if self.enable_lora:
211
- lora_requests = set(seq.lora_request
212
- for seq in self.seq_group_metadata_list
213
- if seq.lora_request is not None)
214
-
215
- lora_mapping = self._prepare_lora_input(
216
- self.seq_group_metadata_list, is_prompt)
217
-
218
- return self.model_input_cls(input_tokens=input_tokens,
219
- input_positions=input_positions,
220
- token_type_ids=token_type_ids,
221
- seq_lens=input_data.seq_lens,
222
- query_lens=input_data.query_lens,
223
- attn_metadata=attn_metadata,
224
- multi_modal_kwargs=multi_modal_kwargs,
225
- lora_mapping=lora_mapping,
226
- lora_requests=lora_requests)
227
-
228
- def _build_input_data(self):
229
- for seq_group_metadata in self.seq_group_metadata_list:
230
- for seq_id, seq_data in seq_group_metadata.seq_data.items():
231
- if seq_group_metadata.is_prompt:
232
- self._compute_prompt_input_tokens(self.input_data,
233
- seq_group_metadata,
234
- seq_data, seq_id)
235
- if seq_group_metadata.multi_modal_data:
236
- self._compute_multi_modal_input(
237
- seq_group_metadata, seq_data)
238
- else:
239
- self._compute_decode_input_tokens(self.input_data,
240
- seq_group_metadata,
241
- seq_data, seq_id)
242
-
243
- def _compute_decode_input_tokens(self, data: ModelInputData,
244
- seq_group_metadata: SequenceGroupMetadata,
245
- seq_data: SequenceData, seq_id: int):
246
- """
247
- Compute decode input tokens, positions, block table and slot mapping.
248
- """
249
- block_size = self.runner.block_size
250
-
251
- block_table = seq_group_metadata.block_tables[seq_id]
252
- seq_len = seq_data.get_len()
253
- context_len = seq_data.get_num_computed_tokens()
254
-
255
- tokens = seq_data.get_last_token_id()
256
- token_positions = seq_len - 1
257
- block_number = block_table[token_positions // block_size]
258
- block_offset = token_positions % block_size
259
- slot = block_number * block_size + block_offset
260
-
261
- # For paged_attention kernel
262
- if self.runner.sliding_window:
263
- start_idx = max(0, seq_len - self.runner.sliding_window)
264
- start_block = start_idx // block_size
265
- start_idx = start_block * block_size
266
- seq_len = seq_len - start_idx
267
- block_table = block_table[start_block:]
268
-
269
- # For MRotaryEmbedding
270
- if seq_data.mrope_position_delta is not None:
271
- next_pos = MRotaryEmbedding.get_next_input_positions(
272
- seq_data.mrope_position_delta,
273
- context_len,
274
- seq_len,
275
- )
276
- for idx in range(3):
277
- data.input_mrope_positions[idx].extend( # type: ignore
278
- next_pos[idx])
279
- else:
280
- data.input_positions.append(token_positions) # type: ignore
281
-
282
- # Update fields
283
- data.input_tokens.append(tokens)
284
- data.max_decode_seq_len = max(data.max_decode_seq_len, seq_len)
285
- data.num_decode_tokens += 1
286
- data.slot_mapping.append(slot)
287
- data.decode_block_tables.append(block_table)
288
- data.query_lens.append(1)
289
- data.seq_lens.append(seq_len)
290
-
291
- def _compute_prompt_input_tokens(self, data: ModelInputData,
292
- seq_group_metadata: SequenceGroupMetadata,
293
- seq_data: SequenceData, seq_id: int):
294
- """
295
- Compute prompt input tokens, positions, block table and slot mapping.
296
- """
297
- token_chunk_size = seq_group_metadata.token_chunk_size
298
- block_size = self.runner.block_size
299
-
300
- block_table = seq_group_metadata.block_tables[seq_id]
301
- seq_len = seq_data.get_len()
302
- context_len = seq_data.get_num_computed_tokens()
303
- seq_len = min(seq_len, context_len + token_chunk_size)
304
-
305
- # For prefix caching
306
- prefix_cache_block_num = len(seq_group_metadata.computed_block_nums)
307
- if prefix_cache_block_num > 0:
308
- prefix_cache_len = (prefix_cache_block_num *
309
- self.runner.block_size)
310
- if prefix_cache_len <= context_len:
311
- # We already passed the cache hit region,
312
- # so do normal computation.
313
- pass
314
- elif context_len < prefix_cache_len < seq_len:
315
- # Partial hit. Compute the missing part.
316
- context_len = prefix_cache_len
317
- token_chunk_size = seq_len - context_len
318
- elif seq_len <= prefix_cache_len:
319
- # Full hit. Only compute the last token to avoid
320
- # erroneous behavior. FIXME: Ideally we should directly
321
- # mark all tokens as computed in the scheduler and do not
322
- # schedule this sequence, so this case should not happen.
323
- context_len = seq_len - 1
324
- token_chunk_size = 1
325
-
326
- tokens = seq_data.get_token_ids()
327
- tokens = tokens[context_len:seq_len]
328
- token_positions = range(context_len, seq_len)
329
- token_types = seq_group_metadata.token_type_ids
330
-
331
- # For encoder-only models, the block_table is None,
332
- # and there is no need to initialize the slot_mapping.
333
- if block_table is not None:
334
- slot_mapping = [_PAD_SLOT_ID] * len(token_positions)
335
- for i, pos in enumerate(token_positions):
336
- block_number = block_table[pos // block_size]
337
- block_offset = pos % block_size
338
- slot = block_number * block_size + block_offset
339
- slot_mapping[i] = slot
340
- data.slot_mapping.extend(slot_mapping)
341
-
342
- # The MROPE positions are prepared in _compute_multi_modal_input
343
- data.input_positions.extend(token_positions)
344
-
345
- if data.token_type_ids is not None:
346
- data.token_type_ids.extend(token_types if token_types else [])
347
-
348
- # Update fields
349
- data.input_tokens.extend(tokens)
350
- data.num_prefills += 1
351
- data.num_prefill_tokens += len(tokens)
352
- data.query_lens.append(len(tokens))
353
- data.prefill_block_tables.append(block_table)
354
- data.seq_lens.append(seq_len)
355
-
356
- def _compute_multi_modal_input(self,
357
- seq_group_metadata: SequenceGroupMetadata,
358
- seq_data: SequenceData):
359
- computed_len = seq_data.get_num_computed_tokens()
360
- seq_len = self.input_data.seq_lens[-1]
361
-
362
- # NOTE: mm_kwargs only includes the subset of multi-modal items that
363
- # intersect with the current prefill positions.
364
- mm_kwargs, placeholder_maps = MultiModalPlaceholderMap.from_seq_group(
365
- seq_group_metadata, range(computed_len, seq_len))
366
-
367
- if not mm_kwargs:
368
- return
369
-
370
- # special processing for mrope position deltas.
371
- if self.runner.model_config.uses_mrope:
372
- assert not self.chunked_prefill, \
373
- "MROPE on CPU does not support chunked-prefill."
374
-
375
- image_grid_thw = mm_kwargs.get("image_grid_thw", None)
376
- video_grid_thw = mm_kwargs.get("video_grid_thw", None)
377
- audio_feature_lengths = mm_kwargs.get("audio_feature_lengths",
378
- None)
379
- assert (
380
- image_grid_thw is not None or video_grid_thw is not None
381
- or audio_feature_lengths is not None), (
382
- "mrope embedding type requires multi-modal input mapper "
383
- "returns 'image_grid_thw' or 'video_grid_thw' or "
384
- "'audio_feature_lengths'.")
385
-
386
- second_per_grid_ts = mm_kwargs.get("second_per_grid_ts", None)
387
- use_audio_in_video = mm_kwargs.get("use_audio_in_video", False)
388
- hf_config = self.runner.model_config.hf_config
389
- token_ids = seq_data.get_token_ids()
390
-
391
- mrope_positions, mrope_position_delta = \
392
- MRotaryEmbedding.get_input_positions(
393
- token_ids,
394
- hf_config=hf_config,
395
- image_grid_thw=image_grid_thw,
396
- video_grid_thw=video_grid_thw,
397
- second_per_grid_ts=second_per_grid_ts,
398
- context_len=computed_len,
399
- audio_feature_lengths=audio_feature_lengths,
400
- use_audio_in_video=use_audio_in_video,
401
- )
402
- seq_data.mrope_position_delta = mrope_position_delta
403
-
404
- for i in range(3):
405
- self.input_data.input_mrope_positions[ # type: ignore
406
- i].extend(mrope_positions[i])
407
-
408
- self.input_data.multi_modal_inputs_list.append(mm_kwargs)
409
- for modality, placeholder_map in placeholder_maps.items():
410
- self.input_data.multi_modal_placeholder_maps[modality].extend(
411
- placeholder_map)
412
-
413
- def _prepare_lora_input(
414
- self, seq_group_metadata_list: List[SequenceGroupMetadata],
415
- is_prefill: bool) -> LoRAMapping:
416
- index_mapping = []
417
- prompt_mapping = []
418
- for seq in seq_group_metadata_list:
419
- lora_id = seq.lora_int_id
420
- query_len = seq.token_chunk_size
421
-
422
- index_mapping += [lora_id] * query_len
423
- prompt_mapping += [lora_id] * (
424
- query_len if seq.sampling_params
425
- and seq.sampling_params.prompt_logprobs is not None else 1)
426
-
427
- return LoRAMapping(index_mapping=tuple(index_mapping),
428
- prompt_mapping=tuple(prompt_mapping),
429
- is_prefill=is_prefill)
430
-
431
-
432
-class CPUModelRunnerBase(ModelRunnerBase[TModelInputForCPU]):
433
- """
434
- Helper class for shared methods between CPU model runners.
435
- """
436
- _model_input_cls: Type[TModelInputForCPU]
437
- _builder_cls: Type[ModelInputForCPUBuilder]
438
- builder: ModelInputForCPUBuilder
439
-
440
- def __init__(
441
- self,
442
- vllm_config: VllmConfig,
443
- kv_cache_dtype: Optional[str] = "auto",
444
- is_driver_worker: bool = False,
445
- return_hidden_states: bool = False,
446
- *args,
447
- **kwargs,
448
- ):
449
- ModelRunnerBase.__init__(self, vllm_config)
450
- model_config = self.model_config
451
- cache_config = self.cache_config
452
-
453
- self.is_driver_worker = is_driver_worker
454
- self.return_hidden_states = return_hidden_states
455
-
456
- self.device = self.device_config.device
457
- self.pin_memory = False
458
-
459
- self.kv_cache_dtype = kv_cache_dtype
460
- self.sliding_window = model_config.get_sliding_window()
461
- self.block_size = cache_config.block_size
462
- num_attn_heads = self.model_config.get_num_attention_heads(
463
- self.parallel_config)
464
- needs_attn_backend = (num_attn_heads != 0
465
- or self.model_config.is_attention_free)
466
- self.attn_backend = get_attn_backend(
467
- self.model_config.get_head_size(),
468
- self.model_config.dtype,
469
- self.kv_cache_dtype,
470
- self.block_size,
471
- self.model_config.is_attention_free,
472
- use_mla=self.model_config.use_mla,
473
- ) if needs_attn_backend else None
474
-
475
- # Lazy initialization.
476
- self.model: nn.Module # Set after init_Model
477
- # Set after load_model.
478
- self.lora_manager: Optional[LRUCacheWorkerLoRAManager] = None
479
- self.sampler = get_sampler()
480
-
481
- if hasattr(self, "_builder_cls"):
482
- # multi-step model runner does not have `_builder_cls`
483
- self.builder = self._builder_cls(weakref.proxy(self))
484
-
485
- def load_model(self) -> None:
486
- self.model = get_model(vllm_config=self.vllm_config)
487
-
488
- if self.lora_config:
489
- assert supports_lora(
490
- self.model
491
- ), f"{self.model.__class__.__name__} does not support LoRA yet."
492
-
493
- if supports_multimodal(self.model):
494
- logger.warning("Regarding multimodal models, vLLM currently "
495
- "only supports adding LoRA to language model.")
496
-
497
- # Use get_text_config() in case of multimodal models
498
- text_config = self.model_config.hf_config.get_text_config()
499
-
500
- self.lora_manager = LRUCacheWorkerLoRAManager(
501
- self.scheduler_config.max_num_seqs,
502
- self.scheduler_config.max_num_batched_tokens,
503
- self.vocab_size,
504
- self.lora_config,
505
- self.device,
506
- self.model.embedding_modules,
507
- self.model.embedding_padding_modules,
508
- max_position_embeddings=text_config.max_position_embeddings,
509
- )
510
- self.model = self.lora_manager.create_lora_manager(self.model)
511
-
512
- def get_model(self) -> nn.Module:
513
- return self.model
514
-
515
- def _prepare_model_input_tensors(
516
- self,
517
- seq_group_metadata_list: List[SequenceGroupMetadata],
518
- finished_requests_ids: Optional[List[str]] = None
519
- ) -> TModelInputForCPU:
520
- """Helper method to prepare the model input based on a given sequence
521
- group. Prepares metadata needed for the base model forward pass but not
522
- metadata for possible additional steps, e.g., sampling.
523
-
524
- """
525
- self.builder.prepare(finished_requests_ids)
526
- self.builder.set_seq_group_list(seq_group_metadata_list)
527
-
528
- return self.builder.build() # type: ignore
529
-
530
- @property
531
- def vocab_size(self) -> int:
532
- return self.model_config.get_vocab_size()
533
-
534
- def remove_all_loras(self):
535
- if not self.lora_manager:
536
- raise RuntimeError("LoRA is not enabled.")
537
- self.lora_manager.remove_all_adapters()
538
-
539
- def set_active_loras(self, lora_requests: Set[LoRARequest],
540
- lora_mapping: LoRAMapping) -> None:
541
- if not self.lora_manager:
542
- raise RuntimeError("LoRA is not enabled.")
543
- self.lora_manager.set_active_adapters(lora_requests, lora_mapping)
544
-
545
- def add_lora(self, lora_request: LoRARequest) -> bool:
546
- if not self.lora_manager:
547
- raise RuntimeError("LoRA is not enabled.")
548
- return self.lora_manager.add_adapter(lora_request)
549
-
550
- def remove_lora(self, lora_id: int) -> bool:
551
- if not self.lora_manager:
552
- raise RuntimeError("LoRA is not enabled.")
553
- return self.lora_manager.remove_adapter(lora_id)
554
-
555
- def pin_lora(self, lora_id: int) -> bool:
556
- if not self.lora_manager:
557
- raise RuntimeError("LoRA is not enabled.")
558
- return self.lora_manager.pin_adapter(lora_id)
559
-
560
- def list_loras(self) -> Set[int]:
561
- if not self.lora_manager:
562
- raise RuntimeError("LoRA is not enabled.")
563
- return self.lora_manager.list_adapters()
564
-
565
-
566
-class CPUModelRunner(CPUModelRunnerBase[ModelInputForCPUWithSamplingMetadata]):
567
- _model_input_cls: Type[ModelInputForCPUWithSamplingMetadata] = (
568
- ModelInputForCPUWithSamplingMetadata)
569
- _builder_cls: Type[ModelInputForCPUBuilder] = ModelInputForCPUBuilder
570
-
571
- def make_model_input_from_broadcasted_tensor_dict(
572
- self,
573
- tensor_dict: Dict[str, Any],
574
- ) -> ModelInputForCPUWithSamplingMetadata:
575
- return ModelInputForCPUWithSamplingMetadata.from_broadcasted_tensor_dict( # noqa: E501
576
- tensor_dict,
577
- attn_backend=self.attn_backend,
578
- )
579
-
580
- def prepare_model_input(
581
- self,
582
- seq_group_metadata_list: List[SequenceGroupMetadata],
583
- virtual_engine: int = 0,
584
- finished_requests_ids: Optional[List[str]] = None
585
- ) -> ModelInputForCPUWithSamplingMetadata:
586
- """Prepare the model input based on a given sequence group, including
587
- metadata for the sampling step.
588
-
589
- """
590
- model_input = self._prepare_model_input_tensors(
591
- seq_group_metadata_list, finished_requests_ids)
592
- # Sampling metadata is only required for the final pp group
593
- generators = self.get_generators(finished_requests_ids)
594
- sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list,
595
- model_input.seq_lens,
596
- model_input.query_lens,
597
- self.device,
598
- pin_memory=False,
599
- generators=generators)
600
-
601
- is_prompt = (seq_group_metadata_list[0].is_prompt
602
- if seq_group_metadata_list else None)
603
- return dataclasses.replace(model_input,
604
- sampling_metadata=sampling_metadata,
605
- virtual_engine=virtual_engine,
606
- is_prompt=is_prompt)
607
-
608
- @torch.no_grad()
609
- def execute_model(
610
- self,
611
- model_input: ModelInputForCPUWithSamplingMetadata,
612
- kv_caches: List[torch.Tensor],
613
- intermediate_tensors: Optional[IntermediateTensors] = None,
614
- num_steps: int = 1,
615
- previous_hidden_states: Optional[torch.Tensor] = None,
616
- ) -> Optional[List[SamplerOutput]]:
617
- if num_steps > 1:
618
- raise ValueError(
619
- "CPU worker does not support multi-step execution.")
620
-
621
- if self.lora_config:
622
- assert model_input.lora_requests is not None
623
- assert model_input.lora_mapping is not None
624
- self.set_active_loras(model_input.lora_requests,
625
- model_input.lora_mapping)
626
-
627
- model_executable = self.model
628
-
629
- multimodal_kwargs = {}
630
- if model_input.multi_modal_kwargs is not None:
631
- multimodal_kwargs = MultiModalKwargs.as_kwargs(
632
- model_input.multi_modal_kwargs,
633
- device=self.device,
634
- )
635
- execute_model_kwargs = {}
636
- if previous_hidden_states is not None:
637
- execute_model_kwargs.update(
638
- {"previous_hidden_states": previous_hidden_states})
639
-
640
- with set_forward_context(model_input.attn_metadata, self.vllm_config,
641
- model_input.virtual_engine):
642
- hidden_states = model_executable(
643
- input_ids=model_input.input_tokens,
644
- positions=model_input.input_positions,
645
- intermediate_tensors=intermediate_tensors,
646
- **execute_model_kwargs,
647
- **multimodal_kwargs,
648
- )
649
-
650
- # Compute the logits.
651
- logits = self.model.compute_logits(hidden_states,
652
- model_input.sampling_metadata)
653
-
654
- # Only perform sampling in the driver worker.
655
- if not self.is_driver_worker:
656
- return []
657
-
658
- # Sample the next token.
659
- output = self.sampler(
660
- logits=logits,
661
- sampling_metadata=model_input.sampling_metadata,
662
- )
663
- if self.return_hidden_states:
664
- # we only need to pass hidden states of most recent token
665
- if model_input.is_prompt:
666
- output.prefill_hidden_states = hidden_states
667
- output.hidden_states = hidden_states
668
- return [output]
669
-
670
- def generate_proposals(self, *args, **kwargs):
671
- return self.model.generate_proposals(*args, **kwargs)
vllm/worker/cpu_pooling_model_runner.py
deleted
-125
@@ -1,125 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
6
-
7
-import torch
8
-
9
-from vllm.forward_context import set_forward_context
10
-from vllm.model_executor.pooling_metadata import PoolingMetadata
11
-from vllm.multimodal import MultiModalKwargs
12
-from vllm.pooling_params import PoolingParams
13
-from vllm.sequence import (IntermediateTensors, PoolerOutput, SequenceData,
14
- SequenceGroupMetadata)
15
-from vllm.worker.cpu_model_runner import (CPUModelRunnerBase, ModelInputForCPU,
16
- ModelInputForCPUBuilder)
17
-
18
-
19
-@dataclasses.dataclass(frozen=True)
20
-class ModelInputForCPUWithPoolingMetadata(ModelInputForCPU):
21
- """
22
- Used by the CPUPoolingModelRunner.
23
- """
24
- pooling_metadata: Optional["PoolingMetadata"] = None
25
-
26
-
27
-class CPUPoolingModelRunner(
28
- CPUModelRunnerBase[ModelInputForCPUWithPoolingMetadata]):
29
- _model_input_cls: Type[ModelInputForCPUWithPoolingMetadata] = (
30
- ModelInputForCPUWithPoolingMetadata)
31
- _builder_cls: Type[ModelInputForCPUBuilder] = ModelInputForCPUBuilder
32
-
33
- @torch.inference_mode()
34
- def execute_model(
35
- self,
36
- model_input: ModelInputForCPUWithPoolingMetadata,
37
- kv_caches: List[torch.Tensor],
38
- intermediate_tensors: Optional[IntermediateTensors] = None,
39
- num_steps: int = 1,
40
- ) -> Optional[Union[List[PoolerOutput], IntermediateTensors]]:
41
- if num_steps > 1:
42
- raise ValueError(
43
- "CPU worker does not support multi-step execution.")
44
-
45
- model_executable = self.model
46
- cross_enc_kwargs = {}
47
- if model_input.token_type_ids is not None:
48
- cross_enc_kwargs["token_type_ids"] = model_input.token_type_ids
49
- execute_model_kwargs = {
50
- "input_ids":
51
- model_input.input_tokens,
52
- "positions":
53
- model_input.input_positions,
54
- **MultiModalKwargs.as_kwargs(
55
- model_input.multi_modal_kwargs or {},
56
- device=self.device,
57
- ),
58
- **cross_enc_kwargs,
59
- "intermediate_tensors":
60
- intermediate_tensors,
61
- }
62
-
63
- with set_forward_context(model_input.attn_metadata, self.vllm_config,
64
- model_input.virtual_engine):
65
- hidden_states = model_executable(**execute_model_kwargs)
66
-
67
- # Only perform pooling in the driver worker.
68
- if not self.is_driver_worker:
69
- return []
70
-
71
- return [
72
- self.model.pooler(hidden_states=hidden_states,
73
- pooling_metadata=model_input.pooling_metadata)
74
- ]
75
-
76
- def make_model_input_from_broadcasted_tensor_dict(
77
- self,
78
- tensor_dict: Dict[str,
79
- Any]) -> ModelInputForCPUWithPoolingMetadata:
80
- return ModelInputForCPUWithPoolingMetadata.from_broadcasted_tensor_dict(
81
- tensor_dict,
82
- attn_backend=self.attn_backend,
83
- )
84
-
85
- def prepare_model_input(
86
- self,
87
- seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
88
- virtual_engine: int = 0,
89
- finished_requests_ids: Optional[List[str]] = None
90
- ) -> ModelInputForCPUWithPoolingMetadata:
91
- assert seq_group_metadata_list is not None
92
- model_input = self._prepare_model_input_tensors(
93
- seq_group_metadata_list, finished_requests_ids)
94
- # Prepare PoolingMetadata.
95
- assert model_input.seq_lens is not None
96
- pooling_metadata = self._prepare_pooling(seq_group_metadata_list,
97
- model_input.seq_lens)
98
-
99
- return dataclasses.replace(model_input,
100
- virtual_engine=virtual_engine,
101
- pooling_metadata=pooling_metadata)
102
-
103
- def _prepare_pooling(
104
- self,
105
- seq_group_metadata_list: List[SequenceGroupMetadata],
106
- prompt_lens: List[int],
107
- ) -> PoolingMetadata:
108
- """Prepare PoolingMetadata for the sequence group metadata list."""
109
- seq_groups: List[Tuple[List[int], PoolingParams]] = []
110
- for i, seq_group_metadata in enumerate(seq_group_metadata_list):
111
- seq_ids = list(seq_group_metadata.seq_data.keys())
112
- pooling_params = seq_group_metadata.pooling_params
113
- seq_groups.append((seq_ids, pooling_params))
114
-
115
- seq_data: Dict[int, SequenceData] = {}
116
- for seq_group_metadata in seq_group_metadata_list:
117
- seq_data.update(seq_group_metadata.seq_data)
118
-
119
- pooling_metadata = PoolingMetadata(
120
- seq_groups=seq_groups,
121
- seq_data=seq_data,
122
- prompt_lens=prompt_lens,
123
- )
124
-
125
- return pooling_metadata
vllm/worker/cpu_worker.py
deleted
-452
@@ -1,452 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""A CPU worker class."""
4
-import os
5
-from importlib import util
6
-from typing import List, Optional, Set, Tuple, Type
7
-
8
-import torch
9
-import torch.distributed
10
-
11
-import vllm.envs as envs
12
-from vllm.attention import get_attn_backend
13
-from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
14
- ParallelConfig, VllmConfig)
15
-from vllm.distributed import (ensure_model_parallel_initialized,
16
- init_distributed_environment)
17
-from vllm.logger import init_logger
18
-from vllm.lora.request import LoRARequest
19
-from vllm.model_executor import set_random_seed
20
-from vllm.sequence import ExecuteModelRequest
21
-from vllm.utils import bind_kv_cache
22
-from vllm.worker.cpu_enc_dec_model_runner import CPUEncoderDecoderModelRunner
23
-from vllm.worker.cpu_model_runner import CPUModelRunner, CPUModelRunnerBase
24
-from vllm.worker.cpu_pooling_model_runner import CPUPoolingModelRunner
25
-from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,
26
- WorkerInput)
27
-
28
-logger = init_logger(__name__)
29
-
30
-
31
-class CPUCacheEngine:
32
- """Manages the KV cache for CPU backend.
33
-
34
- This class is responsible for initializing and managing CPU KV
35
- caches. It also provides methods for performing KV cache operations, such
36
- as copying.
37
- """
38
-
39
- def __init__(self, cache_config: CacheConfig, model_config: ModelConfig,
40
- parallel_config: ParallelConfig,
41
- device_config: DeviceConfig) -> None:
42
- assert device_config.device_type == "cpu"
43
- self.cache_config = cache_config
44
- self.model_config = model_config
45
- self.parallel_config = parallel_config
46
-
47
- self.head_size = model_config.get_head_size()
48
- self.num_layers = model_config.get_num_layers(parallel_config)
49
- self.num_heads = model_config.get_num_kv_heads(parallel_config)
50
-
51
- self.block_size = cache_config.block_size
52
- # Note: In CacheConfig, num_gpu_blocks actual is num_cpu_blocks
53
- # for CPU backend, because we want to reuse KV cache management
54
- # in the scheduler.
55
- self.num_cpu_blocks = cache_config.num_gpu_blocks
56
-
57
- self.dtype = CPUCacheEngine.get_kv_cache_dtype(cache_config,
58
- model_config)
59
-
60
- # Get attention backend.
61
- self.attn_backend = get_attn_backend(
62
- self.model_config.get_head_size(),
63
- self.model_config.dtype,
64
- cache_config.cache_dtype,
65
- self.block_size,
66
- self.model_config.is_attention_free,
67
- use_mla=self.model_config.use_mla,
68
- )
69
-
70
- # Initialize the cache.
71
- self.cpu_cache = self._allocate_kv_cache(self.num_cpu_blocks)
72
-
73
- def _allocate_kv_cache(
74
- self,
75
- num_blocks: int,
76
- ) -> List[torch.Tensor]:
77
- """Allocates KV cache on CPU."""
78
- kv_cache_shape = self.attn_backend.get_kv_cache_shape(
79
- num_blocks, self.block_size, self.num_heads, self.head_size)
80
- kv_cache: List[torch.Tensor] = []
81
- for _ in range(self.num_layers):
82
- kv_cache.append(
83
- torch.empty(kv_cache_shape, dtype=self.dtype, device="cpu"))
84
- return kv_cache
85
-
86
- def swap_in(self, src_to_dst: torch.Tensor) -> None:
87
- raise NotImplementedError("Swap is not supported in CPUCacheEngine.")
88
-
89
- def swap_out(self, src_to_dst: torch.Tensor) -> None:
90
- raise NotImplementedError("Swap is not supported in CPUCacheEngine.")
91
-
92
- def copy(self, src_to_dsts: torch.Tensor) -> None:
93
- self.attn_backend.copy_blocks(self.cpu_cache, src_to_dsts)
94
-
95
- @staticmethod
96
- def get_kv_cache_dtype(cache_config: CacheConfig,
97
- model_config: ModelConfig):
98
- if cache_config.cache_dtype == "auto":
99
- return model_config.dtype
100
- elif cache_config.cache_dtype in ["fp8", "fp8_e5m2"]:
101
- return torch.float8_e5m2
102
- else:
103
- raise NotImplementedError(f"Unsupported KV cache type "
104
- f"{cache_config.cache_dtype}.")
105
-
106
- @staticmethod
107
- def get_cache_block_size(
108
- cache_config: CacheConfig,
109
- model_config: ModelConfig,
110
- parallel_config: ParallelConfig,
111
- ) -> int:
112
- head_size = model_config.get_head_size()
113
- num_heads = model_config.get_num_kv_heads(parallel_config)
114
- num_layers = model_config.get_num_layers(parallel_config)
115
-
116
- key_cache_block = cache_config.block_size * num_heads * head_size
117
- value_cache_block = key_cache_block if not model_config.use_mla else 0
118
- total = num_layers * (key_cache_block + value_cache_block)
119
- dtype = CPUCacheEngine.get_kv_cache_dtype(cache_config, model_config)
120
- dtype_size = torch.tensor([], dtype=dtype).element_size()
121
- return dtype_size * total
122
-
123
-
124
-class CPUWorker(LocalOrDistributedWorkerBase):
125
- """A worker class that executes (a partition of) the model on a CPU socket.
126
-
127
- Each worker is associated with a single CPU socket. The worker is
128
- responsible for maintaining the KV cache and executing the model on the
129
- CPU. In case of distributed inference, each worker is assigned a partition
130
- of the model.
131
- """
132
-
133
- def __init__(
134
- self,
135
- vllm_config: VllmConfig,
136
- local_rank: int,
137
- rank: int,
138
- distributed_init_method: str,
139
- kv_cache_dtype: Optional[str] = "auto",
140
- is_driver_worker: bool = False,
141
- model_runner_cls: Optional[Type[CPUModelRunner]] = None,
142
- ) -> None:
143
- WorkerBase.__init__(self, vllm_config=vllm_config)
144
-
145
- self.local_rank = local_rank
146
- self.rank = rank
147
- vllm_config.parallel_config.rank = rank
148
-
149
- self.distributed_init_method = distributed_init_method
150
-
151
- self.is_driver_worker = is_driver_worker
152
- if self.is_driver_worker:
153
- assert self.rank == 0, "The driver worker must have rank 0."
154
-
155
- if self.model_config.trust_remote_code:
156
- # note: lazy import to avoid importing torch before initializing
157
- from vllm.utils import init_cached_hf_modules
158
- init_cached_hf_modules()
159
-
160
- # Setup OpenMP threads affinity.
161
- omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
162
- self.local_omp_cpuid = "all"
163
- if omp_cpuids == "auto":
164
- self.local_omp_cpuid = self.get_cpus_id_binding_based_on_numa_nodes(
165
- )
166
- else:
167
- self.local_omp_cpuid = omp_cpuids.split("|")[rank]
168
-
169
- # Return hidden states from target model if the draft model is an
170
- # mlp_speculator
171
- speculative_config = self.speculative_config
172
- model_config = self.model_config
173
- speculative_args = {} if speculative_config is None \
174
- or (speculative_config.draft_model_config.model ==
175
- model_config.model) \
176
- or (speculative_config.draft_model_config.hf_config.model_type
177
- not in ["medusa", "mlp_speculator", "eagle"]) \
178
- else {"return_hidden_states": True}
179
- ModelRunnerClass: Type[CPUModelRunnerBase] = CPUModelRunner
180
- if self.model_config.runner_type == "pooling":
181
- ModelRunnerClass = CPUPoolingModelRunner
182
- elif self.model_config.is_encoder_decoder:
183
- ModelRunnerClass = CPUEncoderDecoderModelRunner
184
- self.model_runner: CPUModelRunnerBase = ModelRunnerClass(
185
- vllm_config=vllm_config,
186
- kv_cache_dtype=kv_cache_dtype,
187
- is_driver_worker=is_driver_worker,
188
- **speculative_args,
189
- )
190
- if model_runner_cls is not None:
191
- self.model_runner = model_runner_cls(self.model_runner)
192
- # Uninitialized cache engine. Will be initialized by
193
- # initialize_cache.
194
- self.cache_engine: List[CPUCacheEngine]
195
- # Initialize cpu_cache as pooling models don't initialize kv_caches
196
- self.cpu_cache: Optional[List[List[torch.Tensor]]] = None
197
-
198
- # Torch profiler. Enabled and configured through env vars:
199
- # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
200
- if envs.VLLM_TORCH_PROFILER_DIR:
201
- torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
202
- logger.info("Profiling enabled. Traces will be saved to: %s",
203
- torch_profiler_trace_dir)
204
- self.profiler = torch.profiler.profile(
205
- activities=[
206
- torch.profiler.ProfilerActivity.CPU,
207
- ],
208
- with_stack=True,
209
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
210
- torch_profiler_trace_dir, use_gzip=True))
211
- else:
212
- self.profiler = None
213
-
214
- def start_profile(self):
215
- if self.profiler is None:
216
- raise RuntimeError("Profiler is not enabled.")
217
- self.profiler.start()
218
-
219
- def stop_profile(self):
220
- if self.profiler is None:
221
- raise RuntimeError("Profiler is not enabled.")
222
- self.profiler.stop()
223
-
224
- def init_device(self) -> None:
225
- if self.local_omp_cpuid != "all":
226
- ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
227
- if ret:
228
- logger.info(ret)
229
-
230
- # Note: unique identifier for creating allreduce shared memory
231
- os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
232
- ":")[-1]
233
- self.device = torch.device("cpu")
234
- self.init_distributed_environment()
235
- # Set random seed.
236
- set_random_seed(self.model_config.seed)
237
-
238
- def load_model(self):
239
- self.model_runner.load_model()
240
-
241
- def determine_num_available_blocks(self) -> Tuple[int, int]:
242
- """Determine the number of blocks available for the KV cache.
243
-
244
- This determines how many KV blocks can fit into the configured CPU
245
- KV cache space.
246
-
247
- Note that since vLLM assumes a block resides on GPU if it can be
248
- modified, we return num_gpu_blocks=num_cpu_blocks and num_cpu_blocks=0.
249
- This allows us to reuse the scheduler of vLLM without generalizing it
250
- to different devices.
251
- """
252
- # For CPU device, the block number will be calculated based on the
253
- # cpu_kvcache_space.
254
- cache_block_size = self.get_cache_block_size_bytes()
255
- num_cpu_blocks = int(self.cache_config.cpu_kvcache_space_bytes //
256
- cache_block_size)
257
- num_cpu_blocks = max(num_cpu_blocks, 0)
258
-
259
- # Note: To reuse the cache management procedure,
260
- # use cpu cache as 'gpu cache'.
261
- num_gpu_blocks = num_cpu_blocks
262
- num_cpu_blocks = 0
263
- return num_gpu_blocks, num_cpu_blocks
264
-
265
- def initialize_cache(self, num_gpu_blocks: int,
266
- num_cpu_blocks: int) -> None:
267
- """Initialize the KV cache. Currently, swappable CPU memory is not
268
- supported.
269
-
270
- Since this worker does not support GPUs, we use the num_gpu_blocks to
271
- determine how many non-swappable CPU blocks to allocate.
272
- """
273
- assert (num_cpu_blocks == 0
274
- ), f"{type(self)} does not support swappable cache"
275
-
276
- # Note: To reuse the cache management procedure,
277
- # use cpu cache as 'gpu cache'.
278
- num_cpu_blocks = num_gpu_blocks
279
-
280
- self._validate_num_cpu_blocks(num_cpu_blocks)
281
- self.cache_config.num_gpu_blocks = num_cpu_blocks
282
- self.cache_config.num_cpu_blocks = 0
283
-
284
- # Initialize the cache.
285
- self._init_cache_engine()
286
-
287
- def add_lora(self, lora_request: LoRARequest) -> bool:
288
- return self.model_runner.add_lora(lora_request)
289
-
290
- def remove_lora(self, lora_id: int) -> bool:
291
- return self.model_runner.remove_lora(lora_id)
292
-
293
- def pin_lora(self, lora_id: int) -> bool:
294
- return self.model_runner.pin_lora(lora_id)
295
-
296
- def list_loras(self) -> Set[int]:
297
- return self.model_runner.list_loras()
298
-
299
- def _validate_num_cpu_blocks(self, num_cpu_blocks: int) -> None:
300
- """Raise errors if the num_cpu_blocks is invalid.
301
- """
302
- if num_cpu_blocks <= 0:
303
- raise ValueError("No available memory for the cache blocks. "
304
- "Try increasing `VLLM_CPU_KVCACHE_SPACE` when "
305
- "initializing the engine.")
306
-
307
- max_seq_len = self.cache_config.block_size * num_cpu_blocks
308
- if self.model_config.max_model_len > max_seq_len:
309
- raise ValueError(
310
- f"The model's max seq len ({self.model_config.max_model_len}) "
311
- "is larger than the maximum number of tokens that can be "
312
- f"stored in KV cache ({max_seq_len}). Try increasing "
313
- "`VLLM_CPU_KVCACHE_SPACE` or decreasing `max_model_len` when "
314
- "initializing the engine.")
315
-
316
- def _init_cache_engine(self) -> None:
317
- self.cache_engine = [
318
- CPUCacheEngine(self.cache_config, self.model_config,
319
- self.parallel_config, self.device_config)
320
- for _ in range(self.parallel_config.pipeline_parallel_size)
321
- ]
322
- self.cpu_cache = [
323
- self.cache_engine[ve].cpu_cache
324
- for ve in range(self.parallel_config.pipeline_parallel_size)
325
- ]
326
- bind_kv_cache(self.compilation_config.static_forward_context,
327
- self.cpu_cache)
328
- self.model_runner.block_size = self.cache_engine[0].block_size
329
-
330
- assert all(
331
- self.cpu_cache[ve] is not None
332
- for ve in range(self.parallel_config.pipeline_parallel_size))
333
-
334
- # Populate the cache to warmup the memory
335
- for ve in range(self.parallel_config.pipeline_parallel_size):
336
- for layer_cache in self.cpu_cache[ve]:
337
- layer_cache.fill_(0)
338
-
339
- @property
340
- def do_metadata_broadcast(self) -> bool:
341
- return self.parallel_config.tensor_parallel_size > 1
342
-
343
- @property
344
- def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:
345
- return self.cpu_cache
346
-
347
- @property
348
- def vocab_size(self) -> int:
349
- return self.model_runner.vocab_size
350
-
351
- @property
352
- def max_model_len(self) -> int:
353
- return self.model_config.max_model_len
354
-
355
- def execute_worker(
356
- self,
357
- worker_input: WorkerInput,
358
- ) -> None:
359
- if (worker_input.blocks_to_copy is not None
360
- and worker_input.blocks_to_copy.numel() > 0):
361
- self.cache_engine[worker_input.virtual_engine].copy(
362
- worker_input.blocks_to_copy)
363
-
364
- @torch.inference_mode()
365
- def prepare_worker_input(
366
- self, execute_model_req: ExecuteModelRequest) -> WorkerInput:
367
- assert execute_model_req is not None
368
- virtual_engine: int = execute_model_req.virtual_engine
369
- num_seq_groups: int = len(execute_model_req.seq_group_metadata_list)
370
- blocks_to_copy = torch.tensor(execute_model_req.blocks_to_copy,
371
- device="cpu",
372
- dtype=torch.int64).view(-1, 2)
373
- assert len(execute_model_req.blocks_to_swap_in) == 0
374
- assert len(execute_model_req.blocks_to_swap_out) == 0
375
- return WorkerInput(
376
- num_seq_groups=num_seq_groups,
377
- blocks_to_copy=blocks_to_copy,
378
- virtual_engine=virtual_engine,
379
- )
380
-
381
- def init_distributed_environment(self) -> None:
382
- """Initialize the distributed environment."""
383
-
384
- parallel_config = self.parallel_config
385
- rank = self.rank
386
- distributed_init_method = self.distributed_init_method
387
- init_distributed_environment(
388
- world_size=parallel_config.world_size,
389
- rank=rank,
390
- distributed_init_method=distributed_init_method,
391
- backend="gloo",
392
- )
393
-
394
- # A small all_reduce for warmup.
395
- torch.distributed.all_reduce(torch.zeros(1).cpu())
396
-
397
- ensure_model_parallel_initialized(
398
- parallel_config.tensor_parallel_size,
399
- parallel_config.pipeline_parallel_size)
400
-
401
- def get_cache_block_size_bytes(self) -> int:
402
- """Return the size in bytes of a single KV cache block.
403
- """
404
- return CPUCacheEngine.get_cache_block_size(self.cache_config,
405
- self.model_config,
406
- self.parallel_config)
407
-
408
- def get_cpus_id_binding_based_on_numa_nodes(self) -> str:
409
- """Return CPUs id binding based on NUMA nodes.
410
- """
411
- rank_to_cpus = self.local_omp_cpuid
412
- # Setup OpenMP thread affinity based on NUMA nodes automatically
413
- world_size = self.vllm_config.parallel_config.world_size
414
- libnuma_found = util.find_spec("numa") is not None
415
- psutil_found = util.find_spec("psutil") is not None
416
- if libnuma_found and psutil_found:
417
- import psutil
418
- from numa import info
419
- cpu_count = psutil.cpu_count(logical=False)
420
- cpus_allow_list = psutil.Process().cpu_affinity()
421
- numa_size = info.get_num_configured_nodes()
422
- cpu_count_per_numa = cpu_count // numa_size
423
- num_of_reserved_cpu = min(envs.VLLM_CPU_NUM_OF_RESERVED_CPU,
424
- cpu_count_per_numa // 2)
425
-
426
- # check allow node_to_cpus list
427
- node_to_cpus = []
428
- for i in range(numa_size):
429
- node_intersect = set(
430
- info.node_to_cpus(i)).intersection(cpus_allow_list)
431
- if bool(node_intersect):
432
- node_to_cpus.append(list(node_intersect))
433
-
434
- if world_size > len(node_to_cpus):
435
- logger.error(
436
- "Auto thread-binding failed due to "
437
- "world size: %d is larger than "
438
- "allowed NUMA nodes number: %d."
439
- "Please try to bind threads manually.", world_size,
440
- len(node_to_cpus))
441
- else:
442
- end = cpu_count_per_numa - num_of_reserved_cpu
443
- rank_to_cpus_list = node_to_cpus[self.rank][:end]
444
- rank_to_cpus = ','.join(str(x) for x in rank_to_cpus_list)
445
- logger.info("auto thread-binding list: %s", rank_to_cpus)
446
- else:
447
- logger.warning(
448
- "Auto thread-binding is not supported due to "
449
- "the lack of package numa and psutil,"
450
- "fallback to no thread-binding. To get better performance,"
451
- "please try to manually bind threads.")
452
- return rank_to_cpus
vllm/worker/multi_step_tpu_worker.py
deleted
-108
@@ -1,108 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-from typing import Dict, Optional, Tuple
6
-
7
-import torch
8
-
9
-from vllm.distributed import broadcast_tensor_dict
10
-from vllm.sequence import ExecuteModelRequest
11
-from vllm.worker.tpu_model_runner import ModelInputForTPU
12
-from vllm.worker.tpu_worker import TPUWorker
13
-from vllm.worker.worker_base import WorkerInput
14
-
15
-
16
-class MultiStepTPUWorker(TPUWorker):
17
-
18
- def __init__(self, *args, **kwargs):
19
- super().__init__(*args, **kwargs)
20
- self.cached_model_input: Optional[ModelInputForTPU] = None
21
-
22
- def _get_driver_input_and_broadcast(
23
- self, execute_model_req: ExecuteModelRequest
24
- ) -> Tuple[ModelInputForTPU, WorkerInput, Dict[str, torch.Tensor]]:
25
- assert self.is_driver_worker
26
- assert execute_model_req.virtual_engine == 0
27
-
28
- is_first_multi_step = execute_model_req.is_first_multi_step
29
- is_last_step = execute_model_req.is_last_step
30
- if is_first_multi_step:
31
- worker_input: WorkerInput = self.prepare_worker_input(
32
- execute_model_req=execute_model_req)
33
- worker_input = dataclasses.replace(
34
- worker_input,
35
- num_steps=execute_model_req.num_lookahead_slots + 1)
36
- model_input: ModelInputForTPU = (
37
- self.model_runner.prepare_model_input(
38
- execute_model_req.seq_group_metadata_list,
39
- execute_model_req.virtual_engine,
40
- execute_model_req.finished_requests_ids))
41
-
42
- if execute_model_req.async_callback:
43
- model_input = dataclasses.replace(
44
- model_input,
45
- async_callback=execute_model_req.async_callback)
46
- else:
47
- assert self.cached_model_input is not None
48
- model_input = self.cached_model_input
49
- worker_input = WorkerInput()
50
- model_input = dataclasses.replace(
51
- model_input,
52
- is_first_multi_step=is_first_multi_step,
53
- is_last_step=is_last_step)
54
-
55
- if self.do_metadata_broadcast:
56
- if is_first_multi_step:
57
- broadcast_data = worker_input.as_broadcastable_tensor_dict()
58
- broadcast_data.update(
59
- model_input.as_broadcastable_tensor_dict())
60
- broadcast_tensor_dict(broadcast_data, src=0)
61
- else:
62
- broadcast_data = {
63
- "is_first_multi_step": is_first_multi_step,
64
- "is_last_step": is_last_step,
65
- }
66
- broadcast_tensor_dict(broadcast_data, src=0)
67
-
68
- # Retuning empty dict here to keep this compatible with
69
- # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast`
70
- return model_input, worker_input, {}
71
-
72
- def prepare_input(
73
- self,
74
- execute_model_req: Optional[ExecuteModelRequest] = None,
75
- ) -> Optional[Tuple[ModelInputForTPU, WorkerInput, Dict[str,
76
- torch.Tensor]]]:
77
- if self.is_driver_worker:
78
- if execute_model_req is None:
79
- if self.do_metadata_broadcast:
80
- broadcast_tensor_dict({}, src=0)
81
- return None
82
-
83
- model_input, worker_input, _ = self._get_driver_input_and_broadcast(
84
- execute_model_req)
85
- if model_input.is_first_multi_step:
86
- self.cached_model_input = model_input
87
- return model_input, worker_input, {}
88
- else:
89
- broadcast_data = broadcast_tensor_dict(src=0)
90
- if not broadcast_data:
91
- return None
92
-
93
- if len(broadcast_data) == 2:
94
- assert self.cached_model_input is not None
95
- self.cached_model_input = dataclasses.replace(
96
- self.cached_model_input,
97
- is_first_multi_step=broadcast_data["is_first_multi_step"],
98
- is_last_step=broadcast_data["is_last_step"])
99
- empty_worker_input = WorkerInput()
100
- return self.cached_model_input, empty_worker_input, {}
101
-
102
- worker_input = WorkerInput.from_broadcasted_tensor_dict(
103
- broadcast_data)
104
- model_input = (
105
- self.model_runner.
106
- make_model_input_from_broadcasted_tensor_dict(broadcast_data))
107
- self.cached_model_input = model_input
108
- return model_input, worker_input, {}
vllm/worker/tpu_model_runner.py
deleted
-909
@@ -1,909 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import enum
5
-import time
6
-from dataclasses import dataclass
7
-from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
8
- Type, Union)
9
-from unittest.mock import patch
10
-
11
-import numpy as np
12
-import torch
13
-import torch.nn as nn
14
-import torch_xla.core.xla_model as xm
15
-import torch_xla.runtime as xr
16
-
17
-from vllm.attention import AttentionMetadata, get_attn_backend
18
-from vllm.config import VllmConfig
19
-from vllm.forward_context import get_forward_context, set_forward_context
20
-from vllm.logger import init_logger
21
-from vllm.model_executor.layers.sampler import SamplerOutput
22
-from vllm.model_executor.model_loader import get_model
23
-from vllm.model_executor.sampling_metadata import SamplingMetadata
24
-from vllm.sequence import (CompletionSequenceGroupOutput, IntermediateTensors,
25
- Logprob, SequenceGroupMetadata, SequenceOutput)
26
-from vllm.worker.model_runner_base import (
27
- ModelRunnerBase, ModelRunnerInputBase,
28
- _add_attn_metadata_broadcastable_dict,
29
- _init_attn_metadata_from_tensor_dict)
30
-
31
-if TYPE_CHECKING:
32
- from vllm.attention.backends.abstract import AttentionBackend
33
-
34
-logger = init_logger(__name__)
35
-
36
-# Here we utilize the behavior that out-of-bound index is ignored.
37
-# FIXME(woosuk): Find a more reliable way to prevent possible bugs.
38
-_PAD_SLOT_ID = 1_000_000_000
39
-# FIXME(woosuk): Temporarily disabled top-p sampling since it's too slow.
40
-_ENABLE_TOP_P = False
41
-# FIXME(woosuk): A temporary hack to support `n > 1`.
42
-# This can significantly affect the performance if too large.
43
-_MAX_NUM_SAMPLES = 128
44
-
45
-
46
-class ExecutionMode(enum.Enum):
47
- PREFILL = enum.auto()
48
- DECODE = enum.auto()
49
- PREFIX_PREFILL = enum.auto()
50
-
51
- def is_prefill(self) -> bool:
52
- return self in (ExecutionMode.PREFILL, ExecutionMode.PREFIX_PREFILL)
53
-
54
-
55
-@dataclass(frozen=True)
56
-class ModelInputForTPU(ModelRunnerInputBase):
57
- token_ids: torch.Tensor
58
- position_ids: torch.Tensor
59
- attn_metadata: AttentionMetadata
60
- input_lens: torch.Tensor
61
- t: torch.Tensor
62
- p: torch.Tensor
63
- num_samples: int
64
- n: List[int]
65
- seq_groups: List[List[int]]
66
- is_first_multi_step: bool = True
67
- is_last_step: bool = True
68
- virtual_engine: int = 0
69
- async_callback: Optional[Callable] = None
70
-
71
- def as_broadcastable_tensor_dict(
72
- self) -> Dict[str, Union[int, torch.Tensor]]:
73
- tensor_dict = {
74
- "token_ids": self.token_ids,
75
- "position_ids": self.position_ids,
76
- "input_lens": self.input_lens,
77
- "t": self.t,
78
- "p": self.p,
79
- "num_samples": self.num_samples,
80
- "n": self.n,
81
- "seq_groups": self.seq_groups,
82
- "is_first_multi_step": self.is_first_multi_step,
83
- "is_last_step": self.is_last_step,
84
- "virtual_engine": self.virtual_engine,
85
- }
86
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
87
- return tensor_dict
88
-
89
- @classmethod
90
- def from_broadcasted_tensor_dict(
91
- cls: Type["ModelInputForTPU"],
92
- tensor_dict: Dict[str, Any],
93
- attn_backend: Optional["AttentionBackend"] = None,
94
- ) -> "ModelInputForTPU":
95
- if attn_backend is not None:
96
- tensor_dict = _init_attn_metadata_from_tensor_dict(
97
- attn_backend, tensor_dict)
98
- return cls(**tensor_dict)
99
-
100
-
101
-class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):
102
-
103
- def __init__(
104
- self,
105
- vllm_config: VllmConfig,
106
- is_driver_worker: bool = False,
107
- ):
108
- ModelRunnerBase.__init__(self, vllm_config=vllm_config)
109
- self.is_driver_worker = is_driver_worker
110
-
111
- self.block_size = self.cache_config.block_size
112
- self.max_num_blocks_per_seq = (self.model_config.max_model_len //
113
- self.block_size)
114
- self.block_tables = np.zeros(
115
- (self.scheduler_config.max_num_seqs, self.max_num_blocks_per_seq),
116
- dtype=np.int32)
117
- self.attn_backend = get_attn_backend(
118
- self.model_config.get_head_size(),
119
- self.model_config.dtype,
120
- self.cache_config.cache_dtype,
121
- self.block_size,
122
- self.model_config.is_attention_free,
123
- False,
124
- )
125
- self.cached_step_outputs: List[torch.Tensor] = []
126
-
127
- smem_size = 512 * 1024
128
- block_table_size = 4 * self.block_tables.size
129
- if block_table_size >= smem_size:
130
- logger.warning(
131
- "The max_model_len (%d) is too large. This may degrade the "
132
- "performance due to the insufficient smem size. Consider "
133
- "setting --max-model-len to a smaller value, like %d.",
134
- self.model_config.max_model_len,
135
- self.model_config.max_model_len /
136
- (block_table_size / smem_size))
137
-
138
- def load_model(self) -> None:
139
- self.device = self.device_config.device
140
-
141
- # NOTE(woosuk): While the executor assigns the TP ranks to the worker
142
- # process, the ranks can be different from the ranks internally assigned
143
- # by the xm runtime. Therefore, there is a mismatch in the rank
144
- # assignment between the gloo (cpu) runtime and the xm (tpu) runtime.
145
- # This is not a problem in linear layers because all-reduce is
146
- # rank-agnostic. However, it matters for all-gather as the ranks
147
- # determine the order of concatenating the output tensors.
148
- # As a workaround, we use the xm's rank assignment only when loading
149
- # the embedding weights.
150
- xm_tp_rank = xr.global_ordinal()
151
- with patch(
152
- "vllm.model_executor.layers.vocab_parallel_embedding."
153
- "get_tensor_model_parallel_rank",
154
- return_value=xm_tp_rank):
155
- model = get_model(vllm_config=self.vllm_config)
156
- model = model.eval()
157
- xm.wait_device_ops()
158
- model = ModelWrapper(model)
159
- self.model = torch.compile(model,
160
- backend="openxla",
161
- fullgraph=True,
162
- dynamic=False)
163
-
164
- def get_model(self) -> nn.Module:
165
- return self.model.model
166
-
167
- def _dummy_run(
168
- self,
169
- batch_size: int,
170
- seq_len: int,
171
- kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
172
- exec_mode: ExecutionMode,
173
- ) -> None:
174
- exec_mode = ExecutionMode(exec_mode)
175
- if exec_mode.is_prefill():
176
- seq_len = (seq_len + 15) // 16 * 16
177
- token_ids = torch.zeros((batch_size, seq_len),
178
- dtype=torch.int32,
179
- device=self.device)
180
- position_ids = torch.zeros((batch_size, seq_len),
181
- dtype=torch.int32,
182
- device=self.device)
183
- slot_mapping = torch.zeros((batch_size, seq_len),
184
- dtype=torch.int64,
185
- device=self.device)
186
- input_lens = torch.ones((batch_size, ),
187
- dtype=torch.int32,
188
- device=self.device)
189
- if exec_mode == ExecutionMode.PREFILL:
190
- attn_metadata = self.attn_backend.make_metadata(
191
- num_prefills=batch_size,
192
- num_prefill_tokens=batch_size * seq_len,
193
- num_decode_tokens=0,
194
- slot_mapping=slot_mapping,
195
- multi_modal_placeholder_index_maps=None,
196
- enable_kv_scales_calculation=False,
197
- block_tables=None,
198
- context_lens=None,
199
- effective_query_lens=None,
200
- )
201
- else:
202
- context_lens = torch.ones((batch_size, ),
203
- dtype=torch.int32,
204
- device=self.device)
205
- block_tables = torch.tensor(self.block_tables[:batch_size],
206
- dtype=torch.int32,
207
- device=self.device)
208
- effective_query_lens = torch.ones_like(context_lens)
209
- attn_metadata = self.attn_backend.make_metadata(
210
- num_prefills=batch_size,
211
- num_prefill_tokens=batch_size * seq_len,
212
- num_decode_tokens=0,
213
- slot_mapping=slot_mapping,
214
- multi_modal_placeholder_index_maps=None,
215
- enable_kv_scales_calculation=False,
216
- block_tables=block_tables,
217
- context_lens=context_lens,
218
- effective_query_lens=effective_query_lens,
219
- )
220
- else:
221
- assert seq_len == 1
222
- token_ids = torch.zeros((batch_size, seq_len),
223
- dtype=torch.int32,
224
- device=self.device)
225
- position_ids = torch.zeros((batch_size, seq_len),
226
- dtype=torch.int32,
227
- device=self.device)
228
- slot_mapping = torch.zeros((batch_size, seq_len),
229
- dtype=torch.int64,
230
- device=self.device)
231
- block_tables = torch.zeros(
232
- (batch_size, self.max_num_blocks_per_seq),
233
- dtype=torch.int32,
234
- device=self.device)
235
- context_lens = torch.ones((batch_size, ),
236
- dtype=torch.int32,
237
- device=self.device)
238
- input_lens = torch.ones((batch_size, ),
239
- dtype=torch.int32,
240
- device=self.device)
241
- attn_metadata = self.attn_backend.make_metadata(
242
- num_prefills=0,
243
- num_prefill_tokens=0,
244
- num_decode_tokens=batch_size * seq_len,
245
- slot_mapping=slot_mapping,
246
- multi_modal_placeholder_index_maps=None,
247
- enable_kv_scales_calculation=False,
248
- block_tables=block_tables,
249
- context_lens=context_lens,
250
- )
251
- t = torch.ones((batch_size, ), dtype=torch.float32, device=self.device)
252
- p = torch.ones((batch_size, ), dtype=torch.float32, device=self.device)
253
- num_samples = _MAX_NUM_SAMPLES if exec_mode.is_prefill() else 1
254
-
255
- # NOTE(woosuk): There are two stages of compilation: torch.compile and
256
- # XLA compilation. Using `mark_dynamic` can reduce the torch.compile
257
- # overhead by reusing the FX graph for different shapes.
258
- # However, the XLA graph will still require static shapes and needs to
259
- # be re-compiled for every different shapes. This overhead is inevitable
260
- # in the first run, but can be skipped afterwards as we cache the XLA
261
- # graphs in the disk (VLLM_XLA_CACHE_PATH).
262
- if exec_mode.is_prefill():
263
- # Prefll
264
- torch._dynamo.mark_dynamic(token_ids, 1)
265
- torch._dynamo.mark_dynamic(position_ids, 1)
266
- torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 1)
267
- else:
268
- # Decode
269
- torch._dynamo.mark_dynamic(token_ids, 0)
270
- torch._dynamo.mark_dynamic(position_ids, 0)
271
- torch._dynamo.mark_dynamic(input_lens, 0)
272
- torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)
273
- torch._dynamo.mark_dynamic(attn_metadata.context_lens, 0)
274
- torch._dynamo.mark_dynamic(attn_metadata.block_tables, 0)
275
- torch._dynamo.mark_dynamic(t, 0)
276
- torch._dynamo.mark_dynamic(p, 0)
277
- # Dummy run.
278
- with set_forward_context(attn_metadata, self.vllm_config, 0):
279
- self.model(token_ids, position_ids, input_lens, t, p, num_samples,
280
- kv_caches)
281
-
282
- def warmup_model(
283
- self,
284
- kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
285
- ) -> None:
286
- # Prefill
287
- logger.info("Compiling the model with different input shapes...")
288
- start = time.time()
289
- for batch_size in [1]:
290
- seq_len = 16
291
- while seq_len <= self.model_config.max_model_len:
292
- self._dummy_run(batch_size,
293
- seq_len,
294
- kv_caches,
295
- exec_mode=ExecutionMode.PREFILL)
296
- xm.wait_device_ops()
297
- logger.info("batch_size: %d, seq_len: %d", batch_size, seq_len)
298
- num_tokens = batch_size * seq_len
299
- if num_tokens >= self.scheduler_config.max_num_batched_tokens:
300
- break
301
- seq_len = seq_len * 2
302
-
303
- end = time.time()
304
- logger.info("Compilation for prefill done in %.2f s.", end - start)
305
-
306
- # Prefix prefill
307
- if self.cache_config.enable_prefix_caching:
308
- logger.info("Compiling the model with different input shapes for "
309
- "prefix prefill...")
310
- start = time.time()
311
- for batch_size in [1]:
312
- seq_len = 16
313
- while seq_len <= self.model_config.max_model_len:
314
- self._dummy_run(batch_size,
315
- seq_len,
316
- kv_caches,
317
- exec_mode=ExecutionMode.PREFIX_PREFILL)
318
- xm.wait_device_ops()
319
- logger.info("batch_size: %d, seq_len: %d", batch_size,
320
- seq_len)
321
- num_tokens = batch_size * seq_len
322
- if (num_tokens
323
- >= self.scheduler_config.max_num_batched_tokens):
324
- break
325
- seq_len = seq_len * 2
326
- end = time.time()
327
- logger.info("Compilation for prefix prefill done in %.2f s.",
328
- end - start)
329
-
330
- # Decode
331
- start = time.time()
332
- seq_len = 1
333
- batch_size = 8 # Must be in sync with _get_padded_batch_size()
334
- while True:
335
- self._dummy_run(batch_size,
336
- seq_len,
337
- kv_caches,
338
- exec_mode=ExecutionMode.DECODE)
339
- xm.wait_device_ops()
340
- logger.info("batch_size: %d, seq_len: %d", batch_size, seq_len)
341
-
342
- if batch_size >= self.scheduler_config.max_num_seqs:
343
- break
344
- batch_size = batch_size + 16 if batch_size >= 16 else batch_size * 2
345
-
346
- end = time.time()
347
- logger.info("Compilation for decode done in %.2f s.", end - start)
348
-
349
- def _prepare_prompt(
350
- self,
351
- seq_group_metadata_list: List[SequenceGroupMetadata],
352
- ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, torch.Tensor]:
353
- assert len(seq_group_metadata_list) > 0
354
- input_tokens: List[int] = []
355
- input_positions: List[int] = []
356
- prompt_lens: List[int] = []
357
- context_lens: List[int] = []
358
- slot_mapping: List[int] = []
359
-
360
- for batch_idx, seq_group_metadata in enumerate(
361
- seq_group_metadata_list):
362
- assert seq_group_metadata.is_prompt
363
- seq_ids = list(seq_group_metadata.seq_data.keys())
364
- assert len(seq_ids) == 1
365
- seq_id = seq_ids[0]
366
-
367
- seq_data = seq_group_metadata.seq_data[seq_id]
368
- # Could include output tokens when a request is preempted.
369
- prompt_tokens = seq_data.get_token_ids()
370
- seq_len = len(prompt_tokens)
371
-
372
- num_computed_blocks = len(seq_group_metadata.computed_block_nums)
373
- num_computed_tokens = num_computed_blocks * self.block_size
374
- if num_computed_tokens > 0:
375
- prompt_tokens = prompt_tokens[num_computed_tokens:]
376
- context_lens.append(seq_len)
377
- else:
378
- context_lens.append(0)
379
-
380
- prompt_len = len(prompt_tokens)
381
- prompt_lens.append(prompt_len)
382
-
383
- input_tokens.extend(prompt_tokens)
384
- input_positions.extend(range(num_computed_tokens, seq_len))
385
-
386
- assert seq_group_metadata.block_tables is not None
387
- block_table = seq_group_metadata.block_tables[seq_id]
388
- for i in range(num_computed_tokens, seq_len):
389
- block_number = block_table[i // self.block_size]
390
- block_offset = i % self.block_size
391
- slot = block_number * self.block_size + block_offset
392
- slot_mapping.append(slot)
393
- if num_computed_tokens > 0:
394
- self.block_tables[batch_idx, :len(block_table)] = block_table
395
-
396
- # Add paddings to EACH prompt to the smallest power of 2 that is
397
- # greater than or equal to the prompt length.
398
- # We pad the seq_len to reduce the compilation overhead.
399
- # We execute each prompt individually (i.e., with batch_size 1)
400
- # because the FlashAttention kernel does not support ragged inputs.
401
- # TODO(woosuk): Use SplashAttention to support ragged inputs.
402
- padded_prompt_len = _get_padded_prefill_len(prompt_len)
403
- num_paddings = padded_prompt_len - prompt_len
404
- input_tokens += [0] * num_paddings
405
- input_positions += [0] * num_paddings
406
- slot_mapping += [_PAD_SLOT_ID] * num_paddings
407
-
408
- assert len(prompt_lens) > 0
409
- num_prefills = len(prompt_lens)
410
- input_tokens = torch.tensor(input_tokens,
411
- dtype=torch.int32,
412
- device="cpu")
413
- input_positions = torch.tensor(input_positions,
414
- dtype=torch.int32,
415
- device="cpu")
416
- slot_mapping = torch.tensor(slot_mapping,
417
- dtype=torch.int64,
418
- device="cpu")
419
- prompt_lens = torch.tensor(prompt_lens,
420
- dtype=torch.int32,
421
- device="cpu")
422
- context_lens = torch.tensor(context_lens,
423
- dtype=torch.int32,
424
- device="cpu")
425
- block_tables = torch.tensor(self.block_tables[:num_prefills],
426
- dtype=torch.int32,
427
- device="cpu")
428
- attn_metadata = self.attn_backend.make_metadata(
429
- num_prefills=num_prefills,
430
- num_prefill_tokens=0, # NOTE: This is not used.
431
- num_decode_tokens=0,
432
- slot_mapping=slot_mapping,
433
- multi_modal_placeholder_index_maps=None,
434
- enable_kv_scales_calculation=False,
435
- block_tables=block_tables,
436
- context_lens=context_lens,
437
- effective_query_lens=prompt_lens,
438
- )
439
- return input_tokens, input_positions, attn_metadata, prompt_lens
440
-
441
- def _prepare_decode(
442
- self,
443
- seq_group_metadata_list: List[SequenceGroupMetadata],
444
- ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, torch.Tensor]:
445
- assert len(seq_group_metadata_list) > 0
446
- input_tokens: List[List[int]] = []
447
- input_positions: List[List[int]] = []
448
- slot_mapping: List[List[int]] = []
449
- context_lens: List[int] = []
450
-
451
- batch_idx = 0
452
- for seq_group_metadata in seq_group_metadata_list:
453
- assert not seq_group_metadata.is_prompt
454
- seq_ids = list(seq_group_metadata.seq_data.keys())
455
- for seq_id in seq_ids:
456
- seq_data = seq_group_metadata.seq_data[seq_id]
457
- generation_token = seq_data.get_last_token_id()
458
- input_tokens.append([generation_token])
459
-
460
- seq_len = seq_data.get_len()
461
- position = seq_len - 1
462
- input_positions.append([position])
463
- context_lens.append(seq_len)
464
-
465
- assert seq_group_metadata.block_tables is not None
466
- block_table = seq_group_metadata.block_tables[seq_id]
467
- self.block_tables[batch_idx, :len(block_table)] = block_table
468
- batch_idx += 1
469
-
470
- block_number = block_table[position // self.block_size]
471
- block_offset = position % self.block_size
472
- slot = block_number * self.block_size + block_offset
473
- slot_mapping.append([slot])
474
-
475
- batch_size = _get_padded_batch_size(batch_idx)
476
- num_paddings = batch_size - batch_idx
477
- input_tokens = input_tokens + [[0]] * num_paddings
478
- input_positions = input_positions + [[0]] * num_paddings
479
- slot_mapping = slot_mapping + [[_PAD_SLOT_ID]] * num_paddings
480
- context_lens = context_lens + [0] * num_paddings
481
-
482
- input_tokens = torch.tensor(input_tokens,
483
- dtype=torch.int32,
484
- device="cpu")
485
- input_positions = torch.tensor(input_positions,
486
- dtype=torch.int32,
487
- device="cpu")
488
- slot_mapping = torch.tensor(slot_mapping,
489
- dtype=torch.int64,
490
- device="cpu")
491
- context_lens = torch.tensor(context_lens,
492
- dtype=torch.int32,
493
- device="cpu")
494
- block_tables = torch.tensor(self.block_tables[:batch_size],
495
- dtype=torch.int32,
496
- device="cpu")
497
- input_lens = torch.tensor([1] * batch_size,
498
- dtype=torch.int32,
499
- device="cpu")
500
- attn_metadata = self.attn_backend.make_metadata(
501
- num_prefills=0,
502
- num_prefill_tokens=0,
503
- num_decode_tokens=batch_size,
504
- slot_mapping=slot_mapping,
505
- multi_modal_placeholder_index_maps=None,
506
- enable_kv_scales_calculation=False,
507
- block_tables=block_tables,
508
- context_lens=context_lens,
509
- )
510
- return input_tokens, input_positions, attn_metadata, input_lens
511
-
512
- def _prepare_sample(
513
- self,
514
- seq_group_metadata_list: List[SequenceGroupMetadata],
515
- padded_batch_size: int,
516
- ) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
517
- assert len(seq_group_metadata_list) > 0
518
- t = []
519
- p = []
520
- n = []
521
- for seq_group_metadata in seq_group_metadata_list:
522
- sampling_params = seq_group_metadata.sampling_params
523
- t.append(sampling_params.temperature)
524
- if sampling_params.top_p != 1 and not _ENABLE_TOP_P:
525
- raise NotImplementedError(
526
- "Top-p sampling is currently disabled for the TPU backend "
527
- "due to performance issues.")
528
- p.append(sampling_params.top_p)
529
- if sampling_params.top_k > 0:
530
- raise NotImplementedError(
531
- "Top-k sampling is currently disabled for the TPU backend "
532
- "due to performance issues.")
533
- if sampling_params.n > _MAX_NUM_SAMPLES:
534
- raise NotImplementedError(
535
- f"Best of > {_MAX_NUM_SAMPLES} is not supported by the TPU "
536
- "backend.")
537
- n.append(sampling_params.n)
538
- if sampling_params.logprobs is not None:
539
- raise NotImplementedError(
540
- "logprobs is not currently supported by the TPU backend.")
541
- if sampling_params.prompt_logprobs is not None:
542
- raise NotImplementedError(
543
- "prompt_logprobs is not currently supported by the TPU "
544
- "backend.")
545
-
546
- # Repeat the sampling params if the seq group has multiple seqs.
547
- num_seqs = len(seq_group_metadata.seq_data)
548
- t += [t[-1]] * (num_seqs - 1)
549
- p += [p[-1]] * (num_seqs - 1)
550
- n += [n[-1]] * (num_seqs - 1)
551
-
552
- num_paddings = padded_batch_size - len(t)
553
- t += [1.0] * num_paddings
554
- p += [1.0] * num_paddings
555
-
556
- t = torch.tensor(t, dtype=torch.float32, device="cpu")
557
- p = torch.tensor(p, dtype=torch.float32, device="cpu")
558
- return t, p, n
559
-
560
- def prepare_model_input(
561
- self,
562
- seq_group_metadata_list: List[SequenceGroupMetadata],
563
- virtual_engine: int = 0,
564
- finished_requests_ids: Optional[List[str]] = None,
565
- ) -> ModelInputForTPU:
566
- del finished_requests_ids # Unused.
567
- assert virtual_engine == 0
568
- assert len(seq_group_metadata_list) > 0
569
- # NOTE: We assume that all sequences in the group are all prompts or
570
- # all decodes.
571
- is_prompt = seq_group_metadata_list[0].is_prompt
572
- if is_prompt:
573
- inputs = self._prepare_prompt(seq_group_metadata_list)
574
- else:
575
- inputs = self._prepare_decode(seq_group_metadata_list)
576
- input_tokens, input_positions, attn_metadata, input_lens = inputs
577
- padded_batch_size = input_tokens.shape[0]
578
- t, p, n = self._prepare_sample(seq_group_metadata_list,
579
- padded_batch_size)
580
- num_samples = _MAX_NUM_SAMPLES if is_prompt else 1
581
-
582
- seq_groups = [
583
- list(metadata.seq_data.keys())
584
- for metadata in seq_group_metadata_list
585
- ]
586
- return ModelInputForTPU(input_tokens, input_positions, attn_metadata,
587
- input_lens, t, p, num_samples, n, seq_groups)
588
-
589
- def make_model_input_from_broadcasted_tensor_dict(
590
- self, tensor_dict: Dict[str, Any]) -> ModelInputForTPU:
591
- model_input = ModelInputForTPU.from_broadcasted_tensor_dict(
592
- tensor_dict, attn_backend=self.attn_backend)
593
- return model_input
594
-
595
- @torch.no_grad()
596
- def execute_model(
597
- self,
598
- model_input: ModelInputForTPU,
599
- kv_caches: Optional[List[Any]],
600
- intermediate_tensors: Optional[IntermediateTensors] = None,
601
- num_steps: int = 1,
602
- ) -> List[SamplerOutput]:
603
- assert intermediate_tensors is None
604
- if not model_input.is_first_multi_step:
605
- if not model_input.is_last_step:
606
- return []
607
-
608
- use_async_out_proc = model_input.async_callback is not None
609
- sampler_outputs = []
610
- num_outputs = len(self.cached_step_outputs)
611
- for i in range(num_outputs):
612
- next_token_ids = self.cached_step_outputs.pop(0)
613
- next_token_ids = next_token_ids.cpu().tolist()
614
- sampler_output = _make_decode_output(next_token_ids,
615
- model_input.seq_groups)
616
- sampler_outputs.append(sampler_output)
617
-
618
- if i < num_outputs - 1 and use_async_out_proc:
619
- assert model_input.async_callback is not None
620
- ctx = model_input.async_callback.keywords[ # type: ignore
621
- "ctx"]
622
- ctx.append_output(
623
- outputs=[sampler_output],
624
- seq_group_metadata_list=ctx.seq_group_metadata_list,
625
- scheduler_outputs=ctx.scheduler_outputs,
626
- is_async=False,
627
- is_last_step=False,
628
- is_first_step_output=i == 0)
629
- model_input.async_callback()
630
- if use_async_out_proc:
631
- return [sampler_outputs[-1]]
632
- else:
633
- return sampler_outputs
634
-
635
- is_prompt = model_input.attn_metadata.num_prefills > 0
636
- if is_prompt:
637
- assert num_steps == 1
638
- # NOTE(woosuk): Since the FlashAttention kernel does not support
639
- # ragged inputs, we split the prompts into different batches and
640
- # process them separately. This is a temporary hack that should be
641
- # optimized by using SplashAttention.
642
- orig_slot_mapping = model_input.attn_metadata.slot_mapping
643
- orig_block_tables = model_input.attn_metadata.block_tables
644
- orig_context_lens = model_input.attn_metadata.context_lens
645
- orig_effective_query_lens = \
646
- model_input.attn_metadata.effective_query_lens
647
- batch_size = model_input.input_lens.shape[0]
648
- start_idx = 0
649
- next_token_ids = []
650
- for i in range(batch_size):
651
- # Get the actual prefill_len.
652
- prefill_len = model_input.input_lens[i:i + 1].item()
653
- prefill_len = _get_padded_prefill_len(prefill_len)
654
- end_idx = start_idx + prefill_len
655
-
656
- token_ids = model_input.token_ids[None, start_idx:end_idx].to(
657
- self.device)
658
- position_ids = model_input.position_ids[None,
659
- start_idx:end_idx].to(
660
- self.device)
661
- attn_metadata = model_input.attn_metadata
662
- attn_metadata.num_prefills = 1
663
- attn_metadata.slot_mapping = orig_slot_mapping[
664
- None, start_idx:end_idx].to(self.device)
665
- if orig_context_lens[i].item() > 0:
666
- attn_metadata.context_lens = orig_context_lens[i:i + 1].to(
667
- self.device)
668
- attn_metadata.block_tables = orig_block_tables[
669
- i].unsqueeze(0).to(self.device)
670
- attn_metadata.effective_query_lens = \
671
- orig_effective_query_lens[i:i + 1].to(self.device)
672
- else:
673
- attn_metadata.context_lens = None
674
- attn_metadata.block_tables = None
675
- attn_metadata.effective_query_lens = None
676
- input_lens = model_input.input_lens[i:i + 1].to(self.device)
677
- t = model_input.t[i:i + 1].to(self.device)
678
- p = model_input.p[i:i + 1].to(self.device)
679
- with set_forward_context(model_input.attn_metadata,
680
- self.vllm_config,
681
- model_input.virtual_engine):
682
- output_token_ids = self.model(token_ids, position_ids,
683
- input_lens, t, p,
684
- model_input.num_samples,
685
- kv_caches)
686
- next_token_ids.append(output_token_ids[0])
687
- start_idx = end_idx
688
-
689
- if model_input.async_callback is not None:
690
- model_input.async_callback()
691
- # Retrieve the outputs to CPU.
692
- next_token_ids = [
693
- output_token_ids.cpu().tolist()
694
- for output_token_ids in next_token_ids
695
- ]
696
-
697
- # NOTE(woosuk): Minimal code to construct the sampler outputs.
698
- # The TPU backend does not reuse the sampler, since the TPU backend
699
- # does not support advanced sampling parameters such as logprobs.
700
- zero_logprob = Logprob(0.0)
701
- sampler_outputs = []
702
- for i, seq_group in enumerate(model_input.seq_groups):
703
- seq_ids = seq_group
704
- assert len(seq_ids) == 1
705
- seq_id = seq_ids[0]
706
- seq_outputs = []
707
- for j in range(model_input.n[i]):
708
- next_token_id = next_token_ids[i][j]
709
- seq_outputs.append(
710
- SequenceOutput(seq_id, next_token_id,
711
- {next_token_id: zero_logprob}))
712
- sampler_outputs.append(
713
- CompletionSequenceGroupOutput(seq_outputs, None))
714
- return [SamplerOutput(sampler_outputs)]
715
- else:
716
- token_ids = model_input.token_ids.to(self.device)
717
- position_ids = model_input.position_ids.to(self.device)
718
- attn_metadata = model_input.attn_metadata
719
- attn_metadata.slot_mapping = attn_metadata.slot_mapping.to(
720
- self.device)
721
- attn_metadata.block_tables = attn_metadata.block_tables.to(
722
- self.device)
723
- attn_metadata.context_lens = attn_metadata.context_lens.to(
724
- self.device)
725
- t = model_input.t.to(self.device)
726
- p = model_input.p.to(self.device)
727
- input_lens = model_input.input_lens.to(self.device)
728
- for i in range(num_steps):
729
- slot_mapping = attn_metadata.slot_mapping
730
- with set_forward_context(model_input.attn_metadata,
731
- self.vllm_config,
732
- model_input.virtual_engine):
733
- output_token_ids = self.model(token_ids, position_ids,
734
- input_lens, t, p,
735
- model_input.num_samples,
736
- kv_caches)
737
- self.cached_step_outputs.append(output_token_ids)
738
-
739
- if i < num_steps - 1:
740
- # Prepare the inputs for the next step.
741
- token_ids = output_token_ids.unsqueeze(dim=1).int()
742
- position_ids = position_ids + 1
743
- attn_metadata.context_lens = attn_metadata.context_lens + 1
744
-
745
- block_tables = attn_metadata.block_tables
746
- block_number = block_tables.gather(
747
- 1,
748
- position_ids.long() // self.block_size)
749
- block_offset = position_ids % self.block_size
750
-
751
- is_padding = slot_mapping == _PAD_SLOT_ID
752
- slot_mapping = block_number * self.block_size + block_offset
753
- slot_mapping = slot_mapping.long()
754
- slot_mapping = torch.where(is_padding, _PAD_SLOT_ID,
755
- slot_mapping)
756
- attn_metadata.slot_mapping = slot_mapping
757
-
758
- if model_input.async_callback is not None:
759
- model_input.async_callback()
760
-
761
- if num_steps > 1:
762
- return []
763
- # Retrieve the outputs to CPU.
764
- next_token_ids = self.cached_step_outputs.pop(0)
765
- next_token_ids = next_token_ids.cpu().tolist()
766
- sampler_output = _make_decode_output(next_token_ids,
767
- model_input.seq_groups)
768
- return [sampler_output]
769
-
770
-
771
-class ModelWrapper(nn.Module):
772
-
773
- def __init__(self, model: nn.Module):
774
- super().__init__()
775
- self.model = model
776
-
777
- def forward(
778
- self,
779
- token_ids: torch.Tensor,
780
- position_ids: torch.Tensor,
781
- input_lens: torch.Tensor,
782
- t: torch.Tensor,
783
- p: torch.Tensor,
784
- num_samples: int,
785
- kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
786
- ) -> torch.Tensor:
787
- """Executes the forward pass of the model and samples the next token.
788
-
789
- Args:
790
- token_ids: The input token IDs of shape [batch_size, seq_len].
791
- position_ids: The input position IDs of shape [batch_size, seq_len].
792
- input_lens: The actual input lengths of shape [batch_size].
793
- t: The sampling temperature of shape [batch_size].
794
- p: The top-p probability of shape [batch_size].
795
- num_samples: Number of samples to draw from each logits vector.
796
- kv_caches: The key and value caches. They can be None during the
797
- memory profiling at initialization.
798
- """
799
- batch_size, seq_len = token_ids.shape
800
- # Calculate the positions to sample from.
801
- start_indices = torch.arange(
802
- batch_size, dtype=torch.int32, device=input_lens.device) * seq_len
803
- logits_indices = start_indices + input_lens - 1
804
- attn_metadata = get_forward_context().attn_metadata
805
-
806
- # FIXME(woosuk): This is a temporary hack to avoid using the existing
807
- # sampler and sampling metadata.
808
- sampling_metadata = SamplingMetadata(
809
- seq_groups=[],
810
- selected_token_indices=logits_indices,
811
- categorized_sample_indices={},
812
- num_prompts=attn_metadata.num_prefills,
813
- )
814
-
815
- # Skip this in memory profiling at initialization.
816
- if kv_caches[0][0].numel() > 0:
817
- # index_copy_(slot_mapping) only works when the inserted dimension
818
- # is 0. However, the KV cache in the Pallas backend has the shape
819
- # [num_kv_heads, num_blocks, block_size, head_size]. To make it
820
- # work, we need to flatten the first three dimensions and modify
821
- # the slot_mapping accordingly.
822
- num_kv_heads, num_blocks, block_size, _ = kv_caches[0][0].shape
823
- slot_mapping = attn_metadata.slot_mapping
824
- slot_mapping = slot_mapping.flatten()
825
- head_indices = torch.arange(0,
826
- num_kv_heads,
827
- device=slot_mapping.device,
828
- dtype=slot_mapping.dtype)
829
- head_indices *= block_size * num_blocks
830
- slot_mapping = slot_mapping.repeat_interleave(num_kv_heads).view(
831
- -1, num_kv_heads)
832
- slot_mapping = slot_mapping + head_indices.view(1, -1)
833
- slot_mapping = slot_mapping.flatten()
834
- attn_metadata.slot_mapping = slot_mapping
835
-
836
- hidden_states = self.model(token_ids, position_ids)
837
- hidden_states = hidden_states.flatten(0, 1)
838
- logits = self.model.compute_logits(hidden_states, sampling_metadata)
839
-
840
- # Argmax sampling.
841
- argmax_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
842
- argmax_token_ids = argmax_token_ids.repeat(1, num_samples)
843
-
844
- # Zero temperature means greedy decoding. Avoid division by zero.
845
- nonzero_t = torch.where(t != 0, t, 1.0)
846
- logits = logits / nonzero_t.unsqueeze(dim=1)
847
- if _ENABLE_TOP_P:
848
- logits = _apply_top_p(logits, p.unsqueeze(dim=1))
849
-
850
- # Random sampling.
851
- probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
852
- sampled_token_ids = torch.multinomial(probs,
853
- num_samples,
854
- replacement=True)
855
- if num_samples == 1:
856
- argmax_token_ids = argmax_token_ids.squeeze(dim=-1)
857
- sampled_token_ids = sampled_token_ids.squeeze(dim=-1)
858
- next_token_ids = torch.where(t != 0, sampled_token_ids,
859
- argmax_token_ids)
860
- return next_token_ids
861
-
862
-
863
-def _get_padded_prefill_len(x: int) -> int:
864
- # NOTE(woosuk): The pallas FlashAttention kernel requires the sequence
865
- # length to be a multiple of 16. We pad the prompt length to the nearest
866
- # multiple of 16. This is also good for performance.
867
- if x <= 16:
868
- return 16
869
- return 1 << (x - 1).bit_length()
870
-
871
-
872
-def _get_padded_batch_size(batch_size: int) -> int:
873
- # The GMM Pallas kernel requires num_tokens * topk to be a multiple of 16.
874
- # To meet this requirement in the simplest way, we set the minimal batch
875
- # size to 8.
876
- if batch_size <= 8:
877
- return 8
878
- else:
879
- return ((batch_size + 15) // 16) * 16
880
-
881
-
882
-def _apply_top_p(logits: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
883
- logits_sorted = torch.sort(logits, dim=-1, descending=True).values
884
- sorted_cum_probs = torch.cumsum(logits_sorted.softmax(dim=-1), dim=-1)
885
- cutoff_index = torch.sum(sorted_cum_probs < p, dim=-1, keepdim=True)
886
- cutoff_logit = torch.gather(logits_sorted, -1, cutoff_index)
887
- logits = logits.masked_fill_(logits < cutoff_logit, -float("inf"))
888
- return logits
889
-
890
-
891
-def _make_decode_output(
892
- next_token_ids: List[int],
893
- seq_groups: List[List[int]],
894
-) -> SamplerOutput:
895
- zero_logprob = Logprob(0.0)
896
- sampler_outputs = []
897
- batch_idx = 0
898
- for seq_group in seq_groups:
899
- seq_ids = seq_group
900
- seq_outputs = []
901
- for seq_id in seq_ids:
902
- next_token_id = next_token_ids[batch_idx]
903
- seq_outputs.append(
904
- SequenceOutput(seq_id, next_token_id,
905
- {next_token_id: zero_logprob}))
906
- batch_idx += 1
907
- sampler_outputs.append(CompletionSequenceGroupOutput(
908
- seq_outputs, None))
909
- return SamplerOutput(sampler_outputs)
vllm/worker/tpu_worker.py
deleted
-337
@@ -1,337 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import os
5
-from typing import List, Optional, Tuple, Union
6
-
7
-import torch
8
-import torch_xla.core.xla_model as xm
9
-import torch_xla.debug.profiler as xp
10
-import torch_xla.runtime as xr
11
-
12
-import vllm.envs as envs
13
-from vllm.config import VllmConfig
14
-from vllm.distributed import (ensure_model_parallel_initialized,
15
- init_distributed_environment)
16
-from vllm.logger import init_logger
17
-from vllm.model_executor import set_random_seed
18
-from vllm.sequence import ExecuteModelRequest
19
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, bind_kv_cache, get_dtype_size
20
-from vllm.worker.tpu_model_runner import ExecutionMode, TPUModelRunner
21
-from vllm.worker.worker_base import (LocalOrDistributedWorkerBase,
22
- LoRANotSupportedWorkerBase, WorkerBase,
23
- WorkerInput)
24
-
25
-logger = init_logger(__name__)
26
-
27
-
28
-class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
29
-
30
- def __init__(
31
- self,
32
- vllm_config: VllmConfig,
33
- local_rank: int,
34
- rank: int,
35
- distributed_init_method: str,
36
- is_driver_worker: bool,
37
- ) -> None:
38
- WorkerBase.__init__(self, vllm_config=vllm_config)
39
- self.parallel_config.rank = rank
40
- self.local_rank = local_rank
41
- self.rank = rank
42
- self.distributed_init_method = distributed_init_method
43
- self.is_driver_worker = is_driver_worker
44
-
45
- assert self.device_config.device_type == "tpu"
46
- if self.cache_config.cache_dtype == "auto":
47
- self.cache_dtype = self.model_config.dtype
48
- else:
49
- self.cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
50
- self.cache_config.cache_dtype]
51
-
52
- self.model_runner: TPUModelRunner = TPUModelRunner(
53
- vllm_config=vllm_config, is_driver_worker=is_driver_worker)
54
-
55
- if self.model_config.seed is None:
56
- self.model_config.seed = 0
57
-
58
- if vllm_config.lora_config is not None:
59
- raise NotImplementedError(
60
- "The V0 TPU backend doesn't support LoRA serving")
61
-
62
- def init_device(self) -> None:
63
- os.environ["PJRT_DEVICE"] = "TPU"
64
- torch.set_grad_enabled(False)
65
- torch.set_default_dtype(self.model_config.dtype)
66
-
67
- # NOTE(woosuk): This is just to initialize the TP group and broadcast
68
- # the input objects on CPU. The all-reduce and all-gather ops on TPU
69
- # are invoked by `xm.all_reduce` and `xm.all_gather` which use their
70
- # own context.
71
- init_distributed_environment(
72
- world_size=self.parallel_config.world_size,
73
- rank=self.rank,
74
- local_rank=self.local_rank,
75
- distributed_init_method=self.distributed_init_method,
76
- backend="gloo",
77
- )
78
- ensure_model_parallel_initialized(
79
- self.parallel_config.tensor_parallel_size,
80
- self.parallel_config.pipeline_parallel_size)
81
-
82
- # Device initialization should happen after initializing the distributed
83
- # runtime.
84
- self.device = xm.xla_device()
85
- self.device_config.device = self.device
86
-
87
- # Set random seed.
88
- set_random_seed(self.model_config.seed)
89
- xm.set_rng_state(self.model_config.seed, self.device)
90
-
91
- # Increase the cache size limit, which is the maximum number of
92
- # dynamo graphs that can be compiled.
93
- # NOTE(woosuk): Usually, we compile 10-15 graphs for prefill and
94
- # 30-40 graphs for decode. 128 is an arbitrary safe number.
95
- torch._dynamo.config.cache_size_limit = 128
96
- # Use persistent cache to avoid XLA recompilation.
97
- # NOTE(woosuk): Set per-rank cache path since different ranks
98
- # can have slightly different XLA graphs.
99
- world_size = self.parallel_config.world_size
100
- rank = xr.global_ordinal()
101
- # The PyTorch/XLA compilation cache uses the Torch IR to generate keys.
102
- # Consequently, changes in optimization flags, which affect compilation
103
- # results, don't change the cache key. This can result in the wrong
104
- # compilation being used. To prevent this, disabling the XLA compilation
105
- # cache during development is recommended.We can disable it by
106
- # `export VLLM_XLA_CACHE_PATH=`
107
- if envs.VLLM_XLA_CACHE_PATH:
108
- per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
109
- f"tp{world_size}_rank{rank}")
110
- xr.initialize_cache(per_rank_path, readonly=False)
111
-
112
- self.profiler = None
113
- if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
114
- # For TPU, we can only have 1 active profiler session for 1 profiler
115
- # server. So we only profile on rank0.
116
- self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
117
- logger.info("Profiling enabled. Traces will be saved to: %s",
118
- self.profile_dir)
119
- self.profiler = xp.start_server(9012)
120
-
121
- def start_profile(self):
122
- if self.rank < 1:
123
- if self.profiler is None:
124
- raise RuntimeError("Profiler is not enabled.")
125
- xp.start_trace(self.profile_dir)
126
-
127
- def stop_profile(self):
128
- if self.rank < 1:
129
- if self.profiler is None:
130
- raise RuntimeError("Profiler is not enabled.")
131
- xp.stop_trace()
132
-
133
- def load_model(self):
134
- self.model_runner.load_model()
135
-
136
- def determine_num_available_blocks(self) -> Tuple[int, int]:
137
- num_layers = self.model_config.get_num_layers(self.parallel_config)
138
- head_size = self.model_config.get_head_size()
139
- num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
140
-
141
- # use an empty tensor instead of `None`` to force Dynamo to pass
142
- # it by reference, rather by specializing on the value ``None``.
143
- # the `dtype` argument does not matter, and we use `float32` as
144
- # a placeholder (it has wide hardware support).
145
- kv_caches = [(torch.tensor([], dtype=torch.float32,
146
- device=self.device),
147
- torch.tensor([], dtype=torch.float32,
148
- device=self.device))
149
- for _ in range(num_layers)]
150
- bind_kv_cache(self.compilation_config.static_forward_context,
151
- [kv_caches])
152
- self.model_runner._dummy_run(
153
- batch_size=1,
154
- seq_len=self.scheduler_config.max_num_batched_tokens,
155
- kv_caches=kv_caches,
156
- exec_mode=ExecutionMode.PREFILL,
157
- )
158
- # Synchronize before measuring the memory usage.
159
- xm.wait_device_ops()
160
-
161
- # Get the maximum amount of memory used by the model weights and
162
- # intermediate activations.
163
- m = xm.get_memory_info(self.device)
164
- total_memory_size = m["bytes_limit"]
165
- profiled = m["peak_bytes_used"] # Weights + intermediate activations.
166
-
167
- # Calculate the TPU KV cache size based on profiling.
168
- usable_memory_size = int(total_memory_size *
169
- self.cache_config.gpu_memory_utilization)
170
- tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0)
171
- dtype_bytes = get_dtype_size(self.cache_dtype)
172
- block_size_bytes = (dtype_bytes * self.cache_config.block_size *
173
- num_layers * 2 * head_size * num_kv_heads)
174
- num_tpu_blocks = tpu_kv_cache_bytes // block_size_bytes
175
- num_tpu_blocks = (num_tpu_blocks // 8) * 8 # Round down to 8.
176
-
177
- # Calculate the CPU KV cache size based on the config.
178
- num_cpu_blocks = int(self.cache_config.swap_space_bytes //
179
- block_size_bytes)
180
- num_cpu_blocks = (num_cpu_blocks // 8) * 8 # Round down to 8.
181
- return num_tpu_blocks, num_cpu_blocks
182
-
183
- def initialize_cache(
184
- self,
185
- num_gpu_blocks: int,
186
- num_cpu_blocks: int,
187
- ) -> None:
188
- self.cache_config.num_gpu_blocks = num_gpu_blocks
189
- self.cache_config.num_cpu_blocks = num_cpu_blocks
190
- self.block_size = self.cache_config.block_size
191
-
192
- dtype = self.cache_dtype
193
- num_layers = self.model_config.get_num_layers(self.parallel_config)
194
- num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
195
- head_size = self.model_config.get_head_size()
196
-
197
- self.cpu_cache: List[Tuple[torch.Tensor, torch.Tensor]] = []
198
- self.tpu_cache: List[Tuple[torch.Tensor, torch.Tensor]] = []
199
- tpu_cache_shape = self.model_runner.attn_backend.get_kv_cache_shape(
200
- num_gpu_blocks, self.block_size, num_kv_heads, head_size)
201
- cpu_cache_shape = self.model_runner.attn_backend.get_kv_cache_shape(
202
- num_cpu_blocks, self.block_size, num_kv_heads, head_size)
203
- for _ in range(num_layers):
204
- tpu_k_cache = torch.zeros(tpu_cache_shape,
205
- dtype=dtype,
206
- device=self.device)
207
- tpu_v_cache = torch.zeros_like(tpu_k_cache)
208
- self.tpu_cache.append((tpu_k_cache, tpu_v_cache))
209
- cpu_k_cache = torch.zeros(cpu_cache_shape,
210
- dtype=dtype,
211
- device="cpu")
212
- cpu_v_cache = torch.zeros_like(cpu_k_cache)
213
- self.cpu_cache.append((cpu_k_cache, cpu_v_cache))
214
- bind_kv_cache(self.compilation_config.static_forward_context,
215
- [self.tpu_cache])
216
- self._warmup_model()
217
-
218
- def _warmup_model(self) -> None:
219
- # FIXME(woosuk): Here we are abusing `enforce_eager` which is defined
220
- # for CUDA graphs. We should refactor this part.
221
- if not self.model_config.enforce_eager:
222
- # Warm up the model with all possible input shapes so that
223
- # compilation never happens during the actual execution.
224
- # This may take ~30 mins for the first run and ~20 mins for the
225
- # subsequent runs.
226
- # If `enforce_eager` is True, the ahead-of-time compilation is
227
- # skipped and the compilation happens during the actual execution,
228
- # which is bad for performance but useful for development.
229
- self.model_runner.warmup_model(self.tpu_cache)
230
-
231
- def get_cache_block_size_bytes(self) -> int:
232
- head_size = self.model_config.get_head_size()
233
- num_heads = self.model_config.get_num_kv_heads(self.parallel_config)
234
- num_layers = self.model_config.get_num_layers(self.parallel_config)
235
-
236
- key_cache_block = self.cache_config.block_size * num_heads * head_size
237
- value_cache_block = key_cache_block
238
- total = num_layers * (key_cache_block + value_cache_block)
239
- dtype_size = get_dtype_size(self.cache_dtype)
240
- return dtype_size * total
241
-
242
- @property
243
- def do_metadata_broadcast(self) -> bool:
244
- return self.parallel_config.tensor_parallel_size > 1
245
-
246
- @property
247
- def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:
248
- # NOTE(woosuk): This assumes virtual_engine == 0, i.e., no pipeline
249
- # parallelism.
250
- return [self.tpu_cache]
251
-
252
- def prepare_worker_input(
253
- self,
254
- execute_model_req: ExecuteModelRequest,
255
- ) -> WorkerInput:
256
- virtual_engine = execute_model_req.virtual_engine
257
- num_seq_groups = len(execute_model_req.seq_group_metadata_list)
258
- blocks_to_swap_in = _make_src_to_dst(
259
- execute_model_req.blocks_to_swap_in, "cpu", self.device)
260
- blocks_to_swap_out = _make_src_to_dst(
261
- execute_model_req.blocks_to_swap_out, self.device, "cpu")
262
- blocks_to_copy = _make_src_to_dst(execute_model_req.blocks_to_copy,
263
- self.device, self.device)
264
- return WorkerInput(
265
- num_seq_groups=num_seq_groups,
266
- blocks_to_swap_in=blocks_to_swap_in,
267
- blocks_to_swap_out=blocks_to_swap_out,
268
- blocks_to_copy=blocks_to_copy,
269
- virtual_engine=virtual_engine,
270
- )
271
-
272
- def execute_worker(self, worker_input: WorkerInput) -> None:
273
- virtual_engine = worker_input.virtual_engine
274
- assert virtual_engine == 0
275
- attn_backend = self.model_runner.attn_backend
276
- num_layers = self.model_config.get_num_layers(self.parallel_config)
277
-
278
- # Issue cache operations.
279
- if worker_input.blocks_to_swap_in is not None:
280
- src_indices, dst_indices = worker_input.blocks_to_swap_in
281
- if src_indices.numel() > 0:
282
- # Swap from CPU to TPU.
283
- for i in range(num_layers):
284
- tpu_k_cache, tpu_v_cache = self.tpu_cache[i]
285
- cpu_k_cache, cpu_v_cache = self.cpu_cache[i]
286
- k = cpu_k_cache[:, src_indices].to(self.device)
287
- v = cpu_v_cache[:, src_indices].to(self.device)
288
- _insert_kv(k, v, dst_indices, tpu_k_cache, tpu_v_cache)
289
-
290
- if worker_input.blocks_to_swap_out is not None:
291
- src_indices, dst_indices = worker_input.blocks_to_swap_out
292
- if src_indices.numel() > 0:
293
- # Swap from TPU to CPU.
294
- for i in range(num_layers):
295
- tpu_k_cache, tpu_v_cache = self.tpu_cache[i]
296
- cpu_k_cache, cpu_v_cache = self.cpu_cache[i]
297
- cpu_k_cache[:, dst_indices] = tpu_k_cache[:, src_indices]
298
- cpu_v_cache[:, dst_indices] = tpu_v_cache[:, src_indices]
299
-
300
- if worker_input.blocks_to_copy is not None:
301
- src_indices, dst_indices = worker_input.blocks_to_copy
302
- if src_indices.numel() > 0:
303
- attn_backend.copy_blocks(self.tpu_cache,
304
- (src_indices, dst_indices))
305
-
306
-
307
-def _make_src_to_dst(
308
- mapping: List[Tuple[int, int]],
309
- src_device: Union[torch.device, str],
310
- dst_device: Union[torch.device, str],
311
-) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
312
- if not mapping:
313
- return None
314
-
315
- src_indices = [i for i, _ in mapping]
316
- dst_indices = [i for _, i in mapping]
317
- src_indices = torch.tensor(src_indices,
318
- device=src_device,
319
- dtype=torch.int64)
320
- dst_indices = torch.tensor(dst_indices,
321
- device=dst_device,
322
- dtype=torch.int64)
323
- return src_indices, dst_indices
324
-
325
-
326
-@torch.compile(backend="openxla")
327
-def _insert_kv(
328
- k: torch.Tensor,
329
- v: torch.Tensor,
330
- indices: torch.Tensor,
331
- tpu_k_cache: torch.Tensor,
332
- tpu_v_cache: torch.Tensor,
333
-) -> None:
334
- torch.ops.xla.dynamo_set_buffer_donor_(tpu_k_cache, True)
335
- torch.ops.xla.dynamo_set_buffer_donor_(tpu_v_cache, True)
336
- tpu_k_cache[:, indices] = k
337
- tpu_v_cache[:, indices] = v
vllm/worker/xpu_model_runner.py
deleted
-606
@@ -1,606 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import dataclasses
5
-import time
6
-import weakref
7
-from collections import defaultdict
8
-from dataclasses import dataclass
9
-from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
10
- Type, TypeVar)
11
-
12
-import torch
13
-import torch.nn as nn
14
-
15
-from vllm.attention import get_attn_backend
16
-from vllm.config import VllmConfig
17
-from vllm.distributed import get_pp_group
18
-from vllm.forward_context import set_forward_context
19
-from vllm.inputs import INPUT_REGISTRY, InputRegistry
20
-from vllm.logger import init_logger
21
-from vllm.model_executor import SamplingMetadataCache
22
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
23
-from vllm.model_executor.model_loader import get_model
24
-from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
25
- MultiModalKwargs, MultiModalPlaceholderMap,
26
- MultiModalRegistry)
27
-from vllm.sampling_params import SamplingParams
28
-from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
29
-from vllm.utils import DeviceMemoryProfiler, GiB_bytes, make_tensor_with_pad
30
-from vllm.worker.model_runner import AttentionMetadata, SamplingMetadata
31
-from vllm.worker.model_runner_base import (
32
- ModelRunnerBase, ModelRunnerInputBase, ModelRunnerInputBuilderBase,
33
- _add_attn_metadata_broadcastable_dict,
34
- _add_sampling_metadata_broadcastable_dict,
35
- _init_attn_metadata_from_tensor_dict,
36
- _init_sampling_metadata_from_tensor_dict)
37
-
38
-if TYPE_CHECKING:
39
- from vllm.attention.backends.abstract import AttentionBackend
40
-
41
-logger = init_logger(__name__)
42
-
43
-_PAD_SLOT_ID = -1
44
-
45
-TModelInputForXPU = TypeVar('TModelInputForXPU', bound="ModelInputForXPU")
46
-
47
-
48
-@dataclass(frozen=True)
49
-class ModelInputForXPU(ModelRunnerInputBase):
50
- """
51
- Used by the NeuronModelRunner.
52
- """
53
- input_tokens: Optional[torch.Tensor] = None
54
- input_positions: Optional[torch.Tensor] = None
55
- attn_metadata: Optional["AttentionMetadata"] = None
56
- multi_modal_kwargs: Optional[BatchedTensorInputs] = None
57
- virtual_engine: Optional[int] = None
58
- seq_lens: Optional[List[int]] = None
59
- query_lens: Optional[List[int]] = None
60
- async_callback: Optional[Callable] = None
61
-
62
- def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
63
- tensor_dict = {
64
- "input_tokens": self.input_tokens,
65
- "input_positions": self.input_positions,
66
- }
67
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
68
-
69
- return tensor_dict
70
-
71
- @classmethod
72
- def from_broadcasted_tensor_dict(
73
- cls: Type[TModelInputForXPU],
74
- tensor_dict: Dict[str, Any],
75
- attn_backend: Optional["AttentionBackend"] = None,
76
- ) -> TModelInputForXPU:
77
- if attn_backend is not None:
78
- tensor_dict = _init_attn_metadata_from_tensor_dict(
79
- attn_backend, tensor_dict)
80
- return cls(**tensor_dict)
81
-
82
-
83
-@dataclass(frozen=True)
84
-class ModelInputForXPUWithSamplingMetadata(ModelInputForXPU):
85
- """
86
- Used by the ModelRunner.
87
- """
88
- sampling_metadata: Optional["SamplingMetadata"] = None
89
-
90
- def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
91
- tensor_dict = {
92
- "input_tokens": self.input_tokens,
93
- "input_positions": self.input_positions,
94
- }
95
- _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
96
- _add_sampling_metadata_broadcastable_dict(tensor_dict,
97
- self.sampling_metadata)
98
- return tensor_dict
99
-
100
- @classmethod
101
- def from_broadcasted_tensor_dict(
102
- cls,
103
- tensor_dict: Dict[str, Any],
104
- attn_backend: Optional["AttentionBackend"] = None,
105
- ) -> "ModelInputForXPUWithSamplingMetadata":
106
- tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)
107
- if attn_backend is not None:
108
- tensor_dict = _init_attn_metadata_from_tensor_dict(
109
- attn_backend, tensor_dict)
110
- return cls(**tensor_dict)
111
-
112
-
113
-class ModelInputForXPUBuilder(ModelRunnerInputBuilderBase[ModelInputForXPU]):
114
-
115
- def __init__(self,
116
- runner: "XPUModelRunner",
117
- finished_requests_ids: Optional[List[str]] = None) -> None:
118
- super().__init__()
119
- self.runner = runner
120
- self.model_input_cls = self.runner._model_input_cls
121
- self.attn_backend = self.runner.attn_backend
122
- self.sliding_window = self.runner.sliding_window
123
- self.block_size = self.runner.block_size
124
- self.device = self.runner.device
125
-
126
- def prepare(self,
127
- finished_requests_ids: Optional[List[str]] = None) -> None:
128
- self.seq_group_metadata_list: List[SequenceGroupMetadata] = []
129
-
130
- def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
131
- self.seq_group_metadata_list.append(seq_group_metadata)
132
-
133
- def build(self) -> ModelInputForXPU:
134
- is_prompt = self.seq_group_metadata_list[0].is_prompt
135
- # Prepare input tensors.
136
- if is_prompt:
137
- (input_tokens, input_positions, attn_metadata, seq_lens,
138
- multi_modal_kwargs) = self._prepare_prompt(
139
- self.seq_group_metadata_list)
140
- else:
141
- (input_tokens, input_positions,
142
- attn_metadata) = self._prepare_decode(
143
- self.seq_group_metadata_list)
144
- seq_lens = None
145
- multi_modal_kwargs = None
146
-
147
- return self.model_input_cls(
148
- input_tokens=input_tokens,
149
- input_positions=input_positions,
150
- attn_metadata=attn_metadata,
151
- multi_modal_kwargs=multi_modal_kwargs,
152
- seq_lens=seq_lens,
153
- query_lens=seq_lens,
154
- )
155
-
156
- def _prepare_prompt(
157
- self,
158
- seq_group_metadata_list: List[SequenceGroupMetadata],
159
- ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, List[int],
160
- BatchedTensorInputs]:
161
- assert len(seq_group_metadata_list) > 0
162
- input_tokens: List[int] = []
163
- input_positions: List[int] = []
164
- slot_mapping: List[int] = []
165
- seq_lens: List[int] = []
166
- multi_modal_kwargs_list: List[MultiModalKwargs] = []
167
- multi_modal_placeholder_maps: Dict[
168
- str,
169
- MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)
170
-
171
- for seq_group_metadata in seq_group_metadata_list:
172
- assert seq_group_metadata.is_prompt
173
- seq_ids = list(seq_group_metadata.seq_data.keys())
174
- assert len(seq_ids) == 1
175
- seq_id = seq_ids[0]
176
-
177
- seq_data = seq_group_metadata.seq_data[seq_id]
178
- prompt_tokens = seq_data.get_token_ids()
179
- computed_len = seq_data.get_num_computed_tokens()
180
- seq_len = len(prompt_tokens)
181
-
182
- seq_lens.append(seq_len) # Prompt token num
183
- input_tokens.extend(prompt_tokens) # Token ids
184
-
185
- # Token position ids
186
- # NOTE(woosuk): Here we assume that the first token in the prompt
187
- # is always the first token in the sequence.
188
- positions_range = range(computed_len, seq_len)
189
- input_positions.extend(list(positions_range))
190
-
191
- if seq_group_metadata.multi_modal_data:
192
- # NOTE: mm_kwargs only includes the subset of multi-modal items
193
- # that intersect with the current prefill positions.
194
- mm_kwargs, placeholder_maps = MultiModalPlaceholderMap \
195
- .from_seq_group(seq_group_metadata, positions_range)
196
-
197
- multi_modal_kwargs_list.append(mm_kwargs)
198
-
199
- for modality, placeholder_map in placeholder_maps.items():
200
- multi_modal_placeholder_maps[modality].extend(
201
- placeholder_map)
202
-
203
- if seq_group_metadata.block_tables is None:
204
- # During memory profiling, the block tables are not initialized
205
- # yet. In this case, we just use a dummy slot mapping.
206
- slot_mapping.extend([_PAD_SLOT_ID] * seq_len)
207
- continue
208
-
209
- # Compute the slot mapping.
210
- block_table = seq_group_metadata.block_tables[seq_id]
211
- # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,
212
- # where start_idx is max(0, seq_len - sliding_window).
213
- # For example, if the prompt len is 10, sliding window is 8, and
214
- # block size is 4, the first two tokens are masked and the slot
215
- # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
216
- start_idx = 0
217
- if self.sliding_window is not None:
218
- start_idx = max(0, seq_len - self.sliding_window)
219
-
220
- for i in range(computed_len, seq_len):
221
- if i < start_idx:
222
- slot_mapping.append(_PAD_SLOT_ID)
223
- continue
224
-
225
- block_number = block_table[i //
226
- self.block_size] # type: ignore
227
- block_offset = i % self.block_size # type: ignore
228
- slot = block_number * self.block_size + block_offset
229
- slot_mapping.append(slot)
230
-
231
- num_prompt_tokens = len(input_tokens)
232
-
233
- input_tokens = torch.tensor(input_tokens,
234
- dtype=torch.long,
235
- device=self.device) # type: ignore
236
- input_positions = torch.tensor(input_positions,
237
- dtype=torch.long,
238
- device=self.device) # type: ignore
239
- slot_mapping = torch.tensor(slot_mapping,
240
- dtype=torch.long,
241
- device=self.device) # type: ignore
242
- placeholder_index_maps = {
243
- modality: placeholder_map.index_map()
244
- for modality, placeholder_map in
245
- multi_modal_placeholder_maps.items()
246
- }
247
-
248
- max_seqlen = max(seq_lens)
249
- tmp = [0]
250
- tmp.extend(seq_lens)
251
- seqlen = torch.tensor(tmp)
252
- seqlen_q = torch.cumsum(seqlen, dim=0).to(device=self.device)
253
-
254
- attn_metadata = self.attn_backend.make_metadata(
255
- is_prompt=True,
256
- slot_mapping=slot_mapping,
257
- multi_modal_placeholder_index_maps=placeholder_index_maps,
258
- enable_kv_scales_calculation=False,
259
- seq_lens=seq_lens,
260
- seqlen_q=seqlen_q,
261
- max_seqlen=max_seqlen,
262
- seq_lens_tensor=torch.tensor([]),
263
- max_decode_seq_len=0,
264
- num_prefills=len(seq_lens),
265
- num_prefill_tokens=num_prompt_tokens,
266
- num_decode_tokens=0,
267
- block_tables=torch.tensor([], device=self.device, dtype=torch.int),
268
- )
269
-
270
- multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
271
-
272
- return (input_tokens, input_positions, attn_metadata, seq_lens,
273
- multi_modal_kwargs)
274
-
275
- def _prepare_decode(
276
- self,
277
- seq_group_metadata_list: List[SequenceGroupMetadata],
278
- ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata]:
279
- assert len(seq_group_metadata_list) > 0
280
- input_tokens: List[int] = []
281
- input_positions: List[int] = []
282
- slot_mapping: List[int] = []
283
- seq_lens: List[int] = []
284
- block_tables: List[List[int]] = []
285
-
286
- for seq_group_metadata in seq_group_metadata_list:
287
- assert not seq_group_metadata.is_prompt
288
- assert seq_group_metadata.token_chunk_size == 1
289
-
290
- seq_ids = list(seq_group_metadata.seq_data.keys())
291
-
292
- for seq_id in seq_ids:
293
- seq_data = seq_group_metadata.seq_data[seq_id]
294
- generation_token = seq_data.get_last_token_id()
295
- input_tokens.append(generation_token)
296
-
297
- seq_len = seq_data.get_len()
298
- position = seq_len - 1
299
- input_positions.append(position)
300
-
301
- seq_len = seq_len if self.sliding_window is None else min(
302
- seq_len, self.sliding_window)
303
- seq_lens.append(seq_len)
304
-
305
- block_table = seq_group_metadata.block_tables[seq_id]
306
- block_number = block_table[position // self.block_size]
307
- block_offset = position % self.block_size
308
- slot = block_number * self.block_size + block_offset
309
- slot_mapping.append(slot)
310
-
311
- if self.sliding_window is not None:
312
- sliding_window_blocks = (self.sliding_window //
313
- self.block_size)
314
- block_table = block_table[-sliding_window_blocks:]
315
- block_tables.append(block_table)
316
-
317
- max_decode_seq_len = max(seq_lens)
318
-
319
- input_tokens = torch.tensor(input_tokens,
320
- dtype=torch.long,
321
- device=self.device)
322
- input_positions = torch.tensor(input_positions,
323
- dtype=torch.long,
324
- device=self.device)
325
- slot_mapping = torch.tensor(slot_mapping,
326
- dtype=torch.long,
327
- device=self.device)
328
- seq_lens_tensor = torch.tensor(seq_lens,
329
- dtype=torch.int,
330
- device=self.device)
331
-
332
- block_tables = make_tensor_with_pad(
333
- block_tables,
334
- pad=0,
335
- dtype=torch.int,
336
- device=self.device,
337
- )
338
-
339
- attn_metadata = self.attn_backend.make_metadata(
340
- is_prompt=False,
341
- slot_mapping=slot_mapping,
342
- multi_modal_placeholder_index_maps=None,
343
- enable_kv_scales_calculation=False,
344
- seq_lens=seq_lens,
345
- seqlen_q=torch.tensor([]),
346
- max_seqlen=0,
347
- seq_lens_tensor=seq_lens_tensor,
348
- max_decode_seq_len=max_decode_seq_len,
349
- num_prefill_tokens=0,
350
- num_decode_tokens=len(input_tokens),
351
- num_prefills=0,
352
- block_tables=block_tables,
353
- )
354
- return (
355
- input_tokens,
356
- input_positions,
357
- attn_metadata,
358
- )
359
-
360
-
361
-class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):
362
- _model_input_cls: Type[ModelInputForXPUWithSamplingMetadata] = (
363
- ModelInputForXPUWithSamplingMetadata)
364
- _builder_cls: Type[ModelInputForXPUBuilder] = ModelInputForXPUBuilder
365
-
366
- def __init__(
367
- self,
368
- vllm_config: VllmConfig,
369
- kv_cache_dtype: Optional[str] = "auto",
370
- is_driver_worker: bool = False,
371
- return_hidden_states: bool = False,
372
- input_registry: InputRegistry = INPUT_REGISTRY,
373
- mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
374
- ):
375
-
376
- ModelRunnerBase.__init__(self, vllm_config=vllm_config)
377
- model_config = self.model_config
378
- cache_config = self.cache_config
379
- self.is_driver_worker = is_driver_worker
380
- self.return_hidden_states = return_hidden_states
381
-
382
- self.device = self.device_config.device
383
-
384
- self.kv_cache_dtype = kv_cache_dtype
385
- self.sliding_window = model_config.get_sliding_window()
386
- self.block_size = cache_config.block_size
387
-
388
- self.attn_backend = get_attn_backend(
389
- self.model_config.get_head_size(),
390
- self.model_config.dtype,
391
- self.kv_cache_dtype,
392
- self.block_size,
393
- self.model_config.is_attention_free,
394
- )
395
-
396
- # Multi-modal data support
397
- self.input_registry = input_registry
398
- self.mm_registry = mm_registry
399
-
400
- # Lazy initialization.
401
- self.model: nn.Module # Set after init_Model
402
- self.sampler = get_sampler()
403
-
404
- self.sampling_metadata_cache: SamplingMetadataCache = \
405
- SamplingMetadataCache() \
406
- if self.parallel_config.pipeline_parallel_size == 1 else None
407
-
408
- self.builder = self._builder_cls(weakref.proxy(self))
409
-
410
- def load_model(self) -> None:
411
- with DeviceMemoryProfiler() as m:
412
- self.model = get_model(vllm_config=self.vllm_config)
413
-
414
- self.model_memory_usage = m.consumed_memory
415
- logger.info("Loading model weights took %.4f GiB",
416
- self.model_memory_usage / GiB_bytes)
417
-
418
- def get_model(self) -> nn.Module:
419
- return self.model
420
-
421
- @property
422
- def vocab_size(self) -> int:
423
- return self.model_config.get_vocab_size()
424
-
425
- @torch.inference_mode()
426
- def profile_run(self) -> None:
427
- # Enable top-k sampling to reflect the accurate memory usage.
428
- sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
429
- max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
430
- max_num_seqs = self.scheduler_config.max_num_seqs
431
-
432
- # Profile memory usage with max_num_sequences sequences and the total
433
- # number of tokens equal to max_num_batched_tokens.
434
- seqs: List[SequenceGroupMetadata] = []
435
- # Additional GPU memory may be needed for multi-modal encoding, which
436
- # needs to be accounted for when calculating the GPU blocks for
437
- # vLLM blocker manager.
438
- # To exercise the worst scenario for GPU memory consumption,
439
- # the number of seqs (batch_size) is chosen to maximize the number
440
- # of images processed.
441
- max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
442
- self.model_config)
443
- if max_mm_tokens > 0:
444
- max_num_seqs_orig = max_num_seqs
445
- max_num_seqs = min(max_num_seqs,
446
- max_num_batched_tokens // max_mm_tokens)
447
- if max_num_seqs < 1:
448
- expr = (f"min({max_num_seqs_orig}, "
449
- f"{max_num_batched_tokens} // {max_mm_tokens})")
450
- logger.warning(
451
- "Computed max_num_seqs (%s) to be less than 1. "
452
- "Setting it to the minimum value of 1.", expr)
453
- max_num_seqs = 1
454
-
455
- batch_size = 0
456
- for group_id in range(max_num_seqs):
457
- seq_len = (max_num_batched_tokens // max_num_seqs +
458
- (group_id < max_num_batched_tokens % max_num_seqs))
459
- batch_size += seq_len
460
-
461
- dummy_data = self.input_registry \
462
- .dummy_data_for_profiling(self.model_config,
463
- seq_len,
464
- self.mm_registry)
465
-
466
- seq = SequenceGroupMetadata(
467
- request_id=str(group_id),
468
- is_prompt=True,
469
- seq_data={group_id: dummy_data.seq_data},
470
- sampling_params=sampling_params,
471
- block_tables=None,
472
- lora_request=None,
473
- multi_modal_data=dummy_data.multi_modal_data,
474
- multi_modal_placeholders=dummy_data.multi_modal_placeholders)
475
- seqs.append(seq)
476
-
477
- finished_requests_ids = [seq.request_id for seq in seqs]
478
- model_input = self.prepare_model_input(
479
- seqs, finished_requests_ids=finished_requests_ids)
480
- intermediate_tensors = None
481
- if not get_pp_group().is_first_rank:
482
- intermediate_tensors = self.model.make_empty_intermediate_tensors(
483
- batch_size=batch_size,
484
- dtype=self.model_config.dtype,
485
- device=self.device)
486
- self.execute_model(model_input, None, intermediate_tensors)
487
- torch.xpu.synchronize()
488
- return
489
-
490
- def make_model_input_from_broadcasted_tensor_dict(
491
- self,
492
- tensor_dict: Dict[str,
493
- Any]) -> ModelInputForXPUWithSamplingMetadata:
494
- return (
495
- ModelInputForXPUWithSamplingMetadata.from_broadcasted_tensor_dict(
496
- tensor_dict,
497
- attn_backend=self.attn_backend,
498
- ))
499
-
500
- def _prepare_model_input_tensors(
501
- self,
502
- seq_group_metadata_list: List[SequenceGroupMetadata],
503
- finished_requests_ids: Optional[List[str]] = None
504
- ) -> ModelInputForXPUWithSamplingMetadata:
505
- """Helper method to prepare the model input based on a given sequence
506
- group. Prepares metadata needed for the base model forward pass but not
507
- metadata for possible additional steps, e.g., sampling.
508
-
509
- """
510
- builder = self.builder
511
- builder.prepare(finished_requests_ids)
512
- for seq_group_metadata in seq_group_metadata_list:
513
- builder.add_seq_group(seq_group_metadata)
514
-
515
- return builder.build() # type: ignore
516
-
517
- def prepare_model_input(
518
- self,
519
- seq_group_metadata_list: List[SequenceGroupMetadata],
520
- virtual_engine: int = 0,
521
- finished_requests_ids: Optional[List[str]] = None
522
- ) -> ModelInputForXPUWithSamplingMetadata:
523
- """Prepare the model input based on a given sequence group, including
524
- metadata for the sampling step.
525
-
526
- """
527
- model_input = self._prepare_model_input_tensors(
528
- seq_group_metadata_list, finished_requests_ids)
529
- # Sampling metadata is only required for the final pp group
530
- generators = self.get_generators(finished_requests_ids)
531
- sampling_metadata = SamplingMetadata.prepare(
532
- seq_group_metadata_list,
533
- model_input.seq_lens,
534
- model_input.query_lens,
535
- self.device,
536
- pin_memory=False,
537
- generators=generators,
538
- cache=self.sampling_metadata_cache)
539
-
540
- return dataclasses.replace(model_input,
541
- sampling_metadata=sampling_metadata,
542
- virtual_engine=virtual_engine)
543
-
544
- @torch.inference_mode()
545
- def execute_model(
546
- self,
547
- model_input: ModelInputForXPUWithSamplingMetadata,
548
- kv_caches: List[torch.Tensor],
549
- intermediate_tensors: Optional[IntermediateTensors] = None,
550
- num_steps: int = 1,
551
- ) -> Optional[List[SamplerOutput]]:
552
- if num_steps > 1:
553
- raise ValueError(
554
- "XPUModelRunner does not support multi-step execution.")
555
-
556
- model_executable = self.model
557
- if (self.observability_config is not None
558
- and self.observability_config.collect_model_forward_time):
559
- model_forward_start_time = time.time()
560
- with set_forward_context(model_input.attn_metadata, self.vllm_config,
561
- model_input.virtual_engine):
562
- hidden_or_intermediate_states = model_executable(
563
- input_ids=model_input.input_tokens,
564
- positions=model_input.input_positions,
565
- intermediate_tensors=intermediate_tensors,
566
- **MultiModalKwargs.as_kwargs(
567
- model_input.multi_modal_kwargs or {},
568
- device=self.device,
569
- ),
570
- )
571
- # Compute the logits in the last pipeline stage.
572
- if not get_pp_group().is_last_rank:
573
- return hidden_or_intermediate_states
574
-
575
- if (self.observability_config is not None
576
- and self.observability_config.collect_model_forward_time):
577
- model_forward_end_time = time.time()
578
-
579
- # Compute the logits.
580
- logits = self.model.compute_logits(hidden_or_intermediate_states,
581
- model_input.sampling_metadata)
582
-
583
- # Only perform sampling in the driver worker.
584
- if not self.is_driver_worker:
585
- return []
586
-
587
- if model_input.async_callback is not None:
588
- model_input.async_callback()
589
-
590
- # Sample the next token.
591
- output: SamplerOutput = self.sampler(
592
- logits=logits,
593
- sampling_metadata=model_input.sampling_metadata,
594
- )
595
- if (self.observability_config is not None
596
- and self.observability_config.collect_model_forward_time
597
- and output is not None):
598
- model_forward_time = (model_forward_end_time -
599
- model_forward_start_time)
600
- # If there are multiple workers, we are still tracking the latency
601
- # from the start time of the driver worker to the end time of the
602
- # driver worker. The model forward time will then end up covering
603
- # the communication time as well.
604
- output.model_forward_time = model_forward_time
605
-
606
- return [output]
vllm/worker/xpu_worker.py
deleted
-186
@@ -1,186 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""A XPU worker class."""
4
-import gc
5
-import os
6
-from typing import List, Optional, Tuple
7
-
8
-import intel_extension_for_pytorch # noqa: F401
9
-import oneccl_bindings_for_pytorch # noqa: F401
10
-import torch
11
-import torch.distributed
12
-
13
-from vllm.config import VllmConfig
14
-from vllm.distributed import (ensure_model_parallel_initialized,
15
- init_distributed_environment)
16
-from vllm.distributed.parallel_state import get_pp_group
17
-from vllm.logger import init_logger
18
-from vllm.model_executor import set_random_seed
19
-from vllm.platforms import current_platform
20
-from vllm.worker.cache_engine import CacheEngine
21
-from vllm.worker.worker import Worker
22
-from vllm.worker.worker_base import LoRANotSupportedWorkerBase, WorkerBase
23
-from vllm.worker.xpu_model_runner import XPUModelRunner
24
-
25
-logger = init_logger(__name__)
26
-
27
-
28
-class XPUWorker(LoRANotSupportedWorkerBase, Worker):
29
- """A worker class that executes (a partition of) the model on a GPU.
30
-
31
- Each worker is associated with a single XPU device. The worker is
32
- responsible for maintaining the KV cache and executing the model on the
33
- XPU. In case of distributed inference, each worker is assigned a partition
34
- of the model.
35
- """
36
-
37
- def __init__(
38
- self,
39
- vllm_config: VllmConfig,
40
- local_rank: int,
41
- rank: int,
42
- distributed_init_method: str,
43
- is_driver_worker: bool = False,
44
- ) -> None:
45
- WorkerBase.__init__(self, vllm_config=vllm_config)
46
- device_config = self.device_config
47
- parallel_config = self.parallel_config
48
- assert device_config.device_type == "xpu"
49
- assert current_platform.is_xpu()
50
-
51
- self.parallel_config.rank = rank
52
-
53
- self.local_rank = local_rank
54
- self.rank = rank
55
- self.distributed_init_method = distributed_init_method
56
- self.is_driver_worker = is_driver_worker
57
- if parallel_config and is_driver_worker:
58
- assert rank % parallel_config.tensor_parallel_size == 0, \
59
- "Driver worker should be rank 0 of tensor parallel group."
60
-
61
- self.model_runner = XPUModelRunner( # type: ignore
62
- vllm_config=vllm_config,
63
- kv_cache_dtype=self.cache_config.cache_dtype,
64
- is_driver_worker=is_driver_worker,
65
- )
66
- # Uninitialized cache engine. Will be initialized by
67
- # initialize_cache.
68
- self.cache_engine: List[CacheEngine]
69
- self.gpu_cache: Optional[List[List[torch.Tensor]]]
70
-
71
- def init_device(self) -> None:
72
- if self.device_config.device.type == "xpu" and current_platform.is_xpu(
73
- ):
74
- self.device = torch.device(f"xpu:{self.local_rank}")
75
- torch.xpu.set_device(self.device)
76
- torch.xpu.empty_cache()
77
- self.init_gpu_memory = torch.xpu.get_device_properties(
78
- self.local_rank).total_memory
79
- else:
80
- raise RuntimeError(
81
- f"Not support device type: {self.device_config.device}")
82
- # Initialize the distributed environment.
83
- self.init_worker_distributed_environment()
84
- # Initialize the model.
85
- set_random_seed(self.model_config.seed)
86
-
87
- # keep this method for `empty_cache` and `synchronize` api
88
- @torch.inference_mode()
89
- def determine_num_available_blocks(self) -> Tuple[int, int]:
90
- """Profiles the peak memory usage of the model to determine how many
91
- KV blocks may be allocated without OOMs.
92
-
93
- The engine will first conduct a profiling of the existing memory usage.
94
- Then, it calculate the maximum possible number of GPU and CPU blocks
95
- that can be allocated with the remaining free memory.
96
-
97
- Tip:
98
- You may limit the usage of GPU memory
99
- by adjusting the `gpu_memory_utilization` parameter.
100
- """
101
- # Profile the memory usage of the model and get the maximum number of
102
- # cache blocks that can be allocated with the remaining free memory.
103
- torch.xpu.empty_cache()
104
-
105
- # Execute a forward pass with dummy inputs to profile the memory usage
106
- # of the model.
107
- self.model_runner.profile_run()
108
-
109
- # Calculate the number of blocks that can be allocated with the
110
- # profiled peak memory.
111
- torch.xpu.synchronize()
112
- used_memory = torch.xpu.memory_allocated()
113
- total_gpu_memory = torch.xpu.get_device_properties(
114
- self.local_rank).total_memory
115
- free_gpu_memory = total_gpu_memory - used_memory
116
-
117
- # NOTE(woosuk): Here we assume that the other processes using the same
118
- # GPU did not change their memory usage during the profiling.
119
- peak_memory = self.init_gpu_memory - free_gpu_memory
120
- assert peak_memory > 0, (
121
- "Error in memory profiling. "
122
- f"Initial free memory {self.init_gpu_memory}, current free memory"
123
- f" {free_gpu_memory}. This happens when the GPU memory was "
124
- "not properly cleaned up before initializing the vLLM instance.")
125
-
126
- cache_block_size = self.get_cache_block_size_bytes()
127
- num_gpu_blocks = int(
128
- (total_gpu_memory * self.cache_config.gpu_memory_utilization -
129
- peak_memory) // cache_block_size)
130
- num_cpu_blocks = int(self.cache_config.swap_space_bytes //
131
- cache_block_size)
132
- num_gpu_blocks = max(num_gpu_blocks, 0)
133
- num_cpu_blocks = max(num_cpu_blocks, 0)
134
- gc.collect()
135
- torch.xpu.empty_cache()
136
- return num_gpu_blocks, num_cpu_blocks
137
-
138
- def _warm_up_model(self) -> None:
139
- # IPEX don't support capture graph yet
140
- pass
141
-
142
- def init_worker_distributed_environment(self) -> None:
143
- """Initialize the distributed environment."""
144
-
145
- parallel_config = self.parallel_config
146
- rank = self.rank
147
- distributed_init_method = self.distributed_init_method
148
-
149
- if torch.distributed.is_initialized():
150
- torch_world_size = torch.distributed.get_world_size()
151
- if torch_world_size != parallel_config.world_size:
152
- raise RuntimeError(
153
- "torch.distributed is already initialized but the torch "
154
- "world size does not match parallel_config.world_size "
155
- f"({torch_world_size} vs. {parallel_config.world_size}).")
156
- elif not distributed_init_method:
157
- raise ValueError(
158
- "distributed_init_method must be set if torch.distributed "
159
- "is not already initialized")
160
- else:
161
- # use sockets as default Level zero IPC exchange backend. By
162
- # default oneccl will use `drmfd` as mechanism which need extra
163
- # dependency (libdrm and drm headers) on your system.
164
- ENV_CCL_ATL_TRANSPORT = os.getenv("CCL_ATL_TRANSPORT", "ofi")
165
- ENV_LOCAL_WORLD_SIZE = os.getenv("LOCAL_WORLD_SIZE",
166
- str(parallel_config.world_size))
167
- os.environ["CCL_ATL_TRANSPORT"] = ENV_CCL_ATL_TRANSPORT
168
- os.environ["LOCAL_WORLD_SIZE"] = ENV_LOCAL_WORLD_SIZE
169
- os.environ["LOCAL_RANK"] = str(self.local_rank)
170
- init_distributed_environment(
171
- world_size=parallel_config.world_size,
172
- rank=rank,
173
- distributed_init_method=distributed_init_method,
174
- local_rank=self.local_rank,
175
- backend="ccl")
176
-
177
- ensure_model_parallel_initialized(
178
- parallel_config.tensor_parallel_size,
179
- parallel_config.pipeline_parallel_size)
180
- # global all_reduce needed for overall oneccl warm up
181
- torch.distributed.all_reduce(torch.zeros(1).xpu())
182
-
183
- if parallel_config.pipeline_parallel_size > 1:
184
- # Add pp group init to avoid
185
- # p2p communication as the first call
186
- get_pp_group().all_reduce(torch.zeros(1).xpu())