cuBLASLt

MatmulPlan(; M, N, K, typeA, typeB, typeD, kwargs...)

An execution plan for D = α ⋅ op(A) ⋅ op(B) + β ⋅ C, where op(A) is M × K, op(B) is K × N, and C/D are M × N (column-major).

The type parameter is the element type of α/β (Float32, or Float16/ Float64 for the :f16/:f64 compute types).

Keywords

• M, N, K: shapes of op(A) ⋅ op(B) = D (required).
• typeA, typeB, typeD, typeC = typeD: Julia types or cudaDataTypes.
• transA, transB: 'N', 'T', or 'C' (default 'N').
• lda, ldb, ldc, ldd: leading dimensions of the stored matrices.
• batch = 1: strided-batch count; strideA/strideB/strideC/strideD are element strides between batch entries (default: densely packed).
• compute = :f32: one of :f32, :tf32, :fastf16, :fastbf16, :bf16x9, :f16, :f64.
• scale_modeA, scale_modeB: scale mode (affects algorithm choice); :none or one of :scalarf32, :vec32ue8m0, :vec16ue4m3, :outervecf32, :vec128f32, :blk128x128f32, :perbatchscalarf32. The scale arrays are apply-time arguments (scaleA/scaleB).
• pointer_mode = :host: :host (α/β are Numbers) or :device (α/β are 0-dimensional CuArrays, read at execution time; graph-capture safe).
• alignA, alignB, alignC, alignD (default 256): minimum byte alignment the corresponding operand pointer is guaranteed to have at call time, as a power of two. Alignment gates algorithm validity — pass the true alignment when operands will be views into larger arrays.
• epilogue = :none: fused epilogue, one of :relu, :reluaux, :bias, :relubias, :reluauxbias, :drelu, :drelubgrad, :gelu, :geluaux, :gelubias, :geluauxbias, :dgelu, :dgelubgrad, :bgrada, :bgradb. Forward: D = act(α ⋅ op(A) ⋅ op(B) + β ⋅ C + bias), with *_aux variants stashing the activation input (GELU: values; ReLU: a sign bit-mask) for the backward pass. Backward: :drelu/:dgelu multiply the matmul result elementwise by act′(aux); the *_bgrad variants additionally reduce the result into a bias-gradient vector, and :bgrada/:bgradb reduce operand A/B over K into one. The bias/bgrad/aux pointers are apply-time arguments.
• bias_type: element type of the bias/bgrad vector (defaults to cuBLASLt's rule: D's type, or BF16/FP16 for FP8 matmuls).
• aux_type, aux_ld, aux_stride: aux-buffer element type (GELU only), leading dimension (elements; bits for the ReLU bit-mask, divisible by 128), and batch stride.
• scale_modeC, scale_modeD: like scale_modeA/B, for the C input and the D output (quantization scale); scaleC/scaleD are apply-time.
• out_scale_modeD: block-scale mode of the kernel-computed output scales (MXFP8/NVFP4 output); the out_scaleD array they are written to is apply-time. Requires cuBLASLt ≥ 12.9, and only supported inside a block-scaled matmul (A/B in matching block modes) — the heuristic rejects it on unscaled or tensor-wide-scaled inputs.
• fast_accum = false: FP8 fast accumulation.
• max_workspace = 32 << 20: workspace-size ceiling for the heuristic; the chosen algorithm's actual requirement is stored as plan.workspace_size.

(plan::MatmulPlan)(D, A, B; α = true, β = false, C = D,
                   scaleA = nothing, scaleB = nothing,
                   scaleC = nothing, scaleD = nothing,
                   out_scaleD = nothing, amaxD = nothing,
                   bias = nothing, bgrad = nothing, aux = nothing,
                   workspace = nothing)

Compute D = α ⋅ op(A) ⋅ op(B) + β ⋅ C according to plan, in place.

Operands are anything cuBLASLt.ltdata/cuBLASLt.ltptr/ cuBLASLt.ltstride accept; Transpose/Adjoint/PermutedDimsArray wrappers on A/B (and NNlib's batched wrappers, with the NNlib extension loaded) are unwrapped and checked against the plan's orientation, raw arrays (and identity permutations) are trusted as the stored matrices. α/β are host Numbers (pointer_mode = :host) or 0-dimensional CuArrays (pointer_mode = :device). scaleA/scaleB are device scale arrays, required iff the plan's corresponding scale mode is not :none and the operand does not carry its own (see cuBLASLt.ltscale). For the block modes (:vec32_ue8m0, :vec16_ue4m3) cuBLASLt reads scales through its tiled layout of 128(outer)×4(inner)-entry tiles: data must be swizzled accordingly and the allocation padded to whole tiles, or the kernel reads out of bounds. Epilogue arguments (mandatory iff the plan's epilogue uses them): bias is the input vector of length M for the forward *_bias epilogues; bgrad the output vector the gradient epilogues reduce into (length M, or N for :bgradb); aux the activation-input stash — written by the forward *_aux epilogues, read by the backward d* ones (GELU: a matrix of pre-activations; ReLU: a bit-mask; both sides of a forward/backward pair must agree on its dtype/ld/stride, which nothing checks for you).

Output quantization: scaleD (with scale_modeD) is the quantization scale applied to a narrow D; out_scaleD (with out_scale_modeD) the array the kernel writes computed block scales to (or carried by a scale-bearing D, see cuBLASLt.ltscale); amaxD an optional 1-element Float32 device array the pre-quantization absmax of D is written to; scaleC dequantizes a narrow C.

workspace is an optional preallocated buffer of at least plan.workspace_size bytes, owned by the caller: nothing is allocated, nothing is freed. By default one is allocated from the stream-ordered pool and freed right after the launch.

Under graph capture, always pass workspace explicitly. The pool default allocates on every apply, and CUDA's allocator may synchronize to reclaim — illegal inside a capture. The plan is the seam that makes this easy: it knows its own workspace_size before any launch, so a caller carves that many bytes from wherever it keeps memory (an arena, a slab, a reused buffer) and hands them to the apply, which is then allocation-free.

plan = plan_matmul(D, A, B)
ws = my_allocator(plan.workspace_size)
plan(D, A, B; workspace = ws)     # no allocation, capture-safe

cuBLASLt.UnsupportedConfigError

The heuristic found no algorithm for a configuration: this device, this cuBLASLt, these types/shapes/modes. Distinct from an ArgumentError — those say a caller spelled something wrong, this one says the library cannot serve a well-formed request, which is a routing fact, not a bug.

Consequently it is the one exception the plan cache memoizes (a negative result is as worth caching as a positive one; the heuristic query is not cheap), and the one cuBLASLt.matmul_supported answers false to instead of throwing.

cuBLASLt.activation_symbol(f) -> Symbol or nothing

The epilogue activation f stands for: :relu or :gelu (cuBLASLt's GELU is the tanh approximation) — those two are the entire menu the library fuses. Returns nothing for anything unfusable, so layer code can branch on fusability (the branch constant-folds: this is dispatch on typeof(f)) instead of catching exceptions; passing an unfusable activation = to plan_matmul/matmul! is what throws. Symbols pass through, identity means :none. Register a function with

cuBLASLt.activation_symbol(::typeof(f)) = :relu   # or :gelu

so callers can pass activation = f (e.g. NNlib.relu).

cuBLASLt.empty_plan_cache!()

Empty the cache used by the planless matmul! methods (benchmarking hygiene; freed plans are reclaimed by the GC).

cuBLASLt.handle() -> cublasLtHandle_t

The cuBLASLt handle for the current CUDA context.

cuBLASLt.ltdata(x) -> storage array

The storage array behind an operand — the thing whose eltype and size describe what cuBLASLt reads. Identity by default; block-scaled container types overload this to hand over their payload without this package knowing their layout. Pointer and stride customization live in cuBLASLt.ltptr and cuBLASLt.ltstride.

cuBLASLt.ltptr(x) -> CuPtr{Cvoid}

The device pointer cuBLASLt reads an operand's storage through. Defaults to reinterpret(CuPtr{Cvoid}, pointer(x)).

cuBLASLt.ltscale(x) -> device scale array or nothing

The device scale array an operand carries, nothing by default. Operand types that bundle their scales overload this so plan_matmul and plan application extract scale pointers without explicit scaleA/scaleB arguments (which, when passed, take precedence).

cuBLASLt.ltstride(x, d) -> Integer

The distance between adjacent elements along dimension d, in logical elements. Defaults to stride(x, d). Packed-storage owners overload this when their logical element stride differs from their physical container.

matmul!(D, A, B; α = true, β = false, C = D,
        scaleA = nothing, scaleB = nothing, scaleC = nothing,
        scaleD = nothing, out_scaleD = nothing, amaxD = nothing,
        activation = nothing, bias = nothing, bgrad = nothing,
        aux = nothing, workspace = nothing, kws...)

Planless convenience: plan_matmul through the plan cache, then apply. Same signature as both; kws override derived plan kwargs (e.g. compute = :tf32, transA = 'T', epilogue = :dgelu_bgrad).

matmul_supported(D, A, B; kws...) -> Bool

Whether cuBLASLt can serve this matmul! — same signature, no launch. Value-returning: an unsupported configuration answers false rather than throwing, so a router can pick between cuBLASLt and a fallback without try/catch.

The probe builds the plan and leaves it in the cache, so a true answer costs the following matmul! nothing, and a false answer is cached too (the heuristic query is not cheap, and "no algorithm exists" is a stable fact about this device and this library).

Malformed arguments still throw: false means cuBLASLt has no algorithm, not you called this wrong.

matmul_supported(D, A, B; compute = :tf32) ? matmul!(D, A, B; compute = :tf32) :
                                             mul!(D, A, B)

plan_candidates(D, A, B; count = 8, kws...) -> Vector{MatmulPlan}
plan_candidates(; count = 8, M, N, K, typeA, typeB, typeD, kws...) -> Vector{MatmulPlan}

The algorithm escape hatch: like plan_matmul (or the MatmulPlan kwargs constructor in the second form), but returns up to count candidate plans, heuristic-ranked. Each is independently callable — benchmark them and keep the winner; that is the canonical cuBLASLt autotuning loop, with the plan cache ownership on the caller's side.

plan_matmul(D, A, B; α = true, β = false, C = D,
            scaleA = nothing, scaleB = nothing, workspace = nothing,
            kws...) -> MatmulPlan

Build a MatmulPlan for D = α ⋅ op(A) ⋅ op(B) + β ⋅ C from prototype arguments — the same signature the returned plan is applied with, so every call-time argument doubles as the prototype its plan-time counterpart is derived from:

• shapes, element types, leading dimensions, batching, and pointer alignments come from the arrays (via cuBLASLt.ltdata);
• Transpose/Adjoint wrappers on A/B set transA/transB, as does PermutedDimsArray with the first two dims swapped — (2,1) or, for a batched transpose, (2,1,3) — which sets 'T' (identity permutations are accepted as 'N'; anything else throws). With the NNlib extension loaded, NNlib.batched_transpose sets 'T' and NNlib.batched_adjoint 'C' the same way;
• the types of α/β set pointer_mode (Numbers → :host, 0-dimensional CuArrays → :device);
• scaleA/scaleB/scaleC/scaleD arrays (or scales carried by the operands, see cuBLASLt.ltscale and cuBLASLt.scale_mode) set the scale modes; out_scaleD (or a scale-carrying D) sets out_scale_modeD — kernel-computed output block scales for quantized output;
• activation (a Symbol or a function registered via cuBLASLt.activation_symbol) composed with the presence of bias/aux prototypes sets the epilogue; the gradient epilogues (:drelu_bgrad, :dgelu_bgrad, :bgrada, :bgradb) are spelled with an explicit epilogue kwarg. Bias/aux prototypes set bias_type, aux_type, aux_ld, and the batch strides;
• workspace sets max_workspace = sizeof(workspace), guaranteeing the chosen algorithm fits the buffer.

Remaining kws override anything derived and are passed through to MatmulPlan (e.g. compute, scale_modeA, epilogue, max_workspace).

Alignment caveat: the plan promises the heuristic the prototypes' pointer alignments, and application checks that later operands keep the promise — plan with your worst-aligned representative.

cuBLASLt.scale_mode(x) -> Symbol

The scale mode implied by an operand's type, :none by default. Block-scaled container types overload this (block shape + scale eltype → one of :scalarf32, :vec32ue8m0, :vec16ue4m3, :outervecf32, :vec128f32, :blk128x128f32, :perbatchscalarf32).

cuBLASLt.version() -> VersionNumber

The version of the loaded libcublasLt.