diff --git a/.clang-format b/.clang-format
new file mode 100644
index 00000000..3cc4fd76
--- /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 00000000..45dedf66
--- /dev/null
+++ b/.github/README.md
@@ -0,0 +1,64 @@
+# tpu-raiden CI/CD
+
+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.
+
+## Runners
+
+- **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
+
+| 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 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
+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` (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. **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).
+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 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
new file mode 100644
index 00000000..4e7504af
--- /dev/null
+++ b/.github/actionlint.yaml
@@ -0,0 +1,21 @@
+# 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 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/clang_format.yml b/.github/workflows/clang_format.yml
new file mode 100644
index 00000000..1f05375b
--- /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/device_test.yml b/.github/workflows/device_test.yml
new file mode 100644
index 00000000..262f8cd3
--- /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/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 00000000..40201a8e
--- /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 00000000..a1675972
--- /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 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 (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 00000000..dd34f306
--- /dev/null
+++ b/.github/workflows/presubmit.yml
@@ -0,0 +1,82 @@
+# 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 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 (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"
+    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/BUILD b/BUILD
index a2353746..5d1de6f2 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([
diff --git a/ci/tools/bazel_test.sh b/ci/tools/bazel_test.sh
new file mode 100755
index 00000000..3840a8fa
--- /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 00000000..2386c44e
--- /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