test: add manual cross-runtime SPMD paged attention coverage - #1817
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:
📝 WalkthroughWalkthroughAdded an A5 MIX SPMD paged-attention kernel with AIC/AIV FIFO processing, 24-block orchestration, scene coverage, and A2A3/A5 host-build-graph tests using zero-query cases. ChangesSPMD paged attention
Estimated code review effort: 4 (Complex) | ~60 minutes Mergeability Score: 🟠 High · up to The PR adds new manual paged-attention runtime paths, but the current implementation can address incorrect memory regions or produce invalid results for certain block sizes and tensor shapes. Merge should be blocked until those configurations are corrected or explicitly rejected. Sequence Diagram(s)sequenceDiagram
participant SceneTest
participant OrchestrationEntry
participant MixedKernels
participant KernelEntry
participant AIVProcessing
SceneTest->>OrchestrationEntry: Submit tensors and runtime scalars
OrchestrationEntry->>MixedKernels: Configure AIC and AIV kernels
MixedKernels->>KernelEntry: Dispatch 24 hardware blocks
KernelEntry->>AIVProcessing: Exchange scores and value results through FIFOs
AIVProcessing-->>SceneTest: Store normalized output
Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
6e0cb3f to
c8ec0d2
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp (1)
157-172: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unreachable prefetch path and the unused accumulator tiles.
The only call sites are lines 325-331, and they never pass
current_loaded,has_next, ornext_kv_block_id. Therefore thehas_nextprefetch blocks never execute, andcurrent_loadedis alwaysfalse. ThebMatTile_*_Bbuffers stay unused becausei % 2selects_Aonly on eveniwhile the_Bload never happens; on oddithe code loads_Band moves_B, so both branches work, but the double-buffer intent no longer applies to the serialized per-block handshake described at lines 318-322.The following AIV tiles are declared and assigned an address but never read or written:
gmND,miNewND,savedLlDN,savedLlRow,llRow.Drop the dead parameters and declarations, or keep them behind a comment that states the pipelined variant is intentionally retained for a follow-up change.
Also applies to: 206-222, 492-499
🤖 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/mix/paged_attention_parallel.cpp` around lines 157 - 172, Remove the unreachable prefetch and double-buffering paths from aic_qk_step and its call sites: delete current_loaded, has_next, next_kv_block_id, the unused bMatTile_QK_B handling, and related dead declarations. Also remove the unused AIV tiles gmND, miNewND, savedLlDN, savedLlRow, and llRow, including their address assignments, while preserving the serialized per-block handshake behavior.tests/st/a5/host_build_graph/spmd_paged_attention/test_spmd_paged_attention.py (1)
20-20: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the HBG cases from every base case.
Line 20 copies only
_TmrBase.CASES[0]. If the tensormap_and_ringbuffer class gains a second case, the host-build-graph class silently keeps one case.♻️ Proposed change
- CASES = [{**_TmrBase.CASES[0], "platforms": ["a5sim", "a5"], "manual": True}] + CASES = [{**case, "platforms": ["a5sim", "a5"], "manual": True} for case in _TmrBase.CASES]🤖 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/host_build_graph/spmd_paged_attention/test_spmd_paged_attention.py` at line 20, Update the CASES definition in the host-build-graph test to derive entries from every case in _TmrBase.CASES, applying the existing platforms and manual overrides to each entry instead of selecting only index 0.
🤖 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/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp`:
- Around line 602-622: Add an early validation in the orchestration path
containing this block and in run_aiv to reject any head_dim value other than 128
before launching AIC/AIV processing. Preserve the existing 128-wide execution
for valid inputs and make the unsupported configuration fail visibly using the
surrounding error-handling mechanism.
- Around line 368-375: Update the partial-tail UB address calculation in the
is_last_partial block to use MAX_BLOCK_SIZE for each alternating slot’s stride,
matching the two-slot allocation used by TPOP with TileSplitAxis::TILE_UP_DOWN;
preserve the existing i%2 slot selection and padding behavior.
In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/orchestration/spmd_paged_attention_orch.cpp`:
- Around line 83-87: Update q_tile calculation in the orchestration setup to use
the documented minimum of num_heads and MAX_Q_TILE, ensuring each logical tile
covers only valid rows. Add explicit precondition checks for supported num_heads
shapes and for block_size values accepted by the kernel dispatch (128 or 64), so
unsupported inputs fail visibly; anchor these changes near
q_loop/total_logical_blocks and the existing CASES configuration.
---
Nitpick comments:
In
`@tests/st/a5/host_build_graph/spmd_paged_attention/test_spmd_paged_attention.py`:
- Line 20: Update the CASES definition in the host-build-graph test to derive
entries from every case in _TmrBase.CASES, applying the existing platforms and
manual overrides to each entry instead of selecting only index 0.
In
`@tests/st/a5/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp`:
- Around line 157-172: Remove the unreachable prefetch and double-buffering
paths from aic_qk_step and its call sites: delete current_loaded, has_next,
next_kv_block_id, the unused bMatTile_QK_B handling, and related dead
declarations. Also remove the unused AIV tiles gmND, miNewND, savedLlDN,
savedLlRow, and llRow, including their address assignments, while preserving the
serialized per-block handshake behavior.
🪄 Autofix
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: 78d19396-38f0-45d4-acab-a0209d00fac3
📒 Files selected for processing (5)
tests/st/a2a3/host_build_graph/spmd_paged_attention/test_spmd_paged_attention.pytests/st/a5/host_build_graph/spmd_paged_attention/test_spmd_paged_attention.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.py
- add A5-specific MIX execution and manual A2/A3 and A5 cases - validate supported A5 tile shapes and use portable FIFO lane lookup - keep serialized handshakes and correct partial-tail addressing - run A5 benchmark BGEMM Case0 on sim and document its fixed tiles
c8ec0d2 to
94ef12c
Compare
|
Addressed the remaining review-summary nitpicks:
Also added A5 sim compatibility by using the portable FIFO overload that derives the current lane via |
2 participants
Summary
All newly added cases are manual because A5 currently only passes the small zero-query flow. Larger multi-block workloads stall, and random-query workloads still have numerical mismatches.
Tracks #1816.