Add: READMEs for the four a2a3 paged_attention variants - #1571
Conversation
Four sibling directories share a name prefix and nothing tells you what
separates them. Diffing establishes it, and each README now leads with the
answer:
- `paged_attention` is the baseline: four tasks per KV block (QK / SF / PV /
UP) inside a plain scope, with TensorMap deriving the ordering from the
tensors the tasks share.
- `paged_attention_manual_scope` has **byte-identical kernels** — all four —
and differs only in orchestration, 45 lines of 292. It is the cheapest way
to see what MANUAL scope costs, and it deliberately uses both dependency
APIs in one file: the primitive `set_dependencies(deps, n)` for tasks with a
fixed predecessor set, and `L0TaskArgsWithDeps<>::add_dep` for `UP`, whose
edges are conditional.
- `paged_attention_unroll_manual_scope` is **not** the manual-scope variant
plus an unroll. Every file differs; `aiv_softmax_prepare.cpp` changes 312
lines against a 156-line original. It batches KV blocks into groups of
`N_UNROLL` and submits four tasks per group rather than per block.
- `paged_attention_ringbuffer` contains **no kernels at all** — it points at
the `tests/st` batch-paged-attention sources and exists to stress ring
rotation at `ring_task_window: 64`, 256x below the 16384 default.
The directory index carried a wrong one-liner for the unroll variant
("manual scope plus loop unrolling"); it now describes a second
implementation, and the manual-scope row states the byte-identical kernels.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Warning Review limit reached
Next review available in: 5 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
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 |
The three a5 examples with no a2a3 counterpart to read instead. `bgemm` keeps its kernel in `kernels/mix/` rather than `kernels/aic/` plus `kernels/aiv/`, and that directory name is the example. One source is registered as two incores with different core_types; `__DAV_CUBE__` compiles the TLOAD/TMATMUL/TPUSH half and `__DAV_VEC__` the TPOP/TADD/TSTORE half, both over the same three-tensor args[]. The matmul product moves between them through VEC_FIFO and never reaches GM — only the accumulator C round-trips. a2a3's `benchmark_bgemm` computes the same thing with separate sources and a GM round trip, so the two are not ports of each other. `urma_deferred_completion_demo` and `sdma_async_completion_demo` run an identical two-rank protocol over different transports; `kernel_consumer.cpp` is byte-identical between them, leaving the transport as the only variable. Both are gated behind a workspace overlay that defaults OFF, and the two overlays are mutually exclusive in one build — `CommContext` exposes a single workSpace/workSpaceSize pair, so the CMakeLists raises a FATAL_ERROR if both are on. A stock build therefore skips both tests even on a5 silicon: a green CI run is not coverage of either path. Each README gives the two-step invocation, since the env var is read once by runtime_builder.py to compile the overlay in and once by the test's skipif. The a5 index gains those facts, and its `paged_attention_unroll_manual_scope` row is corrected the same way the a2a3 index was in #1571: the kernels are rewritten (softmax_prepare changes 269 lines against a 149-line original), so it is a second implementation rather than manual scope plus an unroll. The manual_scope row now records that its four kernels are byte-identical to the baseline's and only the orchestration differs, 52 lines of 288.
1 participant
Summary
Four sibling directories share a name prefix and nothing in the tree said what separates them. Diffing settles it, and each README now leads with the answer.
paged_attentionPTO2_SCOPE(), ordering derived by TensorMap from shared tensors.paged_attention_manual_scopepaged_attention_orch.cppdiffers, 45 lines of 292.paged_attention_unroll_manual_scopeaiv_softmax_prepare.cppchanges 312 lines against a 156-line original. Batches KV blocks into groups ofN_UNROLL— four tasks per group, not per block.paged_attention_ringbufferCALLABLEattests/st/.../batch_paged_attention/kernelsand exists to stress ring rotation atring_task_window: 64— 256× below the 16384 default (PTO2_TASK_WINDOW_SIZE).Two details worth surfacing that are otherwise only visible by reading the C++:
manual_scopedemonstrates both dependency APIs in one file, and the choice is principled: the primitivePTO2TaskId deps[]+set_dependencies(deps, n)forSFandPV, whose predecessor set is fixed;L0TaskArgsWithDeps<>::add_depforUP, whose edges are conditional (previous group'sUPonly when one exists; the allocation task only on the last iteration, and only to keep scratch alive). It also shows an edge deliberately left out —UPreadsSF'smi/li, butSF → PV → UPalready orders it.manual_scope'sCase1needsring_task_window: 32768. Long-context cases put more than 16384 in-flight tasks in one MANUAL scope and can wedge the orchestrator withFLOW_CONTROL_DEADLOCK(code 3) at the default.Correction to the index
examples/a2a3/tensormap_and_ringbuffer/README.md(merged in #1564) described the unroll variant as "Manual scope plus loop unrolling." That is wrong — the kernels are rewritten. The row now says so, and the manual-scope row records the byte-identical kernels.Testing
markdownlint-cli2 --config tests/lint/.markdownlint.yaml— 0 errors on all five filesdiff <baseline>/<file> <variant>/<file> | grep -c '^[<>]'PTO2_TASK_WINDOW_SIZE 16384confirmed atsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h:66; case names andruntime_envvalues read from eachCASESlistDocs only; no code, test, or build files touched.