Mohammad Miadh AngkadandMohammad Angkad
a711785475
Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
import pytest
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.quantization.unquant import (
|
|
_BF16_SPLITK_TUNED_TACTICS,
|
|
Bf16GemmBackend,
|
|
should_enable_bf16_splitk_gemm,
|
|
use_bf16_splitk_gemm,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
|
|
|
|
|
@pytest.mark.parametrize("m,n,k", _BF16_SPLITK_TUNED_TACTICS)
|
|
def test_splitk_selects_tuned_oakhaven_shape(m: int, n: int, k: int):
|
|
assert use_bf16_splitk_gemm(m, n, k)
|
|
|
|
|
|
@pytest.mark.parametrize("m", [0, 33, 64])
|
|
@pytest.mark.parametrize("n,k", [(256, 8192), (512, 8192), (2304, 8192), (2560, 8192)])
|
|
def test_splitk_keeps_large_m_on_existing_path(m: int, n: int, k: int):
|
|
assert not use_bf16_splitk_gemm(m, n, k)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"shape",
|
|
[
|
|
(1, 1024, 2048),
|
|
(3, 256, 8192),
|
|
(16, 8192, 4096),
|
|
(32, 4096, 8192),
|
|
],
|
|
)
|
|
def test_splitk_rejects_unmeasured_shapes(shape: tuple[int, int, int]):
|
|
assert not use_bf16_splitk_gemm(*shape)
|
|
|
|
|
|
def test_bf16_splitk_is_enabled_by_default():
|
|
assert envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.default is True
|
|
with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(True):
|
|
assert should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL)
|
|
|
|
|
|
def test_bf16_splitk_env_kill_switch():
|
|
with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(False):
|
|
assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL)
|
|
|
|
|
|
def test_bf16_splitk_does_not_override_torch_backend():
|
|
with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(True):
|
|
assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.TORCH)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main([__file__, "-v"]))
|