From 08f14fd996d61cadda024f5a81f2268e0ee2f757 Mon Sep 17 00:00:00 2001 From: hanrui22 Date: Fri, 21 Aug 2026 18:17:52 +0800 Subject: [PATCH 1/3] upload moe hc pre sinkhorn --- .../docs/perf-integrate/SYNC_VERIFY_REPORT.md | 71 ++++ .../perf-integrate/my_xllm_20cases_perf.json | 182 +++++++++ .../docs/perf-integrate/profile_20cases.py | 114 ++++++ .../docs/perf-integrate/run_20cases.py | 108 ++++++ .../hc_pre_sinkhorn/op_host/CMakeLists.txt | 1 + .../op_kernel/hc_pre_sinkhorn_regbase_base.h | 366 +++++++++++++++++- .../op_kernel/hc_pre_sinkhorn_regbase_perf.h | 128 +++++- 7 files changed, 955 insertions(+), 15 deletions(-) create mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md create mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json create mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py create mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md new file mode 100644 index 0000000..1bb926f --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md @@ -0,0 +1,71 @@ +# hc_pre_sinkhorn 同步到 my-xllm-ops — 精度与性能验证 + +日期：2026-08-21 +来源：`/home/h00801112/codex/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` +目标：`/home/h00801112/codex/my-xllm-ops/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` + +## 1. 同步内容 + +| 文件 | 变更 | md5 | +|------|------|-----| +| `op_kernel/hc_pre_sinkhorn_regbase_base.h` | SoA CombFrag、MultiCopy 转置、ldva TransDataTo5HD + `S_V`/`V_S` | `f152f40d187453990e9c9c882aa38358`（两边一致） | +| `op_kernel/hc_pre_sinkhorn_regbase_perf.h` | `useCombSoa` 分发；`hcScale` 仍 `GetValue()` | `8754c1d245ccdc4f21a389d501e4ce6f`（两边一致） | +| `op_host/CMakeLists.txt` | opc 增加 `-Wno-constant-conversion` | 已同步 | + +未改：`hc_pre_sinkhorn.cpp`、tiling、def、proto、membase kernel。 + +## 2. 编译说明 + +`my-xllm-ops/xllm-ops` 全新 `build.sh` 会重新下载 `json` / `abseil`，gitcode 返回 403，cmake 失败。 +本次用 dest 源码经 opc 编出 `HcPreSinkhorn_20e241df….o`，装入 `${ASCEND_HOME_PATH}/opp/vendors/custom_xllm_math`。host/tiling/aclnn 与源仓一致，无需重编。 + +## 3. 精度 + +| 套件 | 结果 | +|------|------| +| `test/python_test/test_hc_pre_sinkhorn.py`（10 func） | **10/10 PASS** | +| cann-bench 20 case（y / post / comb） | **20/20 PASS** | + +阈值：y bf16 rtol/atol=1e-2；post fp32 1e-4；comb fp32 1e-3。 + +SoA 路径最大 abs：y ≤ 1.5e-5，post ≤ 2.4e-7，comb ≤ 2.1e-7。与源仓 xllm-ops 合入结果一致。 + +## 4. 性能（HcPreSinkhorn kernel µs） + +设备：Ascend950PR。口径：torch_npu profiler。 +对照：源仓合入实测、A5 `mhc_sinkhorn` baseline。 + +| case | bs | M | iter | my-xllm us | xllm-ops us | A5 us | vs A5 | +|------|----|---|------|------------|-------------|-------|-------| +| 1 | 1 | 4 | 20 | 3.142 | 3.138 | 5.36 | 1.71x | +| 2 | 64 | 4 | 20 | 4.546 | 4.650 | 5.82 | 1.28x | +| 3 | 512 | 4 | 20 | 5.512 | 5.505 | 8.47 | 1.54x | +| 4 | 1024 | 4 | 20 | 5.407 | 5.418 | 10.00 | 1.85x | +| 5 | 4096 | 4 | 20 | 6.978 | 6.944 | 10.00 | 1.43x | +| 6 | 16384 | 4 | 20 | 12.884 | 12.828 | 15.90 | 1.23x | +| 7 | 1024 | 2 | 20 | 4.640 | 4.653 | 10.00 | 2.16x | +| 8 | 1024 | 3 | 20 | 5.323 | 5.310 | 10.00 | 1.88x | +| 9 | 1024 | 6 | 20 | 18.812 | 18.768 | 10.00 | 0.53x | +| 10 | 1024 | 8 | 20 | 20.786 | 20.850 | 10.00 | 0.48x | +| 11 | 1024 | 12 | 20 | 24.255 | 24.337 | 10.00 | 0.41x | +| 12 | 1024 | 16 | 20 | 29.423 | 29.455 | 15.50 | 0.53x | +| 13 | 1024 | 4 | 1 | 4.635 | 4.646 | 4.17 | 0.90x | +| 14 | 1024 | 4 | 5 | 4.521 | 4.548 | 5.65 | 1.25x | +| 15 | 1024 | 4 | 40 | 6.498 | 6.578 | 10.00 | 1.54x | +| 16 | 1024 | 4 | 20 | 5.760 | 5.745 | 10.00 | 1.74x | +| 17 | 1024 | 4 | 20 | 5.376 | 5.376 | 10.00 | 1.86x | +| 18 | 1 | 2 | 20 | 3.070 | 3.053 | 10.00 | 3.26x | +| 19 | 1024 | 4 | 20 | 5.524 | 5.427 | 10.00 | 1.81x | +| 20 | 1024 | 4 | 20 | 5.724 | 5.763 | 10.00 | 1.75x | + +- 与源仓 xllm-ops：**数值对齐**（差 < 0.15 µs，测量抖动） +- vs A5：**15/20 更快**，大 batch M≤4 与源仓同档（case 6 ≈ 12.9 µs） +- M≥6 仍走 AoS，与 A5 比偏慢（算子还含 pre/y/post，A5 仅 Sinkhorn） + +墙钟（aclnn 下发）约 260–400 µs，不能当 kernel 时间。 + +原始数据：`docs/perf-integrate/my_xllm_20cases_perf.json` + +## 5. 结论 + +代码已同步到 my-xllm-ops，精度 10/10 + 20/20 通过，性能与源仓合入结果一致。 diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json new file mode 100644 index 0000000..2825ee2 --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json @@ -0,0 +1,182 @@ +[ + { + "case_id": 1, + "bs": 1, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 3.142, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 2, + "bs": 64, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 4.546, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 3, + "bs": 512, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.512, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 4, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.407, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 5, + "bs": 4096, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 6.978, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 6, + "bs": 16384, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 12.884, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 7, + "bs": 1024, + "M": 2, + "d": 2, + "iters": 20, + "kernel_us": 4.64, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 8, + "bs": 1024, + "M": 3, + "d": 3, + "iters": 20, + "kernel_us": 5.323, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 9, + "bs": 1024, + "M": 6, + "d": 6, + "iters": 20, + "kernel_us": 18.812, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 10, + "bs": 1024, + "M": 8, + "d": 8, + "iters": 20, + "kernel_us": 20.786, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 11, + "bs": 1024, + "M": 12, + "d": 12, + "iters": 20, + "kernel_us": 24.255, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 12, + "bs": 1024, + "M": 16, + "d": 16, + "iters": 20, + "kernel_us": 29.423, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 13, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 1, + "kernel_us": 4.635, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 14, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 5, + "kernel_us": 4.521, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 15, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 40, + "kernel_us": 6.498, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 16, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.76, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 17, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.376, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 18, + "bs": 1, + "M": 2, + "d": 2, + "iters": 20, + "kernel_us": 3.07, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 19, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.524, + "kernel": "HcPreSinkhorn" + }, + { + "case_id": 20, + "bs": 1024, + "M": 4, + "d": 4, + "iters": 20, + "kernel_us": 5.724, + "kernel": "HcPreSinkhorn" + } +] \ No newline at end of file diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py new file mode 100644 index 0000000..78c84e5 --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""On-device kernel time for the 20 cann-bench cases via torch_npu profiler.""" +from __future__ import annotations + +import csv +import glob +import json +import shutil +import sys +from collections import defaultdict +from pathlib import Path + +import torch +import torch_npu + +ROOT = Path("/home/h00801112/codex/my-xllm-ops/xllm-ops") +sys.path.insert(0, str(ROOT / "test/python_test")) +import custom_ops # noqa: E402 + +CASES_CSV = Path("/home/h00801112/codex/bestline4/operators/hc_pre_sinkhorn/round2/cases.csv") +OUT = ROOT / "xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json" +WORK = Path("/tmp/hc_pre_sinkhorn_prof_my_xllm") + + +def parse_kernel_csv(work: Path) -> dict: + stats = defaultdict(lambda: [0, 0.0]) + candidates = glob.glob(str(work / "**" / "kernel_details.csv"), recursive=True) + if not candidates: + candidates = glob.glob(str(work / "**" / "op_statistic*.csv"), recursive=True) + for path in candidates: + with open(path, newline="", encoding="utf-8") as fh: + for row in csv.DictReader(fh): + name = (row.get("Name") or row.get("OP Type") or "").strip() + dur = row.get("Duration(us)") or row.get("Total Time(us)") or "0" + if not name: + continue + try: + val = float(dur) + except ValueError: + continue + stats[name][0] += 1 + stats[name][1] += val + return {k: {"count": v[0], "total_us": round(v[1], 3), + "avg_us": round(v[1] / max(v[0], 1), 3)} + for k, v in stats.items()} + + +def pick(stats: dict, needle: str): + for name, val in stats.items(): + if needle.lower() in name.lower(): + return {"kernel": name, **val} + return None + + +def profile_once(fn, tag: str, reps: int = 20) -> dict: + work = WORK / tag + if work.exists(): + shutil.rmtree(work) + work.mkdir(parents=True, exist_ok=True) + exp_cfg = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + ) + for _ in range(5): + fn() + torch.npu.synchronize() + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU], + experimental_config=exp_cfg, + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(str(work)), + ) as prof: + for _ in range(reps): + fn() + prof.step() + torch.npu.synchronize() + return parse_kernel_csv(work) + + +def main(): + torch_npu.npu.set_device(0) + rows = list(csv.DictReader(CASES_CSV.open())) + results = [] + print(f"{'case':>4} {'bs':>6} {'M':>3} {'iter':>5} {'kernel_us':>10} {'kernel':>24}") + for row in rows: + cid = int(row["case_id"]) + bs, m, d = int(row["bs"]), int(row["hc_mult"]), int(row["d"]) + iters, eps = int(row["hc_sinkhorn_iters"]), float(row["hc_eps"]) + hc_mix = 2 * m + m * m + g = torch.Generator().manual_seed(2026 + cid) + mixes = torch.empty(bs, hc_mix).uniform_(-1, 1, generator=g).npu() + rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g).npu() + hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g).npu() + hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g).npu() + x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16).npu() + + def fn(): + return custom_ops.hc_pre_sinkhorn_npu( + mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps) + + stats = profile_once(fn, f"c{cid}", reps=20) + hit = pick(stats, "HcPreSinkhorn") or pick(stats, "hc_pre_sinkhorn") + us = hit["avg_us"] if hit else float("nan") + name = hit["kernel"] if hit else "n/a" + results.append({"case_id": cid, "bs": bs, "M": m, "d": d, "iters": iters, + "kernel_us": us, "kernel": name}) + print(f"{cid:>4} {bs:>6} {m:>3} {iters:>5} {us:>10} {name:>24}") + sys.stdout.flush() + + OUT.write_text(json.dumps(results, indent=2)) + print(f"\nwritten {OUT}") + + +if __name__ == "__main__": + main() diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py new file mode 100644 index 0000000..d38d751 --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Precision + wall-clock check of the 20 cann-bench cases via custom_ops.""" +from __future__ import annotations + +import csv +import sys +import time +from pathlib import Path + +import torch + +torch_npu = __import__("torch_npu") +ROOT = Path("/home/h00801112/codex/my-xllm-ops/xllm-ops") +sys.path.insert(0, str(ROOT / "test/python_test")) +import custom_ops # noqa: E402 +from test_hc_pre_sinkhorn import hc_pre_sinkhorn_golden # noqa: E402 + +CASES_CSV = Path("/home/h00801112/codex/bestline4/operators/hc_pre_sinkhorn/round2/cases.csv") + + +def make_inputs(row): + bs = int(row["bs"]) + m = int(row["hc_mult"]) + d = int(row["d"]) + iters = int(row["hc_sinkhorn_iters"]) + eps = float(row["hc_eps"]) + hc_mix = 2 * m + m * m + vr = row.get("value_range", "[-1,1]") + g = torch.Generator().manual_seed(2026 + int(row["case_id"])) + + if "nan" in vr: + mixes = torch.full((bs, hc_mix), float("nan")) + rsqrt = torch.full((bs,), float("nan")) + hc_scale = torch.full((3,), float("nan")) + hc_base = torch.full((hc_mix,), float("nan")) + x = torch.full((bs, m, d), float("nan"), dtype=torch.bfloat16) + elif "inf" in vr: + mixes = torch.empty(bs, hc_mix).uniform_(-2, 2, generator=g) + mixes.view(-1)[0] = float("inf") + mixes.view(-1)[1] = float("-inf") + rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g) + hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g) + hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g) + x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16) + else: + mixes = torch.empty(bs, hc_mix).uniform_(-1, 1, generator=g) + rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g) + hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g) + hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g) + x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16) + return mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps + + +def max_err(a, b): + if not torch.isfinite(a).any() and not torch.isfinite(b).any(): + return 0.0, 0.0 + diff = (a.float() - b.float()).abs() + denom = b.float().abs().clamp_min(1e-8) + return float(diff.max()), float((diff / denom).max()) + + +def time_op(fn, warmup=10, iters=30): + for _ in range(warmup): + fn() + torch.npu.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.npu.synchronize() + return (time.perf_counter() - t0) * 1e6 / iters + + +def main(): + rows = list(csv.DictReader(CASES_CSV.open())) + print(f"{'case':>4} {'bs':>6} {'M':>3} {'d':>5} {'iter':>5} {'wall_us':>9} {'y':>8} {'post':>8} {'comb':>8} {'ok':>5}") + fails = 0 + recs = [] + for row in rows: + mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps = make_inputs(row) + y_g, post_g, comb_g = hc_pre_sinkhorn_golden( + mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps) + + def launch(): + return custom_ops.hc_pre_sinkhorn_npu( + mixes.npu(), rsqrt.npu(), hc_scale.npu(), hc_base.npu(), x.npu(), + m, iters, eps) + + y_n, post_n, comb_n = launch() + ye, yr = max_err(y_n.cpu(), y_g) + pe, pr = max_err(post_n.cpu(), post_g) + ce, cr = max_err(comb_n.cpu(), comb_g) + y_ok = ye <= 1e-2 or yr <= 1e-2 or (not torch.isfinite(y_g).any()) + p_ok = pe <= 1e-4 or pr <= 1e-4 or (not torch.isfinite(post_g).any()) + c_ok = ce <= 1e-3 or cr <= 1e-3 or (not torch.isfinite(comb_g).any()) + ok = y_ok and p_ok and c_ok + fails += int(not ok) + us = time_op(launch) + print(f"{row['case_id']:>4} {row['bs']:>6} {m:>3} {row['d']:>5} {iters:>5} " + f"{us:>9.2f} {ye:>8.1e} {pe:>8.1e} {ce:>8.1e} {'PASS' if ok else 'FAIL':>5}") + recs.append({"case_id": int(row["case_id"]), "ok": ok, "wall_us": us, + "y_abs": ye, "post_abs": pe, "comb_abs": ce}) + sys.stdout.flush() + print(f"\n{len(rows) - fails}/{len(rows)} PASS") + return 0 if fails == 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt b/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt index 54d9747..74f3ca8 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt @@ -52,6 +52,7 @@ add_ops_compile_options( OP_NAME HcPreSinkhorn OPTIONS --cce-auto-sync=off -Wno-deprecated-declarations + -Wno-constant-conversion -Werror -mllvm -cce-aicore-hoist-movemask=false --op_relocatable_kernel_binary=true diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h index 5261122..fa35ab5 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h @@ -26,6 +26,10 @@ using AscendC::MicroAPI::RegTensor; using AscendC::MicroAPI::UnalignReg; constexpr int32_t BLOCK_SIZE = 32; constexpr int32_t VL_FP32 = 64; +// Number of comb-matrix rows held in registers by the four-unfold fast path. +constexpr int32_t COMB_UNFOLD_NUM = 4; +// Below this many rows per tile the plane-major transpose overhead outweighs the gain. +constexpr int32_t COMB_SOA_MIN_ROWS = 4; __aicore__ inline int32_t CeilDiv(int32_t a, int b) { @@ -302,7 +306,350 @@ __aicore__ inline void VFProcessCombFragRLessVL( } } -__aicore__ inline void VFProcessIteration(RegTensor& sum0, RegTensor& sum1, RegTensor& mix, float eps, MaskReg pregLoop) +// [vec-05 + state_resident] Register-resident Sinkhorn iteration for generic M. +// Phase 1 (initial softmax + column norm) is identical to RLessVL (uses UB staging). +// Phase 2 (Sinkhorn iterations) loads M rows into RegTensor, iterates entirely in +// registers with no UB load/store and no LocalMemBar, then stores M rows back once. +// This eliminates iters * dim0 * M * 2 UB load/store + iters * dim0 LocalMemBar. +// +// Note: RegTensor arrays (RegTensor mix[M]) are NOT supported by bisheng backend +// ("Unsupported Inst must be hoisted"). Use individual variables + if constexpr chains, +// consistent with the FourUnfold path (mix1..mix4). The compiler eliminates unused +// variables and dead branches for each template instantiation. + +// Helper macros for compile-time unrolling over individual RegTensor variables +#define REG_LOAD_N(N) \ + if constexpr (M > N) { \ + LoadInputData(mix##N, combFragLocalAddr, pregLoop, \ + i * dim1 * dim2Align + (N) * dim2Align); \ + } + +#define REG_ROW_NORM_N(N) \ + if constexpr (M > N) { \ + ReduceSum(sumR, mix##N, pregLoop); \ + Duplicate(sumR, sumR, pregLoop); \ + Adds(sumR, sumR, eps, pregLoop); \ + Div(mix##N, mix##N, sumR, pregLoop); \ + Add(sumC, sumC, mix##N, pregLoop); \ + } + +#define REG_COL_NORM_N(N) \ + if constexpr (M > N) { \ + Div(mix##N, mix##N, sumC, pregLoop); \ + } + +#define REG_STORE_N(N) \ + if constexpr (M > N) { \ + StoreOutputData(combFragLocalAddr, mix##N, pregLoop, \ + i * dim1 * dim2Align + (N) * dim2Align); \ + } + +// Expand all 16 slots +#define REG_LOAD_ALL \ + REG_LOAD_N(0) REG_LOAD_N(1) REG_LOAD_N(2) REG_LOAD_N(3) \ + REG_LOAD_N(4) REG_LOAD_N(5) REG_LOAD_N(6) REG_LOAD_N(7) \ + REG_LOAD_N(8) REG_LOAD_N(9) REG_LOAD_N(10) REG_LOAD_N(11) \ + REG_LOAD_N(12) REG_LOAD_N(13) REG_LOAD_N(14) REG_LOAD_N(15) + +#define REG_ROW_NORM_ALL \ + REG_ROW_NORM_N(0) REG_ROW_NORM_N(1) REG_ROW_NORM_N(2) REG_ROW_NORM_N(3) \ + REG_ROW_NORM_N(4) REG_ROW_NORM_N(5) REG_ROW_NORM_N(6) REG_ROW_NORM_N(7) \ + REG_ROW_NORM_N(8) REG_ROW_NORM_N(9) REG_ROW_NORM_N(10) REG_ROW_NORM_N(11) \ + REG_ROW_NORM_N(12) REG_ROW_NORM_N(13) REG_ROW_NORM_N(14) REG_ROW_NORM_N(15) + +#define REG_COL_NORM_ALL \ + REG_COL_NORM_N(0) REG_COL_NORM_N(1) REG_COL_NORM_N(2) REG_COL_NORM_N(3) \ + REG_COL_NORM_N(4) REG_COL_NORM_N(5) REG_COL_NORM_N(6) REG_COL_NORM_N(7) \ + REG_COL_NORM_N(8) REG_COL_NORM_N(9) REG_COL_NORM_N(10) REG_COL_NORM_N(11) \ + REG_COL_NORM_N(12) REG_COL_NORM_N(13) REG_COL_NORM_N(14) REG_COL_NORM_N(15) + +#define REG_STORE_ALL \ + REG_STORE_N(0) REG_STORE_N(1) REG_STORE_N(2) REG_STORE_N(3) \ + REG_STORE_N(4) REG_STORE_N(5) REG_STORE_N(6) REG_STORE_N(7) \ + REG_STORE_N(8) REG_STORE_N(9) REG_STORE_N(10) REG_STORE_N(11) \ + REG_STORE_N(12) REG_STORE_N(13) REG_STORE_N(14) REG_STORE_N(15) + +template +__aicore__ inline void VFProcessCombFragRegResident( + const LocalTensor& combFragLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t iters, uint16_t dim0, uint16_t dim1, + uint16_t dim2) +{ + __local_mem__ float* combFragLocalAddr = (__local_mem__ float*)combFragLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint32_t dim2Align = RoundUp(dim2); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor rsqrt; + RegTensor max; + RegTensor sum; + RegTensor sum1; + uint32_t sreg = dim2; + // [vec-09] mask creation done once, outside all loops + MaskReg pregLoop = UpdateMask(sreg); + + // Phase 1a: Initial softmax per row + accumulate column sum + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputData(mix, mixLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + // Phase 1b: Column normalization + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + + // Phase 2: Register-resident Sinkhorn iterations + // [vec-05] Load M rows into individual RegTensors, iterate without UB staging. + // Declare up to 16 individual RegTensor variables; compiler eliminates unused. + LocalMemBar(); + RegTensor mix0, mix1, mix2, mix3, mix4, mix5, mix6, mix7; + RegTensor mix8, mix9, mix10, mix11, mix12, mix13, mix14, mix15; + RegTensor sumR; // row sum (reused per row) + RegTensor sumC; // column sum accumulator + + for (uint16_t i = 0; i < dim0; i++) { + // Load M rows into registers (one-time UB load per dim0 iteration) + REG_LOAD_ALL + + // Sinkhorn iterations: all in registers, no UB load/store, no LocalMemBar + for (uint16_t iter = 0; iter < iters; iter++) { + Duplicate(sumC, static_cast(0), pregLoop); + // Row normalization + accumulate column sum + REG_ROW_NORM_ALL + // Column normalization + Adds(sumC, sumC, eps, pregLoop); + REG_COL_NORM_ALL + } + + // Store M rows back to UB (one-time UB store per dim0 iteration) + REG_STORE_ALL + } + } +} + +// =================================================================================== +// [round3 / vec-10] SoA (plane-major) comb_frag path for hcMult <= 4. +// +// Stores the M*M matrix entries as separate UB planes of length aAlign, so both +// Sinkhorn reductions become elementwise ops at full lane occupancy. +// MTE2 : NDDMA MultiCopy does (A, M*M) -> (M*M, A) on the fly. +// VEC : register-resident planes, no UB traffic per iteration. +// MTE3 : TransDataTo5HD (ldva address list) converts back to GM layout. +// =================================================================================== +constexpr int32_t COMB_SOA_MAX_MULT = 4; +constexpr int32_t TRANS_BLOCK_HALF = 8; +constexpr int32_t TRANS_BLOCK_FULL = 16; +// GM [A, hcMix] -> UB planes [planes, aAlign] via NDDMA MultiCopy. +// opc for this op is compiled with -Wno-constant-conversion so default +// NdDmaConfig::unsetPad survives -Werror. +__aicore__ inline void CopyInCombTransposed( + const GlobalTensor& mixesGm, const LocalTensor& combT, uint32_t curA, uint32_t aAlign, + uint32_t hcMix, uint32_t planes) +{ + MultiCopyLoopInfo<2> loopInfo; + loopInfo.loopSrcStride[0] = 1; + loopInfo.loopDstStride[0] = aAlign; + loopInfo.loopSize[0] = planes; + loopInfo.loopSrcStride[1] = hcMix; + loopInfo.loopDstStride[1] = 1; + loopInfo.loopSize[1] = curA; + MultiCopyParams params = {loopInfo, 0.0f}; + DataCopy(combT, mixesGm, params); +} + +#define SOA_HAS(R, C) ((R) < M && (C) < M) +#define SOA_PLANE(R, C) (((R) * M + (C)) * aAlign + off) + +#define SOA_AFFINE(R, C) \ + if constexpr (SOA_HAS(R, C)) { \ + DataCopy(c##R##C, combAddr + SOA_PLANE(R, C)); \ + Mul(c##R##C, c##R##C, rsq, preg); \ + Muls(c##R##C, c##R##C, scale, preg); \ + LoadInputDataWithBrc(bas, baseAddr, preg, (R) * M + (C)); \ + Add(c##R##C, c##R##C, bas, preg); \ + } + +#define SOA_AFFINE_ROW(R) SOA_AFFINE(R, 0) SOA_AFFINE(R, 1) SOA_AFFINE(R, 2) SOA_AFFINE(R, 3) + +#define SOA_REDUCE_ROW(OP, R) \ + OP(acc, c##R##0, c##R##1, preg); \ + if constexpr (M > 2) { OP(acc, acc, c##R##2, preg); } \ + if constexpr (M > 3) { OP(acc, acc, c##R##3, preg); } + +#define SOA_APPLY_ROW(OP, R) \ + OP(c##R##0, c##R##0, acc, preg); \ + OP(c##R##1, c##R##1, acc, preg); \ + if constexpr (M > 2) { OP(c##R##2, c##R##2, acc, preg); } \ + if constexpr (M > 3) { OP(c##R##3, c##R##3, acc, preg); } + +#define SOA_EXP_ROW(R) \ + Exp(c##R##0, c##R##0, preg); \ + Exp(c##R##1, c##R##1, preg); \ + if constexpr (M > 2) { Exp(c##R##2, c##R##2, preg); } \ + if constexpr (M > 3) { Exp(c##R##3, c##R##3, preg); } + +#define SOA_ADDS_ROW(R) \ + Adds(c##R##0, c##R##0, eps, preg); \ + Adds(c##R##1, c##R##1, eps, preg); \ + if constexpr (M > 2) { Adds(c##R##2, c##R##2, eps, preg); } \ + if constexpr (M > 3) { Adds(c##R##3, c##R##3, eps, preg); } + +#define SOA_SOFTMAX_ROW(R) \ + if constexpr ((R) < M) { \ + SOA_REDUCE_ROW(Max, R) \ + SOA_APPLY_ROW(Sub, R) \ + SOA_EXP_ROW(R) \ + SOA_REDUCE_ROW(Add, R) \ + SOA_APPLY_ROW(Div, R) \ + SOA_ADDS_ROW(R) \ + } + +#define SOA_ROW_NORM(R) \ + if constexpr ((R) < M) { \ + SOA_REDUCE_ROW(Add, R) \ + Adds(acc, acc, eps, preg); \ + SOA_APPLY_ROW(Div, R) \ + } + +#define SOA_COL_NORM(C) \ + if constexpr ((C) < M) { \ + Add(acc, c0##C, c1##C, preg); \ + if constexpr (M > 2) { Add(acc, acc, c2##C, preg); } \ + if constexpr (M > 3) { Add(acc, acc, c3##C, preg); } \ + Adds(acc, acc, eps, preg); \ + Div(c0##C, c0##C, acc, preg); \ + Div(c1##C, c1##C, acc, preg); \ + if constexpr (M > 2) { Div(c2##C, c2##C, acc, preg); } \ + if constexpr (M > 3) { Div(c3##C, c3##C, acc, preg); } \ + } + +#define SOA_STORE(R, C) \ + if constexpr (SOA_HAS(R, C)) { \ + DataCopy(combAddr + SOA_PLANE(R, C), c##R##C, preg); \ + } + +#define SOA_STORE_ROW(R) SOA_STORE(R, 0) SOA_STORE(R, 1) SOA_STORE(R, 2) SOA_STORE(R, 3) + +template +__aicore__ inline void VFProcessCombFragSoA( + const LocalTensor& combT, const LocalTensor& hcBaseLocal, const LocalTensor& rsqrtLocal, + float scale, float eps, uint16_t iters, uint16_t curA, uint32_t aAlign) +{ + __local_mem__ float* combAddr = (__local_mem__ float*)combT.GetPhyAddr(); + __local_mem__ float* baseAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint16_t chunkCount = CeilDiv(curA, VL_FP32); + __VEC_SCOPE__ + { + RegTensor c00, c01, c02, c03; + RegTensor c10, c11, c12, c13; + RegTensor c20, c21, c22, c23; + RegTensor c30, c31, c32, c33; + RegTensor rsq; + RegTensor bas; + RegTensor acc; + uint32_t remain = curA; + MaskReg preg; + for (uint16_t chunk = 0; chunk < chunkCount; chunk++) { + preg = UpdateMask(remain); + uint32_t off = chunk * VL_FP32; + DataCopy(rsq, rsqrtAddr + off); + SOA_AFFINE_ROW(0) + SOA_AFFINE_ROW(1) + SOA_AFFINE_ROW(2) + SOA_AFFINE_ROW(3) + SOA_SOFTMAX_ROW(0) + SOA_SOFTMAX_ROW(1) + SOA_SOFTMAX_ROW(2) + SOA_SOFTMAX_ROW(3) + SOA_COL_NORM(0) + SOA_COL_NORM(1) + SOA_COL_NORM(2) + SOA_COL_NORM(3) + for (uint16_t iter = 0; iter < iters; iter++) { + SOA_ROW_NORM(0) + SOA_ROW_NORM(1) + SOA_ROW_NORM(2) + SOA_ROW_NORM(3) + SOA_COL_NORM(0) + SOA_COL_NORM(1) + SOA_COL_NORM(2) + SOA_COL_NORM(3) + } + SOA_STORE_ROW(0) + SOA_STORE_ROW(1) + SOA_STORE_ROW(2) + SOA_STORE_ROW(3) + } + } +} + +__aicore__ inline void TransposeCombSoAToAoS( + const LocalTensor& dst, const LocalTensor& src, const LocalTensor& vaAddr, + uint16_t curA, uint32_t aAlign, int32_t planes) +{ + int32_t groups = CeilDiv(planes, TRANS_BLOCK_HALF); + uint8_t aRepeat = static_cast(CeilDiv(curA, TRANS_BLOCK_FULL)); + TransDataTo5HDParams params; + params.repeatTimes = aRepeat; + params.srcRepStride = (aRepeat == 1) ? 0 : 2; + params.dstRepStride = (aRepeat == 1) ? 0 : TRANS_BLOCK_FULL * (TRANS_BLOCK_FULL / TRANS_BLOCK_HALF); + uint64_t dstBase = (uint64_t)(__ubuf__ float*)dst.GetPhyAddr(); + uint64_t srcBase = (uint64_t)(__ubuf__ float*)src.GetPhyAddr(); + LocalTensor dstAddr = vaAddr; + LocalTensor srcAddr = vaAddr[TRANS_BLOCK_FULL]; + // xllm-ops compiles with --cce-auto-sync=off. SetValue is scalar-pipe UB + // writes; TransDataTo5HD consumes the address list on the vector pipe. + PipeBarrier(); + event_t evSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + event_t evVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + for (int32_t i = 0; i < groups; i++) { + for (int32_t j = 0; j < TRANS_BLOCK_HALF; j++) { + int32_t plane = i * TRANS_BLOCK_HALF + j; + uint64_t srcOff = (uint64_t)(plane < planes ? plane : 0) * aAlign * sizeof(float); + srcAddr.SetValue(j, srcBase + srcOff); + srcAddr.SetValue(j + TRANS_BLOCK_HALF, srcBase + srcOff + TRANS_BLOCK_HALF * sizeof(float)); + dstAddr.SetValue( + j * 2, dstBase + (uint64_t)(i * TRANS_BLOCK_HALF + j * TRANS_BLOCK_FULL) * sizeof(float)); + dstAddr.SetValue( + j * 2 + 1, + dstBase + + (uint64_t)(i * TRANS_BLOCK_HALF + (j + TRANS_BLOCK_HALF) * TRANS_BLOCK_FULL) * sizeof(float)); + } + SetFlag(evSV); + WaitFlag(evSV); + AscendC::TransDataTo5HD(dstAddr, srcAddr, params); + SetFlag(evVS); + WaitFlag(evVS); + } +} + +__aicore__ inline void VFProcessIteration(RegTensor& sum0, RegTensor& sum1, RegTensor& mix, float eps, MaskReg pregLoop) { ReduceSum(sum1, mix, pregLoop); Duplicate(sum1, sum1, pregLoop); @@ -506,6 +853,23 @@ __aicore__ inline void CopyOut( DataCopyPad(outputGm, outputTensor, dataCopyParams); } +__aicore__ inline void CopyOutCombSoA( + const LocalTensor& outputTensor, const GlobalTensor& outputGm, uint32_t curA, int32_t planes) +{ + if (planes == TRANS_BLOCK_FULL) { + CopyOut(outputTensor, outputGm, 1, curA * TRANS_BLOCK_FULL); + return; + } + constexpr int32_t rowBytes = TRANS_BLOCK_FULL * sizeof(float); + int32_t burstBytes = planes * sizeof(float); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = static_cast(curA); + dataCopyParams.blockLen = burstBytes; + dataCopyParams.srcStride = (rowBytes - CeilAlign(burstBytes, ONE_BLK_SIZE)) / ONE_BLK_SIZE; + dataCopyParams.dstStride = 0; + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + } // namespace HCPreSinkhorn #endif \ No newline at end of file diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h index 56e9c52..25091f1 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h @@ -43,12 +43,20 @@ class HcPreSinkhornPerf { postGm.SetGlobalBuffer((__gm__ float*)post); combFragGm.SetGlobalBuffer((__gm__ float*)combFrag); + useCombSoa = (tilingData->hcMult <= COMB_SOA_MAX_MULT) && (tilingData->rowFactor >= COMB_SOA_MIN_ROWS); + combPlanes = tilingData->hcMult * tilingData->hcMult; + combAAlign = CeilAlign(static_cast(tilingData->rowFactor), VL_FP32); + // InQue int64_t mixesQue01Size = tilingData->rowFactor * tilingData->hcMultAlign * 2 * sizeof(float); pipe->InitBuffer(mixesQue01, 2, mixesQue01Size); - pipe->InitBuffer( - mixesQue2, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); - pipe->InitBuffer(rsqrtQue, 2, RoundUp(tilingData->rowFactor) * sizeof(float)); + int64_t combInSize = useCombSoa ? + (int64_t)combPlanes * combAAlign * sizeof(float) : + tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float); + pipe->InitBuffer(mixesQue2, 2, combInSize); + int64_t rsqrtSize = useCombSoa ? combAAlign * sizeof(float) : + RoundUp(tilingData->rowFactor) * sizeof(float); + pipe->InitBuffer(rsqrtQue, 2, rsqrtSize); pipe->InitBuffer( xQue, 2, tilingData->rowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); @@ -56,8 +64,8 @@ class HcPreSinkhornPerf { pipe->InitBuffer( yQue, 2, tilingData->rowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); pipe->InitBuffer(postQue, 2, tilingData->rowFactor * tilingData->hcMultAlign * sizeof(float)); - pipe->InitBuffer( - combFragQue, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + int64_t combOutSize = useCombSoa ? (int64_t)TRANS_BLOCK_FULL * combAAlign * sizeof(float) : combInSize; + pipe->InitBuffer(combFragQue, 2, combOutSize); // TBuf pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); @@ -67,6 +75,11 @@ class HcPreSinkhornPerf { hcBase0Local = hcBaseBuf0.Get(); hcBase1Local = hcBaseBuf1.Get(); hcBase2Local = hcBaseBuf2.Get(); + + if (useCombSoa) { + pipe->InitBuffer(vaAddrBuf, 2 * TRANS_BLOCK_FULL * sizeof(uint64_t)); + vaAddrLocal = vaAddrBuf.Get(); + } } __aicore__ inline void Process() @@ -81,7 +94,11 @@ class HcPreSinkhornPerf { CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); - CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + if (useCombSoa) { + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, 1, tilingData->hcMult * tilingData->hcMult); + } else { + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + } event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); SetFlag(eventId); WaitFlag(eventId); @@ -142,23 +159,101 @@ class HcPreSinkhornPerf { // combFrag mixes2Local = mixesQue2.AllocTensor(); - CopyInWithLoopMode( - mixesGm - [mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + tilingData->hcMult * 2], - mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + if (useCombSoa) { + CopyInCombTransposed( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + + tilingData->hcMult * 2], + mixes2Local, curRowFactor, combAAlign, tilingData->hcMix, combPlanes); + } else { + CopyInWithLoopMode( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + + tilingData->hcMult * 2], + mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + } mixesQue2.EnQue(mixes2Local); mixes2Local = mixesQue2.DeQue(); combFragLocal = combFragQue.AllocTensor(); - VFProcessCombFragRLessVLUseFourUnfold( - combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, - tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + // Branch-S (hcMult<=4, enough rows): plane-major SoA, full lane occupancy + // Branch-A (hcMult==4): FourUnfold fast path (already register-resident) + // Branch-B (hcMult in {2,3,6,8,12,16}): RLessVL register-resident via template + // Branch-C (other): fallback to original RLessVL UB-staging path + if (useCombSoa) { + switch (tilingData->hcMult) { + case 2: + VFProcessCombFragSoA<2>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + case 3: + VFProcessCombFragSoA<3>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + default: + VFProcessCombFragSoA<4>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + } + TransposeCombSoAToAoS(combFragLocal, mixes2Local, vaAddrLocal, curRowFactor, combAAlign, combPlanes); + } else if (tilingData->hcMult == COMB_UNFOLD_NUM) { + VFProcessCombFragRLessVLUseFourUnfold( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + } else { + // [vec-05 + state_resident] Register-resident Sinkhorn iteration + switch (tilingData->hcMult) { + case 2: + VFProcessCombFragRegResident<2>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 3: + VFProcessCombFragRegResident<3>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 6: + VFProcessCombFragRegResident<6>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 8: + VFProcessCombFragRegResident<8>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 12: + VFProcessCombFragRegResident<12>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 16: + VFProcessCombFragRegResident<16>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + default: + VFProcessCombFragRLessVL( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + } + } mixesQue2.FreeTensor(mixes2Local); rsqrtQue.FreeTensor(rsqrtLocal); combFragQue.EnQue(combFragLocal); combFragLocal = combFragQue.DeQue(); - CopyOut(combFragLocal, combFragGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult], curRowFactor * tilingData->hcMult, tilingData->hcMult); + int64_t combGmOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult; + if (useCombSoa) { + CopyOutCombSoA(combFragLocal, combFragGm[combGmOffset], curRowFactor, combPlanes); + } else { + CopyOut( + combFragLocal, combFragGm[combGmOffset], curRowFactor * tilingData->hcMult, tilingData->hcMult); + } combFragQue.FreeTensor(combFragLocal); } } @@ -166,6 +261,9 @@ class HcPreSinkhornPerf { private: TPipe* pipe; const HcPreSinkhornTilingData* tilingData; + bool useCombSoa = false; + int32_t combPlanes = 0; + int32_t combAAlign = 0; GlobalTensor mixesGm; GlobalTensor rsqrtGm; GlobalTensor hcScaleGm; @@ -186,6 +284,7 @@ class HcPreSinkhornPerf { TBuf hcBaseBuf0; TBuf hcBaseBuf1; TBuf hcBaseBuf2; + TBuf vaAddrBuf; LocalTensor mixes01Local; LocalTensor mixes2Local; @@ -197,6 +296,7 @@ class HcPreSinkhornPerf { LocalTensor hcBase0Local; LocalTensor hcBase1Local; LocalTensor hcBase2Local; + LocalTensor vaAddrLocal; }; } // namespace HCPreSinkhorn From 91d108caff4d9f2b0e37ca8ee9a6ab1ae5a1abfe Mon Sep 17 00:00:00 2001 From: hanrui22 Date: Fri, 21 Aug 2026 18:46:07 +0800 Subject: [PATCH 2/3] delete doc --- .../docs/perf-integrate/BEFORE_AFTER_PERF.md | 47 +++++ .../docs/perf-integrate/SYNC_VERIFY_REPORT.md | 71 ------- .../perf-integrate/my_xllm_20cases_perf.json | 182 ------------------ .../docs/perf-integrate/profile_20cases.py | 114 ----------- .../docs/perf-integrate/run_20cases.py | 108 ----------- 5 files changed, 47 insertions(+), 475 deletions(-) create mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md delete mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md delete mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json delete mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py delete mode 100644 xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md new file mode 100644 index 0000000..af8ea51 --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md @@ -0,0 +1,47 @@ +# hc_pre_sinkhorn 20 case 优化前后性能 + +日期：2026-08-21 +口径：aclnn + torch_npu profiler，`HcPreSinkhorn` kernel 时间（µs） +优化前：Round 2 RegResident vendor（合入 SoA 前同机实测） +优化后：Round 3 SoA，源码在 `my-xllm-ops/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` + +## 结论 + +- 几何平均加速 **2.15x** +- 14/20 case 加速比 ≥ 1.10x；M≥6（9–12）未走 SoA，约 1.00x +- 最大：case 6（bs=16384, M=4）**204.48 → 12.88 µs，15.87x** +- 小 batch（1/2/18）基本持平；case 18 微回归 3.01 → 3.07 µs + +## 逐 case + +| case | bs | M | iters | 优化前 µs | 优化后 µs | 加速比 | 路径 | +|------|----|---|-------|-----------|-----------|--------|------| +| 1 | 1 | 4 | 20 | 3.155 | 3.142 | 1.00x | 小 batch | +| 2 | 64 | 4 | 20 | 4.602 | 4.546 | 1.01x | 小 batch | +| 3 | 512 | 4 | 20 | 10.780 | 5.512 | 1.96x | SoA | +| 4 | 1024 | 4 | 20 | 16.730 | 5.407 | 3.09x | SoA | +| 5 | 4096 | 4 | 20 | 53.710 | 6.978 | 7.70x | SoA | +| 6 | 16384 | 4 | 20 | 204.479 | 12.884 | **15.87x** | SoA | +| 7 | 1024 | 2 | 20 | 14.609 | 4.640 | 3.15x | SoA | +| 8 | 1024 | 3 | 20 | 16.120 | 5.323 | 3.03x | SoA | +| 9 | 1024 | 6 | 20 | 18.805 | 18.812 | 1.00x | AoS | +| 10 | 1024 | 8 | 20 | 20.854 | 20.786 | 1.00x | AoS | +| 11 | 1024 | 12 | 20 | 24.392 | 24.255 | 1.01x | AoS | +| 12 | 1024 | 16 | 20 | 29.433 | 29.423 | 1.00x | AoS | +| 13 | 1024 | 4 | 1 | 5.199 | 4.635 | 1.12x | SoA | +| 14 | 1024 | 4 | 5 | 6.498 | 4.521 | 1.44x | SoA | +| 15 | 1024 | 4 | 40 | 30.409 | 6.498 | 4.68x | SoA | +| 16 | 1024 | 4 | 20 | 17.030 | 5.760 | 2.96x | SoA | +| 17 | 1024 | 4 | 20 | 16.805 | 5.376 | 3.13x | SoA | +| 18 | 1 | 2 | 20 | 3.014 | 3.070 | 0.98x | 小 batch | +| 19 | 1024 | 4 | 20 | 16.809 | 5.524 | 3.04x | SoA | +| 20 | 1024 | 4 | 20 | 17.056 | 5.724 | 2.98x | SoA | +| **geomean** | | | | | | **2.15x** | | + +## 怎么读 + +- **SoA**（M≤4 且每核行数够）：CombFrag 改成 plane-major，行/列归一化变成逐元素运算，大 batch 和深迭代收益最大。 +- **AoS**（M≥6）：寄存器装不下 M² 个 plane，路径未改，时间不变。 +- **小 batch**：rowFactor<4 不走 SoA，时间在测量噪声内。 + +墙钟（含 aclnn 下发）约 260–400 µs，不能当 kernel 时间。 diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md deleted file mode 100644 index 1bb926f..0000000 --- a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/SYNC_VERIFY_REPORT.md +++ /dev/null @@ -1,71 +0,0 @@ -# hc_pre_sinkhorn 同步到 my-xllm-ops — 精度与性能验证 - -日期：2026-08-21 -来源：`/home/h00801112/codex/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` -目标：`/home/h00801112/codex/my-xllm-ops/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` - -## 1. 同步内容 - -| 文件 | 变更 | md5 | -|------|------|-----| -| `op_kernel/hc_pre_sinkhorn_regbase_base.h` | SoA CombFrag、MultiCopy 转置、ldva TransDataTo5HD + `S_V`/`V_S` | `f152f40d187453990e9c9c882aa38358`（两边一致） | -| `op_kernel/hc_pre_sinkhorn_regbase_perf.h` | `useCombSoa` 分发；`hcScale` 仍 `GetValue()` | `8754c1d245ccdc4f21a389d501e4ce6f`（两边一致） | -| `op_host/CMakeLists.txt` | opc 增加 `-Wno-constant-conversion` | 已同步 | - -未改：`hc_pre_sinkhorn.cpp`、tiling、def、proto、membase kernel。 - -## 2. 编译说明 - -`my-xllm-ops/xllm-ops` 全新 `build.sh` 会重新下载 `json` / `abseil`，gitcode 返回 403，cmake 失败。 -本次用 dest 源码经 opc 编出 `HcPreSinkhorn_20e241df….o`，装入 `${ASCEND_HOME_PATH}/opp/vendors/custom_xllm_math`。host/tiling/aclnn 与源仓一致，无需重编。 - -## 3. 精度 - -| 套件 | 结果 | -|------|------| -| `test/python_test/test_hc_pre_sinkhorn.py`（10 func） | **10/10 PASS** | -| cann-bench 20 case（y / post / comb） | **20/20 PASS** | - -阈值：y bf16 rtol/atol=1e-2；post fp32 1e-4；comb fp32 1e-3。 - -SoA 路径最大 abs：y ≤ 1.5e-5，post ≤ 2.4e-7，comb ≤ 2.1e-7。与源仓 xllm-ops 合入结果一致。 - -## 4. 性能（HcPreSinkhorn kernel µs） - -设备：Ascend950PR。口径：torch_npu profiler。 -对照：源仓合入实测、A5 `mhc_sinkhorn` baseline。 - -| case | bs | M | iter | my-xllm us | xllm-ops us | A5 us | vs A5 | -|------|----|---|------|------------|-------------|-------|-------| -| 1 | 1 | 4 | 20 | 3.142 | 3.138 | 5.36 | 1.71x | -| 2 | 64 | 4 | 20 | 4.546 | 4.650 | 5.82 | 1.28x | -| 3 | 512 | 4 | 20 | 5.512 | 5.505 | 8.47 | 1.54x | -| 4 | 1024 | 4 | 20 | 5.407 | 5.418 | 10.00 | 1.85x | -| 5 | 4096 | 4 | 20 | 6.978 | 6.944 | 10.00 | 1.43x | -| 6 | 16384 | 4 | 20 | 12.884 | 12.828 | 15.90 | 1.23x | -| 7 | 1024 | 2 | 20 | 4.640 | 4.653 | 10.00 | 2.16x | -| 8 | 1024 | 3 | 20 | 5.323 | 5.310 | 10.00 | 1.88x | -| 9 | 1024 | 6 | 20 | 18.812 | 18.768 | 10.00 | 0.53x | -| 10 | 1024 | 8 | 20 | 20.786 | 20.850 | 10.00 | 0.48x | -| 11 | 1024 | 12 | 20 | 24.255 | 24.337 | 10.00 | 0.41x | -| 12 | 1024 | 16 | 20 | 29.423 | 29.455 | 15.50 | 0.53x | -| 13 | 1024 | 4 | 1 | 4.635 | 4.646 | 4.17 | 0.90x | -| 14 | 1024 | 4 | 5 | 4.521 | 4.548 | 5.65 | 1.25x | -| 15 | 1024 | 4 | 40 | 6.498 | 6.578 | 10.00 | 1.54x | -| 16 | 1024 | 4 | 20 | 5.760 | 5.745 | 10.00 | 1.74x | -| 17 | 1024 | 4 | 20 | 5.376 | 5.376 | 10.00 | 1.86x | -| 18 | 1 | 2 | 20 | 3.070 | 3.053 | 10.00 | 3.26x | -| 19 | 1024 | 4 | 20 | 5.524 | 5.427 | 10.00 | 1.81x | -| 20 | 1024 | 4 | 20 | 5.724 | 5.763 | 10.00 | 1.75x | - -- 与源仓 xllm-ops：**数值对齐**（差 < 0.15 µs，测量抖动） -- vs A5：**15/20 更快**，大 batch M≤4 与源仓同档（case 6 ≈ 12.9 µs） -- M≥6 仍走 AoS，与 A5 比偏慢（算子还含 pre/y/post，A5 仅 Sinkhorn） - -墙钟（aclnn 下发）约 260–400 µs，不能当 kernel 时间。 - -原始数据：`docs/perf-integrate/my_xllm_20cases_perf.json` - -## 5. 结论 - -代码已同步到 my-xllm-ops，精度 10/10 + 20/20 通过，性能与源仓合入结果一致。 diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json deleted file mode 100644 index 2825ee2..0000000 --- a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json +++ /dev/null @@ -1,182 +0,0 @@ -[ - { - "case_id": 1, - "bs": 1, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 3.142, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 2, - "bs": 64, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 4.546, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 3, - "bs": 512, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.512, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 4, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.407, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 5, - "bs": 4096, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 6.978, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 6, - "bs": 16384, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 12.884, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 7, - "bs": 1024, - "M": 2, - "d": 2, - "iters": 20, - "kernel_us": 4.64, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 8, - "bs": 1024, - "M": 3, - "d": 3, - "iters": 20, - "kernel_us": 5.323, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 9, - "bs": 1024, - "M": 6, - "d": 6, - "iters": 20, - "kernel_us": 18.812, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 10, - "bs": 1024, - "M": 8, - "d": 8, - "iters": 20, - "kernel_us": 20.786, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 11, - "bs": 1024, - "M": 12, - "d": 12, - "iters": 20, - "kernel_us": 24.255, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 12, - "bs": 1024, - "M": 16, - "d": 16, - "iters": 20, - "kernel_us": 29.423, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 13, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 1, - "kernel_us": 4.635, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 14, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 5, - "kernel_us": 4.521, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 15, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 40, - "kernel_us": 6.498, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 16, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.76, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 17, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.376, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 18, - "bs": 1, - "M": 2, - "d": 2, - "iters": 20, - "kernel_us": 3.07, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 19, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.524, - "kernel": "HcPreSinkhorn" - }, - { - "case_id": 20, - "bs": 1024, - "M": 4, - "d": 4, - "iters": 20, - "kernel_us": 5.724, - "kernel": "HcPreSinkhorn" - } -] \ No newline at end of file diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py deleted file mode 100644 index 78c84e5..0000000 --- a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/profile_20cases.py +++ /dev/null @@ -1,114 +0,0 @@ -#!/usr/bin/env python3 -"""On-device kernel time for the 20 cann-bench cases via torch_npu profiler.""" -from __future__ import annotations - -import csv -import glob -import json -import shutil -import sys -from collections import defaultdict -from pathlib import Path - -import torch -import torch_npu - -ROOT = Path("/home/h00801112/codex/my-xllm-ops/xllm-ops") -sys.path.insert(0, str(ROOT / "test/python_test")) -import custom_ops # noqa: E402 - -CASES_CSV = Path("/home/h00801112/codex/bestline4/operators/hc_pre_sinkhorn/round2/cases.csv") -OUT = ROOT / "xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/my_xllm_20cases_perf.json" -WORK = Path("/tmp/hc_pre_sinkhorn_prof_my_xllm") - - -def parse_kernel_csv(work: Path) -> dict: - stats = defaultdict(lambda: [0, 0.0]) - candidates = glob.glob(str(work / "**" / "kernel_details.csv"), recursive=True) - if not candidates: - candidates = glob.glob(str(work / "**" / "op_statistic*.csv"), recursive=True) - for path in candidates: - with open(path, newline="", encoding="utf-8") as fh: - for row in csv.DictReader(fh): - name = (row.get("Name") or row.get("OP Type") or "").strip() - dur = row.get("Duration(us)") or row.get("Total Time(us)") or "0" - if not name: - continue - try: - val = float(dur) - except ValueError: - continue - stats[name][0] += 1 - stats[name][1] += val - return {k: {"count": v[0], "total_us": round(v[1], 3), - "avg_us": round(v[1] / max(v[0], 1), 3)} - for k, v in stats.items()} - - -def pick(stats: dict, needle: str): - for name, val in stats.items(): - if needle.lower() in name.lower(): - return {"kernel": name, **val} - return None - - -def profile_once(fn, tag: str, reps: int = 20) -> dict: - work = WORK / tag - if work.exists(): - shutil.rmtree(work) - work.mkdir(parents=True, exist_ok=True) - exp_cfg = torch_npu.profiler._ExperimentalConfig( - aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, - profiler_level=torch_npu.profiler.ProfilerLevel.Level1, - ) - for _ in range(5): - fn() - torch.npu.synchronize() - with torch_npu.profiler.profile( - activities=[torch_npu.profiler.ProfilerActivity.NPU], - experimental_config=exp_cfg, - on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(str(work)), - ) as prof: - for _ in range(reps): - fn() - prof.step() - torch.npu.synchronize() - return parse_kernel_csv(work) - - -def main(): - torch_npu.npu.set_device(0) - rows = list(csv.DictReader(CASES_CSV.open())) - results = [] - print(f"{'case':>4} {'bs':>6} {'M':>3} {'iter':>5} {'kernel_us':>10} {'kernel':>24}") - for row in rows: - cid = int(row["case_id"]) - bs, m, d = int(row["bs"]), int(row["hc_mult"]), int(row["d"]) - iters, eps = int(row["hc_sinkhorn_iters"]), float(row["hc_eps"]) - hc_mix = 2 * m + m * m - g = torch.Generator().manual_seed(2026 + cid) - mixes = torch.empty(bs, hc_mix).uniform_(-1, 1, generator=g).npu() - rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g).npu() - hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g).npu() - hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g).npu() - x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16).npu() - - def fn(): - return custom_ops.hc_pre_sinkhorn_npu( - mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps) - - stats = profile_once(fn, f"c{cid}", reps=20) - hit = pick(stats, "HcPreSinkhorn") or pick(stats, "hc_pre_sinkhorn") - us = hit["avg_us"] if hit else float("nan") - name = hit["kernel"] if hit else "n/a" - results.append({"case_id": cid, "bs": bs, "M": m, "d": d, "iters": iters, - "kernel_us": us, "kernel": name}) - print(f"{cid:>4} {bs:>6} {m:>3} {iters:>5} {us:>10} {name:>24}") - sys.stdout.flush() - - OUT.write_text(json.dumps(results, indent=2)) - print(f"\nwritten {OUT}") - - -if __name__ == "__main__": - main() diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py deleted file mode 100644 index d38d751..0000000 --- a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/run_20cases.py +++ /dev/null @@ -1,108 +0,0 @@ -#!/usr/bin/env python3 -"""Precision + wall-clock check of the 20 cann-bench cases via custom_ops.""" -from __future__ import annotations - -import csv -import sys -import time -from pathlib import Path - -import torch - -torch_npu = __import__("torch_npu") -ROOT = Path("/home/h00801112/codex/my-xllm-ops/xllm-ops") -sys.path.insert(0, str(ROOT / "test/python_test")) -import custom_ops # noqa: E402 -from test_hc_pre_sinkhorn import hc_pre_sinkhorn_golden # noqa: E402 - -CASES_CSV = Path("/home/h00801112/codex/bestline4/operators/hc_pre_sinkhorn/round2/cases.csv") - - -def make_inputs(row): - bs = int(row["bs"]) - m = int(row["hc_mult"]) - d = int(row["d"]) - iters = int(row["hc_sinkhorn_iters"]) - eps = float(row["hc_eps"]) - hc_mix = 2 * m + m * m - vr = row.get("value_range", "[-1,1]") - g = torch.Generator().manual_seed(2026 + int(row["case_id"])) - - if "nan" in vr: - mixes = torch.full((bs, hc_mix), float("nan")) - rsqrt = torch.full((bs,), float("nan")) - hc_scale = torch.full((3,), float("nan")) - hc_base = torch.full((hc_mix,), float("nan")) - x = torch.full((bs, m, d), float("nan"), dtype=torch.bfloat16) - elif "inf" in vr: - mixes = torch.empty(bs, hc_mix).uniform_(-2, 2, generator=g) - mixes.view(-1)[0] = float("inf") - mixes.view(-1)[1] = float("-inf") - rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g) - hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g) - hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g) - x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16) - else: - mixes = torch.empty(bs, hc_mix).uniform_(-1, 1, generator=g) - rsqrt = torch.empty(bs).uniform_(0.5, 1.0, generator=g) - hc_scale = torch.empty(3).uniform_(0.5, 1.0, generator=g) - hc_base = torch.empty(hc_mix).uniform_(-0.1, 0.1, generator=g) - x = torch.randn(bs, m, d, generator=g).to(torch.bfloat16) - return mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps - - -def max_err(a, b): - if not torch.isfinite(a).any() and not torch.isfinite(b).any(): - return 0.0, 0.0 - diff = (a.float() - b.float()).abs() - denom = b.float().abs().clamp_min(1e-8) - return float(diff.max()), float((diff / denom).max()) - - -def time_op(fn, warmup=10, iters=30): - for _ in range(warmup): - fn() - torch.npu.synchronize() - t0 = time.perf_counter() - for _ in range(iters): - fn() - torch.npu.synchronize() - return (time.perf_counter() - t0) * 1e6 / iters - - -def main(): - rows = list(csv.DictReader(CASES_CSV.open())) - print(f"{'case':>4} {'bs':>6} {'M':>3} {'d':>5} {'iter':>5} {'wall_us':>9} {'y':>8} {'post':>8} {'comb':>8} {'ok':>5}") - fails = 0 - recs = [] - for row in rows: - mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps = make_inputs(row) - y_g, post_g, comb_g = hc_pre_sinkhorn_golden( - mixes, rsqrt, hc_scale, hc_base, x, m, iters, eps) - - def launch(): - return custom_ops.hc_pre_sinkhorn_npu( - mixes.npu(), rsqrt.npu(), hc_scale.npu(), hc_base.npu(), x.npu(), - m, iters, eps) - - y_n, post_n, comb_n = launch() - ye, yr = max_err(y_n.cpu(), y_g) - pe, pr = max_err(post_n.cpu(), post_g) - ce, cr = max_err(comb_n.cpu(), comb_g) - y_ok = ye <= 1e-2 or yr <= 1e-2 or (not torch.isfinite(y_g).any()) - p_ok = pe <= 1e-4 or pr <= 1e-4 or (not torch.isfinite(post_g).any()) - c_ok = ce <= 1e-3 or cr <= 1e-3 or (not torch.isfinite(comb_g).any()) - ok = y_ok and p_ok and c_ok - fails += int(not ok) - us = time_op(launch) - print(f"{row['case_id']:>4} {row['bs']:>6} {m:>3} {row['d']:>5} {iters:>5} " - f"{us:>9.2f} {ye:>8.1e} {pe:>8.1e} {ce:>8.1e} {'PASS' if ok else 'FAIL':>5}") - recs.append({"case_id": int(row["case_id"]), "ok": ok, "wall_us": us, - "y_abs": ye, "post_abs": pe, "comb_abs": ce}) - sys.stdout.flush() - print(f"\n{len(rows) - fails}/{len(rows)} PASS") - return 0 if fails == 0 else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) From 3f3ea1feb21a60ffb174db681089785260119693 Mon Sep 17 00:00:00 2001 From: hanrui22 Date: Sat, 22 Aug 2026 00:31:17 +0800 Subject: [PATCH 3/3] upload dequant_swiglu_quant --- .../dequant_swiglu_quant_tiling_base.cpp | 65 +++ .../op_kernel/dequant_swiglu_quant.cpp | 24 + .../op_kernel/dequant_swiglu_quant_vecrow.hpp | 479 ++++++++++++++++++ .../tiling_base/error_log.h | 4 + 4 files changed, 572 insertions(+) create mode 100644 xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp b/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp index 612ff81..3b60b51 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp +++ b/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp @@ -62,6 +62,19 @@ constexpr uint32_t PERFORMANCE_ROW_LEN = 128; constexpr uint32_t MIN_CORE = 12; const int64_t DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE = 30013; +// Vectorised-row fast path (see op_kernel/dequant_swiglu_quant_vecrow.hpp). +// It removes the per-row V_S sync + GetValue/SetValue round-trip the other +// dynamic classes pay, which only dominates once there are many rows: measured +// at colLen=256, rows<=512 is a wash or a small loss (bf16 256 rows +31%) while +// rows>=1024 wins steadily (int32 2048 rows -32%, bf16 8192 rows -21%). +// The win also shrinks as colLen grows, because the path spills the fp32 SwiGLU +// result to UB and reads it back -- by colLen=1024 it is a wash. +const int64_t DYNAMIC_VECROW_INT32 = 30020; +const int64_t DYNAMIC_VECROW_FLOAT16 = 30021; +const int64_t DYNAMIC_VECROW_BFLOAT16 = 30022; +constexpr uint32_t VECROW_COL_LEN = 512; +constexpr uint32_t VECROW_MIN_ROW_LEN = 1024; + // Tiling优选参数 struct GluSingleTilingOptParam { // Maximum amount of data that can be transferred by an operator UB at a time. Unit:element @@ -146,6 +159,8 @@ class DequantSwigluQuantTiling : public TilingBaseClass { bool isPerformanceBranch(); + bool isVecRowBranch(); + int64_t getTilingKeyStatic( const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const; @@ -170,6 +185,7 @@ class DequantSwigluQuantTiling : public TilingBaseClass { ge::DataType xInputDataType; bool isPerfBranch = false; + bool isVecRow = false; ge::DataType biasDataType = ge::DT_FLOAT; uint64_t quantScaleShapeSize = 0; @@ -590,6 +606,7 @@ ge::graphStatus DequantSwigluQuantTiling::DoOpTiling() return ge::GRAPH_FAILED; } isPerfBranch = isPerformanceBranch(); + isVecRow = isVecRowBranch(); return ge::GRAPH_SUCCESS; } @@ -647,12 +664,18 @@ int64_t DequantSwigluQuantTiling::getTilingKeyDynamic( if (scaleSize == 1) { return DYNAMIC_FLOAT16_X; } else { + if (isVecRow) { + return DYNAMIC_VECROW_FLOAT16; + } return DYNAMIC_FLOAT16_XD; } } else { if (scaleSize == 1) { return DYNAMIC_BFLOAT16_X; } else { + if (isVecRow) { + return DYNAMIC_VECROW_BFLOAT16; + } return DYNAMIC_BFLOAT16_XD; } } @@ -671,6 +694,12 @@ int64_t DequantSwigluQuantTiling::getTilingKeyDynamic( if (biasType == ge::DT_INT32) { return DYNAMIC_INT_X_INT_BIAS_QUANT_D; } else if (biasType == ge::DT_FLOAT) { + // isVecRowBranch() checks the bias input directly, so the flag being + // set already means no bias is present (biasType defaults to + // DT_FLOAT when the optional input is absent). + if (isVecRow) { + return DYNAMIC_VECROW_INT32; + } if(isPerfBranch) { return DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE; } @@ -695,6 +724,42 @@ bool DequantSwigluQuantTiling::isPerformanceBranch() { return false; } +// The vecrow kernel handles one whole row per VF iteration and keeps the +// per-row max in a vector register, so it needs the row fully loadable and no +// bias / quant_offset / group_index to fold in. Beyond that it is gated on +// shape: colLen must be small enough that spilling the fp32 SwiGLU result to UB +// still pays, and there must be enough rows to amortise its setup. +bool DequantSwigluQuantTiling::isVecRowBranch() { + if (tilingData.get_is32BAligned() != 1) { + return false; + } + // biasIsEmpty is only populated on the int32 branch (checkWeightBiasActivate + // runs under `xDataType == DT_INT32`), so it reads 0 for bf16/fp16 whether or + // not a bias exists. Ask the context directly instead of trusting the field. + if (context_->GetOptionalInputShape(INDEX_IN_BIAS) != nullptr) { + return false; + } + // quant_offset is not folded in by this kernel. + if (context_->GetOptionalInputShape(INDEX_IN_QUANT_OFFSET) != nullptr) { + return false; + } + // Row must be fully loadable: the kernel indexes a row in one shot. + if (tilingData.get_baseColLen() != tilingData.get_colLen()) { + return false; + } + uint64_t colLen = tilingData.get_colLen(); + if (colLen > VECROW_COL_LEN) { + return false; + } + // Below this row count the fast path's fixed setup is not amortised and it + // loses to the row-wise kernel (measured at colLen=256: 256 rows +31%, + // 1024 rows -10%, 8192 rows -21%). + if (tilingData.get_rowLen() < VECROW_MIN_ROW_LEN) { + return false; + } + return true; +} + uint64_t DequantSwigluQuantTiling::GetTilingKey() const { if (quantMode == 0) { // static diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp index fb645f4..a52d391 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp +++ b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp @@ -26,6 +26,7 @@ #include "dequant_swiglu_quant_dynamic_bias_int32.hpp" #include "dequant_swiglu_quant_dynamic_bias_float.hpp" #include "dequant_swiglu_quant_dynamic_performance.hpp" +#include "dequant_swiglu_quant_vecrow.hpp" using namespace AscendC; @@ -58,6 +59,11 @@ using namespace AscendC; #define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_QS_GR 110000100 #define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_QS_GR 110000200 +// Vectorised-row fast path; keys assigned in dequant_swiglu_quant_tiling_base.cpp. +#define DEQUANT_SWIGLU_QUANT_VECROW_INT32 30020 +#define DEQUANT_SWIGLU_QUANT_VECROW_FLOAT16 30021 +#define DEQUANT_SWIGLU_QUANT_VECROW_BFLOAT16 30022 + extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR weightSscaleGM, GM_ADDR activationScaleGM, GM_ADDR biasGM, GM_ADDR quantScaleGM, GM_ADDR quantOffsetGM, @@ -312,6 +318,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_INT32)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) // ORIG_DTYPE_BIAS == DT_BF16 @@ -375,6 +387,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_FLOAT16)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #endif #if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) && (ORIG_DTYPE_X == DT_BF16) @@ -428,6 +446,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_BFLOAT16)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #endif } \ No newline at end of file diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp new file mode 100644 index 0000000..a4279bd --- /dev/null +++ b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp @@ -0,0 +1,479 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * Copyright 2026 The xLLM Authors. All Rights Reserved. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file dequant_swiglu_quant_vecrow.hpp + * \brief Vectorised-row fast path for dynamic per-token quant. + * + * The existing dynamic classes process one row at a time and, per row, pay a + * ReduceMax + V_S sync + GetValue/SetValue + Muls(scalar) round-trip (see + * DequantSwigluQuantDynamicBase::dynamicMultiColMax, which ends in + * maxTempLocal.GetValue(rowId)). With rows in the tens of thousands that scalar + * traffic dominates: three shapes moving identical bytes took 47/70/123us purely + * because their row counts differ. + * + * Here the whole dequant -> SwiGLU -> abs-max -> quant chain runs in MicroAPI + * inside two __VEC_SCOPE__ blocks per tile. The per-row max never leaves a + * vector register: ReduceMax lands it in lane 0 and a broadcast load feeds it + * straight back into the divide. Input tiles are double-buffered so MTE2 + * overlaps the VF work. + * + * Scope: dynamic quant, no bias / quant_offset / group_index. Everything else + * keeps using the original classes. + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP + +#include "kernel_operator.h" + +namespace DequantSwigluQuant { +using namespace AscendC; + +namespace VecRow { + +constexpr uint32_t UB_BLOCK_BYTES = 32; +// Row strides are padded to a whole vector register so every chunk load in the +// VF loops sits on a natural boundary. +constexpr uint32_t VREG_BYTES = 256; +constexpr uint32_t VL_FP32 = 64; +constexpr uint32_t MAX_TILE_ROWS = 64; +constexpr uint32_t UB_BUDGET_BYTES = 184 * 1024; + +// resolves to the round-only cast; CAST_NONE is the plain +// widening the original int32 path uses. +constexpr Reg::CastTrait castS32ToF32 = { + Reg::RegLayout::UNKNOWN, Reg::SatMode::UNKNOWN, Reg::MaskMergeMode::ZEROING, RoundMode::CAST_NONE}; +constexpr Reg::CastTrait castB16ToF32 = { + Reg::RegLayout::ZERO, Reg::SatMode::UNKNOWN, Reg::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; +// fp32 -> int8 goes fp32 -> s16 -> fp16 -> s8, the same chain adv_api's +// TransRegForS8 uses (quantize_impl.h:145). Casting straight to int8 rounds and +// packs lanes wrongly -- it yields a correct scale but a garbage y. +constexpr Reg::CastTrait castF32ToS16 = { + Reg::RegLayout::ZERO, Reg::SatMode::SAT, Reg::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + +__aicore__ inline uint32_t CeilDivU(uint32_t a, uint32_t b) +{ + return (b == 0) ? 0 : (a + b - 1) / b; +} + +// Pad `count` elements of `sz` bytes up to a whole `toBytes` boundary. +__aicore__ inline uint32_t AlignElTo(uint32_t count, uint32_t sz, uint32_t toBytes) +{ + uint32_t bytes = count * sz; + return (bytes + toBytes - 1) / toBytes * toBytes / sz; +} + +} // namespace VecRow + +/** + * InType: int32_t (dequant path) or half/bfloat16_t (direct path). + * hasWeightScale: int32 path multiplies by weight_scale[2H]. + */ +template +class DequantSwigluQuantVecRow { +public: + __aicore__ inline DequantSwigluQuantVecRow() {} + __aicore__ inline ~DequantSwigluQuantVecRow() {} + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, + GM_ADDR quant_scale_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_) + { + using namespace VecRow; + pipe = pipe_; + curBlockIdx = GetBlockIdx(); + + H = static_cast(tilingData->colLen); + rowNum = static_cast(tilingData->rowLen); + actLeft = (tilingData->activateLeft != 0); + hasQs = (tilingData->quantScaleIsEmpty == 0); + hasActScale = (tilingData->activateScaleIsEmpty == 0); + + useCoreNum = tilingData->usedCoreNum; + if (rowNum < useCoreNum) { + useCoreNum = rowNum; + } + + // Same even split the other classes use, so core assignment matches. + uint32_t perRoundCnt = useCoreNum == 0 ? 0 : rowNum / useCoreNum; + uint32_t remainCnt = rowNum - useCoreNum * perRoundCnt; + numRound = perRoundCnt; + if (curBlockIdx < remainCnt) { + numRound = perRoundCnt + 1; + biasOffset = curBlockIdx * (perRoundCnt + 1); + } else { + biasOffset = (perRoundCnt + 1) * remainCnt + (curBlockIdx - remainCnt) * perRoundCnt; + } + + xGm.SetGlobalBuffer((__gm__ InType*)x_gm); + yGm.SetGlobalBuffer((__gm__ int8_t*)y_gm); + scaleGm.SetGlobalBuffer((__gm__ float*)scale_gm); + if constexpr (isInt32) { + weightScaleGm.SetGlobalBuffer((__gm__ float*)weight_scale_gm); + if (hasActScale) { + actScaleGm.SetGlobalBuffer((__gm__ float*)activation_scale_gm); + } + } + if (hasQs) { + quantScaleGm.SetGlobalBuffer((__gm__ float*)quant_scale_gm); + } + + strideIn = AlignElTo(H, sizeof(InType), VREG_BYTES); + strideF = AlignElTo(H, sizeof(float), VREG_BYTES); + strideI8 = AlignElTo(H, 1, VREG_BYTES); + + uint32_t perRowBytes = 2 * strideIn * sizeof(InType) + strideF * sizeof(float) + strideI8; + uint32_t sharedBytes = 2 * strideF * sizeof(float) + (hasQs ? strideF * sizeof(float) : 0) + + 3 * MAX_TILE_ROWS * sizeof(float); + uint32_t avail = (UB_BUDGET_BYTES > sharedBytes) ? (UB_BUDGET_BYTES - sharedBytes) : 0; + tileRows = (perRowBytes > 0) ? (avail / perRowBytes) : 1; + if (tileRows > MAX_TILE_ROWS) tileRows = MAX_TILE_ROWS; + if (tileRows > numRound) tileRows = numRound; + if (tileRows < 1) tileRows = 1; + // Two input sets so the next tile's MTE2 overlaps this tile's VEC work. + if (tileRows > 1) tileRows /= 2; + + pipe->InitBuffer(bufA, 2 * tileRows * strideIn * sizeof(InType)); + pipe->InitBuffer(bufB, 2 * tileRows * strideIn * sizeof(InType)); + pipe->InitBuffer(bufSwi, tileRows * strideF * sizeof(float)); + pipe->InitBuffer(bufY, tileRows * strideI8); + pipe->InitBuffer(bufMax, MAX_TILE_ROWS * sizeof(float)); + pipe->InitBuffer(bufScale, MAX_TILE_ROWS * sizeof(float)); + pipe->InitBuffer(bufWsA, strideF * sizeof(float)); + pipe->InitBuffer(bufWsB, strideF * sizeof(float)); + pipe->InitBuffer(bufAct, 2 * MAX_TILE_ROWS * sizeof(float)); + if (hasQs) { + pipe->InitBuffer(bufQs, strideF * sizeof(float)); + } + } + + __aicore__ inline void Process() + { + using namespace VecRow; + if (curBlockIdx >= useCoreNum || numRound == 0) { + return; + } + + LocalTensor aLocal = bufA.template Get(); + LocalTensor bLocal = bufB.template Get(); + LocalTensor swiLocal = bufSwi.Get(); + LocalTensor yLocal = bufY.Get(); + LocalTensor maxLocal = bufMax.Get(); + LocalTensor scaleLocal = bufScale.Get(); + LocalTensor wsALocal = bufWsA.Get(); + LocalTensor wsBLocal = bufWsB.Get(); + LocalTensor actLocal = bufAct.Get(); + + // weight_scale / quant_scale are row-invariant: load once per core. + if constexpr (isInt32) { + DataCopyExtParams cpWs{1, static_cast(H * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams ppWs{false, 0, 0, 0}; + DataCopyPad(wsALocal, weightScaleGm[0], cpWs, ppWs); + DataCopyPad(wsBLocal, weightScaleGm[H], cpWs, ppWs); + } + if (hasQs) { + DataCopyExtParams cpQs{1, static_cast(H * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams ppQs{false, 0, 0, 0}; + DataCopyPad(bufQs.Get(), quantScaleGm[0], cpQs, ppQs); + } + if (isInt32 || hasQs) { + event_t eid = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eid); + WaitFlag(eid); + } + + __local_mem__ InType* aAddr = (__local_mem__ InType*)aLocal.GetPhyAddr(); + __local_mem__ InType* bAddr = (__local_mem__ InType*)bLocal.GetPhyAddr(); + __local_mem__ float* swiAddr = (__local_mem__ float*)swiLocal.GetPhyAddr(); + __local_mem__ int8_t* yAddr = (__local_mem__ int8_t*)yLocal.GetPhyAddr(); + __local_mem__ float* maxAddr = (__local_mem__ float*)maxLocal.GetPhyAddr(); + __local_mem__ float* scaleAddr = (__local_mem__ float*)scaleLocal.GetPhyAddr(); + __local_mem__ float* wsAAddr = (__local_mem__ float*)wsALocal.GetPhyAddr(); + __local_mem__ float* wsBAddr = (__local_mem__ float*)wsBLocal.GetPhyAddr(); + __local_mem__ float* actAddr = (__local_mem__ float*)actLocal.GetPhyAddr(); + __local_mem__ float* qsAddr = hasQs + ? (__local_mem__ float*)bufQs.Get().GetPhyAddr() : nullptr; + + // activate_left picks which half drives the sigmoid. Matches + // DequantSwigluQuantDynamicBase::BaseProcess: activateLeft==0 -> gate is + // the second half. Resolved here so the VF loops stay branch-free. + __local_mem__ InType* gateAddr = actLeft ? aAddr : bAddr; + __local_mem__ InType* upAddr = actLeft ? bAddr : aAddr; + __local_mem__ float* wsGateAddr = actLeft ? wsAAddr : wsBAddr; + __local_mem__ float* wsUpAddr = actLeft ? wsBAddr : wsAAddr; + + uint32_t dstStrideIn = (strideIn - H) * sizeof(InType) / UB_BLOCK_BYTES; + uint32_t dstStrideY = (strideI8 - H) / UB_BLOCK_BYTES; + uint32_t inHalf = tileRows * strideIn; + uint32_t actHalf = VecRow::MAX_TILE_ROWS; + + DataCopyPadExtParams ppIn{false, 0, 0, 0}; + DataCopyPadExtParams ppAct{false, 0, 0, 0}; + + event_t eidM2V[2] = { + static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)), + static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V))}; + event_t eidM3M2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + event_t eidV2M3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + +#define DSQ_VECROW_ISSUE_LOAD(rIdx, halfIdx) \ + do { \ + uint32_t _rows = tileRows; \ + if (_rows > numRound - (rIdx)) _rows = numRound - (rIdx); \ + uint32_t _start = biasOffset + (rIdx); \ + DataCopyExtParams _cpIn{static_cast(_rows), \ + static_cast(H * sizeof(InType)), \ + static_cast(H * sizeof(InType)), \ + static_cast(dstStrideIn), 0}; \ + DataCopyPad(aLocal[(halfIdx) * inHalf], xGm[_start * 2 * H], _cpIn, ppIn); \ + DataCopyPad(bLocal[(halfIdx) * inHalf], xGm[_start * 2 * H + H], _cpIn, ppIn); \ + if constexpr (isInt32) { \ + if (hasActScale) { \ + DataCopyExtParams _cpAct{1, static_cast(_rows * sizeof(float)), \ + 0, 0, 0}; \ + DataCopyPad(actLocal[(halfIdx) * actHalf], actScaleGm[_start], _cpAct, \ + ppAct); \ + } \ + } \ + } while (0) + + DSQ_VECROW_ISSUE_LOAD(0u, 0u); + SetFlag(eidM2V[0]); + + uint32_t half = 0; + for (uint32_t r = 0; r < numRound; r += tileRows) { + uint32_t curRows = tileRows; + if (curRows > numRound - r) curRows = numRound - r; + uint32_t rowStart = biasOffset + r; + uint32_t nextR = r + tileRows; + uint32_t nextHalf = half ^ 1; + + // Start the next tile's load before touching this one, so MTE2 runs + // underneath the VF work below. The other half is free: its compute + // finished an iteration ago. + if (nextR < numRound) { + DSQ_VECROW_ISSUE_LOAD(nextR, nextHalf); + SetFlag(eidM2V[nextHalf]); + } + + WaitFlag(eidM2V[half]); + + ComputeSwigluAndMax(gateAddr + half * inHalf, upAddr + half * inHalf, + wsGateAddr, wsUpAddr, qsAddr, actAddr + half * actHalf, + swiAddr, maxAddr, curRows); + + QuantizeRows(swiAddr, maxAddr, scaleAddr, yAddr, curRows); + + SetFlag(eidV2M3); + WaitFlag(eidV2M3); + + DataCopyExtParams cpY{static_cast(curRows), static_cast(H), + dstStrideY, 0, 0}; + DataCopyPad(yGm[rowStart * H], yLocal, cpY); + + DataCopyExtParams cpScale{1, static_cast(curRows * sizeof(float)), 0, 0, 0}; + DataCopyPad(scaleGm[rowStart], scaleLocal, cpScale); + + // swi/y/scale are single-buffered: the next iteration's VF writes + // must not start before these stores drain. + SetFlag(eidM3M2); + WaitFlag(eidM3M2); + + half = nextHalf; + } +#undef DSQ_VECROW_ISSUE_LOAD + } + +private: + /** + * Pass 1: dequant + SwiGLU for every row in the tile, spilling the fp32 + * result to `swi` and leaving each row's abs-max in maxBuf[i]. + * + * The abs-max accumulation deliberately uses a full mask while the Abs that + * feeds it uses the tail mask: ZEROING zeroes the inactive lanes of the Abs + * result, and 0 is the identity for a max over absolute values. Masking the + * Max itself would instead zero the lanes already holding the running max + * from earlier chunks. + */ + __aicore__ inline void ComputeSwigluAndMax( + __local_mem__ InType* gateIn, __local_mem__ InType* upIn, + __local_mem__ float* wsGate, __local_mem__ float* wsUp, + __local_mem__ float* qs, __local_mem__ float* actBuf, + __local_mem__ float* swi, __local_mem__ float* maxBuf, uint32_t rows) + { + using namespace VecRow; + const uint16_t chunks = static_cast(CeilDivU(H, VL_FP32)); + const uint16_t rowCnt = static_cast(rows); + const uint32_t hLocal = H; + const uint32_t strideInLocal = strideIn; + const uint32_t strideFLocal = strideF; + const bool hasQsLocal = hasQs; + const bool hasActLocal = hasActScale; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vGate, vUp, vScale, vTmp, vOne, vSig, vMax, vAbs; + MicroAPI::RegTensor vRaw; + MicroAPI::MaskReg m; + MicroAPI::MaskReg mAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg mOne = MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < rowCnt; i++) { + uint32_t sreg = hLocal; + MicroAPI::Duplicate(vMax, 0.0f); + MicroAPI::Duplicate(vOne, 1.0f, mAll); + + if constexpr (isInt32) { + // activation_scale is per row: one broadcast load feeds all + // chunks, replacing a per-row Muls(scalar). + if (hasActLocal) { + MicroAPI::DataCopy(vScale, actBuf + i); + } else { + MicroAPI::Duplicate(vScale, 1.0f, mAll); + } + } + + for (uint16_t c = 0; c < chunks; c++) { + m = MicroAPI::UpdateMask(sreg); + uint32_t inOff = i * strideInLocal + c * VL_FP32; + uint32_t fOff = i * strideFLocal + c * VL_FP32; + + if constexpr (isInt32) { + MicroAPI::DataCopy(vRaw, gateIn + inOff); + MicroAPI::Cast(vGate, vRaw, m); + MicroAPI::DataCopy(vTmp, wsGate + c * VL_FP32); + MicroAPI::Mul(vGate, vGate, vTmp, m); + MicroAPI::Mul(vGate, vGate, vScale, m); + + MicroAPI::DataCopy(vRaw, upIn + inOff); + MicroAPI::Cast(vUp, vRaw, m); + MicroAPI::DataCopy(vTmp, wsUp + c * VL_FP32); + MicroAPI::Mul(vUp, vUp, vTmp, m); + MicroAPI::Mul(vUp, vUp, vScale, m); + } else { + MicroAPI::DataCopy(vRaw, gateIn + inOff); + MicroAPI::Cast(vGate, vRaw, m); + MicroAPI::DataCopy(vRaw, upIn + inOff); + MicroAPI::Cast(vUp, vRaw, m); + } + + // SiLU(gate) = gate / (1 + exp(-gate)) -- the same op + // sequence adv_api's Sigmoid uses on this arch, so results + // match the original path bit for bit. + MicroAPI::Muls(vTmp, vGate, -1.0f, m); + MicroAPI::Exp(vTmp, vTmp, m); + MicroAPI::Adds(vTmp, vTmp, 1.0f, m); + MicroAPI::Div(vSig, vOne, vTmp, m); + MicroAPI::Mul(vGate, vGate, vSig, m); + + MicroAPI::Mul(vGate, vGate, vUp, m); + + if (hasQsLocal) { + MicroAPI::DataCopy(vTmp, qs + c * VL_FP32); + MicroAPI::Mul(vGate, vGate, vTmp, m); + } + + MicroAPI::DataCopy(swi + fOff, vGate, m); + + MicroAPI::Abs(vAbs, vGate, m); + MicroAPI::Max(vMax, vMax, vAbs, mAll); + } + + MicroAPI::ReduceMax(vMax, vMax, mAll); + MicroAPI::DataCopy(maxBuf + i, vMax, mOne); + } + } + } + + /** + * Pass 2: turn each row's max into its scale and quantise. The max is pulled + * back as a broadcast vector load, so scale and 1/scale are computed in-lane + * and the divide never touches the scalar unit. + */ + __aicore__ inline void QuantizeRows( + __local_mem__ float* swi, __local_mem__ float* maxBuf, + __local_mem__ float* scaleBuf, __local_mem__ int8_t* yOut, uint32_t rows) + { + using namespace VecRow; + const uint16_t chunks = static_cast(CeilDivU(H, VL_FP32)); + const uint16_t rowCnt = static_cast(rows); + const uint32_t hLocal = H; + const uint32_t strideFLocal = strideF; + const uint32_t strideI8Local = strideI8; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vVal, vScale, vInv, vOne; + MicroAPI::RegTensor vHalf; + MicroAPI::RegTensor vQ; + MicroAPI::MaskReg m; + MicroAPI::MaskReg mAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg mOne = MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < rowCnt; i++) { + uint32_t sreg = hLocal; + + MicroAPI::DataCopy(vScale, maxBuf + i); + MicroAPI::Muls(vScale, vScale, 1.0f / 127.0f, mAll); + MicroAPI::Maxs(vScale, vScale, 1e-12f, mAll); + MicroAPI::DataCopy(scaleBuf + i, vScale, mOne); + + MicroAPI::Duplicate(vOne, 1.0f, mAll); + MicroAPI::Div(vInv, vOne, vScale, mAll); + + for (uint16_t c = 0; c < chunks; c++) { + m = MicroAPI::UpdateMask(sreg); + MicroAPI::DataCopy( + vVal, swi + i * strideFLocal + c * VL_FP32); + MicroAPI::Mul(vVal, vVal, vInv, m); + // SAT on the s16 step clamps to [-128,127] once the value + // lands in int8, so the explicit clamp is free. + MicroAPI::Cast((MicroAPI::RegTensor&)vHalf, vVal, m); + MicroAPI::Cast(vHalf, (MicroAPI::RegTensor&)vHalf, m); + MicroAPI::Cast(vQ, vHalf, m); + MicroAPI::DataCopy( + yOut + i * strideI8Local + c * VL_FP32, vQ, m); + } + } + } + } + +private: + TPipe* pipe = nullptr; + uint32_t curBlockIdx = 0; + uint32_t H = 0; + uint32_t rowNum = 0; + uint32_t useCoreNum = 0; + uint32_t numRound = 0; + uint32_t biasOffset = 0; + uint32_t tileRows = 1; + uint32_t strideIn = 0; + uint32_t strideF = 0; + uint32_t strideI8 = 0; + bool actLeft = false; + bool hasQs = false; + bool hasActScale = false; + + GlobalTensor xGm; + GlobalTensor weightScaleGm; + GlobalTensor actScaleGm; + GlobalTensor quantScaleGm; + GlobalTensor yGm; + GlobalTensor scaleGm; + + TBuf bufA, bufB, bufSwi, bufY; + TBuf bufWsA, bufWsB, bufQs, bufAct, bufMax, bufScale; +}; + +} // namespace DequantSwigluQuant + +#endif // CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP diff --git a/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h b/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h index 1b09d20..9e2b4e4 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h +++ b/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h @@ -10,6 +10,10 @@ #pragma once +// tiling_base.h calls CheckLogLevel(static_cast(OP), ...); the OP +// module id lives in base/log_types.h, which slog.h pulls in. log/log.h alone +// does not, so include it here as moe_gating_top_k/tiling_base/error_log.h does. +#include "toolchain/slog.h" #include "log/log.h" #ifndef OP_LOGE_FOR_INVALID_DTYPE