minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

enable ut test for xpu devices (#11712)

Browse Source 
Co-authored-by: jundu <jun.du@intel.com>
Co-authored-by: Gao, Pengfei <pengfei.gao@intel.com>
This commit is contained in:
DiweiSun
2026-02-03 11:15:14 -08:00
committed by GitHub
co-authored by jundu Gao, Pengfei
parent 0a6925639b
commit 495290aefd
20 changed files with 237 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_create_kvindices.py
+9 -8
View File
@@ -4,6 +4,7 @@ import numpy as np
import torch
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -15,30 +16,28 @@ register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
class TestCreateKvIndices(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
torch.set_default_device("cuda")
torch.set_default_device(get_device())
def _run_test(self, batch, max_batch, max_context_len):
req_to_token = torch.arange(
max_batch * max_context_len, dtype=torch.int32, device="cuda"
max_batch * max_context_len, dtype=torch.int32, device=get_device()
).reshape((max_batch, max_context_len))
req_pool_indices = torch.tensor(
torch.from_numpy(
np.random.choice(range(max_batch), size=batch, replace=False)
),
dtype=torch.int32,
device="cuda",
device=get_device(),
)
paged_kernel_lens = torch.tensor(
torch.from_numpy(
np.random.choice(range(max_context_len), size=batch, replace=False)
),
dtype=torch.int32,
device="cuda",
device=get_device(),
)
kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32, device="cuda")
kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32, device=get_device())
kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
# ref
@@ -53,7 +52,9 @@ class TestCreateKvIndices(CustomTestCase):
).contiguous()
# triton
kv_indices_triton = torch.empty(kv_indptr[-1], dtype=torch.int32, device="cuda")
kv_indices_triton = torch.empty(
kv_indptr[-1], dtype=torch.int32, device=get_device()
)
create_flashinfer_kv_indices_triton[(batch,)](
req_to_token,
req_pool_indices,
test/registered/attention/test_wave_attention_kernels.py
+64 -39
View File
@@ -21,6 +21,7 @@ from sglang.srt.layers.attention.wave_ops.extend_attention import extend_attenti
from sglang.srt.layers.attention.wave_ops.prefill_attention import (
prefill_attention_wave,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci
# Wave attention kernel unit tests (AMD only - requires wave_lang)
@@ -47,24 +48,24 @@ class TestWaveAttention(unittest.TestCase):
extend_seq_len = 1024
b_seq_len_prefix = torch.full(
(B,), N_CTX // B, dtype=torch.int32, device="cuda"
(B,), N_CTX // B, dtype=torch.int32, device=get_device()
)
b_seq_len_extend = torch.full(
(B,), extend_seq_len, dtype=torch.int32, device="cuda"
(B,), extend_seq_len, dtype=torch.int32, device=get_device()
)
b_seq_len = b_seq_len_prefix + b_seq_len_extend
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
b_req_idx = torch.arange(B, dtype=torch.int32, device="cuda")
b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda")
b_req_idx = torch.arange(B, dtype=torch.int32, device=get_device())
b_start_loc = torch.zeros((B,), dtype=torch.int32, device=get_device())
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda")
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=get_device())
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=get_device())
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
kv_indices = torch.zeros(
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device="cuda"
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=get_device()
)
for i in range(B):
@@ -75,15 +76,21 @@ class TestWaveAttention(unittest.TestCase):
total_token_num = torch.sum(b_seq_len).item()
extend_token_num = torch.sum(b_seq_len_extend).item()
k_buffer = torch.empty(
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
(total_token_num, H_KV, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
v_buffer = torch.empty(
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
(total_token_num, H_KV, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
k_extend = torch.empty(
(extend_token_num, H_KV, D), dtype=dtype, device=get_device()
)
v_extend = torch.empty(
(extend_token_num, H_KV, D), dtype=dtype, device=get_device()
)
q_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
for i in range(B):
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
@@ -96,20 +103,22 @@ class TestWaveAttention(unittest.TestCase):
extend_start_in_buffer:extend_end_in_buffer
]
q_extend[extend_start:extend_end] = torch.empty(
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda"
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device=get_device()
).normal_(mean=0.1, std=0.2)
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
o_extend_mask = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
o_redundant = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
b_seq_len_extend = b_seq_len - b_seq_len_prefix
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=get_device())
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
custom_mask = None
@@ -129,7 +138,9 @@ class TestWaveAttention(unittest.TestCase):
is_causal = True
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_extend = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
extend_attention_fwd(
q_extend,
k_extend,
@@ -146,7 +157,9 @@ class TestWaveAttention(unittest.TestCase):
max_len_extend,
)
o_wave = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
o_wave = torch.empty(
(extend_token_num, H_Q, D), dtype=dtype, device=get_device()
)
extend_attention_wave(
q_extend,
k_extend,
@@ -181,33 +194,37 @@ class TestWaveAttention(unittest.TestCase):
total_tokens = B * seq_len
sm_scale = 1.0 / (D**0.5)
max_kv_splits = 8
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda")
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=get_device())
# q represents the new token being generated, one per batch
q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda")
q = torch.randn(B, H_Q, D, dtype=dtype, device=get_device())
# k_buffer and v_buffer represent all previous tokens
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device="cuda")
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=get_device())
v_buffer = torch.randn(
total_tokens, H_KV, D_V, dtype=dtype, device=get_device()
)
# o will have the same shape as q
o_triton = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
o_triton = torch.zeros(B, H_Q, D_V, dtype=dtype, device=get_device())
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=get_device())
req_to_token = torch.arange(total_tokens, device="cuda", dtype=torch.int32)
b_req_idx = torch.zeros(B + 1, device="cuda", dtype=torch.int32)
b_seq_len = torch.full((B,), seq_len, device="cuda", dtype=torch.int32)
req_to_token = torch.arange(
total_tokens, device=get_device(), dtype=torch.int32
)
b_req_idx = torch.zeros(B + 1, device=get_device(), dtype=torch.int32)
b_seq_len = torch.full((B,), seq_len, device=get_device(), dtype=torch.int32)
b_req_idx[1 : B + 1] = torch.cumsum(b_seq_len, dim=0)
attn_logits = torch.empty(
(B, H_Q, max_kv_splits, D_V + 1),
dtype=torch.float32,
device="cuda",
device=get_device(),
)
attn_lse = torch.empty(
(B, H_Q, max_kv_splits),
dtype=torch.float32,
device="cuda",
device=get_device(),
)
logit_cap = 0.0
@@ -233,13 +250,13 @@ class TestWaveAttention(unittest.TestCase):
attn_logits = torch.empty(
attn_logits_shape,
dtype=torch.float32,
device="cuda",
device=get_device(),
)
attn_logits_max = torch.empty(
attn_logits_max_shape,
dtype=torch.float32,
device="cuda",
device=get_device(),
)
decode_attention_wave(
@@ -288,17 +305,25 @@ class TestWaveAttention(unittest.TestCase):
max_seq_len = max(seq_lens)
# Create random input tensors
q = torch.randn(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
k = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
v = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
q = torch.randn(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
k = torch.randn(
sum(seq_lens), kv_heads, head_dim, dtype=dtype, device=get_device()
)
v = torch.randn(
sum(seq_lens), kv_heads, head_dim, dtype=dtype, device=get_device()
)
o_triton = torch.zeros(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda"
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
o = torch.zeros(
sum(seq_lens), num_heads, head_dim, dtype=dtype, device=get_device()
)
o = torch.zeros(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
# Create b_start_loc and b_seq_len tensors
b_start_loc = torch.tensor([0, seq_lens[0]], device="cuda")
b_seq_len = torch.tensor(seq_lens, device="cuda")
b_start_loc = torch.tensor([0, seq_lens[0]], device=get_device())
b_seq_len = torch.tensor(seq_lens, device=get_device())
context_attention_fwd(
q, k, v, o_triton, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal