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
[XPU] support Triton Attention backend on Intel GPU (#24149)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Sep 4, 2025 at 20:41 UTC
16ded21eeb578cdd7060a32d21fe213d8f31cd0c
5 files changed
+49
-15
.buildkite/scripts/hardware_ci/run-xpu-test.sh
+5
-4
@@ -30,10 +30,11 @@ docker run \
30
bash -c '
31
set -e
32
echo $ZE_AFFINITY_MASK
33
- VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
34
- VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -O.cudagraph_mode=NONE
35
- VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
36
- VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
33
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
34
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -O.cudagraph_mode=NONE
35
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
36
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
37
+ VLLM_ATTENTION_BACKEND=TRITON_ATTN_VLLM_V1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
38
cd tests
39
pytest -v -s v1/core
40
pytest -v -s v1/engine
vllm/_ipex_ops.py
+2
-3
@@ -242,10 +242,9 @@ class ipex_ops:
242
k_scale_float: float = 1.0,
243
v_scale_float: float = 1.0,
244
) -> None:
245
- assert kv_cache_dtype == "auto"
246
- # TODO: support FP8 kv cache.
245
ipex.llm.modules.PagedAttention.reshape_and_cache_flash(
248
- key, value, key_cache, value_cache, slot_mapping)
246
+ key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype,
247
+ k_scale_float, v_scale_float)
248
249
@staticmethod
250
def flash_attn_varlen_func(
vllm/attention/ops/paged_attn.py
+6
-1
@@ -6,9 +6,14 @@ from typing import List, Optional, Tuple
6
7
import torch
8
9
-from vllm import _custom_ops as ops
9
+from vllm.platforms import current_platform
10
from vllm.triton_utils import HAS_TRITON
11
12
+if current_platform.is_cuda_alike():
13
+ from vllm import _custom_ops as ops
14
+elif current_platform.is_xpu():
15
+ from vllm._ipex_ops import ipex_ops as ops
16
+
17
if HAS_TRITON:
18
from vllm.attention.ops.prefix_prefill import context_attention_fwd
19
vllm/platforms/xpu.py
+26
-2
@@ -37,14 +37,38 @@ class XPUPlatform(Platform):
37
dtype: torch.dtype, kv_cache_dtype: Optional[str],
38
block_size: int, use_v1: bool, use_mla: bool,
39
has_sink: bool) -> str:
40
- if selected_backend is not None and selected_backend != _Backend.IPEX:
41
- logger.info("Cannot use %s backend on XPU.", selected_backend)
40
use_v1 = envs.VLLM_USE_V1
41
if not use_v1:
42
raise ValueError("XPU backend only supports V1.")
43
+ TRITON_ATTN_VLLM_V1 = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" # noqa: E501
44
+ FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend" # noqa: E501
45
+ if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
46
+ logger.info_once("Using Triton backend on V1 engine.")
47
+ return TRITON_ATTN_VLLM_V1
48
+ elif selected_backend == _Backend.FLASH_ATTN:
49
+ logger.info_once("Using Flash Attention backend on V1 engine.")
50
+ return FLASH_ATTN_V1
51
+ elif selected_backend:
52
+ raise ValueError(
53
+ f"Invalid attention backend for {cls.device_name}, "
54
+ f"with use_v1: {use_v1} use_mla: {use_mla}")
55
+
56
logger.info("Using Flash Attention backend on V1 engine.")
57
return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
58
59
+ @classmethod
60
+ def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,
61
+ model_config: "ModelConfig") -> bool:
62
+ """
63
+ Check if the kv_cache_dtype is supported.
64
+ XPU only support fp8 kv cache with triton backend.
65
+ """
66
+ if envs.is_set("VLLM_ATTENTION_BACKEND") and \
67
+ envs.VLLM_ATTENTION_BACKEND == "TRITON_ATTN_VLLM_V1":
68
+ return kv_cache_dtype in ["fp8_e4m3", "fp8_e5m2", "fp8"]
69
+
70
+ return False
71
+
72
@classmethod
73
def set_device(cls, device: torch.device) -> None:
74
"""
vllm/v1/attention/backends/triton_attn.py
+10
-5
@@ -7,7 +7,6 @@ from typing import ClassVar, Optional
7
8
import torch
9
10
-from vllm import _custom_ops as ops
10
from vllm import envs
11
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12
AttentionMetadata, AttentionType)
@@ -23,6 +22,11 @@ from vllm.v1.attention.backends.utils import (AttentionCGSupport,
22
CommonAttentionMetadata)
23
from vllm.v1.kv_cache_interface import AttentionSpec
24
25
+if current_platform.is_cuda_alike():
26
+ from vllm import _custom_ops as ops
27
+elif current_platform.is_xpu():
28
+ from vllm._ipex_ops import ipex_ops as ops
29
+
30
logger = init_logger(__name__)
31
32
@@ -337,7 +341,7 @@ class TritonAttentionImpl(AttentionImpl):
341
layer._v_scale,
342
)
343
else:
340
- torch.ops._C_cache_ops.reshape_and_cache_flash(
344
+ ops.reshape_and_cache_flash(
345
key,
346
value,
347
key_cache,
@@ -354,9 +358,10 @@ class TritonAttentionImpl(AttentionImpl):
358
num_tokens, num_heads, head_size = query.shape
359
assert layer._q_scale == 1.0, \
360
"A non 1.0 q_scale is not currently supported."
357
- if not current_platform.is_rocm():
358
- # Skip Q quantization on ROCm, since dequantizing back to
359
- # f32 in the attention kernel is not supported.
361
+ if current_platform.is_cuda():
362
+ # Skip Q quantization on ROCm and XPU, enable this on cuda
363
+ # only, since dequantizing back to f32 in the attention kernel
364
+ # is not supported.
365
query, _ = ops.scaled_fp8_quant(
366
query.reshape(
367
(num_tokens, num_heads * head_size)).contiguous(),