[PCP] Move the host-side preprocessing into runner/pcp_utils.py - #3514
Draft
bhuvanpkaruturi wants to merge 11 commits into
Draft
[PCP] Move the host-side preprocessing into runner/pcp_utils.py#3514bhuvanpkaruturi wants to merge 11 commits into
bhuvanpkaruturi wants to merge 11 commits into
Conversation
Prepares the metadata contract for fusing several prefill requests into one RPA launch. Three additions, all inert while num_reqs == 1: - kv_new_starts: base offset of each fused seq's current-KV block inside the all-gathered new-KV buffer. With one request every block starts at 0, which is the kernel's implicit base, so None keeps the old path. - kv_token_order: permutation taking the all-gathered current K/V from rank order to request-major token order. The single-request path lets the kernel remap addresses itself via pcp_chunk_size; that remap assumes one request, so several requests need the reorder done in JAX. - num_reqs: static, so the two paths compile separately rather than branching on a traced value. Also adds pcp_token_layout(), the per-request zigzag chunking shared by the runner and the precompile path. Padding slots get a 1-token chunk rather than a zero-length one: the cache phase derives its sequence boundaries from these offsets, and two slots sharing an offset produce a zero-length sequence, which the kernel does not terminate on. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Two scalar-prefetch refs replace an assumption that only one request is in flight: - kv_new_starts_ref: per-seq base offset into the new-KV buffer. Without it every sequence reads its current K/V from offset 0, which is only correct when the launch carries a single request. - kv_write_seq_mask_ref: which sequences perform the cache write. This replaces the write_last_seq_only static flag, which could express "the last of one request's two fused seqs" but not "the tail seq of each of R requests". write_last_seq_only is removed rather than kept alongside: it is a special case of the mask, and leaving both would let callers set contradictory values. Validation rejects kv_new_starts together with pcp_chunk_size — the chunk-size remap rewrites addresses on the assumption of one request, so the two are mutually exclusive by construction. Tests cover per-request KV writes landing on the right pages exactly once, and per-request current-phase output against a reference. One test is skipped by default: it asserts the kernel tolerates a zero-length sequence, which it does not — it hangs the device rather than raising, so running it needs PCP_RUN_HANG_REPRO=1 under an external timeout. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Cache phase (in-kernel ring): one seq per REQUEST spanning its head+tail run. Q stays local under the ring, so each rank's buffer is its request slots back to back and the description is just the current-phase arrays with the head/tail duplication undone: cu_q_lens, kv_lens, kv_cache_lens and the block-table rows at even indices, and half the seq count from request_distribution. No tiling, no output collective. R == 1 keeps the ring PR's whole-buffer seq verbatim, which is valid for any current-phase cu_q_lens including the clipped-tail, per-rank one older callers build. The R > 1 construction relies on the ring's lock-step: remote copies and semaphores are matched by loop position, so every rank must run the identical (seq, tile, block, round) schedule. request_distribution is replicated, cu_q_lens is rank-invariant because the runner gives both halves the full chunk length, and the block count comes from the replicated global cache length. Current phase: the all-gathered K/V arrives in rank order. With one request the kernel remaps addresses itself via pcp_chunk_size; that remap assumes a single request, so several requests reorder into request-major order in JAX via kv_token_order and pass kv_new_starts instead. Each request's tail seq performs the fused KV write via kv_write_seq_mask. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Lifts the one-request-per-step restriction. Each request is zigzag- chunked independently, so rank r holds chunk r and chunk 2P-1-r of every request, and request i's head/tail pair occupies a fixed-width slot in every rank's region. The pieces: - Token permutation from natural order to rank order, and its inverse for the K/V reorder. Slots past a request's end stay -1 and are zeroed. - Attention metadata over 2R fused seqs, request i at seqs 2i and 2i+1. Both halves get the full chunk length, so cu_q_lens is rank-invariant and only q_pos_offsets varies by rank. - logits_indices pointing at each request's last real token, wherever the zigzag put it. - Block-table rows duplicated per fused seq, since the KV write is done by the tail seq and page addressing is per-seq. Buffer sizing changes with it: attention seq slots double under PCP, and the token bucket ladder gains 2*P*max_num_seqs of headroom. Rounding each chunk up independently means P*S can exceed the scheduler's token budget, and the scheduler admits by sum(n_i), so without that headroom a legal batch has no bucket to land in. A step with exactly ONE request keeps the single-request layout: chunk t_pad / 2P, head and tail each filling half of every rank's slice, i.e. exactly the layout upstream's runner produces. pcp_forward's R == 1 current phase (the kernel-side rank-order remap when the chunk is page-aligned, to_token_order otherwise) derives the chunk from the padded buffer width -- pcp_chunk_size is a static kernel argument -- so a lone request packed at ceil(n / 2P) would have its current K/V reordered against the wrong chunk whenever the token bucket is not exact. Multi- request steps carry their own permutation (kv_token_order) and do not depend on the buffer width. The extra tail rows past the request's last token are padding rows like any other: they attend only real keys, nothing reads their output, and the KV write is driven by kv_lens - kv_cache_lens, not by the row count. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Adds num_reqs to the warm-up ladder and builds well-formed dummy PCP
metadata for it.
The dummy metadata has to be a batch the runtime could actually produce.
An all-zero query_start_loc satisfies every shape check but describes
zero-length sequences, which hang the kernel on device during warm-up —
the failure looks like a deadlock at startup with no traceback.
The ladder is capped at {1, max_num_seqs} rather than every power of two,
and combinations the runtime cannot reach are skipped: cache_pages only
matters to the gather-KV path, which multi-request does not take, and a
batch of R requests cannot have fewer than 2*P*R tokens. Without these
the graph count roughly quadrupled and warm-up dominated startup.
Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
…{0, max}
Both attention phases now iterate the live seq count from
request_distribution, so the slots that pad the request count up to the
static num_reqs bucket never reach the kernel: drop their 1-token chunks
from the layout and the max(1, ...) floor in pcp_token_layout.
With the ring, cache_pages only decides whether the cache phase is
elided (== 0); every nonzero rung compiled the same kernel. Two rungs
halve the PCP precompile set.
Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Kernel (tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py): - fused current phase with several requests: per-request kv_new_starts, kv_write_seq_mask, every request's strided KV share written exactly once, ragged lengths, mixed cached/uncached; - ring cache phase over three seqs in one launch with different cached lengths (one zero), row counts and page lists, P = 2 and 4, f32 and bf16, against per-seq full-cache reference (per-seq block counts, one cross-rank chain across seq boundaries, lock-step). Interface (tests/layers/common/test_pcp_attention_interface.py): - pcp_forward with R requests: ragged lengths, mixed cached/uncached, live count below the static num_reqs bucket, per-request KV cache write, a request spanning several ring query tiles, the batch shape from a 32k end-to-end run in the headroom token bucket, and ~1M cached tokens next to a short request; - single-request regression cases unchanged. All need >= P TPU devices; run through cdk. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
Review follow-ups for vllm-project#3425. Branching: `kv_new_starts` and `kv_token_order` are now required fields of PCPMetadata and are built by the runner, the precompile dummies, the interface tests and the pcp-vs-tp benchmark for every request count. That removes the optional-argument plumbing in pcp_forward (`*pcp_extra`, `extra_specs`/`extra_args`, the None guard) and the `if multi else None` in the runner and precompile. The single-request `to_token_order`/`inv_row` reshuffle is gone too: the runner's `kv_token_order` for one request is exactly that permutation, so the non-page-aligned path is one `jnp.take` for any request count. The page-aligned single-request fast path (kernel-side remap via `pcp_chunk_size`) and the single-request cache-phase span are kept as they were; merging them into the general layout is the follow-up PR. `cache_pages` becomes `has_cached_kv: bool`. After the in-kernel ring the value only decided whether the cache phase runs, so the two-rung page ladder, its bucket constant and the round-up helper are deleted; precompile loops over {False, True}. Removed the multi-request batch logger and its counter, which were debugging aids. Comments cut to what the code does not say itself. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
bhuvanpkaruturi
force-pushed
the
pcp-preprocessing-module
branch
from
September 4, 2026 22:58
3beabfc to
230a62f
Compare
Pure code motion, no behaviour change. The PCP block in `_prepare_inputs` (prefill-only check, zigzag chunking, token permutation into rank order, per-seq attention metadata, logits indices, `PCPMetadata` assembly) moves to `PCPPreprocessor.prepare_inputs` in the new `runner/pcp_utils.py`, along with the pure layout helpers from `layers/common/attention_metadata.py`, which keeps only the `PCPMetadata` dataclass. Why: `tpu_runner.py` is the most contended file in the tree and the PCP block was one of the largest additions to `_prepare_inputs`; the runner folder already keeps speculative decoding, structured decoding and multimodal preprocessing in their own modules. It also makes the layout logic testable without a runner. `PCPPreprocessor.metadata_to_device` is the one place that places host arrays as a `PCPMetadata`; the compilation manager builds its precompile dummy through it instead of carrying a second copy of the sharding recipe. The token-bucket headroom uses `pcp_max_buffer_tokens`, the closed-form bound of `pcp_buffer_tokens`, so the bound lives next to the function it bounds. The logits indices are read off `kv_token_order` rather than re-deriving the zigzag formula. tests/runner/test_pcp_utils.py covers the layout, the permutation and its inverse, the per-seq arrays, the headroom bound, the logits slots, and `prepare_inputs` end to end on a mesh (skipped with fewer devices than pcp_size). The mock-runner DP tests set `pcp_preprocessor = None`: the runner now gates PCP on that attribute, and a MagicMock auto-attribute is never None. Signed-off-by: Bhuvan Karuturi <bkaruturi@google.com>
bhuvanpkaruturi
force-pushed
the
pcp-preprocessing-module
branch
from
September 4, 2026 23:03
230a62f to
e6cc94b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Pure code motion, no behaviour change. Follows a review suggestion on #3425 to move the CP preprocessing logic into its own file in the runner folder.
Today the per-step PCP work that
_prepare_inputsdoes inline — validating the batch as prefill-only, zigzag-chunking each request, permuting the token buffers into rank order, building the per-seq attention metadata and logits indices, and assemblingPCPMetadata— is one of the largest single blocks intpu_runner.py, and the pure layout helpers it uses sit inlayers/common/attention_metadata.pynext to the dataclass they build.After this PR:
tpu_inference/runner/pcp_utils.pyholds all of it.PCPPreprocessor.prepare_inputsis the old inline block;_prepare_inputsreduces to a single call. The layout helpers (pcp_token_layout,pcp_seq_arrays,pcp_cache_page_buckets,round_up_pcp_cache_pages) move here fromattention_metadata.py, which keeps only thePCPMetadatadataclass — the layer-side contract stays underlayers/, everything that builds it is runner-side. This mirrors how speculative decoding, structured decoding and multimodal preprocessing already live in their own runner modules.pcp_token_permutationandpcp_last_token_slots, andpcp_batch_layoutfolds the buffer-fit assert together with the single-request chunk override, so each step of the layout can be checked on its own without a runner.PCPPreprocessoris constructed in_init_inputsafter the request-count ladder and owns the multi-request log counter. The token-bucket sizing usespcp_buffer_tokens. The compilation manager and the interface test only change their imports.git diff --color-movedshows the move; the runner-side call site is the only logic that is new.Tests
tests/runner/test_pcp_utils.py(39 checks, numpy plus one mesh test): every real token lands in exactly one slot,kv_token_orderundoes the permutation on the live rows, the zigzag rank assignment,cu_q_lens/q_pos_offsets/kv_new_startsper seq, the logits slots, andprepare_inputsend to end on a mesh (skipped with fewer devices thanpcp_size).tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py(38 passed, 1 skipped) andtests/layers/common/test_pcp_attention_interface.py(25 passed) — identical counts to [PCP] Multi-request prefill context parallelism on the ring cache phase #3425 without this change — plus the new runner suite (39 passed).Checklist
Before submitting this PR, please make sure: