docs: update wheel install instructions

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

Author: mckornfield

docs/tutorials/safe-synthesizer-101.ipynb

@@ -35,15 +35,18 @@
       "cell_type": "code",
       "execution_count": null,
       "id": "bb7b0bdd",
-      "metadata": {},
+      "metadata": {
+        "vscode": {
+          "languageId": "shellscript"
+        }
+      },
       "outputs": [],
       "source": [
-        "%%capture\n",
         "# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
         "# SPDX-License-Identifier: Apache-2.0\n",
         "\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",
+        "%%capture\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 datasets\n"
       ]
     },
@@ -203,7 +206,7 @@
   ],
   "metadata": {
     "kernelspec": {
-      "display_name": ".venv",
+      "display_name": "Python 3",
       "language": "python",
       "name": "python3"
     },

docs/user-guide/getting-started.md

@@ -20,17 +20,68 @@ 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 \
+          --index-strategy unsafe-best-match
+        ```
+
+        !!! info "Why `--index-strategy unsafe-best-match`"
+            FlashInfer publishes wheels to flashinfer.ai, but `flashinfer-python`
+            also appears on the PyTorch index at older versions. uv's default
+            `first-match` strategy stops at the first index that contains a
+            package name, so it picks up the wrong version from the PyTorch
+            index and fails to resolve. `--index-strategy unsafe-best-match`
+            tells uv to consider all indexes and pick the best matching version.
 
 === "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 +105,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"},
   ]