InferLLM 框架中 GPU 内核实现分析

InferLLM 框架中的 GPU 内核实现主要由 kernel_gpu.h 和 kernel_gpu.cu 两个文件组成，它们提供了大语言模型在 GPU 上运行所需的各种计算操作的实现。

1. 整体架构

1.1 文件结构

• kernel_gpu.h：声明 GPU 内核函数和注册内核
• kernel_gpu.cu：实现 GPU 内核函数

1.2 核心组件

1. cudaHandle 结构体：管理 CUDA 资源

   struct cudaHandle {
       cudaStream_t stream{nullptr};
       cublasHandle_t cublas_handle{nullptr};
   };
   

2. 内核函数：实现各种计算操作

   • 嵌入查找
   • 元素级操作
   • 归一化
   • Softmax
   • 矩阵乘法
   • 注意力计算
   • 位置编码
   • 掩码操作

3. 内核注册机制：使用宏注册内核函数

   #define PartialImplementKernel(kernel_id, fun)               \
       template <typename... Args>                              \
       struct Comp<KernelID::kernel_id, Args...> {              \
           static void exec(Args... args, cudaHandle* handle) { \
               return fun(std::forward<Args>(args)..., handle); \
           }                                                    \
       };
   

2. 主要功能模块分析

2.1 嵌入查找（Embedding Lookup）

void llm_embedding_get_int4_float(
        const void* weights, const uint32_t* index, float* dst, uint32_t len_seq,
        uint32_t embd, cudaHandle* handle);
void llm_embedding_get_float_float(
        const float* weights, const uint32_t* index, float* dst, uint32_t len_seq,
        uint32_t embd, cudaHandle* handle);

这些函数在 GPU 上实现嵌入查找操作：

• llm_embedding_get_int4_float：从 4 位整数量化的嵌入表中查找
• llm_embedding_get_float_float：从浮点数嵌入表中查找

实现中使用 CUDA 内核函数并行处理多个序列位置：

__global__ void llm_embedding_get_int4_float_gpu(
        const void* weights, const uint32_t* index, float* dst, uint32_t len_seq,
        uint32_t embd, const int weight_stride) {
    int seq_id = blockIdx.y;
    int thread_id = threadIdx.x + blockIdx.x * blockDim.x;
    if (thread_id < embd / 2) {
        uint32_t row = index[seq_id];
        dst = dst + seq_id * embd;
        const void* src = (static_cast<const char*>(weights) + row * weight_stride);
        int q40_block_id = thread_id * 2 / QK40;
        int block_offset = thread_id % (QK40 / 2);
        BlockQ40* q40_block = (BlockQ40*)src + q40_block_id;
        float scale = q40_block->d;
        uint8_t value = q40_block->qs[block_offset];
        const int8_t v1 = value & 0xf;
        const int8_t v2 = value >> 4;
        dst[thread_id * 2] = (v1 - 8) * scale;
        dst[thread_id * 2 + 1] = (v2 - 8) * scale;
    }
}

2.2 元素级操作（Elementwise Operations）

void llm_elemwise_compute_float(
        InData<float> srcs, float* dst, size_t len, ElemMode mode, cudaHandle* handle);
void llm_elemwise_compute_float_scale(
        float* src, float* dst, size_t len, float scale, cudaHandle* handle);
void llm_elemwise_broadcast_dim0_src1_compute_float(
        const float* src0, const float* src1, float* dst, uint32_t len0, uint32_t len1,
        ElemMode mode, cudaHandle* handle);

这些函数实现了各种元素级操作：

• 加法、乘法
• Silu、Gelu 激活函数
• 缩放操作
• 广播操作

实现中使用函数对象（Functor）和模板元编程简化代码：

struct SiluFunctor {
    __device__ float operator()(uint32_t i, const float* input) const {
        float src = input[i];
        return src / (1.0 + exp(-src));
    }
};

struct GeluFunctor {
    __device__ float operator()(uint32_t i, const float* input) const {
        float src = input[i];
        return 0.5 * src * (1 + tanh(sqrt(2.0 / PI) * (src + PGELU * src * src * src)));
    }
};

template <typename Function, typename... Args>
__global__ void ApplyFunction(Function functor, int64_t n, float* ret, Args... args) {
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < n) {
        ret[tid] = functor(tid, args...);
    }
}

2.3 归一化（Normalization）

void llm_norm_compute_float(
        const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps,
        cudaHandle* handle);
void llm_rms_norm_compute_float(
        const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps,
        cudaHandle* handle);

这些函数实现了层归一化和 RMS 归一化：

• llm_norm_compute_float：层归一化，计算均值和方差
• llm_rms_norm_compute_float：RMS 归一化，只计算均方根

实现中使用 warp-level 归约优化性能：

__global__ void rms_norm_f32(const float* x, float* dst, const int ncols, float eps) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    const int WARP_SIZE = blockDim.x;

    float tmp = 0.0f;  // partial sum for thread in warp
    for (int i = 0; i < ncols; i += WARP_SIZE) {
        const int col = i + tid;
        const float xi = x[row * ncols + col];
        tmp += xi * xi;
    }

    // sum up partial sums
    __syncthreads();
#pragma unroll
    for (int mask = 16; mask > 0; mask >>= 1) {
        tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32);
    }

    const float mean = tmp / ncols;
    const float scale = 1.0f / sqrtf(mean + eps);

    for (int i = 0; i < ncols; i += WARP_SIZE) {
        const int col = i + tid;
        dst[row * ncols + col] = scale * x[row * ncols + col];
    }
}

2.4 Softmax

void llm_softmax_compute_float(
        const float* src, float* dst, uint32_t len_row, uint32_t col,
        cudaHandle* handle);

实现 Softmax 函数，使用 warp-level 归约优化性能：

__global__ void softmax_f32_cuda(const float* x, float* dst, const int cols) {
    const int row = blockDim.y * blockIdx.y + threadIdx.y;
    const int block_size = blockDim.x;
    const int tid = threadIdx.x;
    const float* src = x + row * cols;
    dst = dst + row * cols;

    float max = -INFINITY;
    for (int col = tid; col < cols; col += block_size) {
        const float val = src[col];
        max = val > max ? val : max;
    }

    // sum up partial sums
    __syncthreads();
#pragma unroll
    for (int mask = 16; mask > 0; mask >>= 1) {
        float temp = __shfl_xor_sync(0xffffffff, max, mask);
        max = max > temp ? max : temp;
    }

    float sum = 0.0;
    for (int col = tid; col < cols; col += block_size) {
        const float val = expf(src[col] - max);
        sum += val;
        dst[col] = val;
    }

    // sum up partial sums
    __syncthreads();
#pragma unroll
    for (int mask = 16; mask > 0; mask >>= 1) {
        sum += __shfl_xor_sync(0xffffffff, sum, mask, 32);
    }

    for (int col = tid; col < cols; col += block_size) {
        dst[col] /= sum;
    }
}

2.5 矩阵乘法（Matrix Multiplication）

void llm_matmul_compute_int4_float(
        float* dst, const void* src0, const float* bias, const float* src1, uint32_t M,
        uint32_t N, uint32_t K, void* workspace, uint32_t size, cudaHandle* handle);
void llm_matmul_compute_float_float(
        float* dst, const float* src0, const float* bias, const float* src1, uint32_t M,
        uint32_t N, uint32_t K, void* workspace, uint32_t size, cudaHandle* handle);

这些函数实现了不同精度的矩阵乘法：

• llm_matmul_compute_int4_float：4 位整数权重与浮点数激活值的矩阵乘法
• llm_matmul_compute_float_float：浮点数矩阵乘法

浮点数矩阵乘法使用 cuBLAS 库优化性能：

void llm_matmul_compute_float_float(
        float* dst, const float* src0, const float* bias, const float* src1, uint32_t M,
        uint32_t N, uint32_t K, void* workspace, uint32_t size, cudaHandle* handle) {
    cudaStream_t stream = handle->stream;
    cublasHandle_t cublas_handle = handle->cublas_handle;
    float alpha = 1.f;
    float beta = 0.f;
    CUBLAS_CHECK(cublasSetStream(cublas_handle, stream));
    CUBLAS_CHECK(cublasSgemm(
            cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K,
            &alpha, src0, K, src1, K, &beta, dst, N));
    if (bias != nullptr) {
        llm_elemwise_broadcast_dim0_src1_compute_float(
                dst, bias, dst, M, N, ElemMode::Add, handle);
    }
}

2.6 注意力计算（Attention Computation）

void llm_matmul_compute_with_head_stride_float(
        float* dst, const float* srck, const float* srcq, uint32_t seqlen,
        uint32_t embd, uint32_t head, uint32_t nr_past, cudaHandle* handle);
void llm_matmul_compute_with_head_strideq_broadcastk_float(
        float* dst, const float* srck, const float* srcq, uint32_t seqlen,
        uint32_t embd, uint32_t head, uint32_t query_group_num, uint32_t nr_past,
        cudaHandle* handle);
void llm_head_batched_matmul_compute_float(
        float* dst, const float* v, const float* qk, uint32_t seqlen, uint32_t embd,
        uint32_t head, uint32_t nr_past, cudaHandle* handle);
void llm_head_batched_matmul_broadcastv_float(
        float* dst, const float* v, const float* qk, uint32_t seqlen, uint32_t embd,
        uint32_t head, uint32_t query_group_num, uint32_t nr_past, cudaHandle* handle);

这些函数实现了自注意力机制所需的矩阵运算：

• 多头注意力
• 多查询注意力（MQA）

实现中使用 cuBLAS 的批处理矩阵乘法优化性能：

void llm_matmul_compute_with_head_stride_float(
        float* dst, const float* srck, const float* srcq, uint32_t seqlen,
        uint32_t embd, uint32_t head, uint32_t nr_past, cudaHandle* handle) {
    uint32_t head_embd = embd / head;
    uint32_t M = seqlen;
    uint32_t N = seqlen + nr_past;
    uint32_t K = head_embd;
    cudaStream_t stream = handle->stream;
    cublasHandle_t cublas_handle = handle->cublas_handle;
    float alpha = 1.f;
    float beta = 0.f;
    CUBLAS_CHECK(cublasSetStream(cublas_handle, stream));
    CUBLAS_CHECK(cublasSgemmStridedBatched(
            cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K, &alpha, srck, embd,
            head_embd, srcq, embd, head_embd, &beta, dst, N, M * N, head));
}

2.7 位置编码（Position Encoding）

void llm_rope_compute_float(
        float* dst, const float* src0, uint32_t n_past, uint32_t n_rot, RotMode m,
        uint32_t N, uint32_t head, uint32_t embd, cudaHandle* handle);
void llm_glm_rope_compute_float(
        float* dst, const float* src0, uint32_t n_past, uint32_t gmask_positon,
        uint32_t seqlen, uint32_t head, uint32_t embd, cudaHandle* handle);

这些函数实现了旋转位置编码（RoPE）和 GLM 模型特定的旋转位置编码：

template <bool halfmode>
__global__ void rope_compute_float(
        float* dst, const float* src, float theta_scale, int position_offset,
        int n_rot, int seqlen, int head, int head_embd) {
    const int seq = blockIdx.y;
    const int h = blockIdx.x;
    int p = threadIdx.x;

    if (seq >= seqlen || h >= head || p >= head_embd / 2)
        return;

    const int position = seq + position_offset;
    const int half_embd = head_embd / 2;
    const int offset = h * head_embd + seq * head * head_embd;

    if (p < n_rot) {
        const float theta = powf(theta_scale, p);
        const float sin_theta = sinf(position * theta);
        const float cos_theta = cosf(position * theta);

        if (halfmode) {
            const float x0 = src[offset + p];
            const float x1 = src[offset + p + half_embd];
            dst[offset + p] = x0 * cos_theta - x1 * sin_theta;
            dst[offset + p + half_embd] = x0 * sin_theta + x1 * cos_theta;
        } else {
            const float x0 = src[offset + 2 * p];
            const float x1 = src[offset + 2 * p + 1];
            dst[offset + 2 * p] = x0 * cos_theta - x1 * sin_theta;
            dst[offset + 2 * p + 1] = x0 * sin_theta + x1 * cos_theta;
        }
    } else {
        if (halfmode) {
            dst[offset + p] = src[offset + p];
            dst[offset + p + half_embd] = src[offset + p + half_embd];
        } else {
            dst[offset + 2 * p] = src[offset + 2 * p];
            dst[offset + 2 * p + 1] = src[offset + 2 * p + 1];
        }
    }
}

GLM 模型的位置编码实现了特殊的位置计算逻辑：

__global__ void glm_rope_compute_float(
        float* dst, const float* src, int32_t n_past, int32_t gmask_positon,
        int32_t seqlen, int32_t head, int32_t embd) {
    const int seq = blockIdx.y;
    const int h = blockIdx.x;
    int p = threadIdx.x;

    if (seq >= seqlen || h >= head || p >= embd / 2)
        return;

    int quart_embd = embd / 4;
    int half_embd = embd / 2;

    int position_id = MIN(seq + n_past, gmask_positon);
    int block_position_id = MAX((n_past + seq) - gmask_positon, 0);

    bool is_second_half = p >= quart_embd;

    position_id = is_second_half ? block_position_id : position_id;

    p = is_second_half ? p - quart_embd : p;

    const float theta = powf(10000.0f, -2.0f * p / quart_embd);
    const float sin_theta = sinf(position_id * theta);
    const float cos_theta = cosf(position_id * theta);

    const int offset = h * embd + seq * head * embd;
    const int rot_offset = is_second_half ? quart_embd : 0;

    const float x0 = src[offset + rot_offset + p];
    const float x1 = src[offset + rot_offset + p + half_embd];
    dst[offset + rot_offset + p] = x0 * cos_theta - x1 * sin_theta;
    dst[offset + rot_offset + p + half_embd] = x0 * sin_theta + x1 * cos_theta;
}

2.8 掩码操作（Masking Operations）

void llm_diag_mask_inf_float(
        float* dst, const float* src0, uint32_t n_past, uint32_t N, uint32_t head,
        cudaHandle* handle);
void llm_glm_gmask_inf_float(
        float* dst, uint32_t n_past, uint32_t seqlen, uint32_t head,
        cudaHandle* handle);
void llm_scale_diag_mask_inf_float(
        float* dst, const float* src0, float scale, uint32_t n_past, uint32_t seqlen,
        uint32_t head, cudaHandle* handle);

这些函数实现了注意力掩码操作：

• llm_diag_mask_inf_float：自回归掩码，将上三角部分设为负无穷
• llm_glm_gmask_inf_float：GLM 模型特定的掩码
• llm_scale_diag_mask_inf_float：先缩放再掩码

__global__ void diag_mask_inf_f32(
        const float* src, float* dst, const int past, const int len,
        const int head_dim) {
    const int head = blockIdx.z;
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row >= len || col >= past + len || head >= head_dim)
        return;

    const int offset = head * len * (past + len) + row * (past + len) + col;
    if (col >= past && col - past > row) {
        dst[offset] = -INFINITY;
    } else {
        dst[offset] = src[offset];
    }
}

GLM 模型的掩码实现了特殊的掩码逻辑：

__global__ void glm_gmask_inf_f32(
        float* dst, const int past, const int seqlen, const int head) {
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    const int head_id = index / seqlen;
    const int row_id = index % seqlen;

    // the laxt row not set -inf
    if (row_id >= seqlen - 1 || head_id >= head)
        return;

    int total_seq = seqlen + past;
    int offset = head_id * seqlen * total_seq + row_id * total_seq + total_seq - 1;
    dst[offset] = -INFINITY;
}

3. 内核注册机制

GPU 内核使用模板和宏实现内核注册机制，将函数与内核 ID 关联起来：

#define PartialImplementKernel(kernel_id, fun)               \
    template <typename... Args>                              \
    struct Comp<KernelID::kernel_id, Args...> {              \
        static void exec(Args... args, cudaHandle* handle) { \
            return fun(std::forward<Args>(args)..., handle); \
        }                                                    \
    };

#define PartialImplementSpace(kernel_id, fun)        \
    template <typename... Args>                      \
    struct Space<KernelID::kernel_id, Args...> {     \
        static size_t get(Args... args) {            \
            return fun(std::forward<Args>(args)...); \
        }                                            \
    };

#define NOImplementKernel(kernel_id)                         \
    template <typename... Args>                              \
    struct Comp<KernelID::kernel_id, Args...> {              \
        static void exec(Args... args, cudaHandle* handle) { \
            INFER_ASSERT(0, "kernel not implement");         \
        }                                                    \
    };

这些宏用于注册已实现的内核和标记未实现的内核：

PartialImplementKernel(ElemwiseFloat, llm_elemwise_compute_float);
PartialImplementKernel(ElemwiseFloatScale, llm_elemwise_compute_float_scale);
// ... 其他已实现的内核

NOImplementKernel(MatmulInt4FloatPacked);
NOImplementKernel(MatmulInt4WeightReorder);
NOImplementKernel(MatmulInt8Float);
NOImplementKernel(EmbeddingGetInt8Float);

4. CUDA 优化技术

4.1 线程块和网格配置

代码中根据不同的计算需求，使用不同的线程块和网格配置：

// 一维网格和线程块
const dim3 block_dims(CUDA_NUM_THREADS, 1, 1);
const dim3 block_nums((len + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS, 1, 1);

// 二维网格
const dim3 block_dims(512, 1, 1);
const dim3 block_nums((ncols + 511) / 512, rows, 1);

// 三维网格
const dim3 block_dims(kBlockSize, kBlockSize, 1);
const dim3 block_nums(block_x, block_y, head);

4.2 Warp-level 归约

代码中使用 warp-level 归约优化归一化和 Softmax 等操作：

// 使用 __shfl_xor_sync 进行 warp 内归约
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
    tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32);
}

4.3 内存访问优化

代码中使用合适的内存访问模式，减少内存访问开销：

// 连续访问内存
for (int i = 0; i < ncols; i += WARP_SIZE) {
    const int col = i + tid;
    const float xi = x[row * ncols + col];
    tmp += xi * xi;
}

4.4 使用 cuBLAS 库

代码中使用 cuBLAS 库优化矩阵乘法和批处理矩阵乘法：

CUBLAS_CHECK(cublasSgemm(
        cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K,
        &alpha, src0, K, src1, K, &beta, dst, N));

CUBLAS_CHECK(cublasSgemmStridedBatched(
        cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K, &alpha, srck, embd,
        head_embd, srcq, embd, head_embd, &beta, dst, N, M * N, head));

5. 与 CPU 实现的比较

GPU 实现与 CPU 实现（naive 模块）的主要区别：

5.1 并行度

• CPU 实现：使用 TaskSet 和多线程并行
• GPU 实现：使用 CUDA 内核函数和数千个线程并行

5.2 内存管理

• CPU 实现：直接访问主内存
• GPU 实现：需要在 GPU 内存和主内存之间传输数据

5.3 优化方法

• CPU 实现：使用 SIMD 指令和缓存优化
• GPU 实现：使用 CUDA 内核函数、warp-level 归约和 cuBLAS 库

5.4 功能覆盖

• CPU 实现：完整实现所有功能
• GPU 实现：部分功能未实现（如 MatmulInt4FloatPacked、MatmulInt8Float 等）

6. 性能考虑

6.1 内存传输开销

GPU 计算需要在 CPU 和 GPU 之间传输数据，这可能成为性能瓶颈。代码中使用 CUDA 流和异步操作减少这种开销：

cudaStream_t stream = handle->stream;

6.2 内核启动开销

CUDA 内核启动有一定开销，代码中尽量减少内核启动次数，将多个小操作合并为一个大操作：

// 一次内核启动处理多个元素
__global__ void llm_elemwise_compute_float_scale_gpu(
        float* src, float* dst, size_t len, float scale) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index < len) {
        dst[index] = src[index] * scale;
    }
}

6.3 线程分配和负载均衡

代码中根据计算需求，合理分配线程和线程块，确保负载均衡：

// 根据计算需求调整线程块大小
const dim3 block_dims(kBlockSize, kBlockSize, 1);
const dim3 block_nums(block_x, block_y, head);

6.4 内存访问模式

代码中优化内存访问模式，减少内存访问开销：

// 连续访问内存
for (int col = tid; col < cols; col += block_size) {
    const float val = src[col];
    max = val > max ? val : max;
}

7. 未实现的功能

代码中使用 NOImplementKernel 宏标记了未实现的功能：

NOImplementKernel(MatmulInt4FloatPacked);
NOImplementKernel(MatmulInt4WeightReorder);
NOImplementKernel(MatmulInt8Float);
NOImplementKernel(EmbeddingGetInt8Float);

这些功能在 CPU 实现中已经实现，但在 GPU 实现中尚未实现。这可能是因为：

1. 这些功能在 GPU 上实现复杂度高
2. 这些功能在 GPU 上性能提升有限
3. 开发资源有限，优先实现更重要

8. 内核注册与调用机制

GPU 内核的注册与调用机制是通过模板特化和宏定义实现的，这种设计使得框架可以在运行时根据内核 ID 选择合适的实现。

8.1 内核注册宏

#define PartialImplementKernel(kernel_id, fun)               \
    template <typename... Args>                              \
    struct Comp<KernelID::kernel_id, Args...> {              \
        static void exec(Args... args, cudaHandle* handle) { \
            return fun(std::forward<Args>(args)..., handle); \
        }                                                    \
    };

#define PartialImplementSpace(kernel_id, fun)        \
    template <typename... Args>                      \
    struct Space<KernelID::kernel_id, Args...> {     \
        static size_t get(Args... args) {            \
            return fun(std::forward<Args>(args)...); \
        }                                            \
    };

#define NOImplementKernel(kernel_id)                         \
    template <typename... Args>                              \
    struct Comp<KernelID::kernel_id, Args...> {              \
        static void exec(Args... args, cudaHandle* handle) { \
            INFER_ASSERT(0, "kernel not implement");         \
        }                                                    \
    };

这些宏定义了三种类型的内核注册：

1. PartialImplementKernel：注册已实现的计算内核
2. PartialImplementSpace：注册工作空间计算函数
3. NOImplementKernel：标记未实现的内核

8.2 内核注册列表

PartialImplementKernel(ElemwiseFloat, llm_elemwise_compute_float);
PartialImplementKernel(ElemwiseFloatScale, llm_elemwise_compute_float_scale);
PartialImplementKernel(
        ElemwiseBroadcastDim0Src1Float, llm_elemwise_broadcast_dim0_src1_compute_float);
PartialImplementKernel(NormFloat, llm_norm_compute_float);
PartialImplementKernel(RmsNormFloat, llm_rms_norm_compute_float);
PartialImplementKernel(EmbeddingGetInt4Float, llm_embedding_get_int4_float);
PartialImplementKernel(EmbeddingGetFloatFloat, llm_embedding_get_float_float);
PartialImplementKernel(SoftmaxFloat, llm_softmax_compute_float);
PartialImplementKernel(MatmulInt4Float, llm_matmul_compute_int4_float);
PartialImplementKernel(MatmulFloatFloat, llm_matmul_compute_float_float);
PartialImplementKernel(
        MatmulWithHeadStrideFloat, llm_matmul_compute_with_head_stride_float);
PartialImplementKernel(HeadBatchedMatmulFloat, llm_head_batched_matmul_compute_float);
PartialImplementKernel(DiagMaskFloat, llm_diag_mask_inf_float);
PartialImplementKernel(RopeFloat, llm_rope_compute_float);
PartialImplementKernel(GlmRopeFloat, llm_glm_rope_compute_float);
PartialImplementKernel(ScaleDiagMaskFloat, llm_scale_diag_mask_inf_float);
PartialImplementKernel(GlmGmask, llm_glm_gmask_inf_float);
PartialImplementKernel(PermuteFloat, llm_permute_compute_float);

//! multi query attention
PartialImplementKernel(
        MatmulWithHeadStrideQBroadCastKFloat,
        llm_matmul_compute_with_head_strideq_broadcastk_float);
PartialImplementKernel(
        HeadBatchedMatmulBroadCastVFloat, llm_head_batched_matmul_broadcastv_float);

PartialImplementSpace(MatmulInt4Float, llm_matmul_get_workspace_float);
PartialImplementSpace(MatmulFloatFloat, llm_matmul_get_workspace_float_float);

NOImplementKernel(MatmulInt4FloatPacked);
NOImplementKernel(MatmulInt4WeightReorder);
NOImplementKernel(MatmulInt8Float);
NOImplementKernel(EmbeddingGetInt8Float);

这些注册语句将内核 ID 与具体的实现函数关联起来，使得框架可以在运行时根据内核 ID 选择合适的实现。

8.3 内核调用机制

内核调用是通过 Comp 和 Space 模板类实现的：

template <KernelID Id, typename... Args>
struct Comp {
    static void exec(Args... args, cudaHandle* handle);
};

template <KernelID Id, typename... Args>
struct Space {
    static size_t get(Args... args);
};

这些模板类提供了统一的接口，而具体实现由模板特化提供。在运行时，框架可以根据内核 ID 选择合适的实现：

// 调用示例
gpu::Comp<KernelID::MatmulFloatFloat>::exec(
        dst, src0, bias, src1, M, N, K, workspace, size, handle);

9. CUDA 资源管理

9.1 cudaHandle 结构体

struct cudaHandle {
    cudaStream_t stream{nullptr};
    cublasHandle_t cublas_handle{nullptr};
};

cudaHandle 结构体管理 CUDA 资源，包括：

1. stream：CUDA 流，用于异步执行 CUDA 操作
2. cublas_handle：cuBLAS 句柄，用于调用 cuBLAS 库函数

这种设计使得框架可以在多个 CUDA 设备和多个 CUDA 流上并行执行计算，提高计算效率。

9.2 CUDA 流使用

代码中使用 CUDA 流执行异步操作，减少 CPU 和 GPU 之间的同步开销：

cudaStream_t stream = handle->stream;
llm_embedding_get_int4_float_gpu<<<grid, DequantizedBlockSize, 0, stream>>>(
        weights, index, dst, len_seq, embd, weight_stride);

9.3 cuBLAS 库使用

代码中使用 cuBLAS 库优化矩阵乘法和批处理矩阵乘法：

cublasHandle_t cublas_handle = handle->cublas_handle;
CUBLAS_CHECK(cublasSetStream(cublas_handle, stream));
CUBLAS_CHECK(cublasSgemm(
        cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K,
        &alpha, src0, K, src1, K, &beta, dst, N));

10. 条件编译

代码使用条件编译确保 GPU 代码只在启用 GPU 支持时编译：

#if ENABLE_GPU
// GPU 代码
#endif

这种设计使得框架可以在不支持 GPU 的平台上编译和运行，只使用 CPU 实现。

11. 错误处理

代码中使用宏和断言处理错误：

#define CUBLAS_CHECK(func)                                                               \
    do {                                                                                 \
        cublasStatus_t status = (func);                                                  \
        if (status != CUBLAS_STATUS_SUCCESS) {                                           \
            printf("CUBLAS API failed at line %d with error: %s (%d)\n", __LINE__,       \
                   _cudaGetErrorEnum(status), status);                                   \
            return;                                                                      \
        }                                                                                \
    } while (0)

#define CUDA_CHECK(func)                                                              \
    do {                                                                              \
        cudaError_t status = (func);                                                  \
        if (status != cudaSuccess) {                                                  \
            printf("CUDA API failed at line %d with error: %s (%d)\n", __LINE__,      \
                   cudaGetErrorString(status), status);                               \
            return;                                                                   \
        }                                                                             \
    } while (0)

这些宏检查 CUDA 和 cuBLAS 函数的返回值，如果发生错误，打印错误信息并返回。

12. 与其他模块的集成

GPU 内核模块与其他模块的集成主要通过以下方式：

12.1 与 kernel.h 的集成

kernel.h 定义了内核系统的接口，GPU 内核模块实现了这些接口，提供了 GPU 上的计算实现。

12.2 与 tensor.h 的集成

tensor.h 定义了张量数据结构，GPU 内核模块使用这些张量作为输入和输出。

12.3 与 naive 模块的集成

当 GPU 实现不可用时，系统会回退到 naive 模块的 CPU 实现。

13. 性能优化总结

GPU 内核模块使用了多种性能优化技术：

13.1 并行计算

• 使用 CUDA 内核函数和数千个线程并行计算
• 使用 cuBLAS 库优化矩阵乘法和批处理矩阵乘法

13.2 内存优化

• 使用合适的内存访问模式，减少内存访问开销
• 使用共享内存和寄存器优化内存访问

13.3 算法优化

• 使用 warp-level 归约优化归一化和 Softmax 等操作
• 使用批处理矩阵乘法优化注意力计算

13.4 资源管理

• 使用 CUDA 流执行异步操作，减少 CPU 和 GPU 之间的同步开销
• 使用 cuBLAS 库优化矩阵乘法和批处理矩阵乘法

14. 未来优化方向

基于当前实现，可以考虑以下优化方向：

14.1 实现未实现的功能

实现当前标记为 NOImplementKernel 的功能：

• MatmulInt4FloatPacked
• MatmulInt4WeightReorder
• MatmulInt8Float
• EmbeddingGetInt8Float

14.2 使用更高效的 GPU 算法

• 使用 Tensor Core 加速矩阵乘法
• 使用 Flash Attention 算法优化注意力计算
• 使用混合精度计算提高性能

14.3 优化内存使用

• 使用内存池减少内存分配和释放的开销
• 使用流水线并行减少内存占用
• 使用张量并行和模型并行处理大型模型

14.4 支持更多 GPU 平台

• 支持 AMD GPU（使用 HIP）
• 支持移动 GPU（使用 OpenCL 或 Vulkan）
• 支持多 GPU 并行计算

总结

InferLLM 框架中的 GPU 内核实现提供了大语言模型在 GPU 上运行所需的各种计算操作的实现。通过使用 CUDA 内核函数、warp-level 归约、cuBLAS 库等技术，GPU 内核模块在 GPU 上实现了高效的计算。内核注册与调用机制使得框架可以在运行时根据内核 ID 选择合适的实现，而条件编译确保 GPU 代码只在启用 GPU 支持时编译。

虽然当前实现已经覆盖了大部分功能，但仍有一些功能未实现，如 MatmulInt4FloatPacked、MatmulInt4WeightReorder 等。未来可以考虑实现这些功能，并使用更高效的 GPU 算法、优化内存使用、支持更多 GPU 平台等方向进行优化。