[Performance] Benchmark high-perf Paged Attention (#899) vs PA unroll on a2a3

[ChaoZheng109]

Platform

a2a3 (Ascend 910B/C hardware)

Runtime Variant

tensormap_and_ringbuffer

Summary

#899 added a high-performance Paged Attention variant
(spmd_paged_attention_highperf). We need to benchmark it against the existing PA unroll
variant (paged_attention_unroll) on a2a3 and determine which is faster (and under which
shapes), so we can decide which to keep / recommend.

Git Commit ID

f257850

Host Platform

Linux (aarch64)

Reproduction

Run both on the same locked device(s) and compare. (Wrap onboard runs in task-submit; gate
the arch with onboard-arch-precheck.)

.claude/skills/onboard-arch-precheck/check.sh a2a3 || exit 1

# High-perf PA (#899)
task-submit --device auto --device-num 1 --run "\
  python tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention_highperf/test_spmd_paged_attention_highperf.py \
    -p a2a3 -d \$TASK_DEVICE --rounds 10 --skip-golden"

# PA unroll (baseline to compare)
task-submit --device auto --device-num 1 --run "\
  python tests/st/a2a3/tensormap_and_ringbuffer/paged_attention_unroll/test_paged_attention_unroll.py \
    -p a2a3 -d \$TASK_DEVICE --rounds 10 --skip-golden"

Expected Performance

Open question — this issue is to produce the numbers. Acceptance: a side-by-side latency
comparison of spmd_paged_attention_highperf vs paged_attention_unroll across the shared
shapes, with a clear verdict on which wins (overall and per-shape if it varies).

Actual Performance

Not yet measured.

Additional Context

• High-perf PA: tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention_highperf/ (from High performance Paged Attention A2A3 ST Test #899)
• PA unroll: tests/st/a2a3/tensormap_and_ringbuffer/paged_attention_unroll/
  (also paged_attention_unroll_4dims)
• Related: [Performance] Runtime performance optimization tracking #545 (runtime performance optimization tracking).

[ChaoZheng109]

Add a third comparison item: spmd_paged_attention

The benchmark should also include the plain SPMD Paged Attention variant
(spmd_paged_attention) as a comparison item, not just the high-perf variant
and PA unroll. This gives the full picture: baseline SPMD PA → high-perf SPMD PA
→ PA unroll.

Item: tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/

Reproduction (same locked device, wrap in task-submit, gate the arch):

.claude/skills/onboard-arch-precheck/check.sh a2a3 || exit 1

# SPMD Paged Attention (baseline SPMD variant) — NEW comparison item
task-submit --device auto --device-num 1 --run "\
  python tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py \
    -p a2a3 -d \$TASK_DEVICE --rounds 10 --skip-golden"

Updated acceptance: a side-by-side latency comparison of the three variants
— spmd_paged_attention vs spmd_paged_attention_highperf vs
paged_attention_unroll — across the shared shapes, with a clear per-variant
(and per-shape if it varies) verdict.

[TaoZQY]

Paged-Attention Benchmark

Date: 2026-06-25 23:10:29 CST

Environment

item	value
Platform	a2a3
Benchmark device	card 1
NPU	Ascend910, 8 cards visible, 2 chips/card
NPU software	npu-smi 26.0.rc1
Device health	OK
CANN	/usr/local/Ascend/cann-9.0.0, toolkit/runtime/compiler 9.0.0
CANN innerversion	V100R001C10SPC001B250
PTO-ISA	b0304ed37a783167dc8fdeb5fd844edce166efb5 (b0304ed3)
PTO-ISA package	set_package(pto-isa VERSION "9.0.0")
Host	liteserver-hps-148e-00001, Huawei Cloud EulerOS 2.0 aarch64, 320 HiSilicon CPUs

Run Command

The benchmark was run through task-submit with --rounds 10 --skip-golden.

task-submit --device 1 --max-time 2400 --run "$(cat <<'RUN'
cd /data/pyptouser/zhangtao/zt/simpler
source /usr/local/Ascend/cann-9.0.0/set_env.sh
eval "$(pypto-setup --export)"
export PTO_ISA_ROOT=$PWD/build/pto-isa

echo "TASK_DEVICE=$TASK_DEVICE"
echo "PTO_ISA_HEAD=$(git -C "$PTO_ISA_ROOT" rev-parse --short HEAD)"
cat build/lib/pto_isa_build.json

COMMON_ARGS="-p a2a3 -d $TASK_DEVICE --rounds 10 --manual include --skip-golden --enable-device-log-timing"

python tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py $COMMON_ARGS
python tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention_highperf/test_spmd_paged_attention_highperf.py $COMMON_ARGS
python tests/st/a2a3/tensormap_and_ringbuffer/paged_attention_unroll/test_paged_attention_unroll.py $COMMON_ARGS
RUN
)"

Note: the three existing test files do not define identical built-in cases.
spmd_paged_attention and paged_attention_unroll have matching BF16 production
cases; spmd_paged_attention_highperf has FP16 decode cases. For the all-three
side-by-side table below, the same classes were run once more with temporary
same-shape BF16 case overrides.

Same-Shape Side-by-Side

Cells are Avg Host us / Avg Device us. Winner and speedup use Avg Device.

shape	spmd_paged_attention	spmd_paged_attention_highperf	paged_attention_unroll	winner	speedup
b1_h16_kv1_s128_bs128_hd128_bf16	3814.5 / 1285.7	11262.2 / 1318.5	4080.4 / 1240.7	paged_attention_unroll	1.04x
b1_h16_kv1_s512_bs128_hd128_bf16	4243.7 / 1237.4	14979.6 / 1249.7	4826.8 / 1278.1	spmd_paged_attention	1.01x

Original Test-Case Summary

group	result
Matching BF16 production cases (spmd vs unroll)	paged_attention_unroll wins both comparable shapes by Avg Device, about 1.02x faster
Highperf FP16 decode cases	all PASS; Avg Device range 1163.1-1385.8 us; not directly comparable with the BF16 production cases
Stability note	spmd_paged_attention passes on PTO-ISA b0304ed3; newer PTO-ISA revisions reproduced AICore stream-sync timeout 507018 on this machine

Conclusion

For #899, spmd_paged_attention_highperf is not a clear replacement for
paged_attention_unroll on a2a3.

On the same-shape BF16 side-by-side run, the result is shape-dependent and the
gap is small:

• s128: paged_attention_unroll is faster than highperf by Avg Device
  (1240.7 us vs 1318.5 us, about 1.06x).
• s512: highperf is slightly faster than paged_attention_unroll by Avg
  Device (1249.7 us vs 1278.1 us, about 1.02x), but the margin is small
  and highperf has much higher Avg Host time.

For the built-in test cases, paged_attention_unroll remains the safer
recommendation for the existing BF16 production PA shapes. The highperf variant
passes its FP16 decode coverage, but those built-in highperf shapes are not the
same as the BF16 production shapes used by paged_attention_unroll, so they do
not prove an overall win.

Recommendation: keep paged_attention_unroll as the default/recommended PA
implementation for now. Keep spmd_paged_attention_highperf only as a
shape-specific experimental/high-perf path until it shows a consistent
same-shape win over paged_attention_unroll on the production shapes that
matter.

[ChaoZheng109]

参数不同，算子根据参数特化，待调整用例参数后再进行测试