Conversation
…est it through symmetric tensor product. The test that we don't rely on blas is over testing.
vbharadwaj-bk
requested changes
Sep 15, 2026
vbharadwaj-bk
left a comment
Member
There was a problem hiding this comment.
Code looks fine, just one or two questions. Perf-wise, do you know how this compares to the original CUDA version?
| elif torch.version.hip: | ||
| torch_libs = library_paths("cuda")[0] | ||
| extra_link_args.append("-Wl,-rpath," + torch_libs) | ||
| extra_link_args.extend("-L" + path for path in library_paths("cuda")) |
Member
There was a problem hiding this comment.
Hm - library_paths("cuda") on the hip branch looks weird, but I'm hoping this does the appropriate thing on AMD (that CUDA is just being used as a standin for hip.
| @@ -1,8 +1,15 @@ | |||
| #define USE_CUDA | |||
Member
There was a problem hiding this comment.
Remind me why this DEFINE is here?
| explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.get_device_index()) {} | ||
| }; | ||
|
|
||
| AtenTensorHandle tensor_handle(Tensor& tensor) { |
Member
There was a problem hiding this comment.
Hmm - actually what is the purpose of this namespace and the wrapper for tensor.get(), not 100% sure what's going on here (namespace collision?). Might be necessary but I don't understand it atm
Member
|
And this is a great change, thanks - good to get the dependence on BLAS eliminated. |
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.
This is a simple PR which removes our dependency on BLAS. Instead we move the existing calls through a torch stable ABI c cuda shim which provides access to in place bmm. This makes torch responsible for managing the different versions of BLAS and providing a consistent interface.
I could not use an aten level op because there is not an out of place stable matmul at the c++ api level. I have requested one.
pytorch/pytorch#196927
This is not a complete mitigation of the problem. By using a cuda shim, we don't have a solution for sycl / other platforms. Ideally we would just use aten ops at not the c++ level. This is a little complicated because we'd have to plumb the information about how to call bmm into that separate op / the aten graph level. Also we might potentially prefer to just do a grouped gemm, the challenge is that grouped gemm's aren't stable at the torch level.
Also this op, symmetric is only support on torch.
The real path forward for improving the support is to move the grouped gemm behavior directly into a custom kernel. I think that is possible. so I'm looking forward to doing that. But I'm looking to simply implement an improvement over us linking to blas. The reason to not link BLAS directly is because it doesn't play with multi device, see:
#206
This should unblock that.