PCP + ring attention on top of DCP batched-RPA (#3255) - #3381
Closed
wenxindongwork wants to merge 12 commits into
Closed
PCP + ring attention on top of DCP batched-RPA (#3255)#3381wenxindongwork wants to merge 12 commits into
wenxindongwork wants to merge 12 commits into
Conversation
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
…uarded cond in non-cp pass Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
…on schedule and rely on existing bref_override infra Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Adds a third cache-phase strategy for prefill context parallelism:
instead of all-gathering Q (and reduce-scattering the output) or
all-gathering the striped KV cache, the kernel rotates the cache
around the PCP ring one block at a time with in-kernel remote DMAs,
attending each block as it arrives. (m, l, acc) accumulate across all
rounds, so the whole cache folds into one online softmax with no
output collective and only two KV blocks resident.
Implementation: process() drives the ring by scaling its KV-block loop
by cp_group_size -- each iteration is one round of one block, falling
through to the shared flash-attention code with the slot and the
originating rank's stripe length swapped in. Prefetch, waits, and
sizing are the plain path's own, conditioned internally. Rotation uses
one DMA semaphore pair plus a single credit semaphore ("the slot you
are about to overwrite is free"), with the block barrier expressed as
the last round's credit; launch-first/launch-last predicates keep the
ledger balanced at zero. Ring block sizing lives in
get_default_block_sizes: one Q block per launch (each extra block
re-streams the whole cache), ~2MB hops to amortize the per-hop fixed
cost, bq halved until the VMEM estimate fits.
The ring wins where TP shrinks the per-device KV: at pcp4 x tp2 it is
the best strategy from 64k context up (1M cumulative TTFT 4425 vs 4861
auto), while at pcp8 x tp1 it loses to the gathers (9616 vs 4571) --
selection is left to the caller via cache_phase until an automatic
dispatch rule lands.
Verified: kernel suite (ring vs full-cache reference at P in {2,4,8},
both dtypes, plus DCP/decode/prefill/mixed regressions), interface
suite on the production 7-axis mesh, and cumulative-TTFT sweeps to 1M
on both shardings.
Signed-off-by: wenxindongwork <wenxindong@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Delete the gather_kv and gather_q cache-phase paths and the comm-volume auto-dispatch: the ring wins end-to-end at every measured input length on Qwen3-235B-A22B-FP8 at pcp2 x tp4 (256k prefill: 4552 tok/s vs 3749 gather-auto, mean TTFT 317s vs 396s; 4k: -16.5% TTFT; 1k: -2%), needs no gathered-KV HBM materialization, and has no output collective. Also clamp the ring Q tile to 8192 token-head rows: the VMEM estimator undercounts the rows x bkv QK^T scratch and the maximizer otherwise picks tiles that OOM compilation at head_dim=128 (estimator says fit, Mosaic OOMs). The clamp reproduces every benchmark-validated tile choice (512 @ 16 q-heads, 256 @ 32, 1024 @ 8). Signed-off-by: wenxindongwork <wenxindong@google.com>
Signed-off-by: wenxindongwork <wenxindong@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
- get_default_block_sizes: apply the ring tile/hop sizing only for RpaCase.MIXED (the only case the ring runs) so DECODE kernels keep their own sizing instead of an unnecessarily large bq_sz. - sharding config validation: reject odd prefill_context_parallelism at server startup instead of failing later inside the kernel. Signed-off-by: wenxindongwork <wenxindong@google.com>
Reconcile cp_attention.py: keep the USE_BATCHED_RPA_KERNEL dispatch and AttentionScope mapping from the DCP PR, and the in-kernel ring cache phase in pcp_forward from the ring PR. PCP under the batched RPA kernel still raises NotImplementedError; the ring runs on the rpa_v3_cp path.
DescriptionStart with a short description of what the PR does and how this is a change from The rest of the description includes relevant details and context, examples:
If the change fixes a Github issue, please include a link, e.g.,: TestsPlease describe how you tested this change, and include any instructions and/or ChecklistBefore submitting this PR, please make sure:
|
Collaborator
Author
|
Superseded by #3382, which is based on the DCP PR's branch so the diff shows only the ring changes. |
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
2 participants
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
Stacked PR: PCP + ring attention (#3277) rebased on top of the DCP batched-RPA PR (#3255).
This branch is
feat/dcp-on-batched-rpa(#3255, @weiyu0824) with the in-kernel PCP ring attention from #3277 merged on top, so the two features can land together without conflicts. Please review only the merge commit; the underlying changes are reviewed in their own PRs.What the merge reconciles
Both PRs rewrite parts of
tpu_inference/layers/common/cp_attention.py:_rpa_cp_callkeeps feat: Enable DCP for batched RPA kernel #3255'sUSE_BATCHED_RPA_KERNELdispatch (mappingskip_cache_attn/skip_current_attnonto the batched kernel'sAttentionScope), and still forwards the ring kwargs (pcp_ring_axis_name,pcp_ring_mesh_axis_names) on therpa_v3_cppath.pcp_forwardtakes PCP + ring attention #3277's ring cache phase (in-kernel KV ring around the PCP axis, one online softmax, no reduce-scatter) and keeps feat: Enable DCP for batched RPA kernel #3255's guard: PCP underUSE_BATCHED_RPA_KERNEL=1still raisesNotImplementedError.dcp_forwardis feat: Enable DCP for batched RPA kernel #3255's version unchanged (SEQ_ON_LANE-aware cache sharding).Everything else merged cleanly: #3277's changes live in
rpa_v3_cp/kernel.py+sharding.py, #3255's inbatched_rpa/.Not in scope
Porting the ring into the batched RPA kernel itself. The batched kernel fetches KV through
emit_pipelinebuffered refs driven by precomputed DMA descriptors (schedule.py/bref_override.py), so an in-kernel ring there needs a ring-aware KV buffered ref and schedule rounds — a follow-up, not a merge.Tests
On a v7x-8 (worktree code first on
PYTHONPATH):python tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py RaggedPagedAttentionPcpTest— 30/30 pass, including the 6test_pcp_ring_cache_phase_matches_full_cachevariants (P=2/4/8, fp32/bf16).python tests/layers/common/test_pcp_attention_interface.py— 16/16 pass.USE_BATCHED_RPA_KERNEL=1 python tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_dcp_test.py— both tests error withRuntimeError: Array has been deletedwhen the test re-reads the donatedkv_cacheinput. The pristinefeat/dcp-on-batched-rpatree fails identically in the same environment, so this is pre-existing on feat: Enable DCP for batched RPA kernel #3255 (donation vs. this jax version), not introduced by the merge.Checklist
Before submitting this PR, please make sure: