[PCP] Move the host-side preprocessing into runner/pcp_utils.py

tests/runner/test_pcp_utils.py

@@ -0,0 +1,186 @@
+# Copyright 2025 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.
+"""Tests for runner/pcp_utils.py; `test_prepare_inputs` needs pcp_size devices."""
+
+import jax
+import numpy as np
+import pytest
+from jax.sharding import Mesh
+from vllm.utils.math_utils import cdiv, next_power_of_2
+
+from tpu_inference.runner.pcp_utils import (PCPPreprocessor, pcp_batch_layout,
+                                            pcp_buffer_tokens,
+                                            pcp_max_buffer_tokens,
+                                            pcp_seq_arrays, pcp_token_layout,
+                                            pcp_token_permutation)
+
+# (pcp_size, scheduled tokens per request)
+LAYOUTS = [
+    (2, [5]),
+    (2, [4096]),
+    (2, [22061, 3000]),
+    (4, [7, 1, 300]),
+    (2, [1, 1]),
+    (8, [2539, 1200, 7, 3000, 64, 1, 500, 999]),
+]
+
+
+def _t_pad(counts, pcp):
+    """A power-of-two bucket that holds the layout and is divisible by 2P."""
+    return next_power_of_2(max(pcp_buffer_tokens(counts, pcp), 2 * pcp))
+
+
+def _layout(counts, pcp):
+    """(t_pad, chunk, off) of a batch in the bucket `_t_pad` picks."""
+    t_pad = _t_pad(counts, pcp)
+    chunk, off = pcp_batch_layout(counts, t_pad, pcp)
+    return t_pad, chunk, off
+
+
+@pytest.mark.parametrize("pcp,counts", LAYOUTS)
+def test_token_layout(pcp, counts):
+    chunk, off, s_live = pcp_token_layout(counts, pcp)
+    assert chunk == [cdiv(n, 2 * pcp) for n in counts]
+    assert off == list(np.cumsum([0] + [2 * c for c in chunk])[:-1])
+    assert s_live == sum(2 * c for c in chunk)
+    assert pcp_buffer_tokens(counts, pcp) == pcp * s_live
+    assert pcp_buffer_tokens(counts, pcp) <= pcp_max_buffer_tokens(
+        sum(counts), len(counts), pcp)
+
+
+@pytest.mark.parametrize("pcp,counts", LAYOUTS)
+def test_batch_layout(pcp, counts):
+    t_pad, chunk, off = _layout(counts, pcp)
+    if len(counts) == 1:
+        # Single request: the chunk comes from the buffer width.
+        assert chunk == [t_pad // (2 * pcp)] and off == [0]
+    else:
+        assert (chunk, off) == pcp_token_layout(counts, pcp)[:2]
+
+
+def test_batch_layout_rejects_short_buffer():
+    with pytest.raises(AssertionError):
+        pcp_batch_layout([100, 100], 64, 2)
+
+
+@pytest.mark.parametrize("pcp,counts", LAYOUTS)
+def test_token_permutation(pcp, counts):
+    t_pad, chunk, off = _layout(counts, pcp)
+    s_pad = t_pad // pcp
+    perm, kv_order = pcp_token_permutation(counts, chunk, off, t_pad, pcp)
+    total = sum(counts)
+    # Every real token lands in exactly one slot; everything else is padding.
+    assert sorted(perm[perm >= 0].tolist()) == list(range(total))
+    src_off = np.cumsum([0] + counts)[:-1]
+    for i, n_i in enumerate(counts):
+        c_i = chunk[i]
+        for tok in range(n_i):
+            slot = kv_order[pcp * off[i] + tok]
+            # kv_order undoes perm on the live rows.
+            assert perm[slot] == src_off[i] + tok
+            # Zigzag: chunk k sits on rank k (head) or 2P-1-k (tail).
+            k = tok // c_i
+            rank = k if k < pcp else 2 * pcp - 1 - k
+            assert slot // s_pad == rank
+    # Slots reserved for a request are distinct across the request.
+    for i in range(len(counts)):
+        lo, hi = pcp * off[i], pcp * off[i] + 2 * pcp * chunk[i]
+        assert len(set(kv_order[lo:hi].tolist())) == hi - lo
+
+
+@pytest.mark.parametrize("pcp,counts", LAYOUTS)
+def test_seq_arrays(pcp, counts):
+    _, chunk, off = _layout(counts, pcp)
+    n_off = 2 * len(counts) + 3
+    cu_row, q_pos, kv_new_starts = pcp_seq_arrays(chunk, off, pcp, n_off)
+    n_seqs = 2 * len(counts)
+    assert cu_row.shape == (n_off + 1, )
+    assert np.all(np.diff(cu_row[:n_seqs + 1]) > 0)
+    assert np.all(cu_row[n_seqs:] == cu_row[n_seqs])
+    for i, c_i in enumerate(chunk):
+        assert cu_row[2 * i] == off[i]
+        assert cu_row[2 * i + 1] - cu_row[2 * i] == c_i
+        assert cu_row[2 * i + 2] - cu_row[2 * i + 1] == c_i
+        for r in range(pcp):
+            assert q_pos[r, 2 * i] == r * c_i
+            assert q_pos[r, 2 * i + 1] == (2 * pcp - 1 - r) * c_i
+        assert kv_new_starts[2 * i] == kv_new_starts[2 * i + 1] == pcp * off[i]
+    assert np.all(q_pos[:, n_seqs:] == 0)
+
+
+@pytest.mark.parametrize("pcp,counts", LAYOUTS)
+def test_prepare_inputs(pcp, counts):
+    if len(jax.devices()) < pcp:
+        pytest.skip(f"needs {pcp} devices")
+    mesh = Mesh(np.array(jax.devices()[:pcp]), ("pcp", ))
+    pre = PCPPreprocessor(pcp, mesh, [1, 8])
+
+    t_pad, chunk, off = _layout(counts, pcp)
+    n_reqs = len(counts)
+    n_off = 2 * 8
+    # Cached prefixes on the longer requests (a cached 1-token request is a
+    # decode, rejected below).
+    computed = [(3 * i) % 40 if n > 1 else 0 for i, n in enumerate(counts)]
+    total = sum(counts)
+    # Natural-order buffers: token g carries id 1000 + g and position g.
+    input_ids = np.zeros(t_pad, np.int32)
+    input_ids[:total] = 1000 + np.arange(total)
+    positions = np.zeros(t_pad, np.int32)
+    positions[:total] = np.arange(total)
+    seq_lens = np.full(n_off, 7, np.int32)
+    request_distribution = np.array([n_reqs, 0, 0], np.int32)
+    logits_indices = np.full(8, 5, np.int32)
+
+    md = pre.prepare_inputs(counts, computed, t_pad, positions, input_ids,
+                            seq_lens, request_distribution, logits_indices)
+
+    perm, kv_order = pcp_token_permutation(counts, chunk, off, t_pad, pcp)
+    live = perm >= 0
+    assert np.array_equal(input_ids[live], 1000 + perm[live])
+    assert np.array_equal(positions[live], perm[live])
+    assert np.all(input_ids[~live] == 0) and np.all(positions[~live] == 0)
+
+    n_seqs = 2 * n_reqs
+    assert np.array_equal(
+        seq_lens[:n_seqs],
+        np.repeat([n + c for n, c in zip(counts, computed)], 2))
+    assert np.all(seq_lens[n_seqs:] == 0)
+    assert request_distribution.tolist() == [0, 0, n_seqs]
+    # Each request's logits slot holds its last real token.
+    src_off = np.cumsum([0] + counts)[:-1]
+    assert np.array_equal(perm[logits_indices[:n_reqs]],
+                          src_off + np.asarray(counts) - 1)
+    assert np.all(logits_indices[n_reqs:] == -1)
+
+    assert md.query_start_loc.shape == (pcp, n_off + 1)
+    assert md.q_pos_offsets.shape == (pcp, n_off)
+    assert np.array_equal(
+        np.asarray(md.kv_cache_lens)[:n_seqs], np.repeat(computed, 2))
+    assert md.has_cached_kv == (max(computed) > 0)
+    assert md.num_reqs == (1 if n_reqs == 1 else 8)
+    assert np.array_equal(np.asarray(md.kv_token_order), kv_order)
+    assert np.array_equal(
+        np.asarray(md.kv_new_starts)[:n_seqs],
+        np.repeat([pcp * o for o in off], 2))
+
+
+def test_prepare_inputs_rejects_decode():
+    if len(jax.devices()) < 2:
+        pytest.skip("needs 2 devices")
+    mesh = Mesh(np.array(jax.devices()[:2]), ("pcp", ))
+    pre = PCPPreprocessor(2, mesh, [1, 8])
+    # Rejected before any buffer is touched, so shapes do not matter.
+    buf = np.zeros(16, np.int32)
+    with pytest.raises(NotImplementedError):
+        pre.prepare_inputs([1], [5], 16, buf, buf, buf, buf, buf)

tests/runner/test_tpu_runner_dp.py

@@ -32,11 +32,12 @@ def setup_method(self):
 
         # Basic DP configuration
         self.runner.dp_size = 2
-        # PCP is off in these DP tests; without this the MagicMock auto-attr
-        # fails the `prefill_cp_size > 1` comparison in _prepare_inputs.
+        # PCP off; a MagicMock auto-attribute would not be None.
         self.runner.vllm_config.sharding_config.prefill_cp_size = 1
+        self.runner.pcp_preprocessor = None
         self.runner.max_num_tokens = 64
         self.runner.max_num_reqs = 8
+        self.runner.attn_max_num_seqs = 8
         self.runner.max_num_blocks_per_req = 8
         self.runner.num_tokens_paddings = [16, 32, 64]
 
@@ -1421,7 +1422,9 @@ def test_prepare_inputs_dp_uses_attn_data_sharding_for_sampling_metadata(
         runner.input_batch.max_decode_tokens = 1
         runner.dp_size = 2
         runner.vllm_config.sharding_config.prefill_cp_size = 1
+        runner.pcp_preprocessor = None
         runner.max_num_reqs = 8
+        runner.attn_max_num_seqs = 8
         runner.max_num_blocks_per_req = 8
         runner.speculative_config = None
         runner.input_batch.num_reqs = 2

tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py

@@ -319,6 +319,8 @@ def _ragged_paged_attention_kernel_loop(
     bkv_update_ids_ref,  # [6 or 8] (bkv_sem_0_seq_idx, bkv_sem_1_seq_idx, bkv_sem_0_offset, bkv_sem_1_offset, bkv_sem_0_sz, bkv_sem_1_sz) (bkv_smem_0_src_start_base, bkv_smem_1_src_start_base)
     cp_rank_ref: jax.Array | None,  # i32[1]
     q_pos_offset_ref: jax.Array | None,  # i32[max_num_seqs]
+    kv_new_starts_ref: jax.Array | None,  # i32[max_num_seqs]
+    kv_write_seq_mask_ref: jax.Array | None,  # i32[max_num_seqs]
     # Input
     q_hbm_ref,  # [actual_num_kv_heads, max_num_tokens, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
     kv_hbm_ref,  # [max_num_tokens, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
@@ -349,7 +351,6 @@ def _ragged_paged_attention_kernel_loop(
     cp_group_size: int | None = None,
     use_causal_mask: bool = True,
     update_kv_cache: bool = True,
-    write_last_seq_only: bool = False,
     skip_kv_mask: bool = False,
     skip_cache_attn: bool = False,
     skip_current_attn: bool = False,
@@ -437,7 +438,11 @@ def get_kv_new_len(seq_idx):
         return cu_q_lens_ref[seq_idx + 1] - cu_q_lens_ref[seq_idx]
 
     def get_kv_new_end(seq_idx):
+        # End of this seq's block inside the new-KV buffer (a read offset into
+        # kv_hbm_ref); under PCP kv_new_starts gives each seq's base.
         if kv_cache_lens_ref is not None:
+            if kv_new_starts_ref is not None:
+                return kv_new_starts_ref[seq_idx] + get_kv_new_len(seq_idx)
             return get_kv_new_len(seq_idx)
         return cu_q_lens_ref[seq_idx + 1]
 
@@ -1383,10 +1388,11 @@ def start_rotate():
                     # PCP fuses a request's head+tail chunks into ONE launch as
                     # two "sequences" that share the same request (same
                     # kv_lens/kv_cache_lens), so each would write the SAME
-                    # strided current KV. Write on exactly one of them.
-                    if write_last_seq_only:
-                        _do_write = jnp.logical_and(_do_write,
-                                                    seq_idx == end_seq_idx - 1)
+                    # strided current KV. Write on exactly one of them --
+                    # kv_write_seq_mask marks which (each request's tail).
+                    if kv_write_seq_mask_ref is not None:
+                        _do_write = jnp.logical_and(
+                            _do_write, kv_write_seq_mask_ref[seq_idx] != 0)
 
                     @pl.when(_do_write)
                     def update_cur_bkv_to_cache():
@@ -1806,6 +1812,9 @@ def static_validate_inputs(
     *,
     kv_cache_lens: jax.Array | None = None,  # i32[max_num_seqs] - PCP
     q_pos_offsets: jax.Array | None = None,  # i32[max_num_seqs] - PCP
+    kv_new_starts: jax.Array | None = None,  # i32[max_num_seqs] - PCP
+    kv_write_seq_mask: jax.Array | None = None,  # i32[max_num_seqs] - PCP
+    pcp_chunk_size: int | None = None,
     cp_group_size: int | None = None,
     cp_rank: jax.Array | int | None = None,
     pcp_ring_axis_name: str | None = None,
@@ -1993,6 +2002,26 @@ def _validate_block_sizes(block_sizes, prefix):
             raise NotImplementedError(
                 "pcp_ring_axis_name does not support sliding_window")
 
+    for name, arr in (("kv_new_starts", kv_new_starts), ("kv_write_seq_mask",
+                                                         kv_write_seq_mask)):
+        if arr is None:
+            continue
+        if arr.dtype != jnp.int32:
+            raise ValueError(
+                f"Expected int32 dtype for {name}, got {arr.dtype}")
+        if arr.shape != (max_num_seqs, ):
+            raise ValueError(
+                f"Expected {name}.shape to be ({max_num_seqs},), got {arr.shape}"
+            )
+
+    if kv_new_starts is not None:
+        if kv_cache_lens is None:
+            raise ValueError("PCP (kv_new_starts) requires kv_cache_lens.")
+        if pcp_chunk_size is not None:
+            raise ValueError(
+                "kv_new_starts and pcp_chunk_size are mutually exclusive: the "
+                "rank-order remap assumes a single request's new-KV buffer.")
+
     # No constraints for the following inputs.
     del sm_scale
     del mask_value
@@ -2152,7 +2181,6 @@ def get_default_block_sizes(
         "disable_bounds_checks",
         "disable_semaphore_checks",
         "update_kv_cache",
-        "write_last_seq_only",
         "cp_group_size",
         "pcp_chunk_size",
         "pcp_ring_axis_name",
@@ -2178,12 +2206,13 @@ def ragged_paged_attention(
     | None = None,  # i32[1] - per-device rank, sharded along the DCP axis
     cp_group_size: int | None = None,
     q_pos_offsets: jax.Array | None = None,  # i32[max_num_seqs]
+    kv_new_starts: jax.Array | None = None,  # i32[max_num_seqs]
+    kv_write_seq_mask: jax.Array | None = None,  # i32[max_num_seqs]
     pcp_chunk_size: int | None = None,
     pcp_ring_axis_name: str | None = None,
     pcp_ring_mesh_axis_names: tuple[str, ...] | None = None,
     use_causal_mask: bool = True,
     update_kv_cache: bool = True,
-    write_last_seq_only: bool = False,
     skip_kv_mask: bool = False,
     skip_cache_attn: bool = False,
     skip_current_attn: bool = False,
@@ -2235,12 +2264,17 @@ def ragged_paged_attention(
       KV cache around this axis.
     pcp_ring_mesh_axis_names: all axis names of the mesh the ring runs on, in
       order. Defaults to a one-axis mesh.
-    use_causal_mask: if true, use causal mask.
-    write_last_seq_only: PCP only. PCP fuses a request's head and tail chunk
-      into one launch as two "sequences" that are really the same request (same
+    kv_new_starts: PCP only. Base offset of each sequence's current-KV block
+      inside the all-gathered new-KV buffer (`keys`/`values`). Needed when that
+      buffer holds more than one request, packed back to back in request order;
+      leave None for a single request, where every block starts at 0.
+    kv_write_seq_mask: PCP only. Nonzero on the sequences that perform the fused
+      strided KV-cache write. PCP fuses a request's head and tail chunk into one
+      launch as two "sequences" that are really the same request (same
       kv_lens/kv_cache_lens), so each of them would redundantly write the same
-      strided current KV to the cache. When true, the write is performed by the
-      tail seq only.
+      strided current KV; the mask selects exactly one per request (its tail).
+      Leave None to let every sequence write, as in the non-PCP path.
+    use_causal_mask: if true, use causal mask.
     skip_kv_mask: only set to true if use_causal_mask=False and each dynamic
       kv_len % bkv_csz == 0. Set to true can improve performance.
     sm_scale: the softmax scale which will be applied to the Q@K^T.
@@ -2295,6 +2329,9 @@ def ragged_paged_attention(
         distribution,
         kv_cache_lens=kv_cache_lens,
         q_pos_offsets=q_pos_offsets,
+        kv_new_starts=kv_new_starts,
+        kv_write_seq_mask=kv_write_seq_mask,
+        pcp_chunk_size=pcp_chunk_size,
         cp_group_size=cp_group_size,
         cp_rank=cp_rank,
         pcp_ring_axis_name=pcp_ring_axis_name,
@@ -2453,7 +2490,9 @@ def run_rpa_kernel(
             init_bo_ids,
             init_bkv_update_ids,
             cp_rank if cp_group_size is not None else None,
-            q_pos_offsets)
+            q_pos_offsets,
+            kv_new_starts,
+            kv_write_seq_mask)
 
         num_scalers = len(scalar_prefetches)
         # None in scalar_prefetches contribute 0 pytree leaves, so
@@ -2497,7 +2536,6 @@ def run_rpa_kernel(
                 cp_group_size=cp_group_size,
                 pcp_ring_axis_name=pcp_ring_axis_name,
                 pcp_ring_mesh_axis_names=pcp_ring_mesh_axis_names,
-                write_last_seq_only=write_last_seq_only,
                 use_causal_mask=use_causal_mask,
                 skip_kv_mask=skip_kv_mask,
                 skip_cache_attn=skip_cache_attn,