feat(python): package CUDA as an optional extension

.github/workflows/ci.yml

@@ -154,9 +154,9 @@ jobs:
         run: |
           # Build docs for the whole workspace (aside from vortex-python that is having hard time with the python-docs syntax),
           # including all private docs.
-          cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python
+          cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python --exclude vortex-python-cuda
           # nextest doesn't support doc tests, so we run it here
-          cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --no-fail-fast
+          cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --exclude vortex-python-cuda --no-fail-fast
 
   build-rust:
     name: "Rust build (${{matrix.config.name}})"
@@ -187,7 +187,7 @@ jobs:
             target: wasm32-unknown-unknown
             env:
               rustflags: "RUSTFLAGS='-A warnings --cfg getrandom_backend=\"unsupported\"'"
-            args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant"
+            args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant"
     steps:
       - uses: runs-on/action@v2
         if: github.repository == 'vortex-data/vortex'
@@ -355,7 +355,7 @@ jobs:
             --exclude vortex-bench `
             --exclude vortex-python --exclude vortex-duckdb `
             --exclude vortex-fuzz --exclude vortex-cuda --exclude vortex-cuda-ffi `
-            --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda `
+            --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda `
             --exclude duckdb-bench `
             --exclude lance-bench --exclude datafusion-bench --exclude random-access-bench `
             --exclude compress-bench --exclude xtask --exclude vortex-datafusion `

.github/workflows/codspeed.yml

@@ -100,10 +100,6 @@ jobs:
             --bench fsst_cuda \
             --bench runend_cuda \
             --profile bench
-      - name: Verify CUDA kernels were built
-        run: test -f vortex-cuda/kernels/gen/release/dynamic_dispatch.ptx
-      - name: Stage CUDA kernels with Codspeed artifacts
-        run: cp -R vortex-cuda/kernels/gen/release target/codspeed/cuda-kernels
       - name: Upload benchmark executables
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7
         with:
@@ -156,7 +152,6 @@ jobs:
         uses: CodSpeedHQ/action@d872884a306dd4853acf0f584f4b706cf0cc72a2
         env:
           CARGO_MANIFEST_DIR: ${{ github.workspace }}/vortex-cuda
-          VORTEX_CUDA_KERNELS_DIR: ${{ github.workspace }}/target/codspeed/cuda-kernels
         with:
           run: cargo codspeed run $(printf -- '--bench %s ' ${{ matrix.benches }})
           token: ${{ secrets.CODSPEED_TOKEN }}

.github/workflows/cuda.yaml

@@ -38,7 +38,13 @@ jobs:
           filters: |
             cuda:
               - "vortex-cuda/**"
+              - "vortex-python/python/vortex/__init__.py"
+              - "vortex-python/pyproject.toml"
+              - "vortex-python/test/test_cuda.py"
+              - "vortex-python-cuda/**"
               - "vortex-test/**"
+              - "pyproject.toml"
+              - "uv.lock"
               - ".github/workflows/**"
 
   cuda-build-lint:
@@ -62,7 +68,7 @@ jobs:
           command: >-
             cargo build --profile ci --locked --all-features --all-targets
             -p vortex-cuda -p vortex-cuda-ffi -p vortex-cub -p vortex-nvcomp
-            -p gpu-scan-cli -p vortex-test-e2e-cuda
+            -p gpu-scan-cli -p vortex-test-e2e-cuda -p vortex-python-cuda
       - name: Clippy CUDA crates
         run: |
           cargo clippy --profile ci --locked --all-features --all-targets \
@@ -72,6 +78,7 @@ jobs:
             -p vortex-nvcomp \
             -p gpu-scan-cli \
             -p vortex-test-e2e-cuda \
+            -p vortex-python-cuda \
             -- -D warnings
 
   cuda-test:
@@ -95,6 +102,10 @@ jobs:
       - uses: ./.github/actions/setup-rust
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Install uv
+        uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63  # 0.18.6
+        with:
+          sync: false
       - name: Install nextest
         uses: taiki-e/install-action@9e1e5806d4a4822de933115878265be9aaa786d9  # v2
         with:
@@ -112,10 +123,19 @@ jobs:
             -p vortex-cub \
             -p vortex-nvcomp \
             -p vortex-test-e2e-cuda \
+            -p vortex-python-cuda \
             --all-features \
             --no-fail-fast \
             --target x86_64-unknown-linux-gnu \
             --verbose
+      - name: Python CUDA extension checks
+        env:
+          MATURIN_PEP517_ARGS: "--profile ci"
+        run: |
+          # --all-packages installs the shared dev tooling (basedpyright) from the
+          # root `dev` group; --extra cuda adds the vortex-data-cuda extension.
+          uv run --all-packages --extra cuda basedpyright vortex-python vortex-python-cuda
+          uv run --all-packages --extra cuda pytest --benchmark-disable vortex-python/test/test_cuda.py vortex-python-cuda/test
 
   cuda-test-sanitizer:
     needs: [changes]

.github/workflows/musl.yml

@@ -66,7 +66,7 @@ jobs:
         run: |
           cargo nextest run --cargo-profile ci --locked --workspace --no-fail-fast \
             --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp \
-            --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda \
+            --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda \
             --exclude vortex-duckdb --exclude duckdb-bench --exclude vortex-sqllogictest \
             --exclude vortex-bench --exclude lance-bench --exclude datafusion-bench --exclude vortex-datafusion \
             --exclude compress-bench --exclude random-access-bench --exclude vortex-bench-server

.github/workflows/package.yml

@@ -50,6 +50,34 @@ jobs:
         run: |
           cargo set-version --workspace ${{ inputs.version }}
 
+      # cargo set-version bumps Cargo.toml crate versions but not the ==0.1.0 pins
+      # between these two wheels in their pyproject.toml. The base and CUDA extension
+      # wheels must require each other at the exact released version (they exchange
+      # Vortex IPC buffers), so rewrite those pins here.
+      - name: Set Python extension version pins
+        env:
+          VERSION: ${{ inputs.version }}
+        run: |
+          python3 - <<'PY'
+          import os
+          import re
+          from pathlib import Path
+
+          version = os.environ["VERSION"]
+          replacements = {
+              Path("vortex-python/pyproject.toml"): "vortex-data-cuda",
+              Path("vortex-python-cuda/pyproject.toml"): "vortex-data",
+          }
+          for path, package in replacements.items():
+              text = path.read_text()
+              text, count = re.subn(
+                  f"{re.escape(package)}==[^\"\\]]+", f"{package}=={version}", text
+              )
+              if count != 1:
+                  raise RuntimeError(f"Expected one {package} version pin in {path}, found {count}")
+              path.write_text(text)
+          PY
+
       - name: Set up Python
         uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6
         # Latest macOS doesn't allow maturin to install stuff into the global Python interpreter

.github/workflows/rust-instrumented.yml

@@ -88,7 +88,7 @@ jobs:
       - name: Rust Tests
         if: ${{ matrix.suite == 'tests' }}
         run: |
-          cargo nextest run --locked --workspace --all-features --no-fail-fast
+          cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-python-cuda
       - name: Generate coverage report
         run: |
           RUST_SYSROOT="$(rustc --print sysroot)"

Cargo.toml

@@ -37,6 +37,7 @@ members = [
     "fuzz",
     "vortex-jni",
     "vortex-python",
+    "vortex-python-cuda",
     "vortex-tui",
     "vortex-web/crate",
     "vortex-sqllogictest",

docs/api/python/runtime.rst

@@ -9,13 +9,16 @@ otherwise to the number of available CPU cores minus one. Use
 .. autosummary::
    :nosignatures:
 
+   ~vortex.cuda_extension_installed
    ~vortex.set_worker_threads
    ~vortex.worker_threads
 
 .. raw:: html
 
    <hr>
 
+.. autofunction:: vortex.cuda_extension_installed
+
 .. autofunction:: vortex.set_worker_threads
 
 .. autofunction:: vortex.worker_threads

pyproject.toml

@@ -37,6 +37,11 @@ members = ["bench-orchestrator", "vortex-python", "docs"]
 [tool.uv.sources]
 bench-orchestrator = { workspace = true }
 vortex-data = { workspace = true }
+# Resolve vortex-data-cuda from its local path so the optional `cuda` extra is
+# lockable before the package is published to PyPI. A path source (not
+# `workspace = true`) keeps the CUDA extension out of the default `--all-packages`
+# set, so it is only built when the `cuda` extra is requested.
+vortex-data-cuda = { path = "vortex-python-cuda" }
 docs = { workspace = true }
 
 [tool.ruff]

vortex-cuda/build.rs

@@ -8,6 +8,7 @@
 use std::env;
 use std::fs::File;
 use std::io;
+use std::io::Write;
 use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
@@ -32,15 +33,6 @@ fn main() {
 
     std::fs::create_dir_all(&kernels_gen).expect("Failed to create kernels/gen directory");
 
-    // Always emit the kernels output directory path as a compile-time env var so any binary
-    // linking against vortex-cuda can find the PTX files. This must be set regardless
-    // of CUDA availability since the code using env!() is always compiled.
-    // At runtime, VORTEX_CUDA_KERNELS_DIR can be set to override this path.
-    println!(
-        "cargo:rustc-env=VORTEX_CUDA_KERNELS_DIR={}",
-        kernels_gen.display()
-    );
-
     println!("cargo:rerun-if-env-changed=PROFILE");
 
     // Regenerate bit_unpack kernels only when the generator changes
@@ -60,6 +52,8 @@ fn main() {
     generate_dynamic_dispatch_bindings(&kernels_src, &out_dir);
     generate_patches_bindings(&kernels_src, &out_dir);
 
+    generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source");
+
     if !is_cuda_available() {
         return;
     }
@@ -97,6 +91,57 @@ fn main() {
             }
         }
     }
+
+    // Refresh embedded_ptx.rs after nvcc has compiled PTX so the binary embeds the latest kernels.
+    generate_embedded_ptx(&out_dir, &kernels_gen).expect("Failed to generate embedded PTX source");
+}
+
+/// Generates `embedded_ptx.rs`, a module-name lookup table of `include_str!`-embedded PTX.
+///
+/// Called before the `nvcc` check to create an empty table for no-toolkit builds, then again after
+/// PTX compilation so runtime loading never depends on build-machine filesystem paths.
+fn generate_embedded_ptx(out_dir: &Path, kernels_gen: &Path) -> io::Result<()> {
+    let mut ptx_files = Vec::new();
+    if let Ok(entries) = std::fs::read_dir(kernels_gen) {
+        for entry in entries.flatten() {
+            let path = entry.path();
+            if path.extension().and_then(|e| e.to_str()) == Some("ptx") {
+                ptx_files.push(path);
+            }
+        }
+    }
+    ptx_files.sort();
+
+    let mut file = File::create(out_dir.join("embedded_ptx.rs"))?;
+    if ptx_files.is_empty() {
+        writeln!(
+            file,
+            "pub(crate) fn embedded_ptx(_module_name: &str) -> Option<&'static str> {{"
+        )?;
+        writeln!(file, "    None")?;
+        writeln!(file, "}}")?;
+        return Ok(());
+    }
+
+    writeln!(
+        file,
+        "pub(crate) fn embedded_ptx(module_name: &str) -> Option<&'static str> {{"
+    )?;
+    writeln!(file, "    match module_name {{")?;
+    for path in ptx_files {
+        let Some(module_name) = path.file_stem().and_then(|stem| stem.to_str()) else {
+            continue;
+        };
+        writeln!(
+            file,
+            "        {module_name:?} => Some(include_str!({:?})),",
+            path.to_string_lossy()
+        )?;
+    }
+    writeln!(file, "        _ => None,")?;
+    writeln!(file, "    }}")?;
+    writeln!(file, "}}")?;
+    Ok(())
 }
 
 fn generate_unpack<T: FastLanes>(output_dir: &Path, thread_count: usize) -> io::Result<PathBuf> {