[Profiling Report] Quantifying the pypto ↔ CCE paged-attention gap (follow-up to the 6× in simpler/pull: #986)

[noabauma]

Summary

At the matched qwen3 attention shape, the remaining pypto-vs-CCE gap is ~2.3x
(on-core attention time) — down from the ~6× in #986.

Comparison

Like-for-like at the qwen3-14B attention shape — batch 16, 40 q-heads, 8 KV
(5:1 GQA), head_dim 128, context 4096, block 128, fp16:

• CCE: spmd_paged_attention_highperf case b16_h40_kv8_s4096_bs128_fp16
  (added to match qwen3; compiles + passes golden).
• pypto: in-layer fa_fused from the qwen3-14B decode layer.

Since fa_fused runs inside the full decode layer (overlapped with the other
kernels), the fair metric is the attention kernels' on-core (AICube/AIVector)
busy time from the TraCR traces, not wall-clock.

metric	CCE b16_h40_kv8_s4096_bs128_fp16	pypto fa_fused + --max-seq	gap
on-core attention time	~548us	~1263us	~2.3x

Profile view

LEFT: CCE b16_h40_kv8_s4096_bs128_fp16 RIGHT: pypto fa_fused + --max-seq

This is the file to visualize on Perfetto:

cce_b16_h40_kv8_s4096_bs128_fp16.json

qwen3_14b_decode.json

How to recreate the plot

To create these profiles. One has to use the Simpler with TraCR branch, and when compiling Simpler, add the Env flag BUILD_TRACR=ON pip install --no-build-isolation -e .. TraCR is a low-level profiler that captures traces on the Ascend device. Also, set _DEFAULT_RUNTIME = "tensormap_and_ringbuffer" inside the pypto/runtime/worker.py as currently, TraCR is only built on top of this Simpler runtime scheduler. PyPTO should also be built based on this Simpler version; otherwise, PyPTO-based examples will not capture TraCR traces. When running PyPTO or Simpler examples, TraCR will produce a tracr_0/ in ~/ascend/. This has to be post-processed by running this command:

./pypto/runtime/build/output/bin/tracr_process ~/ascend/tracr_0/

It will generate a tracr_0/perfetto.json file, which can be viewed in Perfetto.

Motivation / Use Case

#986
showed the standalone pypto paged-attention example running ~6× slower than
the optimized CCE kernel. With the optimized attention from
pypto-lib#607, the gap
should be much smaller — the question is how much: 20%? 50%? This is the setup
and baseline for answering that quantitatively against the optimized CCE
implementation in simpler#899
(spmd_paged_attention_highperf), profiled on a2a3 (910B) with the simpler/TraCR
runtime profiler.

[lyfne123]

One correction on the comparison itself before reading anything into the 2.3×: the two sides aren't measuring the same scope, so the ~548µs vs ~1263µs isn't an attention-vs-attention number.

• CCE b16_h40_kv8_s4096_bs128_fp16 is the standalone spmd_paged_attention_highperf kernel — paged attention only (QK → softmax → SV → reduce).

• The ~1263µs on the pypto side is the whole qwen3-14B decode layer's on-core busy, not the attention kernels. In this build the decode layer lowers to ~35 kernels, of which only fa_fused and online_softmax are the attention compute. Everything else is a full transformer layer's worth of non-attention work:

  x_gamma / dcr_xgamma / rms_recip / post_rms_reduce   (RMSNorms)
  q_proj / k_proj / v_proj / qk_norm / rope_qkv         (QKV projection + qk-norm + RoPE)
  fa_fused / online_softmax                             (← the actual attention)
  out_proj / residual_rms_cast                          (O-proj + residual + post-norm)
  gate_proj / up_proj / silu / down_proj                (MLP)
  + the *_seed / fa_work_build / copy_* scaffolding
  

So the table is comparing PA-only (CCE) against a full decode layer (pypto) — the projections, RoPE, qk-norm, out-proj and the entire MLP are folded into that 1263µs.

When you isolate just the attention kernels (fa_fused + online_softmax) on-core (AICube + AIVector) busy at this matched shape, it's ~350µs, i.e. below the CCE PA kernel's ~548µs. So at the qwen3-14B attention shape the in-layer attention is already on par with / faster than the optimized CCE PA — the 2.3× is an artifact of the scope mismatch, not a real attention gap.

To make the row apples-to-apples, the pypto side should sum only fa_fused + online_softmax AICube+AIVector busy (exclude q/k/v/out/gate/up/down proj, the RMSNorms, RoPE and qk-norm), or alternatively compare the CCE number against a CCE full decode layer. Happy to share the per-kernel busy breakdown from the TraCR trace so we can drop in the corrected ~350µs and re-baseline.

[noabauma]

Corrected apples-to-apples: qwen3-14B attention vs CCE paged attention (best-of-both)

Re-measured on the matched current set — pypto main (~/src/pypto) + its
bundled runtime (simpler @09f1546f) on pto-isa 016396b5, both sides on the
same runtime. AICore on-core time via the runtime's L2 swimlane
(cross-checked against the device-log Total, which agrees within 2–3%).
Shape b16 / h40 / kv8 / d128 / s4096, BF16 (dtype-invariant — see caveats).
a2a3 / 910B; HBM peak = 1.5 TB/s.

Results — each at its best tiling

kernel	time	eff. HBM BW	% of peak
pypto attention (fa_fused + online_softmax)	~326 µs	0.82 TB/s	55%
CCE #899 PA — best (batch-head tiling)	~343 µs	0.78 TB/s	52%
CCE #899 PA — runtime default (long-seq tiling)	~462 µs	0.58 TB/s	39%

Best-of-both ratio: pypto / CCE-best ≈ 0.95× → essentially on par (pypto ~5%
ahead). Effective traffic = Q+K+V+Out ≈ 256 MiB (KV-dominated; internal GM
workspace excluded).

Conclusion (updated)

• With both kernels at their best tiling, pypto attention (~326 µs) and the
  CCE #899 port (~343 µs) are on par — both ~0.78–0.82 TB/s (~52–55% of HBM
  peak). There is no meaningful gap at this shape.
• The earlier “~29% faster” was an artifact of CCE being mis-tiled: the
  runtime’s heuristic routes b16/s4096 to a slow long-seq KV-split path
  (~462 µs / 0.58 TB/s). Forcing the batch-head tiling recovers CCE to ~343 µs.
• The original 2.3× was a scope artifact (whole decode layer ≈ 1118 µs vs
  attention-only) — that part still stands; attention is ~29% of the layer.
• Neither reaches the production ATB op (~1.0–1.1 TB/s) — a different,
  production-optimized kernel — so there is real bandwidth headroom for both.

The tiling story (key finding)

CCE’s number for this shape is highly tiling-sensitive, spanning:
0.58 (runtime default, long-seq) → 0.78 (batch-head, best on this runtime) →
~0.95 (stock runtime, same tiling, better orchestration) → 1.0–1.1 (ATB).

• The slow default comes from pa_tiling.py’s is_long_seq threshold
  (max_kv ≥ 4096), which mis-routes high-batch mid-context shapes onto the
  KV-split path. A batch-aware threshold (or the old 6144) auto-picks the fast
  batch-head path — recovering ~27% for free.
• Sweep: block_dim is capped at 24 (physical AICube cores); BN
  (batch-head) beats BNS (KV-split) for batch ≥ 16; head-grouping tuned.
  Fastest tiling = BN @ block_dim 24.

Engine breakdown

• CCE PA on the default tiling is AIV(vector)-bound (AIV 454 vs AIC
  380 µs); on the best tiling it’s balanced (AIC ≈ AIV ≈ 332 µs).
• pypto fa_fused is balanced (cube ≈ vec ≈ 294 µs) + a short online_softmax
  reduction (~32 µs). fa_fused is one mixed cube+vec root, so its two halves
  overlap (phase ≈ 294 µs, not their sum).

Scaling

• Seq length (s16384): pypto ~1269 µs vs CCE default ~1659 µs (~0.77×);
  16k routes both heuristics to the long-seq path, so no boundary ambiguity.
• Batch (b64): CCE scales cleanly to ~1720 µs / 0.62 TB/s (still AIV-bound on
  the default tiling); pypto attention at b64 could not be measured — the
  qwen3-14B decode kernel is hardwired to BATCH=16 (needs retuning).

Method / caveats

• On-core swimlane Exec; CCE device-log Total agrees within 2–3%
  (b16: 466 vs 458 µs), so the metric choice doesn’t change the result.
• Dtype-invariant: CCE BF16 (468 µs) ≈ FP16 (462 µs); same cube/vector MAC
  rate on Ascend.
• The CCE-best (343 µs) requires forcing the batch-head tiling; the runtime’s
  default heuristic picks the slow path.
• pypto attention is in-layer (overlapped); Exec is the kernels’ compute time
  regardless of overlap.
• pypto = 5-round (±~2%); the best-CCE point is single-shot.
• Unlocked device (no per-device lock), single 910B die.