diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2650bcd0..ab72c997 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -45,7 +45,6 @@ jobs:
       NVIDIA_DRIVER_CAPABILITIES: all
       NVIDIA_VISIBLE_DEVICES: all
       NVIDIA_DISABLE_REQUIRE: 1
-      DOCS_MAX_VERSIONS: "4"  # Max number of release versions to keep
     container: *container_template
     steps:
       - uses: actions/checkout@v4
@@ -59,16 +58,6 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-pip-docs-
 
-      - name: Restore previous docs output
-        if: github.event_name == 'push'
-        uses: actions/cache@v4
-        with:
-          path: docs/build/html
-          key: docs-output-${{ github.repository }}-${{ github.ref_name }}
-          restore-keys: |
-            docs-output-${{ github.repository }}-${{ github.ref_name }}-
-            docs-output-${{ github.repository }}-
-
       - name: Build docs
         shell: bash
         run: |
@@ -82,41 +71,17 @@ jobs:
           if [[ "${GITHUB_REF}" == refs/tags/v* ]]; then
             VERSION="${GITHUB_REF_NAME}"
             echo "Building docs for release tag ${VERSION}..."
-
-            # Build only this version into its own subdirectory
             sphinx-build source build/html/${VERSION}
-
-            cd build/html
-
-            # Prune old release versions beyond the window
-            mapfile -t TAG_DIRS < <(ls -d v*/ 2>/dev/null | sort -V)
-            while [[ ${#TAG_DIRS[@]} -gt ${DOCS_MAX_VERSIONS} ]]; do
-              echo "Pruning old version: ${TAG_DIRS[0]}"
-              rm -rf "${TAG_DIRS[0]}"
-              TAG_DIRS=("${TAG_DIRS[@]:1}")
-            done
-
-            # Generate versions.json and root index.html
-            python3 ${GITHUB_WORKSPACE}/docs/scripts/generate_versions_json.py \
-              --build-dir .
-
           else
             echo "Building dev docs for main branch..."
-            # Build only main/ — don't touch existing version directories
-            rm -rf build/html/main
             sphinx-build source build/html/main
-
-            cd build/html
-
-            # Generate versions.json and root index.html
-            python3 ${GITHUB_WORKSPACE}/docs/scripts/generate_versions_json.py \
-              --build-dir .
           fi
 
       - name: Upload docs artifact
         if: github.event_name == 'push'
-        uses: actions/upload-pages-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
+          name: docs-build
           path: ${{ github.workspace }}/docs/build/html
 
   test:
@@ -143,18 +108,77 @@ jobs:
   publish:
     if: github.event_name == 'push'
     needs: build
-    runs-on: Linux
+    runs-on: ubuntu-latest
     permissions:
-      pages: write
-      id-token: write
+      contents: write  # Required to push to gh-pages branch
+    env:
+      DOCS_MAX_VERSIONS: "4"  # Max number of release versions to keep
     steps:
+      - name: Checkout source repo (for scripts)
+        uses: actions/checkout@v4
+
+      - name: Determine deploy directory
+        id: vars
+        run: |
+          if [[ "${GITHUB_REF}" == refs/tags/v* ]]; then
+            echo "deploy_dir=${GITHUB_REF_NAME}" >> $GITHUB_OUTPUT
+            echo "is_tag=true" >> $GITHUB_OUTPUT
+          else
+            echo "deploy_dir=main" >> $GITHUB_OUTPUT
+            echo "is_tag=false" >> $GITHUB_OUTPUT
+          fi
+
       - name: Download docs artifact
         uses: actions/download-artifact@v4
         with:
-          name: github-pages
+          name: docs-build
+          path: docs-build
+
+      # Deploy only the specific version subdirectory to gh-pages.
+      # Using target-folder ensures other version dirs are never touched.
+      - name: Deploy docs subdir to gh-pages
+        uses: JamesIves/github-pages-deploy-action@v4
+        with:
+          branch: gh-pages
+          folder: docs-build/${{ steps.vars.outputs.deploy_dir }}
+          target-folder: ${{ steps.vars.outputs.deploy_dir }}
+          clean: true
+          token: ${{ secrets.GITHUB_TOKEN }}
+          commit-message: "docs: update ${{ steps.vars.outputs.deploy_dir }}"
+
+      - name: Checkout gh-pages for metadata update
+        uses: actions/checkout@v4
+        with:
+          ref: gh-pages
+          path: gh-pages
+
+      - name: Prune old release versions and regenerate metadata
+        env:
+          DEPLOY_DIR: ${{ steps.vars.outputs.deploy_dir }}
+          IS_TAG: ${{ steps.vars.outputs.is_tag }}
+        run: |
+          cd gh-pages
+
+          # Remove outdated release versions when a new tag is pushed
+          if [[ "${IS_TAG}" == "true" ]]; then
+            mapfile -t TAG_DIRS < <(ls -d v*/ 2>/dev/null | sort -V)
+            while [[ ${#TAG_DIRS[@]} -gt ${DOCS_MAX_VERSIONS} ]]; do
+              echo "Pruning old version: ${TAG_DIRS[0]}"
+              rm -rf "${TAG_DIRS[0]}"
+              TAG_DIRS=("${TAG_DIRS[@]:1}")
+            done
+          fi
+
+          # Regenerate versions.json and root index.html from whatever dirs exist
+          python3 $GITHUB_WORKSPACE/docs/scripts/generate_versions_json.py --build-dir .
 
-      - name: Deploy GitHub Pages
-        uses: actions/deploy-pages@v4
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
+          git add -A
+          git diff --staged --quiet \
+            && echo "No metadata changes to commit" \
+            || git commit -m "docs: update metadata for ${DEPLOY_DIR}"
+          git push origin gh-pages
 
 
   release-build:
diff --git a/.github/workflows/tests/test_docs_publish.yml b/.github/workflows/tests/test_docs_publish.yml
new file mode 100644
index 00000000..ad44e5ef
--- /dev/null
+++ b/.github/workflows/tests/test_docs_publish.yml
@@ -0,0 +1,220 @@
+name: Test docs publish logic
+
+on:
+  workflow_dispatch:
+    inputs:
+      scenario:
+        description: "Test scenario: main_push or tag_push"
+        required: true
+        default: main_push
+
+jobs:
+  # -----------------------------------------------------------------------
+  # Scenario A: push to main branch — existing v0.1.0, v0.2.0 must survive
+  # -----------------------------------------------------------------------
+  test-main-push:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up fake docs-build artifact (main branch build output)
+        run: |
+          mkdir -p docs-build/main
+          echo "<html><body>main docs</body></html>" > docs-build/main/index.html
+
+      - name: Set up fake gh-pages with existing versioned dirs
+        run: |
+          mkdir -p gh-pages/v0.1.0 gh-pages/v0.2.0 gh-pages/main
+          echo "<html>v0.1.0</html>" > gh-pages/v0.1.0/index.html
+          echo "<html>v0.2.0</html>" > gh-pages/v0.2.0/index.html
+          echo "<html>old main</html>" > gh-pages/main/index.html
+
+      - name: Simulate publish step — update main subdir only
+        run: |
+          DEPLOY_DIR=main
+          IS_TAG=false
+          DOCS_MAX_VERSIONS=4
+
+          # Replace only the main subdir (mirrors the JamesIves deploy + clean)
+          rm -rf gh-pages/${DEPLOY_DIR}
+          cp -r docs-build/${DEPLOY_DIR} gh-pages/${DEPLOY_DIR}
+
+          # No pruning for non-tag builds
+          if [[ "${IS_TAG}" == "true" ]]; then
+            cd gh-pages
+            mapfile -t TAG_DIRS < <(ls -d v*/ 2>/dev/null | sort -V)
+            while [[ ${#TAG_DIRS[@]} -gt ${DOCS_MAX_VERSIONS} ]]; do
+              echo "Pruning old version: ${TAG_DIRS[0]}"
+              rm -rf "${TAG_DIRS[0]}"
+              TAG_DIRS=("${TAG_DIRS[@]:1}")
+            done
+            cd ..
+          fi
+
+          # Regenerate metadata
+          python3 docs/scripts/generate_versions_json.py --build-dir gh-pages
+
+      - name: Assert — v0.1.0 and v0.2.0 still present
+        run: |
+          echo "=== gh-pages structure ==="
+          find gh-pages -maxdepth 1 | sort
+
+          [ -d gh-pages/v0.1.0 ] || (echo "FAIL: v0.1.0 was removed!" && exit 1)
+          [ -d gh-pages/v0.2.0 ] || (echo "FAIL: v0.2.0 was removed!" && exit 1)
+          [ -f gh-pages/main/index.html ] || (echo "FAIL: main/index.html missing!" && exit 1)
+          grep -q "main docs" gh-pages/main/index.html || (echo "FAIL: main/index.html not updated!" && exit 1)
+          [ -f gh-pages/versions.json ] || (echo "FAIL: versions.json missing!" && exit 1)
+          [ -f gh-pages/index.html ] || (echo "FAIL: root index.html missing!" && exit 1)
+
+          echo "=== versions.json ==="
+          cat gh-pages/versions.json
+          python3 -c "
+          import json, sys
+          data = json.load(open('gh-pages/versions.json'))
+          names = [v['name'] for v in data['versions']]
+          assert 'v0.1.0' in names, f'v0.1.0 missing from versions.json: {names}'
+          assert 'v0.2.0' in names, f'v0.2.0 missing from versions.json: {names}'
+          assert 'main' in names, f'main missing from versions.json: {names}'
+          print('PASS: versions.json contains all expected versions')
+          "
+          echo "PASS: main_push scenario — existing versions preserved"
+
+  # -----------------------------------------------------------------------
+  # Scenario B: push of tag v0.3.0 — v0.3.0 added, old dirs still present
+  # -----------------------------------------------------------------------
+  test-tag-push:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up fake docs-build artifact (tag build output)
+        run: |
+          mkdir -p docs-build/v0.3.0
+          echo "<html><body>v0.3.0 docs</body></html>" > docs-build/v0.3.0/index.html
+
+      - name: Set up fake gh-pages with existing versioned dirs
+        run: |
+          mkdir -p gh-pages/v0.1.0 gh-pages/v0.2.0 gh-pages/main
+          echo "<html>v0.1.0</html>" > gh-pages/v0.1.0/index.html
+          echo "<html>v0.2.0</html>" > gh-pages/v0.2.0/index.html
+          echo "<html>main</html>" > gh-pages/main/index.html
+
+      - name: Simulate publish step — add v0.3.0 subdir
+        run: |
+          DEPLOY_DIR=v0.3.0
+          IS_TAG=true
+          DOCS_MAX_VERSIONS=4
+
+          rm -rf gh-pages/${DEPLOY_DIR}
+          cp -r docs-build/${DEPLOY_DIR} gh-pages/${DEPLOY_DIR}
+
+          if [[ "${IS_TAG}" == "true" ]]; then
+            cd gh-pages
+            mapfile -t TAG_DIRS < <(ls -d v*/ 2>/dev/null | sort -V)
+            while [[ ${#TAG_DIRS[@]} -gt ${DOCS_MAX_VERSIONS} ]]; do
+              echo "Pruning old version: ${TAG_DIRS[0]}"
+              rm -rf "${TAG_DIRS[0]}"
+              TAG_DIRS=("${TAG_DIRS[@]:1}")
+            done
+            cd ..
+          fi
+
+          python3 docs/scripts/generate_versions_json.py --build-dir gh-pages
+
+      - name: Assert — all versions present, latest is v0.3.0
+        run: |
+          echo "=== gh-pages structure ==="
+          find gh-pages -maxdepth 1 | sort
+
+          [ -d gh-pages/v0.1.0 ] || (echo "FAIL: v0.1.0 was removed!" && exit 1)
+          [ -d gh-pages/v0.2.0 ] || (echo "FAIL: v0.2.0 was removed!" && exit 1)
+          [ -d gh-pages/v0.3.0 ] || (echo "FAIL: v0.3.0 was missing!" && exit 1)
+          [ -d gh-pages/main ]   || (echo "FAIL: main was removed!" && exit 1)
+
+          echo "=== versions.json ==="
+          cat gh-pages/versions.json
+          python3 -c "
+          import json, sys
+          data = json.load(open('gh-pages/versions.json'))
+          names = [v['name'] for v in data['versions']]
+          assert 'v0.3.0' in names, f'v0.3.0 missing: {names}'
+          assert 'v0.1.0' in names, f'v0.1.0 missing: {names}'
+          assert 'v0.2.0' in names, f'v0.2.0 missing: {names}'
+          assert 'main' in names, f'main missing: {names}'
+          assert data['latest'] == 'v0.3.0', f'latest should be v0.3.0, got {data[\"latest\"]}'
+          print('PASS: versions.json correct, latest =', data['latest'])
+          "
+          echo "PASS: tag_push scenario — new version added, others preserved"
+
+  # -----------------------------------------------------------------------
+  # Scenario C: pruning kicks in when tag count exceeds DOCS_MAX_VERSIONS
+  # -----------------------------------------------------------------------
+  test-prune-old-versions:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up fake docs-build for v0.5.0 (5th tag)
+        run: |
+          mkdir -p docs-build/v0.5.0
+          echo "<html>v0.5.0</html>" > docs-build/v0.5.0/index.html
+
+      - name: Set up fake gh-pages with 4 existing tag dirs (at the limit)
+        run: |
+          for v in v0.1.0 v0.2.0 v0.3.0 v0.4.0; do
+            mkdir -p gh-pages/${v}
+            echo "<html>${v}</html>" > gh-pages/${v}/index.html
+          done
+          mkdir -p gh-pages/main
+          echo "<html>main</html>" > gh-pages/main/index.html
+
+      - name: Simulate publish step for v0.5.0 (triggers prune)
+        run: |
+          DEPLOY_DIR=v0.5.0
+          IS_TAG=true
+          DOCS_MAX_VERSIONS=4
+
+          rm -rf gh-pages/${DEPLOY_DIR}
+          cp -r docs-build/${DEPLOY_DIR} gh-pages/${DEPLOY_DIR}
+
+          if [[ "${IS_TAG}" == "true" ]]; then
+            cd gh-pages
+            mapfile -t TAG_DIRS < <(ls -d v*/ 2>/dev/null | sort -V)
+            while [[ ${#TAG_DIRS[@]} -gt ${DOCS_MAX_VERSIONS} ]]; do
+              echo "Pruning old version: ${TAG_DIRS[0]}"
+              rm -rf "${TAG_DIRS[0]}"
+              TAG_DIRS=("${TAG_DIRS[@]:1}")
+            done
+            cd ..
+          fi
+
+          python3 docs/scripts/generate_versions_json.py --build-dir gh-pages
+
+      - name: Assert — oldest v0.1.0 pruned, max 4 tags kept
+        run: |
+          echo "=== gh-pages structure ==="
+          find gh-pages -maxdepth 1 | sort
+
+          [ ! -d gh-pages/v0.1.0 ] || (echo "FAIL: v0.1.0 should have been pruned!" && exit 1)
+          [ -d gh-pages/v0.2.0 ] || (echo "FAIL: v0.2.0 was over-pruned!" && exit 1)
+          [ -d gh-pages/v0.3.0 ] || (echo "FAIL: v0.3.0 was over-pruned!" && exit 1)
+          [ -d gh-pages/v0.4.0 ] || (echo "FAIL: v0.4.0 was over-pruned!" && exit 1)
+          [ -d gh-pages/v0.5.0 ] || (echo "FAIL: v0.5.0 was not added!" && exit 1)
+          [ -d gh-pages/main ]   || (echo "FAIL: main was removed by pruning!" && exit 1)
+
+          TAG_COUNT=$(ls -d gh-pages/v*/ 2>/dev/null | wc -l)
+          [ "${TAG_COUNT}" -le 4 ] || (echo "FAIL: ${TAG_COUNT} tag dirs exceed DOCS_MAX_VERSIONS=4" && exit 1)
+
+          echo "=== versions.json ==="
+          cat gh-pages/versions.json
+          python3 -c "
+          import json
+          data = json.load(open('gh-pages/versions.json'))
+          names = [v['name'] for v in data['versions']]
+          assert 'v0.1.0' not in names, f'v0.1.0 should be pruned from versions.json: {names}'
+          assert data['latest'] == 'v0.5.0', f'latest should be v0.5.0, got {data[\"latest\"]}'
+          tag_count = sum(1 for v in data['versions'] if v['type'] == 'tag')
+          assert tag_count <= 4, f'Too many tags in versions.json: {tag_count}'
+          print('PASS: pruning correct, latest =', data['latest'], ', tag count =', tag_count)
+          "
+          echo "PASS: prune scenario — oldest version removed, within limit"