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 13, 23:15 UTC
[Multimodal][XPU]Enable vision attn backend for xpu platform (#27525)
Signed-off-by: Yan Ma <yan.ma@intel.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: Yejing Lai <yejing.lai@intel.com> Co-authored-by: Guancheng Fu <110874468+gc-fu@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yan Ma committed
Nov 1, 2025 at 12:45 UTC
7e2729b57e5c7420e945b6cf21850374195984c7
6 files changed
+88
-51
vllm/_ipex_ops.py
+59
-25
@@ -270,21 +270,23 @@ class ipex_ops:
270
271
@staticmethod
272
def flash_attn_varlen_func(
273
- out: torch.Tensor,
273
q: torch.Tensor,
274
k: torch.Tensor,
275
v: torch.Tensor,
276
cu_seqlens_q: torch.Tensor,
278
- seqused_k: torch.Tensor, # we don't support this in ipex kernel
277
max_seqlen_q: int,
278
max_seqlen_k: int,
281
- softmax_scale: float,
282
- causal: bool,
283
- block_table: torch.Tensor,
284
- alibi_slopes: torch.Tensor | None,
279
+ softmax_scale: float | None = None,
280
+ causal: bool = False,
281
+ out: torch.Tensor | None = None,
282
+ block_table: torch.Tensor | None = None,
283
+ alibi_slopes: torch.Tensor | None = None,
284
window_size: list[int] | None = None,
285
softcap: float | None = 0.0,
286
+ seqused_k: torch.Tensor | None = None,
287
cu_seqlens_k: torch.Tensor | None = None,
288
+ # passed in qwen vl
289
+ dropout_p: float = 0.0,
290
# The following parameters are not used in ipex kernel currently,
291
# we keep API compatible to CUDA's.
292
scheduler_metadata=None,
@@ -295,31 +297,63 @@ class ipex_ops:
297
num_splits=0,
298
s_aux: torch.Tensor | None = None,
299
):
300
+ if out is None:
301
+ out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
302
real_window_size: tuple[int, int]
303
if window_size is None:
304
real_window_size = (-1, -1)
305
else:
306
assert len(window_size) == 2
307
real_window_size = (window_size[0], window_size[1])
304
- return ipex.llm.modules.PagedAttention.flash_attn_varlen_func(
305
- out,
306
- q.contiguous(),
307
- k,
308
- v,
309
- cu_seqlens_q,
310
- seqused_k,
311
- max_seqlen_q,
312
- max_seqlen_k,
313
- softmax_scale,
314
- causal,
315
- block_table,
316
- alibi_slopes,
317
- softcap=softcap,
318
- window_size_left=real_window_size[0],
319
- window_size_right=real_window_size[1],
320
- k_scale=1.0,
321
- v_scale=1.0,
322
- )
308
+
309
+ if block_table is None:
310
+ assert cu_seqlens_k is not None, (
311
+ "cu_seqlens_k can't be None when calling varlen_attention."
312
+ )
313
+ if softmax_scale is None:
314
+ softmax_scale = q.shape[-1] ** (-0.5)
315
+ ipex_ops.varlen_attention(
316
+ q.contiguous(),
317
+ k.contiguous(),
318
+ v.contiguous(),
319
+ out,
320
+ cu_seqlens_q,
321
+ cu_seqlens_k,
322
+ None,
323
+ max_seqlen_q,
324
+ max_seqlen_k,
325
+ 0.0,
326
+ softmax_scale,
327
+ False,
328
+ causal,
329
+ False,
330
+ None,
331
+ real_window_size[0],
332
+ real_window_size[1],
333
+ -1,
334
+ )
335
+ return out
336
+ else:
337
+ return ipex.llm.modules.PagedAttention.flash_attn_varlen_func(
338
+ out,
339
+ q.contiguous(),
340
+ k,
341
+ v,
342
+ cu_seqlens_q,
343
+ seqused_k,
344
+ max_seqlen_q,
345
+ max_seqlen_k,
346
+ softmax_scale,
347
+ causal,
348
+ block_table,
349
+ alibi_slopes,
350
+ sink=s_aux,
351
+ softcap=softcap,
352
+ window_size_left=real_window_size[0],
353
+ window_size_right=real_window_size[1],
354
+ k_scale=1.0,
355
+ v_scale=1.0,
356
+ )
357
358
@staticmethod
359
def get_scheduler_metadata(
vllm/attention/layer.py
+18
-17
@@ -123,6 +123,11 @@ def maybe_get_vit_flash_attn_backend(
123
):
124
attn_backend = _Backend.FLASH_ATTN
125
use_upstream_fa = True
126
+ elif current_platform.is_xpu():
127
+ assert attn_backend == _Backend.FLASH_ATTN, (
128
+ "XPU platform only supports FLASH_ATTN as vision attention backend."
129
+ )
130
+ use_upstream_fa = False
131
else:
132
return _Backend.TORCH_SDPA, None
133
@@ -133,7 +138,7 @@ def maybe_get_vit_flash_attn_backend(
138
if use_upstream_fa:
139
from flash_attn import flash_attn_varlen_func
140
else:
136
- from vllm.vllm_flash_attn import flash_attn_varlen_func
141
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
142
else:
143
flash_attn_varlen_func = None
144
@@ -521,22 +526,18 @@ class MultiHeadAttention(nn.Module):
526
# If vllm native fa is selected, we use it directly.
527
use_upstream_fa = False
528
524
- if current_platform.is_xpu():
525
- # currently, only torch_sdpa is supported on xpu
526
- self.attn_backend = _Backend.TORCH_SDPA
527
- else:
528
- self.attn_backend = (
529
- backend
530
- if backend
531
- in {
532
- _Backend.TORCH_SDPA,
533
- _Backend.XFORMERS,
534
- _Backend.PALLAS,
535
- _Backend.ROCM_AITER_FA,
536
- _Backend.FLASH_ATTN,
537
- }
538
- else _Backend.TORCH_SDPA
539
- )
529
+ self.attn_backend = (
530
+ backend
531
+ if backend
532
+ in {
533
+ _Backend.TORCH_SDPA,
534
+ _Backend.XFORMERS,
535
+ _Backend.PALLAS,
536
+ _Backend.ROCM_AITER_FA,
537
+ _Backend.FLASH_ATTN,
538
+ }
539
+ else _Backend.TORCH_SDPA
540
+ )
541
542
self.attn_backend, self._flash_attn_varlen_func = (
543
maybe_get_vit_flash_attn_backend(
vllm/attention/ops/vit_attn_wrappers.py
+1
-1
@@ -70,7 +70,7 @@ def flash_attn_maxseqlen_wrapper(
70
if use_upstream_fa:
71
from flash_attn import flash_attn_varlen_func
72
else:
73
- from vllm.vllm_flash_attn import flash_attn_varlen_func
73
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
74
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
75
output = flash_attn_varlen_func(
76
q,
vllm/model_executor/models/qwen2_5_vl.py
+3
-4
@@ -364,6 +364,8 @@ class Qwen2_5_VisionAttention(nn.Module):
364
365
if current_platform.is_rocm() and self.attn_backend == _Backend.FLASH_ATTN:
366
self.use_upstream_fa = True
367
+ if current_platform.is_xpu():
368
+ self.use_upstream_fa = False
369
self.is_flash_attn_backend = self.attn_backend in {
370
_Backend.FLASH_ATTN,
371
_Backend.ROCM_AITER_FA,
@@ -856,10 +858,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
858
) -> tuple[torch.Tensor, torch.Tensor]:
859
max_seqlen = torch.zeros([], device=cu_seqlens.device)
860
seqlens = torch.zeros(1, device=cu_seqlens.device)
859
- if (
860
- self.attn_backend == _Backend.FLASH_ATTN
861
- or self.attn_backend == _Backend.ROCM_AITER_FA
862
- ):
861
+ if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
862
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
863
elif self.attn_backend == _Backend.XFORMERS:
864
seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
vllm/model_executor/models/qwen2_vl.py
+1
-4
@@ -789,10 +789,7 @@ class Qwen2VisionTransformer(nn.Module):
789
self, cu_seqlens: torch.Tensor
790
) -> tuple[int | None, list[int] | None]:
791
max_seqlen, seqlens = None, None
792
- if (
793
- self.attn_backend == _Backend.FLASH_ATTN
794
- or self.attn_backend == _Backend.ROCM_AITER_FA
795
- ):
792
+ if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
793
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
794
elif self.attn_backend == _Backend.XFORMERS:
795
seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
vllm/platforms/xpu.py
+6
@@ -115,6 +115,12 @@ class XPUPlatform(Platform):
115
device_props = torch.xpu.get_device_properties(device_id)
116
return device_props.total_memory
117
118
+ @classmethod
119
+ def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:
120
+ from vllm.attention.backends.registry import _Backend
121
+
122
+ return _Backend.FLASH_ATTN
123
+
124
@classmethod
125
def inference_mode(cls):
126
return torch.no_grad()