Add: benchmark paged attention unroll on host build graph

[doraemonmj]

Summary

• Add the host_build_graph counterpart of paged_attention_unroll Case1 and Case2.
• Keep the TMR parameters, callable signatures, tolerances, argument generation, and golden unchanged.
• Reuse the existing TMR orchestration and incore kernel sources directly so both runtimes execute the same graph.
• Leave Case3 out because its shared head_dim=256 kernels fail golden on both runtimes.

This is one workload-sized part of #1727.

Performance

Measured on one locked a2a3 NPU through task-submit: eight adjacent TMR/HBG pairs, alternating order, ten rounds per process and case, with each case's first round excluded (72 steady samples per runtime/case).

Case	Host total: TMR -> HBG	Change	Device wall: TMR -> HBG	Change
Case1	54.237 -> 73.281 ms	+35.11%	1.386 -> 1.268 ms	-8.54%
Case2	18.206 -> 32.775 ms	+80.03%	0.766 -> 0.670 ms	-12.62%

Paired device-wall changes (95% CI): Case1 -8.53% (-10.50% to -6.51%); Case2 -12.55% (-16.01% to -8.96%). HBG is consistently faster on device; the measured host increase is in bind and validation.

These measurements used simpler base 7a1b9b11 and PTO-ISA pin 0cefc9a5a1c24c62655cc345d408559595a8af32. Main later merged #1763, which removes redundant HBG arena initialization and explicitly leaves device time unchanged; the device comparison remains representative, while the host/bind values above are conservative pre-#1763 measurements.

Performance job: task_20260811_003357_606081246 (NPU 3, exit 0).

Testing

• HBG Case1 and Case2 onboard golden: task_20260811_003020_385776328006
• Eight-pair onboard performance comparison
• All pre-commit hooks for the changed file

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f2872f10-9a01-45b0-99af-da93fbcf050a

📥 Commits

Reviewing files that changed from the base of the PR and between 9a8b9d0 and ef5b943.

📒 Files selected for processing (1)

• tests/st/a2a3/host_build_graph/paged_attention_unroll/test_paged_attention_unroll.py

---

📝 Walkthrough

Walkthrough

Adds an A2A3 host-build-graph paged-attention unroll benchmark. The test registers kernel bindings, defines two bfloat16 cases, builds task arguments, computes golden results, and executes the module.

Changes

Paged-attention unroll benchmark

Layer / File(s)	Summary
Benchmark definition tests/st/a2a3/host_build_graph/paged_attention_unroll/test_paged_attention_unroll.py	Adds the scene test, host-build-graph orchestration binding, four AIC/AIV kernel bindings, metadata, and two A2A3 bfloat16 cases.
Argument, golden, and execution flow tests/st/a2a3/host_build_graph/paged_attention_unroll/test_paged_attention_unroll.py	Constructs tensor and scalar task arguments, computes expected results with shared paged-attention utilities, and runs the module.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SceneTestCase
  participant generate_args
  participant TaskArgsBuilder
  participant paged_attention_utilities
  SceneTestCase->>generate_args: create benchmark arguments
  generate_args->>TaskArgsBuilder: add TensorArg and Scalar values
  SceneTestCase->>paged_attention_utilities: compute golden result
  paged_attention_utilities-->>SceneTestCase: return expected tensor values
  SceneTestCase->>SceneTestCase: run_module

Loading

Possibly related PRs

• hw-native-sys/simpler#1771: Adds a similar A2A3 host-build-graph benchmark with matching test orchestration patterns.

Poem

A rabbit checks the tensors in a row,
While kernel bindings start the show.
Golden values hop into place,
Two bfloat16 cases race.
The module runs; ears point up—
“Paged attention fills my cup!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the addition of a paged-attention unroll benchmark for the host build graph.
Description check	✅ Passed	The description directly explains the Case1 and Case2 host build graph changes, exclusions, performance results, and testing.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}