feat(phyai): radix prefix-cache bridge for AR attention#22
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a radix-cache to AR attention bridge, adding RadixAttentionPlanner and RadixSequence to support prefix-reusing attention. It also updates the EagerARBackend to gather non-contiguous KV slots, removing the previous contiguous-slab constraint. The review feedback highlights two key improvement opportunities: optimizing the lazy import mechanism in ar/__init__.py to prevent recursive __getattr__ calls, and avoiding an expensive GPU-to-CPU synchronization in the planner by performing the slot ID boundary check on the host-side list rather than on the GPU tensor.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if name in ("RadixAttentionPlanner", "RadixSequence"): | ||
| from phyai.layers.attention.ar import radix | ||
|
|
||
| value = getattr(radix, name) | ||
| globals()[name] = value | ||
| return value |
There was a problem hiding this comment.
Importing the classes directly from the submodule phyai.layers.attention.ar.radix is cleaner and avoids potentially triggering __getattr__ recursively for 'radix'. Additionally, populating both RadixAttentionPlanner and RadixSequence in globals() on the first access avoids subsequent __getattr__ calls for the other class.
| if name in ("RadixAttentionPlanner", "RadixSequence"): | |
| from phyai.layers.attention.ar import radix | |
| value = getattr(radix, name) | |
| globals()[name] = value | |
| return value | |
| if name in ("RadixAttentionPlanner", "RadixSequence"): | |
| from phyai.layers.attention.ar.radix import RadixAttentionPlanner, RadixSequence | |
| globals()["RadixAttentionPlanner"] = RadixAttentionPlanner | |
| globals()["RadixSequence"] = RadixSequence | |
| return globals()[name] |
📝 WalkthroughWalkthroughThis PR adds a radix-cache-to-AR-attention bridge under phyai/src/phyai/layers/attention/ar/radix, comprising a RadixSequence dataclass, a RadixAttentionPlanner (plan/commit/release with rollback semantics), and lazy re-exports from the ar package. CPU and CUDA end-to-end tests validate planning, prefix reuse, and correctness. ChangesRadix-cache AR attention bridge
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant RadixAttentionPlanner
participant PrefixCache
participant KVCachePool
Caller->>RadixAttentionPlanner: plan(sequences)
RadixAttentionPlanner->>PrefixCache: match_prefix, lock node
RadixAttentionPlanner->>KVCachePool: allocate suffix units
RadixAttentionPlanner-->>Caller: ARAttnMetadata
Caller->>RadixAttentionPlanner: commit(sequences)
RadixAttentionPlanner->>PrefixCache: insert fresh suffix units
Caller->>RadixAttentionPlanner: release(sequences)
RadixAttentionPlanner->>PrefixCache: unlock node_ref
RadixAttentionPlanner->>KVCachePool: free uncommitted units
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
RadixAttentionPlanner maps phyai_ext.radix_cache prefix matches into ARAttnMetadata (reuse cached-prefix slots, allocate only the suffix) for the flashinfer paged backend; the attention layer is untouched. Lazily re-exported from phyai.layers.attention.ar so the base package needn't pull the phyai-ext extra.
Glorhop
force-pushed
the
feat/radix-cache-ar
branch
from
7470c96 to
e0e4f77
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (4)
phyai/tests/layers/attention/ar/test_radix_flashinfer.py (2)
110-119: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCall
planner.release()after each scenario to match the planner lifecycle.Same as the gather test:
release()drops prefix-reuse locks and frees uncommitted suffix units. Addingplanner.release([a])after the seed phase andplanner.release([c])after the reuse phase makes the lifecycle explicit and avoids masking leak-detection regressions.🤖 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 `@phyai/tests/layers/attention/ar/test_radix_flashinfer.py` around lines 110 - 119, The test is missing explicit planner lifecycle cleanup after each scenario. Update the radix flashinfer test to call planner.release() for the seeded sequence and the reuse sequence using the relevant RadixSequence symbols (planner, plan, commit, release, a, and c), so prefix-reuse locks and uncommitted suffix units are released and the lifecycle matches the other planner tests.
98-108: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
meta.modeinstead of hardcodingAttnMode.PREFILLin therunhelper.The
run()function hardcodesmode=AttnMode.PREFILLinARAttnCtx, butmetaalready carries the mode fromplanner.plan(). Ifplan()is ever called with a differentmodeargument, the context would be inconsistent with the metadata. Usemode=meta.modefor robustness.♻️ Proposed fix
ctx = ARAttnCtx( backend=backend, plan=plan, - mode=AttnMode.PREFILL, + mode=meta.mode, layout=AttnLayout.RAGGED_3D, kv_pool=pool, write_indices=meta.write_indices, )🤖 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 `@phyai/tests/layers/attention/ar/test_radix_flashinfer.py` around lines 98 - 108, The run helper hardcodes the attention mode in ARAttnCtx, which can make the context inconsistent with the metadata returned by planner.plan(). Update the run(meta, qsub, ksub, vsub) helper to pass meta.mode into ARAttnCtx instead of AttnMode.PREFILL, keeping the mode aligned with the metadata produced by init_forward_metadata and planner.plan().phyai/tests/layers/attention/ar/test_radix_gather_e2e.py (1)
29-76: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCall
planner.release()after each scenario to match the planner lifecycle.The planner docs describe
release()as the step that drops prefix-reuse locks and frees uncommitted suffix units. While the test objects are garbage-collected, explicitly callingrelease([a]),release([d]), andrelease([c])at the end of each phase makes the lifecycle explicit and avoids masking any leak-detection regressions in future test runs within the same process.🤖 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 `@phyai/tests/layers/attention/ar/test_radix_gather_e2e.py` around lines 29 - 76, The radix attention test is missing explicit planner cleanup, so add planner.release() for each RadixSequence scenario to mirror the intended lifecycle. Use the RadixAttentionPlanner release path after the seed sequence a, the decoy sequence d, and the reuse sequence c, so prefix-reuse locks and uncommitted suffix units are dropped deterministically. Keep the calls near the end of each phase in test_radix_reuse_gathers_exact_prefix_plus_suffix to make the test resilient to future leak-detection checks.phyai/src/phyai/layers/attention/ar/radix/planner.py (1)
312-314: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd type annotations to
_assemble(and consider_match_prefix's return type).
_assembleis fully untyped while the rest of the module is typed. Annotating the signature keeps the module typed and documents the-> ARAttnMetadatacontract;Balso reads better as a lowercasebatch_size.♻️ Suggested signature
- def _assemble( - self, B, suffix_lens, total_lens, indices_parts, write_parts, pos_parts, mode - ): + def _assemble( + self, + batch_size: int, + suffix_lens: list[int], + total_lens: list[int], + indices_parts: list[torch.Tensor], + write_parts: list[torch.Tensor], + pos_parts: list[torch.Tensor], + mode: AttnMode, + ) -> ARAttnMetadata:Update the internal
Breferences accordingly.As per coding guidelines: "keep modules typed where practical".
🤖 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 `@phyai/src/phyai/layers/attention/ar/radix/planner.py` around lines 312 - 314, _add type annotations to the untyped planner helpers in the radix attention planner: update `_assemble` in `ARPlanner` to use typed parameters (including renaming `B` to `batch_size` for clarity) and an explicit `-> ARAttnMetadata` return type, and make sure any internal references to `B` are updated accordingly. Also review `_match_prefix` in the same class/module and add its return type if it is still implicit so the module remains consistently typed.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@phyai/src/phyai/layers/attention/ar/radix/planner.py`:
- Around line 312-314: _add type annotations to the untyped planner helpers in
the radix attention planner: update `_assemble` in `ARPlanner` to use typed
parameters (including renaming `B` to `batch_size` for clarity) and an explicit
`-> ARAttnMetadata` return type, and make sure any internal references to `B`
are updated accordingly. Also review `_match_prefix` in the same class/module
and add its return type if it is still implicit so the module remains
consistently typed.
In `@phyai/tests/layers/attention/ar/test_radix_flashinfer.py`:
- Around line 110-119: The test is missing explicit planner lifecycle cleanup
after each scenario. Update the radix flashinfer test to call planner.release()
for the seeded sequence and the reuse sequence using the relevant RadixSequence
symbols (planner, plan, commit, release, a, and c), so prefix-reuse locks and
uncommitted suffix units are released and the lifecycle matches the other
planner tests.
- Around line 98-108: The run helper hardcodes the attention mode in ARAttnCtx,
which can make the context inconsistent with the metadata returned by
planner.plan(). Update the run(meta, qsub, ksub, vsub) helper to pass meta.mode
into ARAttnCtx instead of AttnMode.PREFILL, keeping the mode aligned with the
metadata produced by init_forward_metadata and planner.plan().
In `@phyai/tests/layers/attention/ar/test_radix_gather_e2e.py`:
- Around line 29-76: The radix attention test is missing explicit planner
cleanup, so add planner.release() for each RadixSequence scenario to mirror the
intended lifecycle. Use the RadixAttentionPlanner release path after the seed
sequence a, the decoy sequence d, and the reuse sequence c, so prefix-reuse
locks and uncommitted suffix units are dropped deterministically. Keep the calls
near the end of each phase in test_radix_reuse_gathers_exact_prefix_plus_suffix
to make the test resilient to future leak-detection checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 3c7e2225-3b17-4c75-8088-b3d0ba74c21f
📒 Files selected for processing (8)
phyai/src/phyai/layers/attention/ar/__init__.pyphyai/src/phyai/layers/attention/ar/radix/__init__.pyphyai/src/phyai/layers/attention/ar/radix/planner.pyphyai/src/phyai/layers/attention/ar/radix/sequence.pyphyai/tests/layers/attention/ar/__init__.pyphyai/tests/layers/attention/ar/test_radix_flashinfer.pyphyai/tests/layers/attention/ar/test_radix_gather_e2e.pyphyai/tests/layers/attention/ar/test_radix_planner.py
1 participant
Summary
Wires the existing radix prefix cache (
phyai_ext.radix_cache.PrefixCache) into the AR (LM-side) attention path so requests reuse the KV of a shared prompt prefix and only attend over the new (suffix) tokens — the standard RadixAttention speedup. Model- and encoding-agnostic; pi0.5'sStaticCachepath is untouched. Foundation for the upcoming cosmos adaptation.What's in it
New components under
phyai/layers/attention/ar/radix/:atomsin; prefix/suffix slot split + radix handles out).PrefixCache, with a 3-call lifecycle.plan(seqs) -> ARAttnMetadata: Matches cached prefix → reuses those slots → allocates suffix-only → buildscu_seqlens_q/paged_kv_indptr/paged_kv_indices/write_indices/position_ids, with query = suffix tokens only.commit(seqs): Seeds a fresh sequence into the tree for future reuse.release(seqs): Drops the reused-prefix lock and frees the uncommitted suffix.paged_kv_indicesradix produces, so the planner just feeds it.phyai.layers.attention.ar, keeping the optionalphyai-extextra off the base attention import.Contract
Prefill-with-prefix: the caller feeds Q/K/V for the suffix tokens only; the flashinfer paged backend scatters the new K/V at
write_indicesand readsprefix_slots ++ suffix_slots. Causal masking forq_len < kv_lenis handled by flashinfer's paged prefill — the suffix query is aligned to the tail of the KV viaqo_indptrvspaged_kv_indptr.Scope / non-goals
page_size == 1(both enforced at construction).Tests
tests/layers/attention/ar/test_radix_planner.py— 21 CPU planner-logic tests (exact-tensor metadata, prefix reuse + non-contiguous indices, commit/release,plan()rollback, no-leak under nested reuse, capacity/validation errors, lazy/star-import safety).test_radix_gather_e2e.py— CPU, backend-free: after seed → commit → decoy → reuse, the gathered[cached prefix ++ fresh suffix]equals the full sequence's KV (asserts the exact non-contiguouspaged_kv_indices). Keeps CI numeric coverage without a kernel.test_radix_flashinfer.py— CUDA paged e2e (auto-skips without CUDA + flashinfer).tests/layers/attention/suite green on CPU: 54 passed, 14 skipped (skips are the CUDA-gated flashinfer / vgpu / norm-backend tests).Notes
Hardened across review rounds: device-only tier, lazy/star-import safety,
plan()atomicity (rollback + validate-before-side-effects), order-independent commit, the KV-pool bounds check, and the NodeRef split-leak fix (commit doesn't pin;plan()pre-splits the tree before locking).