From 6736155498265c66e2f962b338221d4c9314f89c Mon Sep 17 00:00:00 2001
From: Yin Lin <yinlin@google.com>
Date: Thu, 18 Jun 2026 05:22:59 +0000
Subject: [PATCH 1/4] Add CPU-only GitHub Actions CI (presubmit + nightly +
 lint).

Mirrors torch_tpu's GitHub Actions setup, scoped to a CPU runner
(linux-x86-n2-32) since google/tpu-raiden is not granted TPU runner
scale-sets. Builds the JAX extension, runs OSS-safe CPU unit tests, builds
the tpu_raiden_jax wheel nightly, and lints license headers + C++ format.
Kokoro CI is kept alongside. See .github/README.md.
---
 .clang-format                      |  25 +++++++
 .github/README.md                  |  65 +++++++++++++++++
 .github/actionlint.yaml            |  20 +++++
 .github/workflows/clang_format.yml |  72 ++++++++++++++++++
 .github/workflows/lint.yml         |  72 ++++++++++++++++++
 .github/workflows/nightly.yml      | 113 +++++++++++++++++++++++++++++
 .github/workflows/presubmit.yml    |  85 ++++++++++++++++++++++
 ci/tools/bazel_test.sh             | 105 +++++++++++++++++++++++++++
 ci/tools/install_clang18.sh        |  41 +++++++++++
 9 files changed, 598 insertions(+)
 create mode 100644 .clang-format
 create mode 100644 .github/README.md
 create mode 100644 .github/actionlint.yaml
 create mode 100644 .github/workflows/clang_format.yml
 create mode 100644 .github/workflows/lint.yml
 create mode 100644 .github/workflows/nightly.yml
 create mode 100644 .github/workflows/presubmit.yml
 create mode 100755 ci/tools/bazel_test.sh
 create mode 100755 ci/tools/install_clang18.sh

diff --git a/.clang-format b/.clang-format
new file mode 100644
index 000000000..3cc4fd76b
--- /dev/null
+++ b/.clang-format
@@ -0,0 +1,25 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This file configures the clang-format tool for C++ code. It enforces the
+# Google C++ Style Guide rules across all modifications:
+# https://google.github.io/styleguide/cppguide.html.
+#
+# Details on configuration formatting and style options can be found at:
+# https://clang.llvm.org/docs/ClangFormatStyleOptions.html
+
+---
+Language:        Cpp
+BasedOnStyle:  Google
+...
diff --git a/.github/README.md b/.github/README.md
new file mode 100644
index 000000000..e228e50d3
--- /dev/null
+++ b/.github/README.md
@@ -0,0 +1,65 @@
+# tpu-raiden CI/CD
+
+GitHub Actions CI for tpu-raiden, modeled on
+[`torch_tpu`](https://github.com/google-pytorch/torch_tpu)'s setup but scoped
+**CPU-only** for now. The existing Kokoro jobs under [`kokoro/`](../kokoro) are
+kept and run alongside these.
+
+## Why CPU-only
+
+torch_tpu runs TPU lanes on self-hosted runner scale-sets
+(`linux-x86-ct5lp/ct6e/tpu7x`), but those are granted **per org/repo** and
+torch_tpu lives in the `google-pytorch` org — that access does **not**
+automatically extend to `google/tpu-raiden`. Until this repo is granted TPU
+runner scale-sets, all GitHub Actions CI here stays on a CPU runner
+(`linux-x86-n2-32`). TPU/device tests continue to run via Kokoro or on a
+developer TPU VM. To add TPU lanes later, grant the repo the scale-set, add the
+label to [`actionlint.yaml`](actionlint.yaml), and reintroduce a privileged TPU
+job (the `git` history has a v5e/v6e/v7 matrix to crib from).
+
+> The `linux-x86-n2-32` label is what we expect to have; if the available CPU
+> scale-set differs, change `runs-on:` in `presubmit.yml` / `nightly.yml` and
+> the label in `actionlint.yaml`.
+
+## Workflows
+
+| Workflow | Trigger | What it does |
+| --- | --- | --- |
+| [`lint.yml`](workflows/lint.yml) | PR | Apache license headers (`addlicense`) + commit-message check. GitHub-hosted. |
+| [`clang_format.yml`](workflows/clang_format.yml) | PR | Google C++ style on changed `.cc/.h` via pinned `clang-format==18`. GitHub-hosted. |
+| [`presubmit.yml`](workflows/presubmit.yml) | PR | Build the JAX extension + run OSS-safe CPU unit tests + import smoke on `linux-x86-n2-32`. |
+| [`nightly.yml`](workflows/nightly.yml) | push to `main` / daily cron 09:00 UTC / dispatch | CPU build + unit tests, build the `tpu_raiden_jax` wheel (twine-checked, uploaded as an artifact), and file a tracking issue on failure. |
+
+Shared logic lives in [`ci/tools/`](../ci/tools): `install_clang18.sh` (the
+clang-18 the build needs) and `bazel_test.sh` (bazel binary + dummy torch_tpu
+override + flags, matching `build.sh`).
+
+## Scope notes
+
+- **JAX-only build.** The JAX path generates a dummy `torch_tpu` Bazel module
+  (same as `build.sh` / `kokoro/.../presubmit.sh`), so no secrets are needed. The
+  **Torch** extension/wheel additionally needs a `torch_tpu` checkout + a local
+  `torch`, and its meaningful tests need a TPU — out of scope here for now.
+- **Unit-test set.** `CPU_TEST_TARGETS` is a curated CPU-safe, OSS-loadable set.
+  `//rpc/...` and `//kv_cache/global_registry/...` (and the `kv_cache_store`
+  tests that depend on them) load Google-internal gRPC rules and are excluded by
+  design — see [`ci/wheel/BUILD.bazel`](../ci/wheel/BUILD.bazel). Expand the list
+  as more packages become CPU-runnable.
+
+## One-time setup to turn this on
+
+1. **CPU runner** — grant `google/tpu-raiden` access to the `linux-x86-n2-32`
+   self-hosted scale-set (or adjust the label to one it already has).
+2. **Bazel remote cache** — bucket `gs://tpu-raiden-bazel-cache` (already used by
+   Kokoro); the runner service account needs object read (and write for the
+   nightly's read-write cache).
+3. **Labels** — create `ci:nightly-failed` for the failure-issue action.
+   `GITHUB_TOKEN` is provided automatically.
+
+## Landing & running
+
+`google/tpu-raiden` is a one-way Copybara mirror of an internal google3 repo
+(commits carry `PiperOrigin-RevId:`), so these files land **internally** and sync
+out — there is no external push/PR path. Once on `main`, the CPU workflows fire on
+PRs/pushes automatically, and `nightly.yml` can be kicked manually with
+`gh workflow run "CI - Nightly (CPU)"` by anyone with write access.
diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
new file mode 100644
index 000000000..4229f1ecc
--- /dev/null
+++ b/.github/actionlint.yaml
@@ -0,0 +1,20 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Declares the self-hosted runner labels used by the workflows so actionlint
+# does not flag them as unknown. CPU-only for now; add TPU scale-set labels here
+# (e.g. linux-x86-tpu7x-224-4tpu) once google/tpu-raiden is granted them.
+self-hosted-runner:
+  labels:
+    - linux-x86-n2-32   # CPU
diff --git a/.github/workflows/clang_format.yml b/.github/workflows/clang_format.yml
new file mode 100644
index 000000000..1f05375b0
--- /dev/null
+++ b/.github/workflows/clang_format.yml
@@ -0,0 +1,72 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Enforces the Google C++ Style Guide on Pull Requests via clang-format, so
+# unformatted C++ cannot be merged into main. Only files changed in the PR are
+# checked, and the clang-format version is pinned to 18 to match the clang-18
+# toolchain raiden builds with (see ci/build_wheel.sh).
+
+name: Clang Format
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+# Cancel old jobs when a contributor pushes new commits to their PR.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+jobs:
+  run_clang_format:
+    name: Clang format check
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out PR HEAD
+        uses: actions/checkout@v5
+
+      - name: Fetch PR base commit
+        run: git fetch --no-tags --depth=1 origin ${{ github.event.pull_request.base.sha }}
+
+      # uv is a fast, single-binary replacement for pip/pipx/virtualenv. Modern
+      # runners enforce PEP 668 ("Externally Managed Environments"), which blocks
+      # global `pip install`; uvx runs the pinned clang-format wheel in an
+      # isolated sandbox without touching the host interpreter.
+      - name: Setup uv
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b  # v8.1.0
+        with:
+          enable-cache: true
+
+      - name: Run clang-format on changed C/C++ files
+        env:
+          BASE_SHA: ${{ github.event.pull_request.base.sha }}
+        run: |
+          # Collect C/C++ sources changed in this PR (added/copied/modified/renamed).
+          mapfile -t FILES < <(
+            git diff --name-only --diff-filter=ACMR "${BASE_SHA}" HEAD -- \
+              '*.c' '*.cc' '*.cpp' '*.cxx' '*.h' '*.hpp' '*.hh'
+          )
+          if [[ ${#FILES[@]} -eq 0 ]]; then
+            echo "No C/C++ files changed; nothing to check."
+            exit 0
+          fi
+          printf 'Checking:\n'; printf '  %s\n' "${FILES[@]}"
+          # clang-format 18.x matches the clang-18 build toolchain. --dry-run
+          # --Werror makes any required reformatting a hard failure.
+          uvx --from 'clang-format==18.1.8' clang-format --dry-run --Werror "${FILES[@]}"
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 000000000..40201a8e3
--- /dev/null
+++ b/.github/workflows/lint.yml
@@ -0,0 +1,72 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: "Lint checks"
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+  pull-requests: read
+  issues: read
+
+# Cancel old jobs when they are replaced by new runs.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  # Don't cancel in-progress jobs for the main branch.
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+jobs:
+  run_lint:
+    name: "Lint check"
+    runs-on: ubuntu-latest
+    steps:
+      - name: "Check out code"
+        uses: "actions/checkout@v5"
+      - name: "Set up Go"
+        uses: "actions/setup-go@v6"
+        with:
+          go-version: "1.24"
+          cache: false
+      - name: "Validate License Headers"
+        shell: "bash"
+        run: |
+          go run github.com/google/addlicense@v1.1.1 -check -s=only .
+      - name: "Validate Commit Messages"
+        if: (success() || failure()) && github.event_name == 'pull_request'
+        shell: "bash"
+        run: |
+          # Fix dubious ownership issue in Actions.
+          git config --global --add safe.directory '*'
+
+          # Get the exact number of commits in the PR to minimize git fetch depth.
+          PR_COMMITS="${{ github.event.pull_request.commits }}"
+          PR_COMMITS="${PR_COMMITS:-0}"
+
+          # Fetch the PR branch head and the base commit to establish the range.
+          git fetch --no-tags --depth=$((PR_COMMITS + 1)) origin "${{ github.event.pull_request.head.sha }}"
+          git fetch --no-tags --depth=1 origin "${{ github.event.pull_request.base.sha }}"
+
+          # Get commit messages between base and head of the pull request.
+          commit_msgs=$(git log "${{ github.event.pull_request.base.sha }}..${{ github.event.pull_request.head.sha }}" --format=%B)
+
+          if echo "$commit_msgs" | grep -qF "[NO DESCRIPTION SET]"; then
+            echo "::error::Found '[NO DESCRIPTION SET]' in one or more commit messages."
+            exit 1
+          fi
+          echo "Commit messages look good."
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
new file mode 100644
index 000000000..f0d0a33c3
--- /dev/null
+++ b/.github/workflows/nightly.yml
@@ -0,0 +1,113 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Nightly CPU build + unit tests. CPU-only (no TPU runners): builds the JAX
+# extension, runs the OSS-safe non-device unit tests, and builds the
+# tpu_raiden_jax wheel (twine-checked) as a packaging smoke test. Also runs on
+# push to main and on demand. Files a tracking issue on failure. Device tests and
+# the torch wheel are intentionally out of scope until TPU runners are available;
+# see .github/README.md.
+
+name: CI - Nightly (CPU)
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+  schedule:
+    - cron: "0 9 * * *"  # 1 AM PST daily. Scheduled workflows run in UTC.
+
+permissions:
+  contents: read
+  issues: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+env:
+  CPU_TEST_TARGETS: >-
+    //kv_cache:logical_block_manager_test
+    //kv_cache:lru_cache_test
+    //tpu_raiden/frameworks/jax:resharding_planner_test
+
+jobs:
+  build_and_test:
+    name: "Nightly build + unit tests (CPU)"
+    runs-on: linux-x86-n2-32
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+    steps:
+      - uses: actions/checkout@v5
+      - name: Fix Git workspace ownership
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      - name: Install clang-18
+        run: ./ci/tools/install_clang18.sh
+
+      - name: Build JAX extension
+        run: ./build.sh jax
+
+      - name: Run CPU unit tests
+        env:
+          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
+          RAIDEN_REMOTE_UPLOAD: "true"
+        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
+
+      - name: Build tpu_raiden_jax wheel (packaging smoke)
+        run: |
+          export WHEEL_VERSION_EXTRAS=".dev$(date +%Y%m%d%H%M%S)"
+          echo "WHEEL_VERSION_EXTRAS=$WHEEL_VERSION_EXTRAS"
+          ./build.sh jax //ci/wheel:raiden_jax_wheel \
+            --repo_env=WHEEL_VERSION_EXTRAS="$WHEEL_VERSION_EXTRAS"
+          mkdir -p dist
+          cp bazel-bin/ci/wheel/tpu_raiden_jax-*.whl dist/
+          ls -lh dist
+
+      - name: Check wheel metadata
+        run: uv run --isolated --with twine twine check dist/*.whl
+
+      - name: Upload wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: tpu_raiden_jax_wheel
+          path: dist/*.whl
+          retention-days: 14
+
+      - name: Upload test logs on failure
+        uses: actions/upload-artifact@v4
+        if: failure()
+        with:
+          name: bazel-testlogs-nightly
+          path: bazel-testlogs/
+          retention-days: 7
+
+  notify_failure:
+    name: Notify failed runs
+    needs: [build_and_test]
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Get current date
+        id: date
+        if: contains(needs.build_and_test.result, 'failure') && github.event_name != 'pull_request'
+        run: echo "date=$(date +'%Y-%m-%d')" >> $GITHUB_OUTPUT
+      - name: File a tracking issue on failure
+        if: contains(needs.build_and_test.result, 'failure') && github.event_name != 'pull_request'
+        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          label-name: "ci:nightly-failed"
+          title-template: "Failed Nightly build (${{ steps.date.outputs.date }})"
diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml
new file mode 100644
index 000000000..398b7bc4b
--- /dev/null
+++ b/.github/workflows/presubmit.yml
@@ -0,0 +1,85 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Per-PR gate. CPU-only: builds the JAX extension and runs the OSS-safe,
+# non-device unit tests + an import smoke check on a CPU runner. There are no TPU
+# lanes here on purpose -- this repo is not (yet) granted the self-hosted TPU
+# runner scale-sets, so all CI is kept on CPU. Device tests stay in Kokoro / on
+# a developer TPU VM until TPU runners are available. See .github/README.md.
+
+name: CI - Presubmit (CPU)
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+  pull-requests: read
+  issues: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+# Curated OSS-safe, CPU-only unit tests (no TPU device, no Google-internal grpc).
+# //rpc/... and //kv_cache/global_registry/... are excluded by design -- see
+# ci/wheel/BUILD.bazel. Expand as more packages become CPU-runnable.
+env:
+  CPU_TEST_TARGETS: >-
+    //kv_cache:logical_block_manager_test
+    //kv_cache:lru_cache_test
+    //tpu_raiden/frameworks/jax:resharding_planner_test
+
+jobs:
+  build_and_test:
+    name: "Presubmit on CPU"
+    # Self-hosted CPU runner. NOTE: confirm this scale-set is granted to
+    # google/tpu-raiden; swap the label if your pool differs.
+    runs-on: linux-x86-n2-32
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+    steps:
+      - uses: actions/checkout@v5
+
+      # Git's CVE-2022-24765 patch blocks root from operating on runner-owned
+      # dirs; trust the ephemeral workspace.
+      - name: Fix Git workspace ownership
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      - name: Install clang-18
+        run: ./ci/tools/install_clang18.sh
+
+      - name: Build JAX extension
+        run: ./build.sh jax
+
+      - name: Run CPU unit tests
+        env:
+          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
+          RAIDEN_REMOTE_UPLOAD: "false"
+        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
+
+      - name: Import smoke check (JAX mocked)
+        run: |
+          export PYTHONPATH="${GITHUB_WORKSPACE}:${GITHUB_WORKSPACE}/bazel-bin:${PYTHONPATH:-}"
+          python3 - <<'PY'
+          import sys
+          from unittest.mock import MagicMock
+          for m in ("jax", "jax.core", "jax.extend", "jax.extend.ffi"):
+              sys.modules[m] = MagicMock()
+          import tpu_raiden  # noqa: F401
+          print("tpu_raiden imported successfully.")
+          PY
diff --git a/ci/tools/bazel_test.sh b/ci/tools/bazel_test.sh
new file mode 100755
index 000000000..3840a8fa0
--- /dev/null
+++ b/ci/tools/bazel_test.sh
@@ -0,0 +1,105 @@
+#!/bin/bash
+
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Runs `bazel test` over a set of raiden targets, reusing the same bazel binary,
+# torch_tpu module override, and flags that build.sh uses. Used by the GitHub
+# Actions CPU workflows so test selection stays in one place rather than
+# duplicated across YAML.
+#
+# Usage:
+#   ci/tools/bazel_test.sh <jax|torch|both> [extra_bazel_flags...] -- <target...>
+#
+# Env:
+#   TORCH_TPU_MODULE_PATH   Path to a torch_tpu checkout (default ../torch_tpu).
+#                           Only needed for `torch`/`both`; `jax` uses a dummy.
+#   RAIDEN_REMOTE_CACHE     If set, used as --remote_cache (e.g.
+#                           https://storage.googleapis.com/tpu-raiden-bazel-cache).
+#   RAIDEN_REMOTE_UPLOAD    "true"/"false" -> --remote_upload_local_results
+#                           (default false; presubmit should stay read-only).
+
+set -exu -o pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "${REPO_ROOT}"
+
+MODE="${1:?usage: bazel_test.sh <jax|torch|both> [flags...] -- <target...>}"
+shift
+
+EXTRA_FLAGS=()
+while [[ "$#" -gt 0 && "$1" != "--" ]]; do
+  EXTRA_FLAGS+=("$1"); shift
+done
+[[ "${1:-}" == "--" ]] && shift
+TARGETS=("$@")
+if [[ ${#TARGETS[@]} -eq 0 ]]; then
+  echo "ERROR: no test targets provided (expected '... -- //pkg:target ...')." >&2
+  exit 1
+fi
+
+# Bootstrap the same standalone Bazel build.sh uses (honors .bazelversion).
+BAZEL_VERSION="$(tr -d '\r\n ' < "${REPO_ROOT}/.bazelversion")"
+BAZEL_BIN="/tmp/bazel-bootstrap-${BAZEL_VERSION}"
+if [[ ! -x "${BAZEL_BIN}" ]]; then
+  curl -Lo "${BAZEL_BIN}" \
+    "https://storage.googleapis.com/bazel/${BAZEL_VERSION}/release/bazel-${BAZEL_VERSION}-linux-x86_64"
+  chmod +x "${BAZEL_BIN}"
+fi
+"${BAZEL_BIN}" --version
+
+# torch_tpu module override: a real checkout for torch builds, otherwise a
+# generated dummy module so JAX-only graph resolution succeeds (mirrors
+# build.sh and kokoro/gcp_ubuntu/presubmit.sh).
+TORCH_TPU_MODULE_PATH="${TORCH_TPU_MODULE_PATH:-${REPO_ROOT}/../torch_tpu}"
+DEFINE_FLAGS=()
+if [[ "${MODE}" == "jax" ]]; then
+  DUMMY_MODULE="$(mktemp -d)/dummy_torch_tpu_module"
+  mkdir -p "${DUMMY_MODULE}"
+  echo 'module(name = "torch_tpu", version = "0.1.1")' > "${DUMMY_MODULE}/MODULE.bazel"
+  MODULE_OVERRIDE="--override_module=torch_tpu=${DUMMY_MODULE}"
+  DEFINE_FLAGS+=(--define with_torch=false)
+else
+  if [[ ! -f "${TORCH_TPU_MODULE_PATH}/MODULE.bazel" ]]; then
+    echo "ERROR: ${MODE} mode needs a torch_tpu checkout at ${TORCH_TPU_MODULE_PATH}." >&2
+    exit 1
+  fi
+  TORCH_TPU_MODULE_PATH="$(cd "${TORCH_TPU_MODULE_PATH}" && pwd)"
+  MODULE_OVERRIDE="--override_module=torch_tpu=${TORCH_TPU_MODULE_PATH}"
+  TORCH_SOURCE="$(python3 -c 'import importlib.util,pathlib;s=importlib.util.find_spec("torch");print(pathlib.Path(next(iter(s.submodule_search_locations))).resolve().parent)')"
+  export TORCH_SOURCE
+  DEFINE_FLAGS+=(--define=TORCH_SOURCE=local --repo_env=TORCH_SOURCE="${TORCH_SOURCE}")
+  [[ "${MODE}" == "torch" ]] && DEFINE_FLAGS+=(--define with_jax=false)
+fi
+
+REMOTE_FLAGS=()
+if [[ -n "${RAIDEN_REMOTE_CACHE:-}" ]]; then
+  REMOTE_FLAGS+=(
+    "--remote_cache=${RAIDEN_REMOTE_CACHE}"
+    "--google_default_credentials"
+    "--remote_upload_local_results=${RAIDEN_REMOTE_UPLOAD:-false}"
+    "--remote_max_connections=25"
+    "--remote_timeout=300s"
+    "--remote_retries=3"
+  )
+fi
+
+"${BAZEL_BIN}" test -c opt --check_visibility=false --verbose_failures \
+  --experimental_repo_remote_exec --incompatible_disallow_empty_glob=false \
+  --repo_env=HERMETIC_PYTHON_VERSION="${HERMETIC_PYTHON_VERSION:-3.12}" \
+  "${MODULE_OVERRIDE}" \
+  "${DEFINE_FLAGS[@]}" \
+  "${REMOTE_FLAGS[@]}" \
+  "${EXTRA_FLAGS[@]}" \
+  "${TARGETS[@]}"
diff --git a/ci/tools/install_clang18.sh b/ci/tools/install_clang18.sh
new file mode 100755
index 000000000..2386c44e1
--- /dev/null
+++ b/ci/tools/install_clang18.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Installs clang-18 into the ml-build container, which raiden's Bazel build
+# requires for XLA's .ll codegen targets (CWG2518 needs clang-18). Factored out
+# of ci/build_wheel.sh so the GitHub Actions jobs and the wheel build share one
+# source of truth. Idempotent: a no-op if clang-18 is already on PATH.
+
+set -exu -o pipefail
+
+if clang --version 2>/dev/null | grep -q 'version 18'; then
+  echo "clang-18 already installed."
+  exit 0
+fi
+
+export DEBIAN_FRONTEND=noninteractive
+apt-get update -qq
+apt-get install -y -qq wget gnupg ca-certificates >/dev/null
+# Add the LLVM jammy-18 apt repo manually (the container's add-apt-repository is
+# broken: python apt_pkg is missing for python3.12).
+wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /usr/share/keyrings/llvm.gpg
+echo "deb [signed-by=/usr/share/keyrings/llvm.gpg] http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main" \
+  > /etc/apt/sources.list.d/llvm18.list
+apt-get update -qq
+apt-get install -y -qq clang-18 >/dev/null
+ln -sf /usr/bin/clang-18 /usr/bin/clang
+ln -sf /usr/bin/clang++-18 /usr/bin/clang++
+clang --version | head -1

From 63e8decdd82575a2efdbf6b47cb981786015df47 Mon Sep 17 00:00:00 2001
From: Yin Lin <yinlin@google.com>
Date: Thu, 18 Jun 2026 05:25:24 +0000
Subject: [PATCH 2/4] Fix root BUILD license header; add tpu7x (v7) runner
 smoke probe.

addlicense flagged the root BUILD file (no header); add the standard Apache
header. Add a pull-request smoke workflow on linux-x86-tpu7x-224-4tpu to test
whether the v7 TPU runner scale-set is available to this repo.
---
 .github/actionlint.yaml           |  3 +-
 .github/workflows/tpu7x_smoke.yml | 61 +++++++++++++++++++++++++++++++
 BUILD                             | 14 +++++++
 3 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/tpu7x_smoke.yml

diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
index 4229f1ecc..28e314cd8 100644
--- a/.github/actionlint.yaml
+++ b/.github/actionlint.yaml
@@ -17,4 +17,5 @@
 # (e.g. linux-x86-tpu7x-224-4tpu) once google/tpu-raiden is granted them.
 self-hosted-runner:
   labels:
-    - linux-x86-n2-32   # CPU
+    - linux-x86-n2-32              # CPU
+    - linux-x86-tpu7x-224-4tpu     # TPU v7 (probe, see tpu7x_smoke.yml)
diff --git a/.github/workflows/tpu7x_smoke.yml b/.github/workflows/tpu7x_smoke.yml
new file mode 100644
index 000000000..b27879825
--- /dev/null
+++ b/.github/workflows/tpu7x_smoke.yml
@@ -0,0 +1,61 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Probe workflow: validates whether the linux-x86-tpu7x-224-4tpu (TPU v7) runner
+# scale-set is available to this repo, and whether the ml-build container can see
+# the v7 chips and run a device test. If v7 is confirmed, TPU lanes can be folded
+# back into presubmit/nightly. Delete this file once that decision is made.
+
+name: CI - tpu7x (v7) smoke probe
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  tpu7x_smoke:
+    name: "Smoke on tpu7x (v7)"
+    runs-on: linux-x86-tpu7x-224-4tpu
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+      options: --privileged
+    steps:
+      - uses: actions/checkout@v5
+      - name: Fix Git workspace ownership
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+      - name: Show runner / container info
+        run: |
+          echo "Runner: ${RUNNER_NAME:-?} on ${RUNNER_OS:-?}"
+          uname -a; nproc
+      - name: Install clang-18
+        run: ./ci/tools/install_clang18.sh
+      - name: Build JAX extension
+        run: ./build.sh jax
+      - name: Verify TPU v7 visible from container
+        run: |
+          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
+      - name: Run one TPU device test
+        env:
+          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
+          RAIDEN_REMOTE_UPLOAD: "false"
+        run: ./ci/tools/bazel_test.sh jax -- //tpu_raiden/frameworks/jax:raw_transfer_test_gl
diff --git a/BUILD b/BUILD
index a2353746e..5d1de6f29 100644
--- a/BUILD
+++ b/BUILD
@@ -1,3 +1,17 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 package(default_visibility = ["//visibility:public"])
 
 exports_files([

From 211ba7b7a7e4f4dfb5d41f8dec1402b4bb025fc7 Mon Sep 17 00:00:00 2001
From: Yin Lin <yinlin@google.com>
Date: Thu, 18 Jun 2026 05:31:15 +0000
Subject: [PATCH 3/4] Switch CI runner to v5e (linux-x86-ct5lp-224-8tpu).

n2-32 and tpu7x jobs stayed queued/unassigned (not granted to this repo), so
run presubmit + nightly on the available v5e TPU host: build, verify TPU
visible, run unit tests + a device test. Drop the tpu7x probe.
---
 .github/README.md                 | 56 +++++++++++++---------------
 .github/actionlint.yaml           |  7 ++--
 .github/workflows/nightly.yml     | 26 ++++++++-----
 .github/workflows/presubmit.yml   | 39 +++++++++++++-------
 .github/workflows/tpu7x_smoke.yml | 61 -------------------------------
 5 files changed, 70 insertions(+), 119 deletions(-)
 delete mode 100644 .github/workflows/tpu7x_smoke.yml

diff --git a/.github/README.md b/.github/README.md
index e228e50d3..2fb4e9e3f 100644
--- a/.github/README.md
+++ b/.github/README.md
@@ -1,25 +1,18 @@
 # tpu-raiden CI/CD
 
 GitHub Actions CI for tpu-raiden, modeled on
-[`torch_tpu`](https://github.com/google-pytorch/torch_tpu)'s setup but scoped
-**CPU-only** for now. The existing Kokoro jobs under [`kokoro/`](../kokoro) are
-kept and run alongside these.
+[`torch_tpu`](https://github.com/google-pytorch/torch_tpu)'s setup. The existing
+Kokoro jobs under [`kokoro/`](../kokoro) are kept and run alongside these.
 
-## Why CPU-only
+## Runner
 
-torch_tpu runs TPU lanes on self-hosted runner scale-sets
-(`linux-x86-ct5lp/ct6e/tpu7x`), but those are granted **per org/repo** and
-torch_tpu lives in the `google-pytorch` org — that access does **not**
-automatically extend to `google/tpu-raiden`. Until this repo is granted TPU
-runner scale-sets, all GitHub Actions CI here stays on a CPU runner
-(`linux-x86-n2-32`). TPU/device tests continue to run via Kokoro or on a
-developer TPU VM. To add TPU lanes later, grant the repo the scale-set, add the
-label to [`actionlint.yaml`](actionlint.yaml), and reintroduce a privileged TPU
-job (the `git` history has a v5e/v6e/v7 matrix to crib from).
-
-> The `linux-x86-n2-32` label is what we expect to have; if the available CPU
-> scale-set differs, change `runs-on:` in `presubmit.yml` / `nightly.yml` and
-> the label in `actionlint.yaml`.
+All build/test jobs run on the **v5e TPU runner `linux-x86-ct5lp-224-8tpu`**.
+This is the scale-set currently granted to `google/tpu-raiden`: jobs targeting
+`linux-x86-n2-32` (CPU) and `linux-x86-tpu7x-224-4tpu` (v7) stayed queued and
+unassigned, so they are not available here yet. v5e is a full host, so it runs
+the JAX build, the CPU unit tests, and TPU device tests. To add other runners
+later, grant the repo the scale-set, add the label to
+[`actionlint.yaml`](actionlint.yaml), and point a job's `runs-on:` at it.
 
 ## Workflows
 
@@ -27,8 +20,8 @@ job (the `git` history has a v5e/v6e/v7 matrix to crib from).
 | --- | --- | --- |
 | [`lint.yml`](workflows/lint.yml) | PR | Apache license headers (`addlicense`) + commit-message check. GitHub-hosted. |
 | [`clang_format.yml`](workflows/clang_format.yml) | PR | Google C++ style on changed `.cc/.h` via pinned `clang-format==18`. GitHub-hosted. |
-| [`presubmit.yml`](workflows/presubmit.yml) | PR | Build the JAX extension + run OSS-safe CPU unit tests + import smoke on `linux-x86-n2-32`. |
-| [`nightly.yml`](workflows/nightly.yml) | push to `main` / daily cron 09:00 UTC / dispatch | CPU build + unit tests, build the `tpu_raiden_jax` wheel (twine-checked, uploaded as an artifact), and file a tracking issue on failure. |
+| [`presubmit.yml`](workflows/presubmit.yml) | PR | Build JAX extension, verify TPU visible, run OSS-safe unit tests + one device test + import smoke, on v5e. |
+| [`nightly.yml`](workflows/nightly.yml) | push to `main` / daily cron 09:00 UTC / dispatch | Build + unit tests + device test, build the `tpu_raiden_jax` wheel (twine-checked, uploaded as an artifact), file a tracking issue on failure. |
 
 Shared logic lives in [`ci/tools/`](../ci/tools): `install_clang18.sh` (the
 clang-18 the build needs) and `bazel_test.sh` (bazel binary + dummy torch_tpu
@@ -39,17 +32,19 @@ override + flags, matching `build.sh`).
 - **JAX-only build.** The JAX path generates a dummy `torch_tpu` Bazel module
   (same as `build.sh` / `kokoro/.../presubmit.sh`), so no secrets are needed. The
   **Torch** extension/wheel additionally needs a `torch_tpu` checkout + a local
-  `torch`, and its meaningful tests need a TPU — out of scope here for now.
-- **Unit-test set.** `CPU_TEST_TARGETS` is a curated CPU-safe, OSS-loadable set.
-  `//rpc/...` and `//kv_cache/global_registry/...` (and the `kv_cache_store`
-  tests that depend on them) load Google-internal gRPC rules and are excluded by
-  design — see [`ci/wheel/BUILD.bazel`](../ci/wheel/BUILD.bazel). Expand the list
-  as more packages become CPU-runnable.
+  `torch` (deploy key) — out of scope here for now.
+- **Test set.** `CPU_TEST_TARGETS` are device-free; `DEVICE_TEST_TARGETS` exercise
+  the v5e. Both are a curated OSS-loadable set. `//rpc/...` and
+  `//kv_cache/global_registry/...` (and the `kv_cache_store` tests that depend on
+  them) load Google-internal gRPC rules and are excluded by design — see
+  [`ci/wheel/BUILD.bazel`](../ci/wheel/BUILD.bazel). Expand the lists as coverage
+  grows; for multiple parallel device tests, add a `--run_under` accelerator-lock
+  helper.
 
 ## One-time setup to turn this on
 
-1. **CPU runner** — grant `google/tpu-raiden` access to the `linux-x86-n2-32`
-   self-hosted scale-set (or adjust the label to one it already has).
+1. **Runner** — `google/tpu-raiden` must be granted the
+   `linux-x86-ct5lp-224-8tpu` scale-set (confirmed picking up jobs).
 2. **Bazel remote cache** — bucket `gs://tpu-raiden-bazel-cache` (already used by
    Kokoro); the runner service account needs object read (and write for the
    nightly's read-write cache).
@@ -60,6 +55,7 @@ override + flags, matching `build.sh`).
 
 `google/tpu-raiden` is a one-way Copybara mirror of an internal google3 repo
 (commits carry `PiperOrigin-RevId:`), so these files land **internally** and sync
-out — there is no external push/PR path. Once on `main`, the CPU workflows fire on
-PRs/pushes automatically, and `nightly.yml` can be kicked manually with
-`gh workflow run "CI - Nightly (CPU)"` by anyone with write access.
+out — there is no external push/PR path for merges. Opening a PR still triggers
+the workflows for validation. Once on `main`, presubmit fires on PRs and
+`nightly.yml` can be kicked with `gh workflow run "CI - Nightly"` by anyone with
+write access.
diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
index 28e314cd8..054f4d675 100644
--- a/.github/actionlint.yaml
+++ b/.github/actionlint.yaml
@@ -13,9 +13,8 @@
 # limitations under the License.
 
 # Declares the self-hosted runner labels used by the workflows so actionlint
-# does not flag them as unknown. CPU-only for now; add TPU scale-set labels here
-# (e.g. linux-x86-tpu7x-224-4tpu) once google/tpu-raiden is granted them.
+# does not flag them as unknown. Only the v5e scale-set is currently available to
+# this repo; n2-32 and tpu7x stayed queued/unassigned.
 self-hosted-runner:
   labels:
-    - linux-x86-n2-32              # CPU
-    - linux-x86-tpu7x-224-4tpu     # TPU v7 (probe, see tpu7x_smoke.yml)
+    - linux-x86-ct5lp-224-8tpu     # TPU v5e, 8 chips
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index f0d0a33c3..6aba2be06 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Nightly CPU build + unit tests. CPU-only (no TPU runners): builds the JAX
-# extension, runs the OSS-safe non-device unit tests, and builds the
+# Nightly build + tests on the v5e TPU runner (linux-x86-ct5lp-224-8tpu): builds
+# the JAX extension, runs the OSS-safe unit tests + a device test, and builds the
 # tpu_raiden_jax wheel (twine-checked) as a packaging smoke test. Also runs on
-# push to main and on demand. Files a tracking issue on failure. Device tests and
-# the torch wheel are intentionally out of scope until TPU runners are available;
-# see .github/README.md.
+# push to main and on demand. Files a tracking issue on failure. The torch wheel
+# is out of scope (needs a torch_tpu checkout); see .github/README.md.
 
-name: CI - Nightly (CPU)
+name: CI - Nightly
 
 on:
   workflow_dispatch:
@@ -42,13 +41,16 @@ env:
     //kv_cache:logical_block_manager_test
     //kv_cache:lru_cache_test
     //tpu_raiden/frameworks/jax:resharding_planner_test
+  DEVICE_TEST_TARGETS: >-
+    //tpu_raiden/frameworks/jax:raw_transfer_test_gl
 
 jobs:
   build_and_test:
-    name: "Nightly build + unit tests (CPU)"
-    runs-on: linux-x86-n2-32
+    name: "Nightly build + tests (v5e)"
+    runs-on: linux-x86-ct5lp-224-8tpu
     container:
       image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+      options: --privileged
     steps:
       - uses: actions/checkout@v5
       - name: Fix Git workspace ownership
@@ -60,11 +62,15 @@ jobs:
       - name: Build JAX extension
         run: ./build.sh jax
 
-      - name: Run CPU unit tests
+      - name: Verify TPU is visible
+        run: |
+          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
+
+      - name: Run unit tests + device test
         env:
           RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
           RAIDEN_REMOTE_UPLOAD: "true"
-        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
+        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS} ${DEVICE_TEST_TARGETS}
 
       - name: Build tpu_raiden_jax wheel (packaging smoke)
         run: |
diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml
index 398b7bc4b..43c9210d2 100644
--- a/.github/workflows/presubmit.yml
+++ b/.github/workflows/presubmit.yml
@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Per-PR gate. CPU-only: builds the JAX extension and runs the OSS-safe,
-# non-device unit tests + an import smoke check on a CPU runner. There are no TPU
-# lanes here on purpose -- this repo is not (yet) granted the self-hosted TPU
-# runner scale-sets, so all CI is kept on CPU. Device tests stay in Kokoro / on
-# a developer TPU VM until TPU runners are available. See .github/README.md.
+# Per-PR gate, running on the v5e TPU runner (linux-x86-ct5lp-224-8tpu): builds
+# the JAX extension, confirms the TPU is visible, runs the OSS-safe unit tests
+# plus one device test, and an import smoke check. The n2-32 / tpu7x runners were
+# not available to this repo (jobs stayed queued unassigned), so v5e is used for
+# everything for now. See .github/README.md.
 
-name: CI - Presubmit (CPU)
+name: CI - Presubmit
 
 on:
   workflow_dispatch:
@@ -35,23 +35,24 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: ${{ github.ref != 'main' }}
 
-# Curated OSS-safe, CPU-only unit tests (no TPU device, no Google-internal grpc).
-# //rpc/... and //kv_cache/global_registry/... are excluded by design -- see
-# ci/wheel/BUILD.bazel. Expand as more packages become CPU-runnable.
+# Curated OSS-safe tests. CPU_TEST_TARGETS run without a device; the device test
+# exercises the v5e. //rpc/... and //kv_cache/global_registry/... are excluded by
+# design (Google-internal grpc) -- see ci/wheel/BUILD.bazel.
 env:
   CPU_TEST_TARGETS: >-
     //kv_cache:logical_block_manager_test
     //kv_cache:lru_cache_test
     //tpu_raiden/frameworks/jax:resharding_planner_test
+  DEVICE_TEST_TARGETS: >-
+    //tpu_raiden/frameworks/jax:raw_transfer_test_gl
 
 jobs:
   build_and_test:
-    name: "Presubmit on CPU"
-    # Self-hosted CPU runner. NOTE: confirm this scale-set is granted to
-    # google/tpu-raiden; swap the label if your pool differs.
-    runs-on: linux-x86-n2-32
+    name: "Presubmit on v5e"
+    runs-on: linux-x86-ct5lp-224-8tpu
     container:
       image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+      options: --privileged
     steps:
       - uses: actions/checkout@v5
 
@@ -66,12 +67,22 @@ jobs:
       - name: Build JAX extension
         run: ./build.sh jax
 
-      - name: Run CPU unit tests
+      - name: Verify TPU is visible
+        run: |
+          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
+
+      - name: Run unit tests
         env:
           RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
           RAIDEN_REMOTE_UPLOAD: "false"
         run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
 
+      - name: Run device test (v5e)
+        env:
+          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
+          RAIDEN_REMOTE_UPLOAD: "false"
+        run: ./ci/tools/bazel_test.sh jax -- ${DEVICE_TEST_TARGETS}
+
       - name: Import smoke check (JAX mocked)
         run: |
           export PYTHONPATH="${GITHUB_WORKSPACE}:${GITHUB_WORKSPACE}/bazel-bin:${PYTHONPATH:-}"
diff --git a/.github/workflows/tpu7x_smoke.yml b/.github/workflows/tpu7x_smoke.yml
deleted file mode 100644
index b27879825..000000000
--- a/.github/workflows/tpu7x_smoke.yml
+++ /dev/null
@@ -1,61 +0,0 @@
-# Copyright 2026 Google LLC.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Probe workflow: validates whether the linux-x86-tpu7x-224-4tpu (TPU v7) runner
-# scale-set is available to this repo, and whether the ml-build container can see
-# the v7 chips and run a device test. If v7 is confirmed, TPU lanes can be folded
-# back into presubmit/nightly. Delete this file once that decision is made.
-
-name: CI - tpu7x (v7) smoke probe
-
-on:
-  workflow_dispatch:
-  pull_request:
-    branches:
-      - main
-
-permissions:
-  contents: read
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  tpu7x_smoke:
-    name: "Smoke on tpu7x (v7)"
-    runs-on: linux-x86-tpu7x-224-4tpu
-    container:
-      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
-      options: --privileged
-    steps:
-      - uses: actions/checkout@v5
-      - name: Fix Git workspace ownership
-        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
-      - name: Show runner / container info
-        run: |
-          echo "Runner: ${RUNNER_NAME:-?} on ${RUNNER_OS:-?}"
-          uname -a; nproc
-      - name: Install clang-18
-        run: ./ci/tools/install_clang18.sh
-      - name: Build JAX extension
-        run: ./build.sh jax
-      - name: Verify TPU v7 visible from container
-        run: |
-          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
-      - name: Run one TPU device test
-        env:
-          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
-          RAIDEN_REMOTE_UPLOAD: "false"
-        run: ./ci/tools/bazel_test.sh jax -- //tpu_raiden/frameworks/jax:raw_transfer_test_gl

From 5e8098acab70e35813efb3f1f9664df340cfd6d0 Mon Sep 17 00:00:00 2001
From: Yin Lin <yinlin@google.com>
Date: Thu, 18 Jun 2026 05:36:23 +0000
Subject: [PATCH 4/4] Split CI: CPU build+tests on n2-32, TPU device tests on
 v5e.

n2-32 and v5e are self-hosted scale-sets that cold-start slowly; jobs may sit
queued for several minutes before pickup. Keep CPU (presubmit/nightly) and TPU
(device_test) in separate workflows so a slow TPU runner never blocks the CPU
gate. v5e is used for device tests since tpu7x is unavailable.
---
 .github/README.md                 | 23 +++++----
 .github/actionlint.yaml           |  5 +-
 .github/workflows/device_test.yml | 77 +++++++++++++++++++++++++++++++
 .github/workflows/nightly.yml     | 28 +++++------
 .github/workflows/presubmit.yml   | 36 +++++----------
 5 files changed, 115 insertions(+), 54 deletions(-)
 create mode 100644 .github/workflows/device_test.yml

diff --git a/.github/README.md b/.github/README.md
index 2fb4e9e3f..45dedf66e 100644
--- a/.github/README.md
+++ b/.github/README.md
@@ -4,15 +4,17 @@ GitHub Actions CI for tpu-raiden, modeled on
 [`torch_tpu`](https://github.com/google-pytorch/torch_tpu)'s setup. The existing
 Kokoro jobs under [`kokoro/`](../kokoro) are kept and run alongside these.
 
-## Runner
+## Runners
 
-All build/test jobs run on the **v5e TPU runner `linux-x86-ct5lp-224-8tpu`**.
-This is the scale-set currently granted to `google/tpu-raiden`: jobs targeting
-`linux-x86-n2-32` (CPU) and `linux-x86-tpu7x-224-4tpu` (v7) stayed queued and
-unassigned, so they are not available here yet. v5e is a full host, so it runs
-the JAX build, the CPU unit tests, and TPU device tests. To add other runners
-later, grant the repo the scale-set, add the label to
-[`actionlint.yaml`](actionlint.yaml), and point a job's `runs-on:` at it.
+- **CPU** — `linux-x86-n2-32` runs the presubmit + nightly build/unit-test jobs.
+- **TPU v5e** — `linux-x86-ct5lp-224-8tpu` runs the device tests (`device_test.yml`).
+  v5e is used because the v7 (`linux-x86-tpu7x-224-4tpu`) scale-set is not
+  available to this repo.
+
+Both are self-hosted scale-sets and can take several minutes to cold-start a
+runner, so a freshly triggered job may sit `queued` for a while before a runner
+picks it up — that is normal, not a failure. CPU and TPU jobs are split across
+separate workflows so a slow TPU runner never blocks the CPU gate.
 
 ## Workflows
 
@@ -20,8 +22,9 @@ later, grant the repo the scale-set, add the label to
 | --- | --- | --- |
 | [`lint.yml`](workflows/lint.yml) | PR | Apache license headers (`addlicense`) + commit-message check. GitHub-hosted. |
 | [`clang_format.yml`](workflows/clang_format.yml) | PR | Google C++ style on changed `.cc/.h` via pinned `clang-format==18`. GitHub-hosted. |
-| [`presubmit.yml`](workflows/presubmit.yml) | PR | Build JAX extension, verify TPU visible, run OSS-safe unit tests + one device test + import smoke, on v5e. |
-| [`nightly.yml`](workflows/nightly.yml) | push to `main` / daily cron 09:00 UTC / dispatch | Build + unit tests + device test, build the `tpu_raiden_jax` wheel (twine-checked, uploaded as an artifact), file a tracking issue on failure. |
+| [`presubmit.yml`](workflows/presubmit.yml) | PR | Build JAX extension + OSS-safe CPU unit tests + import smoke, on `n2-32`. |
+| [`nightly.yml`](workflows/nightly.yml) | push to `main` / daily cron 09:00 UTC / dispatch | CPU build + unit tests + `tpu_raiden_jax` wheel (twine-checked, uploaded as an artifact), file a tracking issue on failure, on `n2-32`. |
+| [`device_test.yml`](workflows/device_test.yml) | PR / daily cron / dispatch | Build JAX extension, verify TPU visible, run device tests on v5e (`ct5lp`). |
 
 Shared logic lives in [`ci/tools/`](../ci/tools): `install_clang18.sh` (the
 clang-18 the build needs) and `bazel_test.sh` (bazel binary + dummy torch_tpu
diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
index 054f4d675..4e7504af4 100644
--- a/.github/actionlint.yaml
+++ b/.github/actionlint.yaml
@@ -13,8 +13,9 @@
 # limitations under the License.
 
 # Declares the self-hosted runner labels used by the workflows so actionlint
-# does not flag them as unknown. Only the v5e scale-set is currently available to
-# this repo; n2-32 and tpu7x stayed queued/unassigned.
+# does not flag them as unknown. CPU jobs run on n2-32; TPU device tests on the
+# v5e scale-set (the v7/tpu7x scale-set is not available to this repo).
 self-hosted-runner:
   labels:
+    - linux-x86-n2-32              # CPU
     - linux-x86-ct5lp-224-8tpu     # TPU v5e, 8 chips
diff --git a/.github/workflows/device_test.yml b/.github/workflows/device_test.yml
new file mode 100644
index 000000000..262f8cd33
--- /dev/null
+++ b/.github/workflows/device_test.yml
@@ -0,0 +1,77 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# TPU device tests on the v5e runner (linux-x86-ct5lp-224-8tpu). Kept separate
+# from the CPU presubmit/nightly so a slow-to-provision TPU runner never blocks
+# the CPU gate. v5e is used because the v7 (tpu7x) scale-set is not available to
+# this repo. Builds the JAX extension, confirms the TPU is visible, and runs the
+# device tests.
+
+name: CI - Device tests (v5e)
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+  schedule:
+    - cron: "0 9 * * *"  # 1 AM PST daily, alongside the nightly CPU build.
+
+permissions:
+  contents: read
+  issues: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+env:
+  DEVICE_TEST_TARGETS: >-
+    //tpu_raiden/frameworks/jax:raw_transfer_test_gl
+
+jobs:
+  device_tests:
+    name: "Device tests on v5e"
+    runs-on: linux-x86-ct5lp-224-8tpu
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+      options: --privileged
+    steps:
+      - uses: actions/checkout@v5
+      - name: Fix Git workspace ownership
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+
+      - name: Install clang-18
+        run: ./ci/tools/install_clang18.sh
+
+      - name: Build JAX extension
+        run: ./build.sh jax
+
+      - name: Verify TPU is visible
+        run: |
+          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
+
+      - name: Run device tests
+        env:
+          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
+          RAIDEN_REMOTE_UPLOAD: "false"
+        run: ./ci/tools/bazel_test.sh jax -- ${DEVICE_TEST_TARGETS}
+
+      - name: Upload test logs on failure
+        uses: actions/upload-artifact@v4
+        if: failure()
+        with:
+          name: bazel-testlogs-device-v5e
+          path: bazel-testlogs/
+          retention-days: 7
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 6aba2be06..a16759729 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Nightly build + tests on the v5e TPU runner (linux-x86-ct5lp-224-8tpu): builds
-# the JAX extension, runs the OSS-safe unit tests + a device test, and builds the
-# tpu_raiden_jax wheel (twine-checked) as a packaging smoke test. Also runs on
-# push to main and on demand. Files a tracking issue on failure. The torch wheel
-# is out of scope (needs a torch_tpu checkout); see .github/README.md.
+# Nightly CPU build + unit tests on linux-x86-n2-32: builds the JAX extension,
+# runs the OSS-safe non-device unit tests, and builds the tpu_raiden_jax wheel
+# (twine-checked) as a packaging smoke test. Also runs on push to main and on
+# demand. Files a tracking issue on failure. TPU device coverage lives in
+# device_test.yml (v5e); the torch wheel needs a torch_tpu checkout and is out of
+# scope. See .github/README.md.
 
-name: CI - Nightly
+name: CI - Nightly (CPU)
 
 on:
   workflow_dispatch:
@@ -41,16 +42,13 @@ env:
     //kv_cache:logical_block_manager_test
     //kv_cache:lru_cache_test
     //tpu_raiden/frameworks/jax:resharding_planner_test
-  DEVICE_TEST_TARGETS: >-
-    //tpu_raiden/frameworks/jax:raw_transfer_test_gl
 
 jobs:
   build_and_test:
-    name: "Nightly build + tests (v5e)"
-    runs-on: linux-x86-ct5lp-224-8tpu
+    name: "Nightly build + unit tests (CPU)"
+    runs-on: linux-x86-n2-32
     container:
       image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
-      options: --privileged
     steps:
       - uses: actions/checkout@v5
       - name: Fix Git workspace ownership
@@ -62,15 +60,11 @@ jobs:
       - name: Build JAX extension
         run: ./build.sh jax
 
-      - name: Verify TPU is visible
-        run: |
-          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
-
-      - name: Run unit tests + device test
+      - name: Run CPU unit tests
         env:
           RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
           RAIDEN_REMOTE_UPLOAD: "true"
-        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS} ${DEVICE_TEST_TARGETS}
+        run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
 
       - name: Build tpu_raiden_jax wheel (packaging smoke)
         run: |
diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml
index 43c9210d2..dd34f3063 100644
--- a/.github/workflows/presubmit.yml
+++ b/.github/workflows/presubmit.yml
@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Per-PR gate, running on the v5e TPU runner (linux-x86-ct5lp-224-8tpu): builds
-# the JAX extension, confirms the TPU is visible, runs the OSS-safe unit tests
-# plus one device test, and an import smoke check. The n2-32 / tpu7x runners were
-# not available to this repo (jobs stayed queued unassigned), so v5e is used for
-# everything for now. See .github/README.md.
+# Per-PR gate. CPU-only: builds the JAX extension and runs the OSS-safe,
+# non-device unit tests + an import smoke check on the CPU runner
+# (linux-x86-n2-32). TPU device coverage lives in device_test.yml (v5e) so a slow
+# TPU runner never blocks the PR gate. See .github/README.md.
 
-name: CI - Presubmit
+name: CI - Presubmit (CPU)
 
 on:
   workflow_dispatch:
@@ -35,24 +34,21 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: ${{ github.ref != 'main' }}
 
-# Curated OSS-safe tests. CPU_TEST_TARGETS run without a device; the device test
-# exercises the v5e. //rpc/... and //kv_cache/global_registry/... are excluded by
-# design (Google-internal grpc) -- see ci/wheel/BUILD.bazel.
+# Curated OSS-safe, CPU-only unit tests (no TPU device, no Google-internal grpc).
+# //rpc/... and //kv_cache/global_registry/... are excluded by design -- see
+# ci/wheel/BUILD.bazel. Expand as more packages become CPU-runnable.
 env:
   CPU_TEST_TARGETS: >-
     //kv_cache:logical_block_manager_test
     //kv_cache:lru_cache_test
     //tpu_raiden/frameworks/jax:resharding_planner_test
-  DEVICE_TEST_TARGETS: >-
-    //tpu_raiden/frameworks/jax:raw_transfer_test_gl
 
 jobs:
   build_and_test:
-    name: "Presubmit on v5e"
-    runs-on: linux-x86-ct5lp-224-8tpu
+    name: "Presubmit on CPU"
+    runs-on: linux-x86-n2-32
     container:
       image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
-      options: --privileged
     steps:
       - uses: actions/checkout@v5
 
@@ -67,22 +63,12 @@ jobs:
       - name: Build JAX extension
         run: ./build.sh jax
 
-      - name: Verify TPU is visible
-        run: |
-          python3 -c "import jax; d=jax.devices(); print('JAX devices:', d); assert any('TPU' in str(x).upper() for x in d), 'no TPU visible'; print('TPU visible OK')"
-
-      - name: Run unit tests
+      - name: Run CPU unit tests
         env:
           RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
           RAIDEN_REMOTE_UPLOAD: "false"
         run: ./ci/tools/bazel_test.sh jax -- ${CPU_TEST_TARGETS}
 
-      - name: Run device test (v5e)
-        env:
-          RAIDEN_REMOTE_CACHE: "https://storage.googleapis.com/tpu-raiden-bazel-cache"
-          RAIDEN_REMOTE_UPLOAD: "false"
-        run: ./ci/tools/bazel_test.sh jax -- ${DEVICE_TEST_TARGETS}
-
       - name: Import smoke check (JAX mocked)
         run: |
           export PYTHONPATH="${GITHUB_WORKSPACE}:${GITHUB_WORKSPACE}/bazel-bin:${PYTHONPATH:-}"