GPU-ускорение

NextStat поддерживает GPU-ускоренную генерацию тоев и батч-вычисление NLL через CUDA (NVIDIA) и Metal (Apple Silicon). GPU-бэкенд выбирается во время выполнения через CLI или Rust API.

Бэкенды

Сборка с поддержкой GPU

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

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

Использование CLI

# 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

Флаг --gpu-sample-toys хранит сгенерированные события прямо на GPU, исключая round-trip D2H+H2D для буфера obs_flat между сэмплером и батч-фиттером.

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

Unbinned GPU WeightSys

Модификатор weightsys теперь lowered в CUDA/Metal ядра с поддержкой интерполяции code0 и code4p. В YAML-спецификации unbinned-модели:

modifiers:
  - type: weightsys
    param: alpha_jet_jes
    lo: 0.95
    hi: 1.05
    interp_code: code4p  # опционально, по умолчанию code0

GPU Flow Evaluation

Flow PDF могут использовать GPU для NLL reduction. Два пути:

• Path 1: CPU flow + GPU NLL — Flow вычисляет log p(x|θ) на CPU (ONNX Runtime), результат загружается на GPU, где CUDA-ядро выполняет extended unbinned likelihood reduction. GpuFlowSession управляет пайплайном автоматически.
• Path 2: CUDA EP + I/O binding — С --features neural-cuda ONNX Runtime использует CUDA Execution Provider. Forward pass flow на GPU, log-prob остаётся device-resident. NLL reduction читает напрямую из GPU-памяти — zero host↔device копирование.

Градиент через центральные конечные разности (ε=1e-4): 2·n_params + 1 NLL вызовов на итерацию.

Ускорение на CPU

Даже без GPU NextStat использует несколько стратегий ускорения на CPU:

• SIMD автовекторизация — Компиляторные векторные инструкции для батч-операций
• Параллелизм Rayon — Work-stealing пул потоков для генерации тоев и сканов параметров
• Apple Accelerate — vDSP и vForce для векторизованной математики на macOS (log, exp и т.д.)

Дифференцируемый слой

CUDA и Metal поддерживают дифференцируемый NLL-слой для интеграции с PyTorch:

Батч-фиттинг тоев

CUDA и Metal поддерживают GPU-ускоренный батч-фиттинг тоев для CLs-проверки гипотез:

Архитектура: фаза A выполняет 3 базовых CPU-фита (свободный, условный при μ_test, условный при μ=0), затем фаза B отправляет на соответствующий GPU-бэкенд для батч-генерации ансамбля тоев.

# CUDA: CLs на тоях
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu cuda

# Metal: CLs на тоях
nextstat hypotest-toys --input workspace.json \
  --mu 1.0 --n-toys 10000 --gpu metal

Бенчмарки производительности

Single-Операция Latency

Батч-тои — большая модель (tHu, 184 params)

GPU lockstep amortizes накладные расходы ядра → sub-linear scaling

Cross-Платформа Итоги (1,000 toys, 184 params)

Батч-тои — малая модель (complex, 8 параметров)

Доминируют накладные расходы запуска ядра → CPU побеждает на обеих платформах

Ключевые наблюдения

• Сходящийся фактор ускорения ~6.3× — CUDA (f64) и Metal (f32) сходятся к одинаковому отношению GPU/CPU несмотря на разную точность, архитектуру и абсолютную производительность. Фундаментальное свойство lockstep-архитектуры батч L-BFGS-B.
• GPU-масштабирование батча сублинейно — 10× тоев → 5.9× времени. Lockstep-выполнение амортизирует накладные расходы запуска ядра по батчу.
• Масштабирование CPU суперлинейно на больших моделях — 10× тоев → 20.3× времени для 184 параметров из-за давления на L3-кеш.
• Кроссовер при ~100 параметрах — GPU побеждает для моделей со ~100+ параметрами. Ниже — доминирует параллелизм Rayon из-за пренебрежимо малых накладных расходов на тоя.

Рекомендация: используйте --gpu cuda или --gpu metal для батч-генерации тоев на моделях со 100+ параметрами. Для малых моделей и одиночных фитов используйте CPU (по умолчанию).

Ограничения Metal

Известные проблемы (исправлены)

• Паника memcpy_dtoh в батч-тоях — cudarc 0.19 требует dst.len() >= src.len() для device-to-host копирования. Когда тои сходятся и n_active < max_batch, хост-буфер был слишком мал. Исправление: выделять хост-буферы размером max_batch и обрезать.
• Сходимость ProfiledDifferentiableSession — допуск L-BFGS-B 1e-6 слишком жёсткий для проекционного градиента вблизи границ параметров. Исправление: допуск 1e-5 + критерий стабильности NLL.

Валидация

# Фит одной модели + паритет градиента (CUDA)
cargo test -p ns-inference --features cuda -- --nocapture

# Батч-тесты Metal
cargo test -p ns-inference --features metal -- --nocapture

# Паритет Python GPU
pytest tests/python/test_gpu_parity.py -v

Интеграционные тесты Rust в 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

Источники допусков: tests/python/_tolerances.py (Python) · crates/ns-inference/src/gpu_single.rs (Rust)