運用 NVIDIA nvmath-python 在規模化環境中執行高效能核心數學運算

NVIDIA nvmath-python 是一套旨在縮短 Python 科學社群與 NVIDIA CUDA-X 數學函式庫之間鴻溝的函式庫。它讓 Python 使用者能夠存取 CUDA-X 效能，處理常見…

NVIDIA Technical Blog · Jul 30, 2026

Run High-Performance Core Math at Scale with NVIDIA nvmath-python

NVIDIA nvmath-python 是一套旨在縮短 Python 科學社群與 NVIDIA CUDA-X 數學函式庫之間鴻溝的函式庫。它讓 Python 使用者能夠存取 CUDA-X 效能，處理常見數學運算而不必中斷既有工作流程。依據 API 不同，運算可在 CPU、支援 CUDA 的 GPU，或是分散式多 GPU、多節點系統上執行。

nvmath-python v1.0 版本釋出

隨著 nvmath-python v1.0 正式推出，本文將探討該函式庫的設計理念，以及從 CPU 或單一 GPU 一直到多 GPU、多節點規模下加速數學運算的獨特功能。nvmath-python 是建構於 CUDA 與 NVPL 數學函式庫（例如 cuFFT、cuBLASLt、cuDSS、cuSPARSE、cuTENSOR、cuBLASMp 等）之上的 Python 抽象層。一種新穎的稀疏性處理方法——通用稀疏張量（UST）——讓使用者能透過領域特定語言建立獨特且應用最佳化的稀疏格式，而無需自行實作程式碼。

快速且彈性的安裝

安裝含有複雜原生依賴的 Python 套件可能既費時又令人沮喪。nvmath-python 安裝迅速，並可依不同環境客製化。

選擇套件管理工具，例如 pip、conda、uv 或 pixi。
可選擇透過套件管理工具的依賴解析系統安裝所有必要依賴，或進行最小安裝（適用於 CI/CD 或僅 CPU 環境）。
挑選 CPU 後端、裝置 API 支援或分散式 API。
選擇搭配的陣列函式庫，例如 NumPy、CuPy 或 PyTorch（亦可全部使用）。詳見安裝指南以了解可用選項。

現有陣列函式庫的實用補充

與 NumPy 等其他數學函式庫類似，nvmath-python 實作了許多工程與科學運算應用中常用的核心數值運算。然而，它並非用來取代通用陣列函式庫，也不提供索引、切片或歸約等傳統功能。

相反地，nvmath-python 專注於在 Python 中完整呈現 CUDA-X 數學函式庫的功能與效能，讓現有陣列函式庫與框架能更輕鬆地使用高度最佳化的 GPU 加速常式，而無需依賴低階 C/C++ 介面。

在下列範例中，nvmath-python 接受 NumPy 陣列，結果亦為 NumPy 陣列。

import numpy as np
import nvmath

m, n, k = 10, 40, 100
a = np.random.randn(m, k)  # a is a NumPy array
b = np.random.randn(k, n)  # b is a NumPy array
c = nvmath.linalg.advanced.matmul(a, b)  # c is also a NumPy array

記憶體與執行空間的選擇

選擇陣列函式庫的彈性同時適用於 GPU 函式庫（如 CuPy）與 CPU 函式庫（如 NumPy）。這是因為 nvmath-python 由下列元件支援：

GPU 函式庫，例如 cuBLAS 與 cuFFT。
CPU 函式庫，例如適用於 NVIDIA Grace 或任何 ARM v8 CPU 的 NVPL，以及適用於 x86 主機的 Intel MKL。
分散式函式庫，例如 cuBLASMp、cuSOLVERMp 或 cuFFTMp。

此支援簡化了 CPU 與 GPU 之間的程式碼遷移，並啟用混合與分散式工作流程，結合 CPU 與 GPU 執行。

下列程式碼說明 nvmath-python 如何支援多種記憶體與執行空間。

import cupy as cp
import numpy as np
import nvmath

N = 2048
a_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
a_cpu = np.random.randn(N) + 1j * np.random.randn(N)
c_gpu = nvmath.fft.fft(a_gpu)
c_cpu = nvmath.fft.fft(a_cpu)

每次呼叫的 fft 執行空間會從其輸入張量（a_gpu 或 a_cpu）推斷而得，儘管可指定不同的執行空間。函式庫的記錄功能會顯示每個運算的執行位置。

通用與專用 API

nvmath-python 內的 API 大致可分為兩類：作為靈活多功能工具（廣泛但淺層）的通用 API，以及作為精準專用工具（狹窄但深入）的專用 API。

通用 API 著重於在各種執行與記憶體空間以及運算元類型之間提供一致的使用者體驗，但其可配置性僅限於廣泛範圍內共用的基準功能。另一方面，專用 API 則針對特定狹窄運算範圍提供全面的功能與配置，可能僅限於特定硬體。

舉例來說，先進矩陣乘法針對 GPU 上的密集運算元實作複合運算 \(\scriptstyle \mathbf{D}=f(\mathbf{A}\mathbf{B}+\mathbf{C})\)，並提供所有必要配置以榨取最高硬體效率。反之，通用矩陣乘法 API 則支援 CPU 與 GPU 執行空間上的密集與結構化運算元，但僅提供適用於其較廣範圍的通用選項子集。

最佳選擇完全取決於特定使用情境：當運算成為效能瓶頸，需要硬體特定最佳化或存取獨特功能時，專用 API 是理想選擇。而通用 API 則更適合非效能關鍵任務，或不需要專用客製化的情境。所有專用 API 皆位於 advanced 子模組中，以與通用 API 區分。

使用 nvmath-python 進行記錄

該函式庫提供與 Python 標準函式庫記錄器（來自logging 模組）的整合，用於在各種層級（除錯、資訊、警告與錯誤）擷取運算細節。

下列範例說明記憶體與執行空間之間的資料流（使用先進 matmul）。

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

m, n, k = 8000, 2000, 4000
a_cpu = np.random.randn(m, k).astype(np.float32)
b_cpu = np.random.randn(k, n).astype(np.float32)
d_cpu = nvmath.linalg.advanced.matmul(a_cpu, b_cpu)

產生的輸出將如下所示：

2025-09-18 14:53:32,166 INFO     = SPECIFICATION PHASE =
2025-09-18 14:53:32,167 INFO     The data type of operand A is 'float32', and that of operand B is 'float32'.
2025-09-18 14:53:32,168 INFO     The input operands' memory space is cpu, and the execution space is on device 0.
...

請注意顯示運算元來源與消耗位置的記錄。這表示記憶體與執行空間之間可能存在昂貴的資料傳輸。接著執行類似實驗，使用通用 API（如 fft）來說明記憶體與執行空間之間的資料流。

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

N = 10000
e_cpu = (np.random.randn(N) + 1j * np.random.randn(N)).astype(np.complex64)
r_cpu = nvmath.fft.fft(e_cpu)

記錄輸出如下：

2025-09-18 15:46:22,295 INFO     The FFT type is C2C.
2025-09-18 15:46:22,295 INFO     The input data type is complex64, and the result data type is complex64.
2025-09-18 15:46:22,296 INFO     The specified FFT axes are (0,).
2025-09-18 15:46:22,297 INFO     The input tensor's memory space is cpu, and the execution space is cpu, with device cpu.
2025-09-18 15:46:22,298 INFO     The specified stream for the FFT ctor is None.
...

請注意執行空間與輸入的記憶體空間相同。只要可能，nvmath-python 會選擇執行空間以最小化資料傳輸負擔。使用者可透過向 API 提供 execution 關鍵字引數，自由選擇所需的執行空間。

為什麼複合運算很重要

在純 NumPy 類 API 中，類似 \(\scriptstyle \mathbf{D}=f(\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{C})\) 的運算在許多使用情境中表現良好。然而，當基礎基本運算具有低算術強度時，將其串聯為一系列呼叫會造成效率低下。一個值得注意的例子是使用 \(\scriptstyle \mathbf{A}\) 為高瘦矩陣計算 GEMM：

\(\scriptstyle \mathbf{D}=\alpha\mathbf{A}\cdot\mathbf{B}+\beta\mathbf{C}\)

下列程式碼說明使用 CuPy 與 nvmath-python 在高瘦矩陣上執行 GEMM。

import cupy as cp
import nvmath

m, n, k = 10_000_000, 40, 10
a = cp.random.randn(m, k, dtype=cp.float32)
b = cp.random.randn(k, n, dtype=cp.float32)
c = cp.random.randn(m, n, dtype=cp.float32)
alpha, beta = 1.5, 0.5

d1 = alpha * cp.matmul(a, b) + beta * c # Multiple kernels
d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta) # Single kernel

圖 1 顯示融合複合運算相較於 NumPy 類 API 帶來可衡量的效益。

Significant performance speedup of fused kernels compared to unfused in low arithmetic intensity composite operations. — *圖 1. 低算術強度複合運算中，核心融合的效果*

nvmath-python 表現更佳，因為其建構於可即時進行核心融合的 cuBLASLt 函式庫。這是提升算術強度的有效技術之一。

使用狀態式 API 分攤準備成本

所有先前的範例皆使用 nvmath-python 的函數式或無狀態 API。這是一種便利的單次呼叫 API，涉及稱為規劃階段的耗時準備邏輯。此外，準備成本可能還包括自動調校的成本。這與在規劃/自動調校後執行所要求數學運算的執行階段不同。

效能注意事項

NVIDIA CUDA-X 數學函式庫採用啟發式方法來決定能產生最佳效能的特定實作。針對特定問題大小、配置或資料類型，可能有多種最佳化核心可供選擇。並非總是顯而易見哪個核心在特定硬體、工作負載與其他因素的組合上表現最佳。自動調校的目標是透過迭代核心選項、測量其效能並選擇最佳者，來覆寫預設核心選擇。因此，自動調校階段可能非常耗時。

在深度學習等工作負載中，相同的運算可能會重複執行，但輸入不同。在執行之間建立並重複使用規劃，可分攤其規劃成本。nvmath-python 的類別式或狀態式 API 支援此工作流程。

下列範例說明在 batch_size 大小的 a 與 b 矩陣批次以及偏差 bias 上使用 matmul 的類別式 API（搭配 RELU_BIAS 尾碼）。前一次矩陣乘法的結果是下一次矩陣乘法的運算元，且有 feed_count 次運算。除了規劃之外，它還會執行自動調校階段。

下列程式碼說明如何使用狀態式 API，將規劃、自動調校與執行視為不同階段。

import nvmath
from nvmath.linalg.advanced import MatmulEpilog
import cupy as cp

feed_count = 10  # The operation feed count.

batch_size = 1024
m, n, k = 1024, 1024, 1024

a = cp.random.rand(batch_size, m, k, dtype=cp.float32)
b = cp.random.rand(batch_size, k, n, dtype=cp.float32)
bias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)

with nvmath.linalg.advanced.Matmul(a, b) as mm:
    # 1. Planning phase
    mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
            epilog_inputs={"bias": bias})

    # 2. Autotuning phase
    mm.autotune(iterations=5)

    # 3. Execution phase.
    for i in range(feed_count):
        d = mm.execute()
        # The result of the previous MM is the operand `a` of the next MM, so use
        # reset_operands_unchecked() to reset the `a` operand.
        mm.reset_operands_unchecked(a=d)

Stateless APIs add a constant overhead to each operation. Stateful APIs quickly amortize one-time initial costs — *圖 2. 狀態式 API 透過多次執行快速分攤規劃與自動調校成本*

圖 2 顯示運算成本如何隨執行次數變化。細虛線代表使用 nvmath-python 無狀態 API 的成本。粗虛線顯示切換至狀態式 API 帶來的成本降低，而虛線-點線則顯示自動調校帶來的額外效能增益。狀態式 API 分攤規格與準備成本，而無狀態 API 在每次執行時都會產生這些成本。自動調校效益可跨工作階段延伸，因為自動調校的規劃可序列化至磁碟並在新的工作階段中載入。

圖 3 顯示，內建啟發式方法通常可以在不進行自動調校的情況下選擇高效能核心。然而，某些問題大小、資料類型、運算元配置、硬體與其他因素的組合會從自動調校中受益。在測試的配置中，NVIDIA RTX A6000 顯示出最大的加速，而 NVIDIA B200 在無需自動調校的情況下即可達到峰值效能。

Bar chart showing the percentage speedup from autotuning on two GPUs: NVIDIA RTX A6000 shows over 250% speedup, while B200 shows minimal gains for this example (indicating that the heuristics do a great job). — *圖 3. 執行階段的自動調校並非總是能帶來顯著加速，儘管 RTX A6000 在此特定問題中獲得 256% 的增益*

使用 nvmath-python 融合自訂核心

nvmath-python 與 Python 編譯器（如 numba-cuda）整合，讓高效能自訂 Python 程式碼能夠即時（JIT）編譯並與 nvmath-python 運算一起使用。

自訂 FFT 回呼

FFT 回呼以具有預先定義簽章的 Python 函數撰寫，並 JIT 編譯為中介表示，稍後用作 nvmath-python 前向或反向 FFT 的自訂前序或尾碼。

Alt text: A sharp, grayscale, overhead photo of a light-colored dog sitting on pavement with its front paws elegantly crossed, casting a dark shadow. Then, the same photo of the dog now heavily blurred to demonstrate the effect of a Gaussian image filter. — *圖 4. 左側為套用高斯濾波器前的原始灰階狗狗影像。右側為結果影像，此影像使用 nvmath-python 的 FFT 搭配自訂 JIT 編譯回呼實作*

高斯濾波器範例

作為說明，我們實作高斯濾波器，對原始影像進行模糊處理。以下程式碼片段使用 PIL 函式庫載入影像，然後轉換為 [0, 1] 的灰階影像，作為 CuPy ndarray。對於影像過濾，我們實作 img → R2C FFT → 高斯濾波器 → C2R iFFT → filtered_img 的鏈結。高斯濾波器為 \(\scriptstyle G(x,y)=\exp\left(-\frac{x^2+y^2}{2\sigma^2}\right)\)，在頻域中亦為高斯函數 \(\scriptstyle H(f_x,f_y)=\exp\left(-2\pi^2\sigma^2(f_x^2+f_y^2)\right)\)。

下列程式碼顯示如何使用 nvmath-python FFT 與自訂回呼函數套用高斯影像濾波器：

from PIL import Image
import nvmath
import cupy as cp

img = cp.asarray(Image.open("your_lovely_dog.jpg").convert("L")) / 255.0 # Gray[0,1]
wh = img.shape[0] * image.shape[1] # We must normalize by the image area
sigma_value = 20.0  # Filter size

# Implement Gaussian filter in the frequency domain
def gaussian_filter(shape, sigma):
    fy = cp.fft.fftfreq(shape[0])[:,None] # Column
    fx = cp.fft.rfftfreq(shape[1])[None,:] # Row
    return = cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))

# Implement FFT epilog wrapper with the pre-defined signature
def epilog_impl(data_out, offset, data, filter_data, unused): # Epilog to be compiled
    data_out[offset] = data * filter_data[offset] / wh

# Compile epilog to LTO-IR targeting the current CUDA device
epilog = nvmath.fft.compile_epilog(epilog_impl, "complex64", "complex64")

# Compute R2C FFT using nvmath-python with the compiled epilog
h_filter = gaussian_filter(img.shape, sigma)
img_fft = nvmath.fft.rfft(image, epilog={"ltoir": epilog, "data": h_filter.data.ptr})

# Compute C2R inverse FFT using nvmath-python
filtered_img = nvmath.fft.irfft(img_fft) # Visualize or save as you want

搭配 nvmath-python 呼叫的自訂 numba-cuda 核心

第二種常見情境是從以 numba-cuda 撰寫的 GPU 核心中呼叫 nvmath-python 裝置 API。nvmath-python 支援 FFT、GEMM、密集直接求解器（LU、Cholesky、QR）與 RNG 的裝置 API。下列範例顯示為蒙地卡羅股價模擬實作幾何布朗運動（GBM）。它使用 nvmath-python 的亂數產生器產生高斯分佈，以及將常態分佈轉換為 GBM 蒙地卡羅路徑的自訂 numba-cuda 程式碼：

from numba import cuda
from nvmath.device import random
import cupy as cp
import math

# Pre-compile the RNGs into IR to use alongside other device code
compiled_rng = random.Compile(cc=None)

# GBM parameters
rng_seed = 7777
n_time_steps, n_paths = 252, 8192
mu, sigma, s0 = 0.003, 0.027, 100.0

# Set up CUDA kernel launch configuration
threads_per_block = 32
blocks = n_paths // threads_per_block + bool(n_paths % threads_per_block)
nthreads = threads_per_block * blocks

# RNG initialization kernel
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def init_rng(states, seed):
    idx = cuda.grid(1)
    random.init(seed, idx, 0, states[idx])

# GBM path generation kernel
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def generate_gbm_paths(states, paths, nsteps, mu, sigma, s0):
    idx = cuda.grid(1)
    if idx >= paths.shape[0]:
        return
    paths[idx, 0] = s0

    # Consume 4 normal variates at a time for better throughput
    for i in range(1, nsteps, 4):
        v = random.normal4(states[idx])  # Returned as float32x4 type
        vals = v.x, v.y, v.z, v.w  # Decompose into a tuple of float32
        for j in range(i, min(i + 4, nsteps)):  # Process a chunk of 4 time steps
            paths[idx, j] = paths[idx, j - 1] * math.exp(mu + sigma * vals[j - i])

# Initialize RNG
states = random.StatesPhilox4_32_10(nthreads)
init_rng[blocks, threads_per_block](states, rng_seed)

# Generate GBM paths on GPU
paths = cp.empty((n_paths, n_time_steps), dtype=cp.float32, order='F')
generate_gbm_paths[blocks, threads_per_block](states, paths, n_time_steps, mu, sigma, s0)

generate_gbm_paths 中的每個運算都具有低算術強度，這使得基於主機 API 的實作效率低下。將這些運算與 numba-cuda 和 nvmath-python 裝置 API 融合至關重要。

開始使用 nvmath-python

nvmath-python 的設計旨在提升生產力而不犧牲效能，重新構思了現代數學函式庫的設計。只需一個簡單指令即可開始使用：

pip install nvmath-python[cu13]

其他資源包括：

包含詳細設定說明的安裝文件
nvmath-python GitHub 儲存庫，包含程式碼範例與深入教學筆記本
NVIDIA 加速運算中心的延伸 Python 訓練教材
先前的文章：使用 nvmath-python 中的通用稀疏張量簡化稀疏深度學習
nvmath-python 產品頁面，了解最新消息與公告

致謝

本函式庫是 NVIDIA 眾多同仁共同努力的成果，包括：

Harun Bayraktar, Becca Zandstein, Lukasz Ligowski, Aart Bik, Yevhenii Havrylko, Juan Galvez, Daniel Ching, Mark Olah, Yang Gao, Szymon Karpinski , Kamil Tokarski , Francesco Rizzi, Jakub Lisowski, Marcin Rogowski, Robbie Jensen , Artem Amogolonov, Sushma Kini, Rachna Pandey, Graham Markall, Michael Yh Wang, Bradley Dice, Liam Zhang, Jack Cui, Chang Liu, Qi Xia, Feng Cheng, Ruilin Tian, Zan Xu, Almog Segal, Kirill Voronin, Evarist Fomenko，以及更多同仁。

Read the original source Next news

#Performance