wehub-resource-sync
365 lines
15 KiB
Plaintext
365 lines
15 KiB
Plaintext
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <cublasLt.h>
|
|
#include <cublas_v2.h>
|
|
#include <cuda.h>
|
|
#include <cuda_fp16.h>
|
|
#include <cuda_fp8.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_runtime_api.h>
|
|
#include <stdlib.h>
|
|
#include <cassert>
|
|
#include <cstdint>
|
|
|
|
#include "paddle/phi/backends/dynload/cublasLt.h"
|
|
#include "paddle/phi/backends/gpu/gpu_info.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
|
|
#include "paddle/phi/api/include/context_pool.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/allocator.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blaslt_gemm_search.h"
|
|
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
|
|
namespace phi {
|
|
|
|
namespace {
|
|
|
|
// Helper function to check if dtype is FP8
|
|
bool IsFp8Dtype(phi::DataType dtype) {
|
|
return dtype == phi::DataType::FLOAT8_E4M3FN ||
|
|
dtype == phi::DataType::FLOAT8_E5M2;
|
|
}
|
|
|
|
// Convert phi::DataType to cudaDataType_t
|
|
cudaDataType_t ScalarTypeToCudaDataType(phi::DataType dtype) {
|
|
switch (dtype) {
|
|
case phi::DataType::FLOAT8_E4M3FN:
|
|
return CUDA_R_8F_E4M3;
|
|
case phi::DataType::FLOAT8_E5M2:
|
|
return CUDA_R_8F_E5M2;
|
|
case phi::DataType::BFLOAT16:
|
|
return CUDA_R_16BF;
|
|
case phi::DataType::FLOAT32:
|
|
return CUDA_R_32F;
|
|
case phi::DataType::FLOAT16:
|
|
return CUDA_R_16F;
|
|
default:
|
|
PADDLE_THROW(common::errors::InvalidArgument("Unsupported data type"));
|
|
}
|
|
}
|
|
|
|
// cuBLAS error checking macro
|
|
#define PADDLE_CUDABLAS_CHECK(func) \
|
|
do { \
|
|
cublasStatus_t status = func; \
|
|
if (status != CUBLAS_STATUS_SUCCESS) { \
|
|
PADDLE_THROW(common::errors::External("cuBLAS error: %d", status)); \
|
|
} \
|
|
} while (0)
|
|
|
|
template <typename Context>
|
|
void cublas_gemm_blockwise_impl(const Context& dev_ctx,
|
|
const DenseTensor& A,
|
|
const DenseTensor& A_decode_scale,
|
|
const DenseTensor& B,
|
|
const DenseTensor& B_decode_scale,
|
|
DenseTensor* D,
|
|
const DenseTensor& bias,
|
|
DenseTensor* pre_gelu_out,
|
|
bool transa,
|
|
bool transb,
|
|
bool grad,
|
|
DenseTensor* workspace,
|
|
bool accumulate,
|
|
bool use_split_accumulator,
|
|
int math_sm_count,
|
|
bool is_A_1d_scaled,
|
|
bool is_B_1d_scaled,
|
|
cudaStream_t stream) {
|
|
// Sanity checks
|
|
PADDLE_ENFORCE_EQ(
|
|
transa,
|
|
true,
|
|
common::errors::InvalidArgument("Only transa == true is supported"));
|
|
PADDLE_ENFORCE_EQ(
|
|
transb,
|
|
false,
|
|
common::errors::InvalidArgument("Only transb == false is supported"));
|
|
PADDLE_ENFORCE_EQ(A.place().GetType(),
|
|
AllocationType::GPU,
|
|
common::errors::InvalidArgument(
|
|
"Input tensor A must be on CUDA device."));
|
|
PADDLE_ENFORCE_EQ(B.place().GetType(),
|
|
AllocationType::GPU,
|
|
common::errors::InvalidArgument(
|
|
"Input tensor B must be on CUDA device."));
|
|
PADDLE_ENFORCE_EQ(D->place().GetType(),
|
|
AllocationType::GPU,
|
|
common::errors::InvalidArgument(
|
|
"Output tensor D must be on CUDA device."));
|
|
PADDLE_ENFORCE_EQ(IsFp8Dtype(A.dtype()),
|
|
true,
|
|
common::errors::InvalidArgument("A must be FP8"));
|
|
PADDLE_ENFORCE_EQ(IsFp8Dtype(B.dtype()),
|
|
true,
|
|
common::errors::InvalidArgument("B must be FP8"));
|
|
PADDLE_ENFORCE_EQ(
|
|
D->dtype() == phi::DataType::BFLOAT16 ||
|
|
D->dtype() == phi::DataType::FLOAT32,
|
|
true,
|
|
common::errors::InvalidArgument("D must be BFloat16 or float"));
|
|
PADDLE_ENFORCE_EQ(
|
|
A_decode_scale.dtype() == phi::DataType::FLOAT32,
|
|
true,
|
|
common::errors::InvalidArgument("A_decode_scale must be float"));
|
|
PADDLE_ENFORCE_EQ(
|
|
B_decode_scale.dtype() == phi::DataType::FLOAT32,
|
|
true,
|
|
common::errors::InvalidArgument("B_decode_scale must be float"));
|
|
PADDLE_ENFORCE_EQ(A.dims().size() == 2,
|
|
true,
|
|
common::errors::InvalidArgument("A must be 2D"));
|
|
PADDLE_ENFORCE_EQ(B.dims().size() == 2,
|
|
true,
|
|
common::errors::InvalidArgument("B must be 2D"));
|
|
PADDLE_ENFORCE_EQ(D->dims().size() == 2,
|
|
true,
|
|
common::errors::InvalidArgument("D must be 2D"));
|
|
|
|
const int m = transa ? A.dims()[0] : A.dims()[1];
|
|
const int k = transa ? A.dims()[1] : A.dims()[0];
|
|
const int n = transb ? B.dims()[1] : B.dims()[0];
|
|
|
|
int lda = k, ldb = k, ldc = m, ldd = m;
|
|
float alpha = 1.0, beta = accumulate ? 1.0 : 0.0;
|
|
|
|
cublasLtHandle_t ltHandle = dev_ctx.cublaslt_handle();
|
|
// Create operation descriptor
|
|
cublasLtMatmulDesc_t operationDesc = nullptr;
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescCreate(
|
|
&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F));
|
|
|
|
#if CUBLAS_VERSION >= 120805 && CUDA_VERSION >= 12080
|
|
// Setup scaling for A and B
|
|
cublasLtMatmulMatrixScale_t A_scale_mode, B_scale_mode;
|
|
// Note: in cuBLAS term, tensor name A and B are swapped.
|
|
if (is_B_1d_scaled && is_A_1d_scaled) {
|
|
A_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F;
|
|
B_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F;
|
|
} else if (!is_B_1d_scaled && is_A_1d_scaled) {
|
|
// So this corresponds to 2Dx1D GEMM.
|
|
A_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F;
|
|
B_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;
|
|
} else if (is_B_1d_scaled && !is_A_1d_scaled) {
|
|
// So this corresponds to 1Dx2D GEMM.
|
|
A_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F;
|
|
B_scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F;
|
|
} else {
|
|
PADDLE_THROW(
|
|
common::errors::InvalidArgument("2Dx2D scaling is not supported"));
|
|
}
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_A_SCALE_MODE,
|
|
&A_scale_mode,
|
|
sizeof(A_scale_mode)));
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_B_SCALE_MODE,
|
|
&B_scale_mode,
|
|
sizeof(B_scale_mode)));
|
|
#else
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Sub-channel FP8 GEMM requires CUDA 12.8 and cuBLAS 12.8.5 or later."));
|
|
#endif
|
|
|
|
// setup transa and transb
|
|
const cublasOperation_t transa_type = transa ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
const cublasOperation_t transb_type = transb ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatmulDescSetAttribute(operationDesc,
|
|
CUBLASLT_MATMUL_DESC_TRANSA,
|
|
&transa_type,
|
|
sizeof(transa_type)));
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatmulDescSetAttribute(operationDesc,
|
|
CUBLASLT_MATMUL_DESC_TRANSB,
|
|
&transb_type,
|
|
sizeof(transb_type)));
|
|
|
|
const void* A_decode_scale_ptr = A_decode_scale.data();
|
|
const void* B_decode_scale_ptr = B_decode_scale.data();
|
|
const cudaDataType_t Atype = ScalarTypeToCudaDataType(A.dtype());
|
|
const cudaDataType_t Btype = ScalarTypeToCudaDataType(B.dtype());
|
|
const cudaDataType_t Dtype = ScalarTypeToCudaDataType(D->dtype());
|
|
|
|
// split_accumulator is always true
|
|
const int8_t fast_accum_mode = 0;
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_FAST_ACCUM,
|
|
&fast_accum_mode,
|
|
sizeof(fast_accum_mode)));
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
|
|
&A_decode_scale_ptr,
|
|
sizeof(A_decode_scale_ptr)));
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
|
|
&B_decode_scale_ptr,
|
|
sizeof(B_decode_scale_ptr)));
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET,
|
|
&math_sm_count,
|
|
sizeof(math_sm_count)));
|
|
|
|
// Setup mat layout descriptors
|
|
cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr,
|
|
Ddesc = nullptr;
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutCreate(
|
|
&Adesc,
|
|
Atype,
|
|
transa_type == CUBLAS_OP_N ? m : k,
|
|
transa_type == CUBLAS_OP_N ? k : m,
|
|
lda));
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutCreate(
|
|
&Bdesc,
|
|
Btype,
|
|
transb_type == CUBLAS_OP_N ? k : n,
|
|
transb_type == CUBLAS_OP_N ? n : k,
|
|
ldb));
|
|
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatrixLayoutCreate(&Cdesc, Dtype, m, n, ldc));
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatrixLayoutCreate(&Ddesc, Dtype, m, n, ldd));
|
|
|
|
// setup epilogue attributes
|
|
cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulDescSetAttribute(
|
|
operationDesc,
|
|
CUBLASLT_MATMUL_DESC_EPILOGUE,
|
|
&epilogue,
|
|
sizeof(epilogue)));
|
|
|
|
// setup preference attributes
|
|
cublasLtMatmulPreference_t preference = nullptr;
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatmulPreferenceCreate(&preference));
|
|
size_t workspace_size = workspace->dims()[0];
|
|
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmulPreferenceSetAttribute(
|
|
preference,
|
|
CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
|
|
&workspace_size,
|
|
sizeof(workspace_size)));
|
|
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatmul(ltHandle,
|
|
operationDesc,
|
|
&alpha,
|
|
A.data(),
|
|
Adesc,
|
|
B.data(),
|
|
Bdesc,
|
|
&beta,
|
|
D->data(),
|
|
Cdesc,
|
|
D->data(),
|
|
Ddesc,
|
|
/*algo*/ nullptr,
|
|
workspace->data(),
|
|
workspace_size,
|
|
stream));
|
|
// Cleanup
|
|
if (preference)
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatmulPreferenceDestroy(preference));
|
|
if (Ddesc)
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutDestroy(Ddesc));
|
|
if (Cdesc)
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutDestroy(Cdesc));
|
|
if (Bdesc)
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutDestroy(Bdesc));
|
|
if (Adesc)
|
|
PADDLE_CUDABLAS_CHECK(phi::dynload::cublasLtMatrixLayoutDestroy(Adesc));
|
|
if (operationDesc)
|
|
PADDLE_CUDABLAS_CHECK(
|
|
phi::dynload::cublasLtMatmulDescDestroy(operationDesc));
|
|
}
|
|
|
|
} // anonymous namespace
|
|
|
|
template <typename T, typename Context>
|
|
void Fp8GemmBlockwiseKernel(const Context& dev_ctx,
|
|
const DenseTensor& A,
|
|
const DenseTensor& A_scale,
|
|
const DenseTensor& B,
|
|
const DenseTensor& B_scale,
|
|
const DenseTensor& input_result,
|
|
const DenseTensor& bias,
|
|
const DenseTensor& pre_gelu,
|
|
const DenseTensor& workspace,
|
|
bool transa,
|
|
bool transb,
|
|
bool grad,
|
|
bool accumulate,
|
|
bool use_split_accumulator,
|
|
int math_sm_count,
|
|
bool is_A_1d_scaled,
|
|
bool is_B_1d_scaled,
|
|
DenseTensor* output,
|
|
DenseTensor* pre_gelu_out,
|
|
DenseTensor* workspace_out) {
|
|
cublas_gemm_blockwise_impl<Context>(dev_ctx,
|
|
A,
|
|
A_scale,
|
|
B,
|
|
B_scale,
|
|
output,
|
|
bias,
|
|
pre_gelu_out,
|
|
transa,
|
|
transb,
|
|
grad,
|
|
workspace_out,
|
|
accumulate,
|
|
use_split_accumulator,
|
|
math_sm_count,
|
|
is_A_1d_scaled,
|
|
is_B_1d_scaled,
|
|
dev_ctx.stream());
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
// Register the kernel
|
|
PD_REGISTER_KERNEL(fp8_gemm_blockwise,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::Fp8GemmBlockwiseKernel,
|
|
phi::bfloat16,
|
|
phi::float8_e4m3fn,
|
|
uint8_t,
|
|
float,
|
|
double) {}
|