feat(cuda): add resilient cuBLASLt GEMM dispatch#1440
Draft
ppodolsky wants to merge 4 commits into
Draft
Conversation
cublasGemmEx with CUBLAS_GEMM_DEFAULT_TENSOR_OP leaves algorithm selection to the legacy heuristic. cublasLt exposes the real heuristic API, so each GEMM shape now gets a cached execution plan (descriptor, layouts, and the heuristic-selected algorithm) plus a 32 MiB workspace so split-K algorithms are eligible - the same machinery PyTorch's linear uses. Workspaces are per CUDA stream since concurrent matmuls must not share scratch. Shapes the heuristic rejects surface as validation errors, which the autotune arbitration already treats as candidate unavailable.
This was referenced Jul 18, 2026
This was referenced Jul 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This draft consolidates the two CubeCL changes required by the downstream
accelerator stack into one upstream PR:
BufferTooBigserver/client boundary
cuBLASLt
zero-K behavior, and execution-stream ownership before launch
per CUDA stream
The allocation retry was previously proposed separately in #1438. Keeping it
here reflects the actual downstream dependency revision and gives CubeCL one
review thread for this accelerator integration.
Behavior and compatibility
The GEMM primitive is optional, not a global dispatch policy. Its trait method
has a default fail-loud implementation, so other backends need no changes.
Consumers check
features.matmul.accelerated_gemm; Burn's existing autotunerthen compares native GEMM against its fused CubeK candidate per shape.
CUDA advertises BF16 only on SM80+. The implementation accepts row-major
matrices or transposed row-major views with strided batches and input
broadcast. Enqueueing remains asynchronous on the CubeCL stream; validation or
launch failures enter the existing stream error queue.
The allocator cleanup path runs only after the underlying storage reports
BufferTooBig. It rechecks the selected pool after cleanup before attemptingone new device allocation, leaving successful allocations and public APIs
unchanged.
Downstream validation required by the PR template
Validation
cargo fmt --all -- --checkcargo test -p cubecl-runtime allocation_oom_reclaims_cached_pages_and_retriescargo check -p cubecl-cuda --exampleson a non-CUDA hostcublas_gemm_check: padding, binding offsets, all transpose pairs,regular and broadcast batches, asynchronous error propagation, and
three-stream producer/GEMM/consumer ordering
dx/dW/dbparity under fusion,autotune, and autodiff
synchronization
Performance
On an A100-SXM4 40GB, BF16 native dispatch beat or matched PyTorch matmul on
14 of 17 transformer-training shapes. In the downstream 115.8M-parameter
transformer workload, letting Burn's autotuner arbitrate native GEMM against fused
CubeK improved steady-state throughput by 15.6% at batch 16 and 14.8% at batch
20, with matching loss and gradient norms.
The PR remains a draft until the upstream GPU matrix has run and dependency
pins can move to merge revisions.