GPU Acceleration

NextStat supports GPU-accelerated toy generation and batch NLL computation via CUDA (NVIDIA) and Metal (Apple Silicon). The GPU backend is selected at runtime through the CLI or Rust API.

Backends

Building with GPU Support

# CUDA (requires nvcc in PATH)
cargo build --workspace --features cuda

# Metal (Apple Silicon, macOS)
cargo build --workspace --features metal

CLI Usage

# NVIDIA GPU (f64)
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu cuda

# Apple Silicon GPU (f32)
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu metal

GPU-Resident Toy Pipeline

The --gpu-sample-toys flag keeps sampled events on the GPU device, eliminating the D2H+H2D round-trip of the large obs_flat buffer between sampler and batch fitter.

nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu cuda --gpu-sample-toys

Unbinned GPU WeightSys

The weightsys rate modifier is now lowered to CUDA/Metal kernels with code0 and code4p interpolation support. In the unbinned model YAML spec:

modifiers:
  - type: weightsys
    param: alpha_jet_jes
    lo: 0.95
    hi: 1.05
    interp_code: code4p  # optional, defaults to code0

GPU Flow Evaluation

Flow PDFs can leverage GPU acceleration for NLL reduction. Two paths are supported:

• Path 1: CPU flow + GPU NLL — The flow evaluates log p(x|θ) on CPU (ONNX Runtime), the result is uploaded to GPU where a dedicated CUDA kernel performs extended unbinned likelihood reduction. GpuFlowSession manages the pipeline automatically.
• Path 2: CUDA EP + I/O binding — With --features neural-cuda, ONNX Runtime uses the CUDA Execution Provider. The flow forward pass runs on GPU, log-prob stays device-resident. NLL reduction reads directly from GPU memory — zero host↔device copies.

Gradients via central finite differences (ε=1e-4): 2·n_params + 1 NLL calls per iteration.

CPU Acceleration

Even without a GPU, NextStat leverages multiple CPU acceleration strategies:

• SIMD auto-vectorization — Compiler-generated vector instructions for batch operations
• Rayon parallelism — Work-stealing thread pool for toy generation and parameter scans
• Apple Accelerate — vDSP and vForce for vectorized math on macOS (log, exp, etc.)

Differentiable Layer

Both CUDA and Metal support the differentiable NLL layer for PyTorch integration:

Batch Toy Fitting

Both CUDA and Metal support GPU-accelerated batch toy fitting for CLs hypothesis testing:

Architecture: Phase A performs 3 baseline CPU fits (free, conditional at μ_test, conditional at μ=0), then Phase B dispatches to the appropriate GPU backend for batch toy ensemble generation.

# CUDA toy-based CLs
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu cuda

# Metal toy-based CLs
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu metal

Performance Benchmarks

Single-Operation Latency

Batch Toys — Large Model (tHu, 184 params)

GPU lockstep amortizes kernel overhead → sub-linear scaling

Cross-Platform Summary (1,000 toys, 184 params)

Batch Toys — Small Model (complex, 8 params)

Kernel launch overhead dominates → CPU wins on both platforms

Key Findings

• Convergent speedup factor ~6.3× — Both CUDA (f64) and Metal (f32) converge to the same GPU/CPU ratio despite different precision, architecture, and absolute performance. This is a fundamental property of the lockstep L-BFGS-B batch architecture.
• GPU batch scaling is sub-linear — 10× toys → 5.9× time. Lockstep execution amortizes kernel launch overhead across the batch.
• CPU scaling is super-linear on large models — 10× toys → 20.3× time for 184 params due to L3 cache pressure.
• Crossover at ~100 parameters — GPU wins for models with ~100+ parameters. Below that, CPU Rayon parallelism dominates due to negligible per-toy overhead.

Recommendation: Use --gpu cuda or --gpu metal for batch toy workloads on models with 100+ parameters. Use CPU (default) for small models and single-model fits.

Metal Limitations

Known Issues (Fixed)

• Batch toys memcpy_dtoh panic — cudarc 0.19 requires dst.len() >= src.len() for device-to-host copies. When toys converge and n_active < max_batch, the host buffer was too small. Fix: allocate host buffers at max_batch size and truncate.
• ProfiledDifferentiableSession convergence — L-BFGS-B tolerance 1e-6 too tight for projected gradient near parameter bounds. Fix: tolerance 1e-5 + NLL stability criterion.

Validation

# Single-model fit + gradient parity (CUDA)
cargo test -p ns-inference --features cuda -- --nocapture

# Metal batch tests
cargo test -p ns-inference --features metal -- --nocapture

# Python GPU parity
pytest tests/python/test_gpu_parity.py -v

Rust integration tests in crates/ns-inference/src/gpu_single.rs:

• test_gpu_nll_matches_cpu
• test_gpu_grad_matches_cpu
• test_gpu_fit_matches_cpu
• test_gpu_session_reuse
• test_gpu_complex_workspace
• test_gpu_nll_and_grad_at_multiple_points

Tolerance source of truth: tests/python/_tolerances.py (Python) · crates/ns-inference/src/gpu_single.rs (Rust)