refactor(gemm): move gemm.cuh/gemm.cu to src/kernels/cuda/common/

include/ is for public-facing interfaces only; gemm primitives are
internal, so relocate them under src/. Update all include paths.
Also rename ctype -> compute_type and add FIXME on bf16 output dtype
promotion hack in linear backward passes.

Author: chen2021673

infini_train/src/kernels/cuda/gemm.cu …ni_train/src/kernels/cuda/common/gemm.cuinfini_train/src/kernels/cuda/gemm.cu renamed to infini_train/src/kernels/cuda/common/gemm.cu infini_train/src/kernels/cuda/gemm.cu renamed to infini_train/src/kernels/cuda/common/gemm.cu

@@ -1,4 +1,4 @@
-#include "infini_train/include/common/cuda/gemm.cuh"
+#include "infini_train/src/kernels/cuda/common/gemm.cuh"
 
 #include <cublas_v2.h>
 
@@ -48,15 +48,15 @@ void GemmCuda(const Device &device, const GemmParams &p) {
     const cudaDataType_t type_c = ToCudaDataType(p.output_dtype);
     // Always use CUBLAS_COMPUTE_32F: required for bf16/fp16 correctness,
     // and fine for fp32 (same compute path).
-    const cublasComputeType_t ctype = CUBLAS_COMPUTE_32F;
+    const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F;
 
     if (p.batch_count == 1) {
         CUBLAS_CHECK(cublasGemmEx(blas_handle, p.trans_a, p.trans_b, p.m, p.n, p.k, &p.alpha, p.A, type_a, p.lda, p.B,
-                                  type_b, p.ldb, &p.beta, p.C, type_c, p.ldc, ctype, CUBLAS_GEMM_DEFAULT));
+                                  type_b, p.ldb, &p.beta, p.C, type_c, p.ldc, compute_type, CUBLAS_GEMM_DEFAULT));
     } else {
         CUBLAS_CHECK(cublasGemmStridedBatchedEx(blas_handle, p.trans_a, p.trans_b, p.m, p.n, p.k, &p.alpha, p.A, type_a,
                                                 p.lda, p.stride_a, p.B, type_b, p.ldb, p.stride_b, &p.beta, p.C, type_c,
-                                                p.ldc, p.stride_c, p.batch_count, ctype, CUBLAS_GEMM_DEFAULT));
+                                                p.ldc, p.stride_c, p.batch_count, compute_type, CUBLAS_GEMM_DEFAULT));
     }
 }
 

infini_train/include/common/cuda/gemm.cuh …i_train/src/kernels/cuda/common/gemm.cuhinfini_train/include/common/cuda/gemm.cuh renamed to infini_train/src/kernels/cuda/common/gemm.cuh infini_train/include/common/cuda/gemm.cuh renamed to infini_train/src/kernels/cuda/common/gemm.cuh

@@ -7,13 +7,13 @@
 #include <cublas_v2.h>
 
 #include "infini_train/include/common/cuda/common_cuda.h"
-#include "infini_train/include/common/cuda/gemm.cuh"
 #include "infini_train/include/common/cuda/kernel_helper.cuh"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dispatcher.h"
 #include "infini_train/include/tensor.h"
 #include "infini_train/src/core/runtime/cuda/cuda_dispatch.h"
 #include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h"
+#include "infini_train/src/kernels/cuda/common/gemm.cuh"
 
 namespace infini_train::kernels::cuda {
 
@@ -165,7 +165,7 @@ std::shared_ptr<Tensor> LinearBackwardInput(const std::shared_ptr<Tensor> &weigh
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // For bf16 compute, accumulate in fp32 to preserve precision.
+    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
     auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
     // No Fill(0) needed: cuBLAS beta=0.0f fully overwrites output.
     auto grad_input = std::make_shared<Tensor>(input_dims, output_dtype, grad_output->GetDevice());
@@ -234,7 +234,7 @@ std::shared_ptr<Tensor> LinearBackwardWeight(const std::shared_ptr<Tensor> &inpu
     auto grad_output_promoted
         = grad_output_dtype == compute_dtype ? grad_output : std::make_shared<Tensor>(grad_output->To(compute_dtype));
 
-    // For bf16 compute, accumulate in fp32 to preserve precision.
+    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
     auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
     const std::vector<int64_t> weight_dims
         = transpose ? std::vector<int64_t>{out_features, in_features} : std::vector<int64_t>{in_features, out_features};
@@ -285,7 +285,7 @@ std::shared_ptr<Tensor> LinearBackwardBias(const std::shared_ptr<Tensor> &grad_o
     const int64_t bs = std::accumulate(dims.rbegin() + 1, dims.rend(), 1, std::multiplies<int64_t>{});
 
     auto compute_dtype = grad_output->Dtype();
-    // For bf16 compute, accumulate in fp32 to preserve precision.
+    // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed.
     auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype;
     auto grad_bias
         = std::make_shared<Tensor>(std::vector<int64_t>{out_features}, output_dtype, grad_output->GetDevice());

infini_train/src/kernels/cuda/linear.cu

@@ -6,9 +6,9 @@
 #include <cublas_v2.h>
 
 #include "infini_train/include/common/cuda/common_cuda.h"
-#include "infini_train/include/common/cuda/gemm.cuh"
 #include "infini_train/include/dispatcher.h"
 #include "infini_train/include/tensor.h"
+#include "infini_train/src/kernels/cuda/common/gemm.cuh"
 
 namespace infini_train::kernels::cuda {
 

infini_train/src/kernels/cuda/matmul.cu

@@ -7,9 +7,9 @@
 #include "glog/logging.h"
 
 #include "infini_train/include/common/cuda/common_cuda.h"
-#include "infini_train/include/common/cuda/gemm.cuh"
 #include "infini_train/include/dispatcher.h"
 #include "infini_train/include/tensor.h"
+#include "infini_train/src/kernels/cuda/common/gemm.cuh"
 
 namespace infini_train::kernels::cuda {