Skip to content

[CUDA] Use GEMM with epilogue instead of AddMM #2569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zcbenz merged 1 commit into from
Sep 9, 2025
Merged

[CUDA] Use GEMM with epilogue instead of AddMM #2569

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 33 additions & 13 deletions mlx/backend/cuda/gemms/cublas_gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,10 +85,10 @@ cublasLtMatrixLayout_t create_matrix_layout(
int32_t batch_count,
int64_t batch_stride) {
cublasLtMatrixLayout_t desc;
if (transposed) {
std::swap(rows, cols);
}
CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutCreate(&desc, type, rows, cols, ld));
cublasLtOrder_t order = transposed ? CUBLASLT_ORDER_COL : CUBLASLT_ORDER_ROW;
CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
desc, CUBLASLT_MATRIX_LAYOUT_ORDER, &order, sizeof(cublasLtOrder_t)));
if (batch_count > 1) {
CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
desc,
Expand Down Expand Up @@ -138,25 +138,34 @@ CublasGemm::CublasGemm(
CUBLASLT_MATMUL_DESC_POINTER_MODE,
&pointer_mode,
sizeof(int32_t)));
cublasOperation_t op = CUBLAS_OP_N;

// In cublasLt matrices use column-major layout, while it is possible to use
// the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
// epilogue does not work with the option. So instead we swap A and B to make
// cublasLt return the row-major result, which works because:
// - the data of a matrix in row-major layout is identical to its transpose in
// column-major layout
// - C^T = (A @ B)^T = B^T @ A^T
cublasOperation_t a_op = b_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_,
CUBLASLT_MATMUL_DESC_TRANSA,
&op,
&a_op,
sizeof(cublasOperation_t)));
cublasOperation_t b_op = a_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_,
CUBLASLT_MATMUL_DESC_TRANSB,
&op,
&b_op,
sizeof(cublasOperation_t)));

auto type = dtype_to_cublas_type(dtype);
a_desc_ = create_matrix_layout(
type, a_rows, a_cols, a_transposed, lda, batch_count, a_batch_stride);
type, b_cols, b_rows, b_transposed, ldb, batch_count, b_batch_stride);
b_desc_ = create_matrix_layout(
type, b_rows, b_cols, b_transposed, ldb, batch_count, b_batch_stride);
type, a_cols, a_rows, a_transposed, lda, batch_count, a_batch_stride);
out_desc_ = create_matrix_layout(
type, a_rows, b_cols, false, b_cols, batch_count, a_rows * b_cols);
type, b_cols, a_rows, false, b_cols, batch_count, a_rows * b_cols);
}

CublasGemm::CublasGemm(
Expand Down Expand Up @@ -191,7 +200,7 @@ CublasGemm::CublasGemm(
b_batch_stride) {
auto type = dtype_to_cublas_type(dtype);
c_desc_ = create_matrix_layout(
type, a_rows, b_cols, false, ldc, batch_count, c_batch_stride);
type, b_cols, a_rows, false, ldc, batch_count, c_batch_stride);
}

CublasGemm::~CublasGemm() {
Expand All @@ -213,14 +222,25 @@ void CublasGemm::set_out(
CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
out_desc_ = create_matrix_layout(
dtype_to_cublas_type(dtype),
rows,
cols,
rows,
transposed,
ld,
batch_count,
batch_stride);
}

void CublasGemm::set_bias(void* bias) {
cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_,
CUBLASLT_MATMUL_DESC_EPILOGUE,
&epilogue,
sizeof(epilogue)));
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
}

void CublasGemm::run(
cu::CommandEncoder& encoder,
array& out,
Expand Down Expand Up @@ -330,9 +350,9 @@ void CublasGemm::execute(
handle_,
matmul_desc_,
&alpha,
a,
b, // a and b are swapped
a_desc_,
b,
a,
b_desc_,
&beta,
c ? c : out,
Expand Down
2 changes: 2 additions & 0 deletions mlx/backend/cuda/gemms/cublas_gemm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ class CublasGemm {
int32_t batch_count,
int64_t batch_stride);

void set_bias(void* bias);

void run(
cu::CommandEncoder& encoder,
array& out,
Expand Down
113 changes: 75 additions & 38 deletions mlx/backend/cuda/matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <numeric>

namespace mlx::core {

namespace {

std::tuple<bool, int64_t, array>
Expand All @@ -28,41 +29,20 @@ check_transpose(cu::CommandEncoder& enc, const Stream& s, const array& arr) {
}
}

} // namespace

void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
nvtx3::scoped_range r("Matmul::eval_gpu");
auto& s = stream();
auto& encoder = cu::get_command_encoder(s);

assert(inputs.size() == 2);
auto& a_pre = inputs[0];
auto& b_pre = inputs[1];
// Return 0s if either input is empty.
if (a_pre.size() == 0 || b_pre.size() == 0) {
array zero(0, a_pre.dtype());
encoder.add_temporary(zero);
fill_gpu(zero, out, s);
return;
}

out.set_data(allocator::malloc(out.nbytes()));

/////////////////////////////////////////////////////////////////////////////
// Init checks and prep

int M = a_pre.shape(-2);
int N = b_pre.shape(-1);
int K = a_pre.shape(-1);

// Keep a vector with copies to be cleared in the completed buffer to release
// the arrays
auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);

/////////////////////////////////////////////////////////////////////////////
void gemm_and_bias(
cu::CommandEncoder& encoder,
int M,
int N,
int K,
bool a_transposed,
int64_t lda,
bool b_transposed,
int64_t ldb,
array& out,
const array& a,
const array& b,
void* bias = nullptr) {
// Check and collapse batch dimensions

auto [batch_shape, a_batch_strides, b_batch_strides] = collapse_batches(a, b);

auto batch_count = out.size() / (M * N);
Expand All @@ -79,7 +59,8 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
batch_shape = {1};
}

if (cu::can_use_gemv(M, N, K, a_transposed, b_transposed)) {
// Use gemmv when possible
if (!bias && cu::can_use_gemv(M, N, K, a_transposed, b_transposed)) {

@awni awni Sep 8, 2025

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.

Am I missing something here or is the return statement removed a bug?

It looks like it will encode the gemv and the full gemm right now?

@zcbenz zcbenz Sep 9, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

It was a mistake, thanks for noticing it!

cu::gemv(
a,
b,
Expand All @@ -95,10 +76,9 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
return;
}

/////////////////////////////////////////////////////////////////////////////
// Invoke cublasLt
CublasGemm gemm(
cu::device(s.device),
encoder.device(),
a.dtype(),
a_transposed,
M,
Expand All @@ -111,9 +91,45 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
batch_shape.back(),
a_batch_strides.back(),
b_batch_strides.back());
if (bias) {
gemm.set_bias(bias);
}
gemm.run(encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides);
}

} // namespace

void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
nvtx3::scoped_range r("Matmul::eval_gpu");
auto& s = stream();
auto& encoder = cu::get_command_encoder(s);

assert(inputs.size() == 2);
auto& a_pre = inputs[0];
auto& b_pre = inputs[1];
// Return 0s if either input is empty.
if (a_pre.size() == 0 || b_pre.size() == 0) {
array zero(0, a_pre.dtype());
encoder.add_temporary(zero);
fill_gpu(zero, out, s);
return;
}

out.set_data(allocator::malloc(out.nbytes()));

int M = a_pre.shape(-2);
int N = b_pre.shape(-1);
int K = a_pre.shape(-1);

// Keep a vector with copies to be cleared in the completed buffer to release
// the arrays
auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);

gemm_and_bias(
encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b);
}

void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
nvtx3::scoped_range r("AddMM::eval_gpu");
auto& s = stream();
Expand All @@ -136,6 +152,27 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);

/////////////////////////////////////////////////////////////////////////////
// Dispatch to GEMM with epilogue or AddMM

if (beta_ == 1 && c.strides(-1) == 1 && c.data_size() == out.shape(-1)) {
out.set_data(allocator::malloc(out.nbytes()));
gemm_and_bias(
encoder,
M,
N,
K,
a_transposed,
lda,
b_transposed,
ldb,
out,
a,
b,
c.data<void>());
return;
Comment on lines +158 to +173

@awni awni Sep 8, 2025

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.

Can we simply setup the right cublas matmul without adding a set_bias by checking ldc?

@zcbenz zcbenz Sep 9, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

The epilogue bias requires the input to be a vector of size out.shape(-1), while the normal addmm interface requires a matrix of same shape with out. So we can't really encapsulate the dispatching code inside CublasGemm because the shape of c is different, unless we move code processing c (e.g. copy_gpu(c, out) and collapse_batches(a, b, c)) into the class, which I think decreases code readability.

}

int64_t ldc;
{
auto stx = c.strides()[c.ndim() - 2];
Expand Down Expand Up @@ -177,7 +214,7 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
}

/////////////////////////////////////////////////////////////////////////////
// Invoke cublasLt
// Invoke cublasLt with AddMM settings

CublasGemm gemm(
cu::device(s.device),
Expand Down