Skip to content

Move bmm calls to stable cshim - #220

Open
asglover wants to merge 18 commits into
mainfrom
move-bmm-calls-to-stable-cshim
Open

asglover wants to merge 18 commits into
mainfrom
move-bmm-calls-to-stable-cshim

Conversation

@asglover

Copy link
Copy Markdown
Collaborator

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.

@asglover asglover mentioned this pull request Sep 14, 2026
@asglover
asglover marked this pull request as ready for review September 15, 2026 04:12
@asglover asglover added the ci-ready Triggers CI checks for a pull request label Sep 15, 2026

@vbharadwaj-bk vbharadwaj-bk left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remind me why this DEFINE is here?

explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.get_device_index()) {}
};

AtenTensorHandle tensor_handle(Tensor& tensor) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@vbharadwaj-bk

Copy link
Copy Markdown
Member

And this is a great change, thanks - good to get the dependence on BLAS eliminated.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci-ready Triggers CI checks for a pull request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants