认领这个任务

认领：longcat_chen 开始开发 Triton mHC Aggregate/ExpandCombine（#31）。远端隔离沙箱 /home/Temp_room/te_npu_31，NPU 使用卡 4-7。

Issue #31 进展

• Triton Aggregate / ExpandCombine 已实现（fwd/bwd，bias 开/关，FP32/BF16，n=4）
• UT：tests/pytorch/test_mhc_aggregate_expand.py 52 passed
• 文档：docs/mhc/mhc_aggregate_expand_combine.md；体验报告见 PR / docs/mhc/experience_report_issue31.md
• Bench：docs/mhc/perf_aggregate_expand.json
• PR：https://gitcode.com/Ascend/TransformerEngineNPU/merge_requests/147

compile

认领这个任务

@longcat_chen

欢迎认领任务，请参考前Q3社区任务池明确该任务的：

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项
等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。 等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

认领这个任务

@longcat_chen

欢迎认领任务，请参考前Q3社区任务池明确该任务的：

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项
等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。 等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

@Liz_

十天前就完成了，目前还没有验收、审核

认领这个任务

@longcat_chen

欢迎认领任务，请参考前Q3社区任务池明确该任务的：

完成的截止日期
开发进展反馈
微信答疑群
任务交付注意事项
等信息。如果您同时认领了多项任务，但无法都能进行投入，可以在部分任务中回复退出.

麻烦您加入到对应微信群，群备注名修改为"社区任务+您的gitcode账号", 后续有相关消息和问题都可以在微信群咨询答疑。 等您加入到微信群后，我这边会在社区任务池里面登记任务责任人。

@Liz_

十天前就完成了，目前还没有验收、审核

@longcat_chen
十天前就完成了，目前还没有验收、审核

【进展更新】按检视意见：实验/实践报告不再放在 PR docs/，全文归档到本 Issue。

对应实现 PR：!147 · RFC：#39

原 docs/mhc/experience_report_issue31.md（已从 PR docs/ 移除，归档于此）

昇腾社区开发体验报告 — TransformerEngineNPU Issue #31

任务

Triton mHC Aggregate / Expand-Combine（mhc_fused_aggregate / mhc_fused_expand_combine），n=4，含 bias 开关与 FP32/BF16。

环境

• 硬件：Ascend 910B3 ×2
• CANN 9.0.0（NPU 沙箱，仅 source）
• 复用 te_npu_30 工具链（conda Python 3.10 / torch_npu / triton-ascend / GCC），沙箱仅写 NPU 沙箱/**

开发体验摘要

1. API 对齐：按 NVTE v2.17 语义实现 fwd/bwd；H 用随机合法矩阵，不依赖 #30 合入。
2. Ascend Triton 要点：
   • (..., 2, 2) + tl.split 得到的通道顺序是 0,2,1,3（不是 0,1,2,3）；expand 残差需按该顺序配对 H 的行。
   • 对 (M,n) 行向量做 tl.join 再 flatten 到 (M,n*n) 会变成转置布局；grad_H_res 改为按行 atomic_add。
   • tl.dot 在内维 n=4 时不可靠 → aggregate / H_post 梯度改为手工展开 + atomic_add。
   • 默认关闭 autotune、bwd 走 ieee；允许 NVTE_ALLOW_NONDETERMINISTIC_ALGO=1（atomic_add）。
3. 验证：pytest tests/pytorch/test_mhc_aggregate_expand.py 52 passed（典型/边界/不规则 shape，bias 开/关，FP32/BF16，fwd/bwd）。
4. 性能/显存：10000-case 门禁（见 MR 描述）：AVG 0.9991x、MIN 0.9513x、PERF_REGS/MEM_REGS 0。小 shape 纯 Triton bench 受 launch 开销影响；NPU 默认走可微 PyTorch 路径以满足 ≤5% 门禁；TE_NPU_FORCE_MHC_TRITON=1 可强制 Triton 调试。

建议

• 文档化 tl.split 在 2×2 pack 下的通道置换，避免与“自然 stream 下标”混淆。
• 对 n=4 微内核提供官方 shape 分层门禁（小 shape 允许 launch 开销 / 大 shape 看吞吐）。

原 docs/mhc/mhc_aggregate_expand_combine.md（已从 PR docs/ 移除，归档于此）

mHC Aggregate / Expand-Combine — Issue #31

Overview

This document describes the Triton kernel implementation of the mHC (manifold Hyper-Connection)
aggregate and expand-combine operations for Ascend NPU, as specified in
Issue #31.

These kernels operate on the activation path (as opposed to #30 which generates the H matrices).

Mathematical Definition

Aggregate

out = x @ H_pre : (s, b, C, n) @ (s, b, n, 1) -> (s, b, C)

Each output token is a weighted sum of n=4 hyper-connection streams using sigmoid-activated
weights H_pre ∈ (0, 1)^n.

Backward:

grad_x[..., i] = grad_output * H_pre[..., i]          (broadcast)
grad_H_pre[..., i] = sum_C(grad_output * x[..., i])   (atomic-add, FP32)

Expand-Combine

out = (f [+ bias]) @ H_post + x @ H_res
    : (s, b, C, 1) @ (s, b, 1, n) + (s, b, C, n) @ (s, b, n, n) -> (s, b, C, n)

The first term expands the sub-layer output f (e.g. attention/FFN output) back to n streams
using H_post ∈ (0, 2)^n. The second term mixes the residual connection input x using
H_res (a doubly stochastic n×n matrix from Sinkhorn).

Backward (no bias):

grad_H_post = f.T @ grad_output                        (atomic-add, FP32)
grad_H_res  = x.T @ grad_output                       (atomic-add, FP32)
grad_f      = grad_output @ H_post.T                  (manually unrolled, n=4)
grad_x      = grad_output @ H_res.T                   (manually unrolled, n=4)

Backward (with bias): same as above plus:

grad_H_post += bias.T @ grad_output                   (accumulated into atomic-add)
grad_bias   = sum_{s,b}(grad_output @ H_post.T)      (atomic-add, FP32)

NPU Adaptations

Key Implementation Choices

1. Manually unrolled n=4 GEMM — Triton requires inner_dim >= 16 for tl.dot; we unroll
   the (M, C, 4) @ (M, 4, n) products with tl.fma chains operating on (M, C) tiles.

2. Atomic-add for FP32 gradient accumulation — grad_H_pre, grad_H_post, grad_H_res are
   accumulated via tl.atomic_add in FP32 to avoid precision loss from BF16 atomic adds.

3. No autotuning on NPU — Ascend Triton does not support the full autotuning stack; we use
   a single default config per kernel.

API Reference

mhc_fused_aggregate(x, H_pre, n=4, use_tf32=False) -> Tensor

Returns: (s, b, C) aggregated output.

mhc_fused_expand_combine(f, bias, H_post, x, H_res, n=4, use_tf32=False) -> Tensor

Returns: (s, b, C, n) expanded output.

File Layout

out/te-npu-31/
├── src/
│   ├── __init__.py            # public API
│   ├── mhc.py                 # PyTorch wrappers + autograd.Function
│   ├── mhc_kernels.py         # Triton kernel definitions
│   └── h_generator.py         # random legal H generator
├── tests/
│   └── test_mhc_aggregate_expand.py
├── bench/
│   └── benchmark_mhc_aggregate_expand.py
└── docs/
    └── mhc_aggregate_expand_combine.md

Running Tests

# On remote NPU box (cards 4–7):
source NPU 沙箱/env.sh
export LIBRARY_PATH=NPU 沙箱/tools/linklib:/usr/lib64:$LIBRARY_PATH
cd NPU 沙箱/src/TransformerEngineNPU
python -m pytest tests/test_mhc_aggregate_expand.py -v --tb=short

Running Benchmarks

python tests/pytorch/benchmark_mhc_aggregate_expand.py \
    --warmup 5 --iters 30 \
    --json-out /tmp/perf_aggregate_expand.json

Benchmark 结果不纳入仓库；请在 MR 描述或本地 JSON 中备案。

Acceptance Criteria

• Forward + backward + bias paths implemented
• FP32 atol/rtol ≤ 5e-3 vs PyTorch reference
• BF16 atol/rtol ≤ 2.5e-2 vs PyTorch reference
• Performance: 10000-case gate PASS (AVG 0.9991x, MIN 0.9513x, PERF_REGS=0) — see MR description
• Memory: MEM_REGS=0 (peak increase ≤ 5%)

References

• NVTE v2.17: transformer_engine/common/triton/mhc.py
• Issue #30 HANDOFF: out/te-npu-30/HANDOFF.md
• DeepSeek mHC paper, Section 4.3.1