docs: update wheel install instructions

Signed-off-by: mkornfield <mkornfield@nvidia.com>
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>

Author: mckornfield

docs/tutorials/safe-synthesizer-101.ipynb

@@ -44,12 +44,15 @@
       "cell_type": "code",
       "execution_count": null,
       "id": "bb7b0bdd",
-      "metadata": {},
+      "metadata": {
+        "vscode": {
+          "languageId": "shellscript"
+        }
+      },
       "outputs": [],
       "source": [
         "%%capture\n",
-        "#TODO: Replace the internal NVIDIA Artifactory endpoint with the public PyPI release when available.\n",
-        "!uv pip install nemo-safe-synthesizer[engine,cu128] --extra-index-url \"https://urm.nvidia.com/artifactory/api/pypi/nv-shared-pypi-local/simple\"\n",
+        "!uv pip install nemo-safe-synthesizer[engine,cu128] --extra-index-url https://download.pytorch.org/whl/cu128 --extra-index-url https://flashinfer.ai/whl/cu128\n",
         "!uv pip install datasets\n"
       ]
     },

docs/user-guide/getting-started.md

@@ -20,17 +20,64 @@ does at each stage.
 
 ### Install the Package
 
+The CUDA and CPU extras depend on packages (PyTorch, FlashInfer) hosted on
+indexes outside PyPI. You must pass the extra index URLs shown below.
+
 === "CUDA 12.8 (Linux with NVIDIA GPU)"
 
-    ```bash
-    pip install "nemo-safe-synthesizer[cu128,engine]"
-    ```
+    === "pip"
+
+        ```bash
+        pip install "nemo-safe-synthesizer[cu128,engine]" \
+          --extra-index-url https://download.pytorch.org/whl/cu128 \
+          --extra-index-url https://flashinfer.ai/whl/cu128
+        ```
+
+    === "uv"
+
+        ```bash
+        uv pip install "nemo-safe-synthesizer[cu128,engine]" \
+          --index https://flashinfer.ai/whl/cu128 \
+          --index https://download.pytorch.org/whl/cu128
+        ```
+
+        !!! info "Why `--index` instead of `--extra-index-url`"
+            `--index` gives the specified indexes higher priority than PyPI, so
+            uv's default `first-match` strategy finds the CUDA-specific wheels
+            on flashinfer.ai and download.pytorch.org before hitting PyPI.
 
 === "CPU (macOS / Linux without GPU)"
 
-    ```bash
-    pip install "nemo-safe-synthesizer[cpu,engine]"
-    ```
+    On **macOS**, PyTorch ships standard wheels on PyPI, so no extra indexes are needed.
+
+    On **Linux**, the CPU-only PyTorch wheels (`+cpu` local version) are hosted
+    on a separate PyTorch index.
+
+    === "pip (macOS)"
+
+        ```bash
+        pip install "nemo-safe-synthesizer[cpu,engine]"
+        ```
+
+    === "pip (Linux)"
+
+        ```bash
+        pip install "nemo-safe-synthesizer[cpu,engine]" \
+          --extra-index-url https://download.pytorch.org/whl/cpu
+        ```
+
+    === "uv (macOS)"
+
+        ```bash
+        uv pip install "nemo-safe-synthesizer[cpu,engine]"
+        ```
+
+    === "uv (Linux)"
+
+        ```bash
+        uv pip install "nemo-safe-synthesizer[cpu,engine]" \
+          --index https://download.pytorch.org/whl/cpu
+        ```
 
     !!! warning "Development use only"
         The CPU install does not support training or generation. Use it to
@@ -54,9 +101,19 @@ does at each stage.
 
 === "Bare package for config definitions"
 
-    ```bash
-    pip install "nemo-safe-synthesizer"
-    ```
+    The bare package has no PyTorch or FlashInfer dependencies.
+
+    === "pip"
+
+        ```bash
+        pip install "nemo-safe-synthesizer"
+        ```
+
+    === "uv"
+
+        ```bash
+        uv pip install "nemo-safe-synthesizer"
+        ```
 
     !!! note "Limited use"
         The bare package includes only the Pydantic configuration models -- no

docs/user-guide/troubleshooting.md

@@ -93,11 +93,9 @@ To diagnose:
 
 1. Verify NVIDIA drivers: `nvidia-smi`
 2. Verify PyTorch CUDA build: `python -c "import torch; print(torch.cuda.is_available())"`
-3. Ensure you installed the CUDA extras, not the CPU-only package:
-
-    ```bash
-    pip install "nemo-safe-synthesizer[cu128,engine]"
-    ```
+3. Ensure you installed the CUDA extras, not the CPU-only package.
+   See [Installation](getting-started.md#install-the-package) for the
+   full command with required index URLs.
 
 Switch to the HuggingFace backend for CPU-only environments (useful for
 development, not recommended for production training).

pyproject.toml

@@ -222,6 +222,12 @@ constraint-dependencies = ["torch==2.9.1", "regex==2025.07.34", "pandas<3"]
     { index = "pytorch-cpu", extra="cpu" },
     { index = "pytorch-cu128", extra="cu128"},
   ]
+  flashinfer-python = [
+    { index = "flashinfer-jit-cache", marker = "sys_platform=='linux'", extra="cu128"},
+  ]
+  flashinfer-cubin = [
+    { index = "flashinfer-jit-cache", marker = "sys_platform=='linux'", extra="cu128"},
+  ]
   flashinfer-jit-cache = [
     { index = "flashinfer-jit-cache", marker = "sys_platform=='linux'", extra="cu128"},
   ]