diff --git a/.github/workflows/pr-test-musa.yml b/.github/workflows/pr-test-musa.yml new file mode 100644 index 000000000..21381c652 --- /dev/null +++ b/.github/workflows/pr-test-musa.yml @@ -0,0 +1,251 @@ +name: PR Test (MUSA) + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + input: + target_stage: + description: "Specific test stage to run (Optional)" + required: false + type: string + default: "" + workflow_call: + inputs: + ref: + description: 'Git ref (branch, tag, or SHA) to test. If not provided, uses the default branch.' + required: false + type: string + default: '' + run_all_tests: + description: "Run all tests (for releasing or testing purpose)" + required: false + type: boolean + default: false + +concurrency: + group: pr-test-musa-${{ inputs.ref || github.ref }} + cancel-in-progress: ${{ github.event_name != 'workflow_call' }} + +jobs: + # ==================== Check Changes ==================== # + check-changes: + runs-on: ubuntu-latest + outputs: + changes_exist: ${{ steps.filter.outputs.main_package == 'true' + || steps.filter.outputs.multimodal_gen == 'true' + || steps.filter.outputs.sgl_kernel == 'true' + || steps.run-mode.outputs.run_all_tests == 'true' + || inputs.target_stage != '' }} + main_package: ${{ steps.filter.outputs.main_package == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} + multimodal_gen: ${{ steps.filter.outputs.multimodal_gen == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} + sgl_kernel: ${{ steps.filter.outputs.sgl_kernel == 'true' || steps.run-mode.outputs.run_all_tests == 'true' }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Determine run mode + id: run-mode + run: | + # Run all tests for workflow_call (when ref input is provided) + # Note: github.event_name is inherited from caller, so we detect workflow_call by checking inputs.ref + if [[ "${{ inputs.run_all_tests }}" == "true" ]]; then + echo "run_all_tests=true" >> $GITHUB_OUTPUT + echo "Run mode: ALL TESTS (run_all_tests=${{ inputs.run_all_tests }})" + else + echo "run_all_tests=false" >> $GITHUB_OUTPUT + echo "Run mode: FILTERED (triggered by ${{ github.event_name }})" + fi + + - name: Detect file changes + id: filter + uses: dorny/paths-filter@v3 + if: steps.run-mode.outputs.run_all_tests != 'true' + with: + filters: | + main_package: + - "python/sglang/!(multimodal_gen)/**" + - "python/pyproject_other.toml" + - "scripts/ci/musa/*" + - "scripts/ci/utils/*" + - "test/**" + - ".github/workflows/pr-test-musa.yml" + multimodal_gen: + - "python/sglang/multimodal_gen/**" + - "python/sglang/cli/**" + - "python/pyproject_other.toml" + sgl_kernel: + - "sgl-kernel/**" + - ".github/workflows/pr-test-musa.yml" + + # ==================== PR Gate ==================== # + pr-gate: + needs: check-changes + if: needs.check-changes.outputs.changes_exist == 'true' + uses: ./.github/workflows/pr-gate.yml + secrets: inherit + + # ==================== Multimodal Gen Tests ==================== # + multimodal-gen-test-1-gpu-musa: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.multimodal_gen == 'true' || inputs.target_stage == 'multimodal-gen-test-1-gpu-musa' + strategy: + fail-fast: false + matrix: + part: [0, 1] + runs-on: s5000-1-gpu-runner + timeout-minutes: 240 + env: + USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + TORCHADA_ENABLE_CPP_OPS: 1 + HF_HUB_CACHE: /hf-cache/hub + steps: + - name: Checkout code + timeout-minutes: 10 + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Install dependencies + run: | + bash scripts/ci/musa/musa_install_dependency.sh + + - name: Run diffusion server tests (1-GPU) + timeout-minutes: 30 + env: + RUNAI_STREAMER_MEMORY_LIMIT: 0 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite_musa.py \ + --suite 1-gpu-musa \ + --partition-id ${{ matrix.part }} \ + --total-partitions 2 + + multimodal-gen-test-2-gpu-musa: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.multimodal_gen == 'true' || inputs.target_stage == 'multimodal-gen-test-2-gpu-musa' + runs-on: s5000-2-gpu-runner + timeout-minutes: 240 + env: + USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + TORCHADA_ENABLE_CPP_OPS: 1 + HF_HUB_CACHE: /hf-cache/hub + steps: + - name: Checkout code + timeout-minutes: 10 + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Install dependencies + run: | + bash scripts/ci/musa/musa_install_dependency.sh + + - name: Run diffusion server tests (2-GPU) + timeout-minutes: 30 + env: + RUNAI_STREAMER_MEMORY_LIMIT: 0 + run: | + cd python + python3 sglang/multimodal_gen/test/run_suite_musa.py \ + --suite 2-gpu-musa + + multimodal-gen-layer-unit-test-musa: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.multimodal_gen == 'true' || inputs.target_stage == 'multimodal-gen-layer-unit-test-musa' + runs-on: s5000-1-gpu-runner + timeout-minutes: 240 + env: + USE_MODELSCOPE: true + SGLANG_IS_IN_CI: true + TORCHADA_ENABLE_CPP_OPS: 1 + HF_HUB_CACHE: /hf-cache/hub + steps: + - name: Checkout code + timeout-minutes: 10 + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + bash scripts/ci/musa/musa_install_dependency.sh + + - name: Run multimodal gen layer unit test + timeout-minutes: 30 + run: | + pytest python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py + pytest python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py + +# =============================================== sgl-kernel ==================================================== + sgl-kernel-unit-test-musa: + needs: [check-changes, pr-gate] + if: needs.check-changes.outputs.sgl_kernel == 'true' || inputs.target_stage == 'sgl-kernel-unit-test-musa' + runs-on: s5000-1-gpu-runner + timeout-minutes: 240 + env: + TORCHADA_ENABLE_CPP_OPS: 1 + HF_HUB_CACHE: /hf-cache/hub + steps: + - name: Checkout code + timeout-minutes: 10 + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.ref }} + + - name: Install dependencies + run: | + bash scripts/ci/musa/musa_install_dependency.sh + + - name: Run sgl-kernel test + timeout-minutes: 20 + run: | + pytest sgl-kernel/tests/test_dsv3_router_gemm.py + pytest sgl-kernel/tests/test_per_token_quant_fp8.py + pytest sgl-kernel/tests/speculative/test_eagle_utils.py + pytest sgl-kernel/tests/speculative/test_ngram_utils.py + pytest sgl-kernel/tests/speculative/test_speculative_sampling.py + pytest sgl-kernel/tests/test_torch_defaults_reset.py + + + pr-test-musa-finish: + needs: + [ + pr-gate, + check-changes, + multimodal-gen-test-1-gpu-musa, + multimodal-gen-test-2-gpu-musa, + multimodal-gen-layer-unit-test-musa, + sgl-kernel-unit-test-musa, + ] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check all dependent job statuses + run: | + # Convert the 'needs' context to a JSON string + json_needs='${{ toJson(needs) }}' + + # Get a list of all job names from the JSON keys + job_names=$(echo "$json_needs" | jq -r 'keys_unsorted[]') + + for job in $job_names; do + # For each job, extract its result + result=$(echo "$json_needs" | jq -r --arg j "$job" '.[$j].result') + + # Print the job name and its result + echo "$job: $result" + + # Check for failure or cancellation and exit if found + if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then + echo "The above jobs failed." + exit 1 + fi + done + # If the loop completes, all jobs were successful + echo "All jobs completed successfully" + exit diff --git a/python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py b/python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py new file mode 100644 index 000000000..2c93e1ad3 --- /dev/null +++ b/python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py b/python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py new file mode 100644 index 000000000..1eaee717c --- /dev/null +++ b/python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/run_suite_musa.py b/python/sglang/multimodal_gen/test/run_suite_musa.py new file mode 100644 index 000000000..b8cd91e22 --- /dev/null +++ b/python/sglang/multimodal_gen/test/run_suite_musa.py @@ -0,0 +1,268 @@ +""" +Test runner for multimodal_gen MUSA suites that manages partitioned execution. + +Usage: + python3 run_suite_musa.py --suite --partition-id --total-partitions + +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() diff --git a/python/sglang/multimodal_gen/test/server/musa/perf_baselines_musa.json b/python/sglang/multimodal_gen/test/server/musa/perf_baselines_musa.json new file mode 100644 index 000000000..464544501 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/musa/perf_baselines_musa.json @@ -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 + } + } +} diff --git a/python/sglang/multimodal_gen/test/server/musa/test_server_2_gpu_a_musa.py b/python/sglang/multimodal_gen/test/server/musa/test_server_2_gpu_a_musa.py new file mode 100644 index 000000000..48909e1a2 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/musa/test_server_2_gpu_a_musa.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/musa/test_server_a_musa.py b/python/sglang/multimodal_gen/test/server/musa/test_server_a_musa.py new file mode 100644 index 000000000..7b4c4bc29 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/musa/test_server_a_musa.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/musa/test_server_b_musa.py b/python/sglang/multimodal_gen/test/server/musa/test_server_b_musa.py new file mode 100644 index 000000000..f961648fa --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/musa/test_server_b_musa.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py b/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py new file mode 100644 index 000000000..ca8363349 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py @@ -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, + ), +] diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 122f5f8cb..a140b3003 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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") +) diff --git a/scripts/ci/musa/musa_install_dependency.sh b/scripts/ci/musa/musa_install_dependency.sh index d3ef53d21..2cb90a324 100755 --- a/scripts/ci/musa/musa_install_dependency.sh +++ b/scripts/ci/musa/musa_install_dependency.sh @@ -1,5 +1,62 @@ #!/bin/bash set -euo pipefail +# Parse command line arguments +OPTIONAL_DEPS="" +SKIP_SGLANG_BUILD="" + +while [[ $# -gt 0 ]]; do + case $1 in + --skip-sglang-build) SKIP_SGLANG_BUILD="1"; shift;; + -h|--help) + echo "Usage: $0 [OPTIONS] [OPTIONAL_DEPS]" + echo "Options:" + echo " --skip-sglang-build Don't build checkout sglang, use what was shipped with the image" + exit 0 + ;; + *) + OPTIONAL_DEPS="$1" + shift + ;; + esac +done + PIP_INSTALL="python3 -m pip install --no-cache-dir" -${PIP_INSTALL} --upgrade pip setuptools torchada +${PIP_INSTALL} --upgrade pip setuptools torchada --user + +WHL_DIR="/sglang-checkout/whl" +if [ -d "$WHL_DIR" ] && compgen -G "${WHL_DIR}"/*.whl > /dev/null; then + echo "Uninstall old packages based on wheel METADATA..." + PKGS=$( + for whl in "${WHL_DIR}"/*.whl; do + meta_file=$(zipinfo -1 "$whl" | awk '/\.dist-info\/METADATA$/ {print; exit}') + [ -n "$meta_file" ] || continue + unzip -p "$whl" "$meta_file" 2>/dev/null | sed -n 's/^Name: //p' | head -n1 + done | sort -u + ) + for pkg in $PKGS; do + echo "Uninstalling $pkg" + pip uninstall -y "$pkg" || true + done + echo "Installing wheel files without dependency resolution..." + ${PIP_INSTALL} "${WHL_DIR}"/*.whl --user +fi + +if [ -n "$SKIP_SGLANG_BUILD" ]; then + echo "Didn't build checkout SGLang" + exit 0 +else + pip uninstall sgl-kernel -y || true + pip uninstall sglang -y || true + # Clear Python cache to ensure latest code is used (works for any env: venv, system, conda) + REPO_ROOT="${GITHUB_WORKSPACE:-$(pwd)}" + find "$REPO_ROOT" -name "*.pyc" -delete 2>/dev/null || true + find "$REPO_ROOT" -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true + + rm -f "${REPO_ROOT}/python/pyproject.toml" && mv "${REPO_ROOT}/python/pyproject_other.toml" "${REPO_ROOT}/python/pyproject.toml" + cd "${REPO_ROOT}" && ${PIP_INSTALL} -v -e "python[dev_musa]" --user + + cd "${REPO_ROOT}/sgl-kernel" + rm -f pyproject.toml && mv pyproject_musa.toml pyproject.toml && MTGPU_TARGET=mp_31 python3 setup_musa.py install --user + echo "$HOME/.local/bin" >> "$GITHUB_PATH" +fi diff --git a/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py index cbc80a607..9fe0e9311 100644 --- a/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py +++ b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py @@ -8,6 +8,23 @@ from sgl_kernel import topk_sigmoid from sglang.utils import is_in_ci +# Optional MUSA import +try: + from sglang.srt.utils import is_musa + + if is_musa(): + from sglang.srt.hardware_backend.musa.kernels.topk import ( + topk_sigmoid as musa_topk_sigmoid, + ) + + MUSA_AVAILABLE = True + else: + musa_topk_sigmoid = None + MUSA_AVAILABLE = False +except ImportError: + musa_topk_sigmoid = None + MUSA_AVAILABLE = False + IS_CI = is_in_ci() @@ -56,6 +73,28 @@ def sglang_topk_sigmoid( return topk_weights, topk_indices +def musa_topk_sigmoid_fn( + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + correction_bias: torch.Tensor = None, +): + num_tokens, num_experts = gating_output.shape + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + musa_topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + renormalize=renormalize, + correction_bias=correction_bias, + ) + + return topk_weights, topk_indices + + def get_topk_sigmoid_input(num_tokens, num_experts): gating_output = torch.randn( (num_tokens, num_experts), dtype=torch.float32, device="cuda" @@ -93,6 +132,28 @@ def calculate_diff(num_tokens, num_experts, topk): f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}" ) + if MUSA_AVAILABLE: + weights_musa, indices_musa = musa_topk_sigmoid_fn( + gating_output.clone(), + topk, + True, + correction_bias.clone(), + ) + weights_diff_musa = torch.abs(weights_sglang - weights_musa).mean().item() + indices_match_musa = torch.equal(indices_sglang, indices_musa) + + if ( + torch.allclose(weights_sglang, weights_musa, atol=1e-3, rtol=1e-3) + and indices_match_musa + ): + print("✅ SGLang and MUSA topk_sigmoid implementations match") + else: + print( + f"❌ MUSA vs SGLang differ: Weights diff={weights_diff_musa}, Indices match={indices_match_musa}" + ) + else: + print("⚠️ MUSA not available, skipping MUSA comparison") + # CI environment uses simplified parameters if IS_CI: @@ -107,11 +168,16 @@ else: configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) -# Filter providers based on vLLM availability +# Filter providers based on availability line_vals = ["sglang", "torch"] line_names = ["SGLang", "Torch"] styles = [("blue", "-"), ("green", "-")] +if MUSA_AVAILABLE: + line_vals.append("musa") + line_names.append("MUSA") + styles.append(("red", "-")) + @triton.testing.perf_report( triton.testing.Benchmark( @@ -144,6 +210,13 @@ def benchmark(num_tokens, num_experts, topk, provider): def fn(): return sglang_topk_sigmoid(gating_output, topk, True, correction_bias) + elif provider == "musa" or provider == "musa1": + if not MUSA_AVAILABLE: + return (0, 0, 0) + + def fn(): + return musa_topk_sigmoid_fn(gating_output, topk, True, correction_bias) + quantiles = [0.5, 0.2, 0.8] ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) diff --git a/sgl-kernel/benchmark/bench_moe_topk_softmax.py b/sgl-kernel/benchmark/bench_moe_topk_softmax.py index 4b2222405..451ae8d80 100644 --- a/sgl-kernel/benchmark/bench_moe_topk_softmax.py +++ b/sgl-kernel/benchmark/bench_moe_topk_softmax.py @@ -17,6 +17,23 @@ except ImportError: vllm_custom_ops = None VLLM_AVAILABLE = False +# Optional MUSA import +try: + from sglang.srt.utils import is_musa + + if is_musa(): + from sglang.srt.hardware_backend.musa.kernels.topk import ( + topk_softmax as musa_topk_softmax, + ) + + MUSA_AVAILABLE = True + else: + musa_topk_softmax = None + MUSA_AVAILABLE = False +except ImportError: + musa_topk_softmax = None + MUSA_AVAILABLE = False + IS_CI = is_in_ci() @@ -61,29 +78,62 @@ def sglang_topk_softmax(gating_output, topk): return topk_weights, topk_indices +def musa_topk_softmax_fn(gating_output, topk): + num_tokens, num_experts = gating_output.shape + + topk_weights = torch.empty( + (num_tokens, topk), device=gating_output.device, dtype=torch.float32 + ) + topk_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device=gating_output.device + ) + + musa_topk_softmax( + topk_weights, + topk_indices, + gating_output, + ) + + return topk_weights, topk_indices + + def calculate_diff(num_tokens, num_experts, topk): gating_output = torch.randn( (num_tokens, num_experts), device="cuda", dtype=torch.float32 ) - weights_vllm, indices_vllm = vllm_topk_softmax(gating_output.clone(), topk) weights_sglang, indices_sglang = sglang_topk_softmax(gating_output.clone(), topk) - weights_diff = torch.abs(weights_vllm - weights_sglang).mean().item() - indices_match = torch.equal(indices_vllm, indices_sglang) + if MUSA_AVAILABLE: + weights_musa, indices_musa = musa_topk_softmax_fn(gating_output.clone(), topk) + weights_diff = torch.abs(weights_sglang - weights_musa).mean().item() + indices_match = torch.equal(indices_sglang, indices_musa) - if not VLLM_AVAILABLE: - print("⚠️ vLLM not available, skipping comparison") - return - - if ( - torch.allclose(weights_vllm, weights_sglang, atol=1e-3, rtol=1e-3) - and indices_match - ): - print("✅ VLLM and SGLang topk_softmax implementations match") + if ( + torch.allclose(weights_sglang, weights_musa, atol=1e-3, rtol=1e-3) + and indices_match + ): + print("✅ SGLang and MUSA topk_softmax implementations match") + else: + print( + f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}" + ) else: - print( - f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}" - ) + print("⚠️ MUSA not available, skipping MUSA comparison") + + if VLLM_AVAILABLE: + weights_vllm, indices_vllm = vllm_topk_softmax(gating_output.clone(), topk) + weights_diff_vllm = torch.abs(weights_vllm - weights_sglang).mean().item() + indices_match_vllm = torch.equal(indices_vllm, indices_sglang) + + if ( + torch.allclose(weights_vllm, weights_sglang, atol=1e-3, rtol=1e-3) + and indices_match_vllm + ): + print("✅ VLLM and SGLang topk_softmax implementations match") + else: + print( + f"❌ VLLM vs SGLang differ: Weights diff={weights_diff_vllm}, Indices match={indices_match_vllm}" + ) # CI environment uses simplified parameters @@ -99,15 +149,20 @@ else: configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) -# Filter providers based on vLLM availability +# Filter providers based on availability +line_vals = ["sglang"] +line_names = ["SGLang"] +styles = [("blue", "-")] + if VLLM_AVAILABLE: - line_vals = ["sglang", "vllm"] - line_names = ["SGLang", "VLLM"] - styles = [("blue", "-"), ("green", "-")] -else: - line_vals = ["sglang"] - line_names = ["SGLang"] - styles = [("blue", "-")] + line_vals.append("vllm") + line_names.append("VLLM") + styles.append(("green", "-")) + +if MUSA_AVAILABLE: + line_vals.append("musa") + line_names.append("MUSA") + styles.append(("red", "-")) @triton.testing.perf_report( @@ -135,6 +190,10 @@ def benchmark(num_tokens, num_experts, topk, provider): fn = lambda: vllm_topk_softmax(gating_output, topk) elif provider == "sglang" or provider == "sglang1": fn = lambda: sglang_topk_softmax(gating_output, topk) + elif provider == "musa" or provider == "musa1": + if not MUSA_AVAILABLE: + return (0, 0, 0) + fn = lambda: musa_topk_softmax_fn(gating_output, topk) quantiles = [0.5, 0.2, 0.8] ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) diff --git a/sgl-kernel/tests/conftest.py b/sgl-kernel/tests/conftest.py index 4aae7ffbb..9d31aac8a 100644 --- a/sgl-kernel/tests/conftest.py +++ b/sgl-kernel/tests/conftest.py @@ -1,6 +1,11 @@ import pytest import torch +from sglang.srt.utils import is_musa + +if is_musa(): + import torchada # noqa: F401 + # This fixture ensures the torch defaults don't get left in modified states between # tests (e.g., when a test fails before restoring the original value), which