cuBLASDx Python Frontends#

cuBLASDx offers a C++ API that’s callable from CUDA C++ kernels. Python users can access selected cuBLASDx-backed functionality through higher-level Python frontends such as NVIDIA Warp or nvmath-python.

Note

These frontends are not a one-to-one binding of every cuBLASDx C++ API. Feature coverage, compiler backend, packaging, and autotuning behavior are defined by the frontend package. Check the frontend documentation for supported operations, versions, and performance guidance.

NVIDIA Warp#

NVIDIA Warp is a Python library that allows developers to write high-performance simulation and graphics code that runs efficiently on both CPUs and NVIDIA GPUs. It uses just-in-time (JIT) compilation to turn Python functions into fast, parallel kernels, making it ideal for tasks like physics simulation, robotics, and geometry processing. Warp also supports differentiable programming, allowing integration with machine learning frameworks for gradient-based optimization-all while keeping the simplicity of Python.

Warp uses cuBLASDx for its matrix multiplication operations in Tile mode.

This is what a simple matmul kernel using Warp looks like:

@wp.kernel
def tile_gemm(A: wp.array2d(dtype=wp.float32), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float64)):
    # output tile index
    i, j = wp.tid()

    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float64)

    _M = A.shape[0]
    _N = B.shape[1]
    K = A.shape[1]

    count = int(K / TILE_K)

    for k in range(0, count):
        a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
        b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))

        # sum += a*b
        wp.tile_matmul(a, b, sum)

    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))

The Warp GitHub repository can be accessed here and offers multiple examples, including one that implements multi-level-perceptron using cuBLASDx for necessary matrix multiplications or how to simulate an N-Body gravitational problem using cuBLASDx for the matrix multiplications.

Warp provides autotuning out of the box through its Tile model of programming, where the user describes the problem on a high level, and then it is autotuned and mapped onto the hardware by Warp.

nvmath-python#

nvmath-python is a Python library that provides high-performance, pythonic access to NVIDIA’s CUDA-X math libraries, enabling accelerated mathematical operations like linear algebra and fast Fourier transforms on both CPUs and NVIDIA GPUs. It integrates seamlessly with popular Python libraries such as CuPy, PyTorch, and NumPy, allowing users to leverage NVIDIA hardware acceleration within familiar workflows without needing C or C++ bindings. With both stateless and stateful APIs, nvmath-python delivers near-native performance for computational tasks in deep learning, data processing, and scientific computing, while supporting advanced features like device kernel fusion and customizable callbacks

This is what a simple matmul kernel using nvmath-python looks like:

MM = matmul(
    size=(m, n, k),
    precision=np.float16,
    data_type="complex",
    transpose_mode=("non_transposed", "transposed"),
    execution="Block",
    compiler="numba",
)

@cuda.jit(link=MM.files)
def f(a, b, c, alpha, beta, output):
    smem = cuda.shared.array(shape=(0,), dtype=value_type)
    smem_a = smem[0:]
    smem_b = smem[a_size:]
    smem_c = smem[a_size + b_size :]

    load_to_shared(a, smem_a, a_dim, lda)
    load_to_shared(b, smem_b, b_dim, ldb)
    load_to_shared(c, smem_c, c_dim, ldc)

    cuda.syncthreads()

    MM(alpha, smem_a, smem_b, beta, smem_c)

    cuda.syncthreads()

    store_from_shared(smem_c, output, c_dim, ldc)

Pipelined GEMM is exposed through the device Matmul frontend. The descriptor enables the pipeline, the Numba kernel is compiled with cuBLASDx pipeline extensions, and the host creates a device pipeline from the global-memory inputs before launching the kernel. The full nvmath-python example is available as cublasdx_device_gemm_performance_pipeline.py.

The pipeline-specific structure looks like this:

from nvmath.device import Matmul
from nvmath.device.common import axpby, make_tensor
from nvmath.device.cublasdx import MAX_ALIGNMENT, DevicePipeline
from nvmath.device.cublasdx_numba import pipeline_extensions

MM = Matmul(
    size=(tile_m, tile_n, tile_k),
    precision=(np.float16, np.float16, np.float32),
    data_type="real",
    arrangement=("row_major", "col_major", "row_major"),
    execution="Block",
    block_size=block_size,
    alignment=MAX_ALIGNMENT,
    with_pipeline=True,
    enable_input_streaming=True,
    static_block_dim=True,
)

@cuda.jit(extensions=pipeline_extensions, launch_bounds=[MM.block_size, 1])
def matmul_kernel(alpha, beta, c, device_pipeline: DevicePipeline):
    smem = cuda.shared.array(shape=(0,), dtype=np.byte, alignment=device_pipeline.buffer_alignment)

    c_tile = c[
        cuda.blockIdx.x * tile_m : (cuda.blockIdx.x + 1) * tile_m,
        cuda.blockIdx.y * tile_n : (cuda.blockIdx.y + 1) * tile_n,
    ]
    gmem_c = make_tensor(c_tile, MM.get_layout_gmem_c(m))

    tile_pipeline = device_pipeline.get_tile(smem, cuda.blockIdx.x, cuda.blockIdx.y)
    accumulator = MM.suggest_accumulator()
    tile_pipeline.execute(accumulator)
    tile_pipeline._del()

    if accumulator.is_thread_active():
        d_frag = accumulator.make_partition_and_copy(gmem_c)
        axpby(alpha, accumulator.get_results(), beta, d_frag)
        accumulator.partition_and_copy(d_frag, gmem_c)

pipeline_depth = 4
assert k // tile_k >= pipeline_depth
device_pipeline = MM.suggest_device_pipeline(pipeline_depth, a_d, b_d)

The device pipeline owns the dynamic shared-memory requirements and the per-block tile state. Inside tile_pipeline.execute(accumulator), cuBLASDx can overlap staged asynchronous global/shared loads with staged asynchronous MMA compute. More than one load stage and more than one compute stage can be in flight; after the K-stage accumulation completes, the Python kernel runs one epilogue stage that updates the output tile. The complete nvmath-python example also derives a maximum pipeline depth from available shared memory, times the kernel with Numba, and compares the result against cuBLASLt.

nvmath-python GitHub repository can be accessed here and offers multiple examples as well, mirroring the examples from the cuBLASDx C++ repository

nvmath-python also offers autotuning by iterating over all possible configurations, measuring their performance and selecting the best one. This can be done in a simple python for loop and does not require any additional code from the user.

On this page