算子融合：RMSNorm、GEMM 与 SwiGLU

分析合并计算与减少显存读写的收益，以及不同形状下的差别。

杜子源2026/10/30公开

对应 01-v2性能瓶颈 — 瓶颈 D（层内算子：Norm / RoPE / MLP）。

在 LLM 实现的过程中，需要很多组件，例如RMSNorm、RoPE、SiLU、Mul等等，这些模块在 v2 版本都是由很多 Pytorch 小的算子拼起来的。但是经过我们之前的分析可以得出来，这种小算子普遍 Kernel 多，中间张量会反复经过显存。

因此在 v3 版本中，我们把它们基本都替换成了高性能专用核。

例如 RoPE、SiLU+Mul 和 RMSNorm 都替换成了 FlashInfer 的实现版本，其中 RoPE/SiLU 和残差融合 RMSNorm 均注册为 custom op 以兼容 torch.compile。

RMSNorm#

1.1 朴素实现#

1
class RMSNorm(nn.Module):
2
    def forward(self, x, residual=None):
3
        x32 = x.float()  # Kernel 1: cast
4
        if residual is not None:
5
            x32 = x32 + residual.float()  # Kernel 2: add
6
        out = (
7
            x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) * weight
8
        )  # pow/mean/rsqrt/mul
9
        return out.to(x.dtype)  # Kernel N: cast

一个 RMSNorm 被拆成 5~8 个独立 kernel，每个 kernel 都走一遍 HBM read → compute → HBM write，中间变量频繁读写。

1.2 FlashInfer 融合#

1
# v3 — compute/layers/layernorm.py
2
def forward(self, x, residual=None):
3
    if residual is None:
4
        return _rmsnorm(x, self.weight, self.eps), x
5
    _fused_add_rmsnorm(x, residual, self.weight, self.eps)
6
    return x, residual

_rmsnorm / _fused_add_rmsnorm 是 FlashInfer 在 ops.py 中注册的 vllm:: custom op。

为什么不是add + F.rms_norm？早期 v3 用 PyTorch 原生 F.rms_norm，但残差加法仍是独立 kernel。实测在 Qwen3-0.6B 形状（hidden=1024）下，FlashInfer 融合比官方快约 1.3x（prefill）/ 2.4x（decode 小批量），因此替换为 FlashInfer 专核，与 RoPE/SiLU 保持一致。
为什么注册成 custom op？fused_add_rmsnorm 是 inplace（改写 x 和 residual），直接放入 torch.compile 会让 reduce-overhead 的 cudagraph 捕获产生空图、prefill 性能暴跌；和 RoPE/SiLU 一样注册为 vllm:: custom op 并声明 mutates_args=("x","residual") 后，编译期作为黑盒处理、别名推断正确，cudagraph 正常捕获。

GEMM 融合#

2.1 QKV 投影合并#

v2 里 q/k/v 三个投影是三个独立 Linear：

1
self.q_proj = nn.Linear(hidden, num_heads * head_dim, bias=False)
2
self.k_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)
3
self.v_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)

三个 GEMM 分别进行读写，而在 v3 中合并为一个 GEMM：

1
# v3 — compute/models/qwen3.py Qwen3Attention
2
self.qkv_proj = nn.Linear(hidden_size, (num_heads + 2 * num_kv_heads) * head_dim, bias=qkv_bias)
3

4

5
def forward(self, positions, hidden_states, attn):
6
    qkv = self.qkv_proj(hidden_states)  # 一次 GEMM 出 q+k+v
7
    q, k, v = self._split_norm_rope(positions, qkv)  # 沿 dim 切分
8
    return self.o_proj(self.attn(q, k, v, attn))

输入只读一次、一次 launch；加载权重时把 HF 的 q_proj/k_proj/v_proj 按偏移拼进 qkv_proj（见 10-模型加载）。

2.2 MLP 的 gate/up 合并与 SiLU 融合#

在 v2 中，SwiGLU MLP 的计算被拆成几个独立步骤。先对输入做一次线性投影得到 gate，并经过 SiLU；再做一次线性投影得到 up；最后将两者逐元素相乘。对应到代码上大致是：

1
gate = F.silu(x @ W_gate)  # 一次 GEMM + 一次 SiLU
2
up = x @ W_up  # 一次 GEMM
3
out = gate * up  # 一次逐元素乘法

这种写法逻辑清楚，但执行时会多出几次 kernel 启动，同时 gate、up 等中间结果也需要额外写入显存。

v3 把 gate_proj 和 up_proj 合并成一个更大的线性层，一次性算出 gate 和 up。权重形状从原来的两个 [intermediate, hidden] 变成 [2 * intermediate, hidden]。这样两次 GEMM 就合并成了一次 GEMM。

1
# v3 — compute/models/qwen3.py Qwen3MLP
2
self.gate_up_proj = nn.Linear(hidden_size, intermediate_size * 2, bias=False)
3
self.act_fn = SiluAndMul()  # torch.ops.vllm.silu_and_mul
4

5

6
def forward(self, x):
7
    x = self.gate_up_proj(x)  # 一次 GEMM，同时得到 gate 和 up
8
    x = self.act_fn(x)  # 单个 kernel 完成 silu(gate) * up
9
    return self.down_proj(x)

这里的 silu_and_mul 会把输入沿最后一维拆成 gate 和 up 两部分，然后在同一个 kernel 里完成 silu(gate) * up。相比原来的写法，它不需要单独启动一次逐元素乘法 kernel，也减少了中间张量的显存写入。

RoPE#

v2 的 RoPE 需要 index_select、chunk、四次 elementwise 和 cat，中间结果反复经过显存。

v3 把 RoPE 换成 FlashInfer 的 apply_rope_with_cos_sin_cache，并注册为 custom op 以兼容 torch.compile：

1
# v3 — compute/layers/ops.py
2
@torch.library.custom_op("vllm::rotary_embedding", mutates_args=())
3
def rotary_embedding(q, k, positions, cos_sin_cache, head_size):
4
    n, n_qh = q.shape[:2]
5
    n_kh = k.shape[1]
6
    qf = q.reshape(n, n_qh * head_size).contiguous()
7
    kf = k.reshape(n, n_kh * head_size).contiguous()
8
    qo, ko = _flashinfer_rope(positions, qf, kf, head_size, cos_sin_cache, is_neox=True)
9
    return qo.view_as(q), ko.view_as(k)

RotaryEmbedding（compute/layers/rotary.py）预计算 cos_sin_cache（[max_pos, head_size]，前一半 cos、后一半 sin），forward 时调用 torch.ops.vllm.rotary_embedding。

为什么这几个 kernel 要注册为 custom op#

torch.compile 对 prefill 做全图编译时，FlashInfer 的 RoPE / SiLU / RMSNorm 内核以及 FlashAttention 无法被 Dynamo trace，必须在 ops.py 中注册为 torch.library.custom_op 并配 fake 实现：

1
@silu_and_mul.register_fake
2
def _silu_and_mul_fake(x):
3
    return torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device)

注册后，编译期把它们当作黑盒（只做形状/dtype 推导），避免 graph break；运行期直接调用底层内核。注意 fused_add_rmsnorm 是 inplace 算子，注册时需在 mutates_args 中声明被改写的参数，否则编译器会误认为输出是新张量，与 CUDA Graph 捕获的实际写入冲突。

小结#

模块	v2	v3	手段
RMSNorm	5~8 个朴素 kernel	FlashInfer `fused_add_rmsnorm` 融合核（残差 add 并入）	换专核
q/k/v 投影	3 个 GEMM	1 个 `qkv_proj` GEMM	GEMM 合并
gate/up 投影	2 个 GEMM + cat	1 个 `gate_up_proj` GEMM	GEMM 合并
SiLU+Mul	`F.silu` + `mul` 两核	FlashInfer `silu_and_mul` 单核	换专核
RoPE	index_select + chunk + rotate	FlashInfer `apply_rope_with_cos_sin_cache` 单核	换专核

联系杜子源源