Repository navigation
CUDA backend: optional cuBLAS for matrix multiplications - #24129
shoumikhin wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24129
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 PendingAs of commit 3865daa with merge base 00f0654 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
1229631 to
89f426f
Compare
SigLIP-like blocks (with bias): 0.847 ms with Triton only, 0.813 ms with ATEN,TRITON (4% faster) Not sure if we want to add .cu path for some small gain in some cases. @Gasoonjia scrubbed them recently. Rationale to keep things in triton or triton-template since (1) you can auto-tune etc. and (2) it also helps with ROCM. |
|
Thanks, fair concern. I measured pi0.5 end to end on a Jetson Thor: 127.5 ms per chunk with Triton only and 113.3 ms with |
| * @return Error::Ok on success, Error::InvalidArgument or Error::NotSupported | ||
| * for operands this shim does not handle, Error::Internal on a cuBLAS failure | ||
| */ | ||
| AOTI_SHIM_EXPORT AOTITorchError |
There was a problem hiding this comment.
if adding more shims you will need to update the windows .lib.
There was a problem hiding this comment.
Thanks. These three shims are Linux only for now. They load cuBLASLt with dlopen and are not ported to Windows yet, so CMake does not build them for Windows and the DLL does not have them. If the .lib listed them, a Windows model would link and then fail to load. So I left the .lib as it is, and its notes now say these names must not be listed. Lowering for Windows also rejects ATEN in max_autotune_gemm_backends. With max_autotune off, Inductor calls these shims for matrix multiplications, and the .lib has never listed these three names, so as far as I can tell that case did not link on Windows before this change either. This change does not alter it. I can add Windows support in a follow-up if you want it.
The CUDA backend runs AOTInductor models without libtorch, so every extern call the generated code makes must be provided by the backend's shim library. Matrix multiplications had no shim, which is why the backend only lets Inductor use Triton for them: a model that picked the ATen (cuBLAS) choice failed at export with a missing fallback kernel. This adds aoti_torch_cuda_mm_out, aoti_torch_cuda_addmm_out and aoti_torch_cuda_bmm_out on top of cuBLASLt: - Row-major and column-major operands, sliced leading dimensions, dimensions of size 1 with any stride, and batches that repeat one matrix (zero batch stride). - addmm runs a bias row through the cuBLASLt bias epilogue, as PyTorch does for a linear layer. Any other self that broadcasts to the result (a column, a scalar, a transposed matrix, or a bias row with beta != 1) is copied into the result first and accumulated there with beta. - Each call runs on the current CUDA stream, and the workspace comes from the backend's stream-ordered allocator, so a call can be captured into a CUDA graph. Handles are per thread and device. - float32 runs in full IEEE precision, PyTorch's default. The shim library does not link cuBLASLt. It loads the library the first time one of these shims runs, so models that use only Triton kernels, and apps that ship only the CUDA runtime, still load the backend without it. If the library cannot be loaded, the shim returns an error that names it. The wheel already treats the cuBLAS libraries as provided by torch. Windows is unchanged: the shims load cuBLASLt with dlopen and are not ported to Windows yet, so they and their test are not built there, and the Windows import library does not list them (its notes now say they must not). A Windows model that calls one fails to link at export, as it did before.
Add a max_autotune_gemm_backends compile spec. TRITON stays the default.
ATEN,TRITON lets Inductor's autotuner compare cuBLAS with its Triton
templates for each matrix multiplication and keep the faster one, and ATEN
uses cuBLAS only. A cuBLAS call is a separate kernel that Inductor cannot
fuse with the ops around it, so it is opt-in.
CudaPartitioner([
CudaBackend.generate_method_name_compile_spec("forward"),
CompileSpec("max_autotune_gemm_backends", b"ATEN,TRITON"),
])
ATEN is refused for ROCm and for Windows lowering, which have no cuBLAS
shims.
89f426f to
3865daa
Compare
|
Personally I prefer to not add this because we would love to lean toward triton kernel as much as possible to make the infra shared between nvidai and amd kernels. Also if we have to add this, we shouldn't change the autotune to use Have you tried to add more autotune config to our triton gemm kernel to better fit to Jetson Nano? If there's only like 11% gap i feel like a better config would be enough, since kernel config is majorly focus on large GPU like A100 or 5090. |
|
Thanks both, that makes sense. I'll close this and try tuning the Triton GEMM autotune configs for Jetson instead, then come back with numbers in a smaller PR if that closes the gap. |
What is wrong today
The CUDA backend runs AOTInductor models without libtorch. Every call the generated code makes outside its own kernels needs a shim in the backend, and matrix multiplications have none. So the backend forces Inductor to use Triton for every matrix multiplication, even for shapes where cuBLAS is faster.
What this changes
mm,addmmandbmm, in float32, float16 and bfloat16. They handle row-major and column-major operands, sliced leading dimensions, dimensions of size 1 with any stride, batches that repeat one matrix, and anyaddmminput that broadcasts to the result. A linear layer's bias runs through the cuBLASLt bias epilogue. Each call runs on the current CUDA stream, and the workspace comes from the backend's stream-ordered allocator, so the shims work under CUDA graph capture.With
ATEN,TRITON, the autotuner times cuBLAS against the Triton templates for each matrix multiplication and keeps the faster one.It is opt-in because it is not a win everywhere: a cuBLAS call cannot be fused with the ops around it. On an A100, bfloat16, four transformer-style blocks, time per call:
ATEN,TRITON(4% faster)ATEN,TRITON(2% slower)Accuracy is the same in both cases.
On a Jetson AGX Thor, a full pi0.5 robot policy (a vision-language-action model) takes 127.5 ms per action chunk with Triton only and 113.3 ms with
ATEN,TRITON, which is 11% faster. Its mean absolute difference from a float32 reference is about the same in both (0.50 with Triton only, 0.53 withATEN,TRITON).The shim library does not link cuBLASLt. It loads it the first time a cuBLAS shim runs, so a model that uses only Triton still loads without it. In Python, torch has already loaded it. A C++ app that runs a model using these shims needs cuBLASLt on its library search path.
The autotuner times PyTorch's own cuBLAS calls, not these shims, so for a given shape the shim's speed can differ from what was timed. The shims also compute float32 in full precision, while with TF32 enabled during lowering the autotuner times cuBLAS in TF32. Either way it can pick cuBLAS where Triton would be faster at run time. This only affects speed, and AOTInductor with libtorch has the same TF32 gap.
ATENis refused for ROCm and for Windows lowering. The shims are Linux only for now: they load cuBLASLt withdlopenand are not ported to Windows yet, so they are not built there, and the Windows import library does not list them. Its notes now say they must not be listed.Limits:
baddbmm,_int_mmand theout_dtypevariants have no shim yet. If autotuning picks one of them, export fails with the existing "missing fallback kernels" error. A matrix multiplication with an empty inner dimension returns an error at run time. Also, on Linux, a model exported withmax_autotuneoff used to fail on any matrix multiplication and now runs it on cuBLAS, in float32, float16 or bfloat16. On Windows it still fails to link, as before.Not in this PR: the backend copies every delegate output on every execute, because AOTInductor always returns its own output buffers. That will be a separate follow-up.
Tested
On an H100 (CUDA 12.8), at this head:
addmminputs that are not a plain bias row (a column, a scalar, a full or transposed matrix, an expanded row, a bias row with beta != 1). The 11 existing shim tests pass too.test_cublas_gemm.py(9 tests): exports a model withATEN, and again withmax_autotuneoff, checks that the generated code calls the three shims (twoaddmmcalls: a bias row and a column), runs it on the ExecuTorch runtime, and compares with eager PyTorch. It also covers the compile spec and its three refusals (an unknown backend, ROCm, Windows).The timings above come from an earlier revision on an A100 (CUDA 13.2) and a Jetson AGX Thor. Since then, only the handling of an
addmminput that is not a bias row and the alignment hint given to cuBLASLt's algorithm search have changed. With the new hint, cuBLASLt picked the same algorithm in every case measured.