diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index f59970f6c..5c05d9640 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -152,6 +152,47 @@ jobs: docker exec -w /sglang-checkout/sgl-kernel/tests/sgl_diffusion ci_sglang python3 -m pytest test_timestep_embedding.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_sigmoid.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_torch_defaults_reset.py + + sgl-kernel-unit-test-2-gpu-amd: + needs: [check-changes] + if: | + always() && + ( + (inputs.target_stage == 'sgl-kernel-unit-test-2-gpu-amd') || + ( + !inputs.target_stage && + needs.check-changes.outputs.sgl_kernel == 'true' + ) + ) + strategy: + fail-fast: false + matrix: + runner: [linux-mi325-gpu-2] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd_ci_install_dependency.sh + + - name: Run test + timeout-minutes: 20 + run: | + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_deterministic_custom_allreduce.py + docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_amd_nccl_allreduce_determinism.py + # =============================================== primary ==================================================== stage-a-test-1-amd: @@ -842,6 +883,7 @@ jobs: check-changes, sgl-kernel-unit-test-amd, + sgl-kernel-unit-test-2-gpu-amd, multimodal-gen-test-1-gpu-amd, multimodal-gen-test-2-gpu-amd, diff --git a/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py b/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py index 10539fdc3..7e3e82c9b 100644 --- a/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py +++ b/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py @@ -12,12 +12,13 @@ This test compares: 2. Deterministic kernel (different batch size) Usage: - python test_amd_deterministic_custom_allreduce.py + pytest test_amd_deterministic_custom_allreduce.py """ import multiprocessing as mp import socket +import pytest import torch import torch.distributed as dist @@ -266,5 +267,14 @@ def main(): p.join() +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="Requires at least 2 CUDA GPUs", +) +def test_deterministic_custom_allreduce(): + """Test that deterministic custom all-reduce produces consistent results.""" + main() + + if __name__ == "__main__": main() diff --git a/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py b/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py index dc79d7455..492d54142 100644 --- a/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py +++ b/sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py @@ -10,12 +10,13 @@ This test compares: 2. Default all-reduce (different batch size) - typically NON-DETERMINISTIC for bfloat16 Usage: - python test_ar.py + pytest test_amd_nccl_allreduce_determinism.py """ import multiprocessing as mp import socket +import pytest import torch import torch.distributed as dist @@ -194,5 +195,14 @@ def main(): p.join() +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="Requires at least 2 CUDA GPUs", +) +def test_nccl_allreduce_determinism(): + """Test NCCL all-reduce determinism behavior with varying batch sizes.""" + main() + + if __name__ == "__main__": main()