PCP + ring attention on top of DCP batched-RPA (#3255) - #3382
Draft
wenxindongwork wants to merge 2 commits into
Draft
PCP + ring attention on top of DCP batched-RPA (#3255)#3382wenxindongwork wants to merge 2 commits into
wenxindongwork wants to merge 2 commits into
Conversation
Apply the in-kernel PCP ring cache phase from #3277 onto the DCP batched-RPA base (#3255). The cache phase streams striped KV cache shards around the pcp axis inside rpa_v3_cp while each rank keeps its local Q, accumulating all rounds in one online softmax; the reduce-scatter epilogue and gather-KV variant are gone. cp_attention.py reconciles the two PRs: the USE_BATCHED_RPA_KERNEL dispatch and AttentionScope mapping stay as in the DCP PR, pcp_forward uses the ring, and PCP under the batched kernel still raises NotImplementedError.
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:
|
Ring rounds are encoded into the schedule's k_idx (block * P + round); rounds > 0 emit zero-size fetches and instead receive the block from the previous rank. Compute reads a private double-buffered ring scratch indexed by round parity: round 0 copies the pipeline's fetched block in locally, and each non-final round sends the resident block to the next rank while computing on it, masked by the source rank's shard length so one online softmax accumulates the full cache across ranks. Synchronization, ported from the rpa_v3_cp in-kernel ring with two hardening changes found on hardware (tpu7x-8): - A startup barrier (get_barrier_semaphore) rendezvouses with both neighbors before any remote traffic, so a first-execution RDMA cannot land on a neighbor still loading the program. - Slot-free credits are granted at the top of the step after the slot's last reader, not at the end of the reading step. The end-of-step release is issued by the scalar core and can outrun the vector core's in-flight reads of the slot, letting the previous rank's next block land mid-read (observed as first-execution corruption entering at the warm/cold device boundary). The ring scratch is zero-initialized like the pipeline KV window, and requires an even ring size (round-parity slot reuse), CACHE_ONLY, update_kv_cache=False, and the HEAD_ALONG_SUBLANE layout. Test: tests/kernels/batched_rpa/pcp_ring_test.py validates the ring against plain CACHE_ONLY over the token-ordered full cache for P=2/4/8, fp32/bf16 on 8 TPU cores.
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
Stacked PR: PCP + ring attention (#3277) on top of the DCP batched-RPA PR (#3255).
The base branch
wxd-dcp-batched-rpa-baseis a snapshot of #3255's head (weiyu0824:feat/dcp-on-batched-rpa@ 9af2d2f), pushed here because GitHub can't target a fork branch as a PR base. So the diff shown is exactly the ring changes on top of the DCP PR. If #3255 gets new commits, the snapshot branch needs a manual update; once #3255 lands, this retargets tomain.Supersedes #3381 (same content, but that one was based on
mainso it showed both PRs' diffs).What this adds on top of #3255
The in-kernel ring cache phase from #3277: for the PCP cache phase, KV cache shards rotate around the pcp axis inside
rpa_v3_cp/kernel.pywhile each rank keeps its local Q, accumulating all rounds in one online softmax (no Q all-gather, no reduce-scatter epilogue, no gather-KV variant).cp_attention.pyreconciles the two PRs:_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 forwards the ring kwargs (pcp_ring_axis_name,pcp_ring_mesh_axis_names) on therpa_v3_cppath.pcp_forwarduses the ring cache phase and keeps feat: Enable DCP for batched RPA kernel #3255's guard: PCP underUSE_BATCHED_RPA_KERNEL=1still raisesNotImplementedError.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 — follow-up work.Tests
On a v7x-8, at this branch's tree:
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 pristine feat: Enable DCP for batched RPA kernel #3255 tree fails identically in the same environment, so this is pre-existing on feat: Enable DCP for batched RPA kernel #3255, not introduced here.Checklist
Before submitting this PR, please make sure: