// Copyright 2024 FP6-LLM authors
//
// 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.
//
// This file is modified from
// https://github.com/usyd-fsalab/fp6_llm/blob/5df6737cca32f604e957e3f63f03ccc2e4d1df0d/fp6_llm/csrc/include/kernel_matmul.cuh
#include "configs.h"
#include "utils_gmem.cuh"
#include "utils_core.cuh"
/************************** Bitwidth of Weight Segments
* ************************/
#define BIT_WIDTH_1 1
#define BIT_WIDTH_2 2
#define BIT_WIDTH_4 4
/*************************** 64*64 Weghts of Weight Matrix
* *********************/
#define WEIGHT_PER_WARP (WARP_M * WARP_K) // 64*64 = 4096
#define SMEM_SIZE_PER_WARP_1BIT \
(WEIGHT_PER_WARP * BIT_WIDTH_1 / \
8) // 512 Bytes, doubleBuffer not taken into consideration
#define SMEM_SIZE_PER_WARP_2BIT \
(WEIGHT_PER_WARP * BIT_WIDTH_2 / \
8) // 1024 Bytes, doubleBuffer not taken into consideration
#define SMEM_SIZE_PER_WARP_4BIT \
(WEIGHT_PER_WARP * BIT_WIDTH_4 / \
8) // 2048 Bytes, doubleBuffer not taken into consideration
#define SMEM_SIZE_PER_TB_1BIT \
(SMEM_SIZE_PER_WARP_1BIT * TilingConfig::BLOCK_WARPS * \
PIPELINE_LEVEL_GMEM) // #WARP=4; Trible-Buffer for 3-level pipeline for A
// = 6 KB; double buffer for 2-level pipeline A= 4
// KB.
#define SMEM_SIZE_PER_TB_2BIT \
(SMEM_SIZE_PER_WARP_2BIT * TilingConfig::BLOCK_WARPS * \
PIPELINE_LEVEL_GMEM) // #WARP=4; Trible-Buffer for 3-level pipeline for A
// = 12 KB; double buffer for 2-level pipeline A= 8
// KB.
#define SMEM_SIZE_PER_TB_4BIT \
(SMEM_SIZE_PER_WARP_4BIT * TilingConfig::BLOCK_WARPS * \
PIPELINE_LEVEL_GMEM) // #WARP=4; Trible-Buffer for 3-level pipeline for A
// = 24 KB; double buffer for 2-level pipeline A= 16
// KB.
#define SMEM_SIZE_PER_TB_A_TILE \
(SMEM_SIZE_PER_TB_1BIT + SMEM_SIZE_PER_TB_2BIT + \
SMEM_SIZE_PER_TB_4BIT) // used in fp6_linear.cu, Kernel_Ex().
/******************** Global Memory Layout For QUANTIZED DATA
* *******************/
#define NUM_INT4_PER_WARP_1BIT (WEIGHT_PER_WARP * BIT_WIDTH_1 / 128) // 32
#define NUM_INT4_PER_WARP_2BIT (WEIGHT_PER_WARP * BIT_WIDTH_2 / 128) // 64
#define NUM_INT4_PER_WARP_4BIT (WEIGHT_PER_WARP * BIT_WIDTH_4 / 128) // 128
/*
* C = A*B
* A: row major with ahead-of-time layout transformation, FP6
* B: col major, FP16
* C: col major, FP16
*/
template <typename TilingConfig, typename OutputDataType, int EXPONENT,
int MANTISSA>
__global__ void QUANT_GEMM_Kernel(const uint4* Weight, const half* Scales,
const half* B, OutputDataType* C,
const size_t M_Global, const size_t N_Global,
const size_t K_Global, int Split_K) {
#ifdef DEBUG_MODE
assert(K_Global % TilingConfig::TILE_K == 0);
assert(M_Global % TilingConfig::TILE_M == 0);
assert(gridDim.y == Split_K * (M_Global / TilingConfig::TILE_M));
#endif
// 1+2+4 weight split
constexpr int BIT_WIDTH = 1 + EXPONENT + MANTISSA;
constexpr int USE_SEG_1BIT = BIT_WIDTH & 1;
constexpr int USE_SEG_2BIT = BIT_WIDTH & 2;
constexpr int USE_SEG_4BIT = BIT_WIDTH & 4;
const uint4* Weight_1bit = Weight;
const uint4* Weight_2bit =
Weight_1bit +
(USE_SEG_1BIT ? M_Global * K_Global * BIT_WIDTH_1 / 128 : 0);
const uint4* Weight_4bit =
Weight_2bit +
(USE_SEG_2BIT ? M_Global * K_Global * BIT_WIDTH_2 / 128 : 0);
// Dynamic shared memory for FP16 A tiles， 128 Bytes aligned
extern __shared__ __align__(128) half smem[];
half(*smem_array)[WARP_K + PADDING_SHARED_MEM_FOR_B_8] =
reinterpret_cast<half(*)[WARP_K + PADDING_SHARED_MEM_FOR_B_8]>(
smem + SMEM_SIZE_PER_TB_A_TILE /
2); // Dynamic shared memory for FP16 B tiles
__shared__ half
QuantScales[64 *
TilingConfig::BLOCK_WARPS]; // static shared memory for
// quantization scales, 64 row
// per warp * 4 warps = 512 Bytes
// Thread Block Mapping, considering SplitK
const size_t BatchID = blockIdx.y / (M_Global / TilingConfig::TILE_M);
const size_t x =
blockIdx.x; // Output Block ID: (BlockID_Row = y; BlockID_Col = x )
const size_t y =
blockIdx.y %
(M_Global / TilingConfig::TILE_M); // Output Block ID: (BlockID_Row = y;
// BlockID_Col = x )
const size_t Tile_Start_M = y * TilingConfig::TILE_M;
const size_t Tile_Start_N = x * TilingConfig::TILE_N;
const size_t NumColumnToCopy =
(N_Global - Tile_Start_N) < TilingConfig::TILE_N
? (N_Global - Tile_Start_N)
: TilingConfig::TILE_N;
const size_t NumBlock_K = K_Global / TilingConfig::TILE_K;
const size_t AverageNumBlock_K = NumBlock_K / Split_K;
const size_t ExtraNumBlock_K = NumBlock_K - AverageNumBlock_K * Split_K;
size_t NumIter = AverageNumBlock_K;
size_t StartBlockID_K = AverageNumBlock_K * BatchID;
if (BatchID < ExtraNumBlock_K) {
NumIter++;
StartBlockID_K += BatchID;
} else
StartBlockID_K += ExtraNumBlock_K;
// Warp ID.
const int warpId = threadIdx.x / WARP_SIZE;
int WARP_i = warpId / TilingConfig::BLOCK_COL_WARPS; // WARP_i: row number;
// WARP_j: column number
// int WARP_j = warpId % TilingConfig::BLOCK_COL_WARPS;
// Global Memory Address for Matrix A (Weight)
// /////////////////////////////////////////////////////////////////////////
// StartPTR for each ThreadBlock(TB)
const uint4* TB_StartGPTR_A_1BIT =
Weight_1bit +
(y * TilingConfig::BLOCK_ROW_WARPS) * NumBlock_K * NUM_INT4_PER_WARP_1BIT;
const uint4* TB_StartGPTR_A_2BIT =
Weight_2bit +
(y * TilingConfig::BLOCK_ROW_WARPS) * NumBlock_K * NUM_INT4_PER_WARP_2BIT;
const uint4* TB_StartGPTR_A_4BIT =
Weight_4bit +
(y * TilingConfig::BLOCK_ROW_WARPS) * NumBlock_K * NUM_INT4_PER_WARP_4BIT;
// StartPTR for each WARP.
const uint4* WARP_StartGPTR_A_1BIT =
TB_StartGPTR_A_1BIT + WARP_i * NumBlock_K * NUM_INT4_PER_WARP_1BIT;
const uint4* WARP_StartGPTR_A_2BIT =
TB_StartGPTR_A_2BIT + WARP_i * NumBlock_K * NUM_INT4_PER_WARP_2BIT;
const uint4* WARP_StartGPTR_A_4BIT =
TB_StartGPTR_A_4BIT + WARP_i * NumBlock_K * NUM_INT4_PER_WARP_4BIT;
// StartPTR for each WARP, considering SplitK
const size_t WARP_Start_UnitID_K = StartBlockID_K;
WARP_StartGPTR_A_1BIT += WARP_Start_UnitID_K * NUM_INT4_PER_WARP_1BIT;
WARP_StartGPTR_A_2BIT += WARP_Start_UnitID_K * NUM_INT4_PER_WARP_2BIT;
WARP_StartGPTR_A_4BIT += WARP_Start_UnitID_K * NUM_INT4_PER_WARP_4BIT;
// Copying A tile from Global to Shared, using double-buffer
// ////////////////////////////////////////////////////////// StartSPTR for
// each ThreadBlock
uint32_t* AFrag_1BIT_SPTR = reinterpret_cast<uint32_t*>(smem);
uint32_t* AFrag_2BIT_SPTR = AFrag_1BIT_SPTR + SMEM_SIZE_PER_TB_1BIT / 4;
uint32_t* AFrag_4BIT_SPTR =
AFrag_2BIT_SPTR +
SMEM_SIZE_PER_TB_2BIT /
4; // 8 buffers including double buffers, 12 for trible buffers
// StartSPTR for each WARP
AFrag_1BIT_SPTR += warpId * SMEM_SIZE_PER_WARP_1BIT / 4;
AFrag_2BIT_SPTR += warpId * SMEM_SIZE_PER_WARP_2BIT / 4;
AFrag_4BIT_SPTR += warpId * SMEM_SIZE_PER_WARP_4BIT / 4;
// Pre-fetch of A tile
for (int i = 0; i < PIPELINE_LEVEL_GMEM - 1; i++) {
if (USE_SEG_1BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_1BIT>(
AFrag_1BIT_SPTR + i * SMEM_SIZE_PER_WARP_1BIT / 4 * 4,
WARP_StartGPTR_A_1BIT);
if (USE_SEG_2BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_2BIT>(
AFrag_2BIT_SPTR + i * SMEM_SIZE_PER_WARP_2BIT / 4 * 4,
WARP_StartGPTR_A_2BIT);
if (USE_SEG_4BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_4BIT>(
AFrag_4BIT_SPTR + i * SMEM_SIZE_PER_WARP_4BIT / 4 * 4,
WARP_StartGPTR_A_4BIT);
WARP_StartGPTR_A_1BIT += SMEM_SIZE_PER_WARP_1BIT / 16;
WARP_StartGPTR_A_2BIT += SMEM_SIZE_PER_WARP_2BIT / 16;
WARP_StartGPTR_A_4BIT += SMEM_SIZE_PER_WARP_4BIT / 16;
}
// Global Memory Address for Matrix A (QuantScale)
// /////////////////////////////////////////////////////////////////////
const half* TB_StartGPTR_A_Scale =
Scales + (y * TilingConfig::BLOCK_ROW_WARPS) * 64;
const half* WARP_StartGPTR_A_Scales = TB_StartGPTR_A_Scale + WARP_i * 64;
CopyFromGlobalToShared_Scales(QuantScales + WARP_i * 64,
WARP_StartGPTR_A_Scales);
// Copying B tile from Global to Shared, considering SplitK
// /////////////////////////////////////////////////////////////
const half* BTile_GPTR =
B + Tile_Start_N * K_Global + StartBlockID_K * TilingConfig::TILE_K;
for (int i = 0; i < PIPELINE_LEVEL_GMEM - 1; i++) {
CopyFromGlobalToShared<TilingConfig::TILE_N, TilingConfig::BLOCK_WARPS>(
smem_array + i * TilingConfig::TILE_N, BTile_GPTR, K_Global,
NumColumnToCopy);
BTile_GPTR += TilingConfig::TILE_K;
}
// Register Allocation for A,B, and C, Initilazed to Zeros
// /////////////////////////////////////////////////////////////////////
constexpr int NumRegSets_a =
WARP_ROW_MMA_TENSORS; // 1 set = 4 registers, containing a 16*16 MMA
// block
constexpr int NumRegSets_b =
(TilingConfig::WARP_COL_MMA_TENSORS == 1)
? 1
: TilingConfig::WARP_COL_MMA_TENSORS /
2; // 1 set = 4 registers, containing a 16*16 MMA block
uint32_t a[NumRegSets_a * PIPELINE_LEVEL_SMEM]
[4]; // double/Trible buffer is used // Registers to store
// decompressed FP6
uint32_t b[NumRegSets_b * PIPELINE_LEVEL_SMEM]
[4]; // double/Triple buffer is used // Register to store FP16 B
// matrix (a slice)
float c[NumRegSets_a * NumRegSets_b][REG_PER_THREAD_C_TENSOR_16_16];
for (int i = 0; i < NumRegSets_a * NumRegSets_b; i++)
for (int j = 0; j < REG_PER_THREAD_C_TENSOR_16_16; j++) c[i][j] = 0.0f;
//
cp_async_wait_all();
__syncthreads();
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
uint32_t Scales_RPTR[4]; // 4 Registers per thread for Quantization Scales
ExtractFromSharedToReg_Scales(Scales_RPTR, QuantScales + WARP_i * 64);
// Initializing the Software Pipeline: writing registers.
// ////////////////////////////////////////////////////////////////////////////////////////////////
initialize_mma_slice<TilingConfig, EXPONENT, MANTISSA>(
a, b, AFrag_1BIT_SPTR, AFrag_2BIT_SPTR, AFrag_4BIT_SPTR, smem_array,
Scales_RPTR);
// The outer loop.
// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
#pragma unroll(1)
for (size_t tile_id_k = 0; tile_id_k < NumIter; tile_id_k++) {
// Trible-Buffer for A Tile
uint32_t* __restrict__ read_SPTR_Frag_1bit =
AFrag_1BIT_SPTR +
((tile_id_k + 0) % PIPELINE_LEVEL_GMEM) * SMEM_SIZE_PER_WARP_1BIT / 4 *
4; // 512 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
uint32_t* __restrict__ read_SPTR_Frag_2bit =
AFrag_2BIT_SPTR +
((tile_id_k + 0) % PIPELINE_LEVEL_GMEM) * SMEM_SIZE_PER_WARP_2BIT / 4 *
4; // 1024 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
uint32_t* __restrict__ read_SPTR_Frag_4bit =
AFrag_4BIT_SPTR +
((tile_id_k + 0) % PIPELINE_LEVEL_GMEM) * SMEM_SIZE_PER_WARP_4BIT / 4 *
4; // 2048 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
uint32_t* __restrict__ read2_SPTR_Frag_1bit =
AFrag_1BIT_SPTR + ((tile_id_k + 1) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_1BIT / 4 * 4;
uint32_t* __restrict__ read2_SPTR_Frag_2bit =
AFrag_2BIT_SPTR + ((tile_id_k + 1) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_2BIT / 4 * 4;
uint32_t* __restrict__ read2_SPTR_Frag_4bit =
AFrag_4BIT_SPTR + ((tile_id_k + 1) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_4BIT / 4 * 4;
uint32_t* __restrict__ write_SPTR_Frag_1bit =
AFrag_1BIT_SPTR +
((tile_id_k + (PIPELINE_LEVEL_GMEM - 1)) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_1BIT / 4 *
4; // 512 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
uint32_t* __restrict__ write_SPTR_Frag_2bit =
AFrag_2BIT_SPTR +
((tile_id_k + (PIPELINE_LEVEL_GMEM - 1)) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_2BIT / 4 *
4; // 1024 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
uint32_t* __restrict__ write_SPTR_Frag_4bit =
AFrag_4BIT_SPTR +
((tile_id_k + (PIPELINE_LEVEL_GMEM - 1)) % PIPELINE_LEVEL_GMEM) *
SMEM_SIZE_PER_WARP_4BIT / 4 *
4; // 2048 (1)*4: 4 WARPs; (2)/4: int*+1 = char*+16
// Trible-Buffer for B Tile
// MODIFICATION NOTE: to support MSVC, half __restrict__ (*read_SPTR ) is
// changed to below. similarly for read2_SPTR and write_SPTR.
half(*__restrict__ read_SPTR)[WARP_K + PADDING_SHARED_MEM_FOR_B_8] =
smem_array +
((tile_id_k + 0) % PIPELINE_LEVEL_GMEM) * TilingConfig::TILE_N;
half(*__restrict__ read2_SPTR)[WARP_K + PADDING_SHARED_MEM_FOR_B_8] =
smem_array +
((tile_id_k + 1) % PIPELINE_LEVEL_GMEM) * TilingConfig::TILE_N;
half(*__restrict__ write_SPTR)[WARP_K + PADDING_SHARED_MEM_FOR_B_8] =
smem_array +
((tile_id_k + (PIPELINE_LEVEL_GMEM - 1)) % PIPELINE_LEVEL_GMEM) *
TilingConfig::TILE_N;
//
bool GlobalCopy = (tile_id_k + PIPELINE_LEVEL_GMEM - 1) < NumIter;
// Copying A tile from Global to Register, Bypassing L1, using double-buffer
if (USE_SEG_1BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_1BIT>(
write_SPTR_Frag_1bit, WARP_StartGPTR_A_1BIT, GlobalCopy);
if (USE_SEG_2BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_2BIT>(
write_SPTR_Frag_2bit, WARP_StartGPTR_A_2BIT, GlobalCopy);
if (USE_SEG_4BIT)
CopyFromGlobalToShared_A<SMEM_SIZE_PER_WARP_4BIT>(
write_SPTR_Frag_4bit, WARP_StartGPTR_A_4BIT, GlobalCopy);
// copying B tile from GlobalMemory to SharedMemory
CopyFromGlobalToShared<TilingConfig::TILE_N, TilingConfig::BLOCK_WARPS>(
write_SPTR, BTile_GPTR, K_Global, NumColumnToCopy, GlobalCopy);
cp_async_group_commit();
core_mma_slice<TilingConfig, EXPONENT, MANTISSA>(
c, a, b, read_SPTR_Frag_1bit, read_SPTR_Frag_2bit, read_SPTR_Frag_4bit,
read_SPTR, Scales_RPTR,
1); // read_SPTR_Frag_2bit, read_SPTR_Frag_4bit are different for each
// WARP; read_SPTR is shared among WARPs
core_mma_slice<TilingConfig, EXPONENT, MANTISSA>(
c, a, b, read_SPTR_Frag_1bit, read_SPTR_Frag_2bit, read_SPTR_Frag_4bit,
read_SPTR, Scales_RPTR, 2);
core_mma_slice<TilingConfig, EXPONENT, MANTISSA>(
c, a, b, read_SPTR_Frag_1bit, read_SPTR_Frag_2bit, read_SPTR_Frag_4bit,
read_SPTR, Scales_RPTR, 3);
// Barriers and Synchronizations
cp_async_wait_group<PIPELINE_LEVEL_GMEM - 2>();
__syncthreads();
core_mma_slice<TilingConfig, EXPONENT, MANTISSA>(
c, a, b, read2_SPTR_Frag_1bit, read2_SPTR_Frag_2bit,
read2_SPTR_Frag_4bit, read2_SPTR, Scales_RPTR, 0);
// Updating global PTRs
WARP_StartGPTR_A_1BIT +=
SMEM_SIZE_PER_WARP_1BIT / 16; // 2KB/16=128 (1)/16: int4*+1 = char*+16
WARP_StartGPTR_A_2BIT +=
SMEM_SIZE_PER_WARP_2BIT / 16; // 4KB/16=256 (1)/16: int4*+1 = char*+16
WARP_StartGPTR_A_4BIT +=
SMEM_SIZE_PER_WARP_4BIT / 16; // 8KB/16=512 (1)/16: int4*+1 = char*+16
BTile_GPTR += TilingConfig::TILE_K;
}
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Store the C fragments to shared memory.
float(*smem_CFrag)[TilingConfig::TILE_M + PADDING_SHARED_MEM_FOR_C_4] =
reinterpret_cast<
float(*)[TilingConfig::TILE_M + PADDING_SHARED_MEM_FOR_C_4]>(smem);
StoreToSharedMemoryFromRegister<TilingConfig>(smem_CFrag, c);
__syncthreads();
// Now that shared memory contains all the D tiles, stream them to global
// memory.
OutputDataType* BlockGlobalPTR = C + BatchID * (M_Global * N_Global) +
Tile_Start_M + Tile_Start_N * M_Global;
for (size_t i = warpId; i < NumColumnToCopy;
i += TilingConfig::BLOCK_WARPS) // i-th column
#pragma unroll
for (size_t j = threadIdx.x % WARP_SIZE; j < TilingConfig::TILE_M;
j += WARP_SIZE) // j-th row
{
if constexpr (std::is_same<OutputDataType, half>::value)
BlockGlobalPTR[j + i * M_Global] = __float2half_rn(smem_CFrag[i][j]);
else
BlockGlobalPTR[j + i * M_Global] = smem_CFrag[i][j];
}
}