Agent BLAS Performance Guide#

Note

This guide is written for AI coding agents: its prose is deliberately terse and imperative. A Markdown mirror for agent consumption is published with the HTML documentation, in the markdown/ folder of the HTML tree (agent index: llms.txt). Paths below are relative to the HTML root:

cuBLASDx performance depends on the kernel it is embedded in, so different idioms win in different cases. This guide is a procedure. For a new kernel it walks the problem from its shape to a tuned kernel (The Path). For a port of an existing kernel, follow Porting An Existing Kernel to distinguish application constraints from implementation choices that can be tuned.

This guide covers GEMM only. TRSM (triangular solve) is documented in the “Using cuBLASDx TRSM” page (using_trsm) of the HTML documentation; its operators and execution model are different and the idioms below do not apply to it.

Note

This guide is self-contained.

It provides workflows for designing a new kernel and porting an existing one, with the decisions and measurements needed at each step. Use the Build Card for build setup and the API card for the computation: Hopper wgmma, Blackwell tcgen05 or supermma. The supermma card covers both the per-tile and pipeline APIs.

The Path (designing a new kernel)#

When the kernel does not exist yet, the whole procedure is three steps. Everything else in the guide is reference material for one of them. For a port, the porting workflow governs and these chapters are reference.

  1. Define the problem and its axes of freedom. Chapter: Finding Problem Shape. Output: the operator answers for a descriptor, plus a record of which axes are fixed by the application and which are free to experiment with (Axes That Can Be Relaxed — the highest-leverage decisions in the guide). Done when: every operator is a decision you made rather than a default you inherited.

  2. Map onto idioms, build the candidates, measure which wins. Chapters: Fitting Into cuBLASDx Idioms (find your rows in the two tables), then Finding Best Coarse Solution (build the matching guide variants, measure, perturb one axis at a time). Done when: the winning combination is within a modest factor of a reference point and further perturbations move it by only a few percent.

  3. Optimize the cuBLASDx use. Chapter: Finetuning Your Program — the low-hanging-fruit checklist first, then profile-directed fixes. Done when: the finetuning checklist passes, including a written-down accuracy cost for every axis relaxed in step 1.

Two preconditions, not steps: the library is installed (below), and every number is taken under Performance Measurement — steps 2 and 3 are meaningless if the measurements lie.

Keep a correctness reference runnable throughout: the naive implementation if you have one, otherwise the first guide variant you adapt in step 2, before any optimization. It is what tells you whether a relaxed axis cost accuracy.

The Library In Brief#

cuBLASDx is a device-extension library from the MathDx package: used inside your kernel, never called from the host. GEMM and TRSM, plus the layout, data-movement and pipelining extensions those operations require. CuTe tensors are the preferred data containers and a key optimization target (with suggested swizzled layouts), but the library is self-contained — you do not need to know CuTe. Three interfaces:

  • Descriptor — composed expression-template operators describing the problem and the resources, statically, so exactly one implementation is compiled:

    using BLAS = decltype(Function<MM>() + Precision<__half, __half, float>() + BlockDim<128>() + ...);

    Executes a per-tile problem (BLAS().execute(...)) and answers queries: traits (BLAS::block_dim), layouts (BLAS::suggest_layout_smem_a()), an accumulator (BLAS::suggest_accumulator()).

  • Accumulator — in register output mode, the results object: data plus knowledge of the transform that produced it. Carries the epilogue methods; can be the accumulation target of a following GEMM, never its A / B operand.

  • Data movement — cublasdx::copy (vectorized global/shared copies; requires cublasdx::copy_wait()) and the pipeline API, the only path to TMA, wgmma and tcgen05.

Getting The Library#

cuBLASDx ships inside the MathDx package:

  1. Download MathDx from the cuBLASDx downloads page.

  2. Unpack it; the contents land in <your_directory>/nvidia/mathdx/yy.mm/.

  3. Point your build at it, either through CMake:

    find_package(mathdx REQUIRED COMPONENTS cublasdx CONFIG)
target_link_libraries(YourProgram mathdx::cublasdx_no_lto)   # header-only, GEMM only

    configured with -Dmathdx_ROOT="<your_directory>/nvidia/mathdx/yy.mm/"; or, for a standalone program, by including directly:

    nvcc -std=c++17 -gencode arch=compute_XYa,code=sm_XYa -O3 -DNDEBUG --expt-relaxed-constexpr \
     -I<mathdx_include_dir> -I<cutlass_include_dir> your_kernel.cu   # XYa: 90a on Hopper, 100a on Blackwell

    --expt-relaxed-constexpr is required: without it the device pass silently drops the cuBLASDx calls — the build succeeds with only a warning (#20013-D) and the kernel runs and writes nothing. The CMake targets carry the flag already.

    Spell architecture-specific targets explicitly, with -gencode arch=compute_90a,code=sm_90a (or the sm_100a equivalent), and confirm with cuobjdump --list-elf that the binary carries an sm_90a cubin.

    The -std=c++17 on this line is for standalone builds only — see the Build Card for PyTorch extensions.

TRSM additionally links a prebuilt LTO fatbin: use mathdx::cublasdx instead of mathdx::cublasdx_no_lto, plus CUDA_SEPARABLE_COMPILATION and INTERPROCEDURAL_OPTIMIZATION on the consuming target (on CUDA Toolkit older than 13.2, link the cublasdx_fatbin target explicitly). GEMM needs none of that.

Build Release with -O3 -DNDEBUG. Debug builds enable runtime asserts throughout cuBLASDx and CuTe, and their measurements tell you nothing.

Build Card#

Building inside a PyTorch extension#

For torch.utils.cpp_extension.load or load_inline, use PyTorch’s default C++ standard. Its headers may require a newer standard than the standalone example above. Do not copy nvcc -std=c++17 into extra_cflags or extra_cuda_cflags; remove an existing C++17 override from both lists.

PyTorch also defines macros that disable CUDA half and bfloat16 conversions and operators. cuBLASDx’s commondx/complex_types.hpp needs those operations. Errors about missing half/bfloat16 operators or conversions can therefore be extension build flags, rather than an error in your GEMM descriptor. Append these undefines to the CUDA flag list passed to the extension builder:

extra_cuda_cflags += [
    "-U__CUDA_NO_HALF_OPERATORS__",
    "-U__CUDA_NO_HALF_CONVERSIONS__",
    "-U__CUDA_NO_HALF2_OPERATORS__",
    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
]

Three more facts about this builder:

  • It does not add -DNDEBUG. Add it to extra_cuda_cflags yourself, or the kernel is a debug build with device-side checks in the wgmma region. These checks can increase register pressure and trigger ptxas advisory C7510, wgmma.mma_async instructions are serialized, and register spills in --resource-usage.

  • Pass the library with extra_include_paths=[f"{CUBLASDX_HOME}/include"]; if your package keeps CUTLASS separately (MathDx: external/cutlass/include), add that directory too.

  • It takes the target from TORCH_CUDA_ARCH_LIST (9.0a for Hopper); confirm a -gencode appears in the emitted command. If a flag contains the substring arch, PyTorch drops every -gencode.

Apply this to the builder actually compiling the extension, including a separate benchmark builder if you wrote one. Inspect its emitted compile command once to confirm the undefines appear after PyTorch’s defines, -DNDEBUG is present and no C++17 override remains.

Release build and the C7510 advisory#

Compile the measured CUDA kernel with -O3 -DNDEBUG, including PyTorch extensions and custom benchmark builders.

If ptxas prints advisory C7510 about serialized wgmma, first check for assert/checking code in the WGMMA region: add -DNDEBUG to the CUDA flags actually compiling that kernel, rebuild, and inspect the emitted command and advisory. Apply the same flag to the reference arm.

Excessive register pressure can force ptxas to serialize wgmma instructions and emit C7510. Once the build flags are correct, address the pressure in the kernel:

  • Reduce the tile size to lower the number of accumulator registers.

  • Shorten register lifetimes in the surrounding code: compute temporary values closer to where they are used, and avoid keeping unnecessary intermediates live across the wgmma region.

Rebuild and check the advisory and register usage after each change.

A cuBLASDx pipeline translation unit takes tens of seconds to compile, so a serial build of many variants is slow. Build variants in parallel (for example make -j<N> over several targets) rather than one after another.

Performance Measurement#

Tensor-core GEMMs throttle at boost clocks, so measurement discipline decides whether a number is real:

  1. N unmeasured warm-up runs, then 10 * N measured ones (the guide uses N = 5). If clocks cannot be locked, interleave the arms from the first measurement rather than adding a second protocol later.

  2. Never measure on zeroed buffers. Operand values change power draw, and power draw changes sustained clocks: a large half-precision GEMM runs measurably faster on zero-filled operands than on random ones, and the apparent throughput can exceed the card’s dense FP16 peak.

    If a measurement beats the datasheet, suspect the data. Fill operands from a distribution resembling the real ones; values that produce denormals, infinities or NaNs mislead in the other direction.

  3. Watch measurement order. The same algorithm measures faster early in a process than after other benchmarks have warmed the device. A configuration always timed after another is always timed hotter.

  4. Compare against a vendor library’s best configuration: cuBLASLt’s heuristic candidates for a given GEMM can differ substantially in runtime, so comparing against only the first one can overstate your kernel’s advantage. Time several candidates and take the fastest.

  5. For reproducible numbers, fix GPU clocks and power limits.

  6. Do not over-read occupancy: GEMMs hide latency with asynchronous instructions, not warp switching.

  7. Check launch geometry before explaining a gap. Grid size against the device’s SM count and tile count against whole waves come first; only then form hypotheses about the kernel.

  8. Verify every hypothesis about a gap before acting on it or writing it down. Attributing a shortfall to a named cause — a technique the reference uses that you lack, a feature the library does not expose — is a claim with an obvious falsification test: disable that technique on the reference (a switch, a parameter, a one-line edit) and re-measure.

    If the reference does not lose the same margin, the hypothesis is dead and the real cause is still unfound. An untested attribution is worse than no explanation: it ends the search one step before the answer.

    Run such a control once, on one shape, after the kernel is final; do not build control arms into every sweep.

  9. Re-measure only when the compiled binary changed. Host-side cleanups, comment edits and notes do not move a kernel; the SASS or an md5sum of the object tells you whether anything did.

  10. Fusion can favor nominally slower instructions: Ampere-style supermma can preprocess elements in registers, where the fastest MMA classes would force a shared-memory store and fence. Confirming that the fastest instructions (tcgen05, TMA) were emitted checks code quality; it does not guarantee the fastest fused kernel.