llama2.c Apple M4 NEON 优化记录

环境信息

设备: Apple M4 (ARM64)
编译器: Apple Clang (gcc 前端)
模型: stories15M.bin (15M 参数)
模型维度: dim=288, hidden_dim=768, n_layers=6, n_heads=6, seq_len=256
编译选项: -O3 -march=armv8.4-a+fp16+dotprod

原始代码分析

三个核心函数

rmsnorm(): RMS 归一化，计算平方和后缩放
softmax(): 数值稳定的 softmax（先减最大值）
matmul(): 矩阵向量乘法 W(d,n) @ x(n)，占 95%+ 时间

原始性能

原始标量版本: 212 tok/s

优化阶段一：NEON 基础向量化

目标函数

rmsnorm(), softmax(), matmul()

采用手段

使用 ARM NEON 128-bit SIMD 指令，每次处理 4 个 float：

rmsnorm

平方和: vmlaq_f32 累加 4 个元素的平方
归一化: vmulq_f32 并行缩放

softmax

求最大值: vmaxq_f32 + vpmax_f32 水平归约
exp+累加: 向量减最大值后，4 个一组调用 expf
归一化: vmulq_f32 乘 1/sum

matmul

内积: vmlaq_f32 每次处理 4 对乘积
归约: vpadd_f32 水平求和

优化效果

NEON 基础版本: 635 tok/s (3.0x 加速)

优化阶段二：matmul 循环展开 + attention 向量化

问题分析

阶段一的 matmul 每次只处理 4 个 float，循环开销大。M4 有 4 个 FMA 流水线，可以并行发射多条指令。

采用手段

matmul 32-float 展开

将内层循环展开为 8 组 NEON 向量（32 个 float），使用 8 个独立累加器：

float32x4_t sum0 = vdupq_n_f32(0.0f);
float32x4_t sum1 = vdupq_n_f32(0.0f);
// ... sum7

for (; j + 32 <= n; j += 32) {
    // 加载 8 组 w 和 x
    sum0 = vmlaq_f32(sum0, w0, x0);
    sum1 = vmlaq_f32(sum1, w1, x1);
    // ... sum7
}
// 8 个累加器归约为 1 个

attention 内积向量化

forward() 中的两个标量内积循环（q·k 和 att·v）也用 NEON 展开到 16 float：

// q·k 点积
float32x4_t qk_sum = vdupq_n_f32(0.0f);
for (; i + 16 <= head_size; i += 16) {
    qk_sum = vmlaq_f32(qk_sum, q0, k0);
    qk_sum = vmlaq_f32(qk_sum, q1, k1);
    // ...
}

// att·v 加权累加（带 vst1q_f32 写回）
float32x4_t av = vdupq_n_f32(a);
for (; i + 16 <= head_size; i += 16) {
    xb0 = vmlaq_f32(xb0, av, v0);
    vst1q_f32(xb + i, xb0);
    // ...
}

优化效果

32-unroll + attention NEON: 1031 tok/s (4.9x 加速)
matmul 耗时: 106ms -> 87ms (部分提升来自 attention 优化减少了竞争)

优化阶段三：权重矩阵转置 + 列并行展开

问题分析

阶段二的 matmul 外层循环遍历输出行 i，内层遍历输入列 j：

每个 i 都要重新加载整个 x 向量
权重 w[i*n + j] 在 i 变化时跨行跳跃，cache 不友好

采用手段

1. 权重矩阵转置

在模型加载时（memory_map_weights 后），为每个权重矩阵创建转置副本：

// TransformerWeights 新增转置指针
float* wq_t; float* wk_t; float* wv_t; float* wo_t;
float* w1_t; float* w2_t; float* w3_t; float* wcls_t;

// 转置函数: src[rows][cols] -> dst[cols][rows]
void transpose_matrix(float* dst, float* src, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            dst[j * rows + i] = src[i * cols + j];
        }
    }
}

关键修正：必须正确传入 (rows, cols) 参数。原始权重存储为 (output_dim, input_dim)，即 rows=output_dim, cols=input_dim。

2. matmul 循环重构

转置后，matmul 改为外层遍历输入列 j，内层更新所有输出：

void matmul(float* xout, float* x, float* w_t, int n, int d) {
    // w_t 形状 (n, d)，即原始 W(d,n) 的转置
    memset(xout, 0, d * sizeof(float));
    for (int j = 0; j < n; j++) {
        float xj = x[j];
        for (int i = 0; i < d; i++) {
            xout[i] += w_t[j * d + i] * xj;
        }
    }
}

优势：

x[j] 只加载一次，广播到所有输出
w_t[j*d + i] 在 i 循环中顺序访问（stride=1）
xout[i] 也顺序访问

3. 4 列并行展开

每次处理 4 个输入列，复用 x 加载：

for (; j + 4 <= n; j += 4) {
    float32x4_t xv0 = vdupq_n_f32(x[j]);
    float32x4_t xv1 = vdupq_n_f32(x[j+1]);
    float32x4_t xv2 = vdupq_n_f32(x[j+2]);
    float32x4_t xv3 = vdupq_n_f32(x[j+3]);
    
    for (; i + 16 <= d; i += 16) {
        float32x4_t o0 = vld1q_f32(xout + i);
        // 加载 4 组 w，每组 4 个向量
        o0 = vmlaq_f32(o0, w00, xv0);
        o0 = vmlaq_f32(o0, w10, xv1);
        o0 = vmlaq_f32(o0, w20, xv2);
        o0 = vmlaq_f32(o0, w30, xv3);
        vst1q_f32(xout + i, o0);
        // ... o1, o2, o3
    }
}

踩坑记录

shared_weights 处理: 当 vocab_size > 0 时权重共享，不能直接用 token_embedding_table 作为 wcls_t，必须也创建转置副本
转置参数传反: w2 的形状是 (hidden_dim, dim)，不是 (dim, hidden_dim)，最初传反导致输出乱码
8 列展开尝试: 展开到 8 列后寄存器压力增大，发生 spill，性能从 1065 降到 961 tok/s

优化效果

转置 + 4列展开: 1065 tok/s (5.0x 加速)
matmul 耗时: 245ms -> 87ms (2.8x)

各阶段性能对比

版本	tok/s	相对原始	matmul 耗时
原始标量	212	1.0x	~246ms
NEON 基础 (4-unroll)	635	3.0x	~113ms
32-unroll + attention NEON	1031	4.9x	~87ms
转置 + 4列展开	1065	5.0x	~87ms
Accelerate CBLAS	943	4.4x	~99ms

Profiling 分析

最终版本的耗时分布（100 tokens）：

模块	耗时	占比
matmul	87.3 ms	95.7%
attention	2.8 ms	3.1%
rope	0.5 ms	0.6%
ffn_act	0.5 ms	0.5%
rmsnorm	0.1 ms	0.1%
residual	0.0 ms	~0%

结论: matmul 仍是绝对瓶颈，但已高度优化。其他模块占比极小，进一步优化收益有限。

优化阶段四：Apple Accelerate CBLAS

背景

尝试使用 SME 指令集时，发现 macOS 用户空间不支持 SME streaming mode（SMSTART 触发 SIGILL）。用户选择使用 Apple Accelerate 框架作为替代。

采用手段

使用 Accelerate 框架的 cblas_sgemv 替换手写 matmul：

#if defined(__APPLE__)
    #include <Accelerate/Accelerate.h>
#endif

void matmul(float* xout, float* x, float* w, int n, int d) {
#if defined(__APPLE__)
    cblas_sgemv(CblasRowMajor, CblasNoTrans, d, n, 1.0f, w, n, x, 1, 0.0f, xout, 1);
#else
    // 保留 NEON fallback...
#endif
}

同时清理了转置权重相关代码（移除 transpose_matrix、转置指针、分配/释放逻辑），减少内存占用和启动时间。

编译命令

gcc -O3 -march=armv8.4-a+fp16+dotprod -DACCELERATE_NEW_LAPACK \
    -o run run.c -lm -framework Accelerate

效果

Accelerate CBLAS: 943 tok/s (4.4x 加速)

对于 15M 小模型，CBLAS 略低于手写 NEON 转置版本（1065 tok/s），原因：

CBLAS 函数调用开销在小矩阵上较明显
小矩阵无法充分发挥 CBLAS 内部调度优势

但 CBLAS 版本的优势：

代码更简洁（一行替代几十行 NEON）
自动利用 Apple Silicon 所有硬件优化（AMX2、NEON、缓存预取）
对大模型（7B+）会显著优于手写 NEON
跨平台兼容（非 Apple 平台自动 fallback 到 NEON）

最终结论

版本	tok/s	推荐场景
原始标量	212	基准对照
NEON 基础	635	通用 ARM64
NEON 转置	1065	Apple M4 小模型最佳性能
Accelerate CBLAS	943	Apple M4 生产环境首选

生产环境推荐 Accelerate CBLAS：代码简洁、维护成本低、自动适配未来硬件、大模型性能更优。

代码简化：移除 rmsnorm/softmax 向量化

背景

Profiling 显示 rmsnorm 和 softmax 合计占比 <1%，向量化对整体性能影响极小。为保持代码简洁，将其还原为原始标量版本。

改动

rmsnorm(): 移除 NEON 平方和与 NEON 缩放，恢复原始标量循环
softmax(): 移除 NEON 最大值、NEON exp+累加、NEON 归一化，恢复原始标量循环
attention 中的 q·k 和 att·v 内积循环：保留 NEON 优化（占比 3.5%，有实际收益）

效果

简化后 CBLAS: ~950 tok/s (与简化前基本一致)

向量化移除后性能无显著变化，验证了 profiling 结论：这两个函数不是瓶颈。

进一步优化方向

更大的模型: 当前 15M 模型太小。CBLAS 在 7B+ 模型上会有更显著优势
量化: 参考 runq.c 的 int8 量化，结合 NEON 的 vdotq_s32 指令
多线程: 对 attention 和 FFN 中的独立操作使用更细粒度的并行

代码变更概要

新增文件/修改

run.c: 主要优化文件
- 添加 #include <Accelerate/Accelerate.h>（Apple 平台）
- rmsnorm(): 标量实现（简洁优先）
- softmax(): 标量实现（简洁优先）
- matmul(): Apple 平台使用 cblas_sgemv，其他平台保留 NEON 32-unroll
- forward(): attention 内积循环保留 NEON 优化

编译命令

# Apple Silicon (推荐)
gcc -O3 -march=armv8.4-a+fp16+dotprod -DACCELERATE_NEW_LAPACK \
    -o run run.c -lm -framework Accelerate

# 通用 ARM64
gcc -O3 -march=armv8.4-a+fp16+dotprod -o run run.c -lm

优化阶段五：跳过 Prompt 阶段的无用 Logits 计算

问题分析

在推理过程中，前 num_prompt_tokens - 1 个 token 的 logits 计算结果被直接丢弃：

// generate 函数中的循环
while (pos < steps) {
    float* logits = forward(transformer, token, pos);  // 总是计算 logits

    if (pos < num_prompt_tokens - 1) {
        // Prompt 阶段：强制使用下一个 prompt token，logits 完全未使用
        next = prompt_tokens[pos + 1];
    } else {
        // 生成阶段：从 logits sample
        next = sample(sampler, logits);
    }
}

被浪费的计算包括：

Final RMSNorm (O(dim))
Classifier matmul (O(dim × vocab_size))

对于 vocab_size ≈ 32000 的模型，classifier matmul 约占每次 forward 计算的 10-15%。

采用手段

1. 修改 forward 函数签名

添加 need_logits 参数：

float* forward(Transformer* transformer, int token, int pos, int need_logits) {
    // ... 所有层计算保持不变（必须执行以更新 KV cache）...

    // 条件执行 final rmsnorm 和 classifier
    if (need_logits) {
        rmsnorm(x, x, w->rms_final_weight, dim);
        matmul(s->logits, x, w->wcls, p->dim, p->vocab_size);
    } else {
        prof_logits_skipped++;
    }

    return need_logits ? s->logits : NULL;
}

2. 修改 generate 函数调用

while (pos < steps) {
    // Prompt 阶段只有最后一个 token 需要 logits
    int need_logits = (pos >= num_prompt_tokens - 1);
    float* logits = forward(transformer, token, pos, need_logits);

    if (pos < num_prompt_tokens - 1) {
        next = prompt_tokens[pos + 1];
    } else {
        next = sample(sampler, logits);
    }
}

3. 修改 chat 函数调用

chat 函数类似处理，确保最后一个 prompt token 能 sample 第一个回复：

int need_logits = (user_idx >= num_prompt_tokens);
float* logits = forward(transformer, token, pos, need_logits);
if (need_logits) {
    next = sample(sampler, logits);
}

性能测试结果

15M 模型（短 prompt: 8 tokens，生成 100 tokens）

指标	基准版本	优化版本	提升
total time	133.8 ms	80.8 ms	-39.5%
matmul time	129.2 ms	76.5 ms	-40.8%
logits skipped	0	7	-

110M 模型（中 prompt: 69 tokens，生成 85 tokens）

指标	基准版本	优化版本	提升
total time	919.1 ms	831.7 ms	-9.5%
matmul time	870.3 ms	784.6 ms	-9.9%
logits skipped	0	68	-
tok/s	167	184	+10%

110M 模型（长 prompt: 351 tokens，生成 47 tokens）

指标	基准版本	优化版本	提升
total time	2506.0 ms	2068.7 ms	-17.5%
matmul time	2215.2 ms	1790.4 ms	-19.2%
logits skipped	0	350	-
tok/s	158.7	191.9	+20.8%

收益分析

节省的计算量公式：

saved_flops = (num_prompt_tokens - 1) × dim × vocab_size

场景	Prompt 占比	total time 节省	tok/s 提升
短 prompt (8 tokens)	7%	-39.5%	-
中 prompt (69 tokens)	44%	-9.5%	+10%
长 prompt (351 tokens)	88%	-17.5%	+20.8%

结论

对于 长 prompt + 短生成 场景（典型的 RAG、文档问答），优化效果显著：

节省大量 classifier matmul 计算
tok/s 提升 20%+

特别适合 prefill-heavy 的应用场景。

附录：关键 NEON 指令说明

指令	功能
`vld1q_f32`	加载 4 个 float 到 128-bit 向量
`vst1q_f32`	存储 128-bit 向量到内存
`vdupq_n_f32`	广播标量到 128-bit 向量
`vmlaq_f32`	向量乘加: `a + b * c`
`vaddq_f32`	向量加法
`vmaxq_f32`	向量取最大值
`vpadd_f32`	64-bit 向量水平相加
`vget_low_f32` / `vget_high_f32`	提取向量的低/高 64-bit
`vget_lane_f32`	提取向量中的单个元素

FilesExpand file tree

optimization.md

Latest commit

History

optimization.md

File metadata and controls

llama2.c Apple M4 NEON 优化记录

环境信息

原始代码分析

三个核心函数

原始性能

优化阶段一：NEON 基础向量化

目标函数

采用手段

rmsnorm

softmax

matmul

优化效果

优化阶段二：matmul 循环展开 + attention 向量化

问题分析

采用手段

matmul 32-float 展开

attention 内积向量化

优化效果

优化阶段三：权重矩阵转置 + 列并行展开

问题分析

采用手段

1. 权重矩阵转置

2. matmul 循环重构

3. 4 列并行展开

踩坑记录

优化效果

各阶段性能对比

Profiling 分析

优化阶段四：Apple Accelerate CBLAS

背景

采用手段

编译命令

效果

最终结论

代码简化：移除 rmsnorm/softmax 向量化

背景

改动

效果

进一步优化方向

代码变更概要

新增文件/修改

编译命令

优化阶段五：跳过 Prompt 阶段的无用 Logits 计算

问题分析

采用手段

1. 修改 forward 函数签名

2. 修改 generate 函数调用

3. 修改 chat 函数调用

性能测试结果

15M 模型（短 prompt: 8 tokens，生成 100 tokens）

110M 模型（中 prompt: 69 tokens，生成 85 tokens）

110M 模型（长 prompt: 351 tokens，生成 47 tokens）

收益分析

结论

附录：关键 NEON 指令说明