[AMD] register kv_canary + mock_model e2e tests to extra-a (1-gpu-small + 2-gpu-large) (#28850)

.github/workflows/pr-test-amd-extra.yml

@@ -12,11 +12,13 @@ name: PR Test Extra (AMD)
 #
 # Stage: extra-a. Each job mirrors the container bring-up of pr-test-amd.yml
 # and dispatches `run_suite.py --hw amd --suite extra-a-test-{config}-amd`:
-#   - 1-gpu-small: mock-model / kv_canary *unit* tests
+#   - 1-gpu-small: mock-model / kv_canary unit + single-GPU canary e2e tests
 #   - 1-gpu-large: single-GPU model e2e tests (quant fp8kv-triton,
-#                  sessions streaming-session, spec standalone triton)
-# The canary *e2e* tests still need the canary JIT kernel ported to ROCm
-# first, so they remain CUDA-only for now.
+#                  sessions streaming-session, spec standalone triton,
+#                  kv_canary self-bench)
+#   - 2-gpu-large: multi-GPU (TP/PP/PD) mock-model + kv_canary e2e tests
+# kv_canary e2e is registered to the same extra-a stage as its CUDA siblings
+# (it now exercises the ROCm canary kernels end-to-end).
 
 on:
   pull_request:
@@ -177,6 +179,36 @@ jobs:
         run: |
           bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite extra-a-test-1-gpu-large-amd --timeout-per-file 1800 ${{ inputs.continue_on_error == true && '--continue-on-error' || '' }}
 
+  # =============================================== extra-a (2-gpu-large) ===============================================
+  # Multi-GPU TP / PP / PD mock-model + kv_canary e2e tests. Mirrors CUDA's
+  # extra-a 2-gpu-large stage; runs on the 2-GPU AMD pool.
+  extra-a-test-2-gpu-large-amd:
+    name: ${{ format('extra-a-test-2-gpu-large-amd{0} (linux-{1}-2gpu-sglang)', inputs.rocm_version && format('-{0}', inputs.rocm_version) || '', inputs.runner_arch || 'mi325') }}
+    needs: [call-gate]
+    if: ${{ !cancelled() && needs.call-gate.result == 'success' }}
+    runs-on: ${{ format('linux-{0}-2gpu-sglang', inputs.runner_arch || 'mi325') }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref || github.sha }}
+
+      - name: Ensure VRAM is clear
+        run: bash scripts/ci/amd/ensure_vram_clear.sh rocm
+
+      - name: Start CI container
+        run: bash scripts/ci/amd/amd_ci_start_container.sh ${{ inputs.rocm_version && format('--rocm-version {0}', inputs.rocm_version) || '' }}
+        env:
+          GITHUB_WORKSPACE: ${{ github.workspace }}
+
+      - name: Install dependencies
+        run: bash scripts/ci/amd/amd_ci_install_dependency.sh
+
+      - name: Run test
+        timeout-minutes: 60
+        run: |
+          bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite extra-a-test-2-gpu-large-amd --timeout-per-file 1800 ${{ inputs.continue_on_error == true && '--continue-on-error' || '' }}
+
   # =============================================== aggregator ====================================================
   # Single fan-in job so branch protection / notifications depend on one job
   # rather than every matrix leg. Fails if any dependent failed or was
@@ -187,6 +219,7 @@ jobs:
         call-gate,
         extra-a-test-1-gpu-small-amd,
         extra-a-test-1-gpu-large-amd,
+        extra-a-test-2-gpu-large-amd,
       ]
     if: always()
     runs-on: ubuntu-latest

test/registered/kv_canary/test_self_e2e_baseline.py

@@ -3,11 +3,12 @@ from __future__ import annotations
 import unittest
 
 from sglang.srt.kv_canary.config import CanaryMode
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=236, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class _BaselineBase(CanaryE2EBase):

test/registered/kv_canary/test_self_e2e_bench_speed.py

@@ -17,6 +17,9 @@ from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER
 
+# CUDA-only: this self-bench asserts a 1.0% kv_canary overhead budget tuned on
+# the CUDA (H100) runner. On ROCm the measured overhead is ~1.26%, so the
+# benchmark is not portable as-is; keep it off AMD CI rather than register-and-skip.
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="1-gpu-large")
 
 

test/registered/kv_canary/test_self_e2e_pd_baseline.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.pd_fixture import CanaryPDFixture
 
 register_cuda_ci(est_time=180, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=106, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class TestPDBaselineMha(CanaryPDFixture):

test/registered/kv_canary/test_self_e2e_pd_perturb.py

@@ -4,10 +4,11 @@ import unittest
 from typing import ClassVar
 
 from sglang.srt.kv_canary.perturb.config import TargetGroupKind
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.pd_fixture import CanaryPDFixture
 
 register_cuda_ci(est_time=180, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=231, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class _PDPerturbBase(CanaryPDFixture):

test/registered/kv_canary/test_self_e2e_perturb_raise.py

@@ -4,10 +4,11 @@ import unittest
 
 from sglang.srt.kv_canary.config import CanaryMode
 from sglang.srt.kv_canary.perturb.config import TargetGroupKind
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=50, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class TestPerturbRaiseMha(CanaryE2EBase):

test/registered/kv_canary/test_self_e2e_perturb_real_kv_unused_cache.py

@@ -5,11 +5,12 @@ from typing import ClassVar
 
 from sglang.srt.kv_canary.config import CanaryMode
 from sglang.srt.kv_canary.perturb.config import TargetGroupKind
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=503, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class _PerturbRealKvUnusedCacheBase(CanaryE2EBase):

test/registered/kv_canary/test_self_e2e_perturb_real_kv_used.py

@@ -5,11 +5,12 @@ from typing import ClassVar
 
 from sglang.srt.kv_canary.config import CanaryMode
 from sglang.srt.kv_canary.perturb.config import TargetGroupKind
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=256, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class _PerturbRealKvUsedBase(CanaryE2EBase):

test/registered/kv_canary/test_self_e2e_perturb_req_to_token.py

@@ -3,11 +3,12 @@ from __future__ import annotations
 import unittest
 
 from sglang.srt.kv_canary.config import CanaryMode
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=175, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class _PerturbReqToTokenBase(CanaryE2EBase):

test/registered/kv_canary/test_self_e2e_pp_baseline.py

@@ -3,10 +3,11 @@ from __future__ import annotations
 import unittest
 
 from sglang.srt.kv_canary.config import CanaryMode
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.pp_fixture import CanaryPPFixture
 
 register_cuda_ci(est_time=220, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=243, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class TestPPBaselineSwa(CanaryPPFixture):

test/registered/kv_canary/test_self_e2e_pp_perturb.py

@@ -5,10 +5,11 @@ from typing import ClassVar
 
 from sglang.srt.kv_canary.config import CanaryMode
 from sglang.srt.kv_canary.perturb.config import TargetGroupKind
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.pp_fixture import CanaryPPFixture
 
 register_cuda_ci(est_time=220, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=298, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class TestPPPerturbSwaSwa(CanaryPPFixture):

test/registered/kv_canary/test_self_e2e_pr_25015.py

@@ -6,10 +6,11 @@ import unittest
 from typing import ClassVar
 
 from sglang.srt.kv_canary.config import CanaryMode
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=101, stage="extra-a", runner_config="1-gpu-small-amd")
 
 _SPEC_EAGLE_TOKEN_ORACLE_ENV = {
     "SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT": "0",

test/registered/kv_canary/test_self_e2e_pr_26329.py

@@ -8,10 +8,11 @@ import unittest
 from typing import ClassVar
 
 from sglang.srt.kv_canary.config import CanaryMode
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.e2e_base import CanaryE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=99, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 _CHUNKED_PREFILL_SIZE = 2048

test/registered/mock_model/test_e2e_pd.py

@@ -7,7 +7,8 @@ from typing import ClassVar, Dict, List
 
 import requests
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kv_canary.violation_log_utils import assert_no_violation_in_log
 from sglang.test.mock_model.utils import (
     MOCK_MODEL_PATH,
@@ -19,6 +20,7 @@ from sglang.test.server_fixtures.disaggregation_fixture import (
 )
 
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=165, stage="extra-a", runner_config="2-gpu-large-amd")
 
 # DO NOT pass --disable-cuda-graph in canary e2e tests.  The canary kernel
 # must run inside the cuda graph alongside the real attn kernel; disabling the
@@ -128,6 +130,13 @@ class TestPdTransferCanaryClean(_MockModelPDBase, unittest.TestCase):
         self.assert_no_canary_violation()
 
 
+@unittest.skipIf(
+    is_hip(),
+    "ROCm: PD full-real-data KV checksum intermittently trips a "
+    "verify_real_kv_hash canary violation on the decode-side transferred prefix "
+    "(see https://github.com/sgl-project/sglang/issues/28971). The baseline PD "
+    "canary test above stays enabled on AMD.",
+)
 class TestPdTransferChecksumFullRealData(_MockModelPDBase, unittest.TestCase):
     """--kv-canary-real-data=all + sweep every step, no perturb, no violation."""
 

test/registered/mock_model/test_e2e_pp.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.mock_model.utils import run_mock_model_bench_serving
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=67, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class TestE2EPipelineParallel(CustomTestCase):

test/registered/mock_model/test_e2e_spec_eagle.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.mock_model.utils import MOCK_MODEL_PATH, run_mock_model_bench_serving
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=77, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class TestE2ESpeculativeEagle(CustomTestCase):

test/registered/mock_model/test_e2e_tp.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.mock_model.utils import run_mock_model_bench_serving
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=167, stage="extra-a", runner_config="2-gpu-large-amd")
 
 
 class TestE2ETensorParallel(CustomTestCase):

test/registered/mock_model/test_self_e2e_perturb_next_token_swap.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import unittest
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.mock_model.perturb_e2e_base import MockModelPerturbE2EBase
 
 register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+register_amd_ci(est_time=131, stage="extra-a", runner_config="1-gpu-small-amd")
 
 
 class TestPerturbNextTokenSwap(MockModelPerturbE2EBase):

test/run_suite.py

@@ -46,15 +46,17 @@ PER_COMMIT_SUITES = {
         # extra-a: label-gated PR opt-in suites in pr-test-amd-extra.yml
         # (mirror of CUDA extra-a; tests stay tagged per-commit but only
         # dispatch when the PR carries the `run-ci-extra` label). 1-gpu-small
-        # carries the mock-model / kv_canary *unit* tests; 1-gpu-large carries
-        # the subset of model e2e tests validated to pass on mi325 (quant
-        # fp8kv-triton, sessions streaming-session EAGLE3, spec standalone
-        # triton-backend variant). The rest of CUDA
-        # extra-a tests fail on ROCm (missing flash_attn.cute/flash_ops
+        # carries the mock-model / kv_canary unit + single-GPU canary e2e
+        # tests; 1-gpu-large carries the subset of model e2e tests validated
+        # to pass on mi325 (quant fp8kv-triton, sessions streaming-session
+        # EAGLE3, spec standalone triton-backend variant); 2-gpu-large carries
+        # the multi-GPU (TP/PP/PD) mock-model + kv_canary e2e tests. The rest
+        # of CUDA extra-a tests fail on ROCm (missing flash_attn.cute/flash_ops
         # kernels, OOM, or accuracy regressions — e.g. gemma4-mtp-31b dips
         # below the gsm8k floor on the topk=3 leg) and stay CUDA-only for now.
         "extra-a-test-1-gpu-small-amd",
         "extra-a-test-1-gpu-large-amd",
+        "extra-a-test-2-gpu-large-amd",
     ],
     HWBackend.MUSA: [],
     HWBackend.CUDA: [