[moe_reduce_fused] optimize H-dimension tiling and C500 tuning #2

cczw · 2026-08-05T10:49:28+08:00

cczw commented

2026-08-05 10:49:28 +08:00

状态：团队 fork 内部评审 PR；尚未向官方 upstream 提交。

Head：feat/moe_reduce_fused_czw_dev（643f3bf4d437f168980d134973c98f7ac61eb80d）
Base：feat/moe_reduce_fused（创建 PR 时远端 head 为 24403acfb3888ad3bd9570e0404c451cc00de8f1）

小组课题信息

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

课题完整简介：完成 moe_reduce_fused 在 TileOPs-Metax 中的 C500 迁移与调优，覆盖 weighted、XSF、FP8 output 和 quantized 四种变体；通过 H 维分块提高 decode workload 的 CTA/wave 覆盖，并提供完整 correctness、benchmark、mcProfiler、Roofline 语义分析和失败实验记录。

PR 简要描述

算子名称：moe_reduce_fused

算子认领 Issue：#9

小组：9 组

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

改动类型：

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

源实现：TileKernels-Metax-source/tile_kernels/moe/reduce_fused_kernel.py，源提交 0266ab740980de7dc03a828b8259cd73d100c2eb

优化 baseline：a409c2d18bf9ddbdc9393a172bed362d2ab001e5

本 PR 测试提交：643f3bf4d437f168980d134973c98f7ac61eb80d

1. 本次 PR 优化方案

优化前的主要问题：原 kernel 每个 token 只启动一个 CTA，单个 CTA 持有完整 H 维 FP32 accumulator。对 T=512 decode，尤其 FP8/quantized 路径，并行覆盖不足且单 CTA 活跃 accumulator 较大。

采用的优化方案：

将 grid 从单一 token 维扩展为 (token, hidden tile)；不同 H tile 写入互不重叠的输出区间，无需 atomic 或跨 CTA reduction。
增加内部编译期参数 block_h，每个 CTA 仅保留 [block_h] FP32 accumulator，同时保留完整 K 维归约顺序。
对 moe.yaml 的 10 个唯一 variant×shape 做完整 divisor sweep；未知 workload 保守回退到 block_h=hidden。
显式校验 block_h 的正值、32 对齐、整除 H 以及与 threads/layout 的兼容关系。

最终 dispatch：

workload	`block_h`
base T32/K2/H256	256
base T512/K8/H3072	1024
base T512/K8/H7168 decode	3584
base T4096/K8/H7168 prefill	7168
XSF tiny / large	128 / 3584
FP8 tiny / large	64 / 512
quantized tiny / large	128 / 512

同时按独立实验验证并回退了 threads=256、sf epilogue 融合、跨 H tile metadata 复用和 split-K 双 accumulator。失败方向均记录在 race_tests/experiment/v02～v05，未混入最终 kernel。

2. 精度验证

对比基准实现：tests/ops/test_moe_reduce_fused.py 中的独立 PyTorch reference。

测试命令：

cd /root/TileOPs-Metax
MACA_PATH=/opt/maca PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH=/opt/tilelang-metax-v0.1.10:/root/TileOPs-Metax \
/opt/conda/bin/python -m pytest -s tests/ops/test_moe_reduce_fused.py

测试 shape/dtype：23 个 correctness/error case；覆盖 FP16/BF16/FP32、四种变体、Manifest 大小 shape、-1 padding、mapping bounds、可选 out 和错误输入。

误差范围：BF16 rtol=atol=1.6e-2；FP16/FP32 常规路径 1e-3；FP8 output rtol=atol=0.1。

验证结果：23/23 passed。最终 Manifest benchmark 文件的 12/12 pytest 节点也全部通过（10 个唯一 variant×shape）。

3. 性能数据【必填】

测试环境：MetaX C500，25% Compute sGPU，Vram Quota 16000 MiB；MX-SMI 2.3.1��driver 3.8.30、MACA 3.7.1.5、Python 3.12.11、PyTorch 2.8.0+metax3.7.1.3、TileLang /opt/tilelang-metax-v0.1.10。

性能复现命令：

cd /root/TileOPs-Metax
MACA_PATH=/opt/maca PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH=/opt/tilelang-metax-v0.1.10:/root/TileOPs-Metax \
/opt/conda/bin/python -m pytest -s benchmarks/ops/bench_moe_reduce_fused.py

协议：显式 warmup；内部 warmup 10；50 repeats × 3 trials；L2 flush；CUPTI kernel-only；取三轮 trial mean 的 median；编译不计时。以下为独立 baseline full12 与当前 full12 对比，baseline 是原 kernel，不是 PyTorch latency：

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

mcProfiler 观测：

base T512/K8/H3072：Achieved/Dispatched waves 从 1024/1024 增至 3072/3072，global read/write 基本不变。
FP8 T512/K8/H7168：waves 从 1024/1024 增至 14336/14336，global read 只增加约 0.024%，标准 latency 改善约 38%。
L2 hit rate 没有改善，收益主要来自 H 维并行覆盖和单 CTA accumulator 压力缓解，而不是缓存命中。

Roofline：报告已记录每个唯一 workload 的语义 F/B/AI；由于缺少可审计的 C500 25% sGPU FP32 峰值和切片 HBM 带宽，不填写虚构的 efficiency。

详细证据：

race_tests/experiment/v01_reduce_fused_h_parallelism.md
race_tests/experiment/v02_threads_sweep.md
race_tests/experiment/v03_fuse_sf_epilogue.md
race_tests/experiment/v04_reuse_metadata_across_h_tiles.md
race_tests/experiment/v05_split_k_accumulator.md

4. 提交自检清单

全部测试用例运行通过，精度达标
已删除临时调试打印、冗余测试代码
性能指标可本地复现，数据真实有效

已知边界

full12 有 12 个 Manifest label，但只有 10 个唯一 variant×shape；三个 base H7168 decode label 参数相同。
benchmark 使用 identity mapping，不能直接外推到所有真实 MoE routing locality。
小于 3% 的变化不宣称稳定收益；tiny case 同时报告绝对量级。
当前 fork base 比本分支 merge-base 多一个队友实验提交；该提交未修改活动 reduce_fused.py，PR 差异仍限定为活动 kernel 和 v01～v05 报告。

> 状态：团队 fork 内部评审 PR；尚未向官方 upstream 提交。 > > Head：`feat/moe_reduce_fused_czw_dev`（`643f3bf4d437f168980d134973c98f7ac61eb80d`） > Base：`feat/moe_reduce_fused`（创建 PR 时远端 head 为 `24403acfb3888ad3bd9570e0404c451cc00de8f1`） ## 小组课题信息课题名称：面向 MetaX C500 的 `moe_reduce_fused` 算子迁移与性能优化课题完整简介：完成 `moe_reduce_fused` 在 TileOPs-Metax 中的 C500 迁移与调优，覆盖 weighted、XSF、FP8 output 和 quantized 四种变体；通过 H 维分块提高 decode workload 的 CTA/wave 覆盖，并提供完整 correctness、benchmark、mcProfiler、Roofline 语义分析和失败实验记录。 ### PR 简要描述算子名称：`moe_reduce_fused` 算子认领 Issue：[#9](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/9) 小组：9 组成员：朱洪彪、古雨、陈占炜、李明洋改动类型： - [ ] `feat`：新增算子或功能 - [x] `optimize`：优化已有实现源实现：`TileKernels-Metax-source/tile_kernels/moe/reduce_fused_kernel.py`，源提交 `0266ab740980de7dc03a828b8259cd73d100c2eb` 优化 baseline：`a409c2d18bf9ddbdc9393a172bed362d2ab001e5` 本 PR 测试提交：`643f3bf4d437f168980d134973c98f7ac61eb80d` ### 1. 本次 PR 优化方案优化前的主要问题：原 kernel 每个 token 只启动一个 CTA，单个 CTA 持有完整 H 维 FP32 accumulator。对 `T=512` decode，尤其 FP8/quantized 路径，并行覆盖不足且单 CTA 活跃 accumulator 较大。采用的优化方案： - 将 grid 从单一 token 维扩展为 `(token, hidden tile)`；不同 H tile 写入互不重叠的输出区间，无需 atomic 或跨 CTA reduction。 - 增加内部编译期参数 `block_h`，每个 CTA 仅保留 `[block_h]` FP32 accumulator，同时保留完整 K 维归约顺序。 - 对 `moe.yaml` 的 10 个唯一 variant×shape 做完整 divisor sweep；未知 workload 保守回退到 `block_h=hidden`。 - 显式校验 `block_h` 的正值、32 对齐、整除 H 以及与 threads/layout 的兼容关系。最终 dispatch： | workload | `block_h` | |---|---:| | base T32/K2/H256 | 256 | | base T512/K8/H3072 | 1024 | | base T512/K8/H7168 decode | 3584 | | base T4096/K8/H7168 prefill | 7168 | | XSF tiny / large | 128 / 3584 | | FP8 tiny / large | 64 / 512 | | quantized tiny / large | 128 / 512 | 同时按独立实验验证并回退了 threads=256、sf epilogue 融合、跨 H tile metadata 复用和 split-K 双 accumulator。失败方向均记录在 `race_tests/experiment/v02`～`v05`，未混入最终 kernel。 ### 2. 精度验证对比基准实现：`tests/ops/test_moe_reduce_fused.py` 中的独立 PyTorch reference。测试命令： ```bash cd /root/TileOPs-Metax MACA_PATH=/opt/maca PYTHONDONTWRITEBYTECODE=1 \ PYTHONPATH=/opt/tilelang-metax-v0.1.10:/root/TileOPs-Metax \ /opt/conda/bin/python -m pytest -s tests/ops/test_moe_reduce_fused.py ``` 测试 shape/dtype：23 个 correctness/error case；覆盖 FP16/BF16/FP32、四种变体、Manifest 大小 shape、`-1` padding、mapping bounds、可选 `out` 和错误输入。误差范围：BF16 `rtol=atol=1.6e-2`；FP16/FP32 常规路径 `1e-3`；FP8 output `rtol=atol=0.1`。验证结果：`23/23 passed`。最终 Manifest benchmark 文件的 `12/12` pytest 节点也全部通过（10 个唯一 variant×shape）。 ### 3. 性能数据【必填】测试环境：MetaX C500，25% Compute sGPU，Vram Quota 16000 MiB；MX-SMI 2.3.1��driver 3.8.30、MACA 3.7.1.5、Python 3.12.11、PyTorch 2.8.0+metax3.7.1.3、TileLang `/opt/tilelang-metax-v0.1.10`。性能复现命令： ```bash cd /root/TileOPs-Metax MACA_PATH=/opt/maca PYTHONDONTWRITEBYTECODE=1 \ PYTHONPATH=/opt/tilelang-metax-v0.1.10:/root/TileOPs-Metax \ /opt/conda/bin/python -m pytest -s benchmarks/ops/bench_moe_reduce_fused.py ``` 协议：显式 warmup；内部 warmup 10；50 repeats × 3 trials；L2 flush；CUPTI kernel-only；取三轮 trial mean 的 median；编译不计时。以下为独立 baseline full12 与当前 full12 对比，baseline 是原 kernel，不是 PyTorch latency： | 唯一 workload | baseline (ms) | 当前 (ms) | 变化 | |---|---:|---:|---:| | base H7168 decode | 0.0526 | 0.0513～0.0514 | 约 -2.3%，按噪声处理 | | base H7168 prefill | 0.3566 | 0.3565 | 无明确变化 | | base H3072 decode | 0.0285 | 0.0263 | -7.72%；三轮确认 -9.02% | | base tiny | 0.0093 | 0.0093 | 无明确变化 | | XSF H7168 | 0.0498 | 0.0489 | -1.81%，按噪声处理 | | XSF tiny | 0.0070 | 0.0069 | 绝对差为亚微秒 | | FP8 H7168 | 0.1014 | 0.0617 | **-39.15%**；三轮确认 -38.28% | | FP8 tiny | 0.0075 | 0.0070 | -6.67%，绝对差很小 | | quantized H7168 | 0.1009 | 0.0633 | **-37.26%**；三轮确认 -37.02% | | quantized tiny | 0.0054 | 0.0047 | -12.96%，绝对差很小 | mcProfiler 观测： - base T512/K8/H3072：Achieved/Dispatched waves 从 `1024/1024` 增至 `3072/3072`，global read/write 基本不变。 - FP8 T512/K8/H7168：waves 从 `1024/1024` 增至 `14336/14336`，global read 只增加约 0.024%，标准 latency 改善约 38%。 - L2 hit rate 没有改善，收益主要来自 H 维并行覆盖和单 CTA accumulator 压力缓解，而不是缓存命中。 Roofline：报告已记录每个唯一 workload 的语义 F/B/AI；由于缺少可审计的 C500 25% sGPU FP32 峰值和切片 HBM 带宽，不填写虚构的 efficiency。详细证据： - `race_tests/experiment/v01_reduce_fused_h_parallelism.md` - `race_tests/experiment/v02_threads_sweep.md` - `race_tests/experiment/v03_fuse_sf_epilogue.md` - `race_tests/experiment/v04_reuse_metadata_across_h_tiles.md` - `race_tests/experiment/v05_split_k_accumulator.md` ### 4. 提交自检清单 - [x] 全部测试用例运行通过，精度达标 - [x] 已删除临时调试打印、冗余测试代码 - [x] 性能指标可本地复现，数据真实有效 ### 已知边界 - full12 有 12 个 Manifest label，但只有 10 个唯一 variant×shape；三个 base H7168 decode label 参数相同。 - benchmark 使用 identity mapping，不能直接外推到所有真实 MoE routing locality。 - 小于 3% 的变化不宣称稳定收益；tiny case 同时报告绝对量级。 - 当前 fork base 比本分支 merge-base 多一个队友实验提交；该提交未修改活动 `reduce_fused.py`，PR 差异仍限定为活动 kernel 和 v01～v05 报告。

cczw added 5 commits 2026-08-05 10:49:29 +08:00

888e8d364c opt(moe): parallelize reduce fused across hidden tiles

Functional: 23/23 moe_reduce_fused correctness cases passed.

Benchmark: all 12 Manifest cases passed; H3072 base -7.7%, FP8 large -39.2%, quant large -37.3%, with no clear regression.

Profile: H3072 waves 3x and FP8 large waves 14x while global bytes stayed effectively flat.

Report: race_tests/experiment/v01_reduce_fused_h_parallelism.md

b181222fd0 docs(experiment): record reduce fused threads sweep

Benchmark: 36/36 legal fixed-block_h thread configurations completed across all 10 unique Manifest workloads.

Conclusion: keep the existing 128-thread large-workload default; 256 threads doubled waves but regressed H3072 latency by 4.3%.

No kernel changes are retained. Also add the required MACA_PATH to the v01 reproduction commands.

16de86f110 docs(experiment): record sf epilogue fusion result

Functional: 23/23 correctness and all 12 Manifest benchmark nodes passed.

Conclusion: FP8 and quantized large workloads showed no gain; generated code already keeps the extra sf pass in thread-local registers.

No kernel changes are retained.

5efd53e1fb docs(experiment): record H-tile metadata reuse result

Benchmark: all 23 legal tiles_per_cta candidates passed output comparison and profiling across the 10 unique workloads.

Profile: halving FP8 H workgroups saved only 13 KiB of reads while waves halved and latency regressed 4.6%.

No kernel changes are retained.

auto-label / apply-label (pull_request) Has been cancelled Details

643f3bf4d4 docs(experiment): record split-K accumulator result

Record 10 unique workload pairs: six K8 candidates remain within -0.96% to +1.68%, while four K2 candidates fail layout lowering. No kernel change is retained.

cczw added 1 commit 2026-08-05 12:10:18 +08:00

6383f0f9e4 docs(profile): establish reduce fused v06 baseline

Functional: 23/23 pass; full benchmark: 12/12 pass
Profile: six fixed-event mcProfiler cases on C500
Conclusion: FP8 H splitting improves wave coverage; base H7168 is already near high memory throughput
Report: race_tests/experiment/v06_profile_v01_bottleneck.md

cczw added 1 commit 2026-08-05 12:26:19 +08:00

ddc3a4c28a docs(experiment): record failed cta acc h decoupling

Functional: 23/23 pass plus 9/9 explicit candidate checks
Benchmark: all acc_h < cta_h candidates regressed by 2.94% to 24.44%
Profile: not collected because focused benchmark stably regressed
Kernel: restored exactly to v06 hash ed860456
Report: race_tests/experiment/v07_decouple_cta_acc_h.md

cczw added 1 commit 2026-08-05 12:39:44 +08:00

82bf9e5e70 docs(experiment): record grid axis order result

Record 23/23 plus focused correctness, three-round interleaved benchmark evidence, early Profile stop, and exact restoration of the retained kernel.

cczw added 1 commit 2026-08-05 13:17:58 +08:00

5f77f65198 docs(experiment): validate routing mapping robustness

Record 23/23 native correctness, 60/60 independent route checks across two implementations, 216 interleaved benchmark rows, and six per-kernel mapping profiles without changing the retained kernel.

cczw added 1 commit 2026-08-05 13:19:11 +08:00

1026ea837c docs(experiment): record k prefetch gate decision

Document why the required scoreboard, sustainable-bandwidth, and register-headroom gates are not simultaneously supported by auditable profiler evidence; no kernel candidate was implemented.

cczw added 1 commit 2026-08-05 13:28:27 +08:00

78fe788858 docs(experiment): close layout hole candidates

Record byte-identical default codegen, 23/23 plus 8/8 correctness, scalarized candidate codegen, and three-round regressions before restoring the retained kernel.

cczw added 1 commit 2026-08-05 14:24:37 +08:00

4a67046121 docs(experiment): finalize reduce fused candidate

Record 23/23 correctness, 30/30 routing robustness, 12/12 official benchmark, 108-row three-version interleaved CUPTI results, final mcProfiler evidence, and semantic Roofline analysis. The tested kernel remains byte-identical to v01.

cczw added 2 commits 2026-08-06 08:11:07 +08:00

0fef79e461 exp(moe): add reduce fused metadata pipeline variants

Functional: 23/23 op tests pass; 4/4 representative official variants pass on 16g-current C500.

Benchmark: not rerun; the default fragment metadata path and pipeline_stages=0 remain unchanged.

Profile: TODO, not collected for the experimental direct/shared metadata and two-stage pipeline paths.

Conclusion: preserve the alternative metadata and look-ahead pipeline implementations as an experiment checkpoint.

bb5e5feac1 test(moe): align reduce fused workloads and baselines

Review: fold the official matrix into workloads/moe.py using FixtureBase and WorkloadBase, and remove the standalone mixed-responsibility workload module.

Functional: 23/23 op tests pass; 4/4 representative official base/x_sf/fp8/quantized variants pass on 16g-current C500.

Contracts: official test and benchmark each collect 48 cases; benchmarks/tests 17/17 and tests/test_ops_manifest.py 7/7 pass; manifest validation passes.

Benchmark: extend the independent vectorized PyTorch baseline to all four variants; absolute timings were not rerun.

Also infer output shape from input shapes and compare FP8 results in FP32.

This pull request has changes conflicting with the target branch.

tileops/kernels/moe/reduce_fused.py
tileops/ops/moe/routed_expert/reduce_fused.py

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 feat/moe_reduce_fused_czw_dev feat/moe_reduce_fused

git pull origin feat/moe_reduce_fused_czw_dev

Step 2:

Merge the changes and update on Gitea.

git checkout feat/moe_reduce_fused

git merge --no-ff feat/moe_reduce_fused_czw_dev

git push origin feat/moe_reduce_fused

Sign in to join this conversation.

No reviewers

No Label

No Milestone

No project

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: ancient_rain/TileOPs-Metax#2