@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 969b4da3a6
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 placeholder attn (#25510)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>

Thomas Parnell committed Sep 24, 2025 at 00:12 UTC 969b4da3a6ab737f72cb33db502b4c0bb70d4139
5 files changed +10 -354
tests/kernels/attention/test_attention_selector.py
+10 -27
@@ -85,8 +85,7 @@ def test_env(
85 if device == "cpu":
86 with patch("vllm.attention.selector.current_platform",
87 CpuPlatform()):
88 - backend = get_attn_backend(16, torch.float16, None, block_size,
89 - False)
88 + backend = get_attn_backend(16, torch.float16, None, block_size)
89 assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
90
91 elif device == "hip":
@@ -106,7 +105,6 @@ def test_env(
105 torch.float16,
106 None,
107 block_size,
109 - False,
108 use_mla=use_mla)
109 assert f"The selected backend, {name}" in str(
110 exc_info.value)
@@ -117,7 +115,6 @@ def test_env(
115 torch.float16,
116 None,
117 block_size,
120 - False,
118 use_mla=use_mla)
119 assert f"The selected backend, {name}" in str(
120 exc_info.value)
@@ -127,7 +124,6 @@ def test_env(
124 torch.float16,
125 None,
126 block_size,
130 - False,
127 use_mla=use_mla)
128 expected = f"{name}_VLLM_V1"
129 assert backend.get_name() == expected
@@ -136,7 +132,6 @@ def test_env(
132 torch.float16,
133 None,
134 block_size,
139 - False,
135 use_mla=use_mla)
136 expected = "TRITON_ATTN_VLLM_V1"
137 assert backend.get_name() == expected
@@ -164,7 +159,6 @@ def test_env(
159 torch.float16,
160 None,
161 block_size,
167 - False,
162 use_mla=use_mla)
163 expected = "CUTLASS_MLA_VLLM_V1"
164 assert backend.get_name() == expected
@@ -179,7 +173,6 @@ def test_env(
173 torch.float16,
174 None,
175 block_size,
182 - False,
176 use_mla=use_mla)
177 expected = "FLASHINFER_MLA"
178 assert backend.get_name() == expected
@@ -199,7 +192,6 @@ def test_env(
192 torch.float16,
193 None,
194 block_size,
202 - False,
195 use_mla=use_mla)
196 expected = f"{name}_VLLM_V1"
197 assert backend.get_name() == expected
@@ -208,7 +200,6 @@ def test_env(
200 torch.float16,
201 None,
202 block_size,
211 - False,
203 use_mla=use_mla)
204 expected = "FLASH_ATTN_MLA"
205 assert backend.get_name() == expected
@@ -218,7 +209,6 @@ def test_env(
209 torch.float16,
210 None,
211 block_size,
221 - False,
212 use_mla=use_mla)
213 expected = "TRITON_MLA_VLLM_V1"
214 assert backend.get_name() == expected
@@ -227,7 +217,6 @@ def test_env(
217 torch.float16,
218 None,
219 block_size,
230 - False,
220 use_mla=use_mla)
221 expected = "FLASHINFER_VLLM_V1"
222 assert backend.get_name() == expected
@@ -236,7 +225,6 @@ def test_env(
225 torch.float16,
226 None,
227 block_size,
239 - False,
228 use_mla=use_mla)
229 expected = "FLASH_ATTN_VLLM_V1"
230 assert backend.get_name() == expected
@@ -245,7 +233,6 @@ def test_env(
233 torch.float16,
234 None,
235 block_size,
248 - False,
236 use_mla=use_mla)
237 assert backend.get_name() == "FLEX_ATTENTION", (
238 "Should fallback to FlexAttention if head size is "
@@ -264,13 +251,13 @@ def test_fp32_fallback(
251 if device == "cpu":
252 with patch("vllm.attention.selector.current_platform",
253 CpuPlatform()):
267 - backend = get_attn_backend(16, torch.float32, None, 16, False)
254 + backend = get_attn_backend(16, torch.float32, None, 16)
255 assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
256
257 elif device == "cuda":
258 with patch("vllm.attention.selector.current_platform",
259 CudaPlatform()):
273 - backend = get_attn_backend(16, torch.float32, None, 16, False)
260 + backend = get_attn_backend(16, torch.float32, None, 16)
261 assert backend.get_name() == "FLEX_ATTENTION"
262
263
@@ -286,29 +273,29 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
273 monkeypatch.setattr(torch.cuda,
274 "get_device_capability",
275 lambda _=None: (7, 5))
289 - backend = get_attn_backend(16, torch.float16, None, 16, False)
276 + backend = get_attn_backend(16, torch.float16, None, 16)
277 assert backend.get_name() != STR_FLASH_ATTN_VAL
278
279 # Reset the monkeypatch for subsequent tests
280 monkeypatch.undo()
281
282 # Unsupported data type
296 - backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16, False)
283 + backend = get_attn_backend(16, torch.float8_e4m3fn, None, 16)
284 assert backend.get_name() != STR_FLASH_ATTN_VAL
285
286 # Unsupported kv cache data type
300 - backend = get_attn_backend(16, torch.float16, "fp8", 16, False)
287 + backend = get_attn_backend(16, torch.float16, "fp8", 16)
288 assert backend.get_name() != STR_FLASH_ATTN_VAL
289
290 # Unsupported block size
304 - backend = get_attn_backend(16, torch.float16, None, 8, False)
291 + backend = get_attn_backend(16, torch.float16, None, 8)
292 assert backend.get_name() != STR_FLASH_ATTN_VAL
293
294 # flash-attn is not installed
295 import sys
296 original_module = sys.modules.get('vllm_flash_attn')
297 monkeypatch.setitem(sys.modules, 'vllm_flash_attn', None)
311 - backend = get_attn_backend(16, torch.float16, None, 16, False)
298 + backend = get_attn_backend(16, torch.float16, None, 16)
299 assert backend.get_name() != STR_FLASH_ATTN_VAL
300
301 # Restore the original module if it existed
@@ -319,11 +306,7 @@ def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
306 monkeypatch.delitem(sys.modules, 'vllm_flash_attn', raising=False)
307
308 # Unsupported head size
322 - backend = get_attn_backend(17, torch.float16, None, 16, False)
323 - assert backend.get_name() != STR_FLASH_ATTN_VAL
324 -
325 - # Attention-free models should bypass env and use PlaceholderAttention
326 - backend = get_attn_backend(16, torch.float16, None, 16, True)
309 + backend = get_attn_backend(17, torch.float16, None, 16)
310 assert backend.get_name() != STR_FLASH_ATTN_VAL
311
312
@@ -336,5 +319,5 @@ def test_invalid_env(monkeypatch: pytest.MonkeyPatch):
319
320 # Should raise ValueError for invalid backend
321 with pytest.raises(ValueError) as exc_info:
339 - get_attn_backend(32, torch.float16, None, 16, False)
322 + get_attn_backend(32, torch.float16, None, 16)
323 assert "Invalid value 'INVALID'" in str(exc_info.value)
vllm/attention/backends/placeholder_attn.py deleted
-314
@@ -1,314 +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 itertools import accumulate
6 -from typing import List, Optional, Tuple, Type
7 -
8 -import torch
9 -
10 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11 - AttentionMetadata,
12 - AttentionMetadataBuilder)
13 -from vllm.attention.backends.utils import CommonAttentionState
14 -from vllm.utils import async_tensor_h2d
15 -
16 -# Placeholder attention backend for models like Mamba and pooling models that
17 -# lack attention.
18 -
19 -
20 -class PlaceholderAttentionBackend(AttentionBackend):
21 - """Placeholder backend for when no attention is needed."""
22 -
23 - @staticmethod
24 - def get_name() -> str:
25 - return "NO_ATTENTION"
26 -
27 - @staticmethod
28 - def get_impl_cls() -> Type["PlaceholderAttentionImpl"]:
29 - return PlaceholderAttentionImpl
30 -
31 - @staticmethod
32 - def get_builder_cls() -> Type["PlaceholderAttentionMetadataBuilder"]:
33 - return PlaceholderAttentionMetadataBuilder
34 -
35 - @staticmethod
36 - def get_metadata_cls() -> Type["PlaceholderAttentionMetadata"]:
37 - return PlaceholderAttentionMetadata
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 (1, 1, 1, 1, 1)
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 - return
59 -
60 - @staticmethod
61 - def copy_blocks(
62 - kv_caches: List[torch.Tensor],
63 - src_to_dists: torch.Tensor,
64 - ) -> None:
65 - return
66 -
67 -
68 -@dataclass
69 -class PlaceholderAttentionMetadata(AttentionMetadata):
70 - """Attention metadata for prefill and decode batched together."""
71 - # (batch_size,). The sequence length per sequence. Sequence length means
72 - # the computed tokens + new tokens None if it is a decoding.
73 - seq_lens: Optional[List[int]]
74 - # seq_lens stored as a tensor.
75 - seq_lens_tensor: Optional[torch.Tensor]
76 -
77 - # Maximum sequence length among prefill batch. 0 if there are decoding
78 - # requests only.
79 - max_prefill_seq_len: int
80 - # Maximum sequence length among decode batch. 0 if there are prefill
81 - # requests only.
82 - max_decode_seq_len: int
83 - # (batch_size,) A tensor of context lengths (tokens that are computed
84 - # so far).
85 - context_lens_tensor: Optional[torch.Tensor]
86 -
87 - # Whether or not if cuda graph is enabled.
88 - # Cuda-graph is currently enabled for decoding only.
89 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
90 - use_cuda_graph: bool
91 -
92 - # Maximum query length in the batch.
93 - max_query_len: Optional[int]
94 -
95 - # Max number of query tokens among request in the batch.
96 - max_decode_query_len: Optional[int]
97 -
98 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in
99 - # the batch, used to index into subquery. E.g., if the subquery length
100 - # is [4, 6], it is [0, 4, 10].
101 - query_start_loc: Optional[torch.Tensor] = None
102 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in
103 - # the batch, used to index into sequence. E.g., if the sequence length is
104 - # [4, 6], it is [0, 4, 10].
105 - seq_start_loc: Optional[torch.Tensor] = None
106 -
107 - # Placeholder.
108 - block_tables: Optional[torch.Tensor] = None
109 -
110 - _cached_prefill_metadata: Optional["PlaceholderAttentionMetadata"] = None
111 - _cached_decode_metadata: Optional["PlaceholderAttentionMetadata"] = None
112 -
113 - @property
114 - def prefill_metadata(self) -> Optional["PlaceholderAttentionMetadata"]:
115 - if self.num_prefills == 0:
116 - return None
117 -
118 - if self._cached_prefill_metadata is not None:
119 - return self._cached_prefill_metadata
120 -
121 - # Compute some attn_metadata fields which default to None
122 - query_start_loc = (None if self.query_start_loc is None else
123 - self.query_start_loc[:self.num_prefills + 1])
124 - seq_lens = (None if self.seq_lens is None else
125 - self.seq_lens[:self.num_prefills])
126 - seq_lens_tensor = (None if self.seq_lens_tensor is None else
127 - self.seq_lens_tensor[:self.num_prefills])
128 - seq_start_loc = (None if self.seq_start_loc is None else
129 - self.seq_start_loc[:self.num_prefills + 1])
130 - context_lens_tensor = (None if self.context_lens_tensor is None else
131 - self.context_lens_tensor[:self.num_prefills])
132 -
133 - # Placeholders
134 - slot_mapping = torch.empty(0)
135 - block_tables = torch.empty(0)
136 -
137 - self._cached_prefill_metadata = PlaceholderAttentionMetadata(
138 - num_prefills=self.num_prefills,
139 - num_prefill_tokens=self.num_prefill_tokens,
140 - num_decode_tokens=0,
141 - slot_mapping=slot_mapping,
142 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,
143 - seq_lens=seq_lens,
144 - seq_lens_tensor=seq_lens_tensor,
145 - max_decode_query_len=0,
146 - max_query_len=self.max_query_len,
147 - max_prefill_seq_len=self.max_prefill_seq_len,
148 - max_decode_seq_len=0,
149 - query_start_loc=query_start_loc,
150 - seq_start_loc=seq_start_loc,
151 - context_lens_tensor=context_lens_tensor,
152 - block_tables=block_tables,
153 - use_cuda_graph=False,
154 - )
155 - return self._cached_prefill_metadata
156 -
157 - @property
158 - def decode_metadata(self) -> Optional["PlaceholderAttentionMetadata"]:
159 - if self.num_decode_tokens == 0:
160 - return None
161 -
162 - if self._cached_decode_metadata is not None:
163 - return self._cached_decode_metadata
164 - assert self.seq_lens_tensor is not None
165 -
166 - # Placeholders
167 - slot_mapping = torch.empty(0)
168 - block_tables = torch.empty(0)
169 - seq_lens_tensor = (None if self.seq_lens_tensor is None else
170 - self.seq_lens_tensor[self.num_prefills:])
171 -
172 - self._cached_decode_metadata = PlaceholderAttentionMetadata(
173 - num_prefills=0,
174 - num_prefill_tokens=0,
175 - num_decode_tokens=self.num_decode_tokens,
176 - slot_mapping=slot_mapping,
177 - enable_kv_scales_calculation=True,
178 - seq_lens=None,
179 - seq_lens_tensor=seq_lens_tensor,
180 - max_decode_query_len=self.max_decode_query_len,
181 - max_query_len=None,
182 - max_prefill_seq_len=0,
183 - max_decode_seq_len=self.max_decode_seq_len,
184 - query_start_loc=(self.query_start_loc[self.num_prefills:] -
185 - self.query_start_loc[self.num_prefills])
186 - if self.query_start_loc is not None else None,
187 - seq_start_loc=self.seq_start_loc[self.num_prefills:]
188 - if self.seq_start_loc is not None else None,
189 - context_lens_tensor=None,
190 - block_tables=block_tables,
191 - use_cuda_graph=self.use_cuda_graph,
192 - )
193 - return self._cached_decode_metadata
194 -
195 -
196 -class PlaceholderAttentionMetadataBuilder(
197 - AttentionMetadataBuilder[PlaceholderAttentionMetadata]):
198 -
199 - def __init__(self, input_builder):
200 -
201 - self.input_builder = input_builder
202 - self.runner = input_builder.runner
203 -
204 - def prepare(self):
205 - self.prefill_seq_lens: List[int] = []
206 - self.context_lens: List[int] = []
207 - self.curr_seq_lens: List[int] = []
208 - self.num_prefills = 0
209 - self.num_prefill_tokens = 0
210 - self.num_decode_tokens = 0
211 -
212 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool):
213 - """Add a sequence group to the metadata. Specifically update/append
214 - 1. context length.
215 - """
216 - is_prompt = inter_data.is_prompt
217 -
218 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
219 - curr_sliding_window_block) in zip(
220 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
221 - inter_data.orig_seq_lens, inter_data.seq_lens,
222 - inter_data.query_lens, inter_data.context_lens,
223 - inter_data.curr_sliding_window_blocks):
224 - self.context_lens.append(context_len)
225 -
226 - if is_prompt:
227 - self.num_prefills += 1
228 - self.num_prefill_tokens += token_len
229 - self.prefill_seq_lens.append(seq_len)
230 - else:
231 - self.num_decode_tokens += query_len
232 - self.curr_seq_lens.append(curr_seq_len)
233 -
234 - def build(self, seq_lens: List[int], query_lens: List[int],
235 - cuda_graph_pad_size: int, batch_size: int):
236 - """Build attention metadata with on-device tensors.
237 -
238 - Args:
239 - seq_lens: The maybe padded sequence lengths of the input sequences.
240 - query_lens: The query lengths of the input sequences.
241 - cuda_graph_pad_size: The padding size for cuda graph.
242 - -1 if cuda graph is not used.
243 - batch_size: The maybe padded batch size.
244 - """
245 -
246 - # Some input builders such as ModelInputForCPUBuilder do not have the
247 - # "inter_data_list" attribute.
248 - # Let's check inter_data_list exists before we reference it.
249 - if hasattr(self.input_builder, "inter_data_list"):
250 - for inter_data in self.input_builder.inter_data_list:
251 - self._add_seq_group(inter_data,
252 - self.input_builder.chunked_prefill_enabled)
253 -
254 - device = self.runner.device
255 - use_captured_graph = cuda_graph_pad_size != -1
256 -
257 - max_query_len = max(query_lens)
258 - decode_query_lens = query_lens[self.num_prefills:]
259 - if len(decode_query_lens) > 0:
260 - max_decode_query_len = max(decode_query_lens)
261 - else:
262 - max_decode_query_len = 1
263 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
264 - max_decode_seq_len = max(self.curr_seq_lens, default=0)
265 - num_decode_tokens = self.num_decode_tokens
266 - query_start_loc = list(accumulate(query_lens, initial=0))
267 - seq_start_loc = list(accumulate(seq_lens, initial=0))
268 -
269 - if use_captured_graph:
270 - num_decode_tokens = batch_size - self.num_prefill_tokens
271 - assert max_query_len > 0, ("query_lens: {}".format(query_lens))
272 -
273 - assert device is not None
274 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
275 - device, self.runner.pin_memory)
276 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
277 - self.runner.pin_memory)
278 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
279 - device,
280 - self.runner.pin_memory)
281 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
282 - device, self.runner.pin_memory)
283 -
284 - # Placeholders
285 - slot_mapping_tensor = torch.empty(0)
286 - block_tables = torch.empty(0)
287 -
288 - return PlaceholderAttentionMetadata(
289 - num_prefills=self.num_prefills,
290 - slot_mapping=slot_mapping_tensor,
291 - enable_kv_scales_calculation=True,
292 - num_prefill_tokens=self.num_prefill_tokens,
293 - num_decode_tokens=num_decode_tokens,
294 - seq_lens=seq_lens,
295 - seq_lens_tensor=seq_lens_tensor,
296 - max_query_len=max_query_len,
297 - max_decode_query_len=max_decode_query_len,
298 - max_prefill_seq_len=max_prefill_seq_len,
299 - max_decode_seq_len=max_decode_seq_len,
300 - query_start_loc=query_start_loc_tensor,
301 - seq_start_loc=seq_start_loc_tensor,
302 - context_lens_tensor=context_lens_tensor,
303 - block_tables=block_tables,
304 - use_cuda_graph=use_captured_graph,
305 - )
306 -
307 -
308 -class PlaceholderAttentionImpl(AttentionImpl):
309 -
310 - def __init__(self, *args, **kwargs) -> None:
311 - return
312 -
313 - def forward(self, *args, **kwargs) -> torch.Tensor:
314 - raise NotImplementedError
vllm/attention/layer.py
-3
@@ -115,12 +115,10 @@ class Attention(nn.Module, AttentionLayerBase):
115 if cache_config is not None:
116 kv_cache_dtype = cache_config.cache_dtype
117 block_size = cache_config.block_size
118 - is_attention_free = cache_config.is_attention_free
118 calculate_kv_scales = cache_config.calculate_kv_scales
119 else:
120 kv_cache_dtype = "auto"
121 block_size = 16
123 - is_attention_free = False
122 calculate_kv_scales = False
123 if num_kv_heads is None:
124 num_kv_heads = num_heads
@@ -185,7 +183,6 @@ class Attention(nn.Module, AttentionLayerBase):
183 dtype,
184 kv_cache_dtype,
185 block_size,
188 - is_attention_free,
186 use_mla=use_mla,
187 has_sink=self.has_sink)
188 else:
vllm/attention/selector.py
-9
@@ -142,7 +142,6 @@ def get_attn_backend(
142 dtype: torch.dtype,
143 kv_cache_dtype: Optional[str],
144 block_size: int,
145 - is_attention_free: bool = False,
145 use_mla: bool = False,
146 has_sink: bool = False,
147 ) -> type[AttentionBackend]:
@@ -156,7 +155,6 @@ def get_attn_backend(
155 dtype=dtype,
156 kv_cache_dtype=kv_cache_dtype,
157 block_size=block_size,
159 - is_attention_free=is_attention_free,
158 use_v1=envs.VLLM_USE_V1,
159 use_mla=use_mla,
160 has_sink=has_sink,
@@ -169,17 +167,10 @@ def _cached_get_attn_backend(
167 dtype: torch.dtype,
168 kv_cache_dtype: Optional[str],
169 block_size: int,
172 - is_attention_free: bool,
170 use_v1: bool = False,
171 use_mla: bool = False,
172 has_sink: bool = False,
173 ) -> type[AttentionBackend]:
177 - # If there are no attention layers (e.g. we are running Mamba),
178 - # use the placeholder NO_ATTENTION
179 - if is_attention_free:
180 - from vllm.attention.backends.placeholder_attn import (
181 - PlaceholderAttentionBackend)
182 - return PlaceholderAttentionBackend
174
175 # Check whether a particular choice of backend was
176 # previously forced.
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-1
@@ -574,7 +574,6 @@ class NixlConnectorWorker:
574 self.model_config.dtype,
575 self.cache_config.cache_dtype,
576 self.block_size,
577 - self.model_config.is_attention_free,
577 use_mla=self.use_mla)
578 self.backend_name = backend.get_name()
579 attn_backend = backend_name_to_enum(self.backend_name)