Skip to content

test(a5/tmr): port a2a3-only tensormap_and_ringbuffer scene/L3 tests to a5 - #1450

Open
yanghaoran29 wants to merge 1 commit into
hw-native-sys:mainfrom
yanghaoran29:port-a2a3-tmr-tests-to-a5
Open

test(a5/tmr): port a2a3-only tensormap_and_ringbuffer scene/L3 tests to a5#1450
yanghaoran29 wants to merge 1 commit into
hw-native-sys:mainfrom
yanghaoran29:port-a2a3-tmr-tests-to-a5

Conversation

@yanghaoran29

@yanghaoran29 yanghaoran29 commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Port the remaining single-device a2a3-only tensormap_and_ringbuffer scene/L3 coverage onto a5 for differential validation.
  • Cover L3 dependency and host-buffer behavior, dynamic registration, batch/multi-round/4D-unrolled paged attention, alternating matmul/add, fan-in lookup, empty-ring heap rebase, and SPMD batch-dispatch bounds.
  • Add the scalar-data, paged-attention ringbuffer, and merge-pipeline-barrier examples plus the a5 InCore orchestration guide.
  • Adapt the ports to current a5 launch, logging, stride, and runtime APIs; make scene compile-cache keys module-aware to prevent same-named test classes from sharing binaries.
  • Exclude multi-device coverage and spmd_paged_attention; spmd_paged_attention_highperf also remains out of scope pending a dedicated dav-c310 port.

Test plan

  • A5Sim: all applicable newly added cases pass.
  • A5 single-device onboard: applicable cases pass when split into tasks under the 30-second task limit.
  • Cache regression unit tests: 3 passed.
  • Pre-commit passed except pyright, explicitly skipped because its npm dependency fetch returns HTTP 403 in the local environment.

@coderabbitai

coderabbitai Bot commented Jul 23, 2026

Copy link
Copy Markdown

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 341dd5e1-9475-4e84-b7c4-257dd7fab23d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The PR adds multiple TensorMap and ringbuffer scene implementations: tiled matmul/add, batched and SPMD paged attention, fan-in and dispatch regression scenarios, plus runtime registration, dependency, host-buffer, and prewarm tests.

Changes

Alternating matmul and add

Layer / File(s) Summary
Tiled kernels and interleaved orchestration
tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/...
Adds AIC matmul and AIV add kernels, per-group tensor views, interleaved task submission, and golden-output validation.

Batched paged attention

Layer / File(s) Summary
Attention kernels and scheduling
tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/...
Adds QK matmul, softmax preparation, PV matmul, online update, and chunked orchestration across batches and heads.
Scene validation
tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py
Adds parameterized execution cases and shared paged-attention golden comparison.

Four-dimensional paged attention

Layer / File(s) Summary
Unrolled kernels and task pipeline
tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/...
Adds pipelined QK, SplitK PV, softmax, online update, and four-stage unrolled orchestration.
Scene validation
tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py
Adds 4D query/output handling and reshaping around the shared golden computation.

SPMD TPUSH/TPOP attention

Layer / File(s) Summary
Mixed FIFO kernel and launch
tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/...
Adds cooperative AIC/AIV FIFO processing, fixed 24-block launch configuration, and runtime tile dispatch.
Scene validation
tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py
Adds callable wiring, parameterized cases, argument construction, and golden-output checks.

Runtime lifecycle and L3 behavior

Layer / File(s) Summary
Dynamic registration
tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py
Adds post-start registration tests covering broadcast, capacity, duplicate handles, unregistering, and re-preparation.
L3 dependencies and grouping
tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py, tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py
Adds dependent subtasks and grouped chip-worker execution scenarios.
Host buffers and prewarm
tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py, tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py
Adds host-buffer zero-copy execution and ring-task-window prewarm initialization tests.

Fan-in and dispatch regression coverage

Layer / File(s) Summary
Fan-in lookup performance
tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/...
Adds explicit producer-to-consumer fan-in dependencies, optional real kernels, and benchmark validation.
SPMD batch dispatch bounds
tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/...
Adds two-task dispatch reproduction kernels and cache-line-position output checks.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

Poem

A rabbit hops through tiles of light,
Matmuls bloom and queues take flight.
FIFOs hum from left to right,
Softmax settles sums just right.
Tests guard each path with care—
“Carrots for kernels everywhere!”

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 17.11% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
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.
Title check ✅ Passed The title accurately summarizes the main change: porting a2a3-only tensormap_and_ringbuffer scene and L3 tests to a5.
Description check ✅ Passed The description matches the changeset, covering the a5 ports, excluded scopes, and test plan.

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

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

🧹 Nitpick comments (4)
tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py (1)

30-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document that scene CALLABLE and CASES metadata must remain read-only.

These mutable class attributes are shared across test cases; any recursive update during path resolution or validation may leak between cases. Mark them read-only/read-only-convention-compliant with typing.ClassVar or a narrow RUF012 suppression if mutation is never intended.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py`
around lines 30 - 66, Mark the shared class attributes CALLABLE and CASES as
read-only metadata by annotating them with typing.ClassVar, or apply a narrow
RUF012 suppression if they must remain mutable by convention. Update both
affected declarations in
tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py:30-66
and :68-134, without changing their contents or runtime behavior.

Source: Linters/SAST tools

tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py (1)

30-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate shared class configuration with ClassVar.

Ruff RUF012 flags both mutable class attributes. These are intentionally class-level scene metadata, so annotate them as ClassVar to make the ownership explicit and clear the warning without changing runtime behavior.

Proposed fix
 import ctypes
+from typing import ClassVar

 ...

-    CALLABLE = {
+    CALLABLE: ClassVar = {
         ...
     }

-    CASES = [
+    CASES: ClassVar = [
         ...
     ]

Also applies to: 48-61

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py`
around lines 30 - 46, Annotate the mutable class-level scene metadata containing
CALLABLE and the corresponding configuration at the other flagged location with
typing.ClassVar. Preserve their existing values and class-level behavior while
clearing Ruff RUF012.

Source: Linters/SAST tools

tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py (1)

38-64: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate mutable class configuration as ClassVar.

CALLABLE and CASES are intentionally class-level configuration, but their mutable types trigger Ruff RUF012. Add ClassVar annotations to clear the warning without moving them to instance state.

Proposed fix
+from typing import ClassVar
+
...
-    CALLABLE = {
+    CALLABLE: ClassVar[dict] = {
...
-    CASES = [
+    CASES: ClassVar[list] = [

Also applies to: 66-73

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py`
around lines 38 - 64, Annotate the class-level mutable configuration attributes
CALLABLE and CASES with typing.ClassVar, preserving their existing values and
keeping them on the class rather than moving them to instance state.

Source: Linters/SAST tools

tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp (1)

64-64: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low value

Align the SPMD paged-attention orchestration entry with the existing AICPU entry convention. All other orchestration kernels under tests/st/a5/tensormap_and_ringbuffer export aicpu_orchestration_entry, but the graph-based multi_round_paged_attention port uses build_paged_attention_graph; pick one entry ABI for this test and keep the kernel export and function_name matching it.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp`
at line 64, Align the SPMD paged-attention entry ABI across the orchestration
kernel and test configuration: choose either aicpu_orchestration_entry or
build_paged_attention_graph, then make the exported function in
spmd_paged_attention_orch.cpp and the function_name in
test_spmd_paged_attention.py use the same symbol, matching the convention
selected for this test.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp`:
- Around line 70-121: Validate matmul_batch and add_batch before the group-count
divisions in the orchestration setup, rejecting zero values and any
configuration where total_matmul_tasks or total_add_tasks is not evenly
divisible by its corresponding batch size. Use the existing runtime assertion or
validation mechanism, and only compute num_matmul_groups and num_add_groups
after these checks so no tasks are silently dropped.

In
`@tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp`:
- Around line 128-143: Case3 dispatches head_dim=256 but the q_tile-64 kernels
only support 128-wide data. Add matching head_dim=256 template dispatches in
aic_qk_matmul.cpp (qk_matmul_batch_impl), aic_pv_matmul.cpp (PV implementation),
and aiv_online_update.cpp (online update), preserving the existing 128 paths for
other inputs; update test_batch_paged_attention.py only if the wide-head-dim
kernels are not implemented, by removing or guarding Case3.

In
`@tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py`:
- Around line 98-113: Disable or remove the manual Case3 configuration in the
test definition, unless the paged-attention orchestration and its q_tile=64 QK,
PV, softmax, and update branches are updated to support head_dim=256. Do not
leave Case3 enabled while only head_dim=128 templates process the workload.

In `@tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py`:
- Around line 25-28: Update the documentation in the host-buffer registration
test to replace stale a2a3/a2a3sim references with a5/a5sim, including the
statement about vector_example kernels, so it accurately matches the test’s
a5sim runtime and KERNELS_BASE configuration.

---

Nitpick comments:
In
`@tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py`:
- Around line 30-46: Annotate the mutable class-level scene metadata containing
CALLABLE and the corresponding configuration at the other flagged location with
typing.ClassVar. Preserve their existing values and class-level behavior while
clearing Ruff RUF012.

In
`@tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py`:
- Around line 30-66: Mark the shared class attributes CALLABLE and CASES as
read-only metadata by annotating them with typing.ClassVar, or apply a narrow
RUF012 suppression if they must remain mutable by convention. Update both
affected declarations in
tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py:30-66
and :68-134, without changing their contents or runtime behavior.

In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py`:
- Around line 38-64: Annotate the class-level mutable configuration attributes
CALLABLE and CASES with typing.ClassVar, preserving their existing values and
keeping them on the class rather than moving them to instance state.

In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp`:
- Line 64: Align the SPMD paged-attention entry ABI across the orchestration
kernel and test configuration: choose either aicpu_orchestration_entry or
build_paged_attention_graph, then make the exported function in
spmd_paged_attention_orch.cpp and the function_name in
test_spmd_paged_attention.py use the same symbol, matching the convention
selected for this test.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 4306e68d-14ef-46f2-a446-2e70cba34302

📥 Commits

Reviewing files that changed from the base of the PR and between c032e07 and 67e6bfa.

📒 Files selected for processing (32)
  • tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpp
  • tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpp
  • tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.py
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.py
  • tests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py
  • tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpp
  • tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.py
  • tests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.py
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpp
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpp
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpp
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpp
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.py
  • tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpp
  • tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpp
  • tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.py
  • tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp
  • tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp
  • tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.py
  • tests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.py
  • tests/st/a5/tensormap_and_ringbuffer/test_l3_group.py
  • tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py
  • tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py

Comment thread tests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py Outdated
@yanghaoran29
yanghaoran29 force-pushed the port-a2a3-tmr-tests-to-a5 branch 2 times, most recently from 92d57d9 to 56b05a5 Compare July 23, 2026 07:18
yanghaoran29 added a commit to yanghaoran29/simpler that referenced this pull request Jul 23, 2026
Two CI failures on PR hw-native-sys#1450 (st-sim-a5 / st-onboard-a5), both arch-drift
from copying a2a3 sources verbatim:

- Kernel compile "reference to 'Stride' is ambiguous" on a5 (Incore
  compile exit 1): a5 kernel headers put another Stride in scope, so bare
  Stride<...> is ambiguous. Qualify to pto::Stride<...> (a5 convention) in
  alternating_matmul_add, batch_paged_attention, paged_attention_unroll_4dims,
  spmd_paged_attention kernels.
- L3 TypeError "missing required keyword-only argument": upstream main made
  Orchestrator.submit_next_level(worker=) / submit_next_level_group(workers=)
  required. Pass worker=0 / workers=[0, 1], matching the a2a3 originals.

Verified: both L3 tests pass on a5sim; all 4 scenes' kernels compile clean
for a5 (device-independent compile check).
@yanghaoran29
yanghaoran29 force-pushed the port-a2a3-tmr-tests-to-a5 branch 8 times, most recently from 8cbef18 to f706ef1 Compare July 28, 2026 08:29
@yanghaoran29
yanghaoran29 force-pushed the port-a2a3-tmr-tests-to-a5 branch 2 times, most recently from d5774be to 5c8863b Compare July 30, 2026 07:36
Port the remaining single-device a2a3 tensormap_and_ringbuffer coverage to a5 for differential validation across architectures.

Add scenes for alternating matmul/add, batch and multi-round paged attention, 4D unroll, dynamic registration, fanin lookup, empty-ring heap rebase, SPMD batch dispatch, and L3 dependency and host-buffer behavior.

Add the scalar-data, paged-attention ringbuffer, and merge-pipeline-barrier examples, plus the a5 InCore orchestration guide. Adapt the ports to current a5 launch, logging, stride, and runtime APIs, and keep scene compile-cache entries distinct across modules.

Exclude multi-device coverage and spmd_paged_attention from this port.
@yanghaoran29
yanghaoran29 force-pushed the port-a2a3-tmr-tests-to-a5 branch from 5c8863b to d772f32 Compare July 31, 2026 03:34
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