Merged
docs: 修复AIDD文档质量扫描报告中的语义与示例问题 #163
xmzAdd on Sep 14
docs: 修复AIDD文档质量扫描报告中的语义与示例问题 #163
Merged
Total: 13 files changed+113-92
| @@ -15,7 +15,7 @@ y = (x * rstd) * gamma + beta | |||
| 15 | ## 函数签名 | 15 | ## 函数签名 |
| 16 | 16 | ||
| 17 | ```python | 17 | ```python |
| 18 | -def add_rms_norm_bias( | 18 | +def npu_add_rms_norm_bias( |
| 19 | x1: torch.Tensor, | 19 | x1: torch.Tensor, |
| 20 | x2: torch.Tensor, | 20 | x2: torch.Tensor, |
| 21 | gamma: torch.Tensor, | 21 | gamma: torch.Tensor, |
| @@ -109,9 +109,11 @@ import torch | |||
| 109 | from mindspeed_ops.api.aclnn.npu_add_rms_norm_bias import npu_add_rms_norm_bias | 109 | from mindspeed_ops.api.aclnn.npu_add_rms_norm_bias import npu_add_rms_norm_bias |
| 110 | 110 | ||
| 111 | # 示例 1：基本使用（无 beta） | 111 | # 示例 1：基本使用（无 beta） |
| 112 | -x1 = torch.randn(4, 256, dtype=torch.float16) | 112 | +dtype = torch.float16 |
| 113 | -x2 = torch.randn(4, 256, dtype=torch.float16) | 113 | +device = torch.device("npu") |
| 114 | -gamma = torch.randn(256, dtype=torch.float16) | 114 | +x1 = torch.randn(4, 256, dtype=dtype).to(device) |
| 115 | +x2 = torch.randn(4, 256, dtype=dtype).to(device) | ||
| 116 | +gamma = torch.randn(256, dtype=dtype).to(device) | ||
| 115 | 117 | ||
| 116 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma) | 118 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma) |
| 117 | 119 | ||
| @@ -122,10 +124,12 @@ print(f"x shape: {x.shape}, dtype: {x.dtype}") | |||
| 122 | 124 | ||
| 123 | ```python | 125 | ```python |
| 124 | # 示例 2：带 beta 偏置 | 126 | # 示例 2：带 beta 偏置 |
| 125 | -x1 = torch.randn(2, 8, 512, dtype=torch.float32) | 127 | +dtype = torch.float32 |
| 126 | -x2 = torch.randn(2, 8, 512, dtype=torch.float32) | 128 | +device = torch.device("npu") |
| 127 | -gamma = torch.randn(512, dtype=torch.float32) | 129 | +x1 = torch.randn(2, 8, 512, dtype=dtype).to(device) |
| 128 | -beta = torch.randn(512, dtype=torch.float32) | 130 | +x2 = torch.randn(2, 8, 512, dtype=dtype).to(device) |
| 131 | +gamma = torch.randn(512, dtype=dtype).to(device) | ||
| 132 | +beta = torch.randn(512, dtype=dtype).to(device) | ||
| 129 | 133 | ||
| 130 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma, beta, epsilon=1e-5) | 134 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma, beta, epsilon=1e-5) |
| 131 | 135 | ||
| @@ -137,9 +141,11 @@ print("AddRmsNormBias operation completed successfully!") | |||
| 137 | 141 | ||
| 138 | ```python | 142 | ```python |
| 139 | # 示例 3：多维 gamma（对最后几维做 RMS Norm） | 143 | # 示例 3：多维 gamma（对最后几维做 RMS Norm） |
| 140 | -x1 = torch.randn(2, 4, 8, 64, dtype=torch.float16) | 144 | +dtype = torch.float16 |
| 141 | -x2 = torch.randn(2, 4, 8, 64, dtype=torch.float16) | 145 | +device = torch.device("npu") |
| 142 | -gamma = torch.randn(8, 64, dtype=torch.float16) # 对最后 2 维做归一化 | 146 | +x1 = torch.randn(2, 4, 8, 64, dtype=dtype).to(device) |
| 147 | +x2 = torch.randn(2, 4, 8, 64, dtype=dtype).to(device) | ||
| 148 | +gamma = torch.randn(8, 64, dtype=dtype).to(device) # 对最后 2 维做归一化 | ||
| 143 | 149 | ||
| 144 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma) | 150 | y, rstd, x = npu_add_rms_norm_bias(x1, x2, gamma) |
| 145 | # rstd shape 为 (2, 4, 1, 1) | 151 | # rstd shape 为 (2, 4, 1, 1) |
| @@ -76,7 +76,7 @@ Python接口返回loss；反向传播时自动使用CATLASS内核生成的 | |||
| 76 | | `hidden_states` | `torch.Tensor` | `[B, S, H]`，最后一维连续且无内部重叠 | `torch.bfloat16` | | 76 | | `hidden_states` | `torch.Tensor` | `[B, S, H]`，最后一维连续且无内部重叠 | `torch.bfloat16` | |
| 77 | | `head_weight` | `torch.Tensor` | `[V, H]`，连续布局 | `torch.bfloat16` | | 77 | | `head_weight` | `torch.Tensor` | `[V, H]`，连续布局 | `torch.bfloat16` | |
| 78 | | `label_chunks` | `Sequence[torch.Tensor]` | 按 `chunk_size` 划分的标签，数量为 `ceil(S / chunk_size)` | 转换为NPU `torch.int32` | | 78 | | `label_chunks` | `Sequence[torch.Tensor]` | 按 `chunk_size` 划分的标签，数量为 `ceil(S / chunk_size)` | 转换为NPU `torch.int32` | |
| 79 | -| `alpha_chunks` | `Sequence[torch.Tensor \| None] \| None` | 每个chunk的归一化因子；省略时使用1 | `torch.float32` 或 `torch.int64` | | 79 | +| `alpha_chunks` | `Sequence[torch.Tensor \| None] \| None` | 每个chunk的归一化因子；需显式传参（接口无默认值），传入 `None` 或某一项为 `None` 时使用1 | `torch.float32` 或 `torch.int64` | |
| 80 | | `chunk_size` | `attr` | 正整数 | `int` | | 80 | | `chunk_size` | `attr` | 正整数 | `int` | |
| 81 | | `ignore_index` | `attr` | 忽略标签，必须在int32范围内 | `int` | | 81 | | `ignore_index` | `attr` | 忽略标签，必须在int32范围内 | `int` | |
| 82 | | `reduction` | `attr` | 支持 `"sum"` 和 `"none"` | `str` | | 82 | | `reduction` | `attr` | 支持 `"sum"` 和 `"none"` | `str` | |
| @@ -11,7 +11,7 @@ | |||
| 11 | ## 函数签名 | 11 | ## 函数签名 |
| 12 | 12 | ||
| 13 | ```python | 13 | ```python |
| 14 | -from mindspeed_ops.api.triton.chunk_gate_delta_rule import chunk_gated_delta_rule | 14 | +from mindspeed_ops.api.triton.chunk_gated_delta_rule import chunk_gated_delta_rule |
| 15 | 15 | ||
| 16 | def chunk_gated_delta_rule( | 16 | def chunk_gated_delta_rule( |
| 17 | q: torch.Tensor, | 17 | q: torch.Tensor, |
| @@ -41,7 +41,7 @@ def chunk_gated_delta_rule( | |||
| 41 | use_qk_l2norm_in_kernel: 是否对 q/k 做 L2 归一化，默认 False | 41 | use_qk_l2norm_in_kernel: 是否对 q/k 做 L2 归一化，默认 False |
| 42 | cu_seqlens: 变长序列累积长度，形状 [N+1]，默认 None | 42 | cu_seqlens: 变长序列累积长度，形状 [N+1]，默认 None |
| 43 | chunk_size: 分块大小，默认 64 | 43 | chunk_size: 分块大小，默认 64 |
| 44 | - head_first: 已弃用，仅支持 False | 44 | + head_first: 已弃用，传 True 仅告警仍可执行，建议使用 False |
| 45 | 45 | ||
| 46 | Returns: | 46 | Returns: |
| 47 | o: 输出张量，形状 [B, T, H, V] | 47 | o: 输出张量，形状 [B, T, H, V] |
| @@ -68,7 +68,7 @@ def chunk_gated_delta_rule( | |||
| 68 | | `use_qk_l2norm_in_kernel` | bool | - | 是否对q/k做L2归一化 | - | | 68 | | `use_qk_l2norm_in_kernel` | bool | - | 是否对q/k做L2归一化 | - | |
| 69 | | `cu_seqlens` | torch.LongTensor | \[N+1] | 变长序列累积长度 | int32, int64 | | 69 | | `cu_seqlens` | torch.LongTensor | \[N+1] | 变长序列累积长度 | int32, int64 | |
| 70 | | `chunk_size` | int | - | 分块大小，默认64 | - | | 70 | | `chunk_size` | int | - | 分块大小，默认64 | - | |
| 71 | -| `head_first` | bool | - | 已弃用，仅支持False | - | | 71 | +| `head_first` | bool | - | 已弃用，传True仅告警仍可执行，建议使用False | - | |
| 72 | 72 | ||
| 73 | ### 约束条件 | 73 | ### 约束条件 |
| 74 | 74 | ||
| @@ -99,7 +99,7 @@ def chunk_gated_delta_rule( | |||
| 99 | 对门控张量 `g`（log空间）沿序列维度按chunk做局部累积求和： | 99 | 对门控张量 `g`（log空间）沿序列维度按chunk做局部累积求和： |
| 100 | 100 | ||
| 101 | ```python | 101 | ```python |
| 102 | -g_cumsum = chunk_local_cumsum(g, chunk_size) | 102 | +g_cumsum = chunk_local_cumsum(g=g, chunk_size=chunk_size) |
| 103 | ``` | 103 | ``` |
| 104 | 104 | ||
| 105 | 这一步将log空间的门控值转换为chunk内的相对衰减因子，后续通过 `exp(g_i - g_j)` 实现因果衰减。 | 105 | 这一步将log空间的门控值转换为chunk内的相对衰减因子，后续通过 `exp(g_i - g_j)` 实现因果衰减。 |
| @@ -107,7 +107,7 @@ g_cumsum = chunk_local_cumsum(g, chunk_size) | |||
| 107 | #### 步骤2：计算KKT矩阵A | 107 | #### 步骤2：计算KKT矩阵A |
| 108 | 108 | ||
| 109 | ```python | 109 | ```python |
| 110 | -A = chunk_scaled_dot_kkt_fwd(k, g=g_cumsum, beta=beta) | 110 | +A = chunk_scaled_dot_kkt_fwd(k=k, g=g_cumsum, beta=beta) |
| 111 | ``` | 111 | ``` |
| 112 | 112 | ||
| 113 | 在每个chunk内计算 `beta * K @ K^T`，并应用下三角掩码和门控差值指数缩放。输出形状为 `[B, T, H, BT]`，其中 `BT = chunk_size`。 | 113 | 在每个chunk内计算 `beta * K @ K^T`，并应用下三角掩码和门控差值指数缩放。输出形状为 `[B, T, H, BT]`，其中 `BT = chunk_size`。 |
| @@ -115,7 +115,7 @@ A = chunk_scaled_dot_kkt_fwd(k, g=g_cumsum, beta=beta) | |||
| 115 | #### 步骤3：矩阵求逆 | 115 | #### 步骤3：矩阵求逆 |
| 116 | 116 | ||
| 117 | ```python | 117 | ```python |
| 118 | -A = solve_tril(A) | 118 | +A = solve_tril(A=A) |
| 119 | ``` | 119 | ``` |
| 120 | 120 | ||
| 121 | 计算 `(I + A)^-1`，将严格下三角的KKT矩阵转换为WY表示所需的下三角因子。`solve_tril` 通过16×16基础分块求逆 + 逐级合并（32×32 / 64×64）实现高效计算。 | 121 | 计算 `(I + A)^-1`，将严格下三角的KKT矩阵转换为WY表示所需的下三角因子。`solve_tril` 通过16×16基础分块求逆 + 逐级合并（32×32 / 64×64）实现高效计算。 |
| @@ -123,7 +123,7 @@ A = solve_tril(A) | |||
| 123 | #### 步骤4：重计算WY表示 | 123 | #### 步骤4：重计算WY表示 |
| 124 | 124 | ||
| 125 | ```python | 125 | ```python |
| 126 | -w, u = recompute_w_u_fwd(k, v, beta, A, g=g_cumsum) | 126 | +w, u = recompute_w_u_fwd(k=k, v=v, beta=beta, A=A, g=g_cumsum) |
| 127 | ``` | 127 | ``` |
| 128 | 128 | ||
| 129 | 通过矩阵A对加权后的键值对进行变换： | 129 | 通过矩阵A对加权后的键值对进行变换： |
| @@ -134,7 +134,7 @@ w, u = recompute_w_u_fwd(k, v, beta, A, g=g_cumsum) | |||
| 134 | #### 步骤5：计算隐藏状态H和残差值V\_new | 134 | #### 步骤5：计算隐藏状态H和残差值V\_new |
| 135 | 135 | ||
| 136 | ```python | 136 | ```python |
| 137 | -h, v_new, final_state = chunk_gated_delta_rule_fwd_h(k, w, u, g, initial_state) | 137 | +h, v_new, final_state = chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, initial_state=initial_state) |
| 138 | ``` | 138 | ``` |
| 139 | 139 | ||
| 140 | 沿序列迭代更新隐藏状态： | 140 | 沿序列迭代更新隐藏状态： |
| @@ -146,7 +146,7 @@ h, v_new, final_state = chunk_gated_delta_rule_fwd_h(k, w, u, g, initial_state) | |||
| 146 | #### 步骤6：计算最终输出O | 146 | #### 步骤6：计算最终输出O |
| 147 | 147 | ||
| 148 | ```python | 148 | ```python |
| 149 | -o = chunk_fwd_o(q, k, v_new, h, g, scale) | 149 | +o = chunk_fwd_o(q=q, k=k, v=v_new, h=h, g=g, scale=scale) |
| 150 | ``` | 150 | ``` |
| 151 | 151 | ||
| 152 | 在每个chunk内计算： | 152 | 在每个chunk内计算： |
| @@ -159,12 +159,12 @@ o = (q @ h + causal(q @ k^T) @ v_new) * scale | |||
| 159 | 159 | ||
| 160 | ### 反向传播 | 160 | ### 反向传播 |
| 161 | 161 | ||
| 162 | -反向传播沿前向的逆序执行，通过以下5个子步骤完成： | 162 | +反向传播沿前向的逆序执行，通过以下7个子步骤完成： |
| 163 | 163 | ||
| 164 | #### 步骤1：重计算WY表示 | 164 | #### 步骤1：重计算WY表示 |
| 165 | 165 | ||
| 166 | ```python | 166 | ```python |
| 167 | -w, u = recompute_w_u_fwd(k, v, beta, A, g) | 167 | +w, u = recompute_w_u_fwd(k=k, v=v, beta=beta, A=A, g=g) |
| 168 | ``` | 168 | ``` |
| 169 | 169 | ||
| 170 | 与前向步骤4相同，重新计算中间变量（checkpointing策略，节省显存）。 | 170 | 与前向步骤4相同，重新计算中间变量（checkpointing策略，节省显存）。 |
| @@ -172,7 +172,7 @@ w, u = recompute_w_u_fwd(k, v, beta, A, g) | |||
| 172 | #### 步骤2：重计算H和V\_new | 172 | #### 步骤2：重计算H和V\_new |
| 173 | 173 | ||
| 174 | ```python | 174 | ```python |
| 175 | -h, v_new, _ = chunk_gated_delta_rule_fwd_h(k, w, u, g, initial_state) | 175 | +h, v_new, _ = chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, initial_state=initial_state) |
| 176 | ``` | 176 | ``` |
| 177 | 177 | ||
| 178 | 与前向步骤5相同，重新计算隐藏状态。 | 178 | 与前向步骤5相同，重新计算隐藏状态。 |
| @@ -180,7 +180,7 @@ h, v_new, _ = chunk_gated_delta_rule_fwd_h(k, w, u, g, initial_state) | |||
| 180 | #### 步骤3：计算dV（局部贡献） | 180 | #### 步骤3：计算dV（局部贡献） |
| 181 | 181 | ||
| 182 | ```python | 182 | ```python |
| 183 | -dv = chunk_bwd_dv_local(q, k, g, do, scale) | 183 | +dv = chunk_bwd_dv_local(q=q, k=k, do=do, g=g, scale=scale) |
| 184 | ``` | 184 | ``` |
| 185 | 185 | ||
| 186 | 计算chunk内注意力对v的梯度：`dv = tril(k @ q^T * scale)^T @ do`。 | 186 | 计算chunk内注意力对v的梯度：`dv = tril(k @ q^T * scale)^T @ do`。 |
| @@ -188,7 +188,9 @@ dv = chunk_bwd_dv_local(q, k, g, do, scale) | |||
| 188 | #### 步骤4：计算dH, dH0, dV | 188 | #### 步骤4：计算dH, dH0, dV |
| 189 | 189 | ||
| 190 | ```python | 190 | ```python |
| 191 | -dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(q, k, w, g, h0, dht, do, dv, scale) | 191 | +dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( |
| 192 | + q=q, k=k, w=w, do=do, dv=dv, g=g, h0=initial_state, dht=dht, scale=scale | ||
| 193 | +) | ||
| 192 | ``` | 194 | ``` |
| 193 | 195 | ||
| 194 | 从最后一个chunk向前迭代，沿隐藏状态递推链反向传播： | 196 | 从最后一个chunk向前迭代，沿隐藏状态递推链反向传播： |
| @@ -200,7 +202,9 @@ dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(q, k, w, g, h0, dht, do, dv, scale) | |||
| 200 | #### 步骤5：计算dQ, dK, dW, dG | 202 | #### 步骤5：计算dQ, dK, dW, dG |
| 201 | 203 | ||
| 202 | ```python | 204 | ```python |
| 203 | -dq, dk, dw, dg = chunk_bwd_dqkwg(q, k, v_new, w, g, h, dv, do, dh, scale) | 205 | +dq, dk, dw, dg = chunk_bwd_dqkwg( |
| 206 | + q=q, k=k, v=v_new, do=do, h=h, dh=dh, g=g, dv=dv, w=w, scale=scale | ||
| 207 | +) | ||
| 204 | ``` | 208 | ``` |
| 205 | 209 | ||
| 206 | 在每个chunk内累积注意力分数梯度 `dS`，并反向传播到 `q`、`k`、`w`、`g`。 | 210 | 在每个chunk内累积注意力分数梯度 `dS`，并反向传播到 `q`、`k`、`w`、`g`。 |
| @@ -208,10 +212,12 @@ dq, dk, dw, dg = chunk_bwd_dqkwg(q, k, v_new, w, g, h, dv, do, dh, scale) | |||
| 208 | #### 步骤6：WY表示反向 | 212 | #### 步骤6：WY表示反向 |
| 209 | 213 | ||
| 210 | ```python | 214 | ```python |
| 211 | -dk2, dv, db, dg2 = prepare_wy_repr_bwd(k, v, beta, g, A, dw, dv) | 215 | +dk2, dv, db, dg2 = prepare_wy_repr_bwd( |
| 216 | + k=k, v=v, g=g, beta=beta, A=A, dw=dw, du=dv | ||
| 217 | +) | ||
| 212 | ``` | 218 | ``` |
| 213 | 219 | ||
| 214 | -将 `dw` 和 `dv` 的梯度通过WY表示反向传播到 `k`、`v`、`β`、`g`。涉及矩阵逆求导链： | 220 | +将 `dw` 与 `du` 的梯度通过WY表示反向传播到 `k`、`v`、`β`、`g`。内核中该参数名为 `du`，本实现传入的是步骤4/5得到的值梯度 `dv`（即调用时 `du=dv`）。涉及矩阵逆求导链： |
| 215 | 221 | ||
| 216 | ```python | 222 | ```python |
| 217 | ∂L/∂M = -A^T @ ∂L/∂A @ A^T | 223 | ∂L/∂M = -A^T @ ∂L/∂A @ A^T |
| @@ -220,7 +226,7 @@ dk2, dv, db, dg2 = prepare_wy_repr_bwd(k, v, beta, g, A, dw, dv) | |||
| 220 | #### 步骤7：门控梯度Cumsum反向 | 226 | #### 步骤7：门控梯度Cumsum反向 |
| 221 | 227 | ||
| 222 | ```python | 228 | ```python |
| 223 | -dg = chunk_local_cumsum(dg, chunk_size, reverse=True) | 229 | +dg = chunk_local_cumsum(g=dg, chunk_size=chunk_size, reverse=True) |
| 224 | ``` | 230 | ``` |
| 225 | 231 | ||
| 226 | 将chunk内的门控梯度通过反向cumsum还原到逐token的梯度。 | 232 | 将chunk内的门控梯度通过反向cumsum还原到逐token的梯度。 |
| @@ -248,7 +254,7 @@ g = F.logsigmoid(torch.rand(B, T, H, device=device, dtype=dtype)) | |||
| 248 | h0 = None | 254 | h0 = None |
| 249 | 255 | ||
| 250 | o, ht = chunk_gated_delta_rule( | 256 | o, ht = chunk_gated_delta_rule( |
| 251 | - q, k, v, g, beta, | 257 | + q=q, k=k, v=v, g=g, beta=beta, |
| 252 | initial_state=h0, | 258 | initial_state=h0, |
| 253 | output_final_state=True, | 259 | output_final_state=True, |
| 254 | ) | 260 | ) |
| @@ -284,7 +290,7 @@ cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) | |||
| 284 | h0 = None | 290 | h0 = None |
| 285 | 291 | ||
| 286 | o, ht = chunk_gated_delta_rule( | 292 | o, ht = chunk_gated_delta_rule( |
| 287 | - q, k, v, g, beta, | 293 | + q=q, k=k, v=v, g=g, beta=beta, |
| 288 | initial_state=h0, | 294 | initial_state=h0, |
| 289 | output_final_state=True, | 295 | output_final_state=True, |
| 290 | cu_seqlens=cu_seqlens, | 296 | cu_seqlens=cu_seqlens, |
| @@ -340,7 +346,7 @@ Gated DeltaNet训练时的典型输入（chunk\_size=64）： | |||
| 340 | 3. **beta维度要求**：`beta` 必须为3维张量 `[B, T, H]`，不支持其他形状。 | 346 | 3. **beta维度要求**：`beta` 必须为3维张量 `[B, T, H]`，不支持其他形状。 |
| 341 | 4. **变长序列约束**：使用 `cu_seqlens` 时，batch size必须为1，且 `initial_state`参数只能传none。 | 347 | 4. **变长序列约束**：使用 `cu_seqlens` 时，batch size必须为1，且 `initial_state`参数只能传none。 |
| 342 | 5. **chunk\_size选择**：默认值为64，需与下游子算子（如 `solve_tril` 仅支持BT=16/32/64）兼容。 | 348 | 5. **chunk\_size选择**：默认值为64，需与下游子算子（如 `solve_tril` 仅支持BT=16/32/64）兼容。 |
| 343 | -6. **head\_first已弃用**：该参数已废弃，仅支持 `head_first=False`。 | 349 | +6. **head\_first已弃用**：该参数已废弃，传入 `head_first=True` 仅触发告警、仍会继续执行（该参数不参与实际计算），建议使用 `head_first=False`。 |
| 344 | 7. **门控g的数值范围**：`g` 为log空间的门控值，内部 `exp(g_diff)` 计算会将其clamp到 `[-9, 0]` 范围，避免指数爆炸。 | 350 | 7. **门控g的数值范围**：`g` 为log空间的门控值，内部 `exp(g_diff)` 计算会将其clamp到 `[-9, 0]` 范围，避免指数爆炸。 |
| 345 | 8. **输入幅度建议**：`k`、`β`、`g` 的输入量级建议控制在 \~0.1范围内，过大的值可能导致矩阵求逆发散。 | 351 | 8. **输入幅度建议**：`k`、`β`、`g` 的输入量级建议控制在 \~0.1范围内，过大的值可能导致矩阵求逆发散。 |
| 346 | 9. **设备要求**：Triton内核需要在NPU或GPU设备上执行，CPU环境不支持。 | 352 | 9. **设备要求**：Triton内核需要在NPU或GPU设备上执行，CPU环境不支持。 |
| @@ -36,7 +36,7 @@ def clipped_swiglu( | |||
| 36 | |------|------|------|------|---------| | 36 | |------|------|------|------|---------| |
| 37 | | `gate_up` | torch.Tensor | `[..., 2K]` | 输入张量，末维为偶数，前K列为gate、后K列为up | 是 | | 37 | | `gate_up` | torch.Tensor | `[..., 2K]` | 输入张量，末维为偶数，前K列为gate、后K列为up | 是 | |
| 38 | | `swiglu_alpha` | float | - | SwiGLU的alpha参数 | 否，默认 `1.702` | | 38 | | `swiglu_alpha` | float | - | SwiGLU的alpha参数 | 否，默认 `1.702` | |
| 39 | -| `swiglu_limit` | float | - | clamp上下限L，gate上限clamp到L，up双向clamp到 `[-L, L]` | 否，默认 `7.0` | | 39 | +| `swiglu_limit` | float | - | clamp上下限L（必须≥0），gate上限clamp到L，up双向clamp到 `[-L, L]` | 否，默认 `7.0` | |
| 40 | 40 | ||
| 41 | ### 前向输出 | 41 | ### 前向输出 |
| 42 | 42 | ||
| @@ -117,34 +117,34 @@ assert y.shape == (B, T, K) | |||
| 117 | 117 | ||
| 118 | ### bf16 | 118 | ### bf16 |
| 119 | 119 | ||
| 120 | -| shape | eager(ms) | triton(ms) | speedup | bw(GB/s) | | 120 | +| shape | eager(ms) | triton(ms) | speedup | |
| 121 | -|---|---|---|---|---| | 121 | +|---|---|---|---| |
| 122 | -| (1, 8192) | 0.056 | 0.029 | 1.96× | 863 | | 122 | +| (1, 8192) | 0.056 | 0.029 | 1.96× | |
| 123 | -| (128, 8192) | 0.066 | 0.029 | 2.27× | 108912 | | 123 | +| (128, 8192) | 0.066 | 0.029 | 2.27× | |
| 124 | -| (1024, 8192) | 0.221 | 0.052 | 4.26× | 484314 | | 124 | +| (1024, 8192) | 0.221 | 0.052 | 4.26× | |
| 125 | -| (8192, 8192) | 1.685 | 0.394 | 4.28× | 511687 | | 125 | +| (8192, 8192) | 1.685 | 0.394 | 4.28× | |
| 126 | -| (1, 4096, 24576) | 2.902 | 0.541 | 5.37× | 558400 | | 126 | +| (1, 4096, 24576) | 2.902 | 0.541 | 5.37× | |
| 127 | -| (1, 8192, 24576) | 5.798 | 1.072 | 5.41× | 563449 | | 127 | +| (1, 8192, 24576) | 5.798 | 1.072 | 5.41× | |
| 128 | -| (4, 4096, 24576) | 11.536 | 2.138 | 5.40× | 565097 | | 128 | +| (4, 4096, 24576) | 11.536 | 2.138 | 5.40× | |
| 129 | 129 | ||
| 130 | ### fp16 | 130 | ### fp16 |
| 131 | 131 | ||
| 132 | -| shape | eager(ms) | triton(ms) | speedup | bw(GB/s) | | 132 | +| shape | eager(ms) | triton(ms) | speedup | |
| 133 | -|---|---|---|---|---| | 133 | +|---|---|---|---| |
| 134 | -| (1024, 8192) | 0.240 | 0.072 | 3.35× | 352026 | | 134 | +| (1024, 8192) | 0.240 | 0.072 | 3.35× | |
| 135 | -| (8192, 8192) | 1.686 | 0.487 | 3.47× | 413849 | | 135 | +| (8192, 8192) | 1.686 | 0.487 | 3.47× | |
| 136 | -| (1, 4096, 24576) | 2.687 | 0.602 | 4.47× | 501783 | | 136 | +| (1, 4096, 24576) | 2.687 | 0.602 | 4.47× | |
| 137 | -| (4, 4096, 24576) | 10.903 | 2.226 | 4.90× | 542669 | | 137 | +| (4, 4096, 24576) | 10.903 | 2.226 | 4.90× | |
| 138 | 138 | ||
| 139 | ### fp32 | 139 | ### fp32 |
| 140 | 140 | ||
| 141 | -| shape | eager(ms) | triton(ms) | speedup | bw(GB/s) | | 141 | +| shape | eager(ms) | triton(ms) | speedup | |
| 142 | -|---|---|---|---|---| | 142 | +|---|---|---|---| |
| 143 | -| (1024, 8192) | 0.296 | 0.117 | 2.54× | 431975 | | 143 | +| (1024, 8192) | 0.296 | 0.117 | 2.54× | |
| 144 | -| (8192, 8192) | 3.790 | 0.542 | 7.00× | 743185 | | 144 | +| (8192, 8192) | 3.790 | 0.542 | 7.00× | |
| 145 | -| (1, 4096, 24576) | 5.596 | 0.747 | 7.49× | 808788 | | 145 | +| (1, 4096, 24576) | 5.596 | 0.747 | 7.49× | |
| 146 | -| (1, 8192, 24576) | 10.569 | 1.283 | 8.24× | 941380 | | 146 | +| (1, 8192, 24576) | 10.569 | 1.283 | 8.24× | |
| 147 | -| (4, 4096, 24576) | 21.150 | 2.576 | 8.21× | 937771 | | 147 | +| (4, 4096, 24576) | 21.150 | 2.576 | 8.21× | |
| 148 | 148 | ||
| 149 | 3D shape `(1, 4096, 24576)` 对应K=12288（LLM hidden常见尺寸），三dtype加速：fp32 8.2× / bf16 5.4× / fp16 4.5×。bf16/fp16的瓶颈为fp32超越函数（sigmoid）计算，fp32已近该算子上限。 | 149 | 3D shape `(1, 4096, 24576)` 对应K=12288（LLM hidden常见尺寸），三dtype加速：fp32 8.2× / bf16 5.4× / fp16 4.5×。bf16/fp16的瓶颈为fp32超越函数（sigmoid）计算，fp32已近该算子上限。 |
| 150 | 150 | ||
| @@ -20,16 +20,16 @@ def chunk_local_cumsum( | |||
| 20 | ) -> torch.Tensor: | 20 | ) -> torch.Tensor: |
| 21 | """ | 21 | """ |
| 22 | Args: | 22 | Args: |
| 23 | - g: 输入张量，shape 为 (B, T, H) or (B, T, H, K) | 23 | + g: 输入张量，shape 为 (B, T, H) 或 (B, T, H, K) |
| 24 | chunk_size: 分块大小，必须是 2 的幂次 | 24 | chunk_size: 分块大小，必须是 2 的幂次 |
| 25 | reverse: 是否在每个 chunk 内反向执行 cumsum，默认 False | 25 | reverse: 是否在每个 chunk 内反向执行 cumsum，默认 False |
| 26 | scale: 输出缩放因子，默认 None | 26 | scale: 输出缩放因子，默认 None |
| 27 | cu_seqlens: 变长序列的累积长度张量，默认 None。提供时 batch size 必须为 1 | 27 | cu_seqlens: 变长序列的累积长度张量，默认 None。提供时 batch size 必须为 1 |
| 28 | - head_first: 是否 head 维度在前，仅支持 False | 28 | + head_first: 是否 head 维度在前，arch32 路径仅支持 False |
| 29 | output_dtype: 输出数据类型，默认 torch.float | 29 | output_dtype: 输出数据类型，默认 torch.float |
| 30 | - chunk_indices: 分块索引张量，shape 为 (B, T // chunk_size), 默认为 None | 30 | + chunk_indices: 分块索引张量，shape 为 (num_chunks, 2)，每行为 (序列编号, 序列内chunk编号)，默认为 None |
| 31 | Returns: | 31 | Returns: |
| 32 | - 分块局部 cumsum 结果，shape 为 (B, T, H) | 32 | + 分块局部 cumsum 结果，shape 为 (B, T, H) 或 (B, T, H, K)，与输入 shape 一致 |
| 33 | """ | 33 | """ |
| 34 | ``` | 34 | ``` |
| 35 | 35 | ||
| @@ -37,19 +37,20 @@ def chunk_local_cumsum( | |||
| 37 | 37 | ||
| 38 | | 参数 | 类型 | 描述 | 支持的数据类型 | | 38 | | 参数 | 类型 | 描述 | 支持的数据类型 | |
| 39 | |------|------|------|---------------| | 39 | |------|------|------|---------------| |
| 40 | -| g | torch.Tensor | 输入张量，shape为 (B, T, H) | float32, float16, bfloat16 | | 40 | +| g | torch.Tensor | 输入张量，shape为 (B, T, H) 或 (B, T, H, K) | float32, float16, bfloat16 | |
| 41 | | chunk_size | int | 分块大小，必须是2的幂次 | - | | 41 | | chunk_size | int | 分块大小，必须是2的幂次 | - | |
| 42 | | reverse | bool | 是否在chunk内反向cumsum，默认False | - | | 42 | | reverse | bool | 是否在chunk内反向cumsum，默认False | - | |
| 43 | | scale | float | 输出缩放因子，默认None | - | | 43 | | scale | float | 输出缩放因子，默认None | - | |
| 44 | | cu_seqlens | torch.Tensor | 变长序列的累积长度，shape为 (num_seqs + 1,)，默认None | int32 | | 44 | | cu_seqlens | torch.Tensor | 变长序列的累积长度，shape为 (num_seqs + 1,)，默认None | int32 | |
| 45 | -| head_first | bool | 是否head维度在前，仅支持False | - | | 45 | +| head_first | bool | 是否head维度在前，arch32路径仅支持False | - | |
| 46 | | output_dtype | torch.dtype | 输出数据类型，默认torch.float | - | | 46 | | output_dtype | torch.dtype | 输出数据类型，默认torch.float | - | |
| 47 | +| chunk_indices | torch.LongTensor | 分块索引张量，shape为 (num_chunks, 2)，每行为 (序列编号, 序列内chunk编号)，默认None；仅在arch32路径的4维输入（vector内核）下生效，3维输入与arch35路径会忽略该参数；为None且提供cu_seqlens时由cu_seqlens自动推导 | int32, int64 | | ||
| 47 | 48 | ||
| 48 | ## 返回值 | 49 | ## 返回值 |
| 49 | 50 | ||
| 50 | | 参数 | 类型 | 描述 | 支持的数据类型 | | 51 | | 参数 | 类型 | 描述 | 支持的数据类型 | |
| 51 | | ------|------|------|---------------| | 52 | | ------|------|------|---------------| |
| 52 | -| output | torch.Tensor | 分块局部cumsum结果，shape为 (B, T, H) | 默认float32，可通过output_dtype指定 | | 53 | +| output | torch.Tensor | 分块局部cumsum结果，shape为 (B, T, H) 或 (B, T, H, K)，与输入shape一致；4维时最后一维K为独立通道维，各通道沿T维独立累加 | 默认float32，可通过output_dtype指定 | |
| 53 | 54 | ||
| 54 | ## 实现原理 | 55 | ## 实现原理 |
| 55 | 56 | ||
| @@ -59,7 +60,7 @@ def chunk_local_cumsum( | |||
| 59 | 2. **Tiling策略**：根据 `reverse` 模式和head维度大小H动态计算分块大小 `BT`，平衡并行度和UB（Unified Buffer）利用率： | 60 | 2. **Tiling策略**：根据 `reverse` 模式和head维度大小H动态计算分块大小 `BT`，平衡并行度和UB（Unified Buffer）利用率： |
| 60 | - `BT = max(chunk_size, next_power_of_2((1 << 11 if reverse else 1 << 12) // H))` | 61 | - `BT = max(chunk_size, next_power_of_2((1 << 11 if reverse else 1 << 12) // H))` |
| 61 | 3. **变长序列处理**：若提供 `cu_seqlens`，调用 `prepare_chunk_indices` 生成chunk索引，支持不同长度的序列独立分块。 | 62 | 3. **变长序列处理**：若提供 `cu_seqlens`，调用 `prepare_chunk_indices` 生成chunk索引，支持不同长度的序列独立分块。 |
| 62 | -4. **启动内核**：调用Triton内核 `chunk_local_cumsum_scalar_kernel`，按 `(NT, B)` 的grid并行执行，每个chunk内独立完成cumsum计算。 | 63 | +4. **启动内核**：arch32路径下，3维输入调用Triton内核 `chunk_local_cumsum_scalar_kernel`，按 `(NT, B)` 的grid并行执行，每个chunk内独立完成cumsum计算；4维输入调用 `chunk_local_cumsum_vector_kernel`，此时 `BT = chunk_size`，按 `(cdiv(S, BS), NT, B * H)` 的grid并行执行（`BS` 固定为32），每个chunk内沿T维独立完成cumsum计算，最后一维S按 `BS` 切分。arch35路径由独立实现承接，3维与4维输入均落到scalar内核（4维先reshape为3维）。 |
| 63 | 5. **后处理**：若指定 `scale`，在内核中对结果进行缩放。 | 64 | 5. **后处理**：若指定 `scale`，在内核中对结果进行缩放。 |
| 64 | 65 | ||
| 65 | ## 使用示例 | 66 | ## 使用示例 |
| @@ -127,8 +128,8 @@ print("varlen chunk_local_cumsum operation exec successfully!") | |||
| 127 | ## 注意事项 | 128 | ## 注意事项 |
| 128 | 129 | ||
| 129 | 1. `chunk_size` 必须是2的幂次，否则会抛出 `ValueError`。 | 130 | 1. `chunk_size` 必须是2的幂次，否则会抛出 `ValueError`。 |
| 130 | -2. 当前算子仅支持3维输入 `(B, T, H)`，不支持4维输入 `(B, T, H, D)`。 | 131 | +2. 算子同时支持3维输入 `(B, T, H)` 与4维输入 `(B, T, H, K)`。arch32路径按输入维度分发：3维走scalar内核，4维走vector内核；arch35路径由独立实现承接，其4维输入会先reshape为3维后复用scalar内核。 |
| 131 | 3. 使用 `cu_seqlens` 时，batch size必须为1，否则会抛出 `ValueError`。 | 132 | 3. 使用 `cu_seqlens` 时，batch size必须为1，否则会抛出 `ValueError`。 |
| 132 | -4. 该算子在arch35平台上暂不支持，会抛出 `NotImplementedError`。 | 133 | +4. 该算子同时支持arch32与arch35架构，运行时自动分发至对应架构的实现。arch32路径下 `head_first=True` 会抛出 `NotImplementedError`。 |
| 133 | 5. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销和Triton下发可能超过计算收益。 | 134 | 5. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销和Triton下发可能超过计算收益。 |
| 134 | 6. `output_dtype` 默认为 `torch.float`，在反向传播和context parallel场景下可有效防止中间结果溢出。 | 135 | 6. `output_dtype` 默认为 `torch.float`，在反向传播和context parallel场景下可有效防止中间结果溢出。 |
| @@ -37,7 +37,7 @@ class FusedCrossEntropyLossFunction(torch.autograd.Function): | |||
| 37 | 37 | ||
| 38 | | 参数 | 类型 | 描述 | 支持的数据类型 | | 38 | | 参数 | 类型 | 描述 | 支持的数据类型 | |
| 39 | | ------ |------|------|---------------| | 39 | | ------ |------|------|---------------| |
| 40 | -| loss | torch.Tensor | [bs * seq_len] | 与输入张量相同的数据类型 | | 40 | +| loss | torch.Tensor | 标量（所有token加权交叉熵之和） | float32（内核以FP32计算输出） | |
| 41 | 41 | ||
| 42 | ## 实现原理 | 42 | ## 实现原理 |
| 43 | 43 | ||
| @@ -54,6 +54,8 @@ import torch | |||
| 54 | from mindspeed_ops.api.triton.fused_cross_entropy_loss import FusedCrossEntropyLossFunction | 54 | from mindspeed_ops.api.triton.fused_cross_entropy_loss import FusedCrossEntropyLossFunction |
| 55 | 55 | ||
| 56 | # 创建测试张量 | 56 | # 创建测试张量 |
| 57 | +dtype = torch.float32 | ||
| 58 | +device = torch.device("npu") | ||
| 57 | logits = torch.randn((1024, 16384), requires_grad=True, dtype=dtype).to(device) | 59 | logits = torch.randn((1024, 16384), requires_grad=True, dtype=dtype).to(device) |
| 58 | loss_weight = torch.randn((1024, ), requires_grad=True, dtype=dtype).to(device) | 60 | loss_weight = torch.randn((1024, ), requires_grad=True, dtype=dtype).to(device) |
| 59 | labels = torch.randint(0, 16384, (1024,), dtype=torch.int32).to(device) | 61 | labels = torch.randint(0, 16384, (1024,), dtype=torch.int32).to(device) |
| @@ -72,10 +72,10 @@ bmm1[b, s, i, d] = h_post[b, s, i] * x[b, s, d] # [B, S, N, D] | |||
| 72 | **分支2（bmm2）**：residual与h_res的矩阵乘法 | 72 | **分支2（bmm2）**：residual与h_res的矩阵乘法 |
| 73 | 73 | ||
| 74 | ```python | 74 | ```python |
| 75 | -bmm2[b, s, i, d] = sum_j(h_res[b, s, i, j] * residual[b, s, j, d]) # [B, S, N, D] | 75 | +bmm2[b, s, i, d] = sum_j(h_res[b, s, j, i] * residual[b, s, j, d]) # [B, S, N, D] |
| 76 | ``` | 76 | ``` |
| 77 | 77 | ||
| 78 | -对N个stream的残差按h_res矩阵进行线性组合。 | 78 | +对N个stream的残差按h_res矩阵进行线性组合。公式中的 `h_res` 采用**用户输入的**排列，并非内部转置后的张量：源码内部前向会先对 `h_res` 执行 `permute(0,1,3,2)` 转置，再参与bmm2计算；反向得到的 `grad_h_res` 同样按用户输出排列书写，源码内部会在返回前转置还原。 |
| 79 | 79 | ||
| 80 | **合并**：逐元素相加 | 80 | **合并**：逐元素相加 |
| 81 | 81 | ||
| @@ -97,12 +97,12 @@ grad_h_post[b, s, i] = sum_d(grad_out[b, s, i, d] * x[b, s, d]) | |||
| 97 | **bmm2反向**： | 97 | **bmm2反向**： |
| 98 | 98 | ||
| 99 | ```python | 99 | ```python |
| 100 | -grad_residual[b, s, j, d] = sum_i(grad_out[b, s, i, d] * h_res[b, s, i, j]) | 100 | +grad_residual[b, s, j, d] = sum_i(grad_out[b, s, i, d] * h_res[b, s, j, i]) |
| 101 | -grad_h_res[b, s, i, j] = sum_d(grad_out[b, s, i, d] * residual[b, s, j, d]) | 101 | +grad_h_res[b, s, i, j] = sum_d(grad_out[b, s, j, d] * residual[b, s, i, d]) |
| 102 | ``` | 102 | ``` |
| 103 | 103 | ||
| 104 | - `grad_h_post` 和 `grad_h_res` 使用 `atomic_add` 累加，支持D维度分块。 | 104 | - `grad_h_post` 和 `grad_h_res` 使用 `atomic_add` 累加，支持D维度分块。 |
| 105 | -- `h_res` 在前向中被permute(0,1,3,2)，反向输出时还原为原始排列。 | 105 | +- 上述公式中的 `h_res` / `grad_h_res` 排列约定见上文bmm2分支说明。 |
| 106 | 106 | ||
| 107 | ### 梯度重计算 | 107 | ### 梯度重计算 |
| 108 | 108 | ||
| @@ -159,7 +159,7 @@ assert result.shape == (B, S, N, D) | |||
| 159 | ## 注意事项 | 159 | ## 注意事项 |
| 160 | 160 | ||
| 161 | 1. `h_post` 最后一维N必须与 `residual` 倒数第二维一致。 | 161 | 1. `h_post` 最后一维N必须与 `residual` 倒数第二维一致。 |
| 162 | -2. `h_res` 形状必须为 [B, S, N, N]，内部会执行 `permute(0,1,3,2)` 转置。 | 162 | +2. `h_res` 形状必须为 [B, S, N, N]，排列约定见上文。 |
| 163 | 3. `residual` 的前两维 (B, S) 必须与 `x` 一致，最后两维必须为 (N, D)。 | 163 | 3. `residual` 的前两维 (B, S) 必须与 `x` 一致，最后两维必须为 (N, D)。 |
| 164 | 4. 内部计算使用FP32精度，确保数值稳定性。 | 164 | 4. 内部计算使用FP32精度，确保数值稳定性。 |
| 165 | 5. 反向计算中 `grad_h_post` 和 `grad_h_res` 使用 `atomic_add` 累加。 | 165 | 5. 反向计算中 `grad_h_post` 和 `grad_h_res` 使用 `atomic_add` 累加。 |
| @@ -55,7 +55,7 @@ def mhc_pre_only_triton( | |||
| 55 | | `norm_gamma` | torch.Tensor | [n_D] | RMS归一化可学习缩放参数 | 是 | BF16, FP32 | | 55 | | `norm_gamma` | torch.Tensor | [n_D] | RMS归一化可学习缩放参数 | 是 | BF16, FP32 | |
| 56 | | `mhc_use_gamma` | bool | - | RMS Norm是否应用gamma | 否，默认True | - | | 56 | | `mhc_use_gamma` | bool | - | RMS Norm是否应用gamma | 否，默认True | - | |
| 57 | | `num_stream` | int | - | stream数量，必须整除n_D | 否，默认4 | - | | 57 | | `num_stream` | int | - | stream数量，必须整除n_D | 否，默认4 | - | |
| 58 | -| `eps` | float | - | RMS归一化防除零常数 | 否，默认1e-6 | BF16, FP32 | | 58 | +| `eps` | float | - | RMS归一化防除零常数 | 否，默认1e-6 | - | |
| 59 | | `recompute_info` | - | - | 梯度重计算信息对象 | 否 | - | | 59 | | `recompute_info` | - | - | 梯度重计算信息对象 | 否 | - | |
| 60 | 60 | ||
| 61 | ### 输出 | 61 | ### 输出 |
| @@ -76,7 +76,7 @@ mhc_pre_only_triton将以下步骤融合为一次前向计算： | |||
| 76 | 1. **RMS归一化**：对输入x进行RMS Norm。 | 76 | 1. **RMS归一化**：对输入x进行RMS Norm。 |
| 77 | 77 | ||
| 78 | ```python | 78 | ```python |
| 79 | - x_norm = x / sqrt(mean(x^2) + eps) * gamma | 79 | + x_norm = x / sqrt(mean(x^2) + eps) * norm_gamma |
| 80 | ``` | 80 | ``` |
| 81 | 81 | ||
| 82 | 2. **线性投影**：将归一化后的输入投影到num_stream维度。 | 82 | 2. **线性投影**：将归一化后的输入投影到num_stream维度。 |
| @@ -4,6 +4,8 @@ | |||
| 4 | 4 | ||
| 5 | recompute\_w\_u\_fwd是一个基于Triton实现的WY表示重计算算子，用于在反向传播过程中重新计算中间变量w和u。该算子是gated delta rule的核心组件，通过矩阵A对键值对进行加权变换，支持可选的门控机制和变长序列处理，充分利用昇腾NPU的并行计算能力实现高性能计算。 | 5 | recompute\_w\_u\_fwd是一个基于Triton实现的WY表示重计算算子，用于在反向传播过程中重新计算中间变量w和u。该算子是gated delta rule的核心组件，通过矩阵A对键值对进行加权变换，支持可选的门控机制和变长序列处理，充分利用昇腾NPU的并行计算能力实现高性能计算。 |
| 6 | 6 | ||
| 7 | +本文中的arch32/arch35为昇腾NPU架构代号（arch32对应Ascend 910B/910_93系列，arch35对应Ascend 950系列），运行时通过 `mindspeed_ops.utils.is_arch35()` 自动判别并分发到对应架构的实现。 | ||
| 8 | + | ||
| 7 | *** | 9 | *** |
| 8 | 10 | ||
| 9 | ## 函数签名 | 11 | ## 函数签名 |
| @@ -91,7 +93,7 @@ def recompute_w_u_fwd( | |||
| 91 | 1. **分块计算**：将工作按 `[BT, BK]` 和 `[BT, BV]` 的tile切分，充分利用NPU Unified Buffer。 | 93 | 1. **分块计算**：将工作按 `[BT, BK]` 和 `[BT, BV]` 的tile切分，充分利用NPU Unified Buffer。 |
| 92 | 2. **持久化kernel**：采用grid-stride loop，每个AI核轮流处理多个tile，减少调度开销。 | 94 | 2. **持久化kernel**：采用grid-stride loop，每个AI核轮流处理多个tile，减少调度开销。 |
| 93 | 3. **门控融合**：将beta、g、gk的指数运算融合到矩阵乘法中，减少中间结果的存储和访问。 | 95 | 3. **门控融合**：将beta、g、gk的指数运算融合到矩阵乘法中，减少中间结果的存储和访问。 |
| 94 | -4. **变长序列支持**：通过 `cu_seqlens` 和 `chunk_indices` 实现对变长序列的高效处理。 | 96 | +4. **变长序列支持**：通过 `cu_seqlens` 推导出chunk索引（`chunk_indices`，非入参，由实现内部自动生成），实现对变长序列的高效处理。 |
| 95 | 97 | ||
| 96 | ### 计算流程 | 98 | ### 计算流程 |
| 97 | 99 | ||
| @@ -183,7 +185,7 @@ w, u = recompute_w_u_fwd(k=k, v=v, beta=beta, A=A, cu_seqlens=cu_seqlens) | |||
| 183 | 185 | ||
| 184 | ## 适用场景 | 186 | ## 适用场景 |
| 185 | 187 | ||
| 186 | -1. **Gated Delta Rule模型训练**：recompute\_w\_u\_fwd是gated delta rule前向计算的核心组件，用于重新计算中间变量。 | 188 | +1. **Gated Delta Rule模型训练**：recompute\_w\_u\_fwd是gated delta rule反向重计算的核心组件，用于重新计算中间变量。 |
| 187 | 2. **变长序列训练**：通过 `cu_seqlens` 支持将不同长度的样本高效打包进同一batch。 | 189 | 2. **变长序列训练**：通过 `cu_seqlens` 支持将不同长度的样本高效打包进同一batch。 |
| 188 | 3. **反向传播优化**：在反向传播过程中重新计算中间结果，减少显存占用。 | 190 | 3. **反向传播优化**：在反向传播过程中重新计算中间结果，减少显存占用。 |
| 189 | 191 | ||
| @@ -49,11 +49,13 @@ import torch | |||
| 49 | from mindspeed_ops.api.triton.rmsnorm_without_weight import RmsnormWithoutWeightFunction | 49 | from mindspeed_ops.api.triton.rmsnorm_without_weight import RmsnormWithoutWeightFunction |
| 50 | 50 | ||
| 51 | # 创建测试张量 | 51 | # 创建测试张量 |
| 52 | -x = torch.randn(1024, 1, 16384) | 52 | +dtype = torch.float32 |
| 53 | +device = torch.device("npu") | ||
| 54 | +x = torch.randn(1024, 1, 16384, dtype=dtype).to(device) | ||
| 53 | norm_eps = 1e-6 | 55 | norm_eps = 1e-6 |
| 54 | 56 | ||
| 55 | # 执行RMSNorm归一化计算 | 57 | # 执行RMSNorm归一化计算 |
| 56 | -triton_result = RmsnormWithoutWeightFunction.apply(x_triton, norm_eps) | 58 | +triton_result = RmsnormWithoutWeightFunction.apply(x, norm_eps) |
| 57 | torch_result = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + norm_eps) | 59 | torch_result = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + norm_eps) |
| 58 | 60 | ||
| 59 | # 验证结果 | 61 | # 验证结果 |
| @@ -26,7 +26,7 @@ def rms_norm_gated( | |||
| 26 | x: 输入张量，形状 [B, T, D] 或 [T, D] | 26 | x: 输入张量，形状 [B, T, D] 或 [T, D] |
| 27 | g: 门控张量，形状与 x 相同 | 27 | g: 门控张量，形状与 x 相同 |
| 28 | weight: 可学习缩放权重，形状 [D] | 28 | weight: 可学习缩放权重，形状 [D] |
| 29 | - bias: 偏置，形状 [D]（可选，当前实现固定为 None） | 29 | + bias: 偏置，形状 [D]（必传参数，无默认值；传 `None` 表示不加偏置） |
| 30 | activation: 门控激活函数类型，支持 'swish'/'silu'/'sigmoid' | 30 | activation: 门控激活函数类型，支持 'swish'/'silu'/'sigmoid' |
| 31 | residual: 残差输入，形状与 x 相同（可选） | 31 | residual: 残差输入，形状与 x 相同（可选） |
| 32 | prenorm: 是否启用 prenorm 模式，启用时返回 (y, residual_out) | 32 | prenorm: 是否启用 prenorm 模式，启用时返回 (y, residual_out) |
| @@ -48,7 +48,7 @@ def rms_norm_gated( | |||
| 48 | | `x` | torch.Tensor | [B, S, D] 或 [T, D] | 输入张量，B为batch大小，S为序列长度，T为BS长度，D为特征维度 | 是 | | 48 | | `x` | torch.Tensor | [B, S, D] 或 [T, D] | 输入张量，B为batch大小，S为序列长度，T为BS长度，D为特征维度 | 是 | |
| 49 | | `g` | torch.Tensor | 同x | 门控张量，与x形状完全一致 | 是 | | 49 | | `g` | torch.Tensor | 同x | 门控张量，与x形状完全一致 | 是 | |
| 50 | | `weight` | torch.Tensor | [D] | 可学习缩放权重（gamma） | 是 | | 50 | | `weight` | torch.Tensor | [D] | 可学习缩放权重（gamma） | 是 | |
| 51 | -| `bias` | torch.Tensor | [D] | 偏置（当前实现固定为None） | 否 | | 51 | +| `bias` | torch.Tensor | [D] | 偏置；必传参数（无默认值），传 `None` 时不加偏置 | 是（无默认值，可传 `None`） | |
| 52 | | `activation` | str | - | 门控激活函数，支持 `'swish'`/`'silu'`/`'sigmoid'` | 否，默认 "swish" | | 52 | | `activation` | str | - | 门控激活函数，支持 `'swish'`/`'silu'`/`'sigmoid'` | 否，默认 "swish" | |
| 53 | | `residual` | torch.Tensor | 同x | 残差输入，在归一化前与x相加 | 否 | | 53 | | `residual` | torch.Tensor | 同x | 残差输入，在归一化前与x相加 | 否 | |
| 54 | | `prenorm` | bool | - | 是否启用prenorm模式 | 否，默认False | | 54 | | `prenorm` | bool | - | 是否启用prenorm模式 | 否，默认False | |
| @@ -107,8 +107,8 @@ def rms_norm_gated( | |||
| 107 | |--------|------| | 107 | |--------|------| |
| 108 | | `x` 与 `g` | 必须相同dtype | | 108 | | `x` 与 `g` | 必须相同dtype | |
| 109 | | `weight` | 可与 `x` 不同dtype，kernel内部自动转换 | | 109 | | `weight` | 可与 `x` 不同dtype，kernel内部自动转换 | |
| 110 | -| `residual` | 可与 `x` 不同dtype，通过 `residual_in_fp32=True` 强制fp32存储 | | 110 | +| `residual` | 可与 `x` 不同dtype，其dtype由传入张量自身决定；`residual_in_fp32` 仅在未传 `residual` 时生效，决定此时 `residual_out` 是否以fp32存储 | |
| 111 | -| 输出 `y` | 与 `x` 相同dtype（除非指定 `out_dtype`） | | 111 | +| 输出 `y` | 恒与 `x` 相同dtype（公共接口 `rms_norm_gated` 不暴露 `out_dtype`） | |
| 112 | 112 | ||
| 113 | ## 实现原理 | 113 | ## 实现原理 |
| 114 | 114 | ||
| @@ -179,13 +179,13 @@ assert y.shape == (B, T, D) | |||
| 179 | 179 | ||
| 180 | ```python | 180 | ```python |
| 181 | residual = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16) | 181 | residual = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16) |
| 182 | -y = model(x, g, residual=residual) | 182 | +y = rms_norm_gated(x, g, weight, bias=None, activation='silu', residual=residual) |
| 183 | ``` | 183 | ``` |
| 184 | 184 | ||
| 185 | ### Prenorm模式 | 185 | ### Prenorm模式 |
| 186 | 186 | ||
| 187 | ```python | 187 | ```python |
| 188 | -y, residual_out = model(x, g, prenorm=True) | 188 | +y, residual_out = rms_norm_gated(x, g, weight, bias=None, activation='silu', prenorm=True) |
| 189 | ``` | 189 | ``` |
| 190 | 190 | ||
| 191 | ## 性能对比 | 191 | ## 性能对比 |
| @@ -215,7 +215,7 @@ y, residual_out = model(x, g, prenorm=True) | |||
| 215 | 1. `x` 与 `g` 必须具有完全相同的形状。 | 215 | 1. `x` 与 `g` 必须具有完全相同的形状。 |
| 216 | 2. `weight` 的长度必须等于x的最后一维D。 | 216 | 2. `weight` 的长度必须等于x的最后一维D。 |
| 217 | 3. 激活函数名称 `swish` 和 `silu` 等价，均表示 `x * sigmoid(x)`。 | 217 | 3. 激活函数名称 `swish` 和 `silu` 等价，均表示 `x * sigmoid(x)`。 |
| 218 | -4. 当前实现中 `bias` 固定为None，如需偏置需扩展实现。 | 218 | +4. `bias` 为必传参数（无默认值），可传 `None` 表示不加偏置；传入 `[D]` 张量时按元素相加并产生 `db` 梯度。 |
| 219 | 5. 反向传播时kernel会重计算前向输出y，这是用计算换显存的设计，在训练显存敏感场景下属于预期行为。 | 219 | 5. 反向传播时kernel会重计算前向输出y，这是用计算换显存的设计，在训练显存敏感场景下属于预期行为。 |
| 220 | 6. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销可能超过计算收益。 | 220 | 6. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销可能超过计算收益。 |
| 221 | -7. 该算子不支持arch35平台。 | 221 | +7. 该算子同时支持arch32与arch35架构，运行时由 `is_arch35()` 自动分发至对应架构的实现，`bias`、`residual`、`prenorm` 在两种架构上均支持；arch35路径额外要求 `D` 不超过 `65536 / x.element_size()`，超出会抛出 `RuntimeError`。 |
| @@ -89,7 +89,7 @@ class SinkhornFunction(torch.autograd.Function): | |||
| 89 | eps: float = 1e-6, | 89 | eps: float = 1e-6, |
| 90 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | 90 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 91 | """ | 91 | """ |
| 92 | - Triton implementation of HC-Split Sinkhorn, optimized for GPU performance | 92 | + Triton implementation of HC-Split Sinkhorn, optimized for Ascend NPU performance |
| 93 | 93 | ||
| 94 | Args: | 94 | Args: |
| 95 | mixes: Input tensor with shape [batch_size, seq_len, (2+hc_mult)*hc_mult] | 95 | mixes: Input tensor with shape [batch_size, seq_len, (2+hc_mult)*hc_mult] |
| @@ -111,11 +111,11 @@ class SinkhornFunction(torch.autograd.Function): | |||
| 111 | 111 | ||
| 112 | | 参数 | 类型 | 描述 | 支持的数据类型 | | 112 | | 参数 | 类型 | 描述 | 支持的数据类型 | |
| 113 | |------|------|------|---------------| | 113 | |------|------|------|---------------| |
| 114 | -| mixes | torch.Tensor | [batch_size, seq_len, (2+hc_mult)*hc_mult] | float32 | | 114 | +| mixes | torch.Tensor | [batch_size, seq_len, (2+hc_mult)*hc_mult] | float32, float16, bfloat16 | |
| 115 | | hc_scale | torch.Tensor | [3] | float32 | | 115 | | hc_scale | torch.Tensor | [3] | float32 | |
| 116 | | hc_base | torch.Tensor | [(2+hc_mult)*hc_mult] | - | | 116 | | hc_base | torch.Tensor | [(2+hc_mult)*hc_mult] | - | |
| 117 | | hc_mult | attr | 立即数，当前仅支持4 | int32 | | 117 | | hc_mult | attr | 立即数，当前仅支持4 | int32 | |
| 118 | -| sinkhorn_iters | attr | 立即数，当前仅支持20 | - | | 118 | +| sinkhorn_iters | attr | 立即数，默认为20，无取值校验（编译期常量） | - | |
| 119 | | eps | attr | 立即数，默认为1e-6 | - | | 119 | | eps | attr | 立即数，默认为1e-6 | - | |
| 120 | 120 | ||
| 121 | ## 输出说明 | 121 | ## 输出说明 |
| @@ -141,12 +141,14 @@ import torch | |||
| 141 | from mindspeed_ops.api.triton.sinkhorn import SinkhornFunction | 141 | from mindspeed_ops.api.triton.sinkhorn import SinkhornFunction |
| 142 | 142 | ||
| 143 | # 创建测试张量 | 143 | # 创建测试张量 |
| 144 | +dtype = torch.float32 | ||
| 145 | +device = torch.device("npu") | ||
| 144 | mixes = torch.randn((1, 2048, 24), requires_grad=True, dtype=dtype).to(device) | 146 | mixes = torch.randn((1, 2048, 24), requires_grad=True, dtype=dtype).to(device) |
| 145 | hc_scale = torch.randn((3,), requires_grad=True, dtype=dtype).to(device) | 147 | hc_scale = torch.randn((3,), requires_grad=True, dtype=dtype).to(device) |
| 146 | hc_base = torch.randn((24,), requires_grad=True, dtype=dtype).to(device) | 148 | hc_base = torch.randn((24,), requires_grad=True, dtype=dtype).to(device) |
| 147 | 149 | ||
| 148 | # 执行Sinkhorn算子 | 150 | # 执行Sinkhorn算子 |
| 149 | -triton_pre, triton_post, triton_comb = SinkhornFunction.apply(triton_mixes, triton_hc_scale, triton_hc_base) | 151 | +triton_pre, triton_post, triton_comb = SinkhornFunction.apply(mixes, hc_scale, hc_base) |
| 150 | 152 | ||
| 151 | # 验证结果 | 153 | # 验证结果 |
| 152 | print("sinkhorn operation executed successfully!") | 154 | print("sinkhorn operation executed successfully!") |
| @@ -41,7 +41,7 @@ def moe_sort_chunks_by_index_with_probs( | |||
| 41 | | 参数 | 类型 | Shape | 支持dtype | 说明 | | 41 | | 参数 | 类型 | Shape | 支持dtype | 说明 | |
| 42 | |------|------|-------|------------|------| | 42 | |------|------|-------|------------|------| |
| 43 | | `inp` | `torch.Tensor` | `[num_tokens, hidden_size]` | `float16`, `bfloat16`, `float32` | 待重排的token表示 | | 43 | | `inp` | `torch.Tensor` | `[num_tokens, hidden_size]` | `float16`, `bfloat16`, `float32` | 待重排的token表示 | |
| 44 | -| `probs` | `torch.Tensor` | `[num_tokens]` | `float16`, `bfloat16`, `float32` | 可选输入，随 `inp` 一起重排 | | 44 | +| `probs` | `torch.Tensor` | `[num_tokens]` | `float16`, `bfloat16`, `float32` | 仅 `moe_sort_chunks_by_index_with_probs` 接口使用，随 `inp` 一起重排 | |
| 45 | | `split_sizes` | `torch.Tensor` | `[num_splits]` | `int32` | 每个chunk在第0维的长度，元素和应等于 `num_tokens` | | 45 | | `split_sizes` | `torch.Tensor` | `[num_splits]` | `int32` | 每个chunk在第0维的长度，元素和应等于 `num_tokens` | |
| 46 | | `sorted_index` | `torch.Tensor` | `[num_splits]` | `int32` | chunk重排索引，应为 `[0, num_splits)` 的合法排列 | | 46 | | `sorted_index` | `torch.Tensor` | `[num_splits]` | `int32` | chunk重排索引，应为 `[0, num_splits)` 的合法排列 | |
| 47 | 47 | ||