Add opt-in CuPy fusion for Gaussian xFit fits

.github/workflows/ci.yml

@@ -157,6 +157,7 @@ jobs:
           tests/xfit/test_models.py::test_cupy_finite_difference_theta_fix_matches_numpy
           tests/xfit/test_models.py::test_fit_dipoles_device_matches_portable_result_without_bulk_d2h
           tests/xfit/test_solver.py::test_cupy_tall_jacobian_diagnostics_match_numpy
+          tests/xfit/test_cupy_fusion.py
       - name: Require all selected tests to execute
         run: |
           uv run --no-sync python - <<'PY'
@@ -166,11 +167,11 @@ jobs:
 
           report = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "xfit-gpu.xml")
           cases = report.findall(".//testcase")
-          assert len(cases) == 6, len(cases)
+          assert len(cases) == 43, len(cases)
           assert not report.findall(".//skipped"), "GPU tests must not skip"
           assert not report.findall(".//failure")
           assert not report.findall(".//error")
-          print("All 6 selected tests passed without skips")
+          print(f"All {len(cases)} selected tests passed without skips")
           PY
       - name: Run native xFit GPU tests
         run: >-

docs/components/xfit.md

@@ -66,6 +66,26 @@ applies one mask or variance plane to every channel of each candidate. If the
 batch size is three, `(3, y, x)` keeps that per-candidate meaning; use the
 explicit `(1, 3, y, x)` shape for per-plane values.
 
+With `backend="cupy"`, `fusion=True` evaluates Gaussian residuals and weighted
+analytic Jacobians in fused CuPy CUDA kernels. The Jacobian is materialized
+for the same CuPy normal-equation contractions used by the default solver.
+It supports float32 and float64, difference and split modes, masks, and
+variances. Final uncertainty diagnostics use the full analytic Jacobian.
+Each fit owns its buffers and runs on the current CuPy stream.
+Select it in the CLI with `--backend cupy --fusion`; the default is
+`fusion=False`. Device pipeline manifests select it with `xfit.fusion: true`.
+Sampled-stamp and finite-difference fits use the default path.
+
+The public synthetic benchmark records the first call separately from every
+warmed sample, including input transfer and result materialization:
+
+```bash
+uv run --locked --extra gpu python examples/xfit/benchmark_gaussian.py \
+  --backend cupy --fusion --batch 16 --height 51 --width 51 --repeat 5
+```
+
+Run the same command without `--fusion` for the default CuPy implementation.
+
 An explicit `backend="cutile"` uses one `cuda.tile` CTA per Gaussian fit to
 form its weighted 8-by-8 normal equations directly. Select this backend
 explicitly and install the `cuphoton[cutile]` extra on Linux with Python

examples/xfit/benchmark_gaussian.py

@@ -59,7 +59,9 @@ def _fixture(
     return images, initial
 
 
-def _initialize(images, initial, mode, workers, counter, barrier, gpu):
+def _initialize(
+    images, initial, mode, workers, counter, barrier, gpu, fusion
+):
     with counter.get_lock():
         index = counter.value
         counter.value += 1
@@ -72,6 +74,7 @@ def _initialize(images, initial, mode, workers, counter, barrier, gpu):
     _worker.initial = initial[start:stop].copy()
     _worker.variance = np.ones_like(_worker.images)
     _worker.mode = mode
+    _worker.fusion = fusion
     _worker.barrier = barrier
     _worker.stream = None
     _worker.pool = None
@@ -113,6 +116,7 @@ def _run(treatment):
             variance=_worker.variance,
             mode=_worker.mode,
             backend=backend,
+            fusion=_worker.fusion,
         )
         seconds = time.perf_counter() - start
     if result.backend != backend:
@@ -198,6 +202,7 @@ def main() -> None:
         help="Restrict all workers together to this many inherited CPUs.",
     )
     parser.add_argument("--batch", type=int, action="append")
+    parser.add_argument("--fusion", action="store_true")
     parser.add_argument(
         "--dtype", choices=("float32", "float64"), default="float64"
     )
@@ -227,7 +232,13 @@ def main() -> None:
         parser.error("workers and image dimensions must be positive")
     if args.compare_native and args.backend:
         parser.error("--compare-native cannot be combined with --backend")
-    backends = args.backend or ["cupy", "cutile"]
+    backends = args.backend or (
+        ["cupy"] if args.fusion else ["cupy", "cutile"]
+    )
+    if args.fusion and (
+        args.compare_native or any(backend != "cupy" for backend in backends)
+    ):
+        parser.error("--fusion requires --backend cupy")
     batches = args.batch or [1, 16, 256, 4096]
     if any(batch < args.workers for batch in batches):
         parser.error("every batch must contain at least --workers candidates")
@@ -274,17 +285,33 @@ def main() -> None:
                 context.Value("i", 0),
                 context.Barrier(args.workers),
                 any(backend != "numpy" for backend in treatments),
+                args.fusion,
             ),
             **executor_options,
         ) as pool:
             _round(pool, None, args.workers, args.timeout_seconds)
             startup = time.perf_counter() - start
             warmups, samples, last = {}, {t: [] for t in treatments}, {}
+            cold = {}
             try:
                 expected, reference_seconds = _round(
                     pool, reference, args.workers, args.timeout_seconds
                 )
                 for treatment in treatments:
+                    if treatment == reference:
+                        cold[treatment] = 1e3 * reference_seconds
+                    else:
+                        results, seconds = _round(
+                            pool,
+                            treatment,
+                            args.workers,
+                            args.timeout_seconds,
+                        )
+                        cold[treatment] = 1e3 * seconds
+                        for actual, baseline in zip(
+                            results, expected, strict=True
+                        ):
+                            _validate(actual[1], baseline[1], dtype)
                     start = time.perf_counter()
                     for _ in range(args.warmup):
                         results, _ = _round(
@@ -345,6 +372,9 @@ def main() -> None:
                 json.dumps(
                     {
                         "backend": treatment,
+                        "fusion": args.fusion,
+                        "cold_milliseconds": cold[treatment],
+                        "milliseconds_samples": wall_ms,
                         "execution": args.execution,
                         "workers": args.workers,
                         "batch": batch,