chore: upgrade vLLM to 0.20.0, cuda to 12.9

.agents/skills/git-worktrees/SKILL.md

@@ -29,10 +29,10 @@ uv sync --frozen
 
 This creates a local `.venv` in the worktree. With uv's cache the install takes ~2-3 seconds on a warm cache.
 
-If you need different extras (e.g. `cu128` vs `cpu`), pass them explicitly:
+If you need different extras (e.g. `cu129` vs `cpu`), pass them explicitly:
 
 ```bash
-uv sync --frozen --extra cu128 --extra engine --group dev
+uv sync --frozen --extra cu129 --extra engine --group dev
 ```
 
 Never run bare `uv sync` without `--frozen` -- it re-locks `uv.lock` and creates dirty state.

.agents/skills/uv-build/SKILL.md

@@ -19,7 +19,7 @@ make bootstrap-tools && make bootstrap-nss cpu
 # Pick a variant:
 make bootstrap-nss dev       # dev tools only (no engine/torch)
 make bootstrap-nss cpu       # + engine + CPU PyTorch
-make bootstrap-nss cu128     # + engine + CUDA 12.8 PyTorch
+make bootstrap-nss cu129     # + engine + CUDA 12.9 PyTorch
 make bootstrap-nss engine    # + engine (no torch)
 ```
 
@@ -30,11 +30,11 @@ Under the hood: `uv sync --frozen --extra <extra> [--extra engine] --group dev`
 | Extra | What it installs |
 |-------|------------------|
 | `cpu` | PyTorch CPU, faiss-cpu, flashinfer (Linux only) |
-| `cu128` | PyTorch+CUDA 12.8, faiss-gpu, flashinfer-jit-cache |
+| `cu129` | PyTorch+CUDA 12.9, faiss-gpu, flashinfer-jit-cache |
 | `engine` | ML pipeline deps (outlines, wandb, tiktoken, etc.) -- no torch |
 | `microservices` | `nemo-microservices` from local path |
 
-`cpu` and `cu128` conflict -- you must pick one, never both. Enforced in `[tool.uv] conflicts`.
+`cpu` and `cu129` conflict -- you must pick one, never both. Enforced in `[tool.uv] conflicts`.
 
 ## Index Management
 
@@ -43,9 +43,9 @@ PyTorch wheels come from dedicated indexes, not PyPI:
 | Index | URL | Used for |
 |-------|-----|----------|
 | `pytorch-cpu` | `download.pytorch.org/whl/cpu` | torch, torchvision (CPU, Linux) |
-| `pytorch-cu128` | `download.pytorch.org/whl/cu128` | torch, torchvision, triton, xformers (CUDA) |
+| `pytorch-cu129` | `download.pytorch.org/whl/cu129` | torch, torchvision, triton (CUDA) |
 | `nv-shared-pypi-local` | NVIDIA Artifactory | Internal NVIDIA packages |
-| `flashinfer-jit-cache` | `flashinfer.ai/whl/cu128` | FlashInfer JIT cache |
+| `flashinfer-jit-cache` | `flashinfer.ai/whl/cu129` | FlashInfer JIT cache |
 | `nvidia-pypi-public` | `pypi.nvidia.com` | Public NVIDIA packages |
 
 All indexes are `explicit = true` (only used when a package is mapped to them in `[tool.uv.sources]`).

.claude/commands/bootstrap.md

@@ -13,9 +13,9 @@ Set up the development environment from scratch.
 2. Install Python dependencies (choose one):
    ```bash
    make bootstrap-nss cpu    # CPU-only (macOS or Linux without GPU)
-   make bootstrap-nss cuda   # CUDA 12.8 (Linux with NVIDIA GPU)
+   make bootstrap-nss cuda   # CUDA 12.9 (Linux with NVIDIA GPU)
    make bootstrap-nss engine # Engine dependencies only (no torch)
    make bootstrap-nss dev    # Minimal dev dependencies only
    ```
 
-Note: `cuda` is an alias for `cu128`. Both are equivalent.
+Note: `cuda` is an alias for `cu129`. Both are equivalent.

.coderabbit.yaml

@@ -356,4 +356,4 @@ knowledge_base:
       - ".cursor/rules/*.mdc"
 
 chat:
-  auto_reply: true
+  auto_reply: true

.cursor/setup-worktree.sh

@@ -16,10 +16,10 @@ fi
 
 # Bare --frozen installs the base environment. For GPU dev work (ty, import
 # checks, GPU tests) run the full command manually after setup:
-#   uv sync --frozen --extra cu128 --extra engine --group dev
+#   uv sync --frozen --extra cu129 --extra engine --group dev
 uv sync --frozen
 echo "Venv ready: $(pwd)/.venv"
-echo "Note: for GPU extras run: uv sync --frozen --extra cu128 --extra engine --group dev"
+echo "Note: for GPU extras run: uv sync --frozen --extra cu129 --extra engine --group dev"
 
 for _envfile in .env .env.local mise.local.toml .local.envrc; do
     if [ -f "$ROOT_WORKTREE_PATH/$_envfile" ]; then

.github/actions/setup-gpu-test-env/action.yml

@@ -27,7 +27,7 @@ inputs:
   cuda-extra:
     description: "CUDA dependency extra to bootstrap"
     required: false
-    default: "cu128"
+    default: "cu129"
 
 runs:
   using: "composite"

AGENTS.md

@@ -24,7 +24,7 @@ Common commands: `make test` (unit tests), `make format` (auto-fix formatting +
 The canonical `uv sync` command for a full GPU/dev environment is:
 
 ```bash
-uv sync --frozen --extra cu128 --extra engine --group dev
+uv sync --frozen --extra cu129 --extra engine --group dev
 ```
 
 Bare `uv sync --frozen` (without extras) installs an incomplete environment -- `ty`, import checks, and GPU tests will fail.

CONTRIBUTING.md

@@ -59,7 +59,7 @@ Please read our [Code of Conduct](CODE_OF_CONDUCT.md) before contributing.
 
    # Install Python dependencies (choose one)
    make bootstrap-nss cpu    # CPU-only (macOS or Linux without GPU)
-   make bootstrap-nss cuda   # CUDA 12.8 (Linux with NVIDIA GPU)
+   make bootstrap-nss cuda   # CUDA 12.9 (Linux with NVIDIA GPU)
    make bootstrap-nss engine # Engine dependencies only
    make bootstrap-nss dev    # Minimal dev dependencies only
   ```

Makefile

@@ -13,7 +13,7 @@ NSS_ROOT_PATH := $(shell pwd)
 # Normalize architecture names
 ifeq ($(ARCH),x86_64)
 	ARCH := amd64
-	PYTORCH_DEPS := cu128
+	PYTORCH_DEPS := cu129
 	export BUILD_ARCH ?= linux/amd64
 endif
 ifeq ($(ARCH),aarch64)
@@ -82,12 +82,12 @@ verify-python-version: ## Verify Python version and install if necessary
 	uv venv --seed --allow-existing --python 3.11
 
 .PHONY: bootstrap-python
-bootstrap-python: .venv ## Bootstrap Python dependencies. Set PYTORCH_DEPS to 'cpu' or 'cu128'. Here mostly for legacy usage.
+bootstrap-python: .venv ## Bootstrap Python dependencies. Set PYTORCH_DEPS to 'cpu' or 'cu129'. Here mostly for legacy usage.
 	uv sync --frozen --extra ${PYTORCH_DEPS} --extra engine --group dev
 
 # Dynamic targets for bootstrap-nss
 # Usage: make bootstrap-nss {dev,engine,cpu,cuda}
-BOOTSTRAP_EXTRAS := dev engine cpu cuda cu128
+BOOTSTRAP_EXTRAS := dev engine cpu cuda cu129
 $(BOOTSTRAP_EXTRAS):
 	@:
 
@@ -97,9 +97,9 @@ bootstrap-nss: .venv ## Bootstrap Python dependencies. Usage: make bootstrap-nss
 	@echo "~~~~~~"
 	@echo "attempting to install nss package with primary extra: $(EXTRA)"
 	@if [ "$(EXTRA)" = "cuda" ]; then \
-		uv sync --frozen --extra cu128 --extra engine --group dev; \
-	elif [ "$(EXTRA)" = "cu128" ]; then \
-		uv sync --frozen --extra cu128 --extra engine --group dev; \
+		uv sync --frozen --extra cu129 --extra engine --group dev; \
+	elif [ "$(EXTRA)" = "cu129" ]; then \
+		uv sync --frozen --extra cu129 --extra engine --group dev; \
 	elif [ "$(EXTRA)" = "cpu" ]; then \
 		uv sync --frozen --extra cpu --extra engine --group dev; \
 	elif [ "$(EXTRA)" = "engine" ]; then \
@@ -489,7 +489,7 @@ NSS_DATASETS := clinc_oos dow_jones_index
 
 define nss_combo_test
 test-nss-$(1)-$(2)-ci: ## Run pytest test for $(shell echo $(1) | tr '_' '-') config with $(shell echo $(2) | tr '_' '-') dataset
-	$(MAKE) bootstrap-nss cu128
+	$(MAKE) bootstrap-nss cu129
 	$(PYTEST_NO_XDIST_CMD) -vv $(PYTEST_CI_OPTS) $(NSS_ROOT_PATH)/tests/e2e/test_dataset_config.py -k "test_$(2)_dataset[$(subst _,-,$(1))]"
 endef
 

README.md

@@ -18,15 +18,17 @@ Read detailed usage below, or jump to the documentation with [Getting Started](h
 
 ```bash
 # With uv (recommended):
-uv pip install "nemo-safe-synthesizer[cu128,engine]" \
-  --index https://flashinfer.ai/whl/cu128 \
-  --index https://download.pytorch.org/whl/cu128 \
+uv pip install "nemo-safe-synthesizer[cu129,engine]" \
+  --index https://flashinfer.ai/whl/cu129 \
+  --index https://download.pytorch.org/whl/cu129 \
+  --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 \
   --index-strategy unsafe-best-match
 
 # With pip:
-pip install "nemo-safe-synthesizer[cu128,engine]" \
-  --extra-index-url https://download.pytorch.org/whl/cu128 \
-  --extra-index-url https://flashinfer.ai/whl/cu128
+pip install "nemo-safe-synthesizer[cu129,engine]" \
+  --extra-index-url https://download.pytorch.org/whl/cu129 \
+  --extra-index-url https://flashinfer.ai/whl/cu129 \
+  --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129
 ```
 
 Or install from source:
@@ -210,7 +212,7 @@ Controls the HuggingFace attention backend used during model loading for trainin
 ```yaml
 # config.yaml
 training:
-  attn_implementation: "kernels-community/vllm-flash-attn3"
+  attn_implementation: "sdpa"
 ```
 
 ```bash
@@ -220,13 +222,14 @@ safe-synthesizer run --training__attn_implementation sdpa --data-source my_data.
 
 | Value | Description | Requires |
 |-------|-------------|----------|
-| `kernels-community/vllm-flash-attn3` | Flash Attention 3 via HuggingFace Kernels Hub (default) | `kernels` pip package |
+| `sdpa` | PyTorch scaled dot product attention (default) | None (built-in) |
+| `eager` | Standard PyTorch attention | None (built-in) |
 | `kernels-community/flash-attn2` | Flash Attention 2 via HuggingFace Kernels Hub | `kernels` pip package |
+| `kernels-community/vllm-flash-attn3` | Flash Attention 3 via HuggingFace Kernels Hub | `kernels` pip package and compatible prebuilt kernel |
 | `flash_attention_2` | Flash Attention 2 (traditional) | `flash-attn` pip package |
-| `sdpa` | PyTorch scaled dot product attention | None (built-in) |
-| `eager` | Standard PyTorch attention | None (built-in) |
+| `flash_attention_3` | Flash Attention 3 (traditional) | `flash-attn-3` support |
 
-If the default `kernels-community/vllm-flash-attn3` is configured but the `kernels` package is not installed, the backend automatically falls back to `sdpa`.
+If a `kernels-community/...` value is configured but the `kernels` package is not installed, the backend automatically falls back to `sdpa`.
 
 ### Generation (`attention_backend`)
 

containers/Dockerfile.cuda

@@ -5,7 +5,7 @@
 #
 # Stages:
 #   tools   -- install mise and all dev tools (uv, ruff, etc.)
-#   deps    -- install Python, uv, and all cu128+engine dependencies
+#   deps    -- install Python, uv, and all cu129+engine dependencies
 #   runtime -- minimal CLI wrapper (ENTRYPOINT entrypoint.sh -> safe-synthesizer)
 #   dev     -- extends runtime with dev tools, tests, and interactive shell
 #
@@ -58,7 +58,7 @@ ARG TARGETARCH
 # ---------------------------------------------------------------------------
 # Build arguments
 # ---------------------------------------------------------------------------
-ARG CUDA_VERSION=12.8.1
+ARG CUDA_VERSION=12.9.1
 ARG UBUNTU_VERSION=22.04
 # Use "runtime" when all deps ship pre-built wheels (current state).
 # Switch to "devel" if a future dependency requires CUDA compilation.
@@ -142,14 +142,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # not when source code changes.
 COPY pyproject.toml uv.lock ./
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --no-install-project --extra cu128 --extra engine --no-group dev
+    uv sync --no-install-project --extra cu129 --extra engine --no-group dev
 
 # -- Project layer: install the project (non-editable so the venv is
 # self-contained and source code is not needed in the runtime image).
 COPY README.md ./
 COPY src/ src/
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --no-editable --extra cu128 --extra engine --no-group dev
+    uv sync --no-editable --extra cu129 --extra engine --no-group dev
 
 # ---------------------------------------------------------------------------
 # Stage 3: runtime -- minimal image wrapping the safe-synthesizer CLI
@@ -233,7 +233,7 @@ COPY . .
 
 # Install dev dependencies into the existing venv.
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --extra cu128 --extra engine --group dev
+    uv sync --extra cu129 --extra engine --group dev
 
 RUN git config --global --add safe.directory /workspace
 

docs/developer-guide/docker.md

@@ -24,7 +24,7 @@ flowchart TD
 
     subgraph stages [Build Stages]
         tools["tools\nInstalls mise + all dev tools\n(.mise.toml is single source of truth)"]
-        deps["deps\nInstalls Python 3.11 via uv\nuv sync cu128+engine"]
+        deps["deps\nInstalls Python 3.11 via uv\nuv sync cu129+engine"]
         runtime["runtime\nCopies venv + Python\nNon-root appuser\ntini + entrypoint.sh"]
         dev["dev\nExtends runtime\nCopies mise tree from tools\nRoot user"]
     end
@@ -44,7 +44,7 @@ flowchart TD
   Uses the [mise Docker cookbook](https://mise.jdx.dev/mise-cookbook/docker.html)
   pattern with `MISE_DATA_DIR=/mise` for stable, copyable paths.
 - deps: copies the uv binary from `tools`, then installs Python and all
-  cu128+engine dependencies. Uses `--mount=type=cache` to avoid
+  cu129+engine dependencies. Uses `--mount=type=cache` to avoid
   re-downloading ~10 GB of PyTorch/CUDA wheels.
 - runtime: copies the venv and uv-managed Python into a fresh CUDA runtime
   base. Runs as non-root `appuser` (uid 1000). GPU access is declared via
@@ -262,8 +262,8 @@ To reduce size:
 
 | Aspect | `Dockerfile.cuda` | `Dockerfile.test_ci` |
 |--------|-------------------|----------------------|
-| Base | `nvidia/cuda:12.8.1-runtime-ubuntu22.04` | `python:3.11-slim` |
-| Extras | `cu128` + `engine` | `cpu` + `engine` |
+| Base | `nvidia/cuda:12.9.1-runtime-ubuntu22.04` | `python:3.11-slim` |
+| Extras | `cu129` + `engine` | `cpu` + `engine` |
 | GPU | Required | Not needed |
 | Stages | `tools` / `deps` / `runtime` / `dev` | `setup` / `install-deps` |
 | Use case | Training, generation, evaluation | CPU-only unit tests and CI checks |

docs/index.md

@@ -18,7 +18,7 @@ NeMo Safe Synthesizer creates private, safe versions of sensitive tabular datase
 - Flexible interfaces -- CLI for scripting, Python SDK for programmatic workflows, YAML configuration
 
 !!! info "System Requirements"
-    NeMo Safe Synthesizer requires a Linux machine with an NVIDIA GPU (A100 80GB+ recommended) and CUDA 12.8+ to run the training and generation pipeline. macOS, Windows, and Apple Silicon are not supported for pipeline execution. A CPU-only install is available for development and configuration validation -- see [Getting Started](user-guide/getting-started.md#install-the-package).
+    NeMo Safe Synthesizer requires a Linux machine with an NVIDIA GPU (A100 80GB+ recommended) and CUDA 12.9+ to run the training and generation pipeline. macOS, Windows, and Apple Silicon are not supported for pipeline execution. A CPU-only install is available for development and configuration validation -- see [Getting Started](user-guide/getting-started.md#install-the-package).
 
 ## Next Steps
 

docs/tutorials/differential-privacy.ipynb

@@ -14,7 +14,7 @@
         "\n",
         "### 🖥️ Prerequisites\n",
         "\n",
-        "This notebook requires a Linux machine with an NVIDIA GPU (H100 recommended, A100 minimum) and CUDA 12.8+. It will not run on macOS, Windows, or Apple Silicon."
+        "This notebook requires a Linux machine with an NVIDIA GPU (H100 recommended, A100 minimum) and CUDA 12.9+. It will not run on macOS, Windows, or Apple Silicon."
       ]
     },
     {
@@ -24,7 +24,7 @@
       "source": [
         "### ⚡ Install Safe Synthesizer\n",
         "\n",
-        "Run the cell below to install NeMo Safe Synthesizer (engine and CUDA 12.8) and kagglehub for the example dataset."
+        "Run the cell below to install NeMo Safe Synthesizer (engine and CUDA 12.9) and kagglehub for the example dataset."
       ]
     },
     {
@@ -39,10 +39,10 @@
         "# SPDX-License-Identifier: Apache-2.0\n",
         "\n",
         "if command -v uv > /dev/null 2>&1; then\n",
-        "    uv pip install \"nemo-safe-synthesizer[engine,cu128]\" --index https://flashinfer.ai/whl/cu128 --index https://download.pytorch.org/whl/cu128 --index-strategy unsafe-best-match\n",
+        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 --index-strategy unsafe-best-match\n",
         "    uv pip install kagglehub\n",
         "else\n",
-        "    pip install \"nemo-safe-synthesizer[engine,cu128]\" --extra-index-url https://flashinfer.ai/whl/cu128 --extra-index-url https://download.pytorch.org/whl/cu128\n",
+        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129\n",
         "    pip install kagglehub\n",
         "fi\n"
       ]
@@ -241,7 +241,7 @@
   ],
   "metadata": {
     "kernelspec": {
-      "display_name": ".venv",
+      "display_name": "Python 3",
       "language": "python",
       "name": "python3"
     },

docs/tutorials/safe-synthesizer-101.ipynb

@@ -16,7 +16,7 @@
         "\n",
         "### 🖥️ Prerequisites\n",
         "\n",
-        "This notebook requires a Linux machine with an NVIDIA GPU (H100 recommended, A100 minimum) and CUDA 12.8+. It will not run on macOS, Windows, or Apple Silicon."
+        "This notebook requires a Linux machine with an NVIDIA GPU (H100 recommended, A100 minimum) and CUDA 12.9+. It will not run on macOS, Windows, or Apple Silicon."
       ]
     },
     {
@@ -26,7 +26,7 @@
       "source": [
         "### ⚡ Install Safe Synthesizer\n",
         "\n",
-        "Run the cell below to install NeMo Safe Synthesizer (engine and CUDA 12.8) and the `datasets` library for the sample dataset."
+        "Run the cell below to install NeMo Safe Synthesizer (engine and CUDA 12.9) and the `datasets` library for the sample dataset."
       ]
     },
     {
@@ -45,10 +45,10 @@
         "# SPDX-License-Identifier: Apache-2.0\n",
         "\n",
         "if command -v uv > /dev/null 2>&1; then\n",
-        "    uv pip install \"nemo-safe-synthesizer[engine,cu128]\" --index https://flashinfer.ai/whl/cu128 --index https://download.pytorch.org/whl/cu128 --index-strategy unsafe-best-match\n",
+        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 --index-strategy unsafe-best-match\n",
         "    uv pip install datasets\n",
         "else\n",
-        "    pip install \"nemo-safe-synthesizer[engine,cu128]\" --extra-index-url https://flashinfer.ai/whl/cu128 --extra-index-url https://download.pytorch.org/whl/cu128\n",
+        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129\n",
         "    pip install datasets\n",
         "fi\n"
       ]

docs/user-guide/configuration.md

@@ -76,7 +76,7 @@ for the full field list.
 | `training.pretrained_model` | `"HuggingFaceTB/SmolLM3-3B"` | HuggingFace model ID or local path | See supported families below; `TinyLlama/TinyLlama-1.1B-Chat-v1.0` for fast CPU/low-VRAM iteration |
 | `training.quantize_model` | `false` | Enable quantization to reduce VRAM usage | Enable if VRAM is limited; 8-bit has lower quality impact than 4-bit |
 | `training.quantization_bits` | `8` | Bit width (4 or 8) when `training.quantize_model` is `true` | Prefer 8 over 4 for quality |
-| `training.attn_implementation` | `"kernels-community/vllm-flash-attn3"` | Attention backend for model loading | Leave at default |
+| `training.attn_implementation` | `"sdpa"` | Attention backend for model loading | Leave at default |
 | `training.rope_scaling_factor` | `"auto"` | Scale the base model's context window via RoPE (`"auto"` or int) | Leave at `"auto"` |
 | `training.validation_ratio` | `0.0` | Fraction of training data held out for validation loss monitoring | Leave at 0.0 unless you specifically want to monitor validation loss |
 | `training.max_vram_fraction` | `0.8` | Fraction of total GPU VRAM to allocate for training. Must be in [0, 1] | Lower if other GPU consumers are active on the same device |

docs/user-guide/docker.md

@@ -13,7 +13,7 @@ for training, generation, and evaluation.
 
 - Docker 20.10+ (BuildKit enabled by default in 23.0+)
 - [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) installed and configured
-- NVIDIA driver compatible with CUDA 12.8
+- NVIDIA driver compatible with CUDA 12.9
 - NVIDIA GPU (A100 or better recommended)
 
 Verify GPU access works: