--- a
+++ b/.github/workflows/bionemo-subpackage-ci.yml
@@ -0,0 +1,226 @@
+name: BioNeMo Sub-Package Workflow
+
+on:
+  # To test or publish sub-packages or adjustments to this workflow that are branched in PR's, manually dispatch this workflow on the PR's branch here: https://github.com/NVIDIA/bionemo-framework/actions/workflows/bionemo-subpackage-ci.yml.
+  workflow_dispatch:
+    inputs:
+      subpackages:
+        description: "BioNeMo sub-packages (comma-separated) to test or publish."
+        required: true
+        type: string
+      test:
+        description: "Test the sub-packages before publishing to PyPI. Strongly recommended for production releases to PyPI. Can be disabled when staging sub-packages on Test PyPI or publishing circular dependencies to PyPI."
+        required: false
+        type: boolean
+        default: true
+      publish:
+        description: "Publish the built package to PyPI. If testing is specified, requires that all sub-package tests succeed based on dependencies published to Test PyPI or PyPI."
+        required: false
+        type: boolean
+        default: false
+      pypi:
+        description: "Publish to PyPI instead of Test PyPI."
+        required: false
+        type: boolean
+        default: false
+      version_overwrite:
+        description: "Overwrite the published version of the sub-package. (Sets skip-existing to False. Requires deleting existing wheels and other artifacts on PyPI.)"
+        required: false
+        type: boolean
+        default: false
+      build_framework:
+        description: "Build framework to use for building and publishing."
+        type: choice
+        options:
+          - "python"
+          - "rust_pyo3_maturin"
+        default: "python"
+        required: true
+      python_version:
+        description: "Python version to use for testing and publishing."
+        required: false
+        type: string
+        default: "3.12"
+      gpu_runner:
+        description: "Specify a GPU runner for testing on NVIDIA GitHub Actions. (For a list of available runners, refer to: https://docs.gha-runners.nvidia.com/runners/)"
+        required: false
+        type: string
+        default: "linux-amd64-gpu-l4-latest-1"
+      cuda_version:
+        description: "NVIDIA CUDA container version to use for testing."
+        required: false
+        type: string
+        default: "nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04"
+
+jobs:
+  configure-workflow-packages:
+    name: "[Configure Workflow Packages] Identify sub-packages for testing and publishing."
+    runs-on: ubuntu-latest
+    outputs:
+      workflow_packages: ${{ steps.parse-dispatch-packages.outputs.dispatch_packages }}
+    steps:
+      - id: parse-dispatch-packages
+        if: ${{ github.event_name == 'workflow_dispatch' }}
+        name: Parse the sub-packages specified in the workflow dispatch.
+        run: |
+          # Send the parsed list of sub-packages to the next job.
+          dispatch_packages=$(echo '${{ github.event.inputs.subpackages }}' | jq -R -c 'split(",")')
+          echo "dispatch_packages=$dispatch_packages" >> "$GITHUB_OUTPUT"
+          echo "[BioNeMo Sub-Package CI] Sub-packages to stage: $dispatch_packages"
+
+  install-and-test:
+    needs: configure-workflow-packages
+    # Check if the previous job has any staged packages to test and publish.
+    if: ${{ needs.configure-workflow-packages.outputs.workflow_packages != '[]' }}
+    strategy:
+      matrix:
+        package: ${{ fromJson(needs.configure-workflow-packages.outputs.workflow_packages) }}
+      fail-fast: false  # Prevent all matrix jobs from failing if one fails.
+    name: "[${{ matrix.package }}] Install and test sub-package."
+    # Use GPU runner only when testing, otherwise use a standard runner.
+    runs-on: ${{ github.event.inputs.test == 'true' && github.event.inputs.gpu_runner || 'ubuntu-latest' }}
+    container:
+      # GPU jobs must run in a container. Use a fresh CUDA base container for package installation and testing.
+      # If testing is disabled, use a lightweight container to quickly skip this job.
+      image: ${{ github.event.inputs.test == 'true' && github.event.inputs.cuda_version || 'ubuntu:latest' }}
+    steps:
+      # Silently skip all steps if testing is disabled, which does not block building or publishing.
+      - name: Install git and system dependencies.
+        if: ${{ github.event.inputs.test == 'true' }}
+        run: |
+          apt-get update
+          apt-get install -qyy git curl lsb-release build-essential
+      - uses: actions/checkout@v4
+        if: ${{ github.event.inputs.test == 'true' }}
+        with:
+          fetch-depth: 0
+          submodules: "recursive"
+      - uses: actions/setup-python@v5
+        if: ${{ github.event.inputs.test == 'true' }}
+        with:
+          python-version: ${{ github.event.inputs.python_version }}
+      - id: install-rust
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Install Rust.
+        run: |
+          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+          . $HOME/.cargo/env
+          rustc --version
+          cargo --version
+          rustup --version
+      - id: install-subpackage-core
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Install sub-package.
+        run: |
+          # Setup environment, i.e. add Rust to PATH and silence pip root user warnings.
+          . $HOME/.cargo/env
+          # Install sub-package and dependencies.
+          pip install --upgrade pip setuptools uv maturin
+          # Install required core & optional [test] dependencies.
+          uv pip install --no-cache --system pytest sub-packages/${{ matrix.package }}[test]
+      - id: install-subpackage-post
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Install sub-package dependencies that need to be installed after the core dependencies.
+        run: |
+          # DEV: Post-install dependencies are configured in [project.optional-dependencies].
+          # `uv pip install --extra <optional-dependency> -r <pyproject.toml>` tracks
+          # post-dependencies in the pyproject.toml and avoids installing core dependencies
+          # redundantly, which causes errors with incompatible --config-setting.
+
+          # TransformerEngine
+          uv pip install --no-cache --no-build-isolation --system --extra te -r sub-packages/${{ matrix.package }}/pyproject.toml || echo "[BioNeMo Sub-Package CI] TE will not be installed."
+
+          # # Apex
+          # # NOTE: --cpp_ext and --cuda_ext are required for building fused Apex kernels.
+          # uv pip install --no-cache --no-build-isolation --system --config-setting="--build-option=--cpp_ext" --config-setting="--build-option=--cuda_ext" --extra apex -r sub-packages/${{ matrix.package }}/pyproject.toml || echo "[BioNeMo Sub-Package CI] Apex will not be installed."
+      - id: test-dispatch-subpackage
+        if: ${{ github.event.inputs.test == 'true' }}
+        name: Test sub-package.
+        run: pytest -vv sub-packages/${{ matrix.package }}
+
+  build-pypi:
+    # Build distributions from either the workflow dispatch or PR.
+    # Validate building before merging or publishing.
+    needs: [configure-workflow-packages, install-and-test]
+    if: ${{ needs.configure-workflow-packages.outputs.workflow_packages != '[]' && github.event.inputs.publish == 'true' }}
+    outputs:
+      staged_packages: ${{ needs.configure-workflow-packages.outputs.workflow_packages }}
+    strategy:
+      matrix:
+        package: ${{ fromJson(needs.configure-workflow-packages.outputs.workflow_packages) }}
+      fail-fast: false  # Prevent all matrix jobs from failing if one fails.
+    name: "[${{ matrix.package }}] Build the sub-package."
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ github.event.inputs.python_version }}
+      - id: build-package
+        name: Build a binary wheel and a source tarball for the sub-package.
+        run: |
+          if [[ "${{ github.event.inputs.test }}" != "true" && "${{ github.event.inputs.version_overwrite }}" != "true" ]]; then
+            # For untested sub-packages, append '-dev' to the version for PyPI.
+            sed -i 's/[[:space:]]*$//' sub-packages/${{ matrix.package }}/VERSION
+            sed -i 's/$/-dev/' sub-packages/${{ matrix.package }}/VERSION
+          fi
+          # Build the sub-package.
+          if [[ "${{ github.event.inputs.build_framework }}" == "python" ]]; then
+            pip install build
+            python -m build sub-packages/${{ matrix.package }}
+          elif [[ "${{ github.event.inputs.build_framework }}" == "rust_pyo3_maturin" ]]; then
+            # Install maturin[zig] to build the Rust sub-package with compatibility for manylinux_X_Y using zig.
+            pip install maturin[zig]
+            maturin build --release --zig -m sub-packages/${{ matrix.package }}/Cargo.toml
+          fi
+      - id: upload-distribution
+        name: Upload distribution packages to the workflow.
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.package }}-build-artifacts
+          path: ${{ github.event.inputs.build_framework == 'rust_pyo3_maturin' && format('sub-packages/{0}/target/wheels', matrix.package) || format('sub-packages/{0}/dist', matrix.package) }}
+
+  publish-to-pypi:
+    needs: [build-pypi, install-and-test]
+    # Require staged sub-package builds for publishing to PyPI.
+    if: ${{ needs.build-pypi.result == 'success' }}
+    strategy:
+      matrix:
+        package: ${{ fromJson(needs.build-pypi.outputs.staged_packages) }}
+      fail-fast: false  # Prevent all matrix jobs from failing if one fails.
+    name: Publish ${{ matrix.package }} to PyPI.
+    runs-on: ubuntu-latest
+    environment:
+      name: ${{ github.event.inputs.pypi && 'pypi' || 'testpypi' }}
+      url: ${{ github.event.inputs.pypi && format('https://pypi.org/p/{0}', matrix.package) || format('https://test.pypi.org/p/{0}', matrix.package) }}
+    permissions:
+      id-token: write
+    steps:
+      - id: download-distribution
+        name: Download the built distribution.
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ matrix.package }}-build-artifacts
+          path: sub-packages/${{ matrix.package }}/dist
+      - id: publish-to-testpypi
+        name: Publish distribution 📦 to Test PyPI for PR.
+        if: ${{ github.event.inputs.pypi == 'false' }}
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          verbose: true
+          packages-dir: sub-packages/${{ matrix.package }}/dist
+          repository-url: https://test.pypi.org/legacy/
+          skip-existing: ${{ github.event.inputs.version_overwrite }}
+      - id: publish-to-pypi
+        name: Publish distribution 📦 to PyPI for Workflow Dispatch.
+        # To require testing before publishing to PyPI, add: ... && needs.install-and-test.result == 'success'
+        # If testing is run but fails, the workflow will fail and not publish to PyPI (or Test PyPI).
+        # We strongly recommend testing when publishing to production PyPI.
+        if: ${{ github.event.inputs.pypi == 'true' }}
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          verbose: true
+          packages-dir: sub-packages/${{ matrix.package }}/dist
+          skip-existing: ${{ github.event.inputs.version_overwrite }}