NVIDIA nvmath-pythonで大規模高性能コア数学を運用

NVIDIA nvmath-pythonは、Python科学コミュニティとNVIDIA CUDA-X数学ライブラリとのギャップを埋めるために設計されたライブラリです。Pythonユーザーに、日常的な…

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は、cuFFT、cuBLASLt、cuDSS、cuSPARSE、cuTENSOR、cuBLASMpなどを含むCUDAおよびNVPL数学ライブラリ上のPythonic抽象化レイヤーです。スパース性の新しいアプローチであるuniversal sparse tensor（UST）により、ユーザーはドメイン固有言語を通じて独自のアプリケーション最適スパース形式を作成でき、コードで実装する必要がありません。

高速で柔軟なインストール

複雑なネイティブ依存関係を持つPythonパッケージのインストールは、時間のかかるフラストレーションのたまる経験になることがあります。nvmath-pythonは迅速にインストールでき、異なる環境に合わせてカスタマイズできます。

pip、conda、uv、pixiなどのパッケージマネージャーを選択できます。
パッケージマネージャーの依存関係解決システムを通じて必要な依存関係をすべてインストールするか、CI/CDやCPUのみの環境などのシナリオに役立つ最小限のインストールを実行するオプションがあります。
CPUバックエンド、デバイスAPIサポート、分散APIを個別に選択できます。
NumPy、CuPy、PyTorchなどのコンパニオン配列ライブラリ（またはすべて）を選択できます。利用可能なオプションについては、詳細なインストールガイドを参照してください。

既存の配列ライブラリに役立つ補完

NumPyなどの他の数学ライブラリと同様に、nvmath-pythonは多くの工学および科学計算アプリケーションで有用なコア数値演算を実装しています。ただし、一般用途の配列ライブラリを置き換えることや、インデックス作成、スライス、削減などの従来の機能を提供することは意図していません。

代わりに、nvmath-pythonはCUDA-X数学ライブラリの全機能とパワーをPythonで公開することに焦点を当て、既存の配列ライブラリやフレームワークが低レベルのC/C++インターフェースに依存することなく、高度に最適化されたGPUアクセラレーションルーチンを使用しやすくします。

以下の例では、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

メモリと実行スペースの選択

配列ライブラリの選択の柔軟性は、CuPyなどのGPUライブラリとNumPyなどのCPUライブラリの両方に適用されます。これは、nvmath-pythonが以下によって支えられているため可能です。

cuBLAS やcuFFTなどのGPUライブラリ。
NVIDIA GraceまたはARM v8 CPU向けのNVPL、x86ホスト向けのIntel MKLなどのCPUライブラリ。
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と専用API

nvmath-python内のAPIは、広く2つのクラスに分けられます。柔軟なマルチツールとして機能する汎用API（幅広いが浅い）と、精密で専用ツールとして設計された専用API（狭く深い）です。

汎用APIは、さまざまな実行・メモリスペースおよびオペランドタイプにわたって統一されたユーザー体験を提供することに焦点を当てていますが、構成可能性は広範なスコープで共有されるベースラインの一般的な機能に制限されます。一方、専用APIは、狭い運用範囲に特化して設計された包括的な機能と構成のセットを提供し、特定のハードウェアに制限される場合があります。

例として、高度な行列乗算は、GPU上の密なオペランドに対して複合演算\(\scriptstyle \mathbf{D}=f(\mathbf{A}\mathbf{B}+\mathbf{C})\)を実装し、ハードウェア効率を最大限に引き出すために必要なすべての構成を提供します。一方、汎用行列乗算APIは、CPUとGPUの実行スペースにわたって密および構造化されたオペランドに対応しますが、より広いスコープに適用可能なオプションの共通サブセットのみを提供します。

最適な選択は、特定のユースケースに完全に依存します。専用APIは、演算が計算のボトルネックとなり、ハードウェア固有の最適化や特定の機能へのアクセスを必要とする場合に理想的です。一方、汎用APIは、パフォーマンスクリティカルでないタスクや、専用カスタマイズが不要な場合に適しています。すべての専用APIは、汎用APIと区別するためにadvancedサブモジュール内に存在します。

nvmath-pythonでのロギング

ライブラリは、さまざまなレベル（デバッグ、情報、警告、エラー）で計算の詳細をキャプチャするためのloggingモジュールからのPython標準ライブラリロガーと統合されています。

以下の例は、メモリと実行スペース間のデータフローを示しています（高度な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.
...

オペランドの出所と消費場所を示すレコードに注意してください。これは、メモリと実行スペース間の潜在的に高コストなデータ転送を示しています。次に、fftのような汎用APIで同様の実験を実行して、メモリと実行スペース間のデータフローを示します。

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}\)がtall-and-skinny行列であるGEMMの計算が注目すべき例です。

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

以下のコードは、CuPyとnvmath-pythonを使用したtall-and-skinny行列上の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は、just-in-timeカーネル融合が可能な基礎のcuBLASLtライブラリにより、はるかに優れたパフォーマンスを発揮します。これは、演算強度を高めるための効果的な手法の一つです。

ステートフルAPIを使用して準備コストを償却

これまでのすべての例は、nvmath-pythonの関数形式、またはステートレスAPIを利用しています。これは便利な単一呼び出しAPIであり、計画フェーズと呼ばれる時間のかかる準備ロジックが含まれます。さらに、準備コストには自動チューニングのコストも含まれる場合があります。これは、計画/自動チューニング後の要求された数学演算を実行する実行フェーズとは異なります。

パフォーマンスに関する注意

NVIDIA CUDA-X数学ライブラリは、最高のパフォーマンスを発揮する特定のインプリメンテーションを決定するためにヒューリスティックを使用します。特定の問題サイズ、レイアウト、データ型に最適化された特殊なカーネルの選択肢が複数存在する場合があります。特定のハードウェア、ワークロード、その他の要因の組み合わせでどのカーネルが最適に実行されるかは、常に明らかではありません。自動チューニングは、カーネルオプションを反復し、パフォーマンスを測定し、最適なものを選択することで、デフォルトのカーネル選択をオーバーライドすることを目的としています。その結果、自動チューニングフェーズは非常に時間のかかるものになる場合があります。

ディープラーニングなどのワークロードでは、同じ演算が異なる入力で繰り返し実行される場合があります。実行間で計画を作成して再利用することで、計画コストを償却できます。nvmath-pythonのクラスベース、またはステートフルAPIは、このワークフローをサポートしています。

以下の例は、batch_sizeサイズの行列aとb、およびバイアスbiasのバッチ上で（RELU_BIASエピローグ付きの）matmulのクラス形式APIの使用を示しています。以前の行列乗算の結果は次の行列乗算のオペランドであり、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は、numba-cudaなどのPythonコンパイラと統合されており、高パフォーマンスのカスタムPythonコードをjust-in-time（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ライブラリを使用し、CuPyのndarrayとしてグレースケール[0, 1]画像に変換します。画像フィルタリングでは、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カーネル

2番目に一般的に使用されるシナリオは、numba-cudaで記述されたGPUカーネル内からnvmath-pythonデバイスAPIを呼び出すことです。nvmath-pythonは、FFT、GEMM、密な直接ソルバー（LU、Cholesky、QR）、RNGのデバイスAPIをサポートしています。以下の例は、モンテカルロ株価シミュレーションのためのGeometric Brownian Motion (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リポジトリ（コードexamplesおよび詳細なチュートリアルnotebooksを含む）
NVIDIA Accelerated Computing Hubの拡張Pythonトレーニング資料
以前の投稿、nvmath-pythonのUniversal Sparse Tensorでスパースディープラーニングを簡素化
最新ニュースと発表のための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, and many more.

Read the original source Next news

#Performance