feat(phyai): radix prefix-cache bridge for AR attention

[Glorhop]

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's StaticCache path is untouched. Foundation for the upcoming cosmos adaptation.

What's in it

• New components under phyai/layers/attention/ar/radix/:
• RadixSequence: Per-request state (pre-encoded atoms in; prefix/suffix slot split + radix handles out).
• RadixAttentionPlanner: The bridge over PrefixCache, featuring a 3-call lifecycle.
• plan(seqs) -> ARAttnMetadata: Matches cached prefix → reuses those slots → allocates suffix-only → builds cu_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.
• AR eager backend → gather-based (KVCachePool.gather_kv): Reads the non-contiguous slot layouts radix reuse produces (diffusion eager stays contiguous-only by design).
• Lazy re-exporting: Bridge is lazily re-exported from phyai.layers.attention.ar, keeping the optional phyai-ext extra off the base attention import.

Contract

Prefill-with-prefix: the caller feeds Q/K/V for the suffix tokens only; the backend scatters the new K/V at write_indices and reads prefix_slots ++ suffix_slots. Causal masking for q_len < kv_len is handled by the existing eager_attn tail alignment.

Scope / non-goals

• pi0.5 keeps StaticCache — this is opt-in, no runner is wired to it yet.
• Device tier only and page_size == 1 (both enforced at construction).
• Decode loop / incremental cache growth, multi-token pages, host/disk tiers: out of scope (would need a C++ insert_suffix_from_node for unit-level pinning).

Tests

• tests/layers/attention/ar/test_radix_planner.py — 22 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_eager_e2e.py — CPU numeric: reused-prefix output allclose a full recompute.
• test_radix_flashinfer.py — CUDA paged e2e (auto-skips without CUDA + flashinfer).
• Full tests/layers/attention/ suite green (61 passed; flashinfer skipped on CPU).

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).

[gemini-code-assist]

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.

[gemini-code-assist]

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.

Suggested change

	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]

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

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: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 41a101a4-75e9-47bd-9ba3-fe7d9e4fc7d7

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}