Update python API of activation, topk, norm and rope and remove vllm dependency (#6614)

Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
Co-authored-by: sdp <sdp@gnr799219.jf.intel.com>
This commit is contained in:
YanbingJiang
2025-06-18 13:11:50 +08:00
committed by GitHub
parent e56685ac1b
commit 094c116f7d
23 changed files with 270 additions and 56 deletions

View File

@@ -21,7 +21,7 @@ class TestActivation(CustomTestCase):
ref_out = SiluAndMul(x)
atol = rtol = precision[ref_out.dtype]
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def test_activation(self):
for params in itertools.product(self.M, self.N, self.dtype):

View File

@@ -60,8 +60,8 @@ class TestGemm(CustomTestCase):
)
atol = rtol = precision[ref.dtype]
self.assertTrue(torch.allclose(ref, out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(ref, out2, atol=atol, rtol=rtol))
torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)
torch.testing.assert_close(ref, out2, atol=atol, rtol=rtol)
def test_bf16_gemm(self):
for params in itertools.product(
@@ -100,13 +100,13 @@ class TestGemm(CustomTestCase):
out = torch.ops.sgl_kernel.int8_scaled_mm_cpu(
Aq2, Bq, As2, Bs, bias if has_bias else None, torch.bfloat16, False
)
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
# test the fused version
fused_out = torch.ops.sgl_kernel.int8_scaled_mm_with_quant(
A, Bq, Bs, bias if has_bias else None, torch.bfloat16, False
)
self.assertTrue(torch.allclose(ref_out, fused_out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, fused_out, atol=atol, rtol=rtol)
def test_int8_gemm(self):
for params in itertools.product(
@@ -165,7 +165,7 @@ class TestGemm(CustomTestCase):
prepack,
)
atol = rtol = precision[ref.dtype]
self.assertTrue(torch.allclose(ref, opt, atol=atol, rtol=rtol))
torch.testing.assert_close(ref, opt, atol=atol, rtol=rtol)
def test_fp8_gemm(self):
for params in itertools.product(

View File

@@ -91,9 +91,7 @@ class TestFusedExperts(CustomTestCase):
fused_output = fused_moe(a, w1, w2, score, topk, renormalize, prepack)
atol = rtol = precision[torch_output.dtype]
self.assertTrue(
torch.allclose(torch_output, fused_output, atol=atol, rtol=rtol)
)
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
def test_bf16_moe(self):
for params in itertools.product(
@@ -171,7 +169,7 @@ class TestFusedExperts(CustomTestCase):
# Increase the tolerance for large input shapes
if M > 35:
atol = rtol = 0.02
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def test_int8_moe(self):
for params in itertools.product(
@@ -235,7 +233,7 @@ class TestFusedExperts(CustomTestCase):
)
atol = rtol = precision[dtype]
self.assertTrue(torch.allclose(ref_out.bfloat16(), out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
def test_fp8_moe(self):
for params in itertools.product(

View File

@@ -47,7 +47,7 @@ class TestNorm(CustomTestCase):
ref_out = self._forward_native(x, weight, variance_epsilon)
atol = rtol = precision[ref_out.dtype]
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
ref_x = x.clone()
residual = torch.randn([m, hidden_size], dtype=dtype)
@@ -61,8 +61,8 @@ class TestNorm(CustomTestCase):
ref_x, weight, variance_epsilon, ref_residual
)
self.assertTrue(torch.allclose(x, ref_x, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(residual, ref_residual, atol=atol, rtol=rtol))
torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
def _l2norm_test(self, m, n, dtype):
@@ -75,7 +75,7 @@ class TestNorm(CustomTestCase):
ref_out = self._forward_native(x, fake_ones_weight, variance_epsilon)
atol = rtol = precision[ref_out.dtype]
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def test_norm(self):
for params in itertools.product(self.M, self.N, self.dtype):

View File

@@ -211,12 +211,12 @@ class TestQKVProjWithROPE(CustomTestCase):
qk_rope_head_dim,
)
atol = rtol = precision[q_ref.dtype]
self.assertTrue(torch.allclose(q_ref, q_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(k_ref, k_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(v_ref, v_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(fused_q_out, q_out))
self.assertTrue(torch.allclose(fused_k_out, k_out))
self.assertTrue(torch.allclose(fused_v_out, v_out))
torch.testing.assert_close(q_ref, q_out, atol=atol, rtol=rtol)
torch.testing.assert_close(k_ref, k_out, atol=atol, rtol=rtol)
torch.testing.assert_close(v_ref, v_out, atol=atol, rtol=rtol)
torch.testing.assert_close(fused_q_out, q_out)
torch.testing.assert_close(fused_k_out, k_out)
torch.testing.assert_close(fused_v_out, v_out)
def test_int8_qkv_proj_with_rope(self):
dtype = torch.bfloat16
@@ -302,12 +302,12 @@ class TestQKVProjWithROPE(CustomTestCase):
qk_rope_head_dim,
)
atol = rtol = precision[q_ref.dtype]
self.assertTrue(torch.allclose(q_ref, q_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(k_ref, k_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(v_ref, v_out, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(fused_q_out, q_out))
self.assertTrue(torch.allclose(fused_k_out, k_out))
self.assertTrue(torch.allclose(fused_v_out, v_out))
torch.testing.assert_close(q_ref, q_out, atol=atol, rtol=rtol)
torch.testing.assert_close(k_ref, k_out, atol=atol, rtol=rtol)
torch.testing.assert_close(v_ref, v_out, atol=atol, rtol=rtol)
torch.testing.assert_close(fused_q_out, q_out)
torch.testing.assert_close(fused_k_out, k_out)
torch.testing.assert_close(fused_v_out, v_out)
def test_fp8_qkv_proj_with_rope(self):
dtype = torch.bfloat16

View File

@@ -75,8 +75,8 @@ class TestROPE(CustomTestCase):
)
atol = rtol = precision[q_pe.dtype]
self.assertTrue(torch.allclose(q_pe, q_pe_clone, atol=atol, rtol=rtol))
self.assertTrue(torch.allclose(k_pe, k_pe_clone, atol=atol, rtol=rtol))
torch.testing.assert_close(q_pe, q_pe_clone, atol=atol, rtol=rtol)
torch.testing.assert_close(k_pe, k_pe_clone, atol=atol, rtol=rtol)
torch.testing.assert_close(k_pe, k_pe_clone)
def test_origin_rope(self):

View File

@@ -71,7 +71,7 @@ class TestSharedExpert(CustomTestCase):
)
atol = rtol = precision[ref.dtype]
self.assertTrue(torch.allclose(ref, res, atol=atol, rtol=rtol))
torch.testing.assert_close(ref, res, atol=atol, rtol=rtol)
def test_bf16_shared_expert(self):
for params in itertools.product(
@@ -129,7 +129,7 @@ class TestSharedExpert(CustomTestCase):
)
atol = rtol = precision[ref2.dtype]
self.assertTrue(torch.allclose(ref2, res2, atol=atol, rtol=rtol))
torch.testing.assert_close(ref2, res2, atol=atol, rtol=rtol)
def test_int8_shared_expert(self):
for params in itertools.product(
@@ -199,7 +199,7 @@ class TestSharedExpert(CustomTestCase):
)
atol = rtol = precision[ref_out.dtype]
self.assertTrue(torch.allclose(ref_out, out, atol=atol, rtol=rtol))
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def test_fp8_shared_expert(self):
for params in itertools.product(

View File

@@ -8,8 +8,8 @@ from utils import precision
from sglang.srt.layers.moe.topk import (
biased_grouped_topk_impl as native_biased_grouped_topk,
)
from sglang.srt.layers.moe.topk import fused_topk_native as native_fused_topk
from sglang.srt.layers.moe.topk import grouped_topk as native_grouped_topk
from sglang.srt.layers.moe.topk import fused_topk_torch_native as native_fused_topk
from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk
from sglang.srt.models.llama4 import Llama4MoE
from sglang.test.test_utils import CustomTestCase

View File

@@ -175,9 +175,11 @@ suites = {
TestFile("cpu/test_decode.py"),
TestFile("cpu/test_extend.py"),
TestFile("cpu/test_gemm.py"),
TestFile("cpu/test_mla.py"),
TestFile("cpu/test_moe.py"),
TestFile("cpu/test_norm.py"),
TestFile("cpu/test_qkv_proj_with_rope.py"),
TestFile("cpu/test_rope.py"),
TestFile("cpu/test_shared_expert.py"),
],
"nightly": [