Skip to content

feat(cuda): add resilient cuBLASLt GEMM dispatch#1440

Draft
ppodolsky wants to merge 4 commits into
tracel-ai:mainfrom
ppodolsky:feat/cublaslt-gemm-dispatch
Draft

feat(cuda): add resilient cuBLASLt GEMM dispatch#1440
ppodolsky wants to merge 4 commits into
tracel-ai:mainfrom
ppodolsky:feat/cublaslt-gemm-dispatch

Conversation

@ppodolsky

@ppodolsky ppodolsky commented Jul 18, 2026

Copy link
Copy Markdown

Summary

This draft consolidates the two CubeCL changes required by the downstream
accelerator stack into one upstream PR:

  • reclaim unused cached pages and retry once when a device allocation reports
    BufferTooBig
  • add a backend-optional accelerated-GEMM descriptor and capability flag at the
    server/client boundary
  • implement asynchronous CUDA BF16 GEMM with FP32 accumulation through
    cuBLASLt
  • validate layouts, dimensions, batch strides, binding ranges, overlap,
    zero-K behavior, and execution-stream ownership before launch
  • cache cuBLASLt descriptors and algorithms by shape, with separate workspaces
    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 autotuner
then 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 attempting
one new device allocation, leaving successful allocations and public APIs
unchanged.

Downstream validation required by the PR template

Validation

  • cargo fmt --all -- --check
  • cargo test -p cubecl-runtime allocation_oom_reclaims_cached_pages_and_retries
  • cargo check -p cubecl-cuda --examples on a non-CUDA host
  • A100 cublas_gemm_check: padding, binding offsets, all transpose pairs,
    regular and broadcast batches, asynchronous error propagation, and
    three-stream producer/GEMM/consumer ordering
  • downstream Burn CUDA linear forward and dx/dW/db parity under fusion,
    autotune, and autodiff
  • 10,000 queued GEMMs followed by one drain confirmed no per-call host
    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.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant