PCP + ring attention on top of DCP batched-RPA (#3255)

tests/kernels/batched_rpa/pcp_ring_test.py

@@ -0,0 +1,210 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PCP ring cache phase on the batched RPA kernel.
+
+The ring streams page-interleaved KV cache shards around the pcp mesh axis
+while every rank attends with its local Q, accumulating all rounds in one
+online softmax. Since softmax is order-invariant, the reference is plain
+CACHE_ONLY attention over the full cache laid out in token order.
+"""
+
+import functools
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest, parameterized
+from jax._src import test_util as jtu
+from jax.experimental.shard_map import shard_map
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as PS
+
+from tpu_inference.kernels.experimental.batched_rpa import \
+    configs as brpa_configs
+from tpu_inference.kernels.experimental.batched_rpa.utils import (
+    align_to, cp_local_cache_len, get_dtype_packing)
+from tpu_inference.kernels.experimental.batched_rpa.wrapper import \
+    ragged_paged_attention
+
+
+def cdiv(a, b):
+    return (a + b - 1) // b
+
+
+def merge_kv(k, v):
+    """[n, nkv, hd] x2 -> [n, nkv_x2 // packing, packing, padded_hd]."""
+    n, nkv, hd = k.shape
+    kv_packing = get_dtype_packing(k.dtype)
+    nkv_x2 = nkv * 2
+    nkv_x2_aligned = align_to(nkv_x2, kv_packing)
+    padded_hd = align_to(hd, 128)
+    kv = jnp.pad(
+        jnp.concat([k, v], axis=-1).reshape(n, nkv_x2, hd),
+        ((0, 0), (0, nkv_x2_aligned - nkv_x2), (0, padded_hd - hd)),
+        constant_values=0,
+    ).reshape(n, nkv_x2_aligned // kv_packing, kv_packing, padded_hd)
+    return kv
+
+
+class BatchedRpaPcpRingTest(jtu.JaxTestCase):
+
+    @parameterized.product(
+        dtype=[jnp.float32, jnp.bfloat16],
+        P=[2, 4, 8],
+    )
+    def test_pcp_ring_cache_phase_matches_full_cache(self, dtype, P):
+        if not jtu.is_device_tpu_at_least(version=4):
+            self.skipTest("Expect TPUv4+")
+        if jax.device_count() < P:
+            self.skipTest(f"Needs {P} devices")
+
+        lp = 128  # local (per-rank) page size
+        nq, nkv, hd = 8, 2, 128
+        q_len = 256
+        # Two full super-pages plus a 200-token remainder so rank shard
+        # lengths differ (rank 0 gets 128 of it, rank 1 gets 72, rest 0).
+        prev_len = 2 * P * lp + 200
+        kv_len = prev_len + q_len
+        max_num_seqs = 8
+        pages_per_seq = cdiv(prev_len, lp)
+        num_pages = pages_per_seq + 8
+        sm = hd**-0.5
+
+        rng = np.random.default_rng(1234)
+
+        def gen(shape):
+            return jnp.array(rng.random(size=shape,
+                                        dtype=np.float32)).astype(dtype)
+
+        q_all = gen((P, q_len, nq, hd))
+        k_prev = gen((prev_len, nkv, hd))
+        v_prev = gen((prev_len, nkv, hd))
+        merged = merge_kv(k_prev, v_prev)
+        dummy_kv = jnp.zeros((q_len, nkv, hd), dtype)
+
+        # Small blocks so both the lane spread (2 q blocks over 2 lanes) and
+        # the multi-block ring (2 blocks of 2 pages) are exercised.
+        blocks = brpa_configs.BlockSizes(bq_sz=128,
+                                         bq_c_sz=128,
+                                         bkv_sz=256,
+                                         batch_size=2,
+                                         n_buffer=3)
+
+        kv_lens = jnp.zeros(max_num_seqs, jnp.int32).at[0].set(kv_len)
+        cu_q_lens = jnp.zeros(max_num_seqs + 1, jnp.int32).at[1:].set(q_len)
+        distribution = jnp.array([0, 0, 1], jnp.int32)
+        page_indices = jnp.zeros(max_num_seqs * pages_per_seq,
+                                 jnp.int32).at[:pages_per_seq].set(
+                                     jnp.arange(pages_per_seq,
+                                                dtype=jnp.int32))
+
+        def cache_from_tokens(tokens):
+            """[n, nkv_x2 // pack, pack, hd] -> (num_pages, lp, ...) cache."""
+            n = tokens.shape[0]
+            npg = max(cdiv(n, lp), 1)
+            padded = jnp.pad(tokens, ((0, npg * lp - n), (0, 0), (0, 0),
+                                      (0, 0)))
+            cache = jnp.zeros((num_pages, lp, *tokens.shape[1:]), dtype)
+            return cache.at[:npg].set(
+                padded.reshape(npg, lp, *tokens.shape[1:]))
+
+        common = dict(
+            sm_scale=sm,
+            attention_scope=brpa_configs.AttentionScope.CACHE_ONLY,
+            update_kv_cache=False,
+            return_lse=True,
+            decode_block_sizes=blocks,
+            prefill_block_sizes=blocks,
+        )
+
+        # Reference: full cache in token order, no CP. kv_cache is donated,
+        # so each call gets its own copy.
+        ref_outs, ref_lses = [], []
+        for r in range(P):
+            o, _, lse = ragged_paged_attention(
+                q_all[r],
+                dummy_kv,
+                dummy_kv,
+                cache_from_tokens(merged),
+                kv_lens,
+                page_indices,
+                cu_q_lens,
+                distribution,
+                **common,
+            )
+            ref_outs.append(np.asarray(o[:q_len], np.float32))
+            ref_lses.append(np.asarray(lse[:q_len], np.float32))
+        ref_out = np.stack(ref_outs)
+        ref_lse = np.stack(ref_lses)
+
+        # Page-interleaved shards: local page j of rank r holds global tokens
+        # [(j * P + r) * lp, +lp), matching cp_local_cache_len.
+        def rank_tokens(r):
+            slices = []
+            for j in range(cdiv(prev_len, P * lp)):
+                start = (j * P + r) * lp
+                end = min(start + lp, prev_len)
+                if start < prev_len:
+                    slices.append(merged[start:end])
+            tokens = (jnp.concatenate(slices, axis=0)
+                      if slices else merged[:0])
+            expected = cp_local_cache_len(jnp.int32(prev_len), P, r, lp)
+            assert tokens.shape[0] == int(expected), (tokens.shape[0],
+                                                      int(expected))
+            return tokens
+
+        cache_sh = jnp.stack([cache_from_tokens(rank_tokens(r))
+                              for r in range(P)])
+
+        mesh = Mesh(np.array(jax.devices()[:P]), ("pcp", ))
+
+        @functools.partial(
+            shard_map,
+            mesh=mesh,
+            in_specs=(PS("pcp"), PS("pcp")),
+            out_specs=(PS("pcp"), PS("pcp")),
+            check_rep=False,
+        )
+        def ring(q_l, c_l):
+            r = jax.lax.axis_index("pcp")
+            o, _, lse = ragged_paged_attention(
+                q_l[0],
+                dummy_kv,
+                dummy_kv,
+                c_l[0],
+                kv_lens,
+                page_indices,
+                cu_q_lens,
+                distribution,
+                cp_rank=jax.lax.reshape(r, (1, )).astype(jnp.int32),
+                cp_group_size=P,
+                pcp_ring_axis_name="pcp",
+                pcp_ring_mesh_axis_names=("pcp", ),
+                **common,
+            )
+            return o[None], lse[None]
+
+        out, lse = jax.jit(ring)(q_all, cache_sh)
+        out = np.asarray(out[:, :q_len], np.float32)
+        lse = np.asarray(lse[:, :q_len], np.float32)
+
+        self.assertTrue(np.all(np.isfinite(out)))
+        self.assertGreater(float(np.abs(out).max()), 0.0)
+        tol = 2e-3 if dtype == jnp.float32 else 2e-2
+        self.assertAllClose(out, ref_out, atol=tol, rtol=tol)
+        self.assertAllClose(lse, ref_lse, atol=tol, rtol=tol)
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py

@@ -565,6 +565,110 @@ def fn(k, v):
         self.assertLess(float((recon == 0).mean()), 0.5)
         self.assertArraysEqual(recon, ref)
 
+    @parameterized.product(dtype=[jnp.float32, jnp.bfloat16], P=[2, 4, 8])
+    def test_pcp_ring_cache_phase_matches_full_cache(self, dtype, P):
+        self._ring_vs_full_cache(dtype, P)
+
+    def _ring_vs_full_cache(self, dtype, P):
+        """In-kernel ring over the striped cache == the same local Q attending
+        the whole un-striped cache.
+
+        This is the property that makes the ring a drop-in cache phase: no Q
+        all-gather and no output collective, so each rank's result must already
+        be the full-cache answer for its own tokens.
+        """
+        if jax.device_count() < P:
+            self.skipTest(f"needs >= {P} devices")
+        self.PAGE = 64
+        # Lprev % P != 0 so the ranks' stripes have different lengths and a
+        # round that uses the wrong originating rank's length is visible.
+        Lprev, C, nq, nkv, hd = 1002, 256, 8, 2, 128
+        Scur = C * P
+        kv_total = Lprev + Scur
+        pps = cdiv(kv_total, self.PAGE)
+        sm = hd**-0.5
+        rng = np.random.default_rng(11)
+        k_prev = self._rand(rng, (Lprev, nkv, hd), dtype)
+        v_prev = self._rand(rng, (Lprev, nkv, hd), dtype)
+        q_all = self._rand(rng, (P, C, nq, hd), dtype)  # rank-major local Q
+
+        kv_lens = self._pad1([kv_total])
+        kv_cache_lens = self._pad1([Lprev])
+        pi = self._pi(pps)
+        dist = jnp.array([0, 0, 1], jnp.int32)
+        cu = self._padcu([0, C])
+        # Small blocks so both the bq loop and the multi-block ring run.
+        blocks = (128, 128, 128, 128)
+        dummy_kv = jnp.zeros((C, nkv, hd), dtype)
+
+        # Reference: plain paged attention over the whole cache, per rank.
+        # kv_cache is donated, so each call needs its own (identical) copy.
+        ref = np.stack([
+            np.asarray(
+                ragged_paged_attention(q_all[r],
+                                       dummy_kv,
+                                       dummy_kv,
+                                       self._cache_from_kv(
+                                           k_prev, v_prev, Lprev, dtype),
+                                       kv_lens,
+                                       pi,
+                                       cu,
+                                       dist,
+                                       kv_cache_lens=kv_cache_lens,
+                                       cp_rank=jnp.array([0], jnp.int32),
+                                       cp_group_size=1,
+                                       skip_current_attn=True,
+                                       use_causal_mask=False,
+                                       update_kv_cache=False,
+                                       return_lse=True,
+                                       sm_scale=sm,
+                                       m_block_sizes=blocks)[0], np.float32)
+            for r in range(P)
+        ])
+
+        # Ring: rank r holds global cache tokens r, r+P, r+2P, ...
+        cache_sh = jnp.stack([
+            self._cache_from_kv(k_prev[r::P], v_prev[r::P],
+                                k_prev[r::P].shape[0], dtype) for r in range(P)
+        ])
+        mesh = Mesh(np.array(jax.devices()[:P]), ("pcp", ))
+        qsp = PS("pcp", None, None, None)
+        csp = PS("pcp", None, None, None, None)
+
+        @partial(shard_map,
+                 mesh=mesh,
+                 in_specs=(qsp, csp),
+                 out_specs=qsp,
+                 check_rep=False)
+        def ring(q_l, c_l):
+            r = jax.lax.axis_index("pcp")
+            o, _, _ = ragged_paged_attention(q_l[0],
+                                             dummy_kv,
+                                             dummy_kv,
+                                             c_l[0],
+                                             kv_lens,
+                                             pi,
+                                             cu,
+                                             dist,
+                                             kv_cache_lens=kv_cache_lens,
+                                             cp_rank=jax.lax.reshape(
+                                                 r, (1, )).astype(jnp.int32),
+                                             cp_group_size=P,
+                                             pcp_ring_axis_name="pcp",
+                                             skip_current_attn=True,
+                                             use_causal_mask=False,
+                                             update_kv_cache=False,
+                                             return_lse=True,
+                                             sm_scale=sm,
+                                             m_block_sizes=blocks)
+            return o[None]
+
+        out = np.asarray(jax.jit(ring)(q_all, cache_sh), np.float32)
+        self.assertTrue(np.all(np.isfinite(out)))
+        self.assertGreater(float(np.abs(out).max()), 0.0)
+        tol = 2e-3 if dtype == jnp.float32 else 2e-2
+        self.assertAllClose(out, ref, atol=tol, rtol=tol)
+
 
 if __name__ == "__main__":
     absltest.main(testLoader=jtu.JaxTestLoader())

tests/layers/common/test_pcp_attention_interface.py

@@ -269,6 +269,20 @@ def _assert_matches(self, out, exp, pcp, C, num_current):
         self.assertAllClose(got, exp, atol=2e-2, rtol=2e-2)
 
     # ------------------------------ tests ------------------------------------
+    @parameterized.product(pcp=[2, 4])
+    def test_ring_cache_phase_matches_reference(self, pcp):
+        """The ring cache phase must reproduce the full-causal reference.
+
+        The ring streams each rank's KV shard around the pcp axis instead of
+        materializing the cache, so any divergence is a synchronization or
+        masking bug in the rotation, not a modelling choice.
+        """
+        if jax.device_count() < pcp:
+            self.skipTest(f"needs >= {pcp} devices")
+        L, S = 512, 128
+        out, _, exp, C = self._run(pcp, L, S, S)
+        self._assert_matches(out, exp, pcp, C, S)
+
     @parameterized.product(pcp=[2, 4])
     def test_chunked_prefill(self, pcp):
         """Wrapper output == full-causal reference, for a chunked prefill: L

tpu_inference/kernels/experimental/batched_rpa/configs.py

@@ -94,6 +94,11 @@ class ServingConfigs:
     attention_scope: AttentionScope = AttentionScope.FULL
     return_lse: bool = False
     update_kv_cache: bool = True
+    # PCP ring cache phase: when set, CACHE_ONLY streams each rank's KV cache
+    # shard around this mesh axis so every rank attends the full cache with
+    # its local Q, accumulating all rounds in one online softmax.
+    pcp_ring_axis_name: str | None = None
+    pcp_ring_mesh_axis_names: tuple[str, ...] | None = None
 
     @property
     def pages_per_seq(self) -> int:
@@ -192,6 +197,11 @@ def n_buffer(self) -> int:
 
     # Define derived values.
 
+    @property
+    def ring_enabled(self) -> bool:
+        return (self.serve.pcp_ring_axis_name is not None
+                and self.serve.attention_scope == AttentionScope.CACHE_ONLY)
+
     @property
     def max_steps_ub(self) -> int:
         """Get maximum upper bound of kernel steps based on SMEM limit."""
@@ -443,3 +453,27 @@ def validate_inputs(
                 raise ValueError(
                     "Context Parallel does not support sliding window right now"
                 )
+
+        if self.serve.pcp_ring_axis_name is not None:
+            if self.serve.cp_group_size is None:
+                raise ValueError(
+                    "pcp_ring_axis_name requires cp_group_size to be set.")
+            if self.serve.cp_group_size % 2 != 0:
+                # The ring double-buffers by round parity; an odd group size
+                # would collide the incoming block with the round-0 fill.
+                raise ValueError(
+                    "pcp_ring_axis_name requires an even cp_group_size, got"
+                    f" {self.serve.cp_group_size}.")
+            if self.serve.attention_scope != AttentionScope.CACHE_ONLY:
+                raise ValueError(
+                    "pcp_ring_axis_name is a cache-phase path and requires"
+                    " AttentionScope.CACHE_ONLY.")
+            if self.serve.update_kv_cache:
+                raise ValueError(
+                    "pcp_ring_axis_name requires update_kv_cache=False; the"
+                    " ring never writes the cache.")
+            if self.serve.kv_layout != KVLayout.HEAD_ALONG_SUBLANE:
+                raise NotImplementedError(
+                    "pcp_ring_axis_name only supports HEAD_ALONG_SUBLANE;"
+                    " SEQ_ALONG_LANE stitches new KV in-place in the block"
+                    " buffer, which would corrupt rotated blocks.")