简介

本文详细介绍了Qwen3-0.6B大模型在CPU上的推理优化方法,包括利用Intel AMX指令集优化矩阵乘法、优化Attention机制、采用查表法加速exp()和silu()函数以及算子融合等技术。通过这些优化,作者使模型推理速度比PyTorch+MKL backend+FlashAttention2的实现提升了约1.5倍,在序列长度为512的情况下,推理时间从0.81秒降至0.56秒。文章提供了详细的代码实现和优化思路,对大模型CPU部署具有实用参考价值。

最近一直在做 Qwen3-0.6B 大模型的 CPU 推理优化。任务是一个语义分类任务,用了 classification header,所以只有 prefill 阶段而没有 decode 阶段,且输入的 token 最大数量为 512。

因为模型很小,所以想试试看在 CPU 上跑能不能成。打算跑在 AWS m7i large 实例上,它配有双核的 Xeon 8488C CPU 和 8G 内存的虚拟化实例,支持 AMX 指令集扩展,花费是 $0.1 每小时。

权重文件是微调训练好的 bf16 的权重。以上是任务概况。

Qwen3-0.6B 模型的基本结构如图所示。可以看到它包含 28 层。主要耗时的部分是 Attention 部分和 Feed forward 部分。

本文主要对这两部份分别优化。我的最终结果比 pytorch+MKL backend+FlashAttention2 的实现快了大约 1.5 倍。

Qwen3-0.6B 模型的基本结构

一、矩阵乘法的优化

首先是对于矩阵乘的优化。这里我将对标 Intel Math Kernal Library(MKL)。

该软件库号称是 intel 平台上最快的矩阵乘法实现。并且该软件库在支持 AMX 的平台上会默认使用 AMX 指令集。AMX 扩展指令集支持多种数据类型。

对于 bf16 类型,tdpbf16ps 能在 52cycles 完成最大 mnk=(16,16,32)的矩阵乘法。相当于 315 flops/cycle。

然而这条指令对于右侧矩阵的行列排列不同于常规的矩阵行列排列方式,需要将数据进行转换才能得到正确的结果。

MKL 的 gemm 为了接口兼容性,会在函数内部进行重排。而我可以在加载权重的时候就重排好。显然这种方式比 MKL gemm 高效一些。

tdpbf16ps 右矩阵(32 行 16 列)数据排列方式

以下是用来对矩阵数据进行转换的代码。值得一提的是,第二个函数我一开始是让 ChatGPT 写的。

它当时写的并不对。而且糟糕的是,它前 8 行 8 列,后 8 行 8 列是对的。这就导致我以为它写的是对的,结果浪费了好几天查 bug。

后来我指出了它的错误,又让它重写了一遍,可还是不对。算了还是自己写吧。

// Convert from linear matrix to amx matrix

其次 amx 可以使用最多 8 个 tmm 寄存器。因此可以同时计算 mnk=(32,32,32)的子矩阵。

这样会节约一部分将数据载入 tmm 寄存器的时间。为了达到最高效率,我限制了输入矩阵的长宽必须是 32 的整数倍。

对于权重部分这总是满足的。而对于 sequence length,我规定输入的 token 必须左 padding 到 32 的整数倍。这样做对 attention 部分也是有好处的。

大矩阵乘法通常为了适配 CPU L1,L2 缓存,会做分块。也就是说对于不同的 CPU 最优的分块方案通常是不同的。

4th Gen Intel® Xeon® 8488C 有 48KB 的 L1,2MB 的 L2 和高达 105MB 的共享 L3。

为此我写了一个测试软件,对于该模型中所有出现过的矩阵大小,使用所有可能的分块方案以及 M,N,K 循环次序方案,最终为每种矩阵乘找出最快的那种方案。

该模型中出现过的的大矩阵乘以及其 MNK 值如下表所列:

  • q_proj(512,2048,1024)
  • k_proj(512,1024,1024)
  • v_proj(512,1024,1024)
  • o_proj(512,1024,2048)
  • up_proj(512,3072,1024)
  • gate_proj(512,3072,1024)
  • down_proj(512,1024,3072)

为了提高精度,在对 K 进行分块时,会使用 float 类型的临时缓冲区。同时我准备了四个个不同的 gemm 函数,在不同的时候会选用其中之一。

// a = (M, K), b = (K, N), C = (M, N), N, M and K should be divisible by 32, c_stride = b_stride

最终,我的代码相比于 MKL 获得了平均 30% 的提速。

二、Attention 的优化

qwen3-0.6B attention 有 16 个 header,每个 header 的 shape 是(seqlen,128)其中我约定了 seqlen 是 32 的整数倍且最大长度为 512。

依次计算 16 个 header,并且每两个 header 共享相同的 k,v 分片,这样对 cache 更友好,也更快。

主要代码如下:

voidqwen3_attn(bf16* x, struct infer_ctx& ctx, struct layer_weight* lw)

attention header optimization

伪代码差不多是这样的:

voidqwen3_attn_header(struct infer_ctx& ctx, int header_id)

三、其他优化

对于 bf16 的 exp() 和 silu() 函数,我试验下来用查表法明显快于泰勒级数展开法。因此我为每个函数 build 了 float[65536] 的表。

虽然表大小为 256KB 但实际访问不了那么多空间。对 Cache 的影响是有限的。

查表可以使用 _mm512_i32gather_ps() 条指令:

__m256i bf16_0 = _mm256_load_si256((__m256i*)(off));

算子融合意思是在一次循环中,尽可能完成多种计算。

最终如果使用阿里的官方推理代码,在 seqlen=512 的情况下,当使用 eager 实现时耗时 1.3s,sdpa 实现时耗时 0.81s,而我的代码耗时 0.56s。都已经看到这里了,不点个赞吗?

四、AI大模型学习和面试资源

我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

这份完整版的大模型 AI 学习和面试资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

在这里插入图片描述

第一阶段: 从大模型系统设计入手,讲解大模型的主要方法;

第二阶段: 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用;

第三阶段: 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统;

第四阶段: 大模型知识库应用开发以LangChain框架为例,构建物流行业咨询智能问答系统;

第五阶段: 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型;

第六阶段: 以SD多模态大模型为主,搭建了文生图小程序案例;

第七阶段: 以大模型平台应用与开发为主,通过星火大模型,文心大模型等成熟大模型构建大模型行业应用。

在这里插入图片描述

👉学会后的收获:👈

• 基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;

• 能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;

• 基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;

• 能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习框架等技术,这些技术的掌握可以提高程序员的编码能力和分析能力,让程序员更加熟练地编写高质量的代码。

在这里插入图片描述

1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集

👉获取方式:
😝有需要的小伙伴,可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓

在这里插入图片描述

Logo

码道开发者社区,聚焦华为云码道 CodeArts 代码智能体,沉淀 Agent、Skill、鸿蒙开发实战内容,供开发者查阅资料、交流技术、分享工程实践

更多推荐