ccf-ai-infra
/
TileOPs-Metax
0
0
Fork 36
Code Issues Pull Requests 12 Actions 4 Packages Projects Releases Wiki Activity

[moe_reduce_fused] feat: 实现并优化 C500 专家输出加权归约算子 #55

Open
ancient_rain wants to merge 17 commits from ancient_rain/TileOPs-Metax:feat/moe_reduce_fused into summer-camp-2026
pull from: ancient_rain/TileOPs-Metax:feat/moe_reduce_fused
merge into: ccf-ai-infra:summer-camp-2026
Conversation 0 Commits 17 Files Changed 14 +2412 -8
ancient_rain commented 2026-08-06 09:37:49 +08:00
Contributor
Copy Link

[moe_reduce_fused] feat: 实现并优化 C500 专家输出加权归约算子

小组课题信息

课题名称：面向 MetaX C500 的 moe_reduce_fused 算子迁移与性能优化

课题完整简介：本组以 MetaX C500 为目标平台，将 TileKernels-Metax 中的 MoE 路由/融合算子迁移至 TileOPs，并建立「Manifest → Op/Kernel → Test → Benchmark」完整信任链。moe_reduce_fused 作为主算子，负责将专家输出按路由权重加权归约回 token 序；本 PR 完成其 TileLang Kernel 实现、4 个 Manifest 变体 Op（Base / WithXsf / Fp8 / Quantized）、29 个正确性测试与 Manifest-driven Benchmark，并在真实 C500（sGPU 切片）上完成精度验证、性能复现与 mcProfiler 瓶颈分析。

PR 简要描述

算子名称：moe_reduce_fused

算子认领 Issue：#9

小组：9 组

成员：朱洪彪、古雨、陈占炜、李明洋

改动类型：

  • feat：新增算子或功能
  • optimize：优化已有实现

源文件参考（TileKernels-Metax dev 分支，源提交 SHA：0266ab740980de7dc03a828b8259cd73d100c2eb）：

  • 源 Kernel：tile_kernels/moe/reduce_fused_kernel.py
  • 源 PyTorch 参考实现：tile_kernels/torch/reduce_fused.py
  • 源测试：tests/moe/test_reduce_fused.py

被测提交 SHA：8568aeb99d23347d2d1152a64b938d615c96ff6f（评测机 /root/review-7770740ac4 工作树，其 tileops/kernels/moe/reduce_fused.py 与该提交内容一致，文件 sha1 d67d5bf7bdcb67ec06118dc5f146da5af0a31783；正确性截图、profile_run_final.log 与 mcProfiler 结果均在该状态下采集）。

当前分支 head 为 ff401ee，相对被测提交仅含三类非 kernel 变更：_infer_output_shapes 适配 manifest --strict 校验；官方矩阵比较对 FP8 输出先转 fp32；官方矩阵按 Review 意见折叠进 workloads/moe.py（WorkloadBase/FixtureBase）并删除独立的 workloads/moe_reduce_fused.py，bench_moe_reduce_fused.py 的 PyTorch 向量化基线扩展至全部 4 个变体。reduce_fused.py 内容未变，上述正确性/性能数据仍然成立。

1. 本次 PR 优化方案

优化前的主要问题：

  • PyTorch 参考实现通过 gather 先生成 [T, K, H] 中间张量再加权求和，访存量被放大 K 倍，且需要 clamp/掩码处理 padding slot；
  • 在 C500 上 PyTorch 向量化基线延迟高：512×8×7168 bf16 约 0.778 ms，4096×8×7168 bf16 约 5.63 ms（同一 ManifestBenchmark 协议测量）；
  • 原始 TileLang kernel（one-block-per-token）每个 token 只启动一个 CTA、单 CTA 持有完整 H 维 FP32 accumulator，对 T=512 decode 并行覆盖不足，FP8/Quant 路径延迟高（TileLang 基线实测：FP8 0.1014 ms、Quant 0.1009 ms、base decode 0.0526 ms）；
  • MACA/TileLang 对 register-only gather 循环的 T.Pipelined 会降级为串行循环，直接流水化无法生效。

采用的优化方案（对应分支 feat/moe_reduce_fused 最新实现）：

  • 单 Kernel 融合 + H 维分块：一个 block 处理 (token, hidden-tile)，按 token_topk_to_pos 索引 x[pos, :] 加权累加，消除 [T, K, H] gather 中间张量；block_h 按 workload 特调（256/512/1024/3584/7168），线程数 min(128, block_h) 向下取 2 的幂，提升 decode 小 T 场景的并行度与占用；
  • 4 级前瞻软件流水线（K≥4）：由于 MACA 对 T.Pipelined 的串行降级，显式展开四 buffer look-ahead pipeline，把后续 slot 的 load 与当前 slot 的 FMA 重叠（K 累加顺序保持不变），gather 的访存延迟被隐藏；K<4 走 fragment 元数据路径；
  • fp32 累加器：K 个 slot 的贡献在 fp32 中累加，最后一次性 cast 到输出 dtype，保证精度；
  • padding 跳过：pos < 0 的 slot 直接跳过（流水线路径对 padding stage 清空 buffer），无无效访存/计算；
  • 变体内联：with_weights/with_sf/with_x_sf 编译期开关内联进同一个 kernel 模板，4 个 Manifest 变体共享一份 kernel 逻辑；按 (T, K, H, in/out dtype, 变体) 特化，每个配置只 JIT 编译一次；
  • 规避 MACA 构建问题：输出直接写全局内存（不使用 out_frag + T.copy 中转，该组合在 MACA/TileLang 构建上会崩溃）；
  • 保留基线对照：新增 tileops/kernels/moe/reduce_fused_base.py（原 one-block-per-token 版本）供回归对照。

关键代码或配置变更：

  • 新增 tileops/kernels/moe/reduce_fused.py（MoeReduceFusedKernel，supported_archs=[80, 86, 89, 90]）；
  • 新增 tileops/kernels/moe/reduce_fused_base.py（基线 kernel 对照版本）；
  • 新增 tileops/ops/moe/routed_expert/reduce_fused.py（MoeReduceFusedFwdOp / MoeReduceFusedWithXsfFwdOp / MoeReduceFusedFp8FwdOp / MoeReduceFusedQuantizedFwdOp）；
  • 注册导出：tileops/kernels/moe/__init__.py、tileops/ops/moe/__init__.py、tileops/ops/moe/routed_expert/__init__.py；
  • 新增 tests/ops/test_moe_reduce_fused.py、benchmarks/ops/bench_moe_reduce_fused.py；官方矩阵按 Review 意见以 MoeReduceFusedOfficialWorkload / MoeReduceFusedOfficialFixture / MoeReduceFusedOfficialBenchmarkFixture 折叠进 workloads/moe.py，对应测试/基准为 tests/ops/test_moe_reduce_fused_official.py、benchmarks/ops/bench_moe_reduce_fused_official.py；
  • tileops/manifest/moe.yaml：MoeReduceFusedFwdOp 及 3 个变体由 spec-only 更新为 implemented；
  • 新增 HANDOFF.md（协作与复现指引）。

2. 精度验证

对比基准实现：独立 PyTorch 参考实现 _ref_reduce_fused（与源仓库 tile_kernels/torch/reduce_fused.py 语义一致：fp32 累加、-1 padding 跳过、可选 x_sf/sf 缩放、输出 cast）；官方 TileKernels 矩阵测试见 tests/ops/test_moe_reduce_fused_official.py（MoeReduceFusedOfficialWorkload 数据生成 + MoeReduceFusedOfficialTest.ref_program 参考 + _assert_equal 按 dtype 容差比较，FP8 先转 fp32）。

测试命令：

python -m pytest -q -m smoke tests/ops/test_moe_reduce_fused.py -v
python -m pytest -q tests/ops/test_moe_reduce_fused.py -v

测试 shape/dtype（共 29 项 = 14 smoke + 15 full）：

  • Base 变体（10 项）：tiny 4×2×256（fp16 / bf16 / fp32）、16×2×512 bf16、32×8×3072 bf16（Qwen3-30B）、32×8×7168 bf16、512×8×7168 bf16（DeepSeek-V3 decode）、1×8×256 bf16（单 token）、8×1×256 bf16（K=1）、无权重路径；
  • 变体（WithXsf / Fp8 / Quantized 各 5 项）：tiny fp16、tiny bf16、32×8×7168 bf16、512×8×7168 bf16（Qwen3-235B）、512×8×3072 bf16（Qwen3-30B）；
  • 错误用例（4 项）：hidden % 256 ≠ 0、topk_weights dtype 错误、shape 不匹配、-1 padding 跳过。

误差范围（atol/rtol）：fp16/fp32 1e-3；bf16 1.6e-2；FP8 输出 0.1（e4m3fn 表示精度限制）。

验证结果：smoke 14 passed；全量 29 passed（MetaX C500 实测，约 137 s，正确性截图与 profile_run_final.log 已归档，复现方式见 HANDOFF.md）。

3. 性能数据【必填】

测试环境：

  • GPU：MetaX C500（sGPU 切片：25% Compute / 16000 MiB Vram Quota，整卡 65536 MiB）
  • Kernel Mode Driver：3.8.30（torch 侧报告 N/A）
  • MACA：3.7.1.5（/opt/maca-3.7.1）
  • PyTorch：2.8.0+metax3.7.1.3（CUDA 11.6）
  • TileLang：0.1.10+cuda.gitf549117c（determine_target('auto') → maca）
  • mcProfiler：3.8.1.4（per-kernel）
  • 被测提交 SHA：8568aeb99d23347d2d1152a64b938d615c96ff6f（评测机 /root/review-7770740ac4 工作树，kernel 内容与上游一致）

性能复现命令（评测机代码目录 /root/review-7770740ac4）：

cd /root/review-7770740ac4
export PYTHONPATH=/opt/tilelang-metax-v0.1.10:$PWD
export MACA_HOME=/opt/maca-3.7.1
export MACA_PATH=/opt/maca
python -m pytest benchmarks/ops/bench_moe_reduce_fused.py -v
# Benchmark report saved to profile_run.log

注意：非交互 shell 需要显式导出 MACA_HOME / MACA_PATH，否则 Triton metax 后端（benchmark harness 依赖）会报找不到 /opt/maca-3.7.1/include。

测量协议：ManifestBenchmark（10 次 warmup + 50 repeats × 3 trials + CUDA event + L2 flush，报告取 trials 中位均值），报告文件 profile_run_final.log（2026-08-05 14:31:10，与评测机 /root/review-7770740ac4/profile_run.log 一致）。2026-08-06 01:27 补充一轮同会话三口径实测（18/18 passed）：TileLang 基线（reduce_fused_base.py）、当前 kernel、PyTorch 向量化基线三者同一协议、同一会话测得，见下文「三口径对照表」与评测机 three_way_bench.json。

优化前（TileLang 基线：原始 one-block-per-token TileLang kernel，即 tileops/kernels/moe/reduce_fused_base.py，与源仓库 TileLang kernel 语义一致；基线提交 a409c2d，同协议、同会话实测，独立 baseline full12）：

唯一 workload TileLang 基线 (ms)
base H7168 decode（512×8, bf16） 0.0526
base H7168 prefill（4096×8, bf16） 0.3566
base H3072 decode（512×8, bf16） 0.0285
base tiny（32×2, fp16） 0.0093
XSF H7168 0.0498
XSF tiny 0.0070
FP8 H7168（FP8 输出） 0.1014
FP8 tiny 0.0075
quantized H7168（FP8 输出） 0.1009
quantized tiny 0.0054

优化后（当前 H 维分块 kernel，同一协议、同一会话实测；测试提交 643f3bf / 本分支 8568aeb 系列）：

唯一 workload 当前 (ms) 变化
base H7168 decode 0.0513 约 -2.3%（按噪声处理）
base H7168 prefill 0.3565 无明确变化
base H3072 decode 0.0263 -7.72%（三轮确认 -9.02%）
base tiny 0.0093 无明确变化
XSF H7168 0.0489 -1.81%（按噪声处理）
XSF tiny 0.0069 绝对差为亚微秒
FP8 H7168 0.0617 -39.15%（三轮确认 -38.28%）
FP8 tiny 0.0070 -6.67%（绝对差很小）
quantized H7168 0.0633 -37.26%（三轮确认 -37.02%）
quantized tiny 0.0047 -12.96%（绝对差很小）

注：「优化后」列为 2026-08-05 14:31 独立三轮对比会话数据；后续跨会话复测（00:53、01:27）FP8/Quant H7168 为 0.0568–0.0608 / 0.0580–0.0588 ms，跨会话差异约 ±5% 噪声量级，TileLang 基线口径的加速结论不变（权威对照见下方「三口径对照表」）。

加速比（TileLang 基线口径）：FP8 H7168 1.64×、Quant H7168 1.59×、H3072 1.08×、base decode 1.03×（噪声级）、base prefill/tiny 无明确变化（已近带宽上限）。

说明：

  • Base/Xsf 路径两者均已接近带宽上限（基线 1.26 TB/s → 当前 1.29 TB/s，实测 BF16 copy 峰值约 1.40 TB/s），因此收益主要集中在 FP8/Quant 路径（-37%～-39%，1.6× 左右），来自 H 维并行覆盖与单 CTA accumulator 压力缓解（mcProfiler：FP8 waves 1024 → 14336，global read 仅 +0.024%）；
  • 作为外部参照：PyTorch 向量化参考实现（同协议，2026-08-06 01:27 同会话实测）H7168 decode：base 0.779 ms（约 14.6×）、XSF 0.791 ms（约 15.6×）、FP8 0.831 ms（约 14.4×）、Quant 0.843 ms（约 14.3×）；详见下方三口径对照表；
  • 协议：显式 warmup + 内部 warmup 10 + 50 repeats × 3 trials + L2 flush + CUPTI kernel-only，取三轮 trial mean 的 median；编译不计时。

三口径对照表（同会话 2026-08-06 01:27 实测，18/18 passed；three_way_bench.json）：

变体 唯一 workload TileLang 基线 (ms) 当前 kernel (ms) PyTorch 基线 (ms) kernel/baseline kernel/PyTorch baseline/PyTorch
Base 512×8 H=7168 bf16（DeepSeek-V3 decode） 0.0521 0.0535 0.7787 0.98×（噪声） 14.6× 14.9×
Base 512×8 H=7168 bf16 prefill（4096 tokens） 0.3564 0.3629 5.6304 0.98×（噪声） 15.5× 15.8×
Base 512×8 H=7168 bf16（Qwen3-235B decode） 0.0520 0.0534 0.7775 0.97×（噪声） 14.6× 14.9×
Base 512×8 H=3072 bf16（Qwen3-30B decode） 0.0281 0.0272 0.3880 1.03× 14.2× 13.8×
Base 512×8 H=7168 bf16（Kimi K2 decode） 0.0521 0.0534 0.7785 0.98×（噪声） 14.6× 14.9×
Base 32×2 H=256 fp16（tiny） 0.0088 0.0088 0.0636 1.00× 7.2× 7.2×
WithXsf 512×8 H=7168 bf16（DeepSeek-V3 decode） 0.0500 0.0508 0.7912 0.98×（噪声） 15.6× 15.8×
WithXsf 512×8 H=7168 bf16（Qwen3-235B decode） 0.0501 0.0510 0.7909 0.98×（噪声） 15.5× 15.8×
WithXsf 512×8 H=3072 bf16（Qwen3-30B decode） 0.0257 0.0259 0.4004 0.99×（噪声） 15.4× 15.6×
WithXsf 32×2 H=256 fp16（tiny） 0.0068 0.0068 0.0740 1.01× 10.9× 10.8×
Fp8（FP8 输出） 512×8 H=7168 bf16（DeepSeek-V3 decode） 0.1019 0.0579 0.8312 1.76× 14.4× 8.2×
Fp8（FP8 输出） 512×8 H=7168 bf16（Qwen3-235B decode） 0.1071 0.0608 0.8321 1.76× 13.7× 7.8×
Fp8（FP8 输出） 512×8 H=3072 bf16（Qwen3-30B decode） 0.0450 0.0421 0.4146 1.07× 9.8× 9.2×
Fp8（FP8 输出） 32×2 H=256 fp16（tiny） 0.0075 0.0070 0.0745 1.08× 10.7× 9.9×
Quantized（FP8 输出） 512×8 H=7168 bf16（DeepSeek-V3 decode） 0.1002 0.0588 0.8430 1.70× 14.3× 8.4×
Quantized（FP8 输出） 512×8 H=7168 bf16（Qwen3-235B decode） 0.0993 0.0580 0.8436 1.71× 14.6× 8.5×
Quantized（FP8 输出） 512×8 H=3072 bf16（Qwen3-30B decode） 0.0442 0.0411 0.4270 1.08× 10.4× 9.7×
Quantized（FP8 输出） 32×2 H=256 fp16（tiny） 0.0052 0.0047 0.0848 1.13× 18.2× 16.2×

三口径同会话结论：base/xsf 已近带宽上限，kernel 相对 TileLang 基线为噪声级（0.97–1.03×）；FP8/Quant H7168 相对基线 1.70–1.76×；全部变体相对 PyTorch 向量化基线 7.2×–18.2×（H7168 decode 13.7×–15.6×）。

Roofline（Manifest 公式，基础变体）：FLOPs = 2*T*K*H；Bytes = (T*K*H + T*H)*elem_bytes + T*K*4 + T*K*4。理论读写字节量 512×8×7168 bf16 为 66.09 MB（58.72 MB 读 + 7.34 MB 写 + 32.8 KB 元数据），benchmark 记录带宽 ≈ 1.25 TB/s（prefill 1.46 TB/s）。

瓶颈分析（mcProfiler 3.8.1.4 per-kernel 观测，_reduce_fused_main_kernel）：

工作负载 Global Read Global Write Memory Access/s AP busy Duty L2C Hit Rate
base H7168 decode（DeepSeek-V3） 58.78 MB 7.34 MB 1.064 TB/s 1.39% 2.06%
base H7168 decode（Qwen3-235B） 58.77 MB 7.34 MB 1.113 TB/s 1.06% 2.77%
base H7168 decode（Kimi K2） 58.77 MB 7.34 MB 1.111 TB/s 1.33% 2.77%
base H7168 prefill（DeepSeek-V3） 470.05 MB 58.72 MB 1.162 TB/s 8.93% 0.47%
base H3072 decode（Qwen3-30B） 25.20 MB 3.15 MB 1.086 TB/s 0.61% 2.79%
base H256 tiny（fp16） 35.3 KB 16.6 KB ~17 GB/s 0.07% 59.0%
fp8 H7168 decode（DeepSeek-V3） 58.82 MB 3.67 MB 0.877 TB/s 1.59% 2.34%
fp8 H7168 decode（Qwen3-235B） 58.82 MB 3.67 MB 0.876 TB/s 1.60% 2.35%
fp8 H3072 decode（Qwen3-30B） 25.23 MB 1.57 MB 0.551 TB/s 1.04% 12.18%
fp8 H256 tiny（fp16） 35.5 KB 8.4 KB ~12 GB/s 0.08% 84.1%
quant H7168 decode（DeepSeek-V3） 58.87 MB 3.67 MB 0.840 TB/s 1.62% 2.94%
quant H7168 decode（Qwen3-235B） 58.86 MB 3.67 MB 0.849 TB/s 1.63% 2.95%
quant H3072 decode（Qwen3-30B） 25.25 MB 1.57 MB 0.508 TB/s 1.20% 12.49%
quant H256 tiny（fp16） 35.7 KB 8.4 KB ~10 GB/s 0.10% 78.5%
xsf H7168 decode（DeepSeek-V3） 58.80 MB 7.34 MB 1.093 TB/s 1.12% 2.30%
xsf H7168 decode（Qwen3-235B） 58.80 MB 7.34 MB 1.088 TB/s 0.35% 2.30%
xsf H3072 decode（Qwen3-30B） 25.23 MB 3.15 MB 1.016 TB/s 0.64% 4.07%
xsf H256 tiny（fp16） 35.6 KB 16.6 KB ~16 GB/s 0.07% 75.5%

结论：

  • 全局读写字节量与理论一致（实测 66.12 MB vs 理论 66.09 MB，约 +24 KB 读放大），访存为纯流式（L2C Hit Rate ≈ 2%），无中间张量放大；
  • Memory Access per Second：base/xsf 约 1.02–1.16 TB/s；fp8/quant 约 0.51–0.88 TB/s（额外 cast/store 与更小的 FP8 输出写），与 benchmark 记录的 roofline 带宽（1.25–1.46 TB/s）同一量级，差距来自启动/同步开销；
  • AP busy Duty：decode 全部 ≤1.63%，prefill 8.93%，tiny <0.1%，ISU stall 以 vls_pipeline_stall 为主，计算单元几乎空闲；
  • 结论：所有变体均为强访存（memory-bound）型；fp8/quant 的带宽略低来自 FP8 cast/store 开销而非访存放大（读写字节与理论一致）；进一步优化方向是提升 HBM 带宽利用率（如继续增大 block_h/线程并行度、减少 padding 分支与 kernel 启动开销）。

4. 提交自检清单

  • 全部测试用例运行通过，精度达标
  • 已删除临时调试打印、冗余测试代码
  • 性能指标可本地复现，数据真实有效
# [moe_reduce_fused] feat: 实现并优化 C500 专家输出加权归约算子 ## 小组课题信息 课题名称：面向 MetaX C500 的 moe_reduce_fused 算子迁移与性能优化 课题完整简介：本组以 MetaX C500 为目标平台，将 TileKernels-Metax 中的 MoE 路由/融合算子迁移至 TileOPs，并建立「Manifest → Op/Kernel → Test → Benchmark」完整信任链。`moe_reduce_fused` 作为主算子，负责将专家输出按路由权重加权归约回 token 序；本 PR 完成其 TileLang Kernel 实现、4 个 Manifest 变体 Op（Base / WithXsf / Fp8 / Quantized）、29 个正确性测试与 Manifest-driven Benchmark，并在真实 C500（sGPU 切片）上完成精度验证、性能复现与 mcProfiler 瓶颈分析。 ### PR 简要描述 算子名称：`moe_reduce_fused` 算子认领 Issue：[#9](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/9) 小组：9 组 成员：朱洪彪、古雨、陈占炜、李明洋 改动类型： - [x] `feat`：新增算子或功能 - [x] `optimize`：优化已有实现 源文件参考（TileKernels-Metax `dev` 分支，源提交 SHA：`0266ab740980de7dc03a828b8259cd73d100c2eb`）： - 源 Kernel：`tile_kernels/moe/reduce_fused_kernel.py` - 源 PyTorch 参考实现：`tile_kernels/torch/reduce_fused.py` - 源测试：`tests/moe/test_reduce_fused.py` 被测提交 SHA：`8568aeb99d23347d2d1152a64b938d615c96ff6f`（评测机 `/root/review-7770740ac4` 工作树，其 `tileops/kernels/moe/reduce_fused.py` 与该提交内容一致，文件 sha1 `d67d5bf7bdcb67ec06118dc5f146da5af0a31783`；正确性截图、`profile_run_final.log` 与 mcProfiler 结果均在该状态下采集）。 当前分支 head 为 `ff401ee`，相对被测提交仅含三类非 kernel 变更：`_infer_output_shapes` 适配 manifest `--strict` 校验；官方矩阵比较对 FP8 输出先转 fp32；官方矩阵按 Review 意见折叠进 `workloads/moe.py`（WorkloadBase/FixtureBase）并删除独立的 `workloads/moe_reduce_fused.py`，`bench_moe_reduce_fused.py` 的 PyTorch 向量化基线扩展至全部 4 个变体。`reduce_fused.py` 内容未变，上述正确性/性能数据仍然成立。 ### 1. 本次 PR 优化方案 优化前的主要问题： - PyTorch 参考实现通过 `gather` 先生成 `[T, K, H]` 中间张量再加权求和，访存量被放大 K 倍，且需要 clamp/掩码处理 padding slot； - 在 C500 上 PyTorch 向量化基线延迟高：512×8×7168 bf16 约 0.778 ms，4096×8×7168 bf16 约 5.63 ms（同一 ManifestBenchmark 协议测量）； - 原始 TileLang kernel（one-block-per-token）每个 token 只启动一个 CTA、单 CTA 持有完整 H 维 FP32 accumulator，对 T=512 decode 并行覆盖不足，FP8/Quant 路径延迟高（TileLang 基线实测：FP8 0.1014 ms、Quant 0.1009 ms、base decode 0.0526 ms）； - MACA/TileLang 对 register-only gather 循环的 `T.Pipelined` 会降级为串行循环，直接流水化无法生效。 采用的优化方案（对应分支 `feat/moe_reduce_fused` 最新实现）： - **单 Kernel 融合 + H 维分块**：一个 block 处理 `(token, hidden-tile)`，按 `token_topk_to_pos` 索引 `x[pos, :]` 加权累加，消除 `[T, K, H]` gather 中间张量；`block_h` 按 workload 特调（256/512/1024/3584/7168），线程数 `min(128, block_h)` 向下取 2 的幂，提升 decode 小 T 场景的并行度与占用； - **4 级前瞻软件流水线**（K≥4）：由于 MACA 对 `T.Pipelined` 的串行降级，显式展开四 buffer look-ahead pipeline，把后续 slot 的 load 与当前 slot 的 FMA 重叠（K 累加顺序保持不变），gather 的访存延迟被隐藏；K<4 走 fragment 元数据路径； - **fp32 累加器**：K 个 slot 的贡献在 fp32 中累加，最后一次性 cast 到输出 dtype，保证精度； - **padding 跳过**：`pos < 0` 的 slot 直接跳过（流水线路径对 padding stage 清空 buffer），无无效访存/计算； - **变体内联**：`with_weights`/`with_sf`/`with_x_sf` 编译期开关内联进同一个 kernel 模板，4 个 Manifest 变体共享一份 kernel 逻辑；按 `(T, K, H, in/out dtype, 变体)` 特化，每个配置只 JIT 编译一次； - **规避 MACA 构建问题**：输出直接写全局内存（不使用 `out_frag` + `T.copy` 中转，该组合在 MACA/TileLang 构建上会崩溃）； - **保留基线对照**：新增 `tileops/kernels/moe/reduce_fused_base.py`（原 one-block-per-token 版本）供回归对照。 关键代码或配置变更： - 新增 `tileops/kernels/moe/reduce_fused.py`（`MoeReduceFusedKernel`，`supported_archs=[80, 86, 89, 90]`）； - 新增 `tileops/kernels/moe/reduce_fused_base.py`（基线 kernel 对照版本）； - 新增 `tileops/ops/moe/routed_expert/reduce_fused.py`（`MoeReduceFusedFwdOp` / `MoeReduceFusedWithXsfFwdOp` / `MoeReduceFusedFp8FwdOp` / `MoeReduceFusedQuantizedFwdOp`）； - 注册导出：`tileops/kernels/moe/__init__.py`、`tileops/ops/moe/__init__.py`、`tileops/ops/moe/routed_expert/__init__.py`； - 新增 `tests/ops/test_moe_reduce_fused.py`、`benchmarks/ops/bench_moe_reduce_fused.py`；官方矩阵按 Review 意见以 `MoeReduceFusedOfficialWorkload` / `MoeReduceFusedOfficialFixture` / `MoeReduceFusedOfficialBenchmarkFixture` 折叠进 `workloads/moe.py`，对应测试/基准为 `tests/ops/test_moe_reduce_fused_official.py`、`benchmarks/ops/bench_moe_reduce_fused_official.py`； - `tileops/manifest/moe.yaml`：`MoeReduceFusedFwdOp` 及 3 个变体由 `spec-only` 更新为 `implemented`； - 新增 `HANDOFF.md`（协作与复现指引）。 ### 2. 精度验证 对比基准实现：独立 PyTorch 参考实现 `_ref_reduce_fused`（与源仓库 `tile_kernels/torch/reduce_fused.py` 语义一致：fp32 累加、`-1` padding 跳过、可选 `x_sf`/`sf` 缩放、输出 cast）；官方 TileKernels 矩阵测试见 `tests/ops/test_moe_reduce_fused_official.py`（`MoeReduceFusedOfficialWorkload` 数据生成 + `MoeReduceFusedOfficialTest.ref_program` 参考 + `_assert_equal` 按 dtype 容差比较，FP8 先转 fp32）。 测试命令： ```bash python -m pytest -q -m smoke tests/ops/test_moe_reduce_fused.py -v python -m pytest -q tests/ops/test_moe_reduce_fused.py -v ``` 测试 shape/dtype（共 29 项 = 14 smoke + 15 full）： - Base 变体（10 项）：tiny `4×2×256`（fp16 / bf16 / fp32）、`16×2×512` bf16、`32×8×3072` bf16（Qwen3-30B）、`32×8×7168` bf16、`512×8×7168` bf16（DeepSeek-V3 decode）、`1×8×256` bf16（单 token）、`8×1×256` bf16（K=1）、无权重路径； - 变体（WithXsf / Fp8 / Quantized 各 5 项）：tiny fp16、tiny bf16、`32×8×7168` bf16、`512×8×7168` bf16（Qwen3-235B）、`512×8×3072` bf16（Qwen3-30B）； - 错误用例（4 项）：hidden % 256 ≠ 0、topk_weights dtype 错误、shape 不匹配、`-1` padding 跳过。 误差范围（atol/rtol）：fp16/fp32 `1e-3`；bf16 `1.6e-2`；FP8 输出 `0.1`（e4m3fn 表示精度限制）。 验证结果：smoke **14 passed**；全量 **29 passed**（MetaX C500 实测，约 137 s，正确性截图与 `profile_run_final.log` 已归档，复现方式见 HANDOFF.md）。 ### 3. 性能数据【必填】 测试环境： - GPU：MetaX C500（sGPU 切片：25% Compute / 16000 MiB Vram Quota，整卡 65536 MiB） - Kernel Mode Driver：3.8.30（torch 侧报告 N/A） - MACA：3.7.1.5（`/opt/maca-3.7.1`） - PyTorch：2.8.0+metax3.7.1.3（CUDA 11.6） - TileLang：0.1.10+cuda.gitf549117c（`determine_target('auto')` → `maca`） - mcProfiler：3.8.1.4（per-kernel） - 被测提交 SHA：`8568aeb99d23347d2d1152a64b938d615c96ff6f`（评测机 `/root/review-7770740ac4` 工作树，kernel 内容与上游一致） 性能复现命令（评测机代码目录 `/root/review-7770740ac4`）： ```bash cd /root/review-7770740ac4 export PYTHONPATH=/opt/tilelang-metax-v0.1.10:$PWD export MACA_HOME=/opt/maca-3.7.1 export MACA_PATH=/opt/maca python -m pytest benchmarks/ops/bench_moe_reduce_fused.py -v # Benchmark report saved to profile_run.log ``` > 注意：非交互 shell 需要显式导出 `MACA_HOME` / `MACA_PATH`，否则 Triton metax 后端（benchmark harness 依赖）会报找不到 `/opt/maca-3.7.1/include`。 测量协议：ManifestBenchmark（10 次 warmup + 50 repeats × 3 trials + CUDA event + L2 flush，报告取 trials 中位均值），报告文件 `profile_run_final.log`（2026-08-05 14:31:10，与评测机 `/root/review-7770740ac4/profile_run.log` 一致）。2026-08-06 01:27 补充一轮同会话三口径实测（18/18 passed）：TileLang 基线（`reduce_fused_base.py`）、当前 kernel、PyTorch 向量化基线三者同一协议、同一会话测得，见下文「三口径对照表」与评测机 `three_way_bench.json`。 优化前（**TileLang 基线**：原始 one-block-per-token TileLang kernel，即 `tileops/kernels/moe/reduce_fused_base.py`，与源仓库 TileLang kernel 语义一致；基线提交 `a409c2d`，同协议、同会话实测，独立 baseline full12）： | 唯一 workload | TileLang 基线 (ms) | | --- | --- | | base H7168 decode（512×8, bf16） | 0.0526 | | base H7168 prefill（4096×8, bf16） | 0.3566 | | base H3072 decode（512×8, bf16） | 0.0285 | | base tiny（32×2, fp16） | 0.0093 | | XSF H7168 | 0.0498 | | XSF tiny | 0.0070 | | FP8 H7168（FP8 输出） | 0.1014 | | FP8 tiny | 0.0075 | | quantized H7168（FP8 输出） | 0.1009 | | quantized tiny | 0.0054 | 优化后（当前 H 维分块 kernel，同一协议、同一会话实测；测试提交 `643f3bf` / 本分支 `8568aeb` 系列）： | 唯一 workload | 当前 (ms) | 变化 | | --- | --- | --- | | base H7168 decode | 0.0513 | 约 -2.3%（按噪声处理） | | base H7168 prefill | 0.3565 | 无明确变化 | | base H3072 decode | 0.0263 | **-7.72%**（三轮确认 -9.02%） | | base tiny | 0.0093 | 无明确变化 | | XSF H7168 | 0.0489 | -1.81%（按噪声处理） | | XSF tiny | 0.0069 | 绝对差为亚微秒 | | FP8 H7168 | 0.0617 | **-39.15%**（三轮确认 -38.28%） | | FP8 tiny | 0.0070 | -6.67%（绝对差很小） | | quantized H7168 | 0.0633 | **-37.26%**（三轮确认 -37.02%） | | quantized tiny | 0.0047 | -12.96%（绝对差很小） | > 注：「优化后」列为 2026-08-05 14:31 独立三轮对比会话数据；后续跨会话复测（00:53、01:27）FP8/Quant H7168 为 0.0568–0.0608 / 0.0580–0.0588 ms，跨会话差异约 ±5% 噪声量级，TileLang 基线口径的加速结论不变（权威对照见下方「三口径对照表」）。 加速比（**TileLang 基线口径**）：FP8 H7168 **1.64×**、Quant H7168 **1.59×**、H3072 **1.08×**、base decode 1.03×（噪声级）、base prefill/tiny 无明确变化（已近带宽上限）。 说明： - Base/Xsf 路径两者均已接近带宽上限（基线 1.26 TB/s → 当前 1.29 TB/s，实测 BF16 copy 峰值约 1.40 TB/s），因此收益主要集中在 **FP8/Quant 路径（-37%～-39%，1.6× 左右）**，来自 H 维并行覆盖与单 CTA accumulator 压力缓解（mcProfiler：FP8 waves 1024 → 14336，global read 仅 +0.024%）； - 作为外部参照：PyTorch 向量化参考实现（同协议，2026-08-06 01:27 同会话实测）H7168 decode：base 0.779 ms（**约 14.6×**）、XSF 0.791 ms（**约 15.6×**）、FP8 0.831 ms（**约 14.4×**）、Quant 0.843 ms（**约 14.3×**）；详见下方三口径对照表； - 协议：显式 warmup + 内部 warmup 10 + 50 repeats × 3 trials + L2 flush + CUPTI kernel-only，取三轮 trial mean 的 median；编译不计时。 三口径对照表（同会话 2026-08-06 01:27 实测，18/18 passed；`three_way_bench.json`）： | 变体 | 唯一 workload | TileLang 基线 (ms) | 当前 kernel (ms) | PyTorch 基线 (ms) | kernel/baseline | kernel/PyTorch | baseline/PyTorch | | --- | --- | --- | --- | --- | --- | --- | --- | | Base | 512×8 H=7168 bf16（DeepSeek-V3 decode） | 0.0521 | 0.0535 | 0.7787 | 0.98×（噪声） | 14.6× | 14.9× | | Base | 512×8 H=7168 bf16 prefill（4096 tokens） | 0.3564 | 0.3629 | 5.6304 | 0.98×（噪声） | 15.5× | 15.8× | | Base | 512×8 H=7168 bf16（Qwen3-235B decode） | 0.0520 | 0.0534 | 0.7775 | 0.97×（噪声） | 14.6× | 14.9× | | Base | 512×8 H=3072 bf16（Qwen3-30B decode） | 0.0281 | 0.0272 | 0.3880 | **1.03×** | 14.2× | 13.8× | | Base | 512×8 H=7168 bf16（Kimi K2 decode） | 0.0521 | 0.0534 | 0.7785 | 0.98×（噪声） | 14.6× | 14.9× | | Base | 32×2 H=256 fp16（tiny） | 0.0088 | 0.0088 | 0.0636 | 1.00× | 7.2× | 7.2× | | WithXsf | 512×8 H=7168 bf16（DeepSeek-V3 decode） | 0.0500 | 0.0508 | 0.7912 | 0.98×（噪声） | **15.6×** | 15.8× | | WithXsf | 512×8 H=7168 bf16（Qwen3-235B decode） | 0.0501 | 0.0510 | 0.7909 | 0.98×（噪声） | 15.5× | 15.8× | | WithXsf | 512×8 H=3072 bf16（Qwen3-30B decode） | 0.0257 | 0.0259 | 0.4004 | 0.99×（噪声） | 15.4× | 15.6× | | WithXsf | 32×2 H=256 fp16（tiny） | 0.0068 | 0.0068 | 0.0740 | 1.01× | 10.9× | 10.8× | | Fp8（FP8 输出） | 512×8 H=7168 bf16（DeepSeek-V3 decode） | 0.1019 | 0.0579 | 0.8312 | **1.76×** | 14.4× | 8.2× | | Fp8（FP8 输出） | 512×8 H=7168 bf16（Qwen3-235B decode） | 0.1071 | 0.0608 | 0.8321 | **1.76×** | 13.7× | 7.8× | | Fp8（FP8 输出） | 512×8 H=3072 bf16（Qwen3-30B decode） | 0.0450 | 0.0421 | 0.4146 | 1.07× | 9.8× | 9.2× | | Fp8（FP8 输出） | 32×2 H=256 fp16（tiny） | 0.0075 | 0.0070 | 0.0745 | 1.08× | 10.7× | 9.9× | | Quantized（FP8 输出） | 512×8 H=7168 bf16（DeepSeek-V3 decode） | 0.1002 | 0.0588 | 0.8430 | **1.70×** | 14.3× | 8.4× | | Quantized（FP8 输出） | 512×8 H=7168 bf16（Qwen3-235B decode） | 0.0993 | 0.0580 | 0.8436 | **1.71×** | 14.6× | 8.5× | | Quantized（FP8 输出） | 512×8 H=3072 bf16（Qwen3-30B decode） | 0.0442 | 0.0411 | 0.4270 | 1.08× | 10.4× | 9.7× | | Quantized（FP8 输出） | 32×2 H=256 fp16（tiny） | 0.0052 | 0.0047 | 0.0848 | 1.13× | 18.2× | 16.2× | > 三口径同会话结论：base/xsf 已近带宽上限，kernel 相对 TileLang 基线为噪声级（0.97–1.03×）；FP8/Quant H7168 相对基线 **1.70–1.76×**；全部变体相对 PyTorch 向量化基线 **7.2×–18.2×**（H7168 decode 13.7×–15.6×）。 Roofline（Manifest 公式，基础变体）：FLOPs = `2*T*K*H`；Bytes = `(T*K*H + T*H)*elem_bytes + T*K*4 + T*K*4`。理论读写字节量 512×8×7168 bf16 为 66.09 MB（58.72 MB 读 + 7.34 MB 写 + 32.8 KB 元数据），benchmark 记录带宽 ≈ 1.25 TB/s（prefill 1.46 TB/s）。 瓶颈分析（mcProfiler 3.8.1.4 per-kernel 观测，`_reduce_fused_main_kernel`）： | 工作负载 | Global Read | Global Write | Memory Access/s | AP busy Duty | L2C Hit Rate | | --- | --- | --- | --- | --- | --- | | base H7168 decode（DeepSeek-V3） | 58.78 MB | 7.34 MB | 1.064 TB/s | 1.39% | 2.06% | | base H7168 decode（Qwen3-235B） | 58.77 MB | 7.34 MB | 1.113 TB/s | 1.06% | 2.77% | | base H7168 decode（Kimi K2） | 58.77 MB | 7.34 MB | 1.111 TB/s | 1.33% | 2.77% | | base H7168 prefill（DeepSeek-V3） | 470.05 MB | 58.72 MB | 1.162 TB/s | 8.93% | 0.47% | | base H3072 decode（Qwen3-30B） | 25.20 MB | 3.15 MB | 1.086 TB/s | 0.61% | 2.79% | | base H256 tiny（fp16） | 35.3 KB | 16.6 KB | ~17 GB/s | 0.07% | 59.0% | | fp8 H7168 decode（DeepSeek-V3） | 58.82 MB | 3.67 MB | 0.877 TB/s | 1.59% | 2.34% | | fp8 H7168 decode（Qwen3-235B） | 58.82 MB | 3.67 MB | 0.876 TB/s | 1.60% | 2.35% | | fp8 H3072 decode（Qwen3-30B） | 25.23 MB | 1.57 MB | 0.551 TB/s | 1.04% | 12.18% | | fp8 H256 tiny（fp16） | 35.5 KB | 8.4 KB | ~12 GB/s | 0.08% | 84.1% | | quant H7168 decode（DeepSeek-V3） | 58.87 MB | 3.67 MB | 0.840 TB/s | 1.62% | 2.94% | | quant H7168 decode（Qwen3-235B） | 58.86 MB | 3.67 MB | 0.849 TB/s | 1.63% | 2.95% | | quant H3072 decode（Qwen3-30B） | 25.25 MB | 1.57 MB | 0.508 TB/s | 1.20% | 12.49% | | quant H256 tiny（fp16） | 35.7 KB | 8.4 KB | ~10 GB/s | 0.10% | 78.5% | | xsf H7168 decode（DeepSeek-V3） | 58.80 MB | 7.34 MB | 1.093 TB/s | 1.12% | 2.30% | | xsf H7168 decode（Qwen3-235B） | 58.80 MB | 7.34 MB | 1.088 TB/s | 0.35% | 2.30% | | xsf H3072 decode（Qwen3-30B） | 25.23 MB | 3.15 MB | 1.016 TB/s | 0.64% | 4.07% | | xsf H256 tiny（fp16） | 35.6 KB | 16.6 KB | ~16 GB/s | 0.07% | 75.5% | 结论： - 全局读写字节量与理论一致（实测 66.12 MB vs 理论 66.09 MB，约 +24 KB 读放大），访存为纯流式（L2C Hit Rate ≈ 2%），无中间张量放大； - Memory Access per Second：base/xsf 约 1.02–1.16 TB/s；fp8/quant 约 0.51–0.88 TB/s（额外 cast/store 与更小的 FP8 输出写），与 benchmark 记录的 roofline 带宽（1.25–1.46 TB/s）同一量级，差距来自启动/同步开销； - AP busy Duty：decode 全部 ≤1.63%，prefill 8.93%，tiny <0.1%，ISU stall 以 `vls_pipeline_stall` 为主，计算单元几乎空闲； - 结论：所有变体均为强访存（memory-bound）型；fp8/quant 的带宽略低来自 FP8 cast/store 开销而非访存放大（读写字节与理论一致）；进一步优化方向是提升 HBM 带宽利用率（如继续增大 block_h/线程并行度、减少 padding 分支与 kernel 启动开销）。 ### 4. 提交自检清单 - [x] 全部测试用例运行通过，精度达标 - [x] 已删除临时调试打印、冗余测试代码 - [x] 性能指标可本地复现，数据真实有效
ancient_rain added 17 commits 2026-08-06 09:37:50 +08:00
d2562e3537 test(moe): add official TileKernels matrix from codex branch 
Add the official TileKernels correctness/benchmark data from
codex/moe-reduce-fused-tests-benchmark so this branch can validate the
moe_reduce_fused implementation against the upstream matrix:

- workloads/moe_reduce_fused.py: official data generation, reference, adapter
- tests/ops/test_moe_reduce_fused_official.py: official correctness matrix
  (renamed to coexist with the branch's own tests)
- benchmarks/ops/bench_moe_reduce_fused_official.py: official benchmark matrix
  (renamed to coexist with the branch's own benchmark)

Adapter changes required by this implementation:
- import the four ops from tileops.ops.moe (routed_expert layout re-exports
  them there)
- map with_weights=False (topk_weights=None) to an all-ones float32 tensor,
  because the Manifest variants declare topk_weights as a required input

Note: tests not executed. The official assert_equal is bitwise and may need
tolerance for bf16/fp8 combos; matrix size is 1152 correctness + 144 bench
params by default.
71e36dd735 fix(moe): tolerance assert, env-gated matrix, and official-matrix test guide 
- workloads/moe_reduce_fused.py: replace bitwise torch.equal with
  dtype-appropriate tolerance comparison (fp32/fp16 1e-3, bf16 1.6e-2,
  fp8 0.1), since the reimplementation is not bit-identical to the
  upstream kernel (accumulation order / fma / fp8 rounding)
- tests/ops/test_moe_reduce_fused_official.py: default to a 48-case
  representative smoke subset; TILEOPS_OFFICIAL_FULL=1 runs the full
  1152-case matrix (TK_FULL_TEST=1 extends to 7489)
- benchmarks/ops/bench_moe_reduce_fused_official.py: default to a
  48-case subset; TILEOPS_OFFICIAL_FULL=1 runs all 144 cases
- docs/moe_reduce_fused_official_matrix_testing.md: environment setup,
  run commands, matrix sizes, and known risks
db9e68d40d Merge feat/moe_reduce_fused_zhb into feat/moe_reduce_fused_gy 
Sync the base branch: add with_weights=False support and its test to the
base variant (dbc355b), plus HANDOFF manifest-workload labels.
a409c2d18b fix(moe): drop torch.compile in official matrix data gen 
torch.compile decoration kills the process on the MetaX/MACA runtime
(SIGKILL at import/decorate time, before any kernel runs). The upstream
markers on generate_topk_idx and elementwise_fma are performance-only;
plain torch keeps identical semantics and lets the matrix run on C500.
2d92c86c15 fix(moe): correct roofline FLOPs for MoeReduceFused scaling variants 
The manifest undercounted FLOPs for the scaled variants: with_x_sf adds one
multiply per routed slot (T*K) and with_sf adds one multiply per output
element (T*H). Update the inline roofline formulas accordingly and clarify
that x_sf is a per-expanded-row (per-routed-slot) scaling factor, not a
per-token one.
0e811542b1 feat(moe): add qwen3 workloads to MoeReduceFused scaled variants 
Add qwen3-235b-decode (H=7168) and qwen3-30b-decode (H=3072) workloads to
the WithXsf, FP8-out, and full-quantized variants so their manifest-driven
benchmarks cover two additional real-model shapes.
7770740ac4 test(moe): cover qwen3-235b/qwen3-30b shapes in reduce-fused variant tests 
Align the WithXsf / FP8 / Quantized op tests with the new manifest
workloads: add qwen3-235b-decode (512x8x7168 bf16) and qwen3-30b-decode
(512x8x3072 bf16) to the shared variant fixture.
auto-label / apply-label (pull_request) Has been cancelled Details
ff401eed69 test(moe): fold official matrix into workloads/moe.py per review 
Sync the workload-generation refactor from feat/moe_reduce_fused_czw_dev
(bb5e5fe) only:
- move the official TileKernels matrix into workloads/moe.py using
  WorkloadBase/FixtureBase (MoeReduceFusedOfficialWorkload / fixtures)
- remove the standalone mixed-responsibility workloads/moe_reduce_fused.py
- update official correctness/benchmark tests to the fixture style and
  extend the vectorized PyTorch baseline to all four variants
- keep fp8 comparison in fp32 (equivalent fix already present)

Kernel, op, manifest and experiment docs are intentionally not synced.
Some checks failed
auto-label / apply-label (pull_request) Has been cancelled
Details
This pull request can be merged automatically.
This branch is out-of-date with the base branch
You are not authorized to merge this pull request.
You can also view command line instructions.

Step 1:

From your project repository, check out a new branch and test the changes.
git checkout -b ancient_rain-feat/moe_reduce_fused summer-camp-2026
git pull feat/moe_reduce_fused

Step 2:

Merge the changes and update on Gitea.
git checkout summer-camp-2026
git merge --no-ff ancient_rain-feat/moe_reduce_fused
git push origin summer-camp-2026
Sign in to join this conversation.
Reviewers
No reviewers
Labels
Clear labels
No items
No Label
Milestone
Clear milestone
No items
No Milestone
Projects
Clear projects
No project
Assignees
Clear assignees
No Assignees
1 Participants
Notifications
Due Date
The due date is invalid or out of range. Please use the format 'yyyy-mm-dd'.

No due date set.

Dependencies

No dependencies set.

Reference: ccf-ai-infra/TileOPs-Metax#55
No description provided.
Delete Branch "ancient_rain/TileOPs-Metax:feat/moe_reduce_fused"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?