PCP + ring attention by wenxindongwork · Pull Request #3277 · vllm-project/tpu-inference

wenxindongwork · 2026-07-27T23:38:29Z

Summary

This PR adds a third PCP cache-phase strategy — cache_phase="ring" in addition to the existing gather_kv/gather_q strategy.

phase	how every rank sees the cache	cost
`gather_q`	all-gather Q, reduce-scatter the output	two collective rounds, comm independent of context
`gather_kv`	all-gather the striped cache into one buffer	comm and peak memory grow with context
ring	rotate KV around the PCP ring inside the kernel	no exposed collective outside of the kernel, ici traffic grows with context

Ring is now the default behavior for PCP since it outcompetes gather_q and gather_kv in almost every setting.

Implementation details

Each rank keeps its local Q and cache stripe; the kernel rotates the cache one block at a time with make_async_remote_copy.

The conceptual change vs the non-ring path is how bkv is prefetched. In the non-ring path, we keep a bkv double buffer and prefetch the next bkv block from HBM. In the ring path, we still keep a bkv double buffer, but prefetch the next block via ICI from the direct neighbor (except for the first ring step which prefetches from HBM).

We also added a custom block size calculation logic for the ring path, which maximizes bq_sz.

PCP strategy performance comparison.

CH=4k, NQ=32 NKV=2 HD=256, bf16:

pcp4 × tp2 (baseline = tp8)

┌─────────┬─────────────┬────────────────────────┬───────────────────────┬────────────────────────┬────────────────────────┬───────────────────────┐
│ Context │ tp8 (8 dev) │      tp4 (4 dev)       │     pcp_ag (auto)     │        pcp_ring        │       gather_kv        │       gather_q        │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│      1k │        0.75 │     0.79 (5.3% slower) │   0.86 (14.7% slower) │    0.85 (13.3% slower) │    0.85 (13.3% slower) │   0.86 (14.7% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│      2k │        0.77 │     0.81 (5.2% slower) │   0.90 (16.9% slower) │    0.89 (15.6% slower) │    0.89 (15.6% slower) │   0.94 (22.1% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│      4k │        0.83 │    0.93 (12.0% slower) │   0.93 (12.0% slower) │    0.92 (10.8% slower) │    0.93 (12.0% slower) │   0.93 (12.0% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│      8k │        1.80 │    2.15 (19.4% slower) │   2.37 (31.7% slower) │    2.69 (49.4% slower) │    2.36 (31.1% slower) │   2.54 (41.1% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│     16k │        4.18 │    5.47 (30.9% slower) │   5.70 (36.4% slower) │    6.55 (1.57x slower) │    5.74 (37.3% slower) │   6.18 (47.8% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│     32k │       10.70 │   15.63 (46.1% slower) │  14.82 (38.5% slower) │   15.75 (47.2% slower) │   14.84 (38.7% slower) │  15.08 (40.9% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│     64k │       31.42 │   49.98 (1.59x slower) │  42.09 (34.0% slower) │   40.05 (27.5% slower) │   42.05 (33.8% slower) │  39.43 (25.5% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│    128k │       99.56 │  174.83 (1.76x slower) │ 117.59 (18.1% slower) │  112.56 (13.1% slower) │  132.02 (32.6% slower) │ 114.41 (14.9% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│    256k │      348.31 │  650.45 (1.87x slower) │  373.72 (7.3% slower) │   353.03 (1.4% slower) │  455.05 (30.6% slower) │  369.94 (6.2% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│    512k │     1299.17 │ 2497.89 (1.92x slower) │ 1309.17 (0.8% slower) │  1220.51 (6.1% faster) │ 1665.95 (28.2% slower) │ 1307.22 (0.6% slower) │
├─────────┼─────────────┼────────────────────────┼───────────────────────┼────────────────────────┼────────────────────────┼───────────────────────┤
│      1M │     4989.68 │ 9766.88 (1.96x slower) │ 4861.91 (2.6% faster) │ 4480.00 (10.2% faster) │ 6348.29 (27.2% slower) │ 4860.66 (2.6% faster) │
└─────────┴─────────────┴────────────────────────┴───────────────────────┴────────────────────────┴────────────────────────┴───────────────────────┘

pcp8 × tp1 (fp8)

┌─────────┬─────────┬───────────┬──────────┬─────────┬─────────┬─────────┐
│ Context │  auto   │ gather_kv │ gather_q │  ring   │   tp8   │   tp4   │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 1k      │    0.74 │      0.75 │     0.69 │    0.75 │    0.76 │    0.78 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 4k      │    0.71 │      0.87 │     0.68 │    0.77 │    0.83 │    0.93 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 8k      │    1.81 │      1.97 │     2.34 │    2.20 │    1.80 │    2.14 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 16k     │    4.46 │      4.62 │     6.05 │    5.11 │    4.18 │    5.45 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 32k     │   11.62 │     11.79 │    15.07 │   12.71 │   11.25 │   15.53 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 64k     │   33.39 │     33.53 │    39.41 │   33.70 │   31.28 │   49.92 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 128k    │  106.89 │    114.08 │   113.34 │  100.03 │   98.97 │  173.62 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 256k    │  355.98 │    411.57 │   362.61 │  330.45 │  345.73 │  644.51 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 512k    │ 1264.58 │   1545.32 │  1272.37 │ 1188.47 │ 1288.43 │ 2473.09 │
├─────────┼─────────┼───────────┼──────────┼─────────┼─────────┼─────────┤
│ 1M      │ 4695.17 │   6015.00 │  4706.93 │ 4471.59 │ 4945.97 │ 9676.35 │
└─────────┴─────────┴───────────┴──────────┴─────────┴─────────┴─────────┘

E2E run (Qwen3 235b fp8)

┌───────────────────────────────┬──────────────┬──────────────────────────────┬──────────────────────────┐
│                               │ TP=4  │ PCP2×TP4 gather-auto (prev PR) │ PCP2×TP4 ring (this PR) │
├───────────────────────────────┼──────────────┼──────────────────────────────┼──────────────────────────┤
│ 256k benchmark duration (s)   │       1411.9 │                        699.2 │                    575.9 │
├───────────────────────────────┼──────────────┼──────────────────────────────┼──────────────────────────┤
│ 256k total throughput (tok/s) │       1856.7 │               3749.2 (2.02×) │           4551.6 (2.45×) │
├───────────────────────────────┼──────────────┼──────────────────────────────┼──────────────────────────┤
│ 256k mean TTFT (s)            │        793.0 │                        396.3 │     317.2 (−20% vs auto) │
├───────────────────────────────┼──────────────┼──────────────────────────────┼──────────────────────────┤
│ 4k mean TTFT (ms)             │       1989.6 │                       1753.5 │          1463.5 (−16.5%) │
├───────────────────────────────┼──────────────┼──────────────────────────────┼──────────────────────────┤
│ 1k mean TTFT (ms)             │        581.9 │                        473.7 │                    466.4 │
└───────────────────────────────┴──────────────┴──────────────────────────────┴──────────────────────────┘

Commands:

NUM_PRECOMPILE_WORKERS=8 VLLM_ENGINE_READY_TIMEOUT_S=7200 \
VLLM_TPU_BUCKET_PADDING_GAP=2048 NEW_MODEL_DESIGN=1 \
/mnt/disks/persist/vllm_conda/bin/vllm serve Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \
  --max-model-len=262144 --max-num-seqs=1 --no-enable-prefix-caching \
  --gpu-memory-utilization=0.95 --tensor-parallel-size=4 \
  --download_dir=/mnt/nvme2 --max-num-batched-tokens=16384 \
  --prefill-context-parallel-size 2

  scripts/vllm/benchmarking/benchmark_serving.py --backend=vllm \
  --model=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 --dataset-name=random \
  --num-prompts=10 --random-input-len=<1024|4096|262134> --random-output-len=1

Validated model quality with gsm 8k

github-actions · 2026-07-27T23:38:44Z

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

why is this change being made,
the problem being solved and any relevant context,
why this is a good solution,
some information about the specific implementation,
shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have made or will make corresponding changes to any relevant documentation.

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>

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>

lk-chen · 2026-08-12T18:41:39Z

@copilot show per-file owner, don't at them, just username

Copilot · 2026-08-12T18:43:07Z

@copilot show per-file owner, don't at them, just username

Per CODEOWNERS for the files touched here:

tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py: kyuyeunk, wenxindongwork, weiyu0824
tpu_inference/layers/common/cp_attention.py: kyuyeunk, lk-chen, jrplatin, gxd3
all touched .buildkite/* files: QiliangCui, yiw-wang, CienetStingLin, yunyao-gg, theminghuang

bhuvanpkaruturi

Left comments on couple of issues (better assertions and reducing VMEM wastage). Other LGTM.

- 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>

wenxindongwork changed the title ~~PCP: in-kernel KV ring for the cache phase~~ PCP: in-kernel KV ring cache phase (kernel + attention_interface) Jul 28, 2026

wenxindongwork force-pushed the wxd-pcp-ring-attn branch from 9631f6e to 327974f Compare July 29, 2026 01:01

wenxindongwork force-pushed the wxd-pcp-gatherkv branch from 9830975 to 833b596 Compare July 31, 2026 05:45

wenxindongwork force-pushed the wxd-pcp-ring-attn branch from 327974f to f9a7a91 Compare July 31, 2026 18:56

wenxindongwork changed the base branch from wxd-pcp-gatherkv to main July 31, 2026 20:47

wenxindongwork force-pushed the wxd-pcp-ring-attn branch from f9a7a91 to a98bb16 Compare July 31, 2026 20:47

wenxindongwork changed the title ~~PCP: in-kernel KV ring cache phase (kernel + attention_interface)~~ PCP: KV ring cache phase Jul 31, 2026

wenxindongwork changed the title ~~PCP: KV ring cache phase~~ PCP + ring attention Aug 7, 2026

wenxindongwork added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 7, 2026

wenxindongwork force-pushed the wxd-pcp-ring-attn branch from 7af52ac to 2a70feb Compare August 7, 2026 23:44

wenxindongwork force-pushed the wxd-pcp-ring-attn branch from 2c03f64 to 351d99f Compare August 7, 2026 23:49

wenxindongwork marked this pull request as ready for review August 8, 2026 00:35

wenxindongwork requested review from gxd3, jrplatin, kyuyeunk, lk-chen and weiyu0824 as code owners August 8, 2026 00:35

wenxindongwork requested a review from bhuvanpkaruturi August 10, 2026 17:25

weiyu0824 reviewed Aug 10, 2026

View reviewed changes

Comment thread tpu_inference/layers/common/cp_attention.py

wenxindongwork mentioned this pull request Aug 10, 2026

PCP: in-kernel ring cache phase, fused with the current phase #3367

Draft

weiyu0824 reviewed Aug 11, 2026

View reviewed changes

Comment thread tpu_inference/layers/common/cp_attention.py

Comment thread tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py

Merge remote-tracking branch 'origin/main' into wxd-pcp-ring-attn

d0a3985

wenxindongwork marked this pull request as draft August 12, 2026 17:17

wenxindongwork added 2 commits August 12, 2026 17:43

Trim ring sizing and cache-phase comments

3029b01

Signed-off-by: wenxindongwork <wenxindong@google.com>

wenxindongwork marked this pull request as ready for review August 12, 2026 18:34

weiyu0824 approved these changes Aug 12, 2026

View reviewed changes

lk-chen approved these changes Aug 12, 2026

View reviewed changes

Copilot started work on behalf of lk-chen August 12, 2026 18:42 View session

Copilot AI requested a review from lk-chen August 12, 2026 18:43

Copilot finished work on behalf of lk-chen August 12, 2026 18:43

wenxindongwork enabled auto-merge (squash) August 12, 2026 18:46

bhuvanpkaruturi reviewed Aug 12, 2026

View reviewed changes

Comment thread tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py Outdated

bhuvanpkaruturi reviewed Aug 12, 2026

View reviewed changes

Comment thread tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py

bhuvanpkaruturi approved these changes Aug 12, 2026

View reviewed changes

wenxindongwork added ready ONLY add when PR is ready to merge/full CI is needed and removed ready ONLY add when PR is ready to merge/full CI is needed labels Aug 12, 2026

wenxindongwork requested review from mrjunwan-lang and sixiang-google as code owners August 12, 2026 22:09

wenxindongwork disabled auto-merge August 12, 2026 22:55

This was referenced Aug 12, 2026

PCP + ring attention on top of DCP batched-RPA (#3255) #3381

Closed

PCP + ring attention on top of DCP batched-RPA (#3255) #3382

Draft

Navigation Menu

Uh oh!

PCP + ring attention - #3277

PCP + ring attention#3277
wenxindongwork wants to merge 5 commits into
mainfrom
wxd-pcp-ring-attn

wenxindongwork commented Jul 27, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 27, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lk-chen commented Aug 12, 2026

Uh oh!

Copilot AI commented Aug 12, 2026

Uh oh!

Uh oh!

Uh oh!

bhuvanpkaruturi left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

wenxindongwork commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation details

PCP strategy performance comparison.

Uh oh!

github-actions Bot commented Jul 27, 2026

Description

Tests

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lk-chen commented Aug 12, 2026

Uh oh!

Copilot AI commented Aug 12, 2026

Uh oh!

Uh oh!

Uh oh!

bhuvanpkaruturi left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

wenxindongwork commented Jul 27, 2026 •

edited

Loading