diff --git a/tests/kernels/batched_rpa/pcp_ring_test.py b/tests/kernels/batched_rpa/pcp_ring_test.py new file mode 100644 index 0000000000..43a9c931ba --- /dev/null +++ b/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()) diff --git a/tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py b/tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py index 66e06db5fa..7dd47a80fe 100644 --- a/tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py +++ b/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()) diff --git a/tests/layers/common/test_pcp_attention_interface.py b/tests/layers/common/test_pcp_attention_interface.py index b7db2c039d..8b1ae0b593 100644 --- a/tests/layers/common/test_pcp_attention_interface.py +++ b/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 diff --git a/tpu_inference/kernels/experimental/batched_rpa/configs.py b/tpu_inference/kernels/experimental/batched_rpa/configs.py index 1d5a0ae909..858196ebc2 100644 --- a/tpu_inference/kernels/experimental/batched_rpa/configs.py +++ b/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.") diff --git a/tpu_inference/kernels/experimental/batched_rpa/kernel.py b/tpu_inference/kernels/experimental/batched_rpa/kernel.py index a37a45c518..7b91f4f8cf 100644 --- a/tpu_inference/kernels/experimental/batched_rpa/kernel.py +++ b/tpu_inference/kernels/experimental/batched_rpa/kernel.py @@ -164,6 +164,10 @@ def rpa_body( acc_scratch_ref: jax.Ref, lse_dma_sem_ref: jax.Ref | None, cp_rank_ref: jax.Array | None, + ring_kv_ref: jax.Ref | None = None, # [2, *kv_vmem_shape] + ring_dma_sems: jax.Ref | None = None, # [2] send/recv + ring_sync_sem: jax.Ref | None = None, + ring_local_sem: jax.Ref | None = None, # [1] *, # Passed refs cu_q_lens_ref: jax.Ref, @@ -174,6 +178,23 @@ def rpa_body( ): step = pl.program_id(0) + ring = cfgs.ring_enabled + if ring: + ring_size = cfgs.serve.cp_group_size + my_ring_id = lax.axis_index(cfgs.serve.pcp_ring_axis_name) + + def ring_device_id(rank): + if cfgs.serve.pcp_ring_mesh_axis_names is None: + return (rank, ) + return tuple( + rank if name == cfgs.serve.pcp_ring_axis_name else lax. + axis_index(name) + for name in cfgs.serve.pcp_ring_mesh_axis_names) + + ring_next_id = ring_device_id(lax.rem(my_ring_id + 1, ring_size)) + ring_prev_id = ring_device_id( + lax.rem(my_ring_id + ring_size - 1, ring_size)) + # Step 1: Fetch metadata. processed_q_len = [] processed_kv_len = [] @@ -185,6 +206,7 @@ def rpa_body( kv_new_start = [] # for new tokens. cache_len = [] # for cache only. int_ty = cfgs.serve.int_ty + ring_pos = jnp.int32(0) # common ring-encoded k_idx across valid lanes for b_idx in range(cfgs.batch_size): s_idx = schedule_ref.s_idx[step, b_idx] is_valid = s_idx != -1 @@ -195,6 +217,15 @@ def rpa_body( safe_s_idx = jnp.maximum(0, s_idx) q_idx = schedule_ref.q_idx[step, b_idx] k_idx = schedule_ref.k_idx[step, b_idx] + if ring: + # k_idx is ring-encoded as block * ring_size + round. All valid + # lanes carry the same encoded value at a given step (a single + # sequence's q blocks all walk the same ring), so the max below + # recovers it while masked lanes contribute 0. + ring_pos = jnp.maximum(ring_pos, + jnp.where(is_valid, k_idx, jnp.int32(0))) + ring_round_b = lax.rem(k_idx, ring_size) + k_idx = k_idx // ring_size k_id = jnp.where(is_valid, k_idx * cfgs.bkv_sz, 0) kv_len = jnp.where(is_valid, kv_lens_ref[safe_s_idx], 0) q_start = jnp.where(is_valid, cu_q_lens_ref[safe_s_idx], 0) @@ -208,6 +239,12 @@ def rpa_body( if (cfgs.serve.cp_group_size is not None): cp_group_size = cfgs.serve.cp_group_size rank = cp_rank_ref[0] + if ring: + # After r hops the resident block originated from rank + # (my_id - r) mod ring_size; mask by that rank's shard + # length instead of our own. + rank = lax.rem(my_ring_id + ring_size - ring_round_b, + ring_size) local_cache_len = utils.cp_local_cache_len( offset, cp_group_size, rank, cfgs.serve.page_size) # We add this to make sure flash_attention.py can mask non-cache tokens out. @@ -243,6 +280,11 @@ def rpa_body( start_k_idx = jnp.maximum(start_k_idx, offset // cfgs.bkv_sz) is_first_k_block = k_idx == start_k_idx + if ring: + # Rounds > 0 of block 0 continue the same online softmax; only + # the very first round of the first block resets the scratch. + is_first_k_block = jnp.logical_and(is_first_k_block, + ring_round_b == 0) reset_cond = jnp.logical_and(is_valid, is_first_k_block) m_scratch_ref[b_idx] = jnp.where(reset_cond, -jnp.inf, m_scratch_ref[b_idx]) @@ -250,6 +292,88 @@ def rpa_body( acc_scratch_ref[b_idx] = jnp.where(reset_cond, 0.0, acc_scratch_ref[b_idx]) + # Ring rotation setup. Every ring step computes on ring_kv_ref[round % 2]: + # round 0 fills slot 0 with a local copy of the pipeline's fetched block, + # and while round r computes, the resident block is sent to the next + # rank's other slot so round r+1's data arrives behind the compute. The + # even ring_size keeps the round parity consistent across block + # boundaries. A regular semaphore carries "your target slot is free" + # credits backwards along the ring (see release conditions below). + if ring: + ring_round = lax.rem(ring_pos, ring_size) + ring_slot = lax.rem(ring_round, 2) + is_last_round = ring_round == ring_size - 1 + is_first_launch = step == 0 + num_steps = jnp.minimum(schedule_ref.actual_steps[0], + cfgs.max_steps_ub) + is_last_launch = step == num_steps - 1 + + # Startup rendezvous with both neighbors before any remote traffic: + # on a cold first execution a fast device's RDMA could otherwise land + # on a neighbor that is still loading the program and lose part of + # the block to its startup initialization. + @pl.when(is_first_launch) + def ring_startup_barrier(): + barrier_sem = pltpu.get_barrier_semaphore() + pltpu.semaphore_signal( + barrier_sem, + 1, + device_id=ring_next_id, + device_id_type=pl.DeviceIdType.MESH, + ) + pltpu.semaphore_signal( + barrier_sem, + 1, + device_id=ring_prev_id, + device_id_type=pl.DeviceIdType.MESH, + ) + pl.semaphore_wait(barrier_sem, 2) + + # Credit release for the slot step t-1 was reading (the target of the + # previous rank's next send). Signaling here — at the top of the + # following step — rather than at the end of step t-1 keeps the + # release behind a full step boundary: the signal is issued by the + # scalar core, which can run ahead of the vector core's still + # in-flight reads of that slot, so an end-of-step release can let the + # incoming block clobber data mid-read. Last-round steps grant no + # credit (the previous rank sends nothing that round). + @pl.when(~is_first_launch & ~is_last_round) + def release_prev_slot(): + pl.semaphore_signal( + ring_sync_sem, + 1, + device_id=ring_prev_id, + device_id_type=pl.DeviceIdType.MESH, + ) + + @pl.when(ring_round == 0) + def fill_ring_slot0(): + cp = pltpu.make_async_copy(kv_in_vref, ring_kv_ref.at[0], + ring_local_sem.at[0]) + cp.start() + cp.wait() + + @pl.when(~is_last_round & ~is_first_launch) + def wait_ring_sync(): + pl.semaphore_wait(ring_sync_sem, 1) + + remote_op = pltpu.make_async_remote_copy( + src_ref=ring_kv_ref.at[ring_slot], + dst_ref=ring_kv_ref.at[1 - ring_slot], + send_sem=ring_dma_sems.at[0], + recv_sem=ring_dma_sems.at[1], + device_id=ring_next_id, + device_id_type=pl.DeviceIdType.MESH, + ) + + @pl.when(~is_last_round) + def start_rotate(): + remote_op.start() + + # Under the ring, compute reads the rotating buffer instead of the + # pipeline's block (which only round 0 touches, via the copy above). + kv_read_vref = ring_kv_ref.at[ring_slot] if ring else kv_in_vref + # Step 2: Fetch inputs. q_p = cfgs.aligned_num_q_heads_per_kv_head // cfgs.serve.packing_q q_ref = q_vref.bitcast(jnp.uint32).reshape(-1, cfgs.aligned_q_head_dim) @@ -315,7 +439,7 @@ def rpa_body( for kv_head_start in range(0, cfgs.model.num_kv_heads, heads_per_load): bkv_lst = strided_load_bkv( - kv_in_vref, + kv_read_vref, b_idx, kv_head_start * 2, cfgs=cfgs, @@ -386,6 +510,15 @@ def rpa_body( ) acc_scratch_ref[:, :, prev_q_slice] = o_next + if ring: + + @pl.when(~is_last_round) + def finish_rotate(): + # Waits both our send and the incoming block for the next round. + # The credit releasing this step's slot to the previous rank is + # granted at the top of the next step (see release_prev_slot). + remote_op.wait() + # Step 4: Write back outputs. calculate_and_store_out( step, @@ -588,6 +721,16 @@ def ragged_paged_attention_pipeline( ), # acc pltpu.SemaphoreType.DMA( (1, )) if return_lse else None, # lse_sem + pltpu.VMEM( + (2, *cfgs.kv_vmem_shape), + dtype=cfgs.serve.dtype_kv, + ) if cfgs.ring_enabled else None, # ring_kv + pltpu.SemaphoreType.DMA( + (2, )) if cfgs.ring_enabled else None, # ring send/recv + pltpu.SemaphoreType.REGULAR + if cfgs.ring_enabled else None, # ring_sync + pltpu.SemaphoreType.DMA( + (1, )) if cfgs.ring_enabled else None, # ring_local ), ) def _run(final_allocs, schedule_ref, dma_sem, scratches): @@ -626,14 +769,26 @@ def _run(final_allocs, schedule_ref, dma_sem, scratches): -1, num_lanes) kv_ref_flat[...] = jnp.zeros_like(kv_ref_flat) + if cfgs.ring_enabled: + # Same zero-init for the ring scratch: whatever a previous + # program left at this address must not be interpretable as + # plausible KV data. + ring_kv_flat = scratches[4].bitcast(jnp.uint32).reshape( + -1, num_lanes) + ring_kv_flat[...] = jnp.zeros_like(ring_kv_flat) + jax.tree.map(lambda x: x.wait(), dma_list) + # Body scratch order: schedule, m, l, acc, lse_sem, cp_rank, + # then the ring refs (ring_kv, ring send/recv, ring_sync, + # ring_local). pipeline_func( (q_hbm_ref, schedule_ref), (kv_cache_hbm_ref, new_kv_hbm_ref, schedule_ref, page_indices_ref), (o_hbm_ref, schedule_ref), - scratches=(schedule_ref, ) + scratches + (cp_rank_ref, ), + scratches=(schedule_ref, ) + scratches[:4] + + (cp_rank_ref, ) + scratches[4:], allocations=final_allocs, ) @@ -677,6 +832,10 @@ def _run(final_allocs, schedule_ref, dma_sem, scratches): compiler_params=pltpu.CompilerParams( vmem_limit_bytes=cfgs.vmem_limit_bytes, disable_bounds_checks=True, + # The ring's startup barrier needs a barrier semaphore. + **({ + "collective_id": 0 + } if cfgs.ring_enabled else {}), ), input_output_aliases=input_output_aliases, name=get_kernel_name(cfgs), diff --git a/tpu_inference/kernels/experimental/batched_rpa/schedule.py b/tpu_inference/kernels/experimental/batched_rpa/schedule.py index 9ccd04794e..a22138abf3 100644 --- a/tpu_inference/kernels/experimental/batched_rpa/schedule.py +++ b/tpu_inference/kernels/experimental/batched_rpa/schedule.py @@ -353,8 +353,17 @@ def k_loop( schedule.dma_q[step, target_lane, 0] = q_src schedule.dma_q[step, target_lane, 1] = q_sz_task - kv_len_start = k_idx * cfgs.bkv_sz - kv_p_start = k_idx * cfgs.bkv_p + if cfgs.ring_enabled: + # k_idx is ring-encoded as block * cp_group_size + round. Only + # round 0 fetches this rank's block from HBM; later rounds receive + # the block from the previous rank over the ring. + ring_block = k_idx // cfgs.serve.cp_group_size + ring_is_round0 = k_idx % cfgs.serve.cp_group_size == 0 + kv_len_start = ring_block * cfgs.bkv_sz + kv_p_start = ring_block * cfgs.bkv_p + else: + kv_len_start = k_idx * cfgs.bkv_sz + kv_p_start = k_idx * cfgs.bkv_p kv_left = k_len - kv_len_start if update_kv_cache: kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0) @@ -373,6 +382,8 @@ def k_loop( dst_vmem = i << cfgs.serve.page_size_log2 dma_sz = kv_left_frm_cache - dst_vmem dma_sz = jnp.clip(dma_sz, 0, cfgs.serve.page_size) + if cfgs.ring_enabled: + dma_sz = jnp.where(ring_is_round0, dma_sz, 0) src_hbm = jnp.minimum(p_offset + i, cfgs.serve.num_page_indices - 1) @@ -565,6 +576,19 @@ def q_loop(q_idx, _, *, s_idx, q_start, q_end, k_len, q_len, num_k): k_len = local_cache_len end_k_idx = jnp.minimum(end_k_idx, pl.cdiv(local_cache_len, cfgs.bkv_sz)) + if cfgs.ring_enabled: + # Ring: every rank must run the same number of steps, so size + # the block loop by rank 0's shard (the longest under + # page-interleaving) and run cp_group_size rounds per block; + # short ranks' tails are masked in the kernel. k_len stays at + # this rank's local length so round-0 fetch sizes clip to the + # pages this rank actually owns. + rank0_cache_len = utils.cp_local_cache_len( + cache_len, cfgs.serve.cp_group_size, 0, + cfgs.serve.page_size) + num_ring_blocks = pl.cdiv(rank0_cache_len, cfgs.bkv_sz) + start_k_idx = 0 + end_k_idx = num_ring_blocks * cfgs.serve.cp_group_size k_loop_fn = functools.partial( k_loop, diff --git a/tpu_inference/kernels/experimental/batched_rpa/wrapper.py b/tpu_inference/kernels/experimental/batched_rpa/wrapper.py index 82ff48f040..4228e7e639 100644 --- a/tpu_inference/kernels/experimental/batched_rpa/wrapper.py +++ b/tpu_inference/kernels/experimental/batched_rpa/wrapper.py @@ -194,6 +194,8 @@ def get_kv_cache_shape( "cp_group_size", "attention_scope", "return_lse", + "pcp_ring_axis_name", + "pcp_ring_mesh_axis_names", ), # Donation of transient inputs can fail for some runtime buffer layouts in # the experimental tuning path. Keep donation only for kv_cache, which is @@ -230,6 +232,8 @@ def ragged_paged_attention( cp_rank: jax.Array | None = None, attention_scope: configs.AttentionScope = configs.AttentionScope.FULL, return_lse: bool = False, + pcp_ring_axis_name: str | None = None, + pcp_ring_mesh_axis_names: tuple[str, ...] | None = None, ) -> tuple[jax.Array, jax.Array] | tuple[jax.Array, jax.Array, jax.Array]: """Perform batched ragged paged attention. @@ -271,6 +275,13 @@ def ragged_paged_attention( attention_scope: Which KV positions to attend to. FULL attends all positions, CACHE_ONLY skips new tokens, NEW_TOKENS_ONLY skips cached tokens. Defaults to FULL. return_lse: If True, return log-sum-exp (lse) values along with the output. Defaults to False. + pcp_ring_axis_name: PCP cache phase only. When set, CACHE_ONLY streams + each rank's KV cache shard around this mesh axis in-kernel so every + rank attends the full cache with its local Q; one online softmax + accumulates all rounds. Requires CACHE_ONLY, an even cp_group_size, + update_kv_cache=False, and the HEAD_ALONG_SUBLANE layout. + pcp_ring_mesh_axis_names: All axis names of the mesh the ring runs on, + in order. Defaults to a one-axis mesh. Returns: out: [max_num_tokens, num_q_heads, head_dim]. Output of self attention. @@ -338,6 +349,8 @@ def ragged_paged_attention( attention_scope=attention_scope, return_lse=return_lse, update_kv_cache=update_kv_cache, + pcp_ring_axis_name=pcp_ring_axis_name, + pcp_ring_mesh_axis_names=pcp_ring_mesh_axis_names, ) q_hbm, new_kv_hbm = prepare_inputs( diff --git a/tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py b/tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py index 389420716e..ac052f00c5 100644 --- a/tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py +++ b/tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py @@ -342,6 +342,8 @@ def _ragged_paged_attention_kernel_loop( m_ref, # [actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, 128], acc_ref, # [actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, head_dim], kv_shuffle_vmem_ref=None, # [bkv_sz // cp_group_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim] + ring_dma_sems=None, # [2] + ring_sync_sem=None, *, # Static kwargs cp_group_size: int | None = None, @@ -367,6 +369,8 @@ def _ragged_paged_attention_kernel_loop( case: RpaCase = RpaCase.MIXED, debug_mode: bool = False, return_lse: bool = False, + pcp_ring_axis_name: str | None = None, + pcp_ring_mesh_axis_names: tuple[str, ...] | None = None, ): assert q_hbm_ref.shape == o_hbm_ref.shape @@ -417,8 +421,11 @@ def _ragged_paged_attention_kernel_loop( q_len = q_end - q_start # Helper functions for context parallelism. - def get_cp_local_size(x): - return (x + cp_group_size - 1 - cp_rank) // cp_group_size + def get_cp_local_size_of_rank(x, rank): + """How much of the `x` cached tokens `rank` owns.""" + return (x + cp_group_size - 1 - rank) // cp_group_size + + ring_enabled = pcp_ring_axis_name is not None def get_kv_new_len(seq_idx): # Under PCP, new KV is all-gathered into token order. @@ -439,9 +446,12 @@ def get_q_pos_offset(seq_idx): return q_pos_offset_ref[seq_idx] return 0 + def get_kv_cache_len_global(seq_idx): + return kv_lens_ref[seq_idx] - get_kv_new_len(seq_idx) + def get_kv_cache_len_local(seq_idx): - global_len = kv_lens_ref[seq_idx] - get_kv_new_len(seq_idx) - return get_cp_local_size(global_len) + return get_cp_local_size_of_rank(get_kv_cache_len_global(seq_idx), + cp_rank) def get_start_bkv_idx(seq_idx): local_cache_len = get_kv_cache_len_local(seq_idx) @@ -453,6 +463,20 @@ def get_start_bkv_idx(seq_idx): start_idx = jnp.maximum(start_idx, local_cache_len // bkv_sz) return start_idx + if ring_enabled: + my_ring_id = lax.axis_index(pcp_ring_axis_name) + + def ring_device_id(rank): + if pcp_ring_mesh_axis_names is None: + return (rank, ) + return tuple( + rank if name == pcp_ring_axis_name else lax.axis_index(name) + for name in pcp_ring_mesh_axis_names) + + ring_next_id = ring_device_id(lax.rem(my_ring_id + 1, cp_group_size)) + ring_prev_id = ring_device_id( + lax.rem(my_ring_id + cp_group_size - 1, cp_group_size)) + if cp_group_size is not None: cp_rank = cp_rank_ref[0] kv_new_len = get_kv_new_len(seq_idx) @@ -590,7 +614,7 @@ def flash_attention_step1_qk_softmax( v = jnp.where(v_span >= kv_cache_len_local_int, v, jnp.array(0.0, dtype=v.dtype)) - if skip_current_attn: + if skip_current_attn and not ring_enabled: kv_cache_len_local_int = kv_cache_len_local.astype(int_ty) mask = mask_and(mask, k_span < kv_cache_len_local_int) v = jnp.where(v_span < kv_cache_len_local_int, v, @@ -674,7 +698,7 @@ def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False): kv_len_start = bkv_idx * bkv_sz kv_p_start = bkv_idx * bkv_p - kv_left = _seq_kv_len_local - kv_len_start + kv_left = jnp.maximum(_seq_kv_len_local - kv_len_start, 0) if update_kv_cache or skip_cache_attn: kv_left_frm_cache = jnp.maximum(kv_left - _seq_kv_new_len, 0) else: @@ -856,8 +880,9 @@ def _update_kv_cache_partial(seq_idx, """ sem = sems.at[3, bkv_sem_idx] - local_offset_start = get_cp_local_size(offset) - local_offset_end = get_cp_local_size(offset + update_sz) + local_offset_start = get_cp_local_size_of_rank(offset, cp_rank) + local_offset_end = get_cp_local_size_of_rank(offset + update_sz, + cp_rank) update_sz = local_offset_end - local_offset_start kv_p_start = local_offset_start // page_size @@ -1167,6 +1192,12 @@ def process(static_q_len=None): actual_bq_csz = min(bq_csz, actual_bq_sz) + if ring_enabled: + global_cache_len = get_kv_cache_len_global(seq_idx) + # Use rank 0 (the longest shard) and mask the short ranks' tail. + max_local_len = get_cp_local_size_of_rank(global_cache_len, 0) + ring_num_bkv = jnp.maximum(cdiv(max_local_len, bkv_sz), 1) + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): next_bq_idx = bq_idx + 1 is_last_bq = next_bq_idx == num_bq @@ -1259,6 +1290,10 @@ def compute_with_bq(bq_idx): end_bkv_idx = jnp.maximum(cdiv(fetch_kv_len, bkv_sz), start_bkv_idx + 1) + if ring_enabled: + # The bkv loop runs cp_group_size rounds per KV block. + end_bkv_idx = ring_num_bkv * cp_group_size + # Prefetch next bq @pl.when(next_seq_idx < end_seq_idx) def prefetch_next_bq(): @@ -1266,30 +1301,77 @@ def prefetch_next_bq(): start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) @pl.loop(start_bkv_idx, end_bkv_idx, unroll=False) - def compute_with_bkv(bkv_idx): + def compute_with_bkv(bkv_idx, effective_kv_len=effective_kv_len): assert bkv_sz % kv_packing == 0 # Get next bkv ids. bkv_sem_idx = sem_ids_ref[1] next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids( seq_idx, bq_idx, bkv_idx, bkv_sem_idx, num_bkv=end_bkv_idx) + if ring_enabled: + round_idx = lax.rem(bkv_idx, cp_group_size) + bkv_idx = bkv_idx // cp_group_size + next_bkv_idx = next_bkv_idx // cp_group_size + bkv_sem_idx = lax.rem(round_idx, 2) + next_bkv_sem_idx = 0 + is_last_round = round_idx == cp_group_size - 1 + else: + round_idx = 0 + is_last_round = jnp.bool_(True) processed_kv_len = bkv_idx * bkv_sz # Prefetch next bkv - @pl.when(next_seq_idx < end_seq_idx) + @pl.when((next_seq_idx < end_seq_idx) & is_last_round) def prefetch_next_bkv(): sem_ids_ref[1] = next_bkv_sem_idx start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx) # Wait for cur bq if not ready yet - @pl.when(bkv_idx == start_bkv_idx) + @pl.when((bkv_idx == start_bkv_idx) & (round_idx == 0)) def wait_cur_bq(): wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) - # Wait for cur bkv - offset, update_sz, src_start_base = wait_fetch_bkv( - seq_idx, bkv_idx, bkv_sem_idx) + if ring_enabled: + + @pl.when(round_idx == 0) + def wait_cur_bkv(): + wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + else: + # Wait for cur bkv + offset, update_sz, src_start_base = wait_fetch_bkv( + seq_idx, bkv_idx, bkv_sem_idx) + + # Send my bkv to next rank in the ring. + if ring_enabled: + src_rank = lax.rem(my_ring_id + cp_group_size - round_idx, + cp_group_size) + effective_kv_len = get_cp_local_size_of_rank( + global_cache_len, src_rank) + next_slot = 1 - bkv_sem_idx + # Chain has no predecessor + is_first_launch = ((seq_idx == start_seq_idx) + & (bq_idx == 0) & (bkv_idx == 0) + & (round_idx == 0)) + + @pl.when(~is_last_round & ~is_first_launch) + def wait_ring_sync(): + pl.semaphore_wait(ring_sync_sem, 1) + + remote_op = pltpu.make_async_remote_copy( + src_ref=bkv_x2_ref.at[bkv_sem_idx, + pl.ds(0, bkv_sz)], + dst_ref=bkv_x2_ref.at[next_slot, + pl.ds(0, bkv_sz)], + send_sem=ring_dma_sems.at[0], + recv_sem=ring_dma_sems.at[1], + device_id=ring_next_id, + device_id_type=pl.DeviceIdType.MESH, + ) + + @pl.when(~is_last_round) + def start_rotate(): + remote_op.start() # Start updating bkv to kv cache if applicable. # Only needed in last bq loop. @@ -1387,6 +1469,39 @@ def attention_loop(idx): acc_ref.at[*prev_lm_slice], ) + if ring_enabled: + + @pl.when(~is_last_round) + def finish_rotate(): + remote_op.wait() + + # Tell the sender its next write target is free. + # Round P-1 doesn't send to the next rank, so it doesn't a signal from round P-2. + @pl.when(round_idx < cp_group_size - 2) + def release_slot_to_sender(): + pl.semaphore_signal( + ring_sync_sem, + 1, + device_id=ring_prev_id, + device_id_type=pl.DeviceIdType.MESH, + ) + + is_last_launch = ((seq_idx == end_seq_idx - 1) + & (bq_idx == num_bq - 1) + & (bkv_idx == ring_num_bkv - 1) + & is_last_round) + + # At round P-1, let the round 0 sender know its next + # write target is free. + @pl.when(is_last_round & ~is_last_launch) + def release_block_to_sender(): + pl.semaphore_signal( + ring_sync_sem, + 1, + device_id=ring_prev_id, + device_id_type=pl.DeviceIdType.MESH, + ) + # Load acc and calculate final output. acc = acc_ref[...] l = broadcast_minor(l_ref[...], acc.shape) # noqa @@ -1692,10 +1807,13 @@ 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 cp_group_size: int | None = None, + cp_rank: jax.Array | int | None = None, + pcp_ring_axis_name: str | None = None, use_causal_mask: bool = True, skip_kv_mask: bool = False, skip_cache_attn: bool = False, skip_current_attn: bool = False, + update_kv_cache: bool = True, sm_scale: float = 1.0, sliding_window: int | None = None, soft_cap: float | None = None, @@ -1861,6 +1979,20 @@ def _validate_block_sizes(block_sizes, prefix): raise NotImplementedError( "PCP does not support sliding_window yet.") + if pcp_ring_axis_name is not None: + if cp_group_size is None or cp_rank is None or cp_group_size % 2 != 0: + raise ValueError( + "pcp_ring_axis_name requires cp_group_size and cp_rank, and " + "cp_group_size must be even.") + if not skip_current_attn or update_kv_cache or use_causal_mask: + raise NotImplementedError( + "pcp_ring_axis_name is a cache-phase path and requires " + "skip_current_attn=True, update_kv_cache=False, and " + "use_causal_mask=False") + if sliding_window is not None: + raise NotImplementedError( + "pcp_ring_axis_name does not support sliding_window") + # No constraints for the following inputs. del sm_scale del mask_value @@ -1882,6 +2014,9 @@ def get_default_block_sizes( pages_per_seq, *, case: RpaCase = RpaCase.MIXED, + pcp_chunk_size: int | None = None, + pcp_ring: bool = False, + vmem_limit_bytes: int | None = None, ): """Get (bq, bkv_sz, bq_csz, bkv_csz) by some heuristic formulas. @@ -1944,13 +2079,54 @@ def get_default_block_sizes( bkv_sz = align_to(bkv_sz, page_size) bkv_sz = max(bkv_csz, (bkv_sz // bkv_csz) * bkv_csz) - return { + bs = { "bq_sz": max(1, bq_sz), "bkv_sz": align_to(bkv_sz, page_size), "bq_csz": max(1, bq_csz), "bkv_csz": align_to(bkv_csz, page_size), } + # PCP current phase (rank-ordered KV remap) needs the prefetch block to + # stay within one head-tail chunk of size C, i.e. bkv_sz <= C. + if pcp_chunk_size is not None and case == RpaCase.MIXED: + bkv_sz = min(bs["bkv_sz"], pcp_chunk_size) + while bkv_sz > page_size and pcp_chunk_size % bkv_sz != 0: + bkv_sz -= page_size + bkv_csz = min(bs["bkv_csz"], bkv_sz) + while bkv_csz > page_size and bkv_sz % bkv_csz != 0: + bkv_csz -= page_size + bs = {**bs, "bkv_sz": bkv_sz, "bkv_csz": bkv_csz} + + if pcp_ring and case == RpaCase.MIXED: + # Ring sizing is the opposite of the default heuristic. The default + # picks small Q tiles because re-streaming KV per tile is nearly free + # from local HBM; the ring re-streams the cache over ICI (~50x slower), + # so its DMA only hides behind compute when the resident tile is as + # large as VMEM allows. + RING_MAX_TILE_ROWS = 8192 + RING_HOP_TARGET_BYTES = 2 * 1024 * 1024 + bytes_per_token = (2 * actual_num_kv_heads * head_dim * + (32 // kv_packing) // 8) + bkv_sz = max( + page_size, + RING_HOP_TARGET_BYTES // bytes_per_token // page_size * page_size) + bq_sz = max( + 1, min(max_q, RING_MAX_TILE_ROWS // max(1, actual_num_q_heads))) + bq_csz = min(bs["bq_csz"], bq_sz) + while bq_csz > 1 and bq_sz % bq_csz != 0: + bq_csz -= 1 + bkv_csz = min(bs["bkv_csz"], bkv_sz) + while bkv_csz > 1 and bkv_sz % bkv_csz != 0: + bkv_csz -= 1 + bs = { + **bs, "bq_sz": bq_sz, + "bq_csz": bq_csz, + "bkv_sz": bkv_sz, + "bkv_csz": bkv_csz + } + + return bs + @jax.jit( static_argnames=( @@ -1979,6 +2155,8 @@ def get_default_block_sizes( "write_last_seq_only", "cp_group_size", "pcp_chunk_size", + "pcp_ring_axis_name", + "pcp_ring_mesh_axis_names", ), donate_argnames="kv_cache", ) @@ -2001,6 +2179,8 @@ def ragged_paged_attention( cp_group_size: int | None = None, q_pos_offsets: 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, @@ -2050,7 +2230,11 @@ def ragged_paged_attention( kv_cache_lens: the number of kv cache tokens that have been computed for each sequence, only needed for PCP. cp_rank: the rank of the current device in the context parallelism group. cp_group_size: the size of the context parallelism group. - q_pos_offsets: the position of the query tokens in the global sequence, only needed for PCP. + q_pos_offsets: the position of the query tokens in the global sequence, only needed for PCP. + pcp_ring_axis_name: PCP only. When set, the cache phase streams the striped + 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 @@ -2112,8 +2296,13 @@ def ragged_paged_attention( kv_cache_lens=kv_cache_lens, q_pos_offsets=q_pos_offsets, cp_group_size=cp_group_size, + cp_rank=cp_rank, + pcp_ring_axis_name=pcp_ring_axis_name, use_causal_mask=use_causal_mask, skip_kv_mask=skip_kv_mask, + skip_cache_attn=skip_cache_attn, + skip_current_attn=skip_current_attn, + update_kv_cache=update_kv_cache, sm_scale=sm_scale, sliding_window=sliding_window, soft_cap=soft_cap, @@ -2245,7 +2434,11 @@ def run_rpa_kernel( l_scratch, m_scratch, acc_scratch, - kv_shuffle_scratch + kv_shuffle_scratch, + pltpu.SemaphoreType.DMA( + (2, )) if pcp_ring_axis_name is not None else None, + pltpu.SemaphoreType.REGULAR + if pcp_ring_axis_name is not None else None, ] scalar_prefetches = ( @@ -2296,6 +2489,8 @@ def run_rpa_kernel( functools.partial( _ragged_paged_attention_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, @@ -2378,6 +2573,9 @@ def _prepare_block_sizes(block_sizes, case): max_num_seqs, pages_per_seq, case=case, + pcp_chunk_size=pcp_chunk_size, + pcp_ring=pcp_ring_axis_name is not None, + vmem_limit_bytes=vmem_limit_bytes, ) else: bs = { @@ -2386,16 +2584,6 @@ def _prepare_block_sizes(block_sizes, case): "bq_csz": block_sizes[2], "bkv_csz": block_sizes[3], } - # PCP current phase (rank-ordered KV remap) needs the prefetch block to - # stay within one head-tail chunk of size C, i.e. bkv_sz <= C. - if pcp_chunk_size is not None and case == RpaCase.MIXED: - bkv_sz = min(bs["bkv_sz"], pcp_chunk_size) - while bkv_sz > page_size and pcp_chunk_size % bkv_sz != 0: - bkv_sz -= page_size - bkv_csz = min(bs["bkv_csz"], bkv_sz) - while bkv_csz > page_size and bkv_sz % bkv_csz != 0: - bkv_csz -= page_size - bs = {**bs, "bkv_sz": bkv_sz, "bkv_csz": bkv_csz} return bs # Decode-only diff --git a/tpu_inference/layers/common/cp_attention.py b/tpu_inference/layers/common/cp_attention.py index 459e0f9feb..08abb4bb01 100644 --- a/tpu_inference/layers/common/cp_attention.py +++ b/tpu_inference/layers/common/cp_attention.py @@ -299,42 +299,6 @@ def _shard_fn(q_local, k_local, v_local, kv_cache_local, kv_lens_local, md.request_distribution, cp_rank_global) -def _pcp_rs_reduce( - o: jax.Array, - lse: jax.Array, - axis: str, - axis_size: int, -) -> tuple[jax.Array, jax.Array]: - """Reduce-scatter across PCP: each rank collects its own token chunk. - - Called inside a PCP shard_map body after the cache phase. - - Input (per rank, all-gathered result from cache phase): - o: [axis_size * chunk, heads, head_dim] - lse: [axis_size * chunk, heads] - Output (per rank, own chunk only): - o: [chunk, heads, head_dim] - lse: [chunk, heads] - """ - chunk = o.shape[0] // axis_size - max_lse = lax.pmax(lse, axis) - max_lse_safe = jnp.where(jnp.isinf(max_lse), 0.0, max_lse) - weights = jnp.exp(lse - max_lse_safe) - o_weighted_sum = lax.psum_scatter(o * weights[..., None].astype(o.dtype), - axis, - scatter_dimension=0, - tiled=True) - denom = lax.psum_scatter(weights, axis, scatter_dimension=0, tiled=True) - max_lse_own = lax.dynamic_slice_in_dim(max_lse_safe, - lax.axis_index(axis) * chunk, chunk, - 0) - denom_safe = jnp.where(denom == 0.0, 1.0, denom)[..., None] - out_merged = o_weighted_sum.astype(denom.dtype) / denom_safe - lse_merged = jnp.where(denom == 0.0, -jnp.inf, - max_lse_own + jnp.log(denom)) - return out_merged, lse_merged - - def pcp_forward( mesh: Mesh, q: jax.Array, @@ -352,11 +316,11 @@ def pcp_forward( """PCP attention forward. Inside the shard_map body: - 1. all_gather Q tokens cache phase needs full sequence view per rank - 2. cache phase attend full Q against this rank's KV cache shard - 3. _pcp_rs_reduce reduce_scatter: each rank collects its token chunk - 4. current phase local Q (head+tail) attends all-gathered current KV - 5. merge_attn_states lse-weighted combine + 1. cache phase in-kernel ring: KV shards rotate around the pcp + axis while each rank attends with its local Q; + one online softmax accumulates all rounds + 2. current phase local Q (head+tail) attends all-gathered current KV + 3. merge_attn_states lse-weighted combine """ if envs.USE_BATCHED_RPA_KERNEL: raise NotImplementedError( @@ -384,18 +348,7 @@ def pcp_forward( k_scale=k_scale, v_scale=v_scale) - # Cache-phase strategy, decided per compile from static comm estimates. - # all_gather and reduce_scatter are duals and move the same (p-1)/p - # fraction, so gather-Q counts both legs; gather-KV moves K and V but has - # no output collective. Volume alone puts the crossover at - # ctx = chunk*NQ/NKV, but gather-Q issues TWO collective rounds (two sync - # points) plus an LSE reweight, so empirically it needs ~2x the raw volume - # advantage before it actually wins. cache_pages = md.pcp.cache_pages - _GATHER_Q_OVERHEAD = 2.0 - comm_q = 2 * padded_q_len * q.shape[1] * q.shape[2] - comm_kv = 2 * (cache_pages * kv_cache.shape[1]) * k.shape[1] * q.shape[2] - use_gather_kv = comm_kv < _GATHER_Q_OVERHEAD * comm_q def _shard_fn(q_local, k_local, v_local, kv_cache_local, kv_lens_local, kv_cache_lens_local, page_indices_local, distribution_local, @@ -412,89 +365,33 @@ def to_token_order(x): # rank-order chunks -> global token order padded_q_len, *x.shape[1:]) # ---- Cache phase -------------------------------------------------- - # Two ways to give every rank what it needs: - # gather-Q : all_gather Q, attend this rank's KV shard, then - # reduce_scatter the partials. Comm ~ 2*chunk*NQ, i.e. - # INDEPENDENT of context length. - # gather-KV: all_gather the striped KV cache into global token order - # so every rank holds the full cache, then attend the - # LOCAL Q against it as ordinary paged attention. No - # output collective at all. Comm ~ ctx*NKV. - # gather-KV wins at short/medium context, gather-Q once ctx is long. - # The cache phase never writes the cache (update_kv_cache=False), so - # the current phase starts from the untouched local shard unless the - # gather-Q path threaded a copy through. - kv_cache_temp = kv_cache_local if cache_pages == 0: # Nothing cached (first chunk of a chunked prefill): the cache # phase would attend an empty cache, be fully masked, and have its # -inf result discarded by merge_attn_states. Skip it outright. context_out = context_lse = None - elif use_gather_kv: - # Compact to this request's live pages, then all_gather with the - # pcp axis INNERMOST: for local page p, offset o, rank j the cache - # holds global token (p*page_l + o)*pcp + j, so flattening - # (p, o, j) is already global token order. Regroup into pages of - # the ORIGINAL width (a pure reshape) and the cache phase becomes - # plain paged attention with cp_group_size=1. - local_q = q_local.shape[0] - cu_kv = jnp.zeros_like(pcp_cu_q_lens_local[0]).at[1:].set(local_q) - max_seqs = kv_lens_local.shape[0] - kv_src = jnp.take(kv_cache_local, - page_indices_local[:cache_pages], - axis=0) - kv_tok = lax.all_gather(kv_src, pcp_axis, axis=2, tiled=False) - n_pages_tok = kv_src.shape[0] * pcp_size - kv_tok = kv_tok.reshape(n_pages_tok, kv_src.shape[1], - *kv_src.shape[2:]) - pi_tok = jnp.tile( - jnp.arange(n_pages_tok, dtype=page_indices_local.dtype), - max_seqs) + else: + cu_ring = jnp.zeros_like(pcp_cu_q_lens_local[0]).at[1:].set( + q_local.shape[0]) context_out, _, context_lse = _rpa_cp_call( q_local, k_local, v_local, - kv_tok, - kv_lens_local, - pi_tok, - cu_kv, - jnp.array([0, 0, 1], jnp.int32), - cp_rank=jnp.zeros((1, ), jnp.int32), - cp_group_size=1, - kv_cache_lens=kv_cache_lens_local, - skip_current_attn=True, - use_causal_mask=False, - update_kv_cache=False, - **common) - else: - # Cache phase: all_gather Q tokens so every rank sees the full - # sequence. PCP local q_local has 2*C tokens (head + tail chunk - # for this rank). After all_gather along the tokens axis: - # pcp_size * 2 * C = padded_q_len. - q_all_tokens = all_gather_tokens(q_local) - cu_cache = jnp.zeros_like( - pcp_cu_q_lens_local[0]).at[1:].set(padded_q_len) - context_out, kv_cache_temp, context_lse = _rpa_cp_call( - q_all_tokens, - k_local, - v_local, kv_cache_local, kv_lens_local, page_indices_local, - cu_cache, + cu_ring, jnp.array([0, 0, 1], jnp.int32), cp_rank=cp_rank, cp_group_size=pcp_size, kv_cache_lens=kv_cache_lens_local, + pcp_ring_axis_name=pcp_axis, + pcp_ring_mesh_axis_names=tuple(mesh.axis_names), skip_current_attn=True, use_causal_mask=False, update_kv_cache=False, **common) - # Rank reduce: reduce_scatter so each rank gets its own 2*C chunk. - context_out, context_lse = _pcp_rs_reduce(context_out, context_lse, - pcp_axis, pcp_size) - # Current phase: local Q (head+tail chunks) attends all-gathered current KV. # pcp_cu_q_lens_local[0] = [0, chunk, chunk+tail_real]; pcp_q_pos_offsets_local[0] = [head_offset, tail_offset]. # remap_kv: if C aligns with page_size, all_gather_tokens() avoids an extra gather-reorder. @@ -508,7 +405,7 @@ def to_token_order(x): # rank-order chunks -> global token order q_local, k_curr, v_curr, - kv_cache_temp, + kv_cache_local, kv_lens_local, page_indices_local, pcp_cu_q_lens_local[0], diff --git a/tpu_inference/layers/common/sharding.py b/tpu_inference/layers/common/sharding.py index bd8197cf3f..245e7e0c48 100644 --- a/tpu_inference/layers/common/sharding.py +++ b/tpu_inference/layers/common/sharding.py @@ -342,6 +342,11 @@ def validate(cls, vllm_config, sharding_strategy): raise ValueError( "Must run Prefill Context Parallelism with NEW_MODEL_DESIGN " "enabled. Please set NEW_MODEL_DESIGN=True") + if sharding_strategy.prefill_context_parallelism % 2 != 0: + raise ValueError( + "Prefill context parallelism size must be even: the ring " + "cache phase rotates KV shards between rank pairs, got " + f"{sharding_strategy.prefill_context_parallelism}.") if sharding_strategy.decode_context_parallelism > 1: raise ValueError( "Only one of prefill/decode context parallelism may be "