Motivation
CubeCL matmuls currently route through generated CubeK kernels. On CUDA, regular BF16 GEMMs can also use cuBLAS/cuBLASLt, whose hardware-specific algorithm selection is materially faster for some transformer projection shapes. A downstream Burn/Hermes integration measured roughly 15% higher end-to-end training throughput when an autotuner could arbitrate the generated kernel against a backend-native GEMM.
Proposed shape
- add a backend-optional
ComputeServer::gemm enqueue operation and an accelerated_gemm element-type capability
- describe logical row-major, transposed-view, strided-batch, and broadcast operands with a small validated descriptor
- implement CUDA BF16 support through cuBLASLt with FP32 accumulation
- enqueue on the CubeCL stream, resolve cross-stream dependencies through the normal server machinery, and report asynchronous validation/runtime failures through the stream error queue
- cache cuBLASLt plans by shape/stride/transpose and keep one 32 MiB workspace per CUDA stream
The operation would be an optional backend primitive rather than a replacement for CubeK: callers check capability, and Burn's existing autotune layer decides which implementation wins for each shape.
Questions
- Does a backend-native GEMM hook belong at the
ComputeServer/ComputeClient layer, or would maintainers prefer it behind a different extension interface?
- Is a set of accelerated element types sufficient capability metadata, or should supported layouts/accumulation types be represented too?
- Should workspace allocation use CUDA directly or be integrated with CubeCL's managed allocation accounting?
I have a working downstream-proven implementation and will link a draft PR so the concrete API and validation contracts are reviewable before treating it as ready.
Motivation
CubeCL matmuls currently route through generated CubeK kernels. On CUDA, regular BF16 GEMMs can also use cuBLAS/cuBLASLt, whose hardware-specific algorithm selection is materially faster for some transformer projection shapes. A downstream Burn/Hermes integration measured roughly 15% higher end-to-end training throughput when an autotuner could arbitrate the generated kernel against a backend-native GEMM.
Proposed shape
ComputeServer::gemmenqueue operation and anaccelerated_gemmelement-type capabilityThe operation would be an optional backend primitive rather than a replacement for CubeK: callers check capability, and Burn's existing autotune layer decides which implementation wins for each shape.
Questions
ComputeServer/ComputeClientlayer, or would maintainers prefer it behind a different extension interface?I have a working downstream-proven implementation and will link a draft PR so the concrete API and validation contracts are reviewable before treating it as ready.