diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000000..b4351a130c
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,56 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - master
+      - 'release/*'
+  pull_request:
+    branches:
+      - master
+      - 'release/*'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+env:
+  PYTHON_VERSIONS: "3.10.13,3.12.8"
+
+jobs:
+  call-build-wheels:
+    strategy:
+      matrix:
+        rocm-version: ["6.3.3"]
+    uses: jax/rocm-jax/.github/workflows/reuseable/build-wheels.yml@master
+    with:
+      python-versions: ${{ env.PYTHON_VERSIONS }}
+      rocm-version: ${{ matrix.rocm-version }}
+  call-build-docker:
+    strategy:
+      matrix:
+        python-version: ["3.10.13,3.12.8"]
+        rocm-version: ["6.3.3"]
+    uses: jax/rocm-jax/.github/workflows/reusable/build-docker.yml@master
+    with:
+      rocm-version: ${{ matrix.python-version }}
+  run-python-unit-tests:
+    strategy:
+      matrix:
+        rocm-version: ["6.3.3"]
+    steps:
+      - name: Checkout plugin repo
+        uses: actions/checkout@v4
+      - name: Checkout JAX repo
+        uses: actions/chekcout@v4
+        with:
+          repository: jax-ml/jax
+          path: jax
+      - name: Run tests
+        env:
+          GPU_COUNT: "8"
+          GFX: "gfx90a"
+          ROCM_VERSION: ${{ matrix.rocm-version }}
+        run: |
+          python3 build/ci_build test "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" --test-cmd "pytest tests/core_test.py tests/linalg_test.py"
+
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
new file mode 100644
index 0000000000..91ad26d3b0
--- /dev/null
+++ b/.github/workflows/nightly.yml
@@ -0,0 +1,50 @@
+name: Nightly
+
+on:
+  schedule:
+    - cron: "0 1 * * *"
+
+concurrency:
+  cancel-in-progress: true
+
+env:
+  PYTHON_VERSIONS: "3.10.13,3.12.8"
+
+jobs:
+  call-build-wheels:
+    strategy:
+      matrix:
+        rocm-version: ["6.3.3"]
+    uses: jax/rocm-jax/.github/workflows/reuseable/build-wheels.yml@master
+    with:
+      python-versions: ${{ env.PYTHON_VERSIONS }}
+      rocm-version: ${{ matrix.rocm-version }}
+  call-build-docker:
+    strategy:
+      matrix:
+        python-version: ["3.10.13,3.12.8"]
+        rocm-version: ["6.3.3"]
+    uses: jax/rocm-jax/.github/workflows/reusable/build-docker.yml@master
+    with:
+      rocm-version: ${{ matrix.python-version }}
+  run-python-unit-tests:
+    strategy:
+      matrix:
+        rocm-version: ["6.3.3"]
+    steps:
+      - name: Checkout plugin repo
+        uses: actions/checkout@v4
+      - name: Checkout JAX repo
+        uses: actions/chekcout@v4
+        with:
+          repository: jax-ml/jax
+          path: jax
+      - name: Run tests
+        env:
+          GPU_COUNT: "8"
+          GFX: "gfx90a"
+          ROCM_VERSION: ${{ matrix.rocm-version }}
+        run: |
+          python3 build/ci_build test "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" --test-cmd "pytest tests"
+
+
diff --git a/.github/workflows/reusable/build-docker.yml b/.github/workflows/reusable/build-docker.yml
new file mode 100644
index 0000000000..893bb306f3
--- /dev/null
+++ b/.github/workflows/reusable/build-docker.yml
@@ -0,0 +1,49 @@
+name: Build docker
+
+on:
+  workflow_call:
+    inputs:
+      rocm-version:
+        required: true
+        type: string
+      artifact-prefix:
+        default: 'plugin_wheels'
+
+jobs:
+  build-docker:
+    runs-on:  mi-250
+    steps:
+      - name: Clean up old runs
+        run: |
+          ls
+          # Make sure that we own all of the files so that we have permissions to delete them
+          docker run -v "./:/rocm-jax" ubuntu /bin/bash -c "chown -R $UID /rocm-jax/* || true"
+          # Remove any old work directories from this machine
+          rm -rf * || true
+          ls
+      - name: Print system info
+        run: |
+          whoami
+          printenv
+          df -h
+      - uses: actions/checkout@v4
+      - name: Download wheel artifacts
+        uses: actions/download-artifacts@v4
+          name: ${{ inputs.artifact-prefix }}_r${{ inputs.rocm-version }}
+          path: ./wheelhouse
+      - name: Build JAX docker image
+        run: |
+          python3 build/ci_build \
+            --rocm-version ${{ inputs.rocm-version }}
+            build_dockers
+      - name: Authenticate to GitHub Container Registry
+        run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+      - name: Push docker images
+        env:
+          ROCM_VERSION: ${{ inputs.rocm-version }}
+        run: |
+          docker tag "jax-ubu22.rocm${ROCM_VERSION//.}" "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}"
+          docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}"
+          docker push "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}"
+          docker push "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}"
+
diff --git a/.github/workflows/reusable/build-wheels.yml b/.github/workflows/reusable/build-wheels.yml
new file mode 100644
index 0000000000..656444d8a2
--- /dev/null
+++ b/.github/workflows/reusable/build-wheels.yml
@@ -0,0 +1,44 @@
+name: Build wheels
+
+on:
+  workflow_call:
+    inputs:
+      python-versions:
+        required: true
+        type: string
+      rocm-version:
+        required: true
+        type: string
+      artifact-prefix:
+        default: 'plugin_wheels'
+
+ build-plugin-docker:
+    runs-on:  mi-250
+    steps:
+      - name: Clean up old runs
+        run: |
+          ls
+          # Make sure that we own all of the files so that we have permissions to delete them
+          docker run -v "./:/rocm-jax" ubuntu /bin/bash -c "chown -R $UID /rocm-jax/* || true"
+          # Remove any old work directories from this machine
+          rm -rf * || true
+          ls
+      - name: Print system info
+        run: |
+          whoami
+          printenv
+          df -h
+      - uses: actions/checkout@v4
+      - name: Build plugin wheels
+        run: |
+          python3 build/ci_build \
+            --compiler clang \
+            --python-versions ${{ inputs.python-versions }} \
+            --rocm-version ${{ inputs.rocm-version }} \
+            dist_wheels
+      - name: Archive plugin wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ inputs.artifact-prefix }}_r${{ env.ROCM_VERSION }}
+          path: ./wheelhouse/*.whl
+