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
[V1] port xformers backend to v1 (#21342)
Signed-off-by: Giancarlo Delfin <gdelfin@meta.com>
Giancarlo Delfin committed
Aug 5, 2025 at 10:04 UTC
469b3ffaaadb6ab15ddbebc47ac11a0c6fddfda2
6 files changed
+438
-1
tests/v1/attention/utils.py
+2
@@ -128,6 +128,8 @@ def get_attention_backend(backend_name: _Backend):
128
"vllm.v1.attention.backends.triton_attn.TritonAttentionBackend",
129
_Backend.TREE_ATTN:
130
"vllm.v1.attention.backends.tree_attn.TreeAttentionBackend",
131
+ _Backend.XFORMERS_VLLM_V1:
132
+ "vllm.v1.attention.backends.xformers.XFormersAttentionBackend",
133
}
134
135
if backend_name not in backend_map:
vllm/engine/arg_utils.py
+1
@@ -1469,6 +1469,7 @@ class EngineArgs:
1469
"TORCH_SDPA_VLLM_V1",
1470
"FLEX_ATTENTION",
1471
"TREE_ATTN",
1472
+ "XFORMERS_VLLM_V1",
1473
]
1474
if (envs.is_set("VLLM_ATTENTION_BACKEND")
1475
and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/cuda.py
+4
@@ -271,6 +271,7 @@ class CudaPlatformBase(Platform):
271
TRITON_ATTN_VLLM_V1 = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501
272
FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
273
TREE_ATTN_V1 = "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend" # noqa: E501
274
+ XFORMERS_V1 = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend" # noqa: E501
275
276
if selected_backend == _Backend.FLASHINFER:
277
logger.info_once("Using FlashInfer backend on V1 engine.")
@@ -291,6 +292,9 @@ class CudaPlatformBase(Platform):
292
elif selected_backend == _Backend.TREE_ATTN:
293
logger.info_once("Using Tree Attention backend on V1 engine.")
294
return TREE_ATTN_V1
295
+ elif selected_backend == _Backend.XFORMERS_VLLM_V1:
296
+ logger.info_once("Using XFormers backend on V1 engine.")
297
+ return XFORMERS_V1
298
299
from vllm.attention.selector import is_attn_backend_supported
300
vllm/platforms/interface.py
+1
@@ -63,6 +63,7 @@ class _Backend(enum.Enum):
63
NO_ATTENTION = enum.auto()
64
FLEX_ATTENTION = enum.auto()
65
TREE_ATTN = enum.auto()
66
+ XFORMERS_VLLM_V1 = enum.auto()
67
68
69
class PlatformEnum(enum.Enum):
vllm/v1/attention/backends/tree_attn.py
-1
@@ -316,7 +316,6 @@ class TreeAttentionImpl(AttentionImpl):
316
logits_soft_cap: Optional[float] = None,
317
attn_type: AttentionType = AttentionType.DECODER,
318
kv_sharing_target_layer_name: Optional[str] = None,
319
- use_irope: bool = False,
319
) -> None:
320
self.num_heads = num_heads
321
self.head_size = head_size
vllm/v1/attention/backends/xformers.py
new
+430
@@ -0,0 +1,430 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Attention layer with XFormersAttention."""
4
+
5
+from dataclasses import dataclass
6
+from typing import TYPE_CHECKING, Optional
7
+
8
+import torch
9
+
10
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
+ AttentionMetadata, AttentionType)
12
+from vllm.attention.ops.triton_unified_attention import unified_attention
13
+from vllm.config import VllmConfig
14
+from vllm.logger import init_logger
15
+from vllm.v1.attention.backends.utils import (
16
+ AttentionMetadataBuilder, CommonAttentionMetadata,
17
+ reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)
18
+from vllm.v1.kv_cache_interface import AttentionSpec
19
+
20
+try:
21
+ from xformers import ops as xops
22
+ from xformers.ops.fmha.attn_bias import (
23
+ AttentionBias, PagedBlockDiagonalCausalWithOffsetPaddedKeysMask)
24
+
25
+ XFORMERS_AVAILABLE = True
26
+except ImportError:
27
+ XFORMERS_AVAILABLE = False
28
+
29
+if TYPE_CHECKING:
30
+ from vllm.v1.core.sched.output import SchedulerOutput
31
+ from vllm.v1.worker.gpu_input_batch import InputBatch
32
+
33
+from vllm import _custom_ops as ops
34
+
35
+logger = init_logger(__name__)
36
+
37
+
38
+class XFormersAttentionBackend(AttentionBackend):
39
+
40
+ accept_output_buffer: bool = True
41
+
42
+ @classmethod
43
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
44
+ return [torch.float16, torch.bfloat16]
45
+
46
+ @classmethod
47
+ def get_supported_head_sizes(cls) -> list[int]:
48
+ return [
49
+ 32,
50
+ 40,
51
+ 48,
52
+ 56,
53
+ 64,
54
+ 72,
55
+ 80,
56
+ 88,
57
+ 96,
58
+ 104,
59
+ 112,
60
+ 120,
61
+ 128,
62
+ 136,
63
+ 144,
64
+ 152,
65
+ 160,
66
+ 168,
67
+ 176,
68
+ 184,
69
+ 192,
70
+ 200,
71
+ 208,
72
+ 216,
73
+ 224,
74
+ 232,
75
+ 240,
76
+ 248,
77
+ 256,
78
+ ]
79
+
80
+ @classmethod
81
+ def validate_head_size(cls, head_size: int) -> None:
82
+ supported_head_sizes = cls.get_supported_head_sizes()
83
+ if head_size not in supported_head_sizes:
84
+ attn_type = cls.__name__.removesuffix("Backend")
85
+ raise ValueError(
86
+ f"Head size {head_size} is not supported by {attn_type}. "
87
+ f"Supported head sizes are: {supported_head_sizes}. "
88
+ "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
89
+ "FlexAttention backend which supports all head sizes.")
90
+
91
+ @staticmethod
92
+ def get_name() -> str:
93
+ return "XFORMERS_VLLM_V1"
94
+
95
+ @staticmethod
96
+ def get_impl_cls() -> type["XFormersAttentionImpl"]:
97
+ return XFormersAttentionImpl
98
+
99
+ @staticmethod
100
+ def get_metadata_cls() -> type["AttentionMetadata"]:
101
+ return XFormersAttentionMetadata
102
+
103
+ @staticmethod
104
+ def get_kv_cache_shape(
105
+ num_blocks: int,
106
+ block_size: int,
107
+ num_kv_heads: int,
108
+ head_size: int,
109
+ ) -> tuple[int, ...]:
110
+ if block_size % 16 != 0:
111
+ raise ValueError("Block size must be a multiple of 16.")
112
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
113
+
114
+ @staticmethod
115
+ def get_builder_cls() -> type["XFormersAttentionMetadataBuilder"]:
116
+ return XFormersAttentionMetadataBuilder
117
+
118
+ @staticmethod
119
+ def use_cascade_attention(*args, **kwargs) -> bool:
120
+ return False
121
+
122
+
123
+@dataclass
124
+class XFormersAttentionMetadata:
125
+ num_actual_tokens: int # Number of tokens excluding padding.
126
+ max_query_len: int
127
+ query_start_loc: torch.Tensor
128
+ max_seq_len: int
129
+ seq_lens: torch.Tensor
130
+ block_table: torch.Tensor
131
+ slot_mapping: torch.Tensor
132
+
133
+ num_prefill_tokens: int = 0
134
+ num_decode_tokens: int = 0
135
+ num_prefills: int = 0
136
+ num_decodes: int = 0
137
+
138
+ # Biases for different attention types.
139
+ attn_bias: Optional["AttentionBias"] = None
140
+
141
+ # Self-attention prefill/decode metadata cache
142
+ _cached_prefill_metadata: Optional["XFormersAttentionMetadata"] = None
143
+ _cached_decode_metadata: Optional["XFormersAttentionMetadata"] = None
144
+
145
+ @property
146
+ def prefill_metadata(self) -> Optional["XFormersAttentionMetadata"]:
147
+ if self.num_prefills == 0:
148
+ return None
149
+
150
+ if self._cached_prefill_metadata is not None:
151
+ # Recover cached prefill-phase attention
152
+ # metadata structure
153
+ return self._cached_prefill_metadata
154
+
155
+ q_start_loc = self.query_start_loc[self.num_decodes:]
156
+ q_seqlens = torch.diff(q_start_loc)
157
+ kv_seqlens = self.seq_lens[self.num_decodes:]
158
+ # Construct & cache prefill-phase attention metadata structure
159
+ self._cached_prefill_metadata = XFormersAttentionMetadata(
160
+ num_actual_tokens=self.num_prefill_tokens,
161
+ max_query_len=int(q_seqlens.max().item()),
162
+ query_start_loc=q_start_loc - q_start_loc[0],
163
+ max_seq_len=int(kv_seqlens.max().item()),
164
+ seq_lens=kv_seqlens,
165
+ block_table=self.block_table[self.num_decodes:],
166
+ slot_mapping=self.slot_mapping[self.num_decode_tokens:],
167
+ )
168
+ return self._cached_prefill_metadata
169
+
170
+ @property
171
+ def decode_metadata(self) -> Optional["XFormersAttentionMetadata"]:
172
+ if self.num_decode_tokens == 0:
173
+ return None
174
+
175
+ if self._cached_decode_metadata is not None:
176
+ # Recover cached decode-phase attention
177
+ # metadata structure
178
+ return self._cached_decode_metadata
179
+
180
+ q_start_loc = self.query_start_loc
181
+ q_seqlens = torch.diff(q_start_loc)
182
+ decode_kv_seqlens = self.seq_lens[:self.num_decodes]
183
+ # Construct & cache decode-phase attention metadata structure
184
+ self._cached_decode_metadata = XFormersAttentionMetadata(
185
+ num_actual_tokens=self.num_decode_tokens,
186
+ max_query_len=int(q_seqlens[:self.num_decodes].max().item()),
187
+ query_start_loc=q_start_loc[:self.num_decodes + 1],
188
+ max_seq_len=int(decode_kv_seqlens.max().item()),
189
+ seq_lens=decode_kv_seqlens,
190
+ block_table=self.block_table[:self.num_decodes],
191
+ slot_mapping=self.slot_mapping[:self.num_decode_tokens],
192
+ attn_bias=self.attn_bias,
193
+ )
194
+ return self._cached_decode_metadata
195
+
196
+
197
+class XFormersAttentionMetadataBuilder(
198
+ AttentionMetadataBuilder[XFormersAttentionMetadata]):
199
+
200
+ def __init__(
201
+ self,
202
+ kv_cache_spec: AttentionSpec,
203
+ layer_names: list[str],
204
+ vllm_config: VllmConfig,
205
+ device: torch.device,
206
+ ):
207
+ assert XFORMERS_AVAILABLE
208
+ self.kv_cache_spec = kv_cache_spec
209
+ self.block_size = kv_cache_spec.block_size
210
+ self._num_decodes = 0
211
+ self._num_decode_tokens = 0
212
+
213
+ def reorder_batch(self, input_batch: "InputBatch",
214
+ scheduler_output: "SchedulerOutput") -> bool:
215
+ return reorder_batch_to_split_decodes_and_prefills(input_batch,
216
+ scheduler_output,
217
+ decode_threshold=1)
218
+
219
+ def build(
220
+ self,
221
+ common_prefix_len: int,
222
+ common_attn_metadata: CommonAttentionMetadata,
223
+ fast_build: bool = False,
224
+ ) -> XFormersAttentionMetadata:
225
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
226
+ split_decodes_and_prefills(common_attn_metadata,
227
+ decode_threshold=1))
228
+
229
+ num_actual_tokens = common_attn_metadata.num_actual_tokens
230
+ q_start_loc = common_attn_metadata.query_start_loc
231
+ q_seqlens = torch.diff(q_start_loc)
232
+ max_query_len = common_attn_metadata.max_query_len
233
+ kv_seqlens = common_attn_metadata.seq_lens
234
+ max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
235
+ block_table = common_attn_metadata.block_table_tensor
236
+ slot_mapping = common_attn_metadata.slot_mapping
237
+
238
+ bias = None
239
+ if num_decodes > 0:
240
+ # Construct the decoder bias.
241
+ decode_q_seqlens = q_seqlens[:num_decodes]
242
+ decode_kv_seqlens = kv_seqlens[:num_decodes]
243
+ bias = (
244
+ PagedBlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens(
245
+ q_seqlen=decode_q_seqlens.tolist(),
246
+ kv_seqlen=decode_kv_seqlens.tolist(),
247
+ page_size=self.block_size,
248
+ block_tables=block_table[:num_decodes],
249
+ device=block_table.device,
250
+ ))
251
+
252
+ return XFormersAttentionMetadata(
253
+ num_actual_tokens=num_actual_tokens,
254
+ num_prefill_tokens=num_prefill_tokens,
255
+ num_decode_tokens=num_decode_tokens,
256
+ num_prefills=num_prefills,
257
+ num_decodes=num_decodes,
258
+ max_query_len=max_query_len,
259
+ query_start_loc=q_start_loc,
260
+ max_seq_len=max_seq_len,
261
+ seq_lens=kv_seqlens,
262
+ block_table=block_table,
263
+ slot_mapping=slot_mapping,
264
+ attn_bias=bias,
265
+ )
266
+
267
+
268
+class XFormersAttentionImpl(AttentionImpl):
269
+
270
+ def __init__(
271
+ self,
272
+ num_heads: int,
273
+ head_size: int,
274
+ scale: float,
275
+ num_kv_heads: int,
276
+ alibi_slopes: Optional[list[float]],
277
+ sliding_window: Optional[int],
278
+ kv_cache_dtype: str,
279
+ logits_soft_cap: Optional[float] = None,
280
+ attn_type: AttentionType = AttentionType.DECODER,
281
+ kv_sharing_target_layer_name: Optional[str] = None,
282
+ ) -> None:
283
+ if kv_sharing_target_layer_name is not None:
284
+ raise NotImplementedError("KV sharing is not supported in V0.")
285
+ if alibi_slopes is not None:
286
+ raise NotImplementedError(
287
+ "XFormers does not support alibi slopes yet.")
288
+ self.num_heads = num_heads
289
+ self.head_size = head_size
290
+ self.scale = float(scale)
291
+ self.num_kv_heads = num_kv_heads
292
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
293
+ self.kv_cache_dtype = kv_cache_dtype
294
+ self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
295
+ if alibi_slopes is not None:
296
+ alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
297
+ self.alibi_slopes = alibi_slopes
298
+ if sliding_window is None:
299
+ self.sliding_window = (-1, -1)
300
+ else:
301
+ self.sliding_window = (sliding_window - 1, 0)
302
+ if logits_soft_cap is None:
303
+ # Setting logits_soft_cap to 0 means no soft cap.
304
+ logits_soft_cap = 0
305
+ self.logits_soft_cap = logits_soft_cap
306
+
307
+ XFormersAttentionBackend.validate_head_size(head_size)
308
+
309
+ if attn_type != AttentionType.DECODER:
310
+ raise NotImplementedError("Encoder self-attention and "
311
+ "encoder/decoder cross-attention "
312
+ "are not implemented for "
313
+ "XFormersAttentionImpl.")
314
+
315
+ def forward(
316
+ self,
317
+ layer: torch.nn.Module,
318
+ query: torch.Tensor,
319
+ key: torch.Tensor,
320
+ value: torch.Tensor,
321
+ kv_cache: torch.Tensor,
322
+ attn_metadata: XFormersAttentionMetadata,
323
+ output: Optional[torch.Tensor] = None,
324
+ output_scale: Optional[torch.Tensor] = None,
325
+ ) -> torch.Tensor:
326
+ """Forward pass with XFormers.
327
+
328
+ Args:
329
+ query: shape = [num_tokens, num_heads, head_size]
330
+ key: shape = [num_tokens, num_kv_heads, head_size]
331
+ value: shape = [num_tokens, num_kv_heads, head_size]
332
+ kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
333
+ attn_metadata: Metadata for attention.
334
+ Returns:
335
+ shape = [num_tokens, num_heads * head_size]
336
+ """
337
+ assert output is not None, "Output tensor must be provided."
338
+
339
+ if output_scale is not None:
340
+ raise NotImplementedError(
341
+ "fused output quantization is not yet supported"
342
+ " for XFormersAttentionImpl")
343
+
344
+ if attn_metadata is None:
345
+ # Profiling run.
346
+ return output
347
+
348
+ # Cache the input KVs.
349
+ key_cache, value_cache = kv_cache.unbind(0)
350
+ if self.kv_sharing_target_layer_name is None:
351
+ # Reshape the input keys and values and store them in the cache.
352
+ # Skip this if sharing KV cache with an earlier attention layer.
353
+ # NOTE(woosuk): Here, key and value are padded while slot_mapping is
354
+ # not padded. However, we don't need to do key[:num_actual_tokens]
355
+ # and value[:num_actual_tokens] because the reshape_and_cache_flash
356
+ # op uses the slot_mapping's shape to determine the number of
357
+ # actual tokens.
358
+ ops.reshape_and_cache_flash(
359
+ key,
360
+ value,
361
+ key_cache,
362
+ value_cache,
363
+ attn_metadata.slot_mapping,
364
+ self.kv_cache_dtype,
365
+ layer._k_scale,
366
+ layer._v_scale,
367
+ )
368
+
369
+ num_actual_tokens = attn_metadata.num_actual_tokens
370
+ num_decode_tokens = attn_metadata.num_decode_tokens
371
+ if prefill_meta := attn_metadata.prefill_metadata:
372
+ descale_shape = (prefill_meta.query_start_loc.shape[0] - 1,
373
+ key.shape[1])
374
+ unified_attention(
375
+ q=query[num_decode_tokens:num_actual_tokens],
376
+ k=key_cache,
377
+ v=value_cache,
378
+ out=output[num_decode_tokens:num_actual_tokens],
379
+ cu_seqlens_q=prefill_meta.query_start_loc,
380
+ max_seqlen_q=prefill_meta.max_query_len,
381
+ seqused_k=prefill_meta.seq_lens,
382
+ max_seqlen_k=prefill_meta.max_seq_len,
383
+ softmax_scale=self.scale,
384
+ causal=True,
385
+ alibi_slopes=self.alibi_slopes,
386
+ window_size=self.sliding_window,
387
+ block_table=prefill_meta.block_table,
388
+ softcap=self.logits_soft_cap,
389
+ q_descale=None, # Not supported
390
+ k_descale=layer._k_scale.expand(descale_shape),
391
+ v_descale=layer._v_scale.expand(descale_shape),
392
+ )
393
+
394
+ if decode_meta := attn_metadata.decode_metadata:
395
+ # Query for decode. KV is not needed because it is already cached.
396
+ decode_query = query[:num_decode_tokens]
397
+ # Reshape query to [1, B_T, G, H, D].
398
+ q = decode_query.view(1, -1, self.num_kv_heads,
399
+ self.num_queries_per_kv, self.head_size)
400
+ # Reshape the k and v caches to [1, Bkv_T, G, H, D]
401
+ cache_k = key_cache.view(1, -1, self.num_kv_heads, 1,
402
+ self.head_size).expand(
403
+ 1,
404
+ -1,
405
+ self.num_kv_heads,
406
+ self.num_queries_per_kv,
407
+ self.head_size,
408
+ )
409
+ cache_v = value_cache.view(1, -1, self.num_kv_heads, 1,
410
+ self.head_size).expand(
411
+ 1,
412
+ -1,
413
+ self.num_kv_heads,
414
+ self.num_queries_per_kv,
415
+ self.head_size,
416
+ )
417
+
418
+ attn_bias = decode_meta.attn_bias
419
+ output[:
420
+ num_decode_tokens] = xops.memory_efficient_attention_forward(
421
+ q,
422
+ cache_k,
423
+ cache_v,
424
+ attn_bias=attn_bias,
425
+ p=0.0,
426
+ scale=self.scale,
427
+ ).view(decode_query.shape)
428
+
429
+ # Reshape the output tensor.
430
+ return output