音乐播放器
sola的小屋
 
文章 标签
20
首页 首页
归档 归档
标签 标签
关于 关于
说说 说说
友人帐 友人帐
Powered by Gridea | Theme: Fog
载入天数...
载入时分秒...
总访问量:  |   访问人数:
阅
2025-04-30
17 min read
大语言模型 大语言模型推理 AI-Infra

[InferLLM大模型推理框架项目](24)kern中optimized模块整体分析(src/kern/optimized)

InferLLM 框架中 optimized 模块分析

InferLLM 框架中的 optimized 模块是针对不同硬件平台的优化实现，主要包括 ARM、RISC-V 向量扩展(RVV)和 x86 平台的优化。这些优化实现利用了各平台特有的指令集和特性，提高了大语言模型推理的性能。

1. 目录结构

optimized 模块按照硬件平台分为三个子目录：

optimized/
├── arm/       # ARM 平台优化
│   ├── kernel.cpp
│   ├── kernel.h
│   ├── optimized.h
│   └── quantize.h
├── rvv/       # RISC-V 向量扩展优化
│   ├── common.cpp
│   ├── common.h
│   ├── kernel.cpp
│   ├── kernel.h
│   ├── optimized.h
│   └── quantize.h
└── x86/       # x86 平台优化
    ├── common.h
    ├── kernel.cpp
    ├── kernel.h
    ├── optimized.h
    └── quantize.h

2. 核心功能模块

2.1 基础向量运算

每个平台都实现了一系列基础向量运算函数，这些函数是高级操作的基础：

2.1.1 元素级操作

// 向量加法
inline void elemwise_vector_add(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z);

// 向量乘法
inline void elemwise_vector_mul(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z);

// 向量缩放
inline void elemwise_vec_scale(
        const int n, const float* __restrict x, float scale, float* __restrict z);

2.1.2 激活函数

// SiLU 激活函数 (x * sigmoid(x))
inline void elemwise_vector_silu(
        const int n, const float* __restrict x, float* __restrict z);

// GELU 激活函数
inline void elemwise_vector_gelu(
        const int n, const float* __restrict x, float* __restrict z);

2.1.3 归约操作

// 求最大值
inline float reduce_max(const int n, const float* __restrict x);

// 求平方和
inline float reduce_square_sum(const int n, const float* __restrict x);

// 减去最大值后求指数和（用于 softmax）
inline float select_sub_max_and_reduce_sum(
        const int n, const float* __restrict x, float* __restrict y, const float max);

2.2 量化计算

每个平台都实现了量化和反量化操作，以及量化数据之间的点积计算：

// 4位整数量化的反量化
inline void dequantize_row_q4_0(const void* x, float* y, int k);

// 8位整数量化的反量化
inline void dequantize_row_q8_0(const void* x, float* y, int k);

// 4位整数量化与8位整数量化的点积
inline float vec_vec_dot_q40_with_q80(const int n, const void* vx, const void* vy);

2.3 矩阵运算

实现了矩阵乘法和注意力计算所需的特殊矩阵运算：

// 带偏移的矩阵乘法（用于注意力计算）
inline void compute_src_offset_embd_matmul(
        const float* __restrict srcq, uint32_t srcq_stride, const float* __restrict srck,
        uint32_t srck_stride, float* __restrict dst, uint32_t dst_rows,
        uint32_t dst_cols, uint32_t embd);

// 带不连续目标的矩阵乘法（用于注意力计算）
inline void comput_matmul_with_dst_uncontinue(
        float* __restrict dst, uint32_t dst_stride, const float* __restrict src0,
        uint32_t src0_stride, const float* __restrict src1, uint32_t src1_rows,
        uint32_t src1_cols, uint32_t embd);

3. 平台特定优化

3.1 ARM 平台优化

ARM 平台的优化主要利用 NEON 指令集进行向量化计算：

// ARM NEON 优化的向量加法
inline void elemwise_vector_add(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    int i = 0;
#if defined(__ARM_NEON)
    for (; i + 15 < n; i += 16) {
        float32x4_t vx0 = vld1q_f32(x + i);
        float32x4_t vy0 = vld1q_f32(y + i);
        float32x4_t vx1 = vld1q_f32(x + i + 4);
        float32x4_t vy1 = vld1q_f32(y + i + 4);
        float32x4_t vx2 = vld1q_f32(x + i + 8);
        float32x4_t vy2 = vld1q_f32(y + i + 8);
        float32x4_t vx3 = vld1q_f32(x + i + 12);
        float32x4_t vy3 = vld1q_f32(y + i + 12);
        float32x4_t vz0 = vaddq_f32(vx0, vy0);
        float32x4_t vz1 = vaddq_f32(vx1, vy1);
        float32x4_t vz2 = vaddq_f32(vx2, vy2);
        float32x4_t vz3 = vaddq_f32(vx3, vy3);
        vst1q_f32(z + i, vz0);
        vst1q_f32(z + i + 4, vz1);
        vst1q_f32(z + i + 8, vz2);
        vst1q_f32(z + i + 12, vz3);
    }
    for (; i + 3 < n; i += 4) {
        float32x4_t vx = vld1q_f32(x + i);
        float32x4_t vy = vld1q_f32(y + i);
        float32x4_t vz = vaddq_f32(vx, vy);
        vst1q_f32(z + i, vz);
    }
#endif
    for (; i < n; i++) {
        z[i] = x[i] + y[i];
    }
}

3.2 RISC-V 向量扩展优化

RISC-V 向量扩展(RVV)的优化利用了 RVV 指令集进行向量化计算：

// RVV 优化的向量加法
inline void elemwise_vector_add(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    vadd(n, x, y, z);
}

// RVV 优化的向量乘法
inline void elemwise_vector_mul(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    vmul(n, x, y, z);
}

RVV 模块还包含了一些辅助函数，如 vadd、vmul、vmax 等，这些函数封装了 RVV 指令集的使用。

3.3 x86 平台优化

x86 平台的优化主要利用 SSE、AVX 和 AVX2 指令集进行向量化计算：

// x86 AVX 优化的向量乘法
inline void elemwise_vector_mul(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    const int nb32 = n / 32;
    const int nb8 = (n - nb32 * 32) / 8;
    const int left = n - nb32 * 32 - nb8 * 8;
    for (int b32 = 0; b32 < nb32; b32++) {
        const float* __restrict x32 = (const float*)x + b32 * 32;
        const float* __restrict y32 = (const float*)y + b32 * 32;
        float* __restrict z32 = z + b32 * 32;
        __m256 vx0 = _mm256_loadu_ps(x32 + 0);
        __m256 vx1 = _mm256_loadu_ps(x32 + 8);
        __m256 vx2 = _mm256_loadu_ps(x32 + 16);
        __m256 vx3 = _mm256_loadu_ps(x32 + 24);
        __m256 vy0 = _mm256_loadu_ps(y32 + 0);
        __m256 vy1 = _mm256_loadu_ps(y32 + 8);
        __m256 vy2 = _mm256_loadu_ps(y32 + 16);
        __m256 vy3 = _mm256_loadu_ps(y32 + 24);

        __m256 vz0 = _mm256_mul_ps(vx0, vy0);
        __m256 vz1 = _mm256_mul_ps(vx1, vy1);
        __m256 vz2 = _mm256_mul_ps(vx2, vy2);
        __m256 vz3 = _mm256_mul_ps(vx3, vy3);

        _mm256_storeu_ps(z32 + 0, vz0);
        _mm256_storeu_ps(z32 + 8, vz1);
        _mm256_storeu_ps(z32 + 16, vz2);
        _mm256_storeu_ps(z32 + 24, vz3);
    }
    // ... 处理剩余元素 ...
}

x86 模块还包含了一些辅助函数，如 exp256_ps 等，这些函数实现了数学函数的向量化版本。

4. 实现策略分析

4.1 分块处理

所有平台的优化实现都采用了分块处理策略，将数据分成多个块进行处理：

  1. 大块：使用最宽的向量指令处理（如 x86 上的 32 个元素一块）
  2. 中块：使用较窄的向量指令处理（如 x86 上的 8 个元素一块）
  3. 剩余元素：使用标量指令处理

这种策略可以最大化利用向量指令的并行性，同时处理所有数据。

4.2 内存访问优化

优化实现注重内存访问模式，尽量使用连续的内存访问：

// 连续加载和存储
__m256 vx0 = _mm256_loadu_ps(x32 + 0);
__m256 vx1 = _mm256_loadu_ps(x32 + 8);
// ...
_mm256_storeu_ps(z32 + 0, vz0);
_mm256_storeu_ps(z32 + 8, vz1);

4.3 计算优化

优化实现使用了各种计算优化技术：

  1. 向量化：使用 SIMD 指令并行处理多个数据
  2. 指令级并行：安排指令顺序，减少依赖
  3. 数学函数优化：使用近似计算或查表法加速数学函数

4.4 回退机制

所有优化实现都包含回退机制，当特定平台的指令集不可用时，使用标量实现：

inline void elemwise_vector_add(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    int i = 0;
#if defined(__ARM_NEON)
    // NEON 优化实现
    // ...
#endif
    // 标量回退实现
    for (; i < n; i++) {
        z[i] = x[i] + y[i];
    }
}

5. 性能优化分析

5.1 ARM 平台性能优化

ARM 平台的性能优化主要体现在：

  1. 使用 NEON 指令集进行向量化计算
  2. 每次处理 16 个元素（4个 float32x4_t 寄存器）
  3. 针对不同数据大小使用不同的处理策略

ARM 平台的优化相对简单，主要依赖 NEON 指令集的基本向量操作。

5.2 RISC-V 向量扩展性能优化

RISC-V 向量扩展的性能优化主要体现在：

  1. 使用 RVV 指令集进行向量化计算
  2. 利用 RVV 的可变长度向量特性，适应不同硬件实现
  3. 封装底层 RVV 指令，提供简洁的接口

RVV 的优化利用了 RISC-V 向量扩展的灵活性，可以适应不同的硬件实现。

5.3 x86 平台性能优化

x86 平台的性能优化最为复杂和全面：

  1. 使用 AVX/AVX2 指令集进行向量化计算
  2. 每次处理 32 个元素（4个 __m256 寄存器）
  3. 实现了复杂数学函数的向量化版本（如 exp256_ps）
  4. 针对不同操作使用不同的优化策略

x86 平台的优化利用了 x86 丰富的 SIMD 指令集和寄存器资源，实现了最全面的优化。

6. 与 naive 实现的比较

optimized 模块与 naive 模块的主要区别：

  1. naive 模块使用标量实现，optimized 模块使用向量化实现
  2. naive 模块适用于所有平台，optimized 模块针对特定平台优化
  3. naive 模块实现简单直观，optimized 模块实现复杂但性能更高
  4. naive 模块作为参考实现和后备方案，optimized 模块作为主要实现

在实际应用中，系统会优先使用 optimized 模块的实现，当特定平台的优化实现不可用时，会回退到 naive 模块的实现。

7. 多线程并行策略

optimized 模块使用 TaskSet 实现多线程并行，每个计算函数都返回一个 TaskSet，包含一个或多个任务及其子任务数量：

TaskSet llm_rms_norm_compute_float(
        const float* src, float* dst, uint32_t seq_len, uint32_t embd, float eps) {
    auto task = [=](const TaskId& id) {
        for (uint32_t i = id.start; i < id.end; i++) {
            // 计算逻辑
        }
    };
    return TaskSet{{task, seq_len}};
}

这种设计使得计算任务可以在多线程环境中高效执行，通过将大型计算任务分解为多个子任务，并分配给不同的线程处理。

7.1 任务分解策略

不同的计算函数使用不同的任务分解策略：

  1. 按序列长度分解：如 llm_rms_norm_compute_float，每个任务处理一个或多个序列位置
  2. 按头数分解：如 llm_matmul_compute_with_head_stride_float，每个任务处理一个或多个注意力头
  3. 按矩阵行列分解：如 llm_matmul_compute_int4_float，使用两个任务集，一个按行分解，一个按列分解

7.2 多阶段任务

某些计算函数使用多阶段任务，如 llm_matmul_compute_int4_float：

TaskSet 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) {
    // 第一阶段：量化输入
    auto task1 = [=](const TaskId& id) {
        // 量化逻辑
    };
    // 第二阶段：计算矩阵乘法
    auto task2 = [=](const TaskId& id) {
        // 矩阵乘法逻辑
    };
    return TaskSet{{task1, M}, {task2, N}};
}

这种设计使得不同阶段的任务可以并行执行，提高计算效率。

8. 内存优化策略

8.1 工作空间管理

optimized 模块使用工作空间进行临时计算，减少内存分配和释放的开销：

size_t llm_matmul_get_workspace_float(uint32_t, uint32_t M, uint32_t N, uint32_t K) {
    return sizeof(float) * K * M;
}

工作空间由调用者分配和管理，计算函数只负责使用。

8.2 量化计算

optimized 模块使用量化计算减少内存占用和计算量：

TaskSet 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) {
    // 量化输入
    auto task1 = [=](const TaskId& id) {
        for (uint32_t m = id.start; m < id.end; m++) {
            BlockQ80* q_src1 = (BlockQ80*)(static_cast<uint8_t*>(workspace) +
                                           m * weight_q80_stride);
            quantize_row_q8_0(src1 + m * K, q_src1, K);
        }
    };
    // 使用量化数据计算
    auto task2 = [=](const TaskId& id) {
        // 使用 vec_vec_dot_q40_with_q80 计算点积
    };
    return TaskSet{{task1, M}, {task2, N}};
}

这种设计减少了内存占用和内存带宽需求，提高了计算效率。

8.3 内存访问模式

optimized 模块优化了内存访问模式，减少内存访问开销：

TaskSet 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) {
    // 按块处理，每次处理 4 列
    auto task2 = [=](const TaskId& id) {
        uint32_t N_len = id.end - id.start;
        uint32_t n_block_4 = N_len / 4;
        uint32_t n_block_4_left = N_len - n_block_4 * 4;
        for (uint32_t block4 = 0; block4 < n_block_4; block4++) {
            uint32_t n = block4 * 4 + id.start;
            // 处理 4 列
            for (uint32_t m = 0; m < M; m++) {
                int8_t* src = q_src + m * weight_q80_stride;
                dst[m * N + n] = vec_vec_dot_q40_with_q80(K, q_weight0, src) + b0;
                dst[m * N + n + 1] = vec_vec_dot_q40_with_q80(K, q_weight1, src) + b1;
                dst[m * N + n + 2] = vec_vec_dot_q40_with_q80(K, q_weight2, src) + b2;
                dst[m * N + n + 3] = vec_vec_dot_q40_with_q80(K, q_weight3, src) + b3;
            }
        }
        // 处理剩余列
    };
}

这种设计提高了缓存命中率，减少了内存访问开销。

9. 注意力计算优化

optimized 模块对注意力计算进行了特殊优化：

9.1 多头注意力

TaskSet 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) {
    uint32_t sub_embd = embd / head;
    uint32_t length = nr_past + seqlen;
    uint32_t line_stride = embd;

    auto task = [=](const TaskId& id) {
        for (uint32_t h = id.start; h < id.end; h++) {
            auto dst_head = dst + h * seqlen * (nr_past + seqlen);
            auto srck_head = srck + h * sub_embd;
            auto srcq_head = srcq + h * sub_embd;
            compute_src_offset_embd_matmul(
                    srcq_head, embd, srck_head, embd, dst_head, seqlen, length,
                    sub_embd);
        }
    };
    return TaskSet{{task, head}};
}

这个函数实现了多头注意力中的 Q 和 K 的矩阵乘法，每个任务处理一个注意力头。

9.2 注意力输出计算

TaskSet 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) {
    uint32_t sub_embd = embd / head;
    uint32_t length = nr_past + seqlen;
    uint32_t line_stride = embd;

    auto task = [=](const TaskId& id) {
        for (uint32_t h = id.start; h < id.end; h++) {
            float* dst_head = dst + h * sub_embd;
            const float* v_head = v + h * sub_embd;
            const float* qk_head = qk + h * seqlen * length;
            comput_matmul_with_dst_uncontinue(
                    dst_head, embd, v_head, embd, qk_head, seqlen, length, sub_embd);
        }
    };
    return TaskSet{{task, head}};
}

这个函数实现了多头注意力中的 QK 和 V 的矩阵乘法，每个任务处理一个注意力头。

10. 矩阵乘法优化

矩阵乘法是大语言模型中最耗时的操作，optimized 模块对矩阵乘法进行了特殊优化：

10.1 量化矩阵乘法

TaskSet 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) {
    // 量化输入
    auto task1 = [=](const TaskId& id) {
        for (uint32_t m = id.start; m < id.end; m++) {
            BlockQ80* q_src1 = (BlockQ80*)(static_cast<uint8_t*>(workspace) +
                                           m * weight_q80_stride);
            quantize_row_q8_0(src1 + m * K, q_src1, K);
        }
    };
    // 使用量化数据计算
    auto task2 = [=](const TaskId& id) {
        // 按块处理，每次处理 4 列
        uint32_t N_len = id.end - id.start;
        uint32_t n_block_4 = N_len / 4;
        uint32_t n_block_4_left = N_len - n_block_4 * 4;
        for (uint32_t block4 = 0; block4 < n_block_4; block4++) {
            // 处理 4 列
        }
        // 处理剩余列
    };
    return TaskSet{{task1, M}, {task2, N}};
}

这个函数实现了 4 位整数权重与浮点数激活值的矩阵乘法，使用量化计算减少内存占用和计算量。

10.2 特殊矩阵乘法

void compute_src_offset_embd_matmul(
        const float* __restrict srcq, uint32_t srcq_stride, const float* __restrict srck,
        uint32_t srck_stride, float* __restrict dst, uint32_t dst_rows,
        uint32_t dst_cols, uint32_t embd);

void comput_matmul_with_dst_uncontinue(
        float* __restrict dst, uint32_t dst_stride, const float* __restrict src0,
        uint32_t src0_stride, const float* __restrict src1, uint32_t src1_rows,
        uint32_t src1_cols, uint32_t embd);

这些函数实现了特殊的矩阵乘法，用于注意力计算中的特殊需求。

11. 未来优化方向

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

11.1 更多硬件平台支持

  • 支持更多 ARM 架构（如 ARMv8.2-A 的 SVE 指令集）
  • 支持更多 RISC-V 扩展（如 RISC-V P 扩展）
  • 支持更多 x86 指令集（如 AVX-512）

11.2 更高效的算法

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

11.3 更多量化方法

  • 支持 3 位、2 位甚至 1 位量化
  • 支持非对称量化
  • 支持组量化

11.4 更高级的并行策略

  • 使用流水线并行减少内存占用
  • 使用张量并行和模型并行处理大型模型
  • 使用异步计算提高计算效率

总结

InferLLM 框架中的 optimized 模块提供了针对不同硬件平台的优化实现，包括 ARM、RISC-V 向量扩展和 x86 平台。这些优化实现利用了各平台特有的指令集和特性，提高了大语言模型推理的性能。

optimized 模块使用多线程并行、量化计算、内存优化和特殊算法等技术，实现了高效的计算。与 naive 模块相比，optimized 模块的实现更复杂，但性能更高。在实际应用中，系统会优先使用 optimized 模块的实现，当特定平台的优化实现不可用时，会回退到 naive 模块的实现。

未来可以考虑支持更多硬件平台、使用更高效的算法、支持更多量化方法和使用更高级的并行策略等方向进行优化，进一步提高大语言模型推理的性能。

赏
  • 文章作者：sola
  • 原文链接：https://solahome.github.io/post/InferLLM-Large-Language-Model-infer-framework-project-24/
  • 版权声明：本博客所有文章除特别声明外，均采用BY-NC-SA 许可协议。转载请注明出处！