PCP: in-kernel ring cache phase, fused with the current phase

[wenxindongwork]

Summary

Adds cache_phase="ring": an in-kernel ring cache phase for rpa_v3_cp that
runs the ring cache rounds AND the causal current phase in ONE kernel launch
with a shared online softmax — no separate current-phase launch and no
merge_attn_states LSE merge. This is the only ring variant this PR ships:
an earlier revision also carried a two-launch ring (cache-only ring launch +
LSE merge), which was strictly slower and has been removed to keep the change
minimal.

Design

The launch keeps the two-seq (head, tail) structure of the current phase, with
both seqs forced to a full C so every rank runs the identical ring schedule
(pad rows are discarded by the caller, like the existing all-pad-tail case).
Per bq block, the bkv loop walks a virtual index v:

• v in [0, ring_num_bkv * P): ring rounds over the rotating cache shards
  (each rank's striped shard travels the ring via in-kernel
  make_async_remote_copy, double-buffered in the existing bkv slots);
• v in [ring_region, end): the current-KV blocks, causal, with the
  skip-cache lower bound applied at runtime.

Region behavior is selected by VALUE, not control flow: zero-size DMAs for the
non-active fetch kind, causal | in_ring for the mask, scalar-selected
effective_kv_len / lower bound — no lax.cond (which measurably regresses
in this kernel). Slot parity stays uniform across the region boundary
(slot = v mod 2; the ring region length is even because P is even).

The region behavior is selected by VALUE inside the shared code paths: the
current-region cache-mask is the existing skip_cache_attn bound made
runtime, fetch gating zeroes DMA sizes, and one in_ring scalar drives all
of it. Three synchronization/layout subtleties, the first handled in the
wrapper and the other two with dedicated regression tests:

1. Token-order current KV. The non-ring current launch takes rank-order
   KV (contiguous only within one head/tail chunk), safe because its block
   size is capped at C. The ring launch keeps ring-sized bkv (~2 MB hops)
   that cross chunks, so the wrapper hands it TOKEN-order KV instead and the
   kernel's fetch stays a plain contiguous copy.
2. Sync-release schedule. Non-last ring blocks release their slot at round
   P-1 (the next writer is the rotation itself). Only the LAST ring block's
   slot stays live through the current region, so only its release is
   deferred to the end of the bq block loop — keeping signals per bq equal to
   waits (a per-bq-only release deadlocks multi-block rings).
3. (seq, bq)-crossing prefetch. The successor of a bq's last block is the
   next bq's ring block 0 at slot 0 — potentially the same slot the last
   current block is still computing from. The crossing prefetch is issued
   after that block's attention instead of before.

Performance

v7x-8, per-request cumulative prefill TTFT (ms), CH=4k, pcp=8 tp=1, NQ=32
NKV=2 HD=256 bf16 ("two-launch ring" is the removed variant, kept as a
reference point):

Context	two-launch ring	ring (this PR)	gather_kv
8k	3.17	2.55	2.77
16k	8.13	6.29	6.72
32k	21.42	17.23	16.19
64k	61.40	53.05	41.32
128k	195.8	180.6	121.1
256k	681.6	661.0	418.0
1M	9616	9769	6033

Fusing removes the fixed per-step cost of the second launch + LSE merge and
hides the first ring hop: 14–23% faster than the two-launch ring at
8k–64k, and the fastest strategy outright at 8k–16k (below gather_kv's
range). At 512k–1M the ring's asymptote is ICI link utilization and per-bq
cache re-streaming, which launch structure does not touch; deeper pipelining
(HBM-spilled (m,l,acc)) remains the follow-up for long-context parity.

ring stays opt-in and is never chosen by the auto heuristic; if the 8k–16k
win holds up on more shapes, wiring it into the auto ladder
(ring < ~16k < gather_kv < ~128k < gather_q) is a natural follow-up.

Testing

• tests/layers/common/test_pcp_attention_interface.py:
  • test_cache_phase_strategies_agree: ring vs gather_kv vs gather_q must
    agree (P=2, 4).
  • test_ring_multi_block: cache large enough that each rank's shard spans
    several ring bkv blocks — reaches the sync-release schedule and the
    crossing prefetch, which small-cache configs cannot.
• Standalone wrapper-level check (vs gather_kv on random data, P=2 and P=4,
  five configs incl. partial/all-pad tails and a boundary block with
  cache_len not bkv-aligned): max |Δ| 3.0–4.9e-4 — the same magnitude as an
  LSE-merge's own noise — negative control (different cache contents) 50–100x
  larger, and the strided KV cache write is BIT-EXACT vs the gather_kv path.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.