CUDA backend

.claude/skills/testing-tsim-correctness/SKILL.md

@@ -46,9 +46,9 @@ handful of T gates; its memory grows exponentially.
   test is a broken test.
 - Stim merges consecutive identical instructions into one with several targets, so
   write circuits that exercise both single- and multi-target paths.
-- GPU tests: `test/conftest.py` keeps every test outside
-  `test/integration/test_metal_sampler.py` on the CPU. In the Metal tests compile the
-  GPU side with `device="metal"`, so a silent fallback cannot make the test pass on the
+- GPU tests: `test/conftest.py` keeps every test outside `test/integration/test_metal_sampler.py`
+  and `test_cuda_sampler.py` on the CPU. In the GPU tests compile the GPU side with
+  `device="metal"` or `device="cuda"`, so a silent fallback cannot make the test pass on the
   CPU, and the reference with `device="cpu"`; compare them with `assert_counts_match`
   (GPU samples never equal CPU samples bit for bit).
 - Keep circuits small (a few qubits, tens of thousands of shots) so the suite stays

AGENTS.md

@@ -68,7 +68,7 @@ more than a handful of T gates: their compile memory grows exponentially.
   a reference. The reference-sample flags change only that.
 - **Seeds:** with `active_statevector` on the CPU, a seed's samples do not depend on the
   thread count or `batch_size`. Do not break this; noise blocks and per-shot draws are seeded
-  by call and shot index (`shots.py`). On a GPU (`statevector/metal/`) they depend on
+  by call and shot index (`shots.py`). On a GPU (`statevector/metal/`, `statevector/cuda/`) they depend on
   `batch_size`, as documented, and differ from the CPU's.
 - **Direct outputs stay direct:** outputs that depend on the noise alone never go through a
   state vector. Most detectors of QEC circuits are direct, which is why the detector sampler is

CHANGELOG.md

@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- **NVIDIA GPU sampling.** With the new `cuda` extra
+  (`pip install "bloqade-tsim[cuda]"`, CUDA 13), Tsim samples on NVIDIA GPUs
+  through CuPy, as it does on Apple GPUs: every component must have at most 11
+  active qubits, `device="auto"` picks the GPU when it can run the circuit, and
+  `device="cuda"` raises an error explaining why it can't. On an RTX 5090,
+  end-to-end cultivation at d=5 samples about 1.6 billion shots per second,
+  postselected, in any output layout.
+- The `jax-cuda12` and `jax-cuda13` extras, for the stabilizer-rank strategies
+  on NVIDIA GPUs with JAX. The `cuda12` and `cuda13` extras of Tsim 1.0 install
+  the same packages.
+
 ## [1.1.0] - 2026-11-12
 
 Tsim 1.1 adds a GPU backend for Apple GPUs. For example, end-to-end

README.md

@@ -24,9 +24,10 @@ uv add bloqade-tsim
 ```
 
 The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra
-needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra:
+needed). To sample on an NVIDIA GPU, install the extra of your strategy:
 ```bash
-pip install "bloqade-tsim[cuda13]"   # or bloqade-tsim[cuda12]
+pip install "bloqade-tsim[cuda]"         # the default strategy, "active_statevector" (CUDA 13)
+pip install "bloqade-tsim[jax-cuda13]"   # the stabilizer-rank strategies, with JAX (or jax-cuda12)
 ```
 
 ### Using Tsim with coding agents
@@ -82,7 +83,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile
 
 | Strategy | Method | Runs on | Cost grows with | Use it for |
 |---|---|---|---|---|
-| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection |
+| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) or NVIDIA GPU (CUDA) | Number of active qubits (at most 30) | Most circuits, especially with postselection |
 | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput |
 
 ```python

docs/index.mdx

@@ -40,10 +40,10 @@ uv add bloqade-tsim
 ```
 
 The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra
-needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra:
+needed). To sample on an NVIDIA GPU with the default strategy, install the `cuda` extra:
 
 ```bash
-uv add "bloqade-tsim[cuda13]"
+uv add "bloqade-tsim[cuda]"
 ```
 
 See [Installation](/install) for more options.
@@ -55,7 +55,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile
 
 | Strategy | Method | Runs on | Cost grows with | Use it for |
 |---|---|---|---|---|
-| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection |
+| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) or NVIDIA GPU (CUDA) | Number of active qubits (at most 30) | Most circuits, especially with postselection |
 | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput |
 
 ```python

docs/install.mdx

@@ -13,55 +13,60 @@ We recommend using [`uv`](https://docs.astral.sh/uv/) for dependency management.
 uv add bloqade-tsim
 ```
 
-For GPU acceleration with CUDA:
-
-```bash
-# For CUDA 13
-uv add "bloqade-tsim[cuda13]"
-
-# For CUDA 12
-uv add "bloqade-tsim[cuda12]"
-```
-
 ## Using pip
 
 ```bash
 pip install bloqade-tsim
 ```
 
-For GPU acceleration with CUDA:
+## GPUs
+
+The default strategy (`"active_statevector"`) samples on an Apple GPU through Metal, or on an
+NVIDIA GPU through CUDA. On an Apple-silicon Mac, the plain install includes PyObjC's Metal
+bindings for it. For NVIDIA GPUs, install an extra; each sampling strategy has its own:
+
+| Extra | Accelerates | Needs |
+|---|---|---|
+| `cuda` | `"active_statevector"` (the default), with CuPy | CUDA 13 (driver 580 or later) |
+| `jax-cuda13` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 13 |
+| `jax-cuda12` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 12 |
 
 ```bash
-pip install "bloqade-tsim[cuda13]"
+uv add "bloqade-tsim[cuda]"
+# or, with both strategies on the GPU:
+uv add "bloqade-tsim[cuda,jax-cuda13]"
 ```
 
-## Apple GPUs
+The `cuda12` and `cuda13` extras of Tsim 1.0 still install the JAX extras. JAX starts only
+when a stabilizer-rank strategy samples, and then reserves 75% of the GPU's memory; to sample
+with both strategies in one process, set `XLA_PYTHON_CLIENT_PREALLOCATE=false` before starting
+Python.
 
-On an Apple-silicon Mac, the default strategy (`"active_statevector"`) samples on the GPU
-through Metal; the plain install includes PyObjC's Metal bindings for it. Where it samples is
-the `device` argument of `compile_sampler` and `compile_detector_sampler`:
+Where `"active_statevector"` samples is the `device` argument of `compile_sampler` and
+`compile_detector_sampler`:
 
 | `device` | Samples |
 | --- | --- |
-| `"auto"` (default) | On the GPU when it can run the circuit, else on the CPU, without a warning |
+| `"auto"` (default) | On a GPU when it can run the circuit, else on the CPU, without a warning |
 | `"cpu"` | On the CPU, in double precision |
-| `"metal"` | On the GPU, or raises an error that says why it cannot |
+| `"metal"` | On an Apple GPU, or raises an error that says why it cannot |
+| `"cuda"` | On an NVIDIA GPU, or raises an error that says why it cannot |
 
 ```python
-sampler = circuit.compile_detector_sampler(device="metal")
+sampler = circuit.compile_detector_sampler(device="cuda")
 ```
 
 The GPU runs circuits whose components have at most 11 active qubits. That includes Clifford
 circuits, whose detectors depend on the noise alone: the GPU samples the noise as well as the
 state vectors. Postselection runs on the GPU, so heavily postselected circuits gain the most.
-A sampler's repr says where it samples (`samples on metal` or `samples on cpu`), and so does
-`sampler.stats.device`.
+A sampler's repr says where it samples (`samples on metal`, `samples on cuda` or
+`samples on cpu`), and so does `sampler.stats.device`.
 
 The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution
 to within float32 rounding. Float32 cannot tell a measurement outcome of probability below
 about 1e-11 from rounding error, so the GPU never draws one; a circuit with a rotation small
 enough to make such an outcome (an angle within about 6e-5 π of a multiple of π/2) samples on
-the CPU (or raises, with `device="metal"`).
+the CPU (or raises, with `device="metal"` or `"cuda"`).
 
 A seed's samples on the GPU differ from the CPU's and depend on `batch_size`, so with
 `device="auto"` a seed reproduces its samples only on the same kind of machine. Use

pyproject.toml

@@ -54,6 +54,12 @@ Repository = "https://github.com/QuEraComputing/tsim"
 Issues = "https://github.com/QuEraComputing/tsim/issues"
 
 [project.optional-dependencies]
+# The default strategy, "active_statevector", on NVIDIA GPUs (CUDA 13).
+cuda = ["cupy-cuda13x[ctk]>=14.0"]
+# The stabilizer-rank strategies ("cat5", "bss", "cutting") on NVIDIA GPUs, with JAX.
+jax-cuda12 = ["jax[cuda12]>=0.4.38"]
+jax-cuda13 = ["jax[cuda13]>=0.7.0"]
+# NOTE: the names of the JAX extras up to 1.0, kept so that they still install.
 cuda12 = ["jax[cuda12]>=0.4.38"]
 cuda13 = ["jax[cuda13]>=0.7.0"]
 

skills/simulating-circuits-with-tsim/SKILL.md

@@ -197,23 +197,26 @@ yourself.
   default. Use `numba.set_num_threads(k)` before `sample` (or `NUMBA_NUM_THREADS=k` before
   importing Tsim) and leave cores free on a shared machine. Do not run several multi-threaded
   Tsim processes side by side; they oversubscribe the CPU.
-- **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every
-  component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. That
+- **GPUs:** on an Apple-silicon Mac (Metal), or on an NVIDIA GPU with the `cuda` extra, the
+  default strategy samples on the GPU when every component has at most 11 active qubits,
+  often 5 to 10 times faster than all CPU cores. That
   includes Clifford circuits, whose detectors come from the noise alone: the GPU samples the
   noise too. It computes outcome probabilities in float32 and never draws an outcome below
   about 1e-11 (circuits with a rotation angle within about 6e-5 π of a multiple of π/2 stay
   on the CPU). Thread settings do not affect GPU sampling.
 - **`device`** (in `compile_sampler` and `compile_detector_sampler`) chooses where it samples:
   `"auto"` (default) uses the GPU when it can and otherwise the CPU, silently; `"cpu"` stays
-  on the CPU in double precision; `"metal"` raises with the reason if the GPU cannot run the
-  circuit (no Apple GPU, width above 11, a tiny rotation, out of GPU memory). Check
-  `sampler.stats.device` (or the repr, `samples on metal`) after compiling, since `"auto"`
-  does not warn. Pass `device="metal"` for long runs that must not drop to the CPU unnoticed,
+  on the CPU in double precision; `"metal"` (Apple) and `"cuda"` (NVIDIA) raise with the
+  reason if the GPU cannot run the circuit (no such GPU, width above 11, a tiny rotation, out
+  of GPU memory). Check `sampler.stats.device` (or the repr, `samples on metal` or `cuda`)
+  after compiling, since `"auto"` does not warn. Pass `device="metal"` or `"cuda"` for long
+  runs that must not drop to the CPU unnoticed,
   and `device="cpu"` to reproduce samples across machines or to compare against the GPU.
 - **Processes:** sample from one process per GPU. Threads sharing a GPU sampler run its calls
   one at a time, and several processes compete for the same GPU. A process forked after its
-  parent used the GPU cannot use it (it samples on the CPU, or raises with `device="metal"`);
-  start workers with the `"spawn"` or `"forkserver"` method (macOS's default is `"spawn"`).
+  parent used the GPU cannot use it (it samples on the CPU, or raises with `device="metal"` or
+  `"cuda"`); start workers with the `"spawn"` or `"forkserver"` method (macOS's default is
+  `"spawn"`).
 - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and
   device. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU,
   `batch_size` does, and GPU samples differ from CPU samples, so with `device="auto"` a seed

src/tsim/circuit.py

@@ -966,24 +966,25 @@ def compile_sampler(
                 depends on the free memory. On a GPU, they do depend on the batch
                 size.
             device: Where "active_statevector" samples. "auto" (the default)
-                samples in float32 on an Apple GPU when there is one that can
-                run the circuit, and otherwise on the CPU, without a warning.
-                "cpu" samples on the CPU in float64, on
+                samples in float32 on an Apple GPU or an NVIDIA GPU when there
+                is one that can run the circuit, and otherwise on the CPU,
+                without a warning. "cpu" samples on the CPU in float64, on
                 ``numba.get_num_threads()`` threads. "metal" samples on an Apple
-                GPU or raises: without one, for a component of more than 11
-                active qubits, a rotation within about 6e-5 π of a multiple of
-                π/2, noise beyond the GPU's tables, or a call that runs out
-                of GPU memory. The stabilizer rank strategies take only "auto".
+                GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or
+                raises: without one, for a component of more than 11 active
+                qubits, a rotation within about 6e-5 π of a multiple of π/2,
+                noise beyond the GPU's tables, or a call that runs out of GPU
+                memory. The stabilizer rank strategies take only "auto".
 
         Returns:
             A CompiledMeasurementSampler that can be used to sample measurements.
 
         Raises:
             ValueError: If ``device`` is unknown, is not "auto" under a
-                stabilizer-rank strategy, or is "metal" for a circuit that the
-                GPU executor cannot run.
-            RuntimeError: If ``device`` is "metal" and this process cannot use
-                an Apple GPU.
+                stabilizer-rank strategy, or is "metal" or "cuda" for a circuit
+                that the GPU executor cannot run.
+            RuntimeError: If ``device`` is "metal" or "cuda" and this process
+                cannot use such a GPU.
 
         """
         from tsim.sampler import CompiledMeasurementSampler
@@ -1021,24 +1022,25 @@ def compile_detector_sampler(
                 settings, since the default batch size depends on the free memory.
                 On a GPU, they do depend on the batch size.
             device: Where "active_statevector" samples. "auto" (the default)
-                samples in float32 on an Apple GPU when there is one that can
-                run the circuit, and otherwise on the CPU, without a warning.
-                "cpu" samples on the CPU in float64, on
+                samples in float32 on an Apple GPU or an NVIDIA GPU when there
+                is one that can run the circuit, and otherwise on the CPU,
+                without a warning. "cpu" samples on the CPU in float64, on
                 ``numba.get_num_threads()`` threads. "metal" samples on an Apple
-                GPU or raises: without one, for a component of more than 11
-                active qubits, a rotation within about 6e-5 π of a multiple of
-                π/2, noise beyond the GPU's tables, or a call that runs out
-                of GPU memory. The stabilizer rank strategies take only "auto".
+                GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or
+                raises: without one, for a component of more than 11 active
+                qubits, a rotation within about 6e-5 π of a multiple of π/2,
+                noise beyond the GPU's tables, or a call that runs out of GPU
+                memory. The stabilizer rank strategies take only "auto".
 
         Returns:
             A CompiledDetectorSampler that can be used to sample detectors and observables.
 
         Raises:
             ValueError: If ``device`` is unknown, is not "auto" under a
-                stabilizer-rank strategy, or is "metal" for a circuit that the
-                GPU executor cannot run.
-            RuntimeError: If ``device`` is "metal" and this process cannot use
-                an Apple GPU.
+                stabilizer-rank strategy, or is "metal" or "cuda" for a circuit
+                that the GPU executor cannot run.
+            RuntimeError: If ``device`` is "metal" or "cuda" and this process
+                cannot use such a GPU.
 
         """
         from tsim.sampler import CompiledDetectorSampler

src/tsim/core/types.py

@@ -28,7 +28,7 @@
 SamplingStrategy = Literal["active_statevector", "cat5", "bss", "cutting"]
 """``"active_statevector"``, or a stabilizer-rank decomposition strategy."""
 
-Device = Literal["auto", "cpu", "metal"]
+Device = Literal["auto", "cpu", "metal", "cuda"]
 """Where ``"active_statevector"`` samples: ``"auto"`` picks a GPU when one can run it."""