Skip to content

Perf: 优化A5双Die任务局部性与SPMD派发 - #9

Open
yanghaoran29 wants to merge 1 commit into
mainfrom
perf/a5-die-locality-spmd-final-20260827
Open

Perf: 优化A5双Die任务局部性与SPMD派发#9
yanghaoran29 wants to merge 1 commit into
mainfrom
perf/a5-die-locality-spmd-final-20260827

Conversation

@yanghaoran29

Copy link
Copy Markdown
Owner

背景

A5由两个Die组成。Main按模4分配cluster,任务通过共享ready queue动态派发;这会让Scheduler访问远端AICore,并使细粒度producer/consumer依赖在两个Die间迁移。

本PR让任务domain、Scheduler负责的物理cluster和依赖链局部性保持一致,降低跨Die完成检测、依赖释放和再次派发开销。

方案

  • 按连续物理cluster分配Scheduler:S0/S1负责Die0,S2/S3负责Die1。
  • 为任务显式标记GLOBALDIE0DIE1,并为AIC/AIV/MIX维护对应ready queue。
  • completion、async-wait和deferred-release传播任务domain;同Die producer/consumer可通过Scheduler-local continuation直接续接。
  • 增加DIE_AFFINEAUTO_DIE_AFFINE scope,分别保持显式依赖和TensorMap自动依赖语义,同时固定局部依赖链所在Die。
  • BGEMM按GEMM+ADD group分Die;PA unroll/manual PA按batch链分Die;Batch PA按chunk分Die;Alternating根任务在两侧均分。
  • Qwen3让Gate/Up、SiLU和Down按MLP分片保持同Die。非sync-start多核SPMD拆成两个逻辑block区间,分别提交到两个Die;24核sync-start attention保持GLOBAL。具有85路fanin的dcr_xgamma保持单个5核任务,避免复制fanin和额外汇合。

性能

设备为Ascend950PR_9579卡1,CANN 9.2.0。非Qwen每个case运行100轮;Qwen运行5轮,按Device删除最快和最慢轮次后取中间3轮。Main使用同卡历史数据,基线commit为3b578e30a2a9e2859d19908b2647393ffdecc543。负数表示本PR更快。

Case Main Device 本PR Device 变化
alternating_matmul_add Case1 1,441.850 us 1,314.125 us -8.858%
benchmark_bgemm Case0 1,458.300 us 1,398.427 us -4.106%
paged_attention_unroll Case1 1,558.700 us 1,540.087 us -1.194%
paged_attention_unroll Case2 901.100 us 881.411 us -2.185%
paged_attention_unroll_manual_scope Case1 1,629.500 us 1,496.358 us -8.171%
paged_attention_unroll_manual_scope Case2 925.550 us 876.091 us -5.344%
batch_paged_attention Case1 7,434.600 us 5,444.691 us -26.766%
qwen3_14b_decode StressBatch16Seq3500 36,083.248 us 35,081.800 us -2.775%
8 case几何平均 -7.787%

8个case均优于历史Main,其中7个改善超过2%。完整数据见outputs/pr-final-all-swimlanes-card1-20260827/comparison.mdsummary.csv

验证

  • 全部8个benchmark运行成功。
  • Qwen拆分SPMD版本完成一次完整golden,结果PASS。
  • 当前PR代码在卡1为8个case各采集一次4级泳道,状态8/8 PASS。
  • 最新泳道位于outputs/pr-final-all-swimlanes-card1-20260827/<case>/merged_swimlane.json,可直接导入Perfetto;同目录保存kernel名称映射。
  • 泳道存在观察者效应,因此性能表使用关闭泳道的多轮benchmark,泳道仅用于检查任务推进和Die分布。

已知限制

  • Main为同卡历史数据,不是同一时间窗口内的交叉A/B。
  • 任务domain需要由样例选择合适粒度;标注不完整可能造成单Die空闲,粒度过细则会增加提交和汇合成本。
  • 多ready queue、domain传播和local-continuation判断会增加固定指令开销。
  • Qwen的dcr_xgamma具有85路fanin,不能像普通SPMD一样直接对半拆分;当前保留为GLOBAL单任务。

基于连续物理cluster分核,为任务引入GLOBAL/DIE0/DIE1 domain队列,并在完成、依赖释放和本地续接路径中保持同Die局部性。为BGEMM与PA系列按依赖链分配Die;Qwen按MLP分片分Die,并将非sync-start多核SPMD拆分到双Die,保留sync-start attention及高fanin dcr_xgamma单任务。

卡1相对历史Main 3b578e3:8个case的Device latency全部改善,几何平均改善7.787%;逐case结果、限制说明以及最终代码生成的8张四级泳道随提交保存。
@yanghaoran29
yanghaoran29 force-pushed the perf/a5-die-locality-spmd-final-20260827 branch 2 times, most recently from 629c7cd to c2b91ca Compare August 27, 2026 08:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant