CI: add an RDMA stage for p2p tests on sci H200 runners

.claude/rules/ci-test-registration.md

@@ -38,6 +38,7 @@ run the test; copying a neighbouring test's `suite=` is not a reason.
 | 4 | `stage-c-2-gpu-h200` | 2× H200 | 2-GPU tests |
 | 5 | `stage-c-4-gpu-h200` | 4× H200 | 4-GPU tests |
 | 6 | `stage-c-8-gpu-h200` or `stage-c-8-gpu-h100` | 8× H200 / 8× H100 | 8-GPU tests |
+| 6 | `stage-c-8-gpu-h200-rdma` | 8× H200 with RDMA between them | 8-GPU tests that need RDMA inside the host, such as p2p weight updates |
 | — | `stage-c-8-gpu-b200` | 8× B200 | preferably tests that cannot run on Hopper (`hardware=["blackwell"]`), any GPU count |
 
 The stage's GPU count equals the count the test requests (`ray start

.github/workflows/pr-test.yml

@@ -284,6 +284,29 @@ jobs:
         ${{ github.event_name == 'workflow_dispatch' && '--match-all-labels' || '' }}
     secrets: inherit
 
+  # Tests that need RDMA between the GPUs of one host (p2p weight updates). Its
+  # runners also carry h200 and 8gpu, so stage-c-8-gpu-h200 runs there too.
+  stage-c-8-gpu-h200-rdma:
+    needs: [resolve-ci-policy, resolve-ci-image, stage-a-cpu]
+    if: |
+      always() && !cancelled() &&
+      needs.resolve-ci-policy.result == 'success' &&
+      needs.resolve-ci-image.result == 'success' &&
+      !contains(fromJSON(needs.resolve-ci-policy.outputs.skipped_stages || '[]'), 'stage-c-8-gpu-h200-rdma') &&
+      (needs.stage-a-cpu.result == 'success' ||
+        needs.resolve-ci-policy.outputs.bypass_fastfail == 'true')
+    uses: ./.github/workflows/_run-ci.yml
+    with:
+      runs_on: '["h200", "8gpu", "rdma"]'
+      container_image: ${{ needs.resolve-ci-image.outputs.container_image }}
+      ref: ${{ inputs.ref || '' }}
+      execute_command: >-
+        python tests/ci/run_suite.py --hw cuda --suite stage-c-8-gpu-h200-rdma
+        --cadence ${{ needs.resolve-ci-policy.outputs.cadence }}
+        --labels ${{ needs.resolve-ci-policy.outputs.raw_labels }}
+        ${{ github.event_name == 'workflow_dispatch' && '--match-all-labels' || '' }}
+    secrets: inherit
+
   stage-c-4-gpu-h200:
     needs: [resolve-ci-policy, resolve-ci-image, stage-a-cpu]
     if: |

docs/developer/ci/00-stage.md

@@ -25,12 +25,15 @@ Stage names follow `stage-<tier>-<gpus>-<hw>` (or `stage-<tier>-<hw>` for CPU, e
 | `stage-c-4-gpu-h200` | 4× H200 | `["h200","4gpu"]` | 5 | both resolvers, `stage-a-cpu` |
 | `stage-c-8-gpu-h100` | 8× H100 | `["h100","8gpu"]` | 1 | both resolvers, `stage-a-cpu` |
 | `stage-c-8-gpu-h200` | 8× H200 | `["h200","8gpu"]` | 1 | both resolvers, `stage-a-cpu` |
+| `stage-c-8-gpu-h200-rdma` | 8× H200, RDMA between the host's GPUs | `["h200","8gpu","rdma"]` | 1 | both resolvers, `stage-a-cpu` |
 | `stage-c-8-gpu-b200` | 8× B200 | `["b200","8gpu"]` | 1 | both resolvers, `stage-a-cpu` |
 | `stage-c-4-gpu-mi350` | 4× MI350 | `["self-hosted","amd","mi350","4gpu"]` | 2 | both resolvers |
 | `nightly-stage-c-2-gpu-mi350` | 2× MI350 | external nightly | — | — |
 | `nightly-stage-c-4-gpu-mi350` | 4× MI350 | external nightly | — | — |
 | `nightly-stage-c-8-gpu-mi350` | 8× MI350 | external nightly | — | — |
 
+`stage-c-8-gpu-h200-rdma` holds the tests that move data between a host's GPUs over RDMA, as p2p weight updates do; only runners whose RDMA works between all of their GPUs carry the `rdma` label. Those runners also carry `h200` and `8gpu`, so `stage-c-8-gpu-h200` uses them as well, and tests without the need never wait for them.
+
 `stage-c-8-gpu-b200` is the only Blackwell stage: the Blackwell fleet is a single unpartitioned host, and an 8-GPU runner cannot share a node with 2/4-GPU runners without both claiming the same physical GPUs. Blackwell tests needing fewer than eight GPUs therefore register against this suite too and leave the surplus idle, which is correct because a test declares its own GPU budget (`ray start --num-gpus`, `torchrun --nproc-per-node`) rather than inferring one from the devices it can see.
 
 In `pr-test.yml`, `tier a` (CPU fast) gates PR-image preparation and the NVIDIA GPU fleet; its GPU stages (`b` / `c`) all depend on both resolvers and `stage-a-cpu`, and run concurrently with each other — the `b` / `c` letters classify role, they are not a sequential pipeline. The MI350 stage has no CPU-test gate.
@@ -78,7 +81,7 @@ That runner pool can include RoCE hosts, and GitHub attaches every job container
 
 **Arch dispatch.** `tests/ci/hardware.py::CUDA_STAGES` is the single source of truth for the CUDA taxonomy: each stage's GPU generation, GPU count, and runner labels. `CI_SUITES` and the `/rerun-test` runner map both derive from it, and a stage's `--suite` already names its generation, so no job passes an arch explicitly.
 
-A test normally runs at its home stage. A [dispatch label](/developer/ci/01-label) can send it to another generation instead, in which case the destination is the smallest stage on that generation with enough GPUs — a test declares its own GPU budget through `ray start --num-gpus` / `torchrun --nproc-per-node` rather than reading the devices it can see, so a wider stage simply leaves the surplus idle. Today every dispatched test lands on `stage-c-8-gpu-b200`, the only Blackwell stage; partitioning a second Blackwell host adds narrower stages and routing follows automatically.
+A test normally runs at its home stage. A [dispatch label](/developer/ci/01-label) can send it to another generation instead, in which case the destination is the smallest stage on that generation with enough GPUs and every capability of the home stage (the runner labels after the GPU model and count, such as `rdma`) — a test declares its own GPU budget through `ray start --num-gpus` / `torchrun --nproc-per-node` rather than reading the devices it can see, so a wider stage simply leaves the surplus idle, but a capability it was homed for cannot be dropped. A generation without such a stage does not run the test. Today every dispatched test lands on `stage-c-8-gpu-b200`, the only Blackwell stage; partitioning a second Blackwell host adds narrower stages and routing follows automatically.
 
 Without a dispatch label nothing leaves its home stage, so scheduled and called runs are unaffected. Whatever stage actually executes a run is also the stage its performance baseline is keyed on, keeping the generations' numbers apart.
 

tests/ci/file_run.py

@@ -24,6 +24,7 @@
     "stage-b-2-gpu-h200": ["h200", "2gpu"],
     "stage-c-8-gpu-h100": ["h100", "8gpu"],
     "stage-c-8-gpu-h200": ["h200", "8gpu"],
+    "stage-c-8-gpu-h200-rdma": ["h200", "8gpu", "rdma"],
     "stage-c-4-gpu-h200": ["h200", "4gpu"],
     "stage-c-2-gpu-h200": ["h200", "2gpu"],
     "stage-c-8-gpu-b200": ["b200", "8gpu"],

tests/ci/hardware.py

@@ -33,13 +33,21 @@
 class CudaStage:
     arch: str
     num_gpus: int
+    # the GPU model and count, then any capability the runners offer beyond the GPUs
     runs_on: tuple[str, ...]
 
+    @property
+    def capabilities(self) -> frozenset[str]:
+        """What a test homed here may rely on besides the GPUs, such as `rdma`."""
+        return frozenset(self.runs_on[2:])
+
 
 CUDA_STAGES: dict[str, CudaStage] = {
     "stage-b-2-gpu-h200": CudaStage("hopper", 2, ("h200", "2gpu")),
     "stage-c-8-gpu-h100": CudaStage("hopper", 8, ("h100", "8gpu")),
     "stage-c-8-gpu-h200": CudaStage("hopper", 8, ("h200", "8gpu")),
+    # RDMA between the GPUs of one host, which p2p weight updates write over
+    "stage-c-8-gpu-h200-rdma": CudaStage("hopper", 8, ("h200", "8gpu", "rdma")),
     "stage-c-4-gpu-h200": CudaStage("hopper", 4, ("h200", "4gpu")),
     "stage-c-2-gpu-h200": CudaStage("hopper", 2, ("h200", "2gpu")),
     "stage-c-8-gpu-b200": CudaStage("blackwell", 8, ("b200", "8gpu")),
@@ -59,10 +67,13 @@ def target_stage(home_suite: str, arch: str) -> str | None:
     """The stage on `arch` that runs a test homed at `home_suite`.
 
     On the test's own arch that is its home stage. On another arch it is the
-    smallest stage with enough GPUs: a test declares its own budget through
-    `ray start --num-gpus` / `torchrun --nproc-per-node` rather than reading the
-    devices it can see, so a larger stage just leaves the surplus idle. None
-    when that arch has no stage big enough.
+    smallest stage with enough GPUs and every capability of the home stage: a
+    test declares its own budget through `ray start --num-gpus` /
+    `torchrun --nproc-per-node` rather than reading the devices it can see, so a
+    larger stage just leaves the surplus idle, but it cannot do without a
+    capability it was homed for. Among those, the one with the fewest
+    capabilities, so no test takes a runner it does not need. None when that
+    arch has no such stage.
 
     Derived rather than tabulated. A hand-written home -> destination map would
     encode a fleet shape that does not exist yet; this rule gives today's
@@ -72,9 +83,15 @@ def target_stage(home_suite: str, arch: str) -> str | None:
     home = CUDA_STAGES[home_suite]
     if home.arch == arch:
         return home_suite
-    fits = [name for name, stage in CUDA_STAGES.items() if stage.arch == arch and stage.num_gpus >= home.num_gpus]
-    # Tie-break on the name so two stages of equal width choose deterministically.
-    return min(fits, key=lambda name: (CUDA_STAGES[name].num_gpus, name), default=None)
+    fits = [
+        name
+        for name, stage in CUDA_STAGES.items()
+        if stage.arch == arch and stage.num_gpus >= home.num_gpus and stage.capabilities >= home.capabilities
+    ]
+    # Tie-break on the name so two equal stages choose deterministically.
+    return min(
+        fits, key=lambda name: (CUDA_STAGES[name].num_gpus, len(CUDA_STAGES[name].capabilities), name), default=None
+    )
 
 
 def dispatch_targets(

tests/ci/stage_selection.py

@@ -19,6 +19,7 @@
         "stage-c-4-gpu-h200",
         "stage-c-8-gpu-h100",
         "stage-c-8-gpu-h200",
+        "stage-c-8-gpu-h200-rdma",
         "stage-c-8-gpu-b200",
         "stage-c-4-gpu-mi350",
     }

tests/ci/test/test_hardware.py

@@ -116,6 +116,25 @@ def test_target_stage_is_none_when_no_stage_is_wide_enough():
     assert target_stage("stage-c-8-gpu-h200", "nonexistent-arch") is None
 
 
+def test_a_capability_of_the_home_stage_is_kept_on_another_arch():
+    """A p2p test on a runner without RDMA between its GPUs fails at connect; no
+    Blackwell stage offers that yet, so the test stays on Hopper."""
+    assert target_stage("stage-c-8-gpu-h200-rdma", "blackwell") is None
+    assert dispatch_targets("stage-c-8-gpu-h200-rdma", ["hopper", "blackwell"], dispatch_arches=BOTH, absorb=True) == {
+        "stage-c-8-gpu-h200-rdma"
+    }
+
+
+def test_a_test_without_a_capability_never_takes_a_stage_that_offers_one():
+    # those runners are the scarce ones the capable tests wait for
+    for home, stage in CUDA_STAGES.items():
+        if stage.capabilities:
+            continue
+        for arch in KNOWN_ARCHES:
+            destination = target_stage(home, arch)
+            assert destination is None or not CUDA_STAGES[destination].capabilities, (home, arch)
+
+
 def test_blackwell_work_never_reaches_a_hopper_stage():
     """`target_stage` alone would route it; `dispatch_targets` never asks.
 

tests/ci/test/test_run_suite.py

@@ -121,6 +121,7 @@ def test_cuda_suites_exact(self):
             "stage-b-2-gpu-h200",
             "stage-c-8-gpu-h100",
             "stage-c-8-gpu-h200",
+            "stage-c-8-gpu-h200-rdma",
             "stage-c-4-gpu-h200",
             "stage-c-2-gpu-h200",
             "stage-c-8-gpu-b200",
@@ -327,7 +328,7 @@ def _reusable_workflow(name: str) -> str:
     def test_every_stage_consumes_resolved_policy(self):
         workflow = self._workflow()
         commands = workflow.split("execute_command:")[1:]
-        assert len(commands) == 8, "stage inventory changed; update this lock test"
+        assert len(commands) == 9, "stage inventory changed; update this lock test"
         for block in commands:
             cmd = block.split("secrets:")[0]
             assert "--cadence ${{ needs.resolve-ci-policy.outputs.cadence }}" in cmd
@@ -356,7 +357,7 @@ def test_non_default_base_pr_needs_run_ci_label(self):
     def test_cpu_and_gpu_stages_use_dedicated_reusable_workflows(self):
         workflow = self._workflow()
         assert workflow.count("uses: ./.github/workflows/_run-cpu-ci.yml") == 2
-        assert workflow.count("uses: ./.github/workflows/_run-ci.yml") == 6
+        assert workflow.count("uses: ./.github/workflows/_run-ci.yml") == 7
         assert workflow.count("uses: ./.github/workflows/_build-pr-ci-image.yml") == 1
         assert "cpu_runner" not in workflow
 
@@ -480,15 +481,15 @@ def test_dispatch_has_no_scope_input_but_runs_all_cuda_domains(self):
         assert "ci_scope" not in dispatch_inputs
         manual_scope = "${{ github.event_name == 'workflow_dispatch' && '--match-all-labels' || '' }}"
         cuda_stages = workflow.split("  stage-b-2-gpu-h200:", 1)[1]
-        assert cuda_stages.count(manual_scope) == 6
+        assert cuda_stages.count(manual_scope) == 7
 
     def test_gpu_gates_consume_shared_bypass_output(self):
         workflow = self._workflow()
         gpu_stages = workflow.split("  stage-b-2-gpu-h200:", 1)[1]
         bypass_gate = "needs.resolve-ci-policy.outputs.bypass_fastfail == 'true'"
-        assert gpu_stages.count(bypass_gate) == 6
-        assert gpu_stages.count("needs.resolve-ci-policy.result == 'success'") == 6
-        assert gpu_stages.count("needs.resolve-ci-image.result == 'success'") == 6
+        assert gpu_stages.count(bypass_gate) == 7
+        assert gpu_stages.count("needs.resolve-ci-policy.result == 'success'") == 7
+        assert gpu_stages.count("needs.resolve-ci-image.result == 'success'") == 7
         assert "needs.stage-a-cpu.result == 'failure'" not in gpu_stages
 
     def test_each_cuda_stage_consumes_the_fail_open_skip_list(self):

tests/e2e/megatron/test_qwen3_30B_A3B_p2p.py

@@ -11,7 +11,7 @@
 
 register_cuda_ci(
     est_time=800,
-    suite="stage-c-8-gpu-h100",
+    suite="stage-c-8-gpu-h200-rdma",
     labels=["megatron", "weight-update"],
     hardware=["hopper", "blackwell"],
 )