GPU paillier client code added, execution context modified to include…

[RovayL]

… GPU code, additional tests for GPU clients have been added, deemed fully functional

---

Note

High Risk
High risk because it introduces new native CUDA extensions and changes cryptography client backend selection and packaging/release workflows, which can affect runtime correctness and distribution/installability across platforms.

Overview
Adds first-class GPU support for Paillier/Paillier-Lookup via in-tree CUDA extensions. Introduces build_gpu_binaries.sh plus new paillier_gpu_ext / paillier_lookup_gpu_ext packages (Makefiles + pybind11 .cu) and updates clients to load these extensions when available.

Changes backend selection from env var to API-level device=. PaillierClient, PaillierLookupClient, and ExecutionContext now take device: "auto"|"cpu"|"gpu" (default "auto") using a shared resolve_device helper, remove DEVICE env handling, normalize GPU cipher formats, and add encode_hamming_server helpers for cross-backend interop.

Updates CI/release to gate and ship GPU builds. CI adds label-gated integration and GPU jobs, builds .so artifacts on a free runner and runs GPU tests on a dedicated runner; release builds per-Python wheels, repairs to manylinux_2_31_x86_64, publishes wheels+sdist, and uses a Hatch build hook/artifacts config to include .so files and tag wheels as non-pure when binaries are present.

Docs/README/CHANGELOG are updated to document manual GPU builds, requirements, and the new device= API, and repo URLs are updated in pyproject.toml.

^{Reviewed by Cursor Bugbot for commit c8fc683. Bugbot is set up for automated code reviews on this repo. Configure here.}

[liwenXtrace]

Unify CPU/GPU Paillier clients

Goal

Delete paillier_gpu_client.py and paillier_lookup_gpu_client.py. paillier_client.py and paillier_lookup_client.py each import the .so directly and expose a unified API. This requires the two C++ pybind modules to expose the same method names/signatures the CPU classes already have.

C++ changes

paillier-GPU-client/paillier_GPU_client.cu — missing API, must add

The pybind module currently exposes only encrypt, encode_hamming_client, encode_hamming_server, decode_hamming_client, get_pk_hex, get_keys_hex. Add the same methods the lookup variant already has (see paillier_GPU_lookup_client.cu:2644-2648 for the pattern):

• std::string stringify_pk() → JSON {"g": ..., "n": ..., "n_squared": ...} (g = n+1)
• std::string stringify_sk() → JSON {"phi": ..., "inv": ...}
• std::string stringify_config() → JSON {"embed_len": ..., "key_len": ...}
• void load_stringified_keys(const std::string& pk_json, const std::string& sk_json)
• void load_config(const py::dict& config)
• pybind .def(...) entries for all five

Also recommended: unify encrypt / decode_hamming_client signatures with CPU. Right now GPU is batch-only with hex-string I/O, CPU is single-only with int I/O — that's why PaillierClient has a hex-conversion shim layer. Either add single-variant methods or standardize on ints.

Drop get_pk_hex / get_keys_hex once stringify_* replaces them, or keep as debug helpers.

paillier-GPU-lookup-client/paillier_GPU_lookup_client.cu — missing table API

Exposes stringify_pk, stringify_sk, stringify_config, load_stringified_keys, load_config — good. But the pybind bindings at paillier_GPU_lookup_client.cu:2612-2650 do not expose:

• dump_tables()
• dump_tables_bytes()
• load_config_bytes()
• the two-arg form load_config(config, tables=...)

PaillierLookupGPU in paillier_lookup_client.py and PaillierLookupClient.__getstate__ both call these methods — they raise AttributeError at runtime, so pickling a GPU lookup client is currently broken, and GPU precomputed-table caching is dead code. Add the missing .def(...) bindings (and extend load_config to accept an optional tables kwarg) so the CPU/GPU API matches.

Packaging changes

Rename the extension directories to valid Python package names so the .so can be imported normally:

• paillier-GPU-client/ → paillier_gpu_ext/ (add __init__.py)
• paillier-GPU-lookup-client/ → paillier_lookup_gpu_ext/ (add __init__.py)

Update the Makefile in each directory and build_gpu_binaries.sh to match the new paths.

Python changes

Delete

• src/xtrace_sdk/x_vec/crypto/paillier_gpu_client.py
• src/xtrace_sdk/x_vec/crypto/paillier_lookup_gpu_client.py
• src/xtrace_sdk/x_vec/crypto/_gpu_loader.py

The custom loader only exists because the old directory names had hyphens. With the rename above, a normal import works. Preserve the "please build the .so" UX at the import site:

try:
    from xtrace_sdk.x_vec.crypto.paillier_gpu_ext import PaillierGPUClient
except ImportError as e:
    raise ImportError(
        "GPU extension not built. Run `./build_gpu_binaries.sh` "
        "(requires Docker, NVIDIA driver >= 550)."
    ) from e

paillier_client.py / paillier_lookup_client.py

• Remove the PaillierGPU / PaillierLookupGPU wrapper classes entirely.

• In PaillierClient.__init__ (and lookup equivalent) when DEVICE=gpu: assign the pybind class directly as self.client:

  self.client = PaillierGPUClient(embed_len=..., key_len=..., skip_key_gen=...)

• Remove the isinstance(self.client, PaillierCPU) dispatch branches in encrypt_vec_one/batch and decode_hamming_client_one/batch once C++ signatures match CPU. stringify_* / load_* methods become trivial passthroughs.

• Once the GPU lookup extension exposes dump_tables / dump_tables_bytes / table-aware load_config, PaillierLookupClient.__getstate__ / __setstate__ can drop their CPU-vs-GPU forks and just call self.client.dump_tables_bytes() / self.client.load_config(config, tables=...) uniformly.

execution_context.py

• Drop the imports of PaillierGPUClient / PaillierLookupGPUClient (lines 17-18).
• Drop the paillier_gpu / paillier_lookup_gpu branches in create() (lines 135-138) — DEVICE=gpu + the base paillier / paillier_lookup type is enough.
• Remove _canonical_homomorphic_type (lines 52-57) — CPU and GPU both report PaillierClient / PaillierLookupClient now, so the canonicalization is a no-op.
• Check SUPPORTED_HOMOMORPHIC_CLIENTS in settings.py and any pickle __getstate__ / __setstate__ for references to the deleted GPU client class names.

Tests

Grep for PaillierGPUClient and PaillierLookupGPUClient — each test that constructs them directly needs to switch to PaillierClient / PaillierLookupClient with DEVICE=gpu. Add a pickle round-trip test for GPU lookup once the dump_tables bindings land (none exists today, which is why the broken delegation wasn't caught).

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c8fc683. Configure here.}

[cursor]

Unused CUDA kernel adds dead code complexity

Low Severity

The decode_hamming_kernel GPU kernel is defined but never invoked by any method in PaillierGPUClient. The decode_hamming_client method uses decrypt_only_kernel (via decrypt_chunks) followed by CPU-side bit counting instead. This dead kernel with its complex parameter list (including debug outputs, mask inputs) adds maintenance burden and confusion about the intended decryption path.

 

^{Reviewed by Cursor Bugbot for commit c8fc683. Configure here.}