Blog Tags GitHub
Back to blog
llm-inference kernel-optimization memory-efficiency kv-cache

FastTPS: An Optimized Method for LLM Token Phase for AI Accelerators

面向AI加速器的LLM Token阶段优化框架FastTPS，包括全局KV缓存管理、序列维度扁平化、Fusion MLP等技术

FastTPS: An Optimized Method for LLM Token Phase for AI Accelerators: 面向AI加速器的LLM Token阶段优化框架FastTPS，包括全局KV缓存管理、序列维度扁平化、Fusion MLP等技术

一、论文概述

属性内容
论文标题FastTPS: An Optimized Method for LLM Token Phase for AI Accelerators
论文编号arXiv:2607.11211
作者Wenzong Yang, Danyang Zhang, Kun Cao, Tejus Siddagangaiah, Rajeev Patwari, Zhanxing Pu, Siyin Kong, Zijiang Yang, Hao Zhu, Varun Sharma, Yue Gao, Tianping Li, Fan Yang, Jicheng Chen, Yushan Chen, Fennian Zhao, Aaron Ng, Elliott Delaye, Ashish Sirasao, Sudip Nag
标签LLM inference, token phase, AI accelerator, KV cache, operator fusion, tiling strategy, memory optimization
图表数8 张
可下载图片是

摘要： The popularity of large language models (LLMs) escalates an ongoing demand for effective inference. However, due to the sequential processing of tokens during the token phase in decoder-only LLMs inference, the inherent low parallelism leads to reduced throughput and suboptimal utilization of the computing units on artificial intelligence (AI) accelerators, particularly when handling long-sequence inputs that impose significant memory overhead. Recently, many reported methods have been developed…

二、核心思想

问题：Decoder-only LLM推理中token阶段（decoding phase）的顺序处理导致AI加速器并行度低、计算单元利用率不足，长序列输入时内存开销尤为严重。现有方法主要关注计算分配优化。方案：FastTPS框架从三个维度优化：(1) 全局KV缓存分层管理（跨L3/L2/L1层次）。2) TPSFLAT tiling策略沿序列维度展平计算。3) Fusion MLP算子融合减少内存流量。

三、技术架构

标准 decoder-only Transformer 的 token 阶段由 attention 块（Section ①）和 MLP 块（Section ②）组成，数据在 L1/L2/L3 三级缓存间频繁往返。

Figure 1. 标准 attention 与 MLP 块的结构概览：Section ① 为 attention 块，Section ② 为 MLP 块，图中标注了数据维度的变化

三个核心组件：

  1. Global KV Cache Management (GKVC)：预分配 L3 静态 KV Cache 内存区域，采用地址跳转策略由 AI 加速器（或软件）直接写入，消除 Concat 算子的内存重组，保证整个推理流水线驻留在加速器上。
  2. TPSFLAT：基于 tiling 的注意力计算方法。将大矩阵分解为 L3-level 和 L2-level 的 tiles，公式：(H/H_b) x (Q_b,K_b) = Tiling_L3_Tile(Q,K)，然后 (H_b) x (Q_i,K_i) = Tiling_L2_Tile(Q_b,K_b)；在 L1 内连续融合 RoPE、QKT、Softmax、PV 四个算子。
  3. Fusion MLP：将 gate/up 投影交错排布融合为一个大 Matmul，再与 SiLU、elewMul 融合，把 I/O 次数从 7 降到 2。

执行流程： Attention 块：L3 tile -> L2 tile -> 逐 block 计算 -> 输出合并 MLP 块：RMSNorm -> 交错 gate-up Matmul -> siLu_Mul -> down projection

核心公式：

  • Eq(1): O = W_down(SiLU(W_gateh_t) * (W_uph_t)) — 标准MLP
  • Eq(2): Step = h*(N_p+N_r)d + N_pd — GKVC 地址跳转步长
  • Eq(3): H/H_b x (Q_b,K_b) = Tiling_L3_Tile(Q,K) — L3级tiling
  • Eq(8): H_b x (Q_i,K_i) = Tiling_L2_Tile(Q_b,K_b) — L2级tiling

GKVC + TPSFLAT 优化后的 attention 块：

Figure 3. GKVC 和 TPSFLAT 优化后的 attention 块结构（红=L3，橙=L2，绿=L1 缓存）。GKVC 将 L1 计算出的 KV 直接写入 L3 预分配位置，无需 Concat；TPSFLAT 在 L1 内连续融合 RoPE、QKT、Softmax、PV

Fusion MLP 工作流：

Figure 4. Fusion MLP 工作流：part1 为 RMSNorm 归一化，part2 为交错 gate-up 权重的核心融合（流水线调度消除算子串行执行的时间跨度），part3 为 down projection 降维

四、核心创新

创新点描述
全局KV缓存管理跨层次内存（L3/L2/L1）的高效KV缓存分配与替换策略
TPSFLAT tiling序列维度展平的tiling策略，多级tile兼顾数值精度与硬件友好性
Fusion MLPRMSNorm + Gate + FFN三合一算子融合，显著减少内存流量
端到端2.5x加速在代表LLM上保持全精度的同时实现显著加速，最大误差<0.22%

五、代码实现分析

TPSFLAT实现：

  • 两级tiling策略：L3-level tiles（大块矩阵操作）和L2-level tiles（细粒度块内操作）
  • L3 tiling: H/H_b x (Q_b, K_b) -> 分布在L3 cache
  • L2 tiling: H_b x (Q_i, K_i) -> 分布到L2 cache
  • V_cache和sin/cos位置编码也按相同tiling分布

Fusion MLP实现：

  • Step 1: RMSNorm作用于输入h
  • Step 2: Gate操作(SiLU) applied to projected tokens
  • Step 3: Fused W_down + element-wise multiply -> output
  • 整个流程整合为单一kernel，极大减少L1-L3间的数据搬运

六、实验结果

在 AMD Ryzen AI 300 系列 NPU 上，对 ChatGLM3-6B、Llama2-7B、Llama3-8B、Llama3.2-1B、Phi3-mini-4k 等模型评测。

Attention Block 加速：

TPSFLAT 消除了 KV Cache Concat（占 66%~82%）与部分计算开销（18%~34%），Llama3(4096) 最大延迟降低 93%，其余模型 80%~92%。

Figure 5. 表 1 中不同模型的 attention（橙）与 TPSFLAT（绿）延迟对比。橙色条中浅橙为 KV Cache &#x27;Concat&#x27; 占比、深橙为其他运算；绿色条上方数字为 TPSFLAT 实现的延迟降低率

延迟与序列长度 NtN_t 呈二次关系（源于 Concat 的 O(2Ntd)O(2N_td) 复杂度）。GKVC 消除 Concat 后，加速比随 NtN_t 单调增长（Llama2 优化 70%~90%）。

Figure 6. Llama2 在不同最大序列长度下的延迟降低率——序列越长，KV Cache &#x27;Concat&#x27; 开销越大，加速收益越明显

MLP Block 加速：

Fusion MLP 将 I/O 从 7 降到 2，Llama3/ChatGLM/Llama2/Phi3 提速超 20%，Llama3.2 降低 68%。

Figure 7. 原始 MLP（橙）与 Fusion MLP（绿）延迟对比，绿色条上方为 Fusion MLP 相对原始 MLP 的提升比例

端到端 FastTPS 加速：

Phi3 达到最高 6.0× 加速，平均加速器资源利用率 86%（Phi3 高达 93%）。

OI（Operational Intensity）分析：

Figure 2. Llama2 中 attention 与 MLP 块的 Roofline 分析（含 GQA 优化）。Ori 为 FastTPS 优化前、Opt 为优化后；attention 块 OI 提升约 3×，Fusion MLP 更接近理论峰值

  • Attention 块通过消除 KV Cache 数据搬运、融合 Softmax/Matmul、合并两次 RoPE，OI 提升约 3×
  • MLP 块经优化后性能显著逼近带宽理论上限（受权重搬运限制）

精度保证：

FastTPS 用 bfloat16 模拟，以 PyTorch float 为基线，精度随序列长度增长反而提升（Softmax 输出趋于均匀）；相比 FlashAttention（~1e-2）高一个数量级（~1e-3）。

Figure 8. TPSFLAT 的最大误差随序列长度变化——误差随序列增长而下降，最大误差小于 0.22%

  • TPSFLAT 最大误差 < 0.22%
  • 全精度保持不变（no numerical degradation）

七、相关工作

  • Attention kernel optimization方法（Flash Attention等）
  • MLP operator fusion技术
  • KV cache management in LLM inference
  • Tiling strategies for GPU/AI accelerators
  • Hierarchical memory optimization in neural network inference

八、总结

贡献：

  • 提出FastTPS框架，专门针对AI加速器上的LLM token阶段进行全面优化
  • 设计Global KV Cache Management实现跨层次内存高效管理
  • 提出TPSFLAT tiling策略，兼顾数值精度与硬件友好性
  • 发明Fusion MLP算子融合技术，大幅减少内存流量

影响：

  • 为AI加速器上的LLM推理提供从内存到算子的全栈优化方案
  • 证明了tiling策略可以有效利用AI加速器的层次化内存结构

局限性：

  • 主要针对decoder-only的token（decoding）阶段，对prefill阶段的增益有限
  • 需要特定硬件支持层级缓存管理

参考资源：

九、参考资源