[PCP] Fewer V4 prefill gathers and fix PCP+TBO accuracy by yitingw1 · Pull Request #2789 · ROCm/ATOM

yitingw1 · 2026-10-28T12:22:18Z

Requires ROCm/aiter#7149 (6-rank groups).

Motivation

At small prefill batches PCP on DeepSeek-V4-Pro is host-bound: a step at -tp 1 -pcp 8 launches ~3,200 GPU ops, and the CPU, not the GPU, sets its length (~190 ms of host time per step, with the GPU idle for up to 126 ms of it). Every attention layer issued 3–5 separate torch.distributed all-gathers, and the PCP reindex blocked on a .item() per ragged field. DSpark gathered three full aux hidden states per token and then projected every row on every rank. Validating this exposed a PCP+TBO accuracy bug on main and a compile-cache gap.

Change	Kind	Effect
Group the attention-side PCP gathers	Perf	397 → 245 collectives per prefill step; TTFT −9.2 % (ISL 2K) / −10.3 % (ISL 8K) at `-pcp 8`, concurrency 1
Reindex from host indptrs	Perf	blocking `.item()` per prefill step: 4 → 0 (16 → 0 under TBO)
DSpark context KV from each rank's own rows	Perf	gathers ~1 KiB/token per stage instead of ~42 KiB/token; projection 1/pcp per rank
Hold TBO comm-stream outputs for the compute stream	Correctness	PCP+TBO gsm8k 0.68 → 0.965
Key the compile cache on the KV cache dtype	Crash fix	bf16 after fp8 DSpark on one machine no longer loads the fp8 drafter graph

Technical Details

1. Grouped gathers

pcp_allgather_rerange_many gathers several tensors in one all_gatherv (aiter's custom all-gather: one launch, one barrier; PyNCCL group past its size limit) and restores token order per tensor. Each layer now issues one gather for the two compressor inputs (projected in _launch_compressors, passed as Compressor.forward(combined=)) and one for the extend K with its positions:

Attention-side collectives per layer (fp8 KV)	`main`	this PR
CSA (30 layers): main + indexer compressor, `k_packed`, `k_rope`, positions	5	2
HCA (31 layers): compressor, `k_packed`, `k_rope`, positions	4	2

MoE's gather / reduce-scatter (2 per layer) are unchanged. Positions are gathered rather than read off the forward context because an MTP draft pass runs at the target's positions + 1. gather_rows (now in pcp_utils, shared with MoE's DP / PCP / comm-fused gathers) carries one-byte dtypes as uint8 and everything else as is: aiter's grouped all_gatherv copies bytes, and only its PyNCCL fallback lacks fp8 / e8m0 (MoE used to dress int32 ids as fp32, a constraint of the per-tensor custom all-gather it no longer calls). At group sizes aiter's custom all-gatherv has no kernel for (6), this PR requires ROCm/aiter#7149, which makes aiter decline them and fall back to NCCL.

2. Reindex without readback

pcp_reindex_ragged sized its output with int(indptr_local[-1].item()). The builder already computes the four prefill indptrs in numpy; it keeps them on the metadata (paged_prefill_indptr_cpu), and pcp_local_indptr derives each rank's indptr on the host, for both the full batch and the PCP+TBO request groups.

3. DSpark context KV

The projection from aux hidden states to window KV (main_proj + norm, then per stage wqkv_a + norm + RoPE + fake-quant) is row by row, so each rank now projects only its own rows at their own positions, gathers the projected KV, and writes as before. DSparkLayer.write_context_kv is split into context_kv_rows + write_context_kv_rows; drafts opt in with context_kv_by_rows (V4 DSpark only; other flavors keep gathering the aux states). pcp_round_robin_split now returns unit strides: .contiguous() left a one-row shard (a 2-token chunk at pcp 8) at row stride 8, which aiter's cached-positions RoPE asserts on.

4. PCP+TBO: comm-stream outputs read on the compute stream

Under TBO a PCP gather runs on the comm stream, so its output is allocated there, while the compute stream reads it. Once the forward frees it, the caching allocator hands that memory straight back to the comm stream, and the comm stream waits only on its own ubatch's compute (ubatching.py, one shared comm and compute stream): the other ubatch's next gather of the same layer can overwrite the extend K before this ubatch's attention has read it.

`-pcp 8 --enable-tbo`, gsm8k 200 q	gsm8k
`main` (two nodes)	0.69 / 0.675
this PR, extend-K gathers on `main`'s path / K only / positions only	0.71 / 0.70 / 0.975
… K + positions on `main`'s path + `record_stream`	0.965
this PR (grouped gathers alone: 0.96)	0.965; final commit 0.97

tbo_on_comm(fn, *args) now runs every TBO comm phase whose output the caching allocator owns: yield, fn on the comm stream, switch back, and record_stream the TBO context's compute stream on every tensor fn returns (nested lists / tuples included). That is the extend K and positions, the compressor inputs, FusedMoE's PCP and DP routed-rows gathers and reduce-scatters, and pure TP's overlapped tbo_all_reduce (-tp 8 --enable-tbo prefill: gsm8k 0.96). Mori keeps its hand-written phases: dispatch / combine return views of each ubatch's own mori buffers, which the allocator never hands out. The grouped gathers passed main's accuracy only by timing (0.96 without the hold), so gsm8k cannot pin the call sites; a test fails instead if any of these writes a comm phase out by hand. Under a graph the V4 compressors read their PCP-gathered inputs on side streams, so those are held for the side streams the same way.

5. Compile cache key

Config.compute_hash gains kv_cache_dtype, the factor #2566 deferred: the DSpark drafter branches on it while tracing (dspark_fp8_planned), and a bf16 run after an fp8 one hashed to the same directory and died in warmup.

Test Plan

pytest tests/test_pcp_ops.py tests/test_tbo_on_comm.py tests/test_tbo_mark_used_on_compute.py \
  tests/model_ops/test_dp_gather_routed_rows.py tests/test_compile_hash_spec.py \
  tests/test_dspark.py tests/attentions/deepseek_v41/test_tbo.py tests/model_ops/test_indexer_rows.py ...

# DeepSeek-V4-Pro(-DSpark), MI355X: -tp 1 -pcp 8, kv fp8, index cache fp4, prefix caching, mbt 512,
# --level 3 --cudagraph-mode FULL; gsm8k 5-shot. Perf: no spec, mbt 16384, random prompts, OSL 1,
# concurrency 1, two main/PR pairs (second in reverse order), warmed up 5 min.
python -m atom.benchmarks.benchmark_serving --dataset-name random --random-input-len 2048 \
  --random-output-len 1 --max-concurrency 1 --num-prompts 1500 --ignore-eos

Test Result

Unit tests: 700 passed, 12 skipped; ruff and black clean. Each new test fails with its change reverted (22 mutations: 17 on the trees before this review round, 5 on the final one). aiter's test_custom_allgatherv.py -t 6 passes with its fix and fails with the world-size gates forced open; -t 8 unchanged.

`-tp 1 -pcp 8` unless noted	Result
DSpark k=5, fp8 KV, 200 q	gsm8k 0.97, acceptance 63.35 % (#2125's smoke: 63.35 %)
DSpark k=5, bf16 KV, 100 q	0.98, acceptance 61.66 %
MTP k=3, 200 q	0.975, acceptance 65.96 %
TBO, 200 q	0.965; final commit 0.97 (`main` 0.69 / 0.675)
`-tp 4 -pcp 2`, 20K-token prompts (PyNCCL fallback for fp8) + 200 q, final commit	4/4 served, 0.97
`-tp 1 -pcp 6` with ROCm/aiter#7149 (no ATOM 6-rank special case), 100 q	0.97
`-tp 8 --enable-tbo prefill`, mbt 16384, TBO from 1024 tokens (64 of 80 prefill steps in a profiled probe ran TBO), 200 q	0.96

Prefill, concurrency 1 (two pairs)	`main`	this PR
Mean TTFT, ISL 2048	367.2 / 366.2 ms	335.0 / 331.3 ms	−9.2 %
Mean TTFT, ISL 8192	389.8 / 388.6 ms	349.8 / 348.6 ms	−10.3 %
Collectives per step (trace)	397	245
Host time per step (trace)	189–193 ms	163–168 ms

Known Limitations

-pcp with speculative decoding and --enable-tbo stays refused at startup ([PCP] Fix sampled-id sync, DSpark aux gather and FP4 indexer padding #2125): PCP+TBO stripes each request group on its own, while DSpark's aux gather and context-KV projection, like the draft prefill split, read the batch as one stripe. Lifting that needs all three to follow the groups.
Under --enable-tbo all the TP all-reduce and DP collectives now record_stream while decode graphs capture. Measured at -tp 8 --enable-tbo all (DSV4-Pro, FULL graphs, 8 of 10 buckets recorded split): 1952 holds per rank during capture, 135 MiB held in total, and the graph pool grows by 0.00 GB with or without the hold; KV blocks unchanged.
The custom all-gather copies each input into its registered buffer (D2D memcpys per step 244 → 518); still a net win at the sizes measured.
Performance was measured only at -tp 1 -pcp 8, concurrency 1, where every attention gather takes the custom all-gather; -tp 4 -pcp 2 was checked for correctness only.
EAGLE3 under PCP on a DSA target gathers the aux states twice (EAGLE3 strips them off the target's output, which DeepseekV2ForCausalLM has already gathered); pre-existing, for a separate fix.
tests/test_fused_compress_ragged.py fails 4 cases on this image on main too; the PCP unit tests import aiter, so a CI runner without it skips them.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

🤖 Generated with Claude Code

Each V4 attention layer under PCP prefill issued 3-5 separate all-gathers (main compressor input, indexer compressor input, extend K -- two under fp8 KV -- and positions), each through torch.distributed. Now it issues two: - `pcp_allgather_rerange_many` gathers several tensors in one grouped `all_gatherv` and restores token order per tensor. One-byte dtypes travel as uint8, since past aiter's custom all-gather limit the call falls back to PyNCCL, which has no fp8 dtype. - The extend K is gathered together with the positions its SWA write addresses the ring by. Positions are gathered rather than read off the forward context: an MTP draft pass runs at the target's positions + 1. - `_launch_compressors` projects both compressors' inputs and gathers them in one collective; `Compressor.forward(combined=)` takes the result. V4-Pro (30 CSA + 31 HCA layers): 397 -> 245 collectives per prefill step. `pcp_reindex_ragged` sized its output with a blocking `.item()`, four times per prefill step (sixteen under PCP+TBO). The builder already computes the four prefill indptrs in numpy; it now keeps them on the metadata (`paged_prefill_indptr_cpu`) and each rank's local indptr is computed on the host (`pcp_local_indptr`), in both the full-batch and the per-request-group paths. `pcp_round_robin_split` returns unit strides: `.contiguous()` leaves a one-row shard at row stride pcp_size, which aiter's cached-positions RoPE rejects. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

After a PCP prefill the aux hooks hold only this rank's round-robin rows. `compute_draft_kv` all-gathered the three aux hidden states (~42 KiB a row) and every rank then projected all rows into window KV, the same work on every rank. The projection (main_proj + norm, then per stage wqkv_a + norm + RoPE + fake-quant) works row by row, so project only the local rows at their own positions, gather the projected KV rows (1 KiB a row a stage, one grouped collective for all stages), and write them as before. `DSparkLayer.write_context_kv` is split into `context_kv_rows` and `write_context_kv_rows`. A draft opts in with `context_kv_by_rows`; only V4 DSpark does, and any other draft keeps gathering the aux states. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…unts aiter's custom all-reduce, which `all_gatherv` tries first, runs at 6 ranks, but its all-gatherv kernel exists for 2, 4 and 8 only and throws otherwise. The attention-side gathers went through torch.distributed before the grouped gather, so PCP 6 worked; take one tensor at a time there again. Also assert that the grouped inputs share a row count (a mismatch threw on the custom path but passed on PyNCCL), and fix two review nits. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

The DSpark drafter branches on the KV cache dtype while tracing (`dspark_fp8_planned`), but `Config.compute_hash` left it out -- the known limitation #2566 deferred. A DSV4-Pro DSpark run with `--kv-cache-dtype bf16` after an fp8 one in the same container hashed to the same directory, loaded the fp8 drafter graph and died in warmup on every rank (`IndexError: list index out of range` on the split graph's outputs). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…eads them Under TBO a PCP gather runs on the comm stream, so its output is allocated there, while the compute stream reads it. When the forward frees it, the caching allocator hands the memory straight back to the comm stream, and the comm stream waits only on its own ubatch's compute: the other ubatch's next collective can overwrite it before this ubatch's kernels have read it. On main this corrupts the gathered extend K. DSV4-Pro, pcp8 + `--enable-tbo`, gsm8k (200 q): main 0.69 / 0.675 on two nodes. With only the extend-K gathers on the main code path: 0.71 (K alone: 0.70, positions alone: 0.975); the same plus `record_stream` on the gathered tensors: 0.965. The grouped gathers of this branch have the same shape, and score 0.96 only by timing. `tbo_mark_used_on_compute` records the compute stream on such tensors after `tbo_switch_to_compute_sync`, for the attention-side gathers (extend K and positions, the compressor inputs) and FusedMoE's PCP routed-rows gather and reduce-scatter. The DP and mori TBO paths have the same shape and are left alone here. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

The previous commit marked the comm-stream outputs by hand at four call sites, and nothing would notice one going missing: this branch scores 0.96 under TBO with or without the marks. `tbo_on_comm(fn, *args)` makes the comm phase one call -- yield, `fn` on the comm stream, switch back, hold every tensor it returns -- so a phase cannot skip the hold, and a test fails if V4 attention or `FusedMoE._run_experts` writes a phase out by hand. `tbo_mark_used_on_compute` now takes the compute stream from the TBO context instead of the current stream, so it no longer depends on being called after the switch back. `pcp_aux_rows` documents that it describes the base layer-hook capture; EAGLE3 strips aux off the target's output instead. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

github-actions · 2026-10-28T12:22:41Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2789 --add-label <label>

valarLip

Thanks — I reviewed HEAD 9d1db12b2. The core changes check out: the grouped gather, host-side indptr, record_stream in tbo_on_comm, and per-row DSpark KV projection. I also confirmed aiter's custom_all_gatherv returns a freshly allocated output (no aliasing with the registered buffer), the GPU/CPU kv_indptr_* copies are written in one place (_build_paged_prefill_meta) and stay consistent, and dummy runs don't enter the new DSpark branch.

Six inline comments. The two most worth acting on:

Under PCP + TBO, DSpark's context KV is reconstructed as one global stripe while the buffer is per-ubatch and padded per group, so draft KV lands on wrong rows (dspark_proposer.py:645).
The comm-stream lifetime hazard this PR fixes via tbo_on_comm remains on the DP / TP-allreduce / MORI paths that still hand-write yield → collective → switch (ubatching.py:615).

valarLip · 2026-10-29T06:08:41Z

+        gather each stage's KV rows (1 KiB a row) instead of the aux hidden
+        states."""
+        pcp_ws = get_pcp_world_size()
+        aux_local = torch.cat([buf[:n_local] for buf in self._aux_buffers], dim=-1)


Under PCP + TBO (per-request grouping) the draft context KV lands on the wrong rows. _write_split_context_kv treats the aux buffer as one round-robin stripe over the whole batch: global pcp_round_robin_split(positions) plus one rerange. But under TBO the hook writes [g0_local | g1_local] by ubatch_token_offset, and each group is padded separately (see model_runner._restore_pcp_balanced_output). Example, pcp=2 with TBO + DSpark and two requests of 3 and 5 tokens: g0 pads to 4 (2 local rows), g1 to 6 (3 local rows), so rank 0's buffer holds global positions [0, 2, 3, 5, 7]; this path instead takes buf[:4] as [0, 2, 4, 6], applies RoPE at those positions and reranges — the draft window KV is written at wrong rows/positions and acceptance drops. drafter.aux_for's older path has the same issue (it predates this PR), and the new path copies it; both should reconstruct per ubatch group.

Agreed on the mechanism, but this combination is refused at startup (llm_engine.py, #2125). Lifting that guard needs the aux gather, context KV and draft prefill split to follow the request groups together, so I left it for a follow-up (Known Limitations).

valarLip · 2026-10-29T06:08:41Z

+            t.record_stream(stream)
+
+
+def tbo_on_comm(fn: Callable, *args):


tbo_on_comm fixes the hazard PR-wide for the call sites that use it (output allocated on the TBO comm stream, read on compute, no record_stream), but the same function family still hand-writes yield → collective → switch without the mark: the DP path in moe.py (dp_gather_routed_rows, ~L5451, and DP reduce_scatterv / reduce_scatter_with_unpadding), module_dispatch_ops.tbo_all_reduce (~L245) and mori_prepare_finalize (~L288/375). With DP attention + TBO, or plain TP + TBO (tbo_all_reduce defaults to overlap mode), ubatch A's gathered hidden states go back to the comm stream as soon as they're freed in forward, and ubatch B's next-layer collective can overwrite them before A's experts / later kernels read them — exactly the mechanism behind the gsm8k 0.69 this PR measured. It's listed under Known Limitations; since the fix is a mechanical switch to tbo_on_comm, it would be better to sweep those sites in this PR.

Done: tbo_all_reduce and the DP gather / reduce-scatter now go through tbo_on_comm. Mori is left as is: its outputs are views of per-ubatch mori buffers, not allocator memory.

valarLip · 2026-10-29T06:08:42Z

+    ]
+    # aiter's custom all-gatherv kernel exists for 2, 4 and 8 ranks only, while
+    # its custom all-reduce (which `all_gatherv` tries first) also runs at 6.
+    if pcp_size not in (2, 4, 8):


The 6-rank fallback lives in this caller (hard-coded (2, 4, 8)) instead of in aiter's custom_all_gatherv, which doesn't check world size — dispatchAllGatherV throws "allgatherv world_size error" in its default branch. Other callers that go straight to device_communicator.all_gatherv remain unprotected: with pcp=6 (or dp=6) and custom all-reduce on, FusedMoE's PCP routed-rows gather (tbo_on_comm(gather) → all_gatherv(ts, heights, pcp_group)) and dp_gather_routed_rows take the custom path, hipMemcpyAsync into the registered pool, then hit the runtime_error and crash the process. The root fix is for aiter's custom_all_gatherv to return None when world_size ∉ {2, 4, 8}, so every caller falls through to PyNCCL and this special case can go.

Agreed. Fixed in ROCm/aiter#7149, which also gates the reduce-scatter and all-gather (they returned garbage at 6 ranks) and fixes an out-of-bounds read in the 6-rank one-stage all-reduce that made -pcp 6 decode fault. I dropped the ATOM special case; this PR now requires #7149 (-pcp 6 gsm8k 0.97).

valarLip · 2026-10-29T06:08:42Z

+    if pcp_size not in (2, 4, 8):
+        return [pcp_allgather_rerange(t, pcp_size) for t in inputs]
+    sent = [
+        t.contiguous().view(torch.uint8) if t.element_size() == 1 else t.contiguous()


This adds its own "one-byte dtypes travel as uint8" wire convention, duplicating fused_moe/routed_rows.gather_rows (which swaps in a float carrier by byte width), and the two disagree about the constraint: gather_rows' comment says PyNCCL and the custom all-gather only accept float dtypes, while this sends uint8 and int64 positions directly. Either that comment is stale for the target aiter, or uint8/int64 fail on the custom/PyNCCL paths here. Suggest reusing gather_rows (or moving it down into pcp_utils) so there's one convention and one statement of the constraint.

The gather_rows comment was stale: grouped all_gatherv copies bytes and PyNCCL accepts uint8 / int32 / int64. Unified into one gather_rows in pcp_utils (one-byte dtypes as uint8, the rest as is), shared with MoE.

valarLip · 2026-10-29T06:08:42Z

-        for compressor, stream in self._compressor_streams():
+        pairs = self._compressor_streams()
+        combined = (
+            self._pcp_gather_compressor_inputs(x, [c for c, _ in pairs])


Under PCP, combined is now allocated on the main stream (both wkv_gate and the gather moved out of Compressor) and then read by the compressors on alt_stream / indexer_stream, without record_stream; before, wkv_gate ran inside the side stream, so its output was allocated there. If _compressors_async() (in hipgraph, non-TBO) and _pcp_active() ever hold together, the combined reference is dropped when _launch_compressors returns, and the main stream's following Q/KV projections can reuse that memory before fused_compress_attn on the side stream reads it. Whether this triggers depends on PCP prefill running under a graph — it should be eager today, so this is plausible rather than confirmed.

Added record_stream for the side streams, with a test.

valarLip · 2026-10-29T06:08:42Z

+
+def test_unmarked_tensor_goes_straight_back_to_the_comm_stream(monkeypatch):
+    """Pins the hazard itself: if this stops holding, the test above no longer
+    shows anything."""


This positive control asserts the caching allocator hands the very same pointer back to the comm stream when unmarked, which depends on the allocator's implementation and config. With PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, a different PyTorch version, or changed block-split/rounding, the unmarked case may not get the same ptr and this would fail spuriously (flaky). Consider asserting the stream-use bookkeeping instead, or skipping when the allocator config differs.

The positive control is now skipped under expandable_segments.

… gathers Review follow-ups on #2789: - tbo_all_reduce and FusedMoE's DP gather / reduce-scatter run through tbo_on_comm, like the PCP comm phases, so their comm-stream outputs are held for the compute stream; the hold walks nested lists and tuples. Mori is unchanged: it returns views of its own per-ubatch buffers. - Under a graph the V4 compressors read their PCP-gathered inputs on side streams; hold them for those streams. - gather_rows moves to pcp_utils and carries one-byte dtypes as uint8 and the rest as is (aiter's all_gatherv copies bytes; only its PyNCCL fallback lacks fp8 / e8m0). MoE's DP, PCP and comm-fused gathers and the PCP grouped gather share it. - The allocator-reuse positive control skips under expandable_segments. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

tbo_all_reduce imports it ahead of the TBO branch, and the non-GPU CI's aiter has no aiter.dist.communication_op. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

ROCm/aiter#7149 makes custom_all_gatherv decline group sizes it has no kernel for, so all_gatherv falls back to NCCL there. The grouped gather now takes the one-collective path at every PCP size, 6 included. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

yitingw1 · 2026-10-30T12:24:45Z

@valar-cop review this pr

valar-cop

some changes needed

any perf number for tp>1 / pp+pcp? and a gpu run of pcp+tbo and pcp+dspark?

ci green.

valar-cop · 2026-10-30T12:32:06Z

+        with record_function(f"draft_kv[bs={scheduled_bs} tok={n}]"):
+            ctx_local = self.model.project_context(aux_local)
+            layers = list(self.model.context_layers)
+            kv_rows = pcp_allgather_rerange_many(


missing tbo_on_comm here? this is the same hazard we fix elsewhere..

Not needed here: compute_draft_kv runs on the main thread after run_model, outside the TBO ubatch threads, so tbo_active() is false and tbo_on_comm would call straight through. PCP + speculative decoding + TBO is also refused at startup (#2125).

valar-cop · 2026-10-30T12:32:06Z

            setattr(attn_metadata, ind_attr, new_indptr)
            setattr(attn_metadata, idx_attr, new_indices)
+            local_indptr_cpu[ind_attr] = local_cpu
+        attn_metadata.paged_prefill_indptr_cpu = local_indptr_cpu


overwrites the full-batch indptr in place, a second call would silently give too-short kv. keep it local instead?

Each metadata object is reindexed once: the non-TBO path reindexes the full batch once, and the PCP+TBO path reindexes each group's fresh ub once while reading the untouched full src. The device indptrs are already replaced in place (as on main). Writing the host copy back keeps the two in agreement; keeping it local would leave a host copy describing the full batch next to rank-local device indptrs.

_apply_pcp_reindex now derives each rank's indptr from the host copy the builder publishes (paged_prefill_indptr_cpu); the test builds its metadata by hand, so it sets that copy beside kv_indptr_extend. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

yitingw1 · 2026-10-30T13:02:06Z

Re the review: perf was measured at -tp 1 -pcp 8 only (Known Limitations); -tp 4 -pcp 2 was checked for correctness (20K-token prompts, gsm8k 0.97), and PP+PCP was not run because there is no PP+PCP deployment yet; we will test it when one comes up. The GPU runs are in the Test Result table: PCP+TBO gsm8k 0.965 (0.97 on the final commit), and PCP+DSpark 0.97 with fp8 KV and 0.98 with bf16 KV. PCP+DSpark+TBO is refused at startup.

valar-cop

LGTM, but please address the pcp 6 validation and dspark collective..

please re-run ATOM Test for fd0893d. any pcp+graph perf number? gsm8k?

ci green.

valar-cop · 2026-10-30T13:34:12Z

            setattr(attn_metadata, ind_attr, new_indptr)
            setattr(attn_metadata, idx_attr, new_indices)
+            local_indptr_cpu[ind_attr] = local_cpu
+        attn_metadata.paged_prefill_indptr_cpu = local_indptr_cpu


this field is documented as the global host copies, now we write local ids into it.. use a separate field?

Reworded in fd912b7: the host copies mirror whichever kv_indptr_* the metadata currently carries; after either reindex (rank or PCP+TBO group) both describe the same rank- or group-local queries.

valar-cop · 2026-10-30T13:34:12Z

+    ]
+    comm = get_pcp_group().device_communicator
+    gathered = gather_rows(
+        lambda ts: comm.all_gatherv(ts, dim=0), [t.contiguous() for t in inputs]


pcp 6 will die inside aiter without #7149, and nothing checks the size at startup. validate at engine init?

Intentional. -pcp 6 depends on ROCm/aiter#7149, which lands first. Restoring the fallback would not make old aiter work either: there the 6-rank FusedMoE all_gatherv throws, and the 6-rank 1-stage custom all-reduce reads out of bounds during decode graph capture, so -pcp 6 with custom AR is already broken on main.

valar-cop · 2026-10-30T13:34:12Z

+        while a compressor pools `ratio` consecutive tokens; its plan and
+        state slots are already full-sequence. Mirrors SGLang's
+        compute_kv_score: gather after the projection, before the pooling."""
+        projected = [c.wkv_gate(x) for c in compressors]


with pcp the gather sits on main stream ahead of the side-stream launches, the overlap is lost. measure or note it?

Not reachable. The compressor side streams are used only when fc.in_hipgraph, and the PCP split is prefill-only. DSV4-Pro never records a prefill graph (supports_token_graph=False, prefill_req_graph_sizes=()), and DSV4.1 refuses context parallel. In the rank-0 traces of our main-vs-PR A/B, every compressor kernel in a prefill step runs on the main stream on both main and this PR.

valar-cop · 2026-10-30T13:34:12Z

+        with record_function(f"draft_kv[bs={scheduled_bs} tok={n}]"):
+            ctx_local = self.model.project_context(aux_local)
+            layers = list(self.model.context_layers)
+            kv_rows = pcp_allgather_rerange_many(


why not route this through tbo_on_comm like the others? lifting the refusal would bring the bug back..

Not needed, even with the startup guard lifted. compute_draft_kv runs on the main thread after run_model returns, outside the TBO ubatch threads, so tbo_active() is false there and the gather never crosses streams. The EAGLE / MTP drafter gathers are in the same position, so I left all of them as is.

The file failed 20 of 22 cases on main with RUN_H2D_GPU_TESTS=1, all from fixtures that fell behind the code: - set_pcp patched pcp_round_robin_query_indices (removed in #2688); the builder reads get_pcp_rank. - make_runner lacked forward_vars["batch_info"], config.max_num_seqs and the builder's csa_layers, max_model_len_idx, _local_heads and index_logits_workspace. - The padded decode case modelled the bucket width through UBatchWrapper._decode_ub_running_bs; since #2692 the split comes from _prepare_ubatch_decode's split_bs, so the test passes the bucket. - prepare_prefill was called with a total_tokens_num the fake batch lacks. 22/22 pass on this branch and on its merge base. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

valar-cop

some changes needed, the pcp 6 fallback removal is risky..

any perf number at tp>1? and the graph-pool footprint from record_stream is still not measured..

ci green.

valar-cop · 2026-10-30T14:32:01Z

+    assert all(t.shape[0] == inputs[0].shape[0] for t in inputs), [
+        t.shape[0] for t in inputs
+    ]
+    comm = get_pcp_group().device_communicator


aiter has no kernel for pcp 6 yet, and we pin no aiter version. keep the per-tensor fallback until #7149 lands?

Intentional. -pcp 6 depends on ROCm/aiter#7149, which lands first. Restoring the fallback would not make old aiter work either: there the 6-rank FusedMoE all_gatherv throws, and the 6-rank 1-stage custom all-reduce reads out of bounds during decode graph capture, so -pcp 6 with custom AR is already broken on main.

valar-cop · 2026-10-30T14:32:01Z

+    queries' segment lengths in `indptr`. Owned ids past its last query are
+    padding queries, whose segments are empty."""
+    lens = np.diff(indptr)
+    lens = np.pad(lens, (0, max(int(owned.max(initial=-1)) + 1 - lens.shape[0], 0)))


nothing checks the host indptr matches the device one. add a cheap shape/last-value assert so a mismatch fails loud instead of giving wrong indices?

Added in fd912b7: a length check at both reindex entries (rank and PCP+TBO group). It reads shapes only; comparing values would need a device read-back, which is the sync this PR removes.

valar-cop · 2026-10-30T14:32:01Z

+        while a compressor pools `ratio` consecutive tokens; its plan and
+        state slots are already full-sequence. Mirrors SGLang's
+        compute_kv_score: gather after the projection, before the pooling."""
+        projected = [c.wkv_gate(x) for c in compressors]


under PCP+graph this loses the compressor overlap, both wkv_gate now serialize on main. keep the projections on side streams? at least note it in known limitations..

Not reachable. The compressor side streams are used only when fc.in_hipgraph, and the PCP split is prefill-only. DSV4-Pro never records a prefill graph (supports_token_graph=False, prefill_req_graph_sizes=()), and DSV4.1 refuses context parallel. In the rank-0 traces of our main-vs-PR A/B, every compressor kernel in a prefill step runs on the main stream on both main and this PR.

Both PCP reindex paths (rank and PCP+TBO group) compute local offsets from the host copy and reindex the GPU copy; a length mismatch would silently reindex against a different query set. Clarify that after either reindex both copies describe the same rank- or group-local queries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

yitingw1 · 2026-11-02T02:54:09Z

Thanks for the review. Pushed fd912b7 (host/device indptr length asserts + docstring); replies inline. On the summary points:

Tests / CI: CI was green on 319ffe1; fd912b7 only adds two asserts and a docstring line. PCP + graph on fd912b7: -tp 1 -pcp 8 gsm8k 5-shot (200 q) 0.975 with DSpark k=5 (acceptance 63.7%), 0.96 with --enable-tbo prefill.
TP>1 perf: not measured on purpose. TP>1 combined with PCP>1 is slower than either alone and is not deployed, so we only keep it correct (-tp 4 -pcp 2 gsm8k 0.97).
Graph pool footprint of record_stream: measured at -tp 8 --enable-tbo all (DSV4-Pro, FULL graphs, 8 of 10 buckets recorded split), same tree with and without the hold. During capture each rank holds 1952 tensors (135 MiB in total); the graph pool grows by 0.00 GB in both arms and KV blocks are identical. Known Limitation 2 updated.

valarLip

Re-read at c889ff9, by reading the code and running the 7 affected CPU test files (163 pass). Thanks for the rounds since 10-29. The tbo_on_comm sweep, the unified gather_rows, the side-stream record_stream, and the graph-pool measurement all read right. I'm not repeating what's been answered: the PCP+DSpark+TBO startup refusal, -pcp 6 on aiter#7149, and the indptr length asserts.

Performance (the three I'd most like addressed)

This PR's point is to stop every rank redoing the whole prefill. Three places still pay per step for work that's already known on the host, or that each rank repeats:

pcp_reindex_ragged uploads with upload_numpy every forward, per field (pcp_utils.py:322). upload_numpy's own docstring says "A buffer uploaded every forward should be a CpuGpuBuffer". That's 4 pageable H2Ds per prefill step, and 12+ under PCP+TBO, plus the arange / searchsorted / gather kernels behind each.
The PCP+TBO group path reindexes a contiguous range with the general machinery (deepseek_v4_attn.py:3333). A group owns [gts, gte). For a contiguous range, the indices are a view indices[cpu[gts]:cpu[gte]] and the indptr is cpu[gts:gte+1] - cpu[gts]: no kernel, no upload. Today each group pays 3–4 kernels and an upload per field.
DSpark gathers bf16 KV rows and then every rank quantizes all n rows (dspark_proposer.py:658). The quantization is per row, so quantizing the local rows first and gathering the 2buff result cuts the wire (by my estimate about 640 B per row instead of 1 KiB) and removes the W-fold repeated quantization.

Items 1 and 2 share a root fix, which is the first inline comment on deepseek_v4_attn.py:406: build the rank- or group-local indptr in the builder's numpy pass and upload it once.

Structure

The rest is inline: the proposer reaching into the draft model, gather_rows' home, the two all-gather paths, Compressor's two modes, and the TBO comm-phase API.

valarLip · 2026-11-02T07:36:38Z

    (same `kv` shared by all 3 ratios; one builder pass)."""
    kv_indptr_extend: torch.Tensor | None = None
    """[total_tokens + 1] int32 GPU — packed cumsum of `extend_count`."""
+    paged_prefill_indptr_cpu: dict[str, np.ndarray] | None = None


paged_prefill_indptr_cpu mirrors the four kv_indptr_* fields in a string-keyed dict, so two copies have to be kept in step by hand. The length asserts added in fd912b7 are there to catch them drifting. "paged" also doesn't fit kv_indptr_extend. The builder already has these indptrs in numpy before upload. Computing the rank-local (or PCP+TBO group-local) indptr there and uploading only that would remove the mirror, the asserts and the per-step re-uploads in pcp_reindex_ragged (performance item 1) in one change.

valarLip · 2026-11-02T07:36:38Z

-    )
-    torch.cumsum(lens, dim=0, out=indptr_local[1:])
-    total = int(indptr_local[-1].item())
+    indptr_local = upload_numpy(indptr_local_cpu, device)


Performance item 1: upload_numpy here runs every prefill forward, for each of the four fields. Its docstring reserves it for one-off and size-varying uploads ("A buffer uploaded every forward should be a CpuGpuBuffer"). With the builder producing the local indptr (comment on deepseek_v4_attn.py:406), this upload and the arange / searchsorted / gather that follow go away. If it stays, it needs a persistent CpuGpuBuffer.

valarLip · 2026-11-02T07:36:38Z

@@ -3317,6 +3331,8 @@ def _build_ubatch_prefill_metadata_balanced(

        # ---- per-token DSV4 fields sliced by the GLOBAL token range [gts,gte) ----
        owned = torch.arange(gts, gte, device=device)


Performance item 2: a group's queries are the contiguous global range [gts, gte), and the host indptr is already here (src_cpu). For a contiguous range, the local indices are indices[src_cpu[gts]:src_cpu[gte]], a view, and the local indptr is src_cpu[gts:gte+1] - src_cpu[gts]. No torch.arange(gts, gte), no pcp_local_indptr, and no pcp_reindex_ragged kernels or upload. The general reindex is only needed for the round-robin rank split.

valarLip · 2026-11-02T07:36:38Z

+                pcp_ws,
+            )
+            for layer, rows in zip(layers, kv_rows):
+                layer.write_context_kv_rows(rows[:n], positions)


Performance item 3: each rank gathers bf16 KV rows (1 KiB a row) for all n rows, then write_context_kv_rows quantizes every one of them (quantize_bf16_to_v4_2buff_triton, deepseek_v4_dspark.py:634). So the quantization runs W times over. It's per row, so quantizing the local n_local rows before the gather and gathering the 2buff tensors through the same pcp_allgather_rerange_many call moves fewer bytes and quantizes each row once.

valarLip · 2026-11-02T07:36:38Z

        ):
            self.model.write_context_kv(main_hidden_all, positions)

+    def _write_split_context_kv(


_write_split_context_kv reads _aux_buffers, walks model.context_layers and calls each layer's two halves (context_kv_rows, write_context_kv_rows) from the proposer. That bypasses DSparkDraftModel.write_context_kv and its dummy-run guard, and adds a context_kv_by_rows switch to pick the path. compute_draft_kv's docstring says this hook stays flavor-agnostic. "Project the local rows, gather, write" belongs in the model's write_context_kv (given the local rows and the split), so the proposer needs neither the layer internals nor the switch. Today _pcp_active keeps dummy runs out, but that is the only thing keeping the bypassed guard from mattering.

valarLip · 2026-11-02T07:36:38Z

+    return _rerange(gathered, pcp_size, input_.shape[0])
+
+
+def gather_rows(gather, rows):


gather_rows now lives in pcp_utils, but its callers include DP (dp_gather_routed_rows) and the TP/DP comm_fused_moe, so comm_fused_moe imports a general gather from the PCP module. It belongs in a general distributed comm module (next to token_transport or similar), with PCP importing it too.

valarLip · 2026-11-02T07:36:38Z

+    return tuple(g.view(t.dtype) for g, t in zip(gathered, rows))
+
+
+def pcp_allgather_rerange_many(


pcp_allgather_rerange (group.all_gather) and pcp_allgather_rerange_many (device_communicator.all_gatherv via gather_rows) do the same thing by two transports, with different dtype handling and fallbacks. The single-tensor one could be _many([x])[0]. Callers that gather several tensors back to back should then use one call:

drafter.py:490, which this PR touches, gathers each aux buffer separately;

attention_mla.py:3079-3081, three gathers;

deepseek_v2.py:1732-1733 and :2857-2861.

Fixing it here and leaving those (Fix-then-sweep) keeps both paths alive.

valarLip · 2026-11-02T07:36:38Z

        state_slot_in: torch.Tensor,  # [bs] int32
        state_slot_out: torch.Tensor,  # [bs] int32
        block_tables: torch.Tensor | None = None,  # [bs, max_blocks_per_seq] int32
+        combined: torch.Tensor | None = None,  # [T_full, 2*coff*head_dim]


Compressor.forward(combined=None) now has two modes: project x itself, or take a projection done outside. And Attention calls c.wkv_gate(x) directly. How a compressor projects its input now lives in two classes. Under PCP, x is only used for a shape assert, and any later change before wkv_gate would silently miss _pcp_gather_compressor_inputs. A Compressor.project(x) that both paths call, with forward taking only the projection, keeps it in one class.

valarLip · 2026-11-02T07:36:38Z

+    tbo_yield_and_switch_from_compute_to_comm()
+    out = fn(*args)
+    tbo_switch_to_compute_sync()
+    tbo_mark_used_on_compute(*(out if isinstance(out, (list, tuple)) else (out,)))


tbo_mark_used_on_compute is used only inside tbo_on_comm, but it's exported from atom.utils.tbo (__init__.py:29, :63). That invites exactly the hand-written "comm, then mark" this PR replaces. Please make it private and drop the export. Since it already walks lists and tuples, this line can be _mark(out).

valarLip · 2026-11-02T07:36:38Z

+    assert log == [(id(c), s) for c, s in zip(combined, streams)]
+
+
+def test_tbo_comm_phases_go_through_tbo_on_comm():


This follows my 10-29 comment on hand-written comm phases. Scanning inspect.getsource of four functions catches the current spelling only: a helper or a renamed call passes, and new call sites aren't scanned. Making tbo_yield_and_switch_from_compute_to_comm private (with MoRI's explicit exception), so a comm phase can only be written through tbo_on_comm, enforces it by the API rather than by text.

valar-cop · 2026-11-02T07:49:22Z

ci triage for c889ff9:

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/41407422465

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/41407422373

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/41407422457

yitingw1 and others added 7 commits October 28, 2026 08:36

Merge remote-tracking branch 'origin/main' into yitingw1/pcp_opt

9d1db12

yitingw1 added the ci:atom label Oct 28, 2026

zufayu requested a review from JiaoliangYu October 29, 2026 01:00

zufayu assigned JiaoliangYu Oct 29, 2026

valarLip reviewed Oct 29, 2026

View reviewed changes

yitingw1 and others added 2 commits October 29, 2026 08:55

test(tbo): stub aiter's communication_op in the TP all-reduce test

e8ff821

tbo_all_reduce imports it ahead of the TBO branch, and the non-GPU CI's aiter has no aiter.dist.communication_op. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

yitingw1 mentioned this pull request Oct 29, 2026

[HIP] fix(custom_ar): make 6-rank custom collectives correct ROCm/aiter#7149

Merged

1 task

valar-cop reviewed Oct 30, 2026

View reviewed changes

Merge branch 'main' into yitingw1/pcp_opt

c889ff9

valarLip reviewed Nov 2, 2026

View reviewed changes

		t.record_stream(stream)


		def tbo_on_comm(fn: Callable, *args):

		@@ -3317,6 +3331,8 @@ def _build_ubatch_prefill_metadata_balanced(

		# ---- per-token DSV4 fields sliced by the GLOBAL token range [gts,gte) ----
		owned = torch.arange(gts, gte, device=device)

		return _rerange(gathered, pcp_size, input_.shape[0])


		def gather_rows(gather, rows):

		return tuple(g.view(t.dtype) for g, t in zip(gathered, rows))


		def pcp_allgather_rerange_many(

		assert log == [(id(c), s) for c, s in zip(combined, streams)]


		def test_tbo_comm_phases_go_through_tbo_on_comm():

Conversation

yitingw1 commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

1. Grouped gathers

2. Reindex without readback

3. DSpark context KV

4. PCP+TBO: comm-stream outputs read on the compute stream

5. Compile cache key

Test Plan

Test Result

Known Limitations

Submission Checklist

Uh oh!

github-actions Bot commented Oct 28, 2026

🏷️ CI Guide

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yitingw1 commented Oct 30, 2026

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yitingw1 commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

yitingw1 commented Oct 28, 2026 •

edited

Loading

yitingw1 commented Oct 30, 2026 •

edited

Loading