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 1, 22:30 UTC
[Kernel][Test] Extend lightning_attn and awq_triton kernel tests to XPU (#43307)
Signed-off-by: Dobrzyniewicz, Agata <agata.dobrzyniewicz@intel.com>
Agata Dobrzyniewicz committed
Jun 4, 2026 at 20:25 UTC
a947f7a420de439a1f403c3f11c0cfdba03715df
3 files changed
+37
-19
tests/kernels/attention/test_lightning_attn.py
+19
-11
@@ -5,8 +5,16 @@ import pytest
5
import torch
6
7
from vllm.model_executor.layers.lightning_attn import linear_decode_forward_triton
8
+from vllm.platforms import current_platform
9
from vllm.utils.torch_utils import set_random_seed
10
11
+DEVICE = current_platform.device_type
12
+
13
+pytestmark = pytest.mark.skipif(
14
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
15
+ reason="Lightning attention Triton kernels require CUDA/ROCm or XPU.",
16
+)
17
+
18
NUM_HEADS = [4, 8]
19
HEAD_SIZES = [64]
20
BATCH_SIZES = [1, 2]
@@ -121,7 +129,7 @@ def test_linear_decode_forward_triton(
129
head_size: int,
130
dtype: torch.dtype,
131
):
124
- torch.set_default_device("cuda")
132
+ torch.set_default_device(DEVICE)
133
set_random_seed(42)
134
base = 0.01
135
q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
@@ -129,16 +137,16 @@ def test_linear_decode_forward_triton(
137
v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
138
139
kv_caches = base * torch.randn(
132
- batch_size, num_heads, head_size, head_size, dtype=dtype, device="cuda"
140
+ batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
141
)
142
143
kv_caches_copy = kv_caches.clone()
144
137
- slope_rate = torch.zeros(num_heads, device="cuda")
145
+ slope_rate = torch.zeros(num_heads, device=DEVICE)
146
for h in range(num_heads):
147
slope_rate[h] = 0.1 * (h + 1)
148
141
- slot_idx = torch.arange(batch_size, device="cuda")
149
+ slot_idx = torch.arange(batch_size, device=DEVICE)
150
151
triton_output = linear_decode_forward_triton(
152
q, k, v, kv_caches, slope_rate, slot_idx
@@ -162,7 +170,7 @@ def test_linear_decode_forward_triton_with_padding(
170
head_size: int,
171
dtype: torch.dtype,
172
):
165
- torch.set_default_device("cuda")
173
+ torch.set_default_device(DEVICE)
174
set_random_seed(42)
175
176
batch_size = 4
@@ -172,16 +180,16 @@ def test_linear_decode_forward_triton_with_padding(
180
v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)
181
182
kv_caches = base * torch.randn(
175
- batch_size, num_heads, head_size, head_size, dtype=dtype, device="cuda"
183
+ batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
184
)
185
186
kv_caches_copy = kv_caches.clone()
187
180
- slope_rate = torch.zeros(num_heads, device="cuda")
188
+ slope_rate = torch.zeros(num_heads, device=DEVICE)
189
for h in range(num_heads):
190
slope_rate[h] = 0.1 * (h + 1)
191
184
- slot_idx = torch.tensor([0, 1, -1, 2], device="cuda")
192
+ slot_idx = torch.tensor([0, 1, -1, 2], device=DEVICE)
193
194
triton_output = linear_decode_forward_triton(
195
q, k, v, kv_caches, slope_rate, slot_idx
@@ -224,7 +232,7 @@ def test_lightning_attention_reference(
232
seq_len: int,
233
dtype: torch.dtype,
234
):
227
- torch.set_default_device("cuda")
235
+ torch.set_default_device(DEVICE)
236
set_random_seed(42)
237
238
base = 0.01
@@ -232,12 +240,12 @@ def test_lightning_attention_reference(
240
k = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
241
v = base * torch.randn(batch_size, num_heads, seq_len, head_size, dtype=dtype)
242
235
- ed = torch.zeros(num_heads, device="cuda")
243
+ ed = torch.zeros(num_heads, device=DEVICE)
244
for h in range(num_heads):
245
ed[h] = 0.1 * (h + 1)
246
247
kv_history = base * torch.randn(
240
- batch_size, num_heads, head_size, head_size, dtype=dtype, device="cuda"
248
+ batch_size, num_heads, head_size, head_size, dtype=dtype, device=DEVICE
249
)
250
251
kv_history_clone = kv_history.clone()
tests/kernels/quantization/test_awq_triton.py
+7
-1
@@ -13,9 +13,15 @@ from vllm.model_executor.layers.quantization.awq_triton import (
13
awq_dequantize_triton,
14
awq_gemm_triton,
15
)
16
+from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
18
18
-device = "cuda"
19
+pytestmark = pytest.mark.skipif(
20
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
21
+ reason="AWQ Triton kernels require CUDA/ROCm or XPU.",
22
+)
23
+
24
+device = current_platform.device_type
25
26
27
def reverse_awq_order(t: torch.Tensor):
vllm/model_executor/layers/lightning_attn.py
+11
-7
@@ -4,6 +4,7 @@
4
import torch
5
from einops import rearrange
6
7
+from vllm.platforms import current_platform
8
from vllm.triton_utils import tl, triton
9
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
10
@@ -403,13 +404,16 @@ class _attention(torch.autograd.Function):
404
v = v.contiguous()
405
s = s.contiguous()
406
406
- # Check CUDA compute capability
407
- capability = torch.cuda.get_device_capability()
408
- if capability[0] < 8:
409
- raise RuntimeError(
410
- "Flash attention currently only supported",
411
- "for compute capability >= 80",
412
- )
407
+ # Check CUDA compute capability (Ampere+ required for flash attention
408
+ # path). Other accelerators (ROCm, XPU) rely on their own Triton
409
+ # backend support and skip this check.
410
+ if current_platform.is_cuda():
411
+ capability = torch.cuda.get_device_capability()
412
+ if capability[0] < 8:
413
+ raise RuntimeError(
414
+ "Flash attention currently only supported",
415
+ "for compute capability >= 80",
416
+ )
417
418
# Get input dimensions
419
b, h, n, d = q.shape