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

[XPU] Adapt device agnostic API usage (#32093)

Browse Source 
Co-authored-by: mingfeima <mingfei.ma@intel.com>
This commit is contained in:
ANSHUMAN TRIPATHY
2026-09-11 10:03:12 +08:00
committed by GitHub
co-authored by mingfeima
parent 67d3a2ea57
commit 2adb2e8485
4 changed files with 115 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/layers/test_fused_gate_sigmoid_mul_add.py
+14 -12
View File
@@ -4,10 +4,12 @@ import pytest
import torch
from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add
from sglang.srt.utils import get_device
DTYPES = [torch.float16, torch.bfloat16]
TOKEN_COUNTS = [1, 2, 4, 8, 16, 64, 512, 1024, 2048, 4096, 8192]
HIDDEN_DIMS = [2048, 3072, 4096, 6144]
DEVICE = get_device()
def _reference(hidden_states, gate_weight, shared_output, final_hidden_states):
@@ -27,10 +29,10 @@ def seed():
def test_correctness(num_tokens, hidden_dim, dtype):
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate_weight = torch.randn(hidden_dim, dtype=dtype, device="cuda")
shared_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
final_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
gate_weight = torch.randn(hidden_dim, dtype=dtype, device=DEVICE)
shared_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
final_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
final_test = final_ref.clone()
_reference(hidden_states, gate_weight, shared_output, final_ref)
@@ -42,10 +44,10 @@ def test_correctness(num_tokens, hidden_dim, dtype):
@pytest.mark.parametrize("dtype", DTYPES)
def test_gate_near_zero(dtype):
num_tokens, hidden_dim = 16, 2048
hs = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gw = torch.zeros(hidden_dim, dtype=dtype, device="cuda")
so = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
f_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
hs = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
gw = torch.zeros(hidden_dim, dtype=dtype, device=DEVICE)
so = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
f_ref = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
f_test = f_ref.clone()
_reference(hs, gw, so, f_ref)
@@ -56,10 +58,10 @@ def test_gate_near_zero(dtype):
def test_inplace_semantics():
num_tokens, hidden_dim = 32, 2048
hs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
gw = torch.randn(hidden_dim, dtype=torch.float16, device="cuda")
so = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
fhs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
hs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device=DEVICE)
gw = torch.randn(hidden_dim, dtype=torch.float16, device=DEVICE)
so = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device=DEVICE)
fhs = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device=DEVICE)
original_ptr = fhs.data_ptr()
fused_gate_sigmoid_mul_add(hs, gw, so, fhs)
test/manual/layers/test_fused_sigmoid_mul.py
+14 -12
View File
@@ -4,11 +4,13 @@ import pytest
import torch
from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul
from sglang.srt.utils import get_device
DTYPES = [torch.float16, torch.bfloat16]
TOKEN_COUNTS = [1, 2, 4, 8, 16, 64, 512, 1024, 2048, 4096, 8192]
HIDDEN_DIMS = [2048, 3072, 4096, 6144]
NUM_HEADS = [1, 28]
DEVICE = get_device()
def _reference(attn_output, gate):
@@ -27,8 +29,8 @@ def seed():
def test_correctness(num_tokens, hidden_dim, dtype):
rtol, atol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (1e-2, 1e-2)
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
gate = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
@@ -46,9 +48,9 @@ def test_3d_shape(num_tokens, num_heads, dtype):
head_dim = 128
attn_output = torch.randn(
num_tokens, num_heads, head_dim, dtype=dtype, device="cuda"
num_tokens, num_heads, head_dim, dtype=dtype, device=DEVICE
)
gate = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device="cuda")
gate = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device=DEVICE)
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
@@ -68,12 +70,12 @@ def test_strided_gate(num_tokens, num_heads, dtype):
# Simulate the real pattern: chunk produces non-contiguous views
q_gate = torch.randn(
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device="cuda"
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device=DEVICE
)
_, gate = torch.chunk(q_gate, 2, dim=-1)
# gate is non-contiguous when num_tokens > 1 or num_heads > 1
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
gate_flat = gate.reshape(num_tokens, hidden_dim)
ref = _reference(attn_output, gate_flat)
@@ -92,10 +94,10 @@ def test_qwen3_5_moe_target_strided_gate(num_tokens, dtype):
hidden_dim = num_heads * head_dim
q_gate = torch.randn(
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device="cuda"
num_tokens, num_heads, 2 * head_dim, dtype=dtype, device=DEVICE
)
_, gate = torch.chunk(q_gate, 2, dim=-1)
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
ref = _reference(attn_output, gate.reshape(num_tokens, hidden_dim))
out = fused_sigmoid_mul(attn_output, gate, inplace=False)
@@ -106,8 +108,8 @@ def test_qwen3_5_moe_target_strided_gate(num_tokens, dtype):
@pytest.mark.parametrize("dtype", DTYPES)
def test_gate_near_zero(dtype):
num_tokens, hidden_dim = 16, 2048
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
gate = torch.zeros(num_tokens, hidden_dim, dtype=dtype, device="cuda")
attn_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
gate = torch.zeros(num_tokens, hidden_dim, dtype=dtype, device=DEVICE)
ref = _reference(attn_output, gate)
out = fused_sigmoid_mul(attn_output, gate)
@@ -118,9 +120,9 @@ def test_gate_near_zero(dtype):
def test_returns_new_tensor():
num_tokens, hidden_dim = 32, 2048
attn_output = torch.randn(
num_tokens, hidden_dim, dtype=torch.float16, device="cuda"
num_tokens, hidden_dim, dtype=torch.float16, device=DEVICE
)
gate = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device="cuda")
gate = torch.randn(num_tokens, hidden_dim, dtype=torch.float16, device=DEVICE)
out = fused_sigmoid_mul(attn_output, gate)