P5916. 轻量化推理引擎中的 Masked Fused RMSNorm 算子

场次：2026年10月14日-AI方向 · 来源：CodeFun2000 P5916

题目描述

在轻量化大模型推理引擎中，Transformer 的隐藏状态会在多个子模块之间不断传递。为减少内存读写、提高推理效率，工程实现里常把若干操作并进同一个算子，例如：

• 残差连接 Residual Add
• RMSNorm 归一化
• 逐元素缩放 Scale
• 逐元素偏置 Bias
• Padding Mask 处理

请按下列步骤，实现一个简化的 Masked Fused RMSNorm。

计算步骤

对第 $k$ 个样本，输入向量记为 $x_k$，同时给出残差向量 $r_k$、残差系数 $\alpha_k$、有效长度 $l_k$、缩放向量 $\gamma$、偏置向量 $\beta$ 和数值稳定项 $\epsilon$。

1. 残差融合

$h_{k,i} = x_{k,i} + \alpha_k \cdot r_{k,i}$

2. 有效部分的 RMS 与 energy

只有前 $l_k$ 个位置是有效位置，其后一律视为 padding，不参与 RMS 与 energy。

$energy_k = \sum_{i=1}^{l_k} h_{k,i}^2$

$RMS_k = \sqrt{\frac{energy_k}{d_k} + \epsilon}$

其中 $d_k = \max(l_k, 1)$。

3. 归一化输出

输出向量 $y_k$ 按位置取值：

• 有效位置（$1 \le i \le l_k$）： $y_{k,i} = \gamma_i \cdot \frac{h_{k,i}}{RMS_k} + \beta_i$
• 无效位置（$l_k < i \le n$）： $y_{k,i} = 0$

任务说明

给定 $m$ 个样本。对每个样本输出三项：

1. Masked Fused RMSNorm 之后的结果向量 $y_k$
2. 该样本有效部分的 $RMS_k$
3. 该样本有效部分的平方和 $energy_k$

输入格式

第一行三个数 $m$ $n$ $\epsilon$：样本个数、每个向量的最大长度、RMSNorm 的数值稳定项。 第二行 $n$ 个浮点数，依次为缩放向量 $\gamma_1, \gamma_2, \dots, \gamma_n$。 第三行 $n$ 个浮点数，依次为偏置向量 $\beta_1, \beta_2, \dots, \beta_n$。 随后 $m$ 行，每行一个样本，格式为： $l$ $\alpha$ $x_1$ $x_2$ $\dots$ $x_n$ $r_1$ $r_2$ $\dots$ $r_n$ 其中：

• $l$ 是当前样本的有效长度
• $\alpha$ 是残差系数
• $x_1 \dots x_n$ 是输入向量
• $r_1 \dots r_n$ 是残差向量

数据范围

• $1 \le m \le 200$
• $1 \le n \le 256$
• $0 \le l_k \le n$
• $0 \le \epsilon \le 1$
• $-10^8 \le x_{k,i}, r_{k,i} \le 10^8$
• $-10^3 \le \gamma_i, \beta_i \le 10^3$
• $-10^3 \le \alpha_k \le 10^3$

输出格式

共 $m$ 行，每行对应一个样本： $y_1$ $y_2$ $\dots$ $y_n$ $|$ $RMS$ $|$ $energy$

要求：

1. 相邻数据之间只用一个空格，| 左右也各有一个空格；
2. 全部浮点数保留两位小数，并按标准四舍五入；
3. 不要输出 -0.00。绝对值小于 0.005 的数一律写成 0.00。

示例

示例 1

输入：

2 3 0.5
1.0 1.0 1.0
0.0 2.0 -1.0
3 1.0 2.0 1.0 2.0 1.0 1.0 0.0
2 1.0 0.0 0.0 9.0 0.0 0.0 8.0

输出：

1.21 2.81 -0.19 | 2.48 | 17.00
0.00 2.00 0.00 | 0.71 | 0.00

说明： 第 1 个样本中，$l=3$，$\alpha=1.0$，残差融合后 $h = x + \alpha r = [3,\ 2,\ 2]$。 因此 $energy = 3^2 + 2^2 + 2^2 = 17$， $RMS = \sqrt{\frac{17}{3} + 0.5} \approx 2.4833$， 保留两位小数为 $2.48$。 再按 $y_i = \gamma_i \cdot \frac{h_i}{RMS} + \beta_i$ 计算，三个分量约是 $1.2081$、$2.8054$、$-0.1946$，保留两位小数后为 1.21 2.81 -0.19。

第 2 个样本中，$l=2$，$\alpha=1.0$。前两个位置融合后都是 $0$，第三个位置属于 padding，不计入 energy。 于是 $energy = 0$， $RMS = \sqrt{0.5} \approx 0.7071$，保留两位小数为 $0.71$。 有效位置上 $h_i = 0$，输出等于对应的 $\beta_i$；无效位置输出 $0$。该行为 0.00 2.00 0.00。

示例 2

输入：

2 3 0.000004
2.0 1.0 0.5
1.0 0.0 -1.0
3 0.0 2.0 2.0 4.0 1.0 3.0 3.0
0 1.0 6.0 6.0 6.0 2.0 2.0 2.0

输出：

2.41 0.71 -0.29 | 2.83 | 24.00
0.00 0.00 0.00 | 0.00 | 0.00

说明： 第 1 个样本的残差系数是 $0$，残差向量不进入融合： $h = [2,\ 2,\ 4]$， $energy = 2^2 + 2^2 + 4^2 = 24$， $RMS = \sqrt{\frac{24}{3} + 0.000004} \approx 2.8284$， 保留两位小数为 $2.83$。 代入缩放与偏置后，三个分量约是 $2.4142$、$0.7071$、$-0.2929$，保留两位小数为 2.41 0.71 -0.29。

第 2 个样本 $l = 0$，没有有效元素，因此 $energy = 0$， $RMS = \sqrt{0.000004} = 0.002$。 保留两位小数后为 $0.00$，输出向量也全部是 $0.00$。

提示

四舍五入

• 统一保留两位小数，不足两位要补零；
• 不要输出 -0.00，应写成 0.00。

输出格式

• 每行依次是 $n$ 个结果、一个 |、$RMS$、一个 |、$energy$；
• 空格和换行按上面的格式书写，不要多输出字符；
• $\gamma$、$\beta$ 按下标 $i$ 取，不要和样本编号 $k$ 混用。

---

解题思路

第一步：理解问题本质

题面把五个算子融合在一起，但对每个样本来说，它仍然只是一次逐元素的数值计算。把公式拆开看，真正需要注意的只有四件事：

1. 残差融合：$h_{k,i}=x_{k,i}+\alpha_k\cdot r_{k,i}$。这里 $\alpha_k$ 是第 $k$ 个样本自己的系数，同一个样本的所有位置共用它。

2. mask 的位置：$l_k$ 决定前多少个位置是有效位置。有效位置参与 energy 统计，也参与归一化输出；其后的位置是 padding。

3. energy 与 RMS 只统计有效部分：

   $$energy_k=\sum_{i=1}^{l_k}h_{k,i}^2,\qquad RMS_k=\sqrt{\frac{energy_k}{d_k}+\epsilon},\qquad d_k=\max(l_k,1)$$

   注意 $\epsilon$ 是在除法之后、开方之前加进去的，它保证分母不会恰为 $0$。

4. 输出分两段：有效位置输出 $\gamma_i\cdot\frac{h_{k,i}}{RMS_k}+\beta_i$，padding 位置直接输出 $0$。

还有一个容易忽略的对应关系：$\gamma_i$、$\beta_i$ 的下标是向量位置 $i$，不是样本编号 $k$。所有样本共用同一组缩放、偏置向量。

由于样本之间完全独立，逐个样本处理即可，不需要把所有样本同时保存在内存里。

第二步：暴力解法

最直白的写法是严格照着公式翻译：算第 $i$ 个输出分量时，现场把 $energy$ 重新累加一遍。

for k in range(m):
    l, alpha = ...
    x, r = ...
    h = [x[i] + alpha * r[i] for i in range(n)]

    for i in range(n):
        if i < l:
            # 每算一个位置，就从头再累加一次 energy
            energy = 0.0
            for j in range(l):
                energy += h[j] * h[j]
            rms = math.sqrt(energy / max(l, 1) + eps)
            y_i = gamma[i] * h[i] / rms + beta[i]
        else:
            y_i = 0.0

这段代码的结果是对的，但 $energy$ 与 $RMS$ 对同一个样本来说是常量：它们只和 $h$ 的前 $l$ 项有关，和当前算的是第几个位置无关。把它放进位置循环里重算，复杂度就变成了 $O(m\cdot n^2)$。本题 $m\le 200$、$n\le 256$，最坏约 $1.3\times 10^7$ 次乘加，虽然还能跑过，但这些计算全是重复的。

第三步：优化解法——先算 energy，再算输出

把「每个样本变化的量」和「每个位置变化的量」分开：

• 与位置无关：$h$、$energy$、$RMS$；
• 与位置有关：$y_i$。

于是每个样本只需要两遍扫描：

# 第一遍：融合 + 累加 energy（只累加前 l 个位置）
h = [0.0] * n
energy = 0.0
for i in range(l):
    h[i] = x[i] + alpha * r[i]
    energy += h[i] * h[i]

# 样本级别的常量只算一次
rms = math.sqrt(energy / max(l, 1) + eps)

# 第二遍：按位置写输出，padding 直接写 0
for i in range(n):
    if i < l:
        values[i] = gamma[i] * h[i] / rms + beta[i]
    else:
        values[i] = 0.0

复杂度降到 $O(m\cdot n)$，并且 $h_i$ 只计算一次。

第四步：最优解法 / 完整实现细节

到第三步为止，算法的复杂度已经是最优的了，剩下的工作都在「实现细节」上，而这道题的判分恰恰卡在细节：

1. 一次性读入：用 sys.stdin.buffer.read().split() 把全部 token 读进来，再用游标顺序取数。浮点数直接用 float(token) 解析，行内空格、换行都不会影响。
2. 融合与累加合并：对 $i<l$ 的循环里同时算 $h_i$ 和累加 $h_i^2$；对 $i\ge l$ 再补一遍 $h_i$（padding 的 $h$ 不参与输出，但保留完整的 $h$ 能让「哪些位置被 mask」更直观）。
3. 分母取 $\max(l,1)$：当 $l=0$ 时没有有效元素，$energy=0$，此时 $RMS=\sqrt{\epsilon}$，仍然要正常输出，不能除以 $0$。
4. padding 输出 $0$ 而不是 $\beta_i$：mask 的语义是「这些位置不产生贡献」，所以偏置也不能加。
5. 统一格式化：所有浮点数（包括 $RMS$ 与 $energy$）都保留两位小数。f"{v:.2f}" 会自动补零，例如 $17$ 输出 17.00。
6. 避免 -0.00：浮点误差会让「数学上是 $0$」的结果带一个负号（如 -1e-17）。题面规定绝对值小于 $0.005$ 的数一律写成 0.00，所以格式化前先做一次就近归零，格式化后再检查一次字符串，双保险。
7. 退化情况：约束允许 $\epsilon=0$。若此时有效部分又恰好全为 $0$，则 $RMS=0$，$h_i/RMS$ 成为 $0/0$。由于有效位置上 $h_i$ 全是 $0$，按极限把它视为 $0$，输出退化为 $\beta_i$，代码不会抛 ZeroDivisionError。

---

完整代码实现

import math
import sys


def format_number(value):
    """保留两位小数，并保证不会输出 "-0.00"。"""
    # 题面规定：绝对值小于 0.005 的数一律写成 0.00
    if abs(value) < 0.005:
        return "0.00"
    text = f"{value:.2f}"
    # 双保险：极端情况下舍入后仍可能拼出 "-0.00"
    return "0.00" if text == "-0.00" else text


def solve():
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    pos = 0
    m = int(tokens[pos])
    n = int(tokens[pos + 1])
    eps = float(tokens[pos + 2])
    pos += 3

    # gamma、beta 按位置取，所有样本共用，只需要读一次
    gamma = [float(tokens[pos + i]) for i in range(n)]
    pos += n
    beta = [float(tokens[pos + i]) for i in range(n)]
    pos += n

    lines = []

    for _ in range(m):
        # 每个样本：l、alpha、x 向量、r 向量
        length = int(tokens[pos])
        alpha = float(tokens[pos + 1])
        pos += 2
        if length < 0:
            length = 0
        elif length > n:
            length = n

        x = [float(tokens[pos + i]) for i in range(n)]
        pos += n
        r = [float(tokens[pos + i]) for i in range(n)]
        pos += n

        # 融合与 energy 合并到同一轮：energy 只累加有效位置
        h = [0.0] * n
        energy = 0.0
        for i in range(length):
            h[i] = x[i] + alpha * r[i]
            energy += h[i] * h[i]
        # padding 位置也补上 h，虽然它不参与输出
        for i in range(length, n):
            h[i] = x[i] + alpha * r[i]

        # 分母取 max(l, 1)：l = 0 时 RMS 仍为 sqrt(eps)
        divisor = length if length > 0 else 1
        rms = math.sqrt(energy / divisor + eps)

        values = []
        for i in range(n):
            if i < length:
                if rms == 0.0:
                    # eps = 0 且有效部分全为 0：把 h/RMS 视为 0
                    values.append(beta[i])
                else:
                    values.append(gamma[i] * h[i] / rms + beta[i])
            else:
                # padding 位置输出 0，注意不是 beta[i]
                values.append(0.0)

        parts = [format_number(value) for value in values]
        parts.append("|")
        parts.append(format_number(rms))
        parts.append("|")
        parts.append(format_number(energy))
        lines.append(" ".join(parts))

    sys.stdout.write("\n".join(lines) + "\n")


if __name__ == "__main__":
    solve()

NumPy 版本

以下版本使用 NumPy。只有在评测环境允许第三方库时才能提交；CodeFun2000 的普通 Python 环境不一定安装 NumPy。

import sys

import numpy as np


def format_number(value):
    """保留两位小数，并保证不会输出 "-0.00"。"""
    # 题面规定：绝对值小于 0.005 的数一律写成 0.00
    if abs(value) < 0.005:
        return "0.00"
    text = f"{value:.2f}"
    # 双保险：极端情况下舍入后仍可能拼出 "-0.00"
    return "0.00" if text == "-0.00" else text


def solve():
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    m = int(tokens[0])
    n = int(tokens[1])
    eps = float(tokens[2])

    pos = 3
    # gamma、beta 按位置取，所有样本共用，各读成一条长度为 n 的向量
    gamma = np.array(tokens[pos:pos + n], dtype=np.float64)
    pos += n
    beta = np.array(tokens[pos:pos + n], dtype=np.float64)
    pos += n

    # 每个样本一行，共 2 + 2n 个数：l、alpha、x_1..x_n、r_1..r_n
    width = 2 + 2 * n
    data = np.array(tokens[pos:pos + m * width], dtype=np.float64).reshape(m, width)

    lengths = np.clip(data[:, 0].astype(np.int64), 0, n)   # (m,)，题面保证 0 <= l <= n
    alpha = data[:, 1:2]                                   # (m, 1)，按行广播
    x = data[:, 2:2 + n]                                   # (m, n)
    r = data[:, 2 + n:2 + 2 * n]                           # (m, n)

    # 第 1 步：残差融合 h = x + alpha * r，alpha 逐样本广播到整行
    h = x + alpha * r

    # 第 2 步：有效位置掩码，第 i 列有效 <=> i < l_k。布尔矩阵直接当 mask
    mask = np.arange(n)[None, :] < lengths[:, None]

    # 第 3 步：energy 只累加有效位置的平方，padding 位置先清零再平方
    # 用 cumsum 取最后一列：它是按列顺序逐项累加的，与「逐位置累加」的
    # 结果逐位相同；np.sum 的成对求和虽然更准，但会在 1e22 量级上与逐项
    # 累加差出一个 ULP，而 energy 要原样打印两位小数
    squares = np.where(mask, h * h, 0.0)
    energy = np.cumsum(squares, axis=1)[:, -1]

    # 第 4 步：RMS = sqrt(energy / max(l, 1) + eps)，eps 在开方里面
    divisor = np.maximum(lengths, 1).astype(np.float64)
    rms = np.sqrt(energy / divisor + eps)

    # 第 5 步：有效位置输出 gamma_i * h_i / RMS + beta_i，padding 位置输出 0
    # rms 为 0（eps = 0 且有效部分全为 0）时分子也全为 0，除法按 1 处理即可
    safe_rms = np.where(rms > 0.0, rms, 1.0)
    y = np.where(
        mask,
        gamma[None, :] * h / safe_rms[:, None] + beta[None, :],
        0.0,
    )

    lines = []
    for k in range(m):
        parts = [format_number(value) for value in y[k]]
        parts.append("|")
        parts.append(format_number(rms[k]))
        parts.append("|")
        parts.append(format_number(energy[k]))
        lines.append(" ".join(parts))

    sys.stdout.write("\n".join(lines) + "\n")


if __name__ == "__main__":
    solve()

这一版把哪些循环换成了数组运算：

• 残差融合 $h=x+\alpha r$ 用 alpha[:, None] 把每个样本自己的系数广播到整行，$m\times n$ 次乘加一次完成；
• 掩码用 np.arange(n)[None, :] < lengths[:, None] 直接生成 $m\times n$ 的布尔矩阵，不需要逐位置判断 i < l；
• energy、$RMS$、$y$ 都是整块算，gamma、beta 沿列广播、$RMS$ 沿行广播，逐元素公式与纯 Python 版一字不差；
• energy 这类「逐项累加出来、又要原样打印」的量用 np.cumsum(...)[:, -1] 求而不是 np.sum：np.sum 的成对求和更精确，但求和顺序与逐项累加不同，在 $h_i^2$ 达到 $10^{22}$ 量级时会差出一个 ULP，而 energy 要求原样保留两位小数；
• 只有最终格式化仍然逐个取值，因为 f"{v:.2f}" 的舍入和 -0.00 归零按数值判断，保持与纯 Python 版完全一致。

---

示例推演

样例 1

$m=2$、$n=3$、$\epsilon=0.5$，$\gamma=[1,1,1]$，$\beta=[0,2,-1]$。

第 1 个样本：$l=3$、$\alpha=1$，$x=[2,1,2]$，$r=[1,1,0]$。

先做残差融合：

$$h=[2+1\times1,\ 1+1\times1,\ 2+1\times0]=[3,2,2]$$

有效长度等于 $n$，energy 统计全部三个位置：

$$energy=3^2+2^2+2^2=9+4+4=17$$

$d=\max(3,1)=3$，于是

$$RMS=\sqrt{\frac{17}{3}+0.5}=\sqrt{6.1667}\approx2.4833\rightarrow2.48$$

逐分量代入：

位置 $i$	$h_i$	$\gamma_i\cdot h_i/RMS$	$\beta_i$	$y_i$	输出
1	3	$3/2.4833=1.2081$	0	1.2081	1.21
2	2	$2/2.4833=0.8054$	2	2.8054	2.81
3	2	$2/2.4833=0.8054$	$-1$	$-0.1946$	-0.19

该行输出 1.21 2.81 -0.19 | 2.48 | 17.00。

第 2 个样本：$l=2$、$\alpha=1$，$x=[0,0,9]$，$r=[0,0,8]$。

融合后 $h=[0,0,17]$。只有前 $2$ 个位置有效，而它们的 $h$ 都是 $0$，所以 $energy=0+0=0$。$d=\max(2,1)=2$：

$$RMS=\sqrt{\frac{0}{2}+0.5}=\sqrt{0.5}\approx0.7071\rightarrow0.71$$

位置 1、2 上 $h_i=0$，输出就是偏置本身：$y_1=0$、$y_2=2$；位置 3 是 padding，输出 $0$。

该行输出 0.00 2.00 0.00 | 0.71 | 0.00。

样例 2

$m=2$、$n=3$、$\epsilon=0.000004$，$\gamma=[2,1,0.5]$，$\beta=[1,0,-1]$。

第 1 个样本：$l=3$、$\alpha=0$，$x=[2,2,4]$，$r=[3,3,1]$。

$\alpha=0$ 时残差项整体为 $0$，所以 $h=x=[2,2,4]$。

$$energy=2^2+2^2+4^2=4+4+16=24$$ $$RMS=\sqrt{\frac{24}{3}+0.000004}=\sqrt{8.000004}\approx2.8284\rightarrow2.83$$

位置 $i$	$\gamma_i\cdot h_i/RMS$	$\beta_i$	$y_i$	输出
1	$2\times2/2.8284=1.4142$	1	2.4142	2.41
2	$1\times2/2.8284=0.7071$	0	0.7071	0.71
3	$0.5\times4/2.8284=0.7071$	$-1$	$-0.2929$	-0.29

该行输出 2.41 0.71 -0.29 | 2.83 | 24.00。

第 2 个样本：$l=0$、$\alpha=1$，$x=[6,6,6]$，$r=[2,2,2]$。

$l=0$ 表示没有任何有效位置：$energy=0$，$d=\max(0,1)=1$，

$$RMS=\sqrt{\frac{0}{1}+0.000004}=0.002\rightarrow0.00$$

三个位置全是 padding，全部输出 $0$。该行输出 0.00 0.00 0.00 | 0.00 | 0.00。

注意这里 $h=[8,8,8]$ 明明不为零，但因为全部被 mask，既不进 $energy$ 也不进输出。这正是 Masked 的语义。

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
暴力：逐位置重算 energy	$O(m\cdot n^2)$	$O(n)$	每个输出位置都从头累加一次平方和
优化：energy、RMS 只算一次	$O(m\cdot n)$	$O(n)$	每个样本两遍扫描
本题实现（融合 + 单次累加）	$O(m\cdot n)$	$O(n)$	融合与累加合并成一轮，常数更小

$m\le 200$、$n\le 256$ 时，$O(mn)$ 只有约 $5\times10^4$ 次浮点乘加，远远低于 2 秒的限制。

---

易错点总结

1. energy 只统计前 $l$ 个位置

padding 位置的 $h_i$ 即使很大也不能进 $energy$。样例 1 第 2 个样本的 $h_3=17$，如果误把它算进去，$energy$ 会变成 $289$，$RMS$ 与整行输出全错。

2. 分母是 $\max(l,1)$，不是 $l$

$l=0$ 时直接写 energy / l 会抛出 ZeroDivisionError。题目已经给出 $d_k=\max(l_k,1)$，照抄即可。

3. $\epsilon$ 在开方里面

正确写法是 $\sqrt{energy/d+\epsilon}$，不是 $\sqrt{energy/d}+\epsilon$。把 $\epsilon$ 提到根号外，数值会明显不同。

4. padding 输出 $0$，不是 $\beta_i$

mask 之后这些位置不产生贡献，缩放与偏置都不该加。样例 1 第 2 个样本的位置 3 若输出 $\beta_3=-1$，就会得到 -1.00 而不是 0.00。

5. $\gamma$、$\beta$ 用位置下标，不用样本编号

所有样本共用同一组 $\gamma$、$\beta$。写成 gamma[k] 在 $m>n$ 时直接越界，在 $m<n$ 时会静默取错值。

6. 不要输出 -0.00

浮点运算会出现 -1e-17 这类「数学上是 0」的结果。题面明确要求绝对值小于 $0.005$ 的数写成 0.00，所以格式化前必须先归零。

7. 保留两位小数要补零

17 要输出 17.00，0 要输出 0.00。使用 f"{value:.2f}" 可以自动补齐；round(value, 2) 之后再 str() 不行，因为 str(round(17.0, 2)) 得到的是 17.0，不会补成 17.00。

8. $\epsilon=0$ 的退化情形

约束允许 $\epsilon=0$。如果有效位置上的 $h$ 全为 $0$，则 $RMS=0$，归一化项是 $0/0$。代码里把这种情况视为 $0$，输出退化为 $\beta_i$，既保证不崩溃，也与「有效位置没有信号」的直觉一致。

9. 输出分隔符

每行是 $n$ 个结果、|、$RMS$、|、$energy$。用 " ".join(...) 拼接可以保证相邻数据之间只有一个空格，| 左右各一个空格。

---

扩展思考

• 真正的融合算子：工程实现里这一步通常写成 kernel，残差写回、RMS 统计、缩放偏置一趟走完，避免把中间张量写回显存。本题的「先算 energy 再算输出」正是这种单遍思路的简化版。
• 数值稳定性：$x_{k,i}$ 最大 $10^8$、$\alpha_k$ 最大 $10^3$，融合后 $h_i$ 可达 $10^{11}$，$h_i^2$ 可达 $10^{22}$。用双精度（float64）累加完全够用；如果改用 float32，累加 256 项时会损失有效精度，通常的做法是先把每个平方项按块累加，或者用补偿求和（Kahan summation）。
• 分组 RMSNorm：真实模型里常把隐藏维度分成若干组，每组单独算 RMS。做法与本题完全一样，只是把「前 $l$ 个位置」换成「某一组的下标区间」。
• 两种 mask 的区别：本题是 padding mask（按有效长度截断），自回归模型还需要 causal mask（第 $i$ 个位置只能看到 $1..i$）。两者可以叠加，此时 $energy$ 的累加上界要取两者的交集。
• 多 batch 的写法：本题样本之间互不影响，所以逐样本处理即可；如果输入改成「先给一个矩阵再给一个 mask 矩阵」，也可以先算每一行的 $energy$，再整块广播计算 $y$，结果相同。

---

相关题目

• 滑动窗口注意力机制
• 整数缩放因子下的全整型矩阵乘积