Split QKV + RMSNorm + RoPE 融合算子

Tue, 26 May 2026 09:50:58 +0800

源代码: vllm-ascend/vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_rope.py

背景#

问题：内存墙#

LLM 推理（尤其是 decode 阶段）是典型的 memory-bound 场景。每一次算子调用都是一次「从 Global Memory 搬数据到片上 → 计算 → 搬回 Global Memory」的循环。如果不融合，Split QKV → RMSNorm → RoPE 这三个步骤各自独立执行：

非融合流程:
  hidden_states ──[load]──> Split Q,K,V ──[store]──> q_in, k_in, v_in
  q_in ──[load]──> RMSNorm ──[store]──> q_normed
  k_in ──[load]──> RMSNorm ──[store]──> k_normed
  q_normed ──[load]──> RoPE ──[store]──> q_out
  k_normed ──[load]──> RoPE ──[store]──> k_out

每一次 load/store 都是一次 GM ↔ UB 的数据搬运。数据在总线上反复进出，但实际的计算量（几个乘加 + 一个开方）极小——这就是 memory-bandwidth bottleneck。

Ascend on Atomage's Blog

Split QKV + RMSNorm + RoPE 融合算子

背景#

问题：内存墙#