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 4, 22:30 UTC
[XPU] enable triton attention test on XPU by removing cuda device binding (#39627)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Apr 20, 2026 at 20:57 UTC
e06de7f0057fe1dfcc4fb039b52a61d39a079c4c
3 files changed
+52
-36
tests/kernels/attention/test_triton_decode_attention.py
+31
-24
@@ -4,9 +4,12 @@
4
import pytest
5
import torch
6
7
+from vllm.platforms import current_platform
8
from vllm.utils.math_utils import cdiv
9
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
10
11
+DEVICE_TYPE = current_platform.device_type
12
+
13
14
@pytest.mark.parametrize("B", [3, 5])
15
@pytest.mark.parametrize("L", [1027, 1025])
@@ -25,33 +28,35 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
28
29
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
30
req_to_page = torch.randint(
28
- 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
31
+ 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
32
)
33
req_to_token = req_to_page * PAGE_SIZE
34
req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
32
- req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1)
35
+ req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
36
+ 1, 1, -1
37
+ )
38
req_to_token = req_to_token.view(B, -1)
39
req_to_token = req_to_token[:, :seq_len].contiguous()
40
41
# q represents the new token being generated, one per batch
37
- q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
42
+ q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
43
44
# k_buffer and v_buffer represent all previous tokens
45
# Page size is 1.
41
- k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")
42
- v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda")
46
+ k_buffer = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
47
+ v_buffer = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
48
49
# o will have the same shape as q
45
- o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
50
+ o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
51
47
- lse = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
52
+ lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
53
49
- b_seq_len = torch.full((B,), seq_len, device="cuda")
54
+ b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
55
56
attn_logits = torch.empty(
57
(B, H_Q, num_kv_splits, D_V + 1),
58
dtype=torch.float32,
54
- device="cuda",
59
+ device=DEVICE_TYPE,
60
)
61
62
# Call the original implementation.
@@ -127,25 +132,27 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
132
133
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
134
req_to_page = torch.randint(
130
- 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"
135
+ 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device=DEVICE_TYPE
136
)
137
req_to_token = req_to_page * PAGE_SIZE
138
req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)
134
- req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1)
139
+ req_to_token = req_to_token + torch.arange(PAGE_SIZE, device=DEVICE_TYPE).view(
140
+ 1, 1, -1
141
+ )
142
req_to_token = req_to_token.view(B, -1)
143
req_to_token = req_to_token[:, :seq_len].contiguous()
144
138
- q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")
145
+ q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
146
147
# Create BF16 K/V as reference
141
- k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")
142
- v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda")
148
+ k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE)
149
+ v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE)
150
151
# --- BF16 reference ---
145
- o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
146
- lse_ref = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
152
+ o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
153
+ lse_ref = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
154
attn_logits = torch.empty(
148
- (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
155
+ (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
156
)
157
158
if PAGE_SIZE == 1:
@@ -156,7 +163,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
163
o_ref,
164
lse_ref,
165
req_to_token,
159
- b_seq_len=torch.full((B,), seq_len, device="cuda"),
166
+ b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
167
attn_logits=attn_logits,
168
num_kv_splits=num_kv_splits,
169
sm_scale=sm_scale,
@@ -171,7 +178,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
178
o_ref,
179
lse_ref,
180
req_to_page,
174
- b_seq_len=torch.full((B,), seq_len, device="cuda"),
181
+ b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
182
attn_logits=attn_logits,
183
num_kv_splits=num_kv_splits,
184
sm_scale=sm_scale,
@@ -182,10 +189,10 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
189
k_fp8, k_scale = _quantize_to_fp8(k_bf16)
190
v_fp8, v_scale = _quantize_to_fp8(v_bf16)
191
185
- o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
186
- lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device="cuda")
192
+ o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
193
+ lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
194
attn_logits_fp8 = torch.empty(
188
- (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"
195
+ (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
196
)
197
198
if PAGE_SIZE == 1:
@@ -196,7 +203,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
203
o_fp8,
204
lse_fp8,
205
req_to_token,
199
- b_seq_len=torch.full((B,), seq_len, device="cuda"),
206
+ b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
207
attn_logits=attn_logits_fp8,
208
num_kv_splits=num_kv_splits,
209
sm_scale=sm_scale,
@@ -213,7 +220,7 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
220
o_fp8,
221
lse_fp8,
222
req_to_page,
216
- b_seq_len=torch.full((B,), seq_len, device="cuda"),
223
+ b_seq_len=torch.full((B,), seq_len, device=DEVICE_TYPE),
224
attn_logits=attn_logits_fp8,
225
num_kv_splits=num_kv_splits,
226
sm_scale=sm_scale,
tests/kernels/attention/test_triton_prefill_attention.py
+17
-10
@@ -5,8 +5,11 @@ import pytest
5
import torch
6
import torch.nn.functional as F
7
8
+from vllm.platforms import current_platform
9
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
10
11
+DEVICE_TYPE = current_platform.device_type
12
+
13
14
def ref_masked_attention(
15
q: torch.Tensor,
@@ -92,17 +95,19 @@ def test_context_attention(
95
torch.manual_seed(42)
96
97
# Generate random sequence lengths for each batch
95
- seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
98
+ seq_lens = torch.randint(
99
+ max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
100
+ )
101
total_tokens = seq_lens.sum().item()
102
103
# Create batch start locations
99
- b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
104
+ b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
105
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
106
107
# Create input tensors
103
- q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
104
- k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
105
- v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
108
+ q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
109
+ k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
110
+ v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
111
o = torch.zeros_like(q)
112
113
# Call Triton kernel
@@ -169,17 +174,19 @@ def test_context_attention_sliding_window(
174
torch.manual_seed(42)
175
176
# Generate random sequence lengths for each batch
172
- seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
177
+ seq_lens = torch.randint(
178
+ max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE
179
+ )
180
total_tokens = seq_lens.sum().item()
181
182
# Create batch start locations
176
- b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
183
+ b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE)
184
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
185
186
# Create input tensors
180
- q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
181
- k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
182
- v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
187
+ q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE)
188
+ k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
189
+ v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE)
190
o = torch.zeros_like(q)
191
192
# Call Triton kernel
tests/kernels/attention/test_triton_unified_attention.py
+4
-2
@@ -10,6 +10,8 @@ from vllm.utils.math_utils import next_power_of_2
10
from vllm.utils.torch_utils import set_random_seed
11
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
12
13
+DEVICE_TYPE = current_platform.device_type
14
+
15
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
16
HEAD_SIZES = [128, 256]
17
BLOCK_SIZES = [16]
@@ -114,7 +116,7 @@ def test_triton_unified_attn(
116
q_dtype: torch.dtype | None,
117
seq_threshold_3D: int,
118
) -> None:
117
- torch.set_default_device("cuda")
119
+ torch.set_default_device(DEVICE_TYPE)
120
121
set_random_seed(0)
122
num_seqs = len(seq_lens)
@@ -249,7 +251,7 @@ def test_triton_unified_attn_fp16_input_fp8_output(
251
seq_threshold_3D: int,
252
) -> None:
253
"""Test with fp16 input and fp8 output using output_scale."""
252
- torch.set_default_device("cuda")
254
+ torch.set_default_device(DEVICE_TYPE)
255
256
set_random_seed(0)
257
num_seqs = len(seq_lens)