ecker/bitsandbytes-rocm
1
0
Fork 1
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

Added fp16 and thread/item template.

Browse Source
This commit is contained in:
Tim Dettmers 2023-04-28 18:26:52 -07:00
parent 3aef78342a
commit f6df4aef6a
6 changed files with 53 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
bitsandbytes/functional.py
View File

@ -1381,9 +1381,9 @@ def cutlass3_gemm(
transposed_A=False, transposed_A=False,
transposed_B=False, transposed_B=False,
): ):
sout = check_matmul(A, B, out, transposed_A, transposed_B, expected_type=torch.float32) sout = check_matmul(A, B, out, transposed_A, transposed_B, expected_type=A.dtype)
if out is None: if out is None:
out = torch.zeros(size=sout, dtype=torch.float32, device=A.device) out = torch.zeros(size=sout, dtype=A.dtype, device=A.device)
sA = A.shape sA = A.shape
sB = B.shape sB = B.shape
@ -1464,7 +1464,12 @@ def cutlass3_gemm(
lda = ct.c_int32(lda) lda = ct.c_int32(lda)
ldb = ct.c_int32(ldb) ldb = ct.c_int32(ldb)
ldc = ct.c_int32(ldc) ldc = ct.c_int32(ldc)
if A.dtype == torch.float32:
lib.cgemm_host_fp32(m, n, k, get_ptr(A), get_ptr(B), get_ptr(out), lda, ldb, ldc) lib.cgemm_host_fp32(m, n, k, get_ptr(A), get_ptr(B), get_ptr(out), lda, ldb, ldc)
elif A.dtype == torch.float16:
lib.cgemm_host_fp16(m, n, k, get_ptr(A), get_ptr(B), get_ptr(out), lda, ldb, ldc)
else:
raise NotImplementedError(f'Matmul not implemented for data type {A.dtype}')
return out return out

39
csrc/kernels.cu
View File

@ -2949,18 +2949,18 @@ template <int FORMAT> __global__ void kExtractOutliers(char *A, int *idx, char *
#define ROWS 2 #define ROWS 2
template <typename T> __global__ void gemm_device(int M, int N, int K, T const* A, T* B, T * out, int lda, int ldb, int ldc) template <typename T, int ITEMS, int THREADS> __global__ void gemm_device(int M, int N, int K, T const* A, T* B, T * out, int lda, int ldb, int ldc)
{ {
// 0. We want to fill a 8x128 tile for a thread block so we have 8x16 tile for each warp // 0. We want to fill a 8x128 tile for a thread block so we have 8x16 tile for each warp
// 1. Load dataB into register // 1. Load dataB into register
// 2. Dequantize B // 2. Dequantize B
// 3. Fetch data from A and multiply // 3. Fetch data from A and multiply
typedef cub::BlockLoad<T, 256 , 4, cub::BLOCK_LOAD_WARP_TRANSPOSE> LoadA; typedef cub::BlockLoad<T, THREADS , ITEMS, cub::BLOCK_LOAD_WARP_TRANSPOSE> LoadA;
//__shared__ typename LoadA::TempStorage loada; //__shared__ typename LoadA::TempStorage loada;
typedef cub::BlockLoad<T, 256 , 4, cub::BLOCK_LOAD_WARP_TRANSPOSE> LoadB; typedef cub::BlockLoad<T, THREADS , ITEMS, cub::BLOCK_LOAD_WARP_TRANSPOSE> LoadB;
//__shared__ typename LoadB::TempStorage loadb; //__shared__ typename LoadB::TempStorage loadb;
typedef cub::BlockReduce<T, 256> BlockReduce; typedef cub::BlockReduce<T, THREADS> BlockReduce;
// Allocate shared memory for BlockReduce // Allocate shared memory for BlockReduce
//__shared__ typename BlockReduce::TempStorage reduce; //__shared__ typename BlockReduce::TempStorage reduce;
@ -2971,15 +2971,13 @@ template <typename T> __global__ void gemm_device(int M, int N, int K, T const*
} temp_storage; } temp_storage;
T dataA[4]; T dataA[ITEMS];
T local_B[4]; T local_B[ITEMS];
T local_accC[ROWS]; T local_accC[ROWS];
int valid_items = 0; int valid_items = 0;
const int warp_id = threadIdx.x/32;
const int warp_lane = threadIdx.x % 32;
const int col_offset = blockIdx.x * 8; const int col_offset = blockIdx.x * 8;
__shared__ T tileA[ROWS*1024]; __shared__ T tileA[ROWS*THREADS*ITEMS];
__shared__ T accumulatorC[ROWS*8]; __shared__ T accumulatorC[ROWS*8];
//#pragma unroll 8 //#pragma unroll 8
@ -2991,17 +2989,17 @@ template <typename T> __global__ void gemm_device(int M, int N, int K, T const*
__syncthreads(); __syncthreads();
for(int inner_idx = 0; inner_idx < K; inner_idx+= 1024) for(int inner_idx = 0; inner_idx < K; inner_idx+= THREADS*ITEMS)
{ {
valid_items = K - inner_idx > 1024 ? 1024 : K - inner_idx; valid_items = K - inner_idx > THREADS*ITEMS ? THREADS*ITEMS : K - inner_idx;
int baserow = 0; int baserow = 0;
for(int row = baserow; row < (baserow+ROWS) && row < N; row++) for(int row = baserow; row < (baserow+ROWS) && row < N; row++)
{ {
LoadA(temp_storage.loada).Load(&(A[(row*K) + inner_idx]), dataA, valid_items, 0.0f); LoadA(temp_storage.loada).Load(&(A[(row*K) + inner_idx]), dataA, valid_items, 0.0f);
#pragma unroll 4 #pragma unroll ITEMS
for(int k = 0; k < 4; k++) for(int k = 0; k < ITEMS; k++)
tileA[row*1024 + threadIdx.x + (k*blockDim.x)] = dataA[k]; tileA[row*THREADS*ITEMS + threadIdx.x + (k*THREADS)] = dataA[k];
__syncthreads(); __syncthreads();
} }
@ -3021,16 +3019,16 @@ template <typename T> __global__ void gemm_device(int M, int N, int K, T const*
local_accC[k] = 0.0f; local_accC[k] = 0.0f;
int base_idxB = ldb*colB; int base_idxB = ldb*colB;
valid_items = K - inner_idx > 1024 ? 1024 : K - inner_idx; valid_items = K - inner_idx > THREADS*ITEMS ? THREADS*ITEMS : K - inner_idx;
LoadB(temp_storage.loadb).Load(&(B[base_idxB + inner_idx]), local_B, valid_items, 0.0f); LoadB(temp_storage.loadb).Load(&(B[base_idxB + inner_idx]), local_B, valid_items, 0.0f);
__syncthreads(); __syncthreads();
for(int row = 0; row < ROWS && row < N; row++) for(int row = 0; row < ROWS && row < N; row++)
{ {
#pragma unroll 4 #pragma unroll ITEMS
for(int k = 0; k < 4; k++) for(int k = 0; k < ITEMS; k++)
{ {
int idxA = row*1024 + threadIdx.x + (blockDim.x*k); int idxA = row*THREADS*ITEMS + threadIdx.x + (THREADS*k);
local_accC[row] += tileA[idxA]*local_B[k]; local_accC[row] += tileA[idxA]*local_B[k];
} }
@ -3124,7 +3122,10 @@ __global__ void with_staging_unified(float const* global_in, float * global_out,
// TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB, // TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB,
// TC * out, CStride dC, CBlockLayout , CThreadLayout tC, // TC * out, CStride dC, CBlockLayout , CThreadLayout tC,
// half alpha, half beta); // half alpha, half beta);
template __global__ void gemm_device<float>(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc); template __global__ void gemm_device<float, 4, 256>(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc);
template __global__ void gemm_device<half, 4, 256>(int M, int N, int K, half const* A, half* B, half * out, int lda, int ldb, int ldc);
template __global__ void gemm_device<float, 8, 256>(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc);
template __global__ void gemm_device<half, 8, 256>(int M, int N, int K, half const* A, half* B, half * out, int lda, int ldb, int ldc);
//template __global__ void kMatmul_inference_4bit<NF4, half, half, half>(half *A, unsigned char *B, half *out, int lda, int ldb, int rowsA, int colsA, int colsB); //template __global__ void kMatmul_inference_4bit<NF4, half, half, half>(half *A, unsigned char *B, half *out, int lda, int ldb, int rowsA, int colsA, int colsB);

2
csrc/kernels.cuh
View File

@ -138,6 +138,6 @@ template <int FORMAT> __global__ void kExtractOutliers(char *A, int *idx, char *
template <size_t stages_count /* Pipeline with stages_count stages */> template <size_t stages_count /* Pipeline with stages_count stages */>
__global__ void with_staging_unified(float const* global_in, float * global_out, size_t size, size_t batch_sz); __global__ void with_staging_unified(float const* global_in, float * global_out, size_t size, size_t batch_sz);
template <typename T> __global__ void gemm_device(int M, int N, int K, T const* A, T* B, T * out, int lda, int ldb, int ldc); template <typename T, int ITEMS, int THREADS> __global__ void gemm_device(int M, int N, int K, T const* A, T* B, T * out, int lda, int ldb, int ldc);
#endif #endif

3
csrc/ops.cu
View File

@ -689,7 +689,7 @@ template <typename T> void gemm_host(int m, int n, int k, T const* A, T* B, T
cout << m << endl; cout << m << endl;
cout << n << endl; cout << n << endl;
cout << k << endl; cout << k << endl;
gemm_device gemm_device<T, 8, 256>
<<< num_blocks, dimBlock, 0, 0 >>> <<< num_blocks, dimBlock, 0, 0 >>>
(m, n, k, (m, n, k,
A, A,
@ -702,6 +702,7 @@ template <typename T> void gemm_host(int m, int n, int k, T const* A, T* B, T
//============================================================== //==============================================================
template void gemm_host<float>(int m, int n, int k, float const* A, float* B, float * out, int lda, int ldb, int ldc); template void gemm_host<float>(int m, int n, int k, float const* A, float* B, float * out, int lda, int ldb, int ldc);
template void gemm_host<half>(int m, int n, int k, half const* A, half* B, half * out, int lda, int ldb, int ldc);
template void extractOutliers<COL_TURING>(char * A, int *idx, char *out, int idx_size, int rows, int cols); template void extractOutliers<COL_TURING>(char * A, int *idx, char *out, int idx_size, int rows, int cols);
template void extractOutliers<COL_AMPERE>(char * A, int *idx, char *out, int idx_size, int rows, int cols); template void extractOutliers<COL_AMPERE>(char * A, int *idx, char *out, int idx_size, int rows, int cols);

5
csrc/pythonInterface.c
View File

@ -22,6 +22,8 @@ void estimateQuantiles_fp16(half *A, float *code, float offset, int n){ estimate
void gemm_host_fp32(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc) void gemm_host_fp32(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc)
{ gemm_host<float>(M, N, K, A, B, out, lda, ldb, ldc); } { gemm_host<float>(M, N, K, A, B, out, lda, ldb, ldc); }
void gemm_host_fp16(int M, int N, int K, half const* A, half* B, half * out, int lda, int ldb, int ldc)
{ gemm_host<half>(M, N, K, A, B, out, lda, ldb, ldc); }
#define MAKE_FUNC32(fname, oname, gtype, gbits) \ #define MAKE_FUNC32(fname, oname, gtype, gbits) \
@ -314,6 +316,9 @@ extern "C"
void cgemm_host_fp32(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc) void cgemm_host_fp32(int M, int N, int K, float const* A, float* B, float * out, int lda, int ldb, int ldc)
{ gemm_host_fp32(M, N, K, A, B, out, lda, ldb, ldc); } { gemm_host_fp32(M, N, K, A, B, out, lda, ldb, ldc); }
void cgemm_host_fp16(int M, int N, int K, half const* A, half* B, half * out, int lda, int ldb, int ldc)
{ gemm_host_fp16(M, N, K, A, B, out, lda, ldb, ldc); }
#endif #endif
void cquantize_blockwise_cpu_fp32(float *code, float *A, float *absmax, unsigned char *out, long long blocksize, long long n){ quantize_cpu(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_cpu_fp32(float *code, float *A, float *absmax, unsigned char *out, long long blocksize, long long n){ quantize_cpu(code, A, absmax, out, blocksize, n); }
void cdequantize_blockwise_cpu_fp32(float *code, unsigned char *A, float *absmax, float *out, long long blocksize, long long n){ dequantize_cpu(code, A, absmax, out, blocksize, n); } void cdequantize_blockwise_cpu_fp32(float *code, unsigned char *A, float *absmax, float *out, long long blocksize, long long n){ dequantize_cpu(code, A, absmax, out, blocksize, n); }

14
tests/test_functional.py
View File

@ -2352,20 +2352,26 @@ def test_normal_map_tree():
print(pivots) print(pivots)
def test_cutlass3_gemm(): #@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=['fp32', 'fp16'])
A = torch.rand(2, 4092).cuda() @pytest.mark.parametrize("dtype", [torch.float16], ids=['fp16'])
B = torch.rand(4*4092, 4092).cuda() def test_cutlass3_gemm(dtype):
for i in range(2):
A = torch.rand(2, 4092, dtype=dtype, device='cuda')
B = torch.rand(4*4092, 4092, dtype=dtype, device='cuda')
#A = torch.rand(2, 4, dtype=dtype, device='cuda')
#B = torch.rand(4, 4, dtype=dtype, device='cuda')
#print('') #print('')
#print(A) #print(A)
#print(B.t()) #print(B.t())
C1 = torch.matmul(A, B.t()) C1 = torch.matmul(A, B.t())
C2 = F.cutlass3_gemm(A, B.t()) C2 = F.cutlass3_gemm(A, B.t())
#print(C1) #print(C1)
#print(C2) #print(C2)
torch.testing.assert_close(C1, C2) #torch.testing.assert_close(C1, C2)
def test_pipeline_func(): def test_pipeline_func():