@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e6750d0b18
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[V0 Deprecation] Remove unused classes in attention (#25541)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>

Woosuk Kwon committed Sep 24, 2025 at 13:24 UTC e6750d0b18e07631bb2ea7f256f7dd444d4936fa
6 files changed +11 -716
vllm/attention/__init__.py
+1 -5
@@ -2,9 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from vllm.attention.backends.abstract import (AttentionBackend,
5 - AttentionMetadata,
6 - AttentionMetadataBuilder,
7 - AttentionState, AttentionType)
5 + AttentionMetadata, AttentionType)
6 from vllm.attention.layer import Attention
7 from vllm.attention.selector import get_attn_backend
8
@@ -13,7 +11,5 @@ __all__ = [
11 "AttentionBackend",
12 "AttentionMetadata",
13 "AttentionType",
16 - "AttentionMetadataBuilder",
17 - "AttentionState",
14 "get_attn_backend",
15 ]
vllm/attention/backends/abstract.py
+3 -142
@@ -2,10 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 -from contextlib import contextmanager
6 -from dataclasses import dataclass, fields
7 -from typing import (Any, Dict, Generic, List, Optional, Protocol, Set, Tuple,
8 - Type, TypeVar)
5 +from typing import Generic, List, Optional, Protocol, Tuple, Type, TypeVar
6
7 import torch
8
@@ -49,18 +46,13 @@ class AttentionBackend(ABC):
46 def get_metadata_cls() -> Type["AttentionMetadata"]:
47 raise NotImplementedError
48
52 - @staticmethod
53 - @abstractmethod
54 - def get_state_cls() -> Type["AttentionState"]:
55 - raise NotImplementedError
56 -
49 @classmethod
50 def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
51 return cls.get_metadata_cls()(*args, **kwargs)
52
53 @staticmethod
54 @abstractmethod
63 - def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
55 + def get_builder_cls(): # -> Type["AttentionMetadataBuilder"]:
56 raise NotImplementedError
57
58 @staticmethod
@@ -77,149 +69,18 @@ class AttentionBackend(ABC):
69 def get_kv_cache_stride_order() -> Tuple[int, ...]:
70 raise NotImplementedError
71
80 - @staticmethod
81 - @abstractmethod
82 - def swap_blocks(
83 - src_kv_cache: torch.Tensor,
84 - dst_kv_cache: torch.Tensor,
85 - src_to_dst: torch.Tensor,
86 - ) -> None:
87 - raise NotImplementedError
88 -
89 - @staticmethod
90 - @abstractmethod
91 - def copy_blocks(
92 - kv_caches: List[torch.Tensor],
93 - src_to_dists: torch.Tensor,
94 - ) -> None:
95 - raise NotImplementedError
96 -
72 @classmethod
73 def full_cls_name(cls) -> tuple[str, str]:
74 return (cls.__module__, cls.__qualname__)
75
76
102 -@dataclass
77 class AttentionMetadata:
104 - """Attention metadata for prefill and decode batched together."""
105 - # Total number of prefill requests.
106 - num_prefills: int
107 - # Number of prefill tokens.
108 - num_prefill_tokens: int
109 - # Number of decode tokens. Note that it is equivalent to the number of
110 - # decode requests.
111 - num_decode_tokens: int
112 - # (num_tokens,). The indices of the token slots that input tokens will be
113 - # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size
114 - # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot
115 - # in block 0, and 1st slot in block 1, respectively.
116 - slot_mapping: torch.Tensor
117 -
118 - # Enable/disable KV scales calculation. This is so that we can disable the
119 - # calculation until after prefill and cuda graph capture.
120 - enable_kv_scales_calculation: bool
121 -
122 - @property
123 - @abstractmethod
124 - def prefill_metadata(self) -> Optional["AttentionMetadata"]:
125 - """Return the attention metadata that's required to run prefill
126 - attention."""
127 - pass
128 -
129 - @property
130 - @abstractmethod
131 - def decode_metadata(self) -> Optional["AttentionMetadata"]:
132 - """Return the attention metadata that's required to run decode
133 - attention."""
134 - pass
135 -
136 - def asdict_zerocopy(self,
137 - skip_fields: Optional[Set[str]] = None
138 - ) -> Dict[str, Any]:
139 - """Similar to dataclasses.asdict, but avoids deepcopying."""
140 - if skip_fields is None:
141 - skip_fields = set()
142 - # Note that if we add dataclasses as fields, they will need
143 - # similar handling.
144 - return {
145 - field.name: getattr(self, field.name)
146 - for field in fields(self) if field.name not in skip_fields
147 - }
78 + pass
79
80
81 T = TypeVar("T", bound=AttentionMetadata)
82
83
153 -class AttentionState(ABC, Generic[T]):
154 - """Holds attention backend-specific objects reused during the
155 - lifetime of the model runner."""
156 -
157 - @abstractmethod
158 - def __init__(self, runner: Any):
159 - ...
160 -
161 - @abstractmethod
162 - @contextmanager
163 - def graph_capture(self, max_batch_size: int):
164 - """Context manager used when capturing CUDA graphs."""
165 - yield
166 -
167 - @abstractmethod
168 - def graph_clone(self, batch_size: int) -> "AttentionState[T]":
169 - """Clone attention state to save in CUDA graph metadata."""
170 - ...
171 -
172 - @abstractmethod
173 - def graph_capture_get_metadata_for_batch(
174 - self,
175 - batch_size: int,
176 - is_encoder_decoder_model: bool = False) -> T:
177 - """Get attention metadata for CUDA graph capture of batch_size."""
178 - ...
179 -
180 - @abstractmethod
181 - def get_graph_input_buffers(
182 - self,
183 - attn_metadata: T,
184 - is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
185 - """Get attention-specific input buffers for CUDA graph capture."""
186 - ...
187 -
188 - @abstractmethod
189 - def prepare_graph_input_buffers(
190 - self,
191 - input_buffers: Dict[str, Any],
192 - attn_metadata: T,
193 - is_encoder_decoder_model: bool = False) -> None:
194 - """In-place modify input buffers dict for CUDA graph replay."""
195 - ...
196 -
197 - @abstractmethod
198 - def begin_forward(self, model_input) -> None:
199 - """Prepare state for forward pass."""
200 - ...
201 -
202 -
203 -class AttentionMetadataBuilder(ABC, Generic[T]):
204 - """Abstract class for attention metadata builders."""
205 -
206 - @abstractmethod
207 - def __init__(self, input_builder) -> None:
208 - """Create the builder, remember some configuration and parameters."""
209 - raise NotImplementedError
210 -
211 - @abstractmethod
212 - def prepare(self) -> None:
213 - """Prepare for one batch."""
214 - raise NotImplementedError
215 -
216 - @abstractmethod
217 - def build(self, seq_lens: List[int], query_lens: List[int],
218 - cuda_graph_pad_size: int, batch_size: int) -> T:
219 - """Build attention metadata with on-device tensors."""
220 - raise NotImplementedError
221 -
222 -
84 class AttentionLayer(Protocol):
85
86 _q_scale: torch.Tensor
vllm/attention/backends/utils.py
+1 -544
@@ -1,559 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention backend utils"""
4 -from contextlib import contextmanager
4 from dataclasses import dataclass
6 -from itertools import accumulate
7 -from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
5 +from typing import Optional
6
9 -import numpy as np
10 -import torch
11 -
12 -from vllm.attention import (AttentionMetadata, AttentionMetadataBuilder,
13 - AttentionState)
14 -from vllm.attention.backends.abstract import AttentionType
7 from vllm.config import ModelConfig
8 from vllm.logger import init_logger
17 -from vllm.utils import async_tensor_h2d, make_tensor_with_pad
9
10 logger = init_logger(__name__)
11
12 PAD_SLOT_ID = -1
13
23 -# Switch to numpy implementation of compute_slot_mapping
24 -# if we have at least this many elements. Could be tuned further.
25 -_COMPUTE_SLOT_MAPPING_NUMPY_NUMEL = 256
26 -
27 -
28 -def is_block_tables_empty(block_tables: Union[None, Dict]):
29 - """
30 - Check if block_tables is None or a dictionary with all None values.
31 - """
32 - if block_tables is None:
33 - return True
34 - return (isinstance(block_tables, dict)
35 - and all(value is None for value in block_tables.values()))
36 -
37 -
38 -def compute_slot_mapping_start_idx(is_prompt: bool, query_len: int,
39 - context_len: int, sliding_window: int):
40 - """
41 - Compute the start index of slot mapping.
42 - """
43 - start_idx = 0
44 - if is_prompt and sliding_window is not None:
45 - start_idx = max(0, query_len - sliding_window)
46 - return start_idx
47 -
48 -
49 -def _compute_slot_mapping_python(slot_mapping: List[int],
50 - block_table: List[int], range_start: int,
51 - range_end: int, block_size: int):
52 - for i in range(range_start, range_end):
53 - block_number = block_table[i // block_size]
54 - block_offset = i % block_size
55 - slot = block_number * block_size + block_offset
56 - slot_mapping.append(slot)
57 -
58 -
59 -def _compute_slot_mapping_numpy(slot_mapping: List[int],
60 - block_table: List[int], range_start: int,
61 - range_end: int, block_size: int):
62 - block_table_array = np.array(block_table)
63 - idx = np.arange(range_start, range_end)
64 - block_offset = idx % block_size
65 - idx //= block_size
66 - seq_slot_mapping_array = block_table_array[idx]
67 - seq_slot_mapping_array *= block_size
68 - seq_slot_mapping_array += block_offset
69 - slot_mapping.extend(seq_slot_mapping_array)
70 -
71 -
72 -def compute_slot_mapping(is_profile_run: bool, slot_mapping: List[int],
73 - seq_id: int, seq_len: int, context_len: int,
74 - start_idx: int, block_size: int,
75 - block_tables: Dict[int, List[int]]):
76 - """
77 - Compute slot mapping.
78 - """
79 - if is_profile_run:
80 - # During memory profiling, the block tables are not
81 - # initialized yet. In this case, we just use a dummy
82 - # slot mapping.
83 - # In embeddings, the block tables are {seq_id: None}.
84 - slot_mapping.extend([PAD_SLOT_ID] * seq_len)
85 - return
86 -
87 - # Mask the [0, start_idx) tokens of the prompt with
88 - # PAD_SLOT_ID, where start_idx is max(0, seq_len -
89 - # sliding_window). For example, if the prompt len is 10,
90 - # sliding window is 8, and block size is 4, the first two
91 - # tokens are masked and the slot mapping will be
92 - # [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
93 - padding_mask_len = max(0, start_idx - context_len)
94 - slot_mapping.extend([PAD_SLOT_ID] * padding_mask_len)
95 -
96 - range_start = max(start_idx, context_len)
97 - range_end = seq_len
98 - numel = range_end - range_start
99 - block_table = block_tables[seq_id]
100 -
101 - # numpy implementation will be faster than python if we have
102 - # many elements, otherwise it will be slower.
103 - if numel < _COMPUTE_SLOT_MAPPING_NUMPY_NUMEL:
104 - _compute_slot_mapping_python(slot_mapping, block_table, range_start,
105 - range_end, block_size)
106 - else:
107 - _compute_slot_mapping_numpy(slot_mapping, block_table, range_start,
108 - range_end, block_size)
109 -
110 -
111 -TAttentionMetadata = TypeVar("TAttentionMetadata", bound='AttentionMetadata')
112 -
113 -
114 -class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):
115 -
116 - _metadata_cls: Type[TAttentionMetadata]
117 -
118 - def __init__(self, input_builder):
119 - self.input_builder = input_builder
120 - self.runner = input_builder.runner
121 -
122 - self.sliding_window = input_builder.sliding_window
123 - self.block_size = input_builder.block_size
124 -
125 - def prepare(self):
126 - self.slot_mapping: List[int] = []
127 - self.prefill_seq_lens: List[int] = []
128 - self.context_lens: List[int] = []
129 - self.block_tables: List[List[int]] = []
130 - self.curr_seq_lens: List[int] = []
131 - self.num_prefills = 0
132 - self.num_prefill_tokens = 0
133 - self.num_decode_tokens = 0
134 -
135 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool):
136 - is_prompt = inter_data.is_prompt
137 - block_tables = inter_data.block_tables
138 -
139 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
140 - curr_sliding_window_block) in zip(
141 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
142 - inter_data.orig_seq_lens, inter_data.seq_lens,
143 - inter_data.query_lens, inter_data.context_lens,
144 - inter_data.curr_sliding_window_blocks):
145 - self.context_lens.append(context_len)
146 - if is_prompt:
147 - self.num_prefills += 1
148 - self.num_prefill_tokens += token_len
149 - self.prefill_seq_lens.append(seq_len)
150 - else:
151 - assert query_len == 1, (
152 - "seq_len: {}, context_len: {}, query_len: {}".format(
153 - seq_len, context_len, query_len))
154 - self.num_decode_tokens += query_len
155 - self.curr_seq_lens.append(curr_seq_len)
156 -
157 - # Compute block table.
158 - # TODO(sang): Combine chunked prefill and prefix caching by
159 - # only allowing multiple of block_size chunk size.
160 - # NOTE: This only works for oooooooxxx style attention.
161 - block_table = []
162 - if inter_data.prefix_cache_hit:
163 - block_table = block_tables[seq_id]
164 - elif ((chunked_prefill_enabled or not is_prompt)
165 - and block_tables is not None):
166 - if curr_sliding_window_block == 0:
167 - block_table = block_tables[seq_id]
168 - else:
169 - block_table = block_tables[seq_id][
170 - -curr_sliding_window_block:]
171 - self.block_tables.append(block_table)
172 -
173 - # Compute slot mapping.
174 - is_profile_run = is_block_tables_empty(block_tables)
175 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
176 - context_len,
177 - self.sliding_window)
178 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
179 - seq_len, context_len, start_idx,
180 - self.block_size, inter_data.block_tables)
181 -
182 - def build(self, seq_lens: List[int], query_lens: List[int],
183 - cuda_graph_pad_size: int, batch_size: int):
184 - """Build attention metadata with on-device tensors.
185 -
186 - Args:
187 - seq_lens: The maybe padded sequence lengths of the input sequences.
188 - query_lens: The query lengths of the input sequences.
189 - cuda_graph_pad_size: The padding size for cuda graph.
190 - -1 if cuda graph is not used.
191 - batch_size: The maybe padded batch size.
192 - """
193 - for inter_data in self.input_builder.inter_data_list:
194 - self._add_seq_group(inter_data,
195 - self.input_builder.chunked_prefill_enabled)
196 -
197 - device = self.runner.device
198 - use_captured_graph = cuda_graph_pad_size != -1
199 -
200 - max_query_len = max(query_lens)
201 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
202 - max_decode_seq_len = max(self.curr_seq_lens, default=0)
203 - num_decode_tokens = self.num_decode_tokens
204 - query_start_loc = list(accumulate(query_lens, initial=0))
205 - seq_start_loc = list(accumulate(seq_lens, initial=0))
206 -
207 - if use_captured_graph:
208 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
209 - self.block_tables.extend([] * cuda_graph_pad_size)
210 - num_decode_tokens = batch_size
211 -
212 - # The shape of graph_block_tables is
213 - # [max batch size, max context len // block size].
214 - input_block_tables = self.runner.graph_block_tables[:batch_size]
215 - for i, block_table in enumerate(self.block_tables):
216 - if block_table:
217 - input_block_tables[i, :len(block_table)] = block_table
218 - block_tables = torch.from_numpy(input_block_tables).to(
219 - device, non_blocking=True)
220 - else:
221 - block_tables = make_tensor_with_pad(
222 - self.block_tables,
223 - pad=0,
224 - dtype=torch.int,
225 - device=device,
226 - )
227 - assert max_query_len > 0, "query_lens: {}".format(query_lens)
228 -
229 - assert device is not None
230 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
231 - device, self.runner.pin_memory)
232 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
233 - self.runner.pin_memory)
234 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,
235 - device, self.runner.pin_memory)
236 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
237 - device,
238 - self.runner.pin_memory)
239 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
240 - device, self.runner.pin_memory)
241 -
242 - return self._metadata_cls( # type: ignore
243 - num_prefills=self.num_prefills,
244 - slot_mapping=slot_mapping_tensor,
245 - enable_kv_scales_calculation=True,
246 - num_prefill_tokens=self.num_prefill_tokens,
247 - num_decode_tokens=num_decode_tokens,
248 - seq_lens=seq_lens,
249 - seq_lens_tensor=seq_lens_tensor,
250 - max_query_len=max_query_len,
251 - max_prefill_seq_len=max_prefill_seq_len,
252 - max_decode_seq_len=max_decode_seq_len,
253 - query_start_loc=query_start_loc_tensor,
254 - seq_start_loc=seq_start_loc_tensor,
255 - context_lens_tensor=context_lens_tensor,
256 - block_tables=block_tables,
257 - use_cuda_graph=use_captured_graph,
258 - )
259 -
260 -
261 -class CommonAttentionState(AttentionState):
262 -
263 - def __init__(self, runner):
264 - self.runner = runner
265 - self._is_graph_capturing = False
266 -
267 - @contextmanager
268 - def graph_capture(self, max_batch_size: int):
269 -
270 - self._is_graph_capturing = True
271 -
272 - self._graph_slot_mapping = torch.full((max_batch_size, ),
273 - PAD_SLOT_ID,
274 - dtype=torch.long,
275 - device=self.runner.device)
276 - self._graph_seq_lens = torch.ones(max_batch_size,
277 - dtype=torch.int32,
278 - device=self.runner.device)
279 - self._graph_block_tables = torch.from_numpy(
280 - self.runner.graph_block_tables).to(device=self.runner.device)
281 -
282 - yield
283 -
284 - self._is_graph_capturing = False
285 - del self._graph_slot_mapping
286 - del self._graph_seq_lens
287 - del self._graph_block_tables
288 -
289 - def graph_clone(self, batch_size: int) -> "CommonAttentionState":
290 - assert self._is_graph_capturing
291 - return self.__class__(self.runner)
292 -
293 - def graph_capture_get_metadata_for_batch(
294 - self, batch_size: int, is_encoder_decoder_model: bool = False):
295 - assert self._is_graph_capturing
296 - attn_metadata = self.runner.attn_backend.make_metadata(
297 - num_prefills=0,
298 - num_prefill_tokens=0,
299 - num_decode_tokens=batch_size,
300 - slot_mapping=self._graph_slot_mapping[:batch_size],
301 - enable_kv_scales_calculation=True,
302 - seq_lens=None,
303 - seq_lens_tensor=self._graph_seq_lens[:batch_size],
304 - max_query_len=1,
305 - max_decode_query_len=1,
306 - max_prefill_seq_len=0,
307 - max_decode_seq_len=self.runner.max_model_len,
308 - query_start_loc=None,
309 - seq_start_loc=None,
310 - context_lens_tensor=None,
311 - block_tables=self._graph_block_tables[:batch_size],
312 - use_cuda_graph=True,
313 - )
314 - if is_encoder_decoder_model:
315 - # The encoder decoder model works only with XFormers and
316 - # Flash Attention backend. Assert the same.
317 - assert self.runner.attn_backend.get_name() in \
318 - ["XFORMERS", "FLASH_ATTN"], \
319 - f"Expected attn_backend name to be either 'XFORMERS' or " \
320 - f"'FLASH_ATTN', but got '{self.runner.attn_backend.get_name()}'"
321 - self._update_captured_metadata_for_enc_dec_model(
322 - batch_size=batch_size, attn_metadata=attn_metadata)
323 -
324 - return attn_metadata
325 -
326 - def get_graph_input_buffers(
327 - self,
328 - attn_metadata,
329 - is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
330 - input_buffers = {
331 - "slot_mapping": attn_metadata.slot_mapping,
332 - "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,
333 - "block_tables": attn_metadata.decode_metadata.block_tables,
334 - }
335 - if is_encoder_decoder_model:
336 - # The encoder decoder model works only with XFormers and
337 - # Flash Attention backend. Assert the same.
338 - assert self.runner.attn_backend.get_name() in \
339 - ["XFORMERS", "FLASH_ATTN"], \
340 - f"Expected attn_backend name to be either 'XFORMERS' or " \
341 - f"'FLASH_ATTN', but got '{self.runner.attn_backend.get_name()}'"
342 - self._add_additional_input_buffers_for_enc_dec_model(
343 - attn_metadata=attn_metadata, input_buffers=input_buffers)
344 - return input_buffers
345 -
346 - def prepare_graph_input_buffers(
347 - self,
348 - input_buffers,
349 - attn_metadata,
350 - is_encoder_decoder_model: bool = False) -> None:
351 - input_buffers["seq_lens_tensor"].copy_(
352 - attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)
353 - input_buffers["block_tables"].copy_(
354 - attn_metadata.decode_metadata.block_tables, non_blocking=True)
355 - if is_encoder_decoder_model:
356 - # The encoder decoder model works only with XFormers and
357 - # Flash Attention backend. Assert the same.
358 - assert self.runner.attn_backend.get_name() in\
359 - ["XFORMERS", "FLASH_ATTN"], \
360 - f"Expected attn_backend name to be either 'XFORMERS' or "\
361 - f"'FLASH_ATTN', but "\
362 - f"got '{self.runner.attn_backend.get_name()}'"
363 - self._prepare_input_buffers_for_enc_dec_model(
364 - attn_metadata, input_buffers)
365 -
366 - def begin_forward(self, model_input) -> None:
367 - return
368 -
369 - def _update_captured_metadata_for_enc_dec_model(self, batch_size: int,
370 - attn_metadata):
371 - """
372 - Updates the attention metadata parameters for CUDA graph capture in an
373 - encoder-decoder model.
374 -
375 - This method modifies attention-related tensors and metadata required
376 - for CUDA graph capture in encoder-decoder models. Specifically, it
377 - updates the cross-attention and encoder sequence tensors in the
378 - AttentionMetadata object.
379 - """
380 - # During decode phase the cross_slot_mapping will be empty. Hence set
381 - # an empty tensor for CUDA Graph capture.
382 - attn_metadata.cross_slot_mapping = torch.tensor(
383 - [], dtype=torch.int).cuda()
384 - attn_metadata.cross_block_tables = torch.full(
385 - (batch_size, self.runner.get_max_block_per_batch()),
386 - 1,
387 - dtype=torch.int).cuda()
388 - attn_metadata.encoder_seq_lens = torch.full((batch_size, ),
389 - 1,
390 - dtype=torch.int).cuda()
391 - attn_metadata.encoder_seq_lens_tensor = torch.full(
392 - (batch_size, ), 1, dtype=torch.int).cuda()
393 - attn_metadata.max_encoder_seq_len = self.runner.max_model_len
394 - attn_metadata.num_encoder_tokens = 0
395 -
396 - def _add_additional_input_buffers_for_enc_dec_model(
397 - self, attn_metadata, input_buffers: Dict[str, Any]):
398 - """
399 - Saves additional input buffers specific to the encoder-decoder model
400 - from the attention metadata.
401 -
402 - This method extracts and stores encoder-decoder related input buffers
403 - from the `attn_metadata` into the `input_buffers` dictionary. The
404 - buffers include encoder sequence lengths, cross-slot mappings, and
405 - cross-block tables, which are essential for the encoder-decoder model
406 - during CUDA graph replay.
407 - """
408 - input_buffers["encoder_seq_lens_tensor"] = (
409 - attn_metadata.decode_metadata.encoder_seq_lens_tensor)
410 - input_buffers["cross_slot_mapping"] = (
411 - attn_metadata.decode_metadata.cross_slot_mapping)
412 - input_buffers["cross_block_tables"] = (
413 - attn_metadata.decode_metadata.cross_block_tables)
414 -
415 - def _prepare_input_buffers_for_enc_dec_model(self, attn_metadata,
416 - input_buffers: Dict[str,
417 - Any]):
418 - """
419 - Populates input buffers with data from the encoder-decoder model's
420 - attention metadata.
421 -
422 - This method fills the input buffers with encoder-decoder specific
423 - tensors. It copies data from the `attn_metadata` and keyword arguments
424 - (`kwargs`) into corresponding buffers in the `input_buffers` dictionary.
425 - The copied data includes attention-related metadata as well as input
426 - IDs and positional information for the encoder.
427 - """
428 - input_buffers["encoder_seq_lens_tensor"].copy_(
429 - attn_metadata.decode_metadata.encoder_seq_lens_tensor,
430 - non_blocking=True)
431 - input_buffers["cross_slot_mapping"].copy_(
432 - attn_metadata.decode_metadata.cross_slot_mapping,
433 - non_blocking=True)
434 - input_buffers["cross_block_tables"].copy_(
435 - attn_metadata.decode_metadata.cross_block_tables,
436 - non_blocking=True)
437 -
438 -
439 -def is_all_encoder_attn_metadata_set(attn_metadata):
440 - '''
441 - All attention metadata required for encoder attention is set.
442 - '''
443 - return ((attn_metadata.encoder_seq_lens is not None)
444 - and (attn_metadata.encoder_seq_lens_tensor is not None)
445 - and (attn_metadata.max_encoder_seq_len is not None))
446 -
447 -
448 -def is_all_cross_attn_metadata_set(attn_metadata):
449 - '''
450 - All attention metadata required for enc/dec cross-attention is set.
451 -
452 - Superset of encoder attention required metadata.
453 - '''
454 - return (attn_metadata.is_all_encoder_attn_metadata_set
455 - and (attn_metadata.cross_slot_mapping is not None)
456 - and (attn_metadata.cross_block_tables is not None))
457 -
458 -
459 -def get_seq_len_block_table_args(
460 - attn_metadata,
461 - is_prompt: bool,
462 - attn_type: str,
463 -) -> tuple:
464 - '''
465 - The particular choice of sequence-length- and block-table-related
466 - attributes which should be extracted from attn_metadata is dependent
467 - on the type of attention operation.
468 -
469 - Decoder attn -> select entirely decoder self-attention-related fields
470 - Encoder/decoder cross-attn -> select encoder sequence lengths &
471 - cross-attn block-tables fields
472 - Encoder attn -> select encoder sequence lengths fields & no block tables
473 -
474 - Arguments:
475 -
476 - * attn_metadata: Attention metadata structure associated with attention op
477 - * is_prompt: True if prefill, False otherwise
478 - * attn_type: encoder attention, decoder self-attention,
479 - encoder/decoder cross-attention
480 -
481 - Returns:
482 -
483 - * Appropriate sequence-lengths tensor
484 - * Appropriate max sequence-length scalar
485 - * Appropriate block tables (or None)
486 - '''
487 -
488 - if attn_type == AttentionType.DECODER:
489 - # Decoder self-attention
490 - # Choose max_seq_len based on whether we are in prompt_run
491 - if is_prompt:
492 - max_seq_len = attn_metadata.max_prefill_seq_len
493 - else:
494 - max_seq_len = attn_metadata.max_decode_seq_len
495 - return (attn_metadata.seq_lens_tensor, max_seq_len,
496 - attn_metadata.block_tables)
497 - elif attn_type == AttentionType.ENCODER_DECODER:
498 - # Enc/dec cross-attention KVs match encoder sequence length;
499 - # cross-attention utilizes special "cross" block tables
500 - return (attn_metadata.encoder_seq_lens_tensor,
501 - attn_metadata.max_encoder_seq_len,
502 - attn_metadata.cross_block_tables)
503 - elif attn_type == AttentionType.ENCODER:
504 - # No block tables associated with encoder attention
505 - return (attn_metadata.encoder_seq_lens_tensor,
506 - attn_metadata.max_encoder_seq_len, None)
507 - else:
508 - raise AttributeError(f"Invalid attention type {str(attn_type)}")
509 -
510 -
511 -def get_num_prefill_decode_query_kv_tokens(
512 - attn_metadata,
513 - attn_type: str,
514 -) -> Tuple[int, int, int]:
515 - """
516 - Calculate the number of prefill and decode tokens for query, key/value
517 - based on the attention metadata and the specified attention type.
518 -
519 - Args:
520 - attn_metadata (AttentionMetadata): Attention Metadata object.
521 - attn_type (AttentionType): The type of attention being used.
522 - Returns:
523 - Tuple[int, int, int]: A tuple containing three integers:
524 - - The number of prefill query tokens.
525 - - The number of prefill key/value tokens.
526 - - The number of decode query tokens.
527 -
528 - Raises:
529 - AssertionError: If the number of encoder tokens in `attn_metadata`
530 - is `None` when required for the calculations.
531 - """
532 - num_prefill_query_tokens = 0
533 - num_decode_query_tokens = 0
534 - num_prefill_kv_tokens = 0
535 - if attn_type == AttentionType.ENCODER:
536 - # Encoder attention is only invoked during prefill phase.
537 - # The same input servers a both query and key.
538 - assert attn_metadata.num_encoder_tokens is not None
539 - num_prefill_query_tokens = attn_metadata.num_encoder_tokens
540 - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens
541 - num_decode_query_tokens = 0
542 - elif attn_type == AttentionType.ENCODER_DECODER:
543 - assert attn_metadata.num_encoder_tokens is not None
544 - num_prefill_query_tokens = attn_metadata.num_prefill_tokens
545 - # The key is the encoder/cross-attention.
546 - num_prefill_kv_tokens = attn_metadata.num_encoder_tokens
547 - num_decode_query_tokens = attn_metadata.num_decode_tokens
548 - else: # attn_type == AttentionType.DECODER or
549 - # attn_type == AttentionType.ENCODER_ONLY
550 - num_prefill_query_tokens = attn_metadata.num_prefill_tokens
551 - num_prefill_kv_tokens = attn_metadata.num_prefill_tokens
552 - num_decode_query_tokens = attn_metadata.num_decode_tokens
553 -
554 - return (num_prefill_query_tokens, num_prefill_kv_tokens,
555 - num_decode_query_tokens)
556 -
14
15 @dataclass
16 class MLADims:
vllm/v1/attention/backends/cpu_attn.py
-15
@@ -11,7 +11,6 @@ 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
14 from vllm.config import VllmConfig
15 from vllm.logger import init_logger
16 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
@@ -65,10 +64,6 @@ class TorchSDPABackend(AttentionBackend):
64 def get_metadata_cls() -> type["AttentionMetadata"]:
65 return TorchSDPAMetadata
66
68 - @staticmethod
69 - def get_state_cls() -> type["CommonAttentionState"]:
70 - return CommonAttentionState
71 -
67 @staticmethod
68 def get_builder_cls() -> type["TorchSDPAMetadataBuilderV1"]:
69 return TorchSDPAMetadataBuilderV1
@@ -835,16 +830,6 @@ class _PagedAttention:
830 blocksparse_head_sliding_step,
831 )
832
838 - @staticmethod
839 - def copy_blocks(
840 - kv_caches: list[torch.Tensor],
841 - src_to_dists: torch.Tensor,
842 - *args,
843 - ) -> None:
844 - key_caches = [kv_cache[0] for kv_cache in kv_caches]
845 - value_caches = [kv_cache[1] for kv_cache in kv_caches]
846 - ops.copy_blocks(key_caches, value_caches, src_to_dists)
847 -
833
834 class _IPEXPagedAttention(_PagedAttention):
835
vllm/v1/attention/backends/pallas.py
-5
@@ -8,7 +8,6 @@ import torch
8
9 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10 AttentionLayer, AttentionType)
11 -from vllm.attention.backends.utils import CommonAttentionState
11 from vllm.config import VllmConfig
12 from vllm.logger import init_logger
13 from vllm.utils import cdiv, next_power_of_2
@@ -97,10 +96,6 @@ class PallasAttentionBackend(AttentionBackend):
96 def get_metadata_cls() -> type["PallasMetadata"]:
97 return PallasMetadata
98
100 - @staticmethod
101 - def get_state_cls() -> type["CommonAttentionState"]:
102 - return CommonAttentionState
103 -
99 @staticmethod
100 def get_kv_cache_shape(
101 num_blocks: int,
vllm/v1/spec_decode/eagle.py
+6 -5
@@ -9,7 +9,6 @@ import numpy as np
9 import torch
10 import torch.nn as nn
11
12 -from vllm.attention.backends.abstract import AttentionMetadataBuilder
12 from vllm.attention.layer import Attention
13 from vllm.config import (CompilationLevel, VllmConfig,
14 get_layers_from_vllm_config)
@@ -25,7 +24,8 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
24 from vllm.v1.attention.backends.tree_attn import (TreeAttentionMetadata,
25 TreeAttentionMetadataBuilder)
26 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
28 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
27 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
28 + CommonAttentionMetadata)
29 from vllm.v1.kv_cache_interface import KVCacheConfig
30 from vllm.v1.sample.metadata import SamplingMetadata
31 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
@@ -184,8 +184,9 @@ class EagleProposer:
184 builder = (self._get_attention_metadata_builder()
185 if self.attn_metadata_builder is None else
186 self.attn_metadata_builder)
187 - attn_metadata = builder.build_for_drafting(
188 - common_attn_metadata=common_attn_metadata, draft_index=0)
187 + attn_metadata = builder.build_for_drafting( # type: ignore
188 + common_attn_metadata=common_attn_metadata,
189 + draft_index=0)
190
191 # At this moment, we assume all eagle layers belong to the same KV
192 # cache group, thus using the same attention metadata.
@@ -319,7 +320,7 @@ class EagleProposer:
320 exceeds_max_model_len, PADDING_SLOT_ID)
321
322 # Rebuild attention metadata
322 - attn_metadata = builder.build_for_drafting(
323 + attn_metadata = builder.build_for_drafting( # type: ignore
324 common_attn_metadata=common_attn_metadata,
325 draft_index=token_index + 1)
326 for layer_name in self.attn_layer_names: