zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[CUDA] Set bias as input when using bias epilogue (#2584)

Browse Source
This commit is contained in:
Cheng
2025-09-11 15:31:09 +09:00
committed by GitHub
parent d6977f2a57
commit 6a3acf2301
4 changed files with 13 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/backend/cuda/gemms/cublas_gemm.cpp
View File

@@ -230,15 +230,20 @@ void CublasGemm::set_out(
batch_stride);
}
void CublasGemm::set_bias(void* bias) {
void CublasGemm::set_bias(cu::CommandEncoder& encoder, const array& bias) {
encoder.set_input_array(bias);
cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_,
CUBLASLT_MATMUL_DESC_EPILOGUE,
&epilogue,
sizeof(epilogue)));
auto* bias_ptr = bias.data<void>();
CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
matmul_desc_, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
matmul_desc_,
CUBLASLT_MATMUL_DESC_BIAS_POINTER,
&bias_ptr,
sizeof(bias_ptr)));
}
void CublasGemm::run(

2
mlx/backend/cuda/gemms/cublas_gemm.h
View File

@@ -55,7 +55,7 @@ class CublasGemm {
int32_t batch_count,
int64_t batch_stride);
void set_bias(void* bias);
void set_bias(cu::CommandEncoder& encoder, const array& bias);
void run(
cu::CommandEncoder& encoder,