test(a5/tmr): port a2a3-only tensormap_and_ringbuffer scene/L3 tests to a5 - #1450
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe 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. ChangesAlternating matmul and add
Batched paged attention
Four-dimensional paged attention
SPMD TPUSH/TPOP attention
Runtime lifecycle and L3 behavior
Fan-in and dispatch regression coverage
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
There was a problem hiding this comment.
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 winDocument that scene
CALLABLEandCASESmetadata 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.ClassVaror 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 winAnnotate shared class configuration with
ClassVar.Ruff RUF012 flags both mutable class attributes. These are intentionally class-level scene metadata, so annotate them as
ClassVarto 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 winAnnotate mutable class configuration as
ClassVar.
CALLABLEandCASESare intentionally class-level configuration, but their mutable types trigger Ruff RUF012. AddClassVarannotations 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 valueAlign the SPMD paged-attention orchestration entry with the existing AICPU entry convention. All other orchestration kernels under
tests/st/a5/tensormap_and_ringbufferexportaicpu_orchestration_entry, but the graph-basedmulti_round_paged_attentionport usesbuild_paged_attention_graph; pick one entry ABI for this test and keep the kernel export andfunction_namematching 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
📒 Files selected for processing (32)
tests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aic/kernel_matmul.cpptests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/aiv/kernel_add.cpptests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/kernels/orchestration/alternating_orch.cpptests/st/a5/tensormap_and_ringbuffer/alternating_matmul_add/test_alternating_matmul_add.pytests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpptests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpptests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpptests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpptests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/kernels/orchestration/paged_attention_orch.cpptests/st/a5/tensormap_and_ringbuffer/batch_paged_attention/test_batch_paged_attention.pytests/st/a5/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.pytests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/aic/kernel_write_const_visible.cpptests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/kernels/orchestration/fanin_lookup_perf_orch.cpptests/st/a5/tensormap_and_ringbuffer/fanin_lookup_perf/test_fanin_lookup_perf.pytests/st/a5/tensormap_and_ringbuffer/multi_round_paged_attention/test_multi_round_paged_attention.pytests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_pv_matmul.cpptests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aic/aic_qk_matmul.cpptests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_online_update.cpptests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/aiv/aiv_softmax_prepare.cpptests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/kernels/orchestration/paged_attention_orch.cpptests/st/a5/tensormap_and_ringbuffer/paged_attention_unroll_4dims/test_paged_attention_unroll_4dims.pytests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aic/kernel_write.cpptests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/aiv/kernel_write.cpptests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/kernels/orchestration/spmd_batch_dispatch_oob_orch.cpptests/st/a5/tensormap_and_ringbuffer/spmd_batch_dispatch_oob/test_spmd_batch_dispatch_oob.pytests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpptests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpptests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/test_spmd_paged_attention.pytests/st/a5/tensormap_and_ringbuffer/test_l3_dependency.pytests/st/a5/tensormap_and_ringbuffer/test_l3_group.pytests/st/a5/tensormap_and_ringbuffer/test_l3_host_buffer_registration.pytests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py
92d57d9 to
56b05a5
Compare
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).
8cbef18 to
f706ef1
Compare
84ea9e5 to
2189dde
Compare
|
Reviewed at The port itself is careful. I verified each of the three adaptations that looked suspicious and all three are correct:
Five things I'd like changed before merge, plus three notes that need no code change. 1. Unreachable guards in the a5 orchestration
In the previous revision these two No case can reach them. I computed all three: This matters because the PR's stated purpose is 1:1 differential validation. That only works if the two arch copies are logically identical — otherwise a future behavioral difference can't be attributed to silicon vs. a branch this port added. Either delete them (restores parity with a2a3), or add a case that actually trips one and asserts the failure — in which case please also state "add input validation" as an explicit goal in the PR body, since nothing in the current description covers it. 2. Half-qualified
|
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 scalar-data, paged-attention ringbuffer, and merge-pipeline-barrier examples with architecture-specific READMEs plus the A5 InCore orchestration guide. Adapt the ports to current A5 APIs and architecture-derived launch defaults. Make scene compile-cache keys module-aware as preventative isolation for same-named test classes in different modules. Exclude multi-device coverage and spmd_paged_attention from this port.
2189dde to
713cec6
Compare
|
@ChaoWao Addressed all five requested changes:
Also clarified in the PR body that the module-aware compile-cache key is preventative, folded that behavior into the commit message, and kept the simulation applicability note explicit in the test plan. Thanks for the careful review. |
Summary
tensormap_and_ringbufferscene/L3 coverage onto a5 for differential validation.spmd_paged_attention;spmd_paged_attention_highperfalso remains out of scope pending a dedicated dav-c310 port.Test plan