[MUSA][17/N] ci: Add MUSA diffusion, sgl-kernel tests, and CI workflow support (#20672)
Co-authored-by: ximin.chen <ximin.chen@mthreads.com> Co-authored-by: R0CKSTAR <xiaodong.ye@mthreads.com>
This commit is contained in:
johnnycxm
co-authored by
ximin.chen
R0CKSTAR
ximin.chen
R0CKSTAR
parent
15e6572f21
commit
cdf5771f91
@@ -0,0 +1,285 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Tests for MUSA-specific RMSNorm custom op.
|
||||
|
||||
These tests call forward_musa directly and compare against forward_native
|
||||
as the reference implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# We need the MUSA platform to be available for these tests.
|
||||
# Skip the entire module if MUSA is not available.
|
||||
_musa_available = hasattr(torch, "musa") and torch.musa.is_available()
|
||||
pytestmark = pytest.mark.skipif(not _musa_available, reason="MUSA device not available")
|
||||
|
||||
# Use a fixed seed for reproducibility
|
||||
SEED = 42
|
||||
|
||||
|
||||
def get_musa_device():
|
||||
return torch.device("musa:0")
|
||||
|
||||
|
||||
class TestRMSNorm:
|
||||
"""Tests for RMSNorm.forward_musa vs forward_native."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self):
|
||||
self.device = get_musa_device()
|
||||
|
||||
def _make_norm(self, hidden_size, eps=1e-6, var_hidden_size=None):
|
||||
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
|
||||
|
||||
norm = RMSNorm(hidden_size, eps=eps, var_hidden_size=var_hidden_size)
|
||||
norm = norm.to(self.device)
|
||||
return norm
|
||||
|
||||
# --- Basic correctness: no residual ---
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_size",
|
||||
[64, 128, 256, 512, 1024, 2048],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_no_residual_matches_native(self, hidden_size, dtype):
|
||||
"""forward_musa without residual should match forward_native."""
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(4, hidden_size, dtype=dtype, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
|
||||
atol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-4
|
||||
rtol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-4
|
||||
torch.testing.assert_close(out_musa, out_native, atol=atol, rtol=rtol)
|
||||
|
||||
# --- With residual ---
|
||||
@pytest.mark.parametrize(
|
||||
"hidden_size",
|
||||
[64, 128, 256, 512, 1024],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_with_residual_matches_native(self, hidden_size, dtype):
|
||||
"""forward_musa with residual should match forward_native."""
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(4, hidden_size, dtype=dtype, device=self.device)
|
||||
residual = torch.randn(4, hidden_size, dtype=dtype, device=self.device)
|
||||
|
||||
# Clone inputs since forward_musa modifies them in-place
|
||||
x_musa, res_musa = x.clone(), residual.clone()
|
||||
x_native, res_native = x.clone(), residual.clone()
|
||||
|
||||
out_musa, res_out_musa = norm.forward_musa(x_musa, res_musa)
|
||||
out_native, res_out_native = norm.forward_native(x_native, res_native)
|
||||
|
||||
atol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-4
|
||||
rtol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-4
|
||||
torch.testing.assert_close(out_musa, out_native, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(res_out_musa, res_out_native, atol=atol, rtol=rtol)
|
||||
|
||||
# --- 3D input shapes ---
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 1, 128),
|
||||
(2, 8, 128),
|
||||
(4, 16, 256),
|
||||
(2, 32, 512),
|
||||
],
|
||||
ids=lambda s: f"shape={'x'.join(map(str, s))}",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
|
||||
def test_3d_input_no_residual(self, shape, dtype):
|
||||
"""forward_musa should handle 3D inputs correctly."""
|
||||
hidden_size = shape[-1]
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(shape, dtype=dtype, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
|
||||
atol = 1e-2 if dtype == torch.float16 else 1e-4
|
||||
rtol = 1e-2 if dtype == torch.float16 else 1e-4
|
||||
torch.testing.assert_close(out_musa, out_native, atol=atol, rtol=rtol)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(2, 8, 128),
|
||||
(4, 16, 256),
|
||||
],
|
||||
ids=lambda s: f"shape={'x'.join(map(str, s))}",
|
||||
)
|
||||
def test_3d_input_with_residual(self, shape):
|
||||
"""forward_musa should handle 3D inputs with residual correctly."""
|
||||
hidden_size = shape[-1]
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
dtype = torch.float32
|
||||
x = torch.randn(shape, dtype=dtype, device=self.device)
|
||||
residual = torch.randn(shape, dtype=dtype, device=self.device)
|
||||
|
||||
x_musa, res_musa = x.clone(), residual.clone()
|
||||
x_native, res_native = x.clone(), residual.clone()
|
||||
|
||||
out_musa, res_out_musa = norm.forward_musa(x_musa, res_musa)
|
||||
out_native, res_out_native = norm.forward_native(x_native, res_native)
|
||||
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(res_out_musa, res_out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
# --- Non-contiguous input ---
|
||||
def test_non_contiguous_input_no_residual(self):
|
||||
"""forward_musa should handle non-contiguous inputs (makes them contiguous)."""
|
||||
hidden_size = 128
|
||||
norm = self._make_norm(hidden_size)
|
||||
# Create non-contiguous tensor via slicing
|
||||
x_base = torch.randn(8, hidden_size, dtype=torch.float32, device=self.device)
|
||||
x = x_base[::2] # shape (4, 128), non-contiguous
|
||||
assert not x.is_contiguous()
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_non_contiguous_input_with_residual(self):
|
||||
"""forward_musa should handle non-contiguous inputs with residual."""
|
||||
hidden_size = 128
|
||||
norm = self._make_norm(hidden_size)
|
||||
x_base = torch.randn(8, hidden_size, dtype=torch.float32, device=self.device)
|
||||
x = x_base[::2] # non-contiguous
|
||||
assert not x.is_contiguous()
|
||||
residual = torch.randn(4, hidden_size, dtype=torch.float32, device=self.device)
|
||||
|
||||
x_musa, res_musa = x.clone(), residual.clone()
|
||||
x_native, res_native = x.clone(), residual.clone()
|
||||
|
||||
out_musa, res_out_musa = norm.forward_musa(x_musa, res_musa)
|
||||
out_native, res_out_native = norm.forward_native(x_native, res_native)
|
||||
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(res_out_musa, res_out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
# --- Output properties ---
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_output_dtype_preserved_no_residual(self, dtype):
|
||||
"""Output dtype should match input dtype when no residual."""
|
||||
norm = self._make_norm(128)
|
||||
x = torch.randn(4, 128, dtype=dtype, device=self.device)
|
||||
out = norm.forward_musa(x)
|
||||
assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_output_dtype_preserved_with_residual(self, dtype):
|
||||
"""Output and residual dtype should match input dtype."""
|
||||
norm = self._make_norm(128)
|
||||
x = torch.randn(4, 128, dtype=dtype, device=self.device)
|
||||
residual = torch.randn(4, 128, dtype=dtype, device=self.device)
|
||||
out, res_out = norm.forward_musa(x, residual)
|
||||
assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
|
||||
assert res_out.dtype == dtype, f"Expected {dtype}, got {res_out.dtype}"
|
||||
|
||||
def test_output_device_preserved(self):
|
||||
"""Output should remain on the same MUSA device."""
|
||||
norm = self._make_norm(128)
|
||||
x = torch.randn(4, 128, dtype=torch.float32, device=self.device)
|
||||
out = norm.forward_musa(x)
|
||||
assert out.device == x.device
|
||||
|
||||
# --- Epsilon sensitivity ---
|
||||
@pytest.mark.parametrize("eps", [1e-5, 1e-6, 1e-8])
|
||||
def test_different_epsilon(self, eps):
|
||||
"""Different epsilon values should produce consistent results."""
|
||||
norm = self._make_norm(128, eps=eps)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(4, 128, dtype=torch.float32, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
# --- Weight initialization ---
|
||||
def test_custom_weight(self):
|
||||
"""RMSNorm with non-default weights should still match native."""
|
||||
norm = self._make_norm(128)
|
||||
# Set custom weights
|
||||
with torch.no_grad():
|
||||
norm.weight.fill_(2.0)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(4, 128, dtype=torch.float32, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_random_weight(self):
|
||||
"""RMSNorm with random weights should still match native."""
|
||||
norm = self._make_norm(256)
|
||||
with torch.no_grad():
|
||||
norm.weight.copy_(torch.randn(256, device=self.device))
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(8, 256, dtype=torch.float32, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
# --- Dispatch test ---
|
||||
def test_dispatch_calls_forward_musa(self):
|
||||
"""On MUSA platform, dispatch_forward should select forward_musa."""
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
if current_platform.is_musa():
|
||||
norm = self._make_norm(128)
|
||||
assert norm._forward_method == norm.forward_musa
|
||||
|
||||
# --- Large hidden size ---
|
||||
def test_large_hidden_size(self):
|
||||
"""Test with a large hidden size (float32)."""
|
||||
hidden_size = 4096
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(2, hidden_size, dtype=torch.float32, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_large_hidden_size_half(self):
|
||||
"""Test with a large hidden size (float16)."""
|
||||
hidden_size = 4096
|
||||
norm = self._make_norm(hidden_size)
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(2, hidden_size, dtype=torch.float16, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# --- Single token ---
|
||||
def test_single_token(self):
|
||||
"""Test with a single token (batch_size=1)."""
|
||||
norm = self._make_norm(128)
|
||||
x = torch.randn(1, 128, dtype=torch.float32, device=self.device)
|
||||
|
||||
out_musa = norm.forward_musa(x.clone())
|
||||
out_native = norm.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_single_token_with_residual(self):
|
||||
"""Test with a single token and residual."""
|
||||
norm = self._make_norm(128)
|
||||
x = torch.randn(1, 128, dtype=torch.float32, device=self.device)
|
||||
residual = torch.randn(1, 128, dtype=torch.float32, device=self.device)
|
||||
|
||||
x_musa, res_musa = x.clone(), residual.clone()
|
||||
x_native, res_native = x.clone(), residual.clone()
|
||||
|
||||
out_musa, res_out_musa = norm.forward_musa(x_musa, res_musa)
|
||||
out_native, res_out_native = norm.forward_native(x_native, res_native)
|
||||
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(res_out_musa, res_out_native, atol=1e-4, rtol=1e-4)
|
||||
@@ -0,0 +1,119 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Tests for MUSA-specific SiluAndMul custom op.
|
||||
|
||||
These tests call forward_musa directly and compare against forward_native
|
||||
as the reference implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# We need the MUSA platform to be available for these tests.
|
||||
# Skip the entire module if MUSA is not available.
|
||||
_musa_available = hasattr(torch, "musa") and torch.musa.is_available()
|
||||
pytestmark = pytest.mark.skipif(not _musa_available, reason="MUSA device not available")
|
||||
|
||||
# Use a fixed seed for reproducibility
|
||||
SEED = 42
|
||||
|
||||
|
||||
def get_musa_device():
|
||||
return torch.device("musa:0")
|
||||
|
||||
|
||||
class TestSiluAndMul:
|
||||
"""Tests for SiluAndMul.forward_musa vs forward_native."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self):
|
||||
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
|
||||
|
||||
self.op = SiluAndMul()
|
||||
self.device = get_musa_device()
|
||||
|
||||
# --- Shape parametrization ---
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[
|
||||
(1, 64), # minimal 2D
|
||||
(4, 128), # small 2D
|
||||
(32, 256), # medium 2D
|
||||
(128, 1024), # large 2D
|
||||
(1, 1, 64), # minimal 3D
|
||||
(2, 8, 128), # small 3D
|
||||
(4, 16, 512), # medium 3D
|
||||
(2, 32, 2048), # large 3D
|
||||
],
|
||||
ids=lambda s: f"shape={'x'.join(map(str, s))}",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_forward_matches_native(self, shape, dtype):
|
||||
"""forward_musa output should match forward_native within tolerance."""
|
||||
torch.manual_seed(SEED)
|
||||
x = torch.randn(shape, dtype=dtype, device=self.device)
|
||||
x_native = x.clone().detach()
|
||||
|
||||
out_musa = self.op.forward_musa(x)
|
||||
out_native = self.op.forward_native(x_native)
|
||||
|
||||
expected_last_dim = shape[-1] // 2
|
||||
assert out_musa.shape == out_native.shape
|
||||
assert out_musa.shape[-1] == expected_last_dim
|
||||
|
||||
atol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-5
|
||||
rtol = 1e-2 if dtype in (torch.float16, torch.bfloat16) else 1e-5
|
||||
torch.testing.assert_close(out_musa, out_native, atol=atol, rtol=rtol)
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
|
||||
def test_output_dtype_preserved(self, dtype):
|
||||
"""Output dtype should match input dtype."""
|
||||
x = torch.randn(4, 128, dtype=dtype, device=self.device)
|
||||
out = self.op.forward_musa(x)
|
||||
assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
|
||||
|
||||
def test_output_device_preserved(self):
|
||||
"""Output should remain on the same MUSA device."""
|
||||
x = torch.randn(4, 128, dtype=torch.float16, device=self.device)
|
||||
out = self.op.forward_musa(x)
|
||||
assert out.device == x.device
|
||||
|
||||
def test_zeros_input(self):
|
||||
"""silu(0) * 0 = 0, so output should be all zeros."""
|
||||
x = torch.zeros(4, 128, dtype=torch.float32, device=self.device)
|
||||
out = self.op.forward_musa(x)
|
||||
torch.testing.assert_close(
|
||||
out, torch.zeros(4, 64, dtype=torch.float32, device=self.device)
|
||||
)
|
||||
|
||||
def test_large_values(self):
|
||||
"""Test with large magnitude inputs to check numerical stability."""
|
||||
x = torch.randn(4, 128, dtype=torch.float32, device=self.device) * 100.0
|
||||
out_musa = self.op.forward_musa(x)
|
||||
out_native = self.op.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-2, rtol=1e-2)
|
||||
|
||||
def test_non_contiguous_input(self):
|
||||
"""forward_musa should handle non-contiguous inputs correctly."""
|
||||
# Create non-contiguous tensor via transpose
|
||||
x_base = torch.randn(128, 4, dtype=torch.float32, device=self.device)
|
||||
x = x_base.t() # shape (4, 128), non-contiguous
|
||||
assert not x.is_contiguous()
|
||||
|
||||
out_musa = self.op.forward_musa(x)
|
||||
out_native = self.op.forward_native(x.clone())
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_single_element_last_dim(self):
|
||||
"""Edge case: last dim = 2 (d=1)."""
|
||||
x = torch.randn(4, 2, dtype=torch.float32, device=self.device)
|
||||
out_musa = self.op.forward_musa(x)
|
||||
out_native = self.op.forward_native(x.clone())
|
||||
assert out_musa.shape[-1] == 1
|
||||
torch.testing.assert_close(out_musa, out_native, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_dispatch_calls_forward_musa(self):
|
||||
"""On MUSA platform, dispatch_forward should select forward_musa."""
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
if current_platform.is_musa():
|
||||
assert self.op._forward_method == self.op.forward_musa
|
||||
@@ -0,0 +1,268 @@
|
||||
"""
|
||||
Test runner for multimodal_gen MUSA suites that manages partitioned execution.
|
||||
|
||||
Usage:
|
||||
python3 run_suite_musa.py --suite <suite_name> --partition-id <id> --total-partitions <num>
|
||||
|
||||
Example:
|
||||
python3 run_suite_musa.py --suite 1-gpu-musa --partition-id 0 --total-partitions 2
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import tabulate
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
SUITES = {
|
||||
"1-gpu-musa": [
|
||||
"musa/test_server_a_musa.py",
|
||||
"musa/test_server_b_musa.py",
|
||||
],
|
||||
"2-gpu-musa": [
|
||||
"musa/test_server_2_gpu_a_musa.py",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Run multimodal_gen MUSA test suite")
|
||||
parser.add_argument(
|
||||
"--suite",
|
||||
type=str,
|
||||
required=True,
|
||||
choices=list(SUITES.keys()),
|
||||
help="The test suite to run (valid names are defined in SUITES)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--partition-id",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Index of the current partition (for parallel execution)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--total-partitions",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Total number of partitions",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base-dir",
|
||||
type=str,
|
||||
default="server",
|
||||
help="Base directory for tests relative to this script's parent",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-k",
|
||||
"--filter",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Pytest filter expression (passed to pytest -k)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--continue-on-error",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Continue running remaining tests even if one fails.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def collect_test_items(files, filter_expr=None):
|
||||
"""Collect test item node IDs from the given files using pytest --collect-only."""
|
||||
cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"]
|
||||
if filter_expr:
|
||||
cmd.extend(["-k", filter_expr])
|
||||
cmd.extend(files)
|
||||
|
||||
print(f"Collecting tests with command: {' '.join(cmd)}")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
if result.returncode not in (0, 5):
|
||||
error_msg = (
|
||||
f"pytest --collect-only failed with exit code {result.returncode}\n"
|
||||
f"Command: {' '.join(cmd)}\n"
|
||||
)
|
||||
if result.stderr:
|
||||
error_msg += f"stderr:\n{result.stderr}\n"
|
||||
if result.stdout:
|
||||
error_msg += f"stdout:\n{result.stdout}\n"
|
||||
logger.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if result.returncode == 5:
|
||||
print(
|
||||
"No tests were collected (exit code 5). This may be expected with filters."
|
||||
)
|
||||
|
||||
test_items = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
line = line.strip()
|
||||
if line and "::" in line and not line.startswith(("=", "-", " ")):
|
||||
test_id = line.split()[0] if " " in line else line
|
||||
if "::" in test_id:
|
||||
test_items.append(test_id)
|
||||
|
||||
print(f"Collected {len(test_items)} test items")
|
||||
return test_items
|
||||
|
||||
|
||||
def run_pytest(files, filter_expr=None, exitfirst=False):
|
||||
if not files:
|
||||
print("No files to run.")
|
||||
return 0
|
||||
|
||||
base_cmd = [sys.executable, "-m", "pytest", "-s", "-v"]
|
||||
if exitfirst:
|
||||
base_cmd.append("-x")
|
||||
|
||||
if filter_expr:
|
||||
base_cmd.extend(["-k", filter_expr])
|
||||
|
||||
max_retries = 6
|
||||
for i in range(max_retries + 1):
|
||||
cmd = list(base_cmd)
|
||||
if i > 0:
|
||||
cmd.append("--last-failed")
|
||||
cmd.extend(files)
|
||||
|
||||
if i > 0:
|
||||
print(
|
||||
f"Performance assertion failed. Retrying ({i}/{max_retries}) with --last-failed..."
|
||||
)
|
||||
|
||||
print(f"Running command: {' '.join(cmd)}")
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
bufsize=0,
|
||||
)
|
||||
|
||||
output_bytes = bytearray()
|
||||
while True:
|
||||
chunk = process.stdout.read(4096)
|
||||
if not chunk:
|
||||
break
|
||||
sys.stdout.buffer.write(chunk)
|
||||
sys.stdout.buffer.flush()
|
||||
output_bytes.extend(chunk)
|
||||
|
||||
process.wait()
|
||||
returncode = process.returncode
|
||||
|
||||
if returncode == 0:
|
||||
return 0
|
||||
|
||||
if returncode == 5:
|
||||
print(
|
||||
"No tests collected (exit code 5). This is expected when filters "
|
||||
"deselect all tests in a partition. Treating as success."
|
||||
)
|
||||
return 0
|
||||
|
||||
full_output = output_bytes.decode("utf-8", errors="replace")
|
||||
is_perf_assertion = (
|
||||
"multimodal_gen/test/server/test_server_utils.py" in full_output
|
||||
and "AssertionError" in full_output
|
||||
)
|
||||
is_flaky_ci_assertion = (
|
||||
"SafetensorError" in full_output
|
||||
or "FileNotFoundError" in full_output
|
||||
or "TimeoutError" in full_output
|
||||
)
|
||||
is_oom_error = (
|
||||
"out of memory" in full_output.lower()
|
||||
or "oom killer" in full_output.lower()
|
||||
)
|
||||
|
||||
if not (is_perf_assertion or is_flaky_ci_assertion or is_oom_error):
|
||||
return returncode
|
||||
|
||||
print("Max retry exceeded")
|
||||
return returncode
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
current_file_path = Path(__file__).resolve()
|
||||
test_root_dir = current_file_path.parent
|
||||
target_dir = test_root_dir / args.base_dir
|
||||
|
||||
if not target_dir.exists():
|
||||
print(f"Error: Target directory {target_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
|
||||
suite_files_rel = SUITES[args.suite]
|
||||
suite_files_abs = []
|
||||
for rel_path in suite_files_rel:
|
||||
abs_path = target_dir / rel_path
|
||||
if not abs_path.exists():
|
||||
print(f"Warning: Test file {rel_path} not found in {target_dir}. Skipping.")
|
||||
continue
|
||||
suite_files_abs.append(str(abs_path))
|
||||
|
||||
if not suite_files_abs:
|
||||
print(f"No valid test files found for suite '{args.suite}'.")
|
||||
sys.exit(0)
|
||||
|
||||
all_test_items = collect_test_items(suite_files_abs, filter_expr=args.filter)
|
||||
if not all_test_items:
|
||||
print(f"No test items found for suite '{args.suite}'.")
|
||||
sys.exit(0)
|
||||
|
||||
my_items = [
|
||||
item
|
||||
for i, item in enumerate(all_test_items)
|
||||
if i % args.total_partitions == args.partition_id
|
||||
]
|
||||
|
||||
partition_info = (
|
||||
f"{args.partition_id + 1}/{args.total_partitions} "
|
||||
f"(0-based id={args.partition_id})"
|
||||
)
|
||||
rows = [[args.suite, partition_info]]
|
||||
msg = (
|
||||
tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql") + "\n"
|
||||
)
|
||||
msg += f"Enabled {len(my_items)} test(s):\n"
|
||||
for item in my_items:
|
||||
msg += f" - {item}\n"
|
||||
print(msg, flush=True)
|
||||
print(
|
||||
f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}"
|
||||
)
|
||||
print(f"Selected {len(suite_files_abs)} files:")
|
||||
for file_path in suite_files_abs:
|
||||
print(f" - {os.path.basename(file_path)}")
|
||||
|
||||
if not my_items:
|
||||
print("No items assigned to this partition. Exiting success.")
|
||||
sys.exit(0)
|
||||
|
||||
print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}")
|
||||
exit_code = run_pytest(my_items, exitfirst=not args.continue_on_error)
|
||||
|
||||
msg = (
|
||||
"\n"
|
||||
+ tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")
|
||||
+ "\n"
|
||||
)
|
||||
msg += f"Executed {len(my_items)} test(s):\n"
|
||||
for item in my_items:
|
||||
msg += f" - {item}\n"
|
||||
print(msg, flush=True)
|
||||
|
||||
sys.exit(exit_code)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,264 @@
|
||||
{
|
||||
"metadata":{
|
||||
"model":"Diffusion Server",
|
||||
"hardware":"CI S5000 pool",
|
||||
"description":"Reference numbers captured from the CI diffusion server baseline run"
|
||||
},
|
||||
"scenarios":{
|
||||
"qwen_image_t2i_musa":{
|
||||
"stages_ms":{
|
||||
"InputValidationStage":0.09,
|
||||
"TextEncodingStage":658.4,
|
||||
"LatentPreparationStage":0.33,
|
||||
"TimestepPreparationStage":24.39,
|
||||
"DenoisingStage":36196.6,
|
||||
"DecodingStage":40.44
|
||||
},
|
||||
"denoise_step_ms":{
|
||||
"0":666.68,
|
||||
"1":732.33,
|
||||
"2":721.29,
|
||||
"3":729.27,
|
||||
"4":725.05,
|
||||
"5":721.71,
|
||||
"6":722.22,
|
||||
"7":725.93,
|
||||
"8":724.94,
|
||||
"9":724.14,
|
||||
"10":730.43,
|
||||
"11":719.92,
|
||||
"12":726.24,
|
||||
"13":722.04,
|
||||
"14":727.68,
|
||||
"15":720.31,
|
||||
"16":721.75,
|
||||
"17":725.65,
|
||||
"18":720.23,
|
||||
"19":724.12,
|
||||
"20":726.35,
|
||||
"21":723.27,
|
||||
"22":731.58,
|
||||
"23":724.97,
|
||||
"24":721.48,
|
||||
"25":722.0,
|
||||
"26":722.37,
|
||||
"27":719.81,
|
||||
"28":721.64,
|
||||
"29":724.81,
|
||||
"30":723.9,
|
||||
"31":725.42,
|
||||
"32":719.86,
|
||||
"33":728.04,
|
||||
"34":728.55,
|
||||
"35":723.13,
|
||||
"36":722.0,
|
||||
"37":730.11,
|
||||
"38":724.06,
|
||||
"39":728.35,
|
||||
"40":728.04,
|
||||
"41":726.62,
|
||||
"42":728.47,
|
||||
"43":728.11,
|
||||
"44":728.59,
|
||||
"45":721.5,
|
||||
"46":724.59,
|
||||
"47":729.26,
|
||||
"48":726.05,
|
||||
"49":721.13
|
||||
},
|
||||
"expected_e2e_ms":37190.98,
|
||||
"expected_avg_denoise_ms":723.72,
|
||||
"expected_median_denoise_ms":724.7
|
||||
},
|
||||
"wan2_1_t2v_1.3b_musa":{
|
||||
"stages_ms":{
|
||||
"InputValidationStage":0.12,
|
||||
"TextEncodingStage":1097.75,
|
||||
"LatentPreparationStage":0.24,
|
||||
"TimestepPreparationStage":5.66,
|
||||
"DenoisingStage":47399.84,
|
||||
"DecodingStage":946.08,
|
||||
"per_frame_generation":null
|
||||
},
|
||||
"denoise_step_ms":{
|
||||
"0":783.06,
|
||||
"1":970.52,
|
||||
"2":939.72,
|
||||
"3":947.58,
|
||||
"4":941.44,
|
||||
"5":955.26,
|
||||
"6":960.39,
|
||||
"7":951.84,
|
||||
"8":959.68,
|
||||
"9":953.33,
|
||||
"10":940.87,
|
||||
"11":958.5,
|
||||
"12":952.7,
|
||||
"13":933.4,
|
||||
"14":952.0,
|
||||
"15":951.6,
|
||||
"16":947.04,
|
||||
"17":939.28,
|
||||
"18":956.88,
|
||||
"19":960.1,
|
||||
"20":949.73,
|
||||
"21":954.77,
|
||||
"22":959.98,
|
||||
"23":947.37,
|
||||
"24":957.51,
|
||||
"25":953.39,
|
||||
"26":953.73,
|
||||
"27":959.57,
|
||||
"28":942.59,
|
||||
"29":958.05,
|
||||
"30":952.76,
|
||||
"31":952.76,
|
||||
"32":950.6,
|
||||
"33":948.76,
|
||||
"34":957.53,
|
||||
"35":940.86,
|
||||
"36":958.11,
|
||||
"37":940.9,
|
||||
"38":949.1,
|
||||
"39":951.81,
|
||||
"40":948.61,
|
||||
"41":957.28,
|
||||
"42":951.41,
|
||||
"43":953.09,
|
||||
"44":955.69,
|
||||
"45":941.93,
|
||||
"46":952.96,
|
||||
"47":953.5,
|
||||
"48":939.25,
|
||||
"49":942.69
|
||||
},
|
||||
"expected_e2e_ms":50007.17,
|
||||
"expected_avg_denoise_ms":947.83,
|
||||
"expected_median_denoise_ms":952.35
|
||||
},
|
||||
"wan2_2_t2v_a14b_2gpu_musa": {
|
||||
"stages_ms": {
|
||||
"per_frame_generation": null,
|
||||
"TimestepPreparationStage": 4.8,
|
||||
"DenoisingStage": 281769.53,
|
||||
"DecodingStage": 1728.76,
|
||||
"InputValidationStage": 0.13,
|
||||
"TextEncodingStage": 1098.56,
|
||||
"LatentPreparationStage": 0.33
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 6262.73,
|
||||
"1": 7103.05,
|
||||
"2": 7039.74,
|
||||
"3": 7052.97,
|
||||
"4": 7046.74,
|
||||
"5": 7059.9,
|
||||
"6": 7048.75,
|
||||
"7": 7054.99,
|
||||
"8": 7052.75,
|
||||
"9": 7059.59,
|
||||
"10": 7056.92,
|
||||
"11": 7053.61,
|
||||
"12": 7063.48,
|
||||
"13": 7045.34,
|
||||
"14": 7051.22,
|
||||
"15": 7051.25,
|
||||
"16": 7048.83,
|
||||
"17": 7053.91,
|
||||
"18": 7060.47,
|
||||
"19": 7056.4,
|
||||
"20": 7055.39,
|
||||
"21": 7052.72,
|
||||
"22": 7054.84,
|
||||
"23": 7058.17,
|
||||
"24": 7052.26,
|
||||
"25": 7057.8,
|
||||
"26": 7362.31,
|
||||
"27": 7053.57,
|
||||
"28": 7044.42,
|
||||
"29": 7044.03,
|
||||
"30": 7056.59,
|
||||
"31": 7045.81,
|
||||
"32": 7051.59,
|
||||
"33": 7048.66,
|
||||
"34": 7050.84,
|
||||
"35": 7048.57,
|
||||
"36": 7056.21,
|
||||
"37": 7056.08,
|
||||
"38": 7053.32,
|
||||
"39": 7047.35
|
||||
},
|
||||
"expected_e2e_ms": 285727.55,
|
||||
"expected_avg_denoise_ms": 7043.96,
|
||||
"expected_median_denoise_ms": 7056.07
|
||||
},
|
||||
"wan2_1_i2v_14b_480P_2gpu_musa": {
|
||||
"stages_ms": {
|
||||
"InputValidationStage": 10.35,
|
||||
"TextEncodingStage": 1098.55,
|
||||
"ImageEncodingStage": 730.6,
|
||||
"ImageVAEEncodingStage": 1188.4,
|
||||
"LatentPreparationStage": 0.5,
|
||||
"TimestepPreparationStage": 7.99,
|
||||
"DenoisingStage": 134355.99,
|
||||
"DecodingStage": 1200.0,
|
||||
"per_frame_generation": null
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 2441.69,
|
||||
"1": 2711.82,
|
||||
"2": 2701.26,
|
||||
"3": 2692.91,
|
||||
"4": 2686.94,
|
||||
"5": 2685.35,
|
||||
"6": 2689.45,
|
||||
"7": 2685.49,
|
||||
"8": 2687.74,
|
||||
"9": 2690.0,
|
||||
"10": 2691.31,
|
||||
"11": 2691.65,
|
||||
"12": 2698.52,
|
||||
"13": 2689.37,
|
||||
"14": 2690.18,
|
||||
"15": 2691.48,
|
||||
"16": 2695.74,
|
||||
"17": 2691.77,
|
||||
"18": 2689.77,
|
||||
"19": 2690.6,
|
||||
"20": 2686.84,
|
||||
"21": 2694.9,
|
||||
"22": 2689.56,
|
||||
"23": 2696.38,
|
||||
"24": 2689.13,
|
||||
"25": 2686.38,
|
||||
"26": 2689.53,
|
||||
"27": 2695.91,
|
||||
"28": 2691.9,
|
||||
"29": 2691.55,
|
||||
"30": 2700.09,
|
||||
"31": 2691.45,
|
||||
"32": 2696.8,
|
||||
"33": 2689.42,
|
||||
"34": 2695.87,
|
||||
"35": 2690.31,
|
||||
"36": 2687.65,
|
||||
"37": 2697.62,
|
||||
"38": 2683.46,
|
||||
"39": 2692.41,
|
||||
"40": 2699.33,
|
||||
"41": 2695.25,
|
||||
"42": 2691.16,
|
||||
"43": 2687.13,
|
||||
"44": 2692.92,
|
||||
"45": 2690.4,
|
||||
"46": 2696.53,
|
||||
"47": 2689.51,
|
||||
"48": 2692.06,
|
||||
"49": 2680.93
|
||||
},
|
||||
"expected_e2e_ms": 138624.98,
|
||||
"expected_avg_denoise_ms": 2686.91,
|
||||
"expected_median_denoise_ms": 2691.24
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
"""
|
||||
MUSA-specific 2-GPU diffusion performance test.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
|
||||
TWO_GPU_MUSA_CASES_A,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401
|
||||
DiffusionServerBase,
|
||||
diffusion_server,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TestDiffusionServerTwoGpuMusaA(DiffusionServerBase):
|
||||
"""Performance tests for 2-GPU diffusion cases on MUSA."""
|
||||
|
||||
@pytest.fixture(params=TWO_GPU_MUSA_CASES_A, ids=lambda c: c.id)
|
||||
def case(self, request) -> DiffusionTestCase:
|
||||
"""Provide a DiffusionTestCase for each 2-GPU MUSA test."""
|
||||
return request.param
|
||||
@@ -0,0 +1,28 @@
|
||||
"""
|
||||
MUSA-specific diffusion performance test (1-GPU).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
|
||||
ONE_GPU_MUSA_CASES_A,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401
|
||||
DiffusionServerBase,
|
||||
diffusion_server,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TestDiffusionServerOneGpuMusaImage(DiffusionServerBase):
|
||||
"""Performance tests for 1-GPU diffusion cases on MUSA"""
|
||||
|
||||
@pytest.fixture(params=ONE_GPU_MUSA_CASES_A, ids=lambda c: c.id)
|
||||
def case(self, request) -> DiffusionTestCase:
|
||||
"""Provide a DiffusionTestCase for each 1-GPU MUSA test."""
|
||||
return request.param
|
||||
@@ -0,0 +1,28 @@
|
||||
"""
|
||||
MUSA-specific diffusion performance test (1-GPU).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
|
||||
ONE_GPU_MUSA_CASES_B,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401
|
||||
DiffusionServerBase,
|
||||
diffusion_server,
|
||||
)
|
||||
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TestDiffusionServerOneGpuMusaVideo(DiffusionServerBase):
|
||||
"""Performance tests for 1-GPU diffusion cases on MUSA"""
|
||||
|
||||
@pytest.fixture(params=ONE_GPU_MUSA_CASES_B, ids=lambda c: c.id)
|
||||
def case(self, request) -> DiffusionTestCase:
|
||||
"""Provide a DiffusionTestCase for each 1-GPU MUSA test."""
|
||||
return request.param
|
||||
@@ -0,0 +1,53 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sglang.multimodal_gen.test.server.testcase_configs import (
|
||||
T2V_PROMPT,
|
||||
DiffusionSamplingParams,
|
||||
DiffusionServerArgs,
|
||||
DiffusionTestCase,
|
||||
T2I_sampling_params,
|
||||
TI2V_sampling_params,
|
||||
)
|
||||
|
||||
ONE_GPU_MUSA_CASES_A: list[DiffusionTestCase] = [
|
||||
DiffusionTestCase(
|
||||
"qwen_image_t2i_musa",
|
||||
DiffusionServerArgs(
|
||||
model_path="Qwen/Qwen-Image",
|
||||
modality="image",
|
||||
),
|
||||
T2I_sampling_params,
|
||||
run_consistency_check=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
ONE_GPU_MUSA_CASES_B: list[DiffusionTestCase] = [
|
||||
DiffusionTestCase(
|
||||
"wan2_1_t2v_1.3b_musa",
|
||||
DiffusionServerArgs(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
modality="video",
|
||||
custom_validator="video",
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=T2V_PROMPT,
|
||||
),
|
||||
run_consistency_check=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
TWO_GPU_MUSA_CASES_A: list[DiffusionTestCase] = [
|
||||
DiffusionTestCase(
|
||||
"wan2_1_i2v_14b_480P_2gpu_musa",
|
||||
DiffusionServerArgs(
|
||||
model_path="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
|
||||
modality="video",
|
||||
custom_validator="video",
|
||||
num_gpus=2,
|
||||
),
|
||||
TI2V_sampling_params,
|
||||
run_consistency_check=False,
|
||||
),
|
||||
]
|
||||
@@ -492,6 +492,8 @@ def _with_default_num_gpus(
|
||||
|
||||
|
||||
# Load global configuration
|
||||
BASELINE_CONFIG = BaselineConfig.load(
|
||||
Path(__file__).with_name("perf_baselines.json")
|
||||
).update(Path(__file__).parent / "ascend" / "perf_baselines_npu.json")
|
||||
BASELINE_CONFIG = (
|
||||
BaselineConfig.load(Path(__file__).with_name("perf_baselines.json"))
|
||||
.update(Path(__file__).parent / "ascend" / "perf_baselines_npu.json")
|
||||
.update(Path(__file__).parent / "musa" / "perf_baselines_musa.json")
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user