Мост между Python и CUDA-X
Выпуск nvmath-python v1.0 знаменует собой важный шаг в интеграции научных вычислений на Python с инфраструктурой NVIDIA. Библиотека выступает в роли абстрактного слоя над мощными инструментами CUDA-X и NVPL, включая такие библиотеки, как cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR и cuBLASMp. Это позволяет использовать оптимизированные под железо операции прямо в Python, сохраняя совместимость с экосистемами NumPy, CuPy и PyTorch.
Универсальные разреженные тензоры (UST)
Ключевой инновацией версии 1.0 стала поддержка Universal Sparse Tensor (UST). Эта технология использует специализированный язык (DSL), позволяющий пользователям создавать собственные оптимальные форматы разреженных данных для конкретных приложений. Это избавляет от необходимости писать сложный код на C++ для настройки форматов хранения данных, что критично для задач машинного обучения и научных симуляций, где эффективность использования памяти играет решающую роль.
Гибкость бэкендов: CPU, GPU и кластеры
Библиотека поддерживает выполнение операций в различных пространствах памяти и вычислений. Разработчики могут выбирать между:
- GPU: cuBLAS, cuFFT, cuSPARSE.
- CPU: NVPL (для архитектур NVIDIA Grace/ARM v8) и Intel MKL (для x86).
- Распределенные системы: cuBLASMp, cuSOLVERMp, cuFFTMp для работы на мульти-нодовых кластерах.
Это упрощает миграцию кода между CPU и GPU, а также позволяет создавать гибридные рабочие процессы. Библиотека интегрируется со стандартным модулем logging Python, что дает детальную трассировку данных: разработчик видит, где находятся операнды (CPU/GPU) и где происходит вычисление, что помогает выявлять узкие места, связанные с пересылкой данных.
Две категории API: Универсальные и Специализированные
Архитектура nvmath-python разделена на два типа интерфейсов, что позволяет балансировать между простотой использования и максимальной производительностью:
| Тип API | Характеристика | Пример использования |
|---|---|---|
| Generic (Универсальные) | Широкий охват, базовая конфигурация, кросс-платформенность (CPU/GPU) | Быстрая реализация прототипов, задачи, не являющиеся узким местом |
| Advanced (Специализированные) | Глубокая настройка, доступ к аппаратным оптимизациям, JIT-компиляция | Высоконагруженные операции (например, fused composite operations), требующие максимальной эффективности |
Специализированные API (находящиеся в подмодулях advanced) позволяют использовать такие функции, как слияние ядер (kernel fusion) через cuBLASLt или JIT-компиляцию FFT-обработчиков через numba-cuda, что позволяет амортизировать затраты на планирование и автонастройку.
Установка и совместимость
Установка библиотеки оптимизирована для скорости и гибкости. Поддерживаются менеджеры пакетов pip, conda, uv, pixi. Доступны варианты установки «все зависимости» или минимальной конфигурации (например, для CI/CD или сред только с CPU). Библиотека не заменяет NumPy или CuPy, а дополняет их, предоставляя доступ к «тяжелой» математике, которую эти библиотеки могут делегировать nvmath-python для ускорения.
Источник: NVIDIA dev blog ↗