今天给大家带来 华为&清华大学尹首一教授团队 最新的大模型论文。
【论文标题】
Tackling the Dynamicity in a Production LLM Serving System with SOTA Optimizations via Hybrid Prefill/Decode/Verify Scheduling on Efficient Meta-kernels
【论文链接】https://arxiv.org/pdf/2412.18106
1
摘要
要满足生产级大型语言模型(LLM)服务系统对低延迟和低成本高效率日益增长的需求,需要集成先进的优化技术。
然而,LLM 的输入输出长度动态且不可预测,再加上这些优化技术,加剧了工作负载的可变性,使得人工智能加速器,尤其是采用tile-based(基于分块)编程模型的 DSA,难以保持高效率。
为了应对这一挑战,本文推出了 XY-Serve,这是一个多功能、昇腾(Ascend) 原生、端到端生产型 LLM 服务系统。
其核心理念是一种抽象机制,通过将计算分解为统一的、硬件友好的、细粒度meta primitives来平滑工作负载的可变性。
对于注意力,本文提出了一种元内核,可以计算具有架构感知瓦片大小的 matmul-softmax-matmul 基本模式。
对于 GEMM,提出了一种虚拟填充方案,它能适应动态形状变化,同时使用具有各种固定块大小的高效 GEMM 基元。
实验结果表明,与目前在 Ascend NPU 上公开的基线相比,端到端吞吐量提高了 89%。
此外,与现有库相比,本文的方法优于现有的 GEMM 内核(平均快 14.6%)和 注意力内核(平均快 21.5%)。
虽然这项工作是 Ascend 本机开发的,但研究人员相信这种方法也能很容易地应用于 SIMT 架构。
2
背景
LLMs在自然语言处理和计算机视觉等领域广泛应用,但随着模型规模和输入序列增长,其推理成本和延迟成为挑战。
为应对这些挑战,多种优化技术应运而生,但也带来新的复杂性,尤其在生产系统中,如何在缺乏友好 SIMT 编程模型的 AI 加速器上高效集成这些优化技术成为关键问题。
优化技术带来的复杂性
-
自动前缀缓存(APC)
新查询可重用缓存的(K/V)数据,减少计算量,但增加了输入提示长度的可变性,因为缓存前缀数量取决于查询历史和实时内存可用性。
-
推测解码(SD)
使用草稿模型使目标模型每步可生成多个tokens,缓解解码阶段的内存瓶颈,但改变了解码阶段处理方式,使其不再一次处理一个token,而是处理动态变化的推测token长度,且注意力掩码变得更复杂,还导致验证阶段缺少特定核,只能使用预填充注意力核作为回退,无法进行特定优化。
-
SplitFuse
将长序列预填充tokens拆分为较小块并与解码tokens一起调度,减少对解码阶段的中断,但将预填充和解码阶段合并为单个批次,使批次内多阶段管理更复杂。
对计算模块的挑战
-
注意力模块
-
新技术引入使注意力形状和掩码结构更多样,降低当前 NPU 注意力内核的 MFU,增加 MBU。
-
不同阶段的注意力操作复杂性不同,处理不同阶段和批次组合的优化工作繁重且不切实际,常见的批处理执行方式虽简单但会引入额外内存开销,降低系统性能和计算资源利用率。
-
通用矩阵乘法(GEMM)模块
LLMs 中 GEMM 操作的维度(M、N、K)动态变化,设计支持任意形状且高性能的矩阵乘法在 AI 加速器上极具挑战性,通用 GEMM 内核性能不稳定,无法在所有操作中实现最优利用。
AI 加速器面临的挑战
以华为 Ascend NPU 为例,其基于 DaVinci 架构,AIC 处理矩阵计算,AIV 负责向量操作,两者无直接数据路径,数据交互依赖 L2 缓存,这在设计混合核时需重点考虑。
与 GPU 相比,NPU 核心粒度更大,负载均衡更关键。且其基于 tile 的编程模型在处理动态工作负载时面临挑战,传统解决方案中的填充/ 解填充操作会导致计算和内存浪费,产生大量开销。
3
贡献
为应对生产级大语言模型服务系统中动态输入输出长度和优化技术带来的挑战,论文提出了 XY-Serve 系统,其主要贡献可总结为以下四点:
-
提出 XY-Serve 系统
设计了一种多功能、Ascend 原生的端到端生产 LLM 服务系统 XY-Serve,通过抽象机制将计算分解为统一、硬件友好的细粒度元原语,以应对 LLM 推理中动态输入输出长度和优化技术导致的工作负载变化问题,提高 AI 加速器效率。
-
优化注意力计算
-
提出元注意力内核,支持自动前缀缓存、分块预填充和推测解码等高级特性,通过基数树管理 K/V 缓存实现tokens级重用,并能处理不同类型推测解码的掩码,提高了注意力计算的灵活性和效率。
-
设计了基于 tile 的 Cube-Vector 编排流水线,根据序列长度优化计算流程,利用掩码稀疏性跳过冗余计算,在不同注意力场景下(如预填充、解码等)均实现了性能提升,平均比现有实现快 21.5%。
- 优化 GEMM 计算
-
针对 GEMM 操作引入虚拟填充方案,结合选择性读写机制,在处理动态形状变化时避免了实际填充开销,使矩阵计算能无缝处理任意形状输入,同时保持固定形状优化的性能优势,平均性能提升 14.6%。
-
通过离线优化确定不同形状的高效任务分配策略,利用静态优化技术(如 swizzling)提升硬件利用率,减少内存操作开销,提高了 GEMM 计算在动态工作负载下的性能。
4
XY-Serve 系统
- Token-wise Scheduling
-
调度流程与策略
用户请求经 APC 模块处理后,未匹配tokens进入调度队列,系统从队列选取固定预算长度tokens组成块,优先预填充请求,可拆分长预填充提示,预留解码和推测tokens,调度以单个token为粒度。
-
动态性应对
队列中提示长度动态,单块内 P/D/V 阶段组成不可预测,通过灵活调度策略处理。
- Task Decomposition(任务分解)
-
总体机制
为解决动态工作负载在 AI 加速器执行问题,提出动态分解机制,将工作负载转换为 tile-based 计算单元,通过 Token-Table 逻辑分解各阶段,计算模块并行处理。
-
注意力分解
Token-Table 记录注意力各阶段关键属性,不同阶段分解为不同数量 tiling 块,每个块含多个 tiling 单元并行计算。
-
线性分解
线性操作中不同阶段tokens共享权重,连接成大张量计算,不区分阶段,对结果矩阵 tiling,不同操作有特定 tile 大小。
- Task Reordering(任务重排序)
-
重排序目的与流程
分解任务后重排序并生成 Task-Table 提升性能,各模块依 Task-Table 中 coreID 检索 tile 单元,简化核设计。
-
注意力重排序
计算 tile 负载并排序,按对称循环分配到 AI 核心平衡负载,调度信息存于 Task-Table 传至注意力模块。
-
线性重排序
离线优化确定线性操作任务分配策略,运行时根据 Token-Table 检索预优化 Task-Table 指导执行。
5
元注意力机制
- Meta-Attention Design(元注意力设计)
-
token级处理
使用基数树管理 K/V 缓存实现token级重用,支持前缀重用和分块预填充,可自动缓存历史 K/V 数据,通过复制写机制处理新数据,保证 K/V 缓存连续性,跟踪最后块tokens数确保准确处理。
-
推测解码掩码处理
针对序列和树型推测解码,分析掩码结构,对树型推测仅生成特定部分掩码传至内核,设计逐行处理掩码方法支持任意推测长度。
- Meta-Attention Optimizations(元注意力优化)
- 基于 tile 的 Cube-Vector 编排
-
提出流水线确保中间数据经 L2 缓存,避免 HBM 访问,根据序列长度采用不同计算流程,充分利用硬件资源,对解码任务优化计算流程,平衡 cube 和 vector 单元负载,提高内存带宽利用率。
-
利用掩码稀疏性
采用掩码感知策略跳过注意力计算中无效区域,在 QK、Softmax、SV 操作中利用掩码稀疏性减少计算量,支持任意形状掩码且无外部掩码生成传递开销,提高注意力机制效率和灵活性。
6
SmoothGEMM 优化
- 虚拟填充机制
-
M 维度虚拟填充
GEMM 中 M 维度与tokens数相关,为避免动态填充内存开销,采用芯片级虚拟填充结合选择性读写机制,按 tile 大小分配芯片内缓冲区,确保计算正确性,提升动态负载处理效率。
-
N 和 K 维度优化
线性操作中 N 和 K 维度由模型结构决定,无需额外填充;注意力 GEMM 中,N 和 K 维度与序列长度有关,由 K/V 缓存块页结构自然处理,无需额外开销,保证矩阵乘法对任意形状的高效支持。
- 动态形状处理与内存优化
-
动态形状处理
通过平滑技术将动态tokens大小转换为固定形状,离线优化特定形状,利用 swizzling 技术提升性能,运行时根据动态负载选择最佳配置,充分发挥硬件潜力。
-
内存开销去除
融合 Split 和 Permute 操作到 GEMM 计算,减少内存操作,对 QK 和 SV 计算采用 tile-based 和 stride-based 读写直接访问和存储数据,消除显式 Permute 操作,降低内存开销。
7
实验结果
元注意力性能
预填充注意力(Prefill Attention):在任意长度前缀重用场景下,随着系统前缀增加,XY-Serve 的元注意力内核处理时间减少,而 PromptFlashAttention(PFA)内核因不支持 PagedAttention 且 K/V 连接时间随重用长度增加而抵消了前缀重用的好处。即使排除 K/V 连接过程比较计算时间,元注意力内核平均性能也优于 PFA 内核 22.4%。
分块预填充长序列(Chunked Prefill with Long Sequences):对于长序列处理,XY-Serve 的 Chunked Prefill 方法在所有序列长度下性能均超越 PFA,纯计算时间比较中,改进高达 22.2%。
推测解码(Speculative Decoding):在不同上下文长度和批量大小的验证内核性能评估中,XY-Serve 的内核始终优于 PFA,随着上下文长度增加,性能提升更显著。排除 PFA 的连接操作后,元注意力内核平均性能仍提高 28.6%。
解码性能(Decode Performance):在解码阶段,XY-Serve 的内核在处理任意上下文长度和批量大小的情况下,相比 IncreFlashAttention(IFA)内核,在所有组合下均实现性能提升,平均改进 12.9%。
SmoothGEMM 性能
在处理用户输入长度从 1 到 4096(采用 Chunked Prefill 策略)的场景中,比较基于 Llama2-7B 和 Qwen2-7B 模型的线性算子(QKV、OProj、GateUp、Down)性能,SmoothGEMM 平均比 torch-npu linear 快 14.6%,且在序列大小超过 1k 时性能稳定,能达到理想 MFU。
端到端评估
vLLM 定期基准测试:使用 vLLM 社区的定期基准测试,以 Qwen2-7B 为模型,对比 Ascend-vLLM,XY-Serve 在未启用高级功能时,实现的 QPS 提升高达 79%,平均 TTFT 降低 64%,平均 TBT 降低 57%;启用预填充分块批处理和 P/D/V 融合优化后,QPS 提升至 89%,平均 TBT 延迟降低 69%,但在高吞吐量压力下的解码重场景中,启用 P/D/V 融合会使 TTFT 略有恶化。
Ascend NPUs 与 GPUs 对比:在 Qwen2-7B 和 Llama2-7B 模型的预填充和解码阶段,比较 XY-Serve 在 Ascend 910B 上与 vLLMv0.6.4.post1 在 Nvidia A800 上的端到端推理 MFU 和 MBU,XY-Serve 的 MFU 与 A800 相当,MBU 方面优势明显,提升高达 17%。
如何学习大模型 AI ?
由于新岗位的生产效率,要优于被取代岗位的生产效率,所以实际上整个社会的生产效率是提升的。
但是具体到个人,只能说是:
“最先掌握AI的人,将会比较晚掌握AI的人有竞争优势”。
这句话,放在计算机、互联网、移动互联网的开局时期,都是一样的道理。
我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。
我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。
第一阶段(10天):初阶应用
该阶段让大家对大模型 AI有一个最前沿的认识,对大模型 AI 的理解超过 95% 的人,可以在相关讨论时发表高级、不跟风、又接地气的见解,别人只会和 AI 聊天,而你能调教 AI,并能用代码将大模型和业务衔接。
- 大模型 AI 能干什么?
- 大模型是怎样获得「智能」的?
- 用好 AI 的核心心法
- 大模型应用业务架构
- 大模型应用技术架构
- 代码示例:向 GPT-3.5 灌入新知识
- 提示工程的意义和核心思想
- Prompt 典型构成
- 指令调优方法论
- 思维链和思维树
- Prompt 攻击和防范
- …
第二阶段(30天):高阶应用
该阶段我们正式进入大模型 AI 进阶实战学习,学会构造私有知识库,扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架,抓住最新的技术进展,适合 Python 和 JavaScript 程序员。
- 为什么要做 RAG
- 搭建一个简单的 ChatPDF
- 检索的基础概念
- 什么是向量表示(Embeddings)
- 向量数据库与向量检索
- 基于向量检索的 RAG
- 搭建 RAG 系统的扩展知识
- 混合检索与 RAG-Fusion 简介
- 向量模型本地部署
- …
第三阶段(30天):模型训练
恭喜你,如果学到这里,你基本可以找到一份大模型 AI相关的工作,自己也能训练 GPT 了!通过微调,训练自己的垂直大模型,能独立训练开源多模态大模型,掌握更多技术方案。
到此为止,大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗?
- 为什么要做 RAG
- 什么是模型
- 什么是模型训练
- 求解器 & 损失函数简介
- 小实验2:手写一个简单的神经网络并训练它
- 什么是训练/预训练/微调/轻量化微调
- Transformer结构简介
- 轻量化微调
- 实验数据集的构建
- …
第四阶段(20天):商业闭环
对全球大模型从性能、吞吐量、成本等方面有一定的认知,可以在云端和本地等多种环境下部署大模型,找到适合自己的项目/创业方向,做一名被 AI 武装的产品经理。
- 硬件选型
- 带你了解全球大模型
- 使用国产大模型服务
- 搭建 OpenAI 代理
- 热身:基于阿里云 PAI 部署 Stable Diffusion
- 在本地计算机运行大模型
- 大模型的私有化部署
- 基于 vLLM 部署大模型
- 案例:如何优雅地在阿里云私有部署开源大模型
- 部署一套开源 LLM 项目
- 内容安全
- 互联网信息服务算法备案
- …
学习是一个过程,只要学习就会有挑战。天道酬勤,你越努力,就会成为越优秀的自己。
如果你能在15天内完成所有的任务,那你堪称天才。然而,如果你能完成 60-70% 的内容,你就已经开始具备成为一名大模型 AI 的正确特征了。