#P5916. 轻量化推理引擎中的 Masked Fused RMSNorm 算子
-
1000ms
Tried: 467
Accepted: 36
Difficulty: 3
所属公司 :
华为
时间 :2026年10月14日-AI方向
轻量化推理引擎中的 Masked Fused RMSNorm 算子
解题思路
对每个样本先做残差融合 hi=xi+αri。energy 只累加前 l 个有效位置的 hi2,分母 d=max(l,1),RMS=energy/d+ϵ。有效位置输出 yi=γi⋅hi/RMS+βi,下标从 0 计时也就是 i<l;无效位置直接输出 0,不能把 γi,βi 加进去。l=0 时没有有效位,energy 为 0,整段 y 都是 0。若有效位上的 hi 全是 0 且 ϵ=0,则 RMS 为 0,此时 hi/RMS 没有定义,但每个有效 hi 必然是 0,归一化项按 0 处理,有效位只保留 βi。
实现时按样本各扫一遍。浮点要保留两位小数,并且是标准四舍五入:恰好满五时远离零。不要用 C++ 的 printf("%.2f")、Python 的 format 或 Java 的 String.format("%.2f"),它们对能精确表示的 0.125、1.125 会走银行家舍入,印成 0.12、1.12。绝对值小于 0.005 的数一律写成 0.00,这样不会出现 -0.00。同一行用单个空格分隔,| 两侧各有一个空格。γ,β 按位置下标取,不要误用样本编号。
复杂度分析
每个样本只扫常数遍长度为 n 的向量,总时间 O(mn),额外空间 O(n)。m≤200,n≤256,线性扫描足够。h 最大约 1.001×1011,平方和约 1024,必须用浮点累加,放进 32 位或 64 位整数会溢出。
在轻量化大模型推理引擎中,Transformer 的隐藏状态会在多个子模块之间不断传递。为减少内存读写、提高推理效率,工程实现里常把若干操作并进同一个算子,例如:
- 残差连接 Residual Add
- RMSNorm 归一化
- 逐元素缩放 Scale
- 逐元素偏置 Bias
- Padding Mask 处理
请按下列步骤,实现一个简化的 Masked Fused RMSNorm。
计算步骤
对第 k 个样本,输入向量记为 xk,同时给出残差向量 rk、残差系数 αk、有效长度 lk、缩放向量 γ、偏置向量 β 和数值稳定项 ϵ。
1. 残差融合
hk,i=xk,i+αk⋅rk,i
2. 有效部分的 RMS 与 energy
只有前 lk 个位置是有效位置,其后一律视为 padding,不参与 RMS 与 energy。
energyk=∑i=1lkhk,i2
RMSk=dkenergyk+ϵ
其中 dk=max(lk,1)。
3. 归一化输出
输出向量 yk 按位置取值:
- 有效位置(1≤i≤lk): yk,i=γi⋅RMSkhk,i+βi
- 无效位置(lk<i≤n): yk,i=0
任务说明
给定 m 个样本。对每个样本输出三项:
- Masked Fused RMSNorm 之后的结果向量 yk
- 该样本有效部分的 RMSk
- 该样本有效部分的平方和 energyk
输入
第一行三个数 m n ϵ:样本个数、每个向量的最大长度、RMSNorm 的数值稳定项。 第二行 n 个浮点数,依次为缩放向量 γ1,γ2,…,γn。 第三行 n 个浮点数,依次为偏置向量 β1,β2,…,βn。 随后 m 行,每行一个样本,格式为: l α x1 x2 … xn r1 r2 … rn 其中:
- l 是当前样本的有效长度
- α 是残差系数
- x1…xn 是输入向量
- r1…rn 是残差向量
数据范围
- 1≤m≤200
- 1≤n≤256
- 0≤lk≤n
- 0≤ϵ≤1
- −108≤xk,i,rk,i≤108
- −103≤γi,βi≤103
- −103≤αk≤103
输出
共 m 行,每行对应一个样本: y1 y2 … yn ∣ RMS ∣ energy
要求:
- 相邻数据之间只用一个空格,
|左右也各有一个空格; - 全部浮点数保留两位小数,并按标准四舍五入;
- 不要输出
-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,α=1.0,残差融合后
h=x+αr=[3, 2, 2]。
因此
energy=32+22+22=17,
RMS=317+0.5≈2.4833,
保留两位小数为 2.48。
再按 yi=γi⋅RMShi+βi 计算,三个分量约是 1.2081、2.8054、−0.1946,保留两位小数后为
1.21 2.81 -0.19。
第 2 个样本中,l=2,α=1.0。前两个位置融合后都是 0,第三个位置属于 padding,不计入 energy。
于是 energy=0,
RMS=0.5≈0.7071,保留两位小数为 0.71。
有效位置上 hi=0,输出等于对应的 β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=22+22+42=24,
RMS=324+0.000004≈2.8284,
保留两位小数为 2.83。
代入缩放与偏置后,三个分量约是 2.4142、0.7071、−0.2929,保留两位小数为
2.41 0.71 -0.29。
第 2 个样本 l=0,没有有效元素,因此 energy=0, RMS=0.000004=0.002。 保留两位小数后为 0.00,输出向量也全部是 0.00。
提示
四舍五入
- 统一保留两位小数,不足两位要补零;
- 不要输出
-0.00,应写成0.00。
输出格式
- 每行依次是 n 个结果、一个
|、RMS、一个|、energy; - 空格和换行按上面的格式书写,不要多输出字符;
- γ、β 按下标 i 取,不要和样本编号 k 混用。