Skip to content

Commit 252e6cd

Browse files
chen2021673chen2021673
committed
refactor(sgemv): extract shared SgemvCuda primitive; add sgemv branches in linear kernels
1 parent 97dabe4 commit 252e6cd

5 files changed

Lines changed: 153 additions & 85 deletions

File tree

‎infini_train/include/common/cuda/gemm.cuh‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,28 @@ struct GemmParams {
6969
*/
7070
void GemmCuda(const GemmParams &p);
7171

72+
/**
73+
* Parameter bundle for a single SGEMV call (fp32 only):
74+
* y = alpha * op(A) * x + beta * y
75+
*
76+
* op(A) is m_phys-by-n_phys when trans==N, or n_phys-by-m_phys when trans==T,
77+
* where m_phys and n_phys are the physical (pre-transpose) row/col counts of A.
78+
*/
79+
struct SgemvParams {
80+
cublasOperation_t trans = CUBLAS_OP_N;
81+
int m = 0;
82+
int n = 0;
83+
const float *A = nullptr;
84+
int lda = 0;
85+
const float *x = nullptr;
86+
int incx = 1;
87+
float *y = nullptr;
88+
int incy = 1;
89+
float alpha = 1.0f;
90+
float beta = 0.0f;
91+
cublasHandle_t blas_handle = nullptr;
92+
};
93+
94+
void SgemvCuda(const SgemvParams &p);
95+
7296
} // namespace infini_train::kernels::cuda

‎infini_train/src/kernels/cpu/matmul.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ std::shared_ptr<Tensor> MatmulBackwardInput(const std::shared_ptr<Tensor> &other
7474
}
7575

7676
auto grad_input = std::make_shared<Tensor>(input_dims, DataType::kFLOAT32);
77-
grad_input->Fill<float>(0.0f);
77+
grad_input->Fill(0.0f);
7878

7979
for (int64_t b = 0; b < bs; ++b) {
8080
for (int64_t i = 0; i < m; ++i) {
@@ -116,7 +116,7 @@ std::shared_ptr<Tensor> MatmulBackwardOther(const std::shared_ptr<Tensor> &input
116116
}
117117

118118
auto grad_other = std::make_shared<Tensor>(other_dims, DataType::kFLOAT32);
119-
grad_other->Fill<float>(0.0f);
119+
grad_other->Fill(0.0f);
120120

121121
for (int64_t b = 0; b < bs; ++b) {
122122
for (int64_t i = 0; i < m; ++i) {

‎infini_train/src/kernels/cuda/gemm.cu‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,10 @@ void GemmCuda(const GemmParams &p) {
7070
}
7171
}
7272

73+
void SgemvCuda(const SgemvParams &p) {
74+
DCHECK(p.blas_handle != nullptr);
75+
CUBLAS_CHECK(
76+
cublasSgemv(p.blas_handle, p.trans, p.m, p.n, &p.alpha, p.A, p.lda, p.x, p.incx, &p.beta, p.y, p.incy));
77+
}
78+
7379
} // namespace infini_train::kernels::cuda

‎infini_train/src/kernels/cuda/linear.cu‎

Lines changed: 101 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include "infini_train/include/common/cuda/kernel_helper.cuh"
1212
#include "infini_train/include/dispatcher.h"
1313
#include "infini_train/include/tensor.h"
14+
#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h"
1415

1516
namespace infini_train::kernels::cuda {
1617

@@ -76,40 +77,56 @@ std::shared_ptr<Tensor> LinearForward(const std::shared_ptr<Tensor> &input, cons
7677
output->Fill(0.0);
7778
}
7879

79-
// TODO(zbl): use cublasSgemv if possible for convenience and simplicity
80-
//
81-
// - if a is transposed:
82-
// weight is [out_features, in_features] here
83-
// output = input * weight.T --> output.T = weight * input.T
84-
// C = output.T[out_features, bs]
85-
// A = weight.T[in_features, out_features]
86-
// B = input.T[in_features, bs]
87-
//
88-
// - if a is not transposed:
89-
// output = input * weight --> output.T = weight.T * input.T
90-
// C = output.T[out_features, bs]
91-
// A = weight.T[out_features, in_features]
92-
// B = input.T[in_features, bs]
93-
GemmParams p;
94-
p.trans_a = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
95-
p.trans_b = CUBLAS_OP_N;
96-
p.m = static_cast<int>(out_features);
97-
p.n = static_cast<int>(bs);
98-
p.k = static_cast<int>(in_features);
99-
p.A = weight->DataPtr();
100-
p.lda = static_cast<int>(transpose ? in_features : out_features);
101-
p.B = input->DataPtr();
102-
p.ldb = static_cast<int>(in_features);
103-
p.C = output->DataPtr();
104-
p.ldc = static_cast<int>(out_features);
105-
p.alpha = 1.0f;
106-
p.beta = 1.0f; // bias already written into output; beta=1 accumulates
107-
p.batch_count = 1;
108-
p.input_dtype = dtype;
109-
p.output_dtype = dtype;
110-
p.blas_handle = GetCublasHandle(device);
111-
112-
GemmCuda(p);
80+
// When bs==1 and fp32, use cublasSgemv (more efficient than GEMM for matrix-vector).
81+
// cublasSgemv does not support bf16, so bf16 falls through to GemmCuda.
82+
if (bs == 1 && dtype == DataType::kFLOAT32) {
83+
SgemvParams p;
84+
p.trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
85+
p.m = static_cast<int>(transpose ? in_features : out_features);
86+
p.n = static_cast<int>(transpose ? out_features : in_features);
87+
p.A = static_cast<const float *>(weight->DataPtr());
88+
p.lda = static_cast<int>(transpose ? in_features : out_features);
89+
p.x = static_cast<const float *>(input->DataPtr());
90+
p.y = static_cast<float *>(output->DataPtr());
91+
p.alpha = 1.0f;
92+
p.beta = 1.0f; // output already initialized with bias or zero above
93+
p.blas_handle = GetCublasHandle(device);
94+
SgemvCuda(p);
95+
} else {
96+
// cuBLAS is colmun-major
97+
// - if a is transposed:
98+
// weight is [out_features, in_features] here
99+
// output = input * weight.T --> output.T = weight * input.T
100+
// C = output.T[out_features, bs]
101+
// A = weight.T[in_features, out_features]
102+
// B = input.T[in_features, bs]
103+
//
104+
// - if a is not transposed:
105+
// output = input * weight --> output.T = weight.T * input.T
106+
// C = output.T[out_features, bs]
107+
// A = weight.T[out_features, in_features]
108+
// B = input.T[in_features, bs]
109+
GemmParams p;
110+
p.trans_a = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
111+
p.trans_b = CUBLAS_OP_N;
112+
p.m = static_cast<int>(out_features);
113+
p.n = static_cast<int>(bs);
114+
p.k = static_cast<int>(in_features);
115+
p.A = weight->DataPtr();
116+
p.lda = static_cast<int>(transpose ? in_features : out_features);
117+
p.B = input->DataPtr();
118+
p.ldb = static_cast<int>(in_features);
119+
p.C = output->DataPtr();
120+
p.ldc = static_cast<int>(out_features);
121+
p.alpha = 1.0f;
122+
p.beta = 1.0f; // bias already written into output; beta=1 accumulates
123+
p.batch_count = 1;
124+
p.input_dtype = dtype;
125+
p.output_dtype = dtype;
126+
p.blas_handle = GetCublasHandle(device);
127+
128+
GemmCuda(p);
129+
}
113130

114131
return output;
115132
}
@@ -152,40 +169,56 @@ std::shared_ptr<Tensor> LinearBackwardInput(const std::shared_ptr<Tensor> &weigh
152169
// No Fill(0) needed: cuBLAS beta=0.0f fully overwrites output.
153170
auto grad_input = std::make_shared<Tensor>(input_dims, output_dtype, grad_output->GetDevice());
154171

155-
// TODO(zbl): use cublasSgemv if possible
156-
// - if transpose:
157-
// weight is [out_features, in_features] here
158-
// d_input = d_output * weight --> d_input.T = weight.T * d_output.T
159-
// C = d_input.T[in_features, bs]
160-
// A = weight.T[in_features, out_features]
161-
// B = d_output.T[out_features, bs]
162-
//
163-
// - if not transpose:
164-
// weight is [in_features, out_features] here
165-
// d_input = d_output * weight.T --> d_input.T = weight * d_output.T
166-
// C = d_input.T[in_features, bs]
167-
// A = weight.T[out_features, in_features]
168-
// B = d_output.T[out_features, bs]
169-
GemmParams p;
170-
p.trans_a = transpose ? CUBLAS_OP_N : CUBLAS_OP_T;
171-
p.trans_b = CUBLAS_OP_N;
172-
p.m = static_cast<int>(in_features);
173-
p.n = static_cast<int>(bs);
174-
p.k = static_cast<int>(out_features);
175-
p.A = weight->DataPtr();
176-
p.lda = static_cast<int>(transpose ? in_features : out_features);
177-
p.B = grad_output_promoted->DataPtr();
178-
p.ldb = static_cast<int>(out_features);
179-
p.C = grad_input->DataPtr();
180-
p.ldc = static_cast<int>(in_features);
181-
p.alpha = 1.0f;
182-
p.beta = 0.0f;
183-
p.batch_count = 1;
184-
p.input_dtype = compute_dtype;
185-
p.output_dtype = output_dtype;
186-
p.blas_handle = GetCublasHandle(grad_output->GetDevice());
187-
188-
GemmCuda(p);
172+
// When bs==1 and fp32, use cublasSgemv (more efficient than GEMM for matrix-vector).
173+
// cublasSgemv does not support bf16, so bf16 falls through to GemmCuda.
174+
if (bs == 1 && compute_dtype == DataType::kFLOAT32) {
175+
SgemvParams p;
176+
p.trans = transpose ? CUBLAS_OP_N : CUBLAS_OP_T;
177+
p.m = static_cast<int>(transpose ? in_features : out_features);
178+
p.n = static_cast<int>(transpose ? out_features : in_features);
179+
p.A = static_cast<const float *>(weight->DataPtr());
180+
p.lda = static_cast<int>(transpose ? in_features : out_features);
181+
p.x = static_cast<const float *>(grad_output_promoted->DataPtr());
182+
p.y = static_cast<float *>(grad_input->DataPtr());
183+
p.alpha = 1.0f;
184+
p.beta = 0.0f;
185+
p.blas_handle = GetCublasHandle(grad_output->GetDevice());
186+
SgemvCuda(p);
187+
} else {
188+
// - if transpose:
189+
// weight is [out_features, in_features] here
190+
// d_input = d_output * weight --> d_input.T = weight.T * d_output.T
191+
// C = d_input.T[in_features, bs]
192+
// A = weight.T[in_features, out_features]
193+
// B = d_output.T[out_features, bs]
194+
//
195+
// - if not transpose:
196+
// weight is [in_features, out_features] here
197+
// d_input = d_output * weight.T --> d_input.T = weight * d_output.T
198+
// C = d_input.T[in_features, bs]
199+
// A = weight.T[out_features, in_features]
200+
// B = d_output.T[out_features, bs]
201+
GemmParams p;
202+
p.trans_a = transpose ? CUBLAS_OP_N : CUBLAS_OP_T;
203+
p.trans_b = CUBLAS_OP_N;
204+
p.m = static_cast<int>(in_features);
205+
p.n = static_cast<int>(bs);
206+
p.k = static_cast<int>(out_features);
207+
p.A = weight->DataPtr();
208+
p.lda = static_cast<int>(transpose ? in_features : out_features);
209+
p.B = grad_output_promoted->DataPtr();
210+
p.ldb = static_cast<int>(out_features);
211+
p.C = grad_input->DataPtr();
212+
p.ldc = static_cast<int>(in_features);
213+
p.alpha = 1.0f;
214+
p.beta = 0.0f;
215+
p.batch_count = 1;
216+
p.input_dtype = compute_dtype;
217+
p.output_dtype = output_dtype;
218+
p.blas_handle = GetCublasHandle(grad_output->GetDevice());
219+
220+
GemmCuda(p);
221+
}
189222

190223
return grad_input;
191224
}

‎infini_train/src/kernels/cuda/outer.cu‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -97,24 +97,29 @@ std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> OuterBackward(const
9797

9898
switch (promoted_type) {
9999
case DataType::kFLOAT32: {
100-
// fp32: use cublasSgemv (specialized matrix-vector kernel, more efficient than GEMM for this shape)
101-
// cublasSgemv does not support bf16, so bf16 falls through to GemmCuda below.
102-
float alpha = 1.0f, beta = 0.0f;
103-
cublasHandle_t handle = GetCublasHandle(device);
104-
105100
// grad_input[M] = grad_output[M, N] × other[N]
106-
// y = grad_input[M], A = grad_output.T[N, M], x = other[N]
107-
CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_T, N, M, &alpha,
108-
static_cast<const float *>(grad_output_promoted->DataPtr()), N,
109-
static_cast<const float *>(other_promoted->DataPtr()), 1, &beta,
110-
static_cast<float *>(grad_input->DataPtr()), 1));
101+
SgemvParams p_input;
102+
p_input.trans = CUBLAS_OP_T;
103+
p_input.m = static_cast<int>(N);
104+
p_input.n = static_cast<int>(M);
105+
p_input.A = static_cast<const float *>(grad_output_promoted->DataPtr());
106+
p_input.lda = static_cast<int>(N);
107+
p_input.x = static_cast<const float *>(other_promoted->DataPtr());
108+
p_input.y = static_cast<float *>(grad_input->DataPtr());
109+
p_input.blas_handle = GetCublasHandle(device);
110+
SgemvCuda(p_input);
111111

112112
// grad_other[N] = grad_output.T[N, M] × input[M]
113-
// y = grad_other[N], A = grad_output.T[N, M], x = input[M]
114-
CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_N, N, M, &alpha,
115-
static_cast<const float *>(grad_output_promoted->DataPtr()), N,
116-
static_cast<const float *>(input_promoted->DataPtr()), 1, &beta,
117-
static_cast<float *>(grad_other->DataPtr()), 1));
113+
SgemvParams p_other;
114+
p_other.trans = CUBLAS_OP_N;
115+
p_other.m = static_cast<int>(N);
116+
p_other.n = static_cast<int>(M);
117+
p_other.A = static_cast<const float *>(grad_output_promoted->DataPtr());
118+
p_other.lda = static_cast<int>(N);
119+
p_other.x = static_cast<const float *>(input_promoted->DataPtr());
120+
p_other.y = static_cast<float *>(grad_other->DataPtr());
121+
p_other.blas_handle = GetCublasHandle(device);
122+
SgemvCuda(p_other);
118123
break;
119124
}
120125
case DataType::kBFLOAT16: {

0 commit comments

Comments
 (0)