[PCP] Fewer V4 prefill gathers and fix PCP+TBO accuracy

atom/config.py

@@ -2822,6 +2822,10 @@ def compute_hash(self) -> str:
         factors.append(bool(getattr(self.dcp_config, "indexer_dcp_only", False)))
         factors.append(self.enable_dp_attention)
         factors.append(self.index_cache_dtype)
+        # The KV cache dtype is fixed at trace time too: the DSpark drafter
+        # branches on it (`dspark_fp8_planned`), so a bf16 run over an fp8 run's
+        # cache loads a graph built for the other window layout.
+        factors.append(self.kv_cache_dtype)
         # DeepSeek-V4.1 decoder SWA bounded replay compiles the backbone as two
         # graphs (early + late layers) instead of one; both compile under the
         # `backbone` tag, so the flag must key the artifacts.

atom/distributed/pcp_utils.py

@@ -24,6 +24,8 @@
     get_prefill_context_model_parallel_world_size,
 )
 
+from atom.utils import upload_numpy
+
 
 class PcpBalGroup(NamedTuple):
     """One request group for PCP+TBO request-boundary split prefill.
@@ -105,7 +107,12 @@ def pcp_round_robin_split(
     # Divisibility by pcp_size is guaranteed upstream by pcp_pad_len (callers
     # pad before splitting); the view below would error if violated.
     rest = tuple(input_.shape[1:])
-    return input_.view(-1, pcp_size, *rest)[:, pcp_rank].contiguous()
+    # Not `.contiguous()`: a one-row shard counts as contiguous while its row
+    # stride is still pcp_size, which kernels that check strides reject (aiter's
+    # cached-positions RoPE).
+    return input_.view(-1, pcp_size, *rest)[:, pcp_rank].clone(
+        memory_format=torch.contiguous_format
+    )
 
 
 def pcp_allgather_rerange(
@@ -127,16 +134,47 @@ def pcp_allgather_rerange(
     group = get_pcp_group()
     # aiter all_gather(dim=0) returns rank-major concat: [pcp*L, *rest].
     gathered = group.all_gather(input_.contiguous(), dim=0)
-    local_len = input_.shape[0]
-    rest = tuple(input_.shape[1:])
+    return _rerange(gathered, pcp_size, input_.shape[0])
+
+
+def gather_rows(gather, rows):
+    """`gather(list of tensors) -> list of tensors` over `rows`, one-byte
+    dtypes carried as uint8: aiter's grouped `all_gatherv` copies bytes
+    whatever the dtype, but its PyNCCL fallback has no fp8 / e8m0 / fp4x2."""
+    gathered = gather(
+        [t.view(torch.uint8) if t.element_size() == 1 else t for t in rows]
+    )
+    return tuple(g.view(t.dtype) for g, t in zip(gathered, rows))
+
+
+def pcp_allgather_rerange_many(
+    inputs: list[torch.Tensor], pcp_size: int | None = None
+) -> list[torch.Tensor]:
+    """`pcp_allgather_rerange` of every tensor in `inputs` (same row count) in
+    one grouped collective."""
+    if pcp_size is None:
+        pcp_size = get_pcp_world_size()
+    if pcp_size <= 1:
+        return list(inputs)
+    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
+    gathered = gather_rows(
+        lambda ts: comm.all_gatherv(ts, dim=0), [t.contiguous() for t in inputs]
+    )
+    return [_rerange(g, pcp_size, t.shape[0]) for g, t in zip(gathered, inputs)]
+
+
+def _rerange(gathered: torch.Tensor, pcp_size: int, local_len: int) -> torch.Tensor:
     # rank-major [pcp, L, *rest] -> transpose -> token-major [L, pcp, *rest]
     # -> flatten to [L*pcp, *rest] == original global order.
-    out = (
+    rest = tuple(gathered.shape[1:])
+    return (
         gathered.view(pcp_size, local_len, *rest)
         .transpose(0, 1)
         .reshape(pcp_size * local_len, *rest)
     )
-    return out
 
 
 # ==== MoE-path PCP collectives (rank-major gather + reduce_scatter) ====
@@ -246,10 +284,22 @@ def pcp_pad_dense(t: torch.Tensor, n_pad: int) -> torch.Tensor:
     return torch.cat([t, t.new_zeros(n_pad, *t.shape[1:])], dim=0)
 
 
+def pcp_local_indptr(indptr: np.ndarray, owned: np.ndarray) -> np.ndarray:
+    """Host side of `pcp_reindex_ragged`: the prefix sum of the `owned`
+    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)))
+    out = np.zeros(owned.shape[0] + 1, dtype=np.int32)
+    np.cumsum(lens[owned], out=out[1:])
+    return out
+
+
 def pcp_reindex_ragged(
     kv_indptr: torch.Tensor,  # [T_global + 1] int32 — global per-query prefix sum
     kv_indices: torch.Tensor,  # [kv_indptr[-1]] — ragged packed values
     owned_q: torch.Tensor,  # [T_local] long — global query ids this rank owns
+    indptr_local_cpu: np.ndarray,  # [T_local + 1] — `pcp_local_indptr`
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Reindex a ragged (indptr, indices) pair down to this rank's queries.
 
@@ -262,17 +312,15 @@ def pcp_reindex_ragged(
     Used to shard the per-query prefill index buffers (kv_indptr/kv_indices
     _prefix_swa / _extend) to 1/W while the values themselves still point into
     the full KV (paged unified_kv) / full extend kv tensor.
+
+    `indptr_local` comes from the host (`indptr_local_cpu`), so sizing the
+    output reads nothing back from the device.
     """
     device = kv_indptr.device
     owned_q = owned_q.to(device)
     starts = kv_indptr[owned_q]  # [T_local]
-    ends = kv_indptr[owned_q + 1]  # [T_local]
-    lens = ends - starts  # [T_local] per-owned-query segment length
-    indptr_local = torch.zeros(
-        owned_q.shape[0] + 1, dtype=kv_indptr.dtype, device=device
-    )
-    torch.cumsum(lens, dim=0, out=indptr_local[1:])
-    total = int(indptr_local[-1].item())
+    indptr_local = upload_numpy(indptr_local_cpu, device)
+    total = int(indptr_local_cpu[-1])
     if total == 0:
         return indptr_local, kv_indices.new_empty(0)
     # Build a gather map: for each output slot, which source index to read.

atom/model_ops/attentions/deepseek_v4_attn.py

@@ -55,6 +55,7 @@
     get_pcp_rank,
     get_pcp_world_size,
     pcp_is_enabled,
+    pcp_local_indptr,
     pcp_owned_rows_by_request,
     pcp_pad_dense,
     pcp_pad_indptr,
@@ -388,6 +389,9 @@ class AttentionMetaData_DSV4(AttentionMetaData):
     (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
+    """Host copies of the four `kv_indptr_*` above, keyed by field name. PCP
+    reindexes from these, so it never reads the GPU copies back."""
 
 
 class DeepseekV4Backend(AttentionBackend):
@@ -2868,10 +2872,13 @@ def _apply_pcp_reindex(
         # on its own pcp pad). Either way the divisor is pcp_size.
         padded_total = pcp_pad_len(total_tokens, pcp_size)
         n_pad = padded_total - total_tokens
-        # This rank's round-robin queries, generated on the GPU.
+        # This rank's round-robin queries, on the GPU and on the host.
         owned_q = torch.arange(get_pcp_rank(), padded_total, pcp_size, device=device)
+        owned_cpu = np.arange(get_pcp_rank(), padded_total, pcp_size)
 
         # --- ragged per-query buffers: pad indptr to padded_total, then 1/W ---
+        indptr_cpu = attn_metadata.paged_prefill_indptr_cpu
+        local_indptr_cpu = {}
         for ind_attr, idx_attr in (
             ("kv_indptr_prefix_swa", "kv_indices_prefix_swa"),
             ("kv_indptr_prefix_csa", "kv_indices_prefix_csa"),
@@ -2883,9 +2890,14 @@ def _apply_pcp_reindex(
             if indptr is None or indices is None:
                 continue
             indptr = pcp_pad_indptr(indptr, n_pad)  # dummy queries: 0-length KV
-            new_indptr, new_indices = pcp_reindex_ragged(indptr, indices, owned_q)
+            local_cpu = pcp_local_indptr(indptr_cpu[ind_attr], owned_cpu)
+            new_indptr, new_indices = pcp_reindex_ragged(
+                indptr, indices, owned_q, local_cpu
+            )
             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
 
         # --- dense per-token fields: pad then round-robin-slice to 1/W ---
         if attn_metadata.skip_prefix_len_csa is not None:
@@ -3184,6 +3196,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)
+        owned_cpu = np.arange(gts, gte)
+        ub.paged_prefill_indptr_cpu = {}
         for ind_attr, idx_attr in (
             ("kv_indptr_prefix_swa", "kv_indices_prefix_swa"),
             ("kv_indptr_prefix_csa", "kv_indices_prefix_csa"),
@@ -3194,7 +3208,11 @@ def _build_ubatch_prefill_metadata_balanced(
             indices = getattr(src, idx_attr, None)
             if indptr is None or indices is None:
                 continue
-            ni, nx = pcp_reindex_ragged(indptr, indices, owned)
+            local_cpu = pcp_local_indptr(
+                src.paged_prefill_indptr_cpu[ind_attr], owned_cpu
+            )
+            ni, nx = pcp_reindex_ragged(indptr, indices, owned, local_cpu)
+            ub.paged_prefill_indptr_cpu[ind_attr] = local_cpu
             # kv_indices_extend are ROW offsets into the per-fwd kv_full tensor.
             # In the full metadata they index the WHOLE sequence's kv_full [0,T);
             # for this group kv_full only holds the group's tokens (global order
@@ -3856,6 +3874,12 @@ def _build_paged_prefill_meta(
         attn_metadata.kv_indptr_prefix_csa = csa_indptr
         attn_metadata.kv_indices_prefix_hca = hca_indices[:hca_total]
         attn_metadata.kv_indptr_prefix_hca = hca_indptr
+        attn_metadata.paged_prefill_indptr_cpu = {
+            "kv_indptr_extend": ext_indptr_np,
+            "kv_indptr_prefix_swa": swa_indptr_np,
+            "kv_indptr_prefix_csa": csa_indptr_np,
+            "kv_indptr_prefix_hca": hca_indptr_np,
+        }
         attn_metadata.skip_prefix_len_csa = skip_csa_gpu
         attn_metadata.envelope_rows = envelope_rows
 

atom/model_ops/fused_moe/comm_fused_moe.py

@@ -17,7 +17,7 @@
 from aiter.ops.flydsl.moe_common import GateMode
 
 from atom.config import get_current_atom_config
-from atom.model_ops.fused_moe.routed_rows import gather_rows
+from atom.distributed.pcp_utils import gather_rows
 from atom.model_ops.topK import reassign_gathered_shared_ids
 from atom.utils import envs
 from atom.utils.custom_register import direct_register_custom_op

atom/model_ops/fused_moe/routed_rows.py

@@ -56,18 +56,6 @@ def pack_routed_rows(
     return rows if x_scale is None else (*rows, x_scale)
 
 
-_CARRIERS = {4: torch.float32, 2: torch.bfloat16, 1: torch.bfloat16}
-
-
-def gather_rows(gather, rows):
-    """`gather(list of tensors) -> list of tensors` over `rows`, each carried
-    as the float dtype of its byte width (int32 ids as fp32; fp8 / e8m0 byte
-    pairs as bf16, so a one-byte row's width must be even): a gather moves
-    bytes, but PyNCCL and the custom all-gather take float dtypes only."""
-    gathered = gather([t.view(_CARRIERS[t.element_size()]) for t in rows])
-    return tuple(g.view(t.dtype) for g, t in zip(gathered, rows))
-
-
 def resolve_routed_rows_wire(
     in_dtype: torch.dtype,
     hidden_dim: int,

atom/model_ops/module_dispatch_ops.py

@@ -220,36 +220,26 @@ def tbo_all_reduce(x: torch.Tensor) -> torch.Tensor:
     if not tbo_active():
         return tensor_model_parallel_all_reduce(x)
 
-    # Default "inline": correct Plan-A baseline (no overlap, never hangs). Only
-    # move the AR onto the comm stream when explicitly opted into "overlap".
+    # "inline": the plain all-reduce on the compute stream (no overlap).
+    # "overlap" (default): on the comm stream, overlapping the partner ubatch.
     if envs.ATOM_TBO_TP_AR_MODE != "overlap":
         return tensor_model_parallel_all_reduce(x)
 
     from atom.utils.tbo.ubatching import (
         tbo_current_ubatch_id,
         tbo_get_comm_stream,
         tbo_get_ubatch_tp_comm,
-        tbo_switch_to_compute_sync,
-        tbo_yield_and_switch_from_compute_to_comm,
+        tbo_on_comm,
     )
 
-    ubatch_id = tbo_current_ubatch_id()
-
-    ub_comm = tbo_get_ubatch_tp_comm(ubatch_id)
+    ub_comm = tbo_get_ubatch_tp_comm(tbo_current_ubatch_id())
     if ub_comm is None:
         # world_size == 1: nothing to reduce.
         return x
 
-    # Hand the CPU baton to the partner ubatch and move onto the comm stream,
-    # so this AR overlaps the partner's compute.
-    tbo_yield_and_switch_from_compute_to_comm()
-    comm_stream = tbo_get_comm_stream()
-    # out-of-place all_reduce on this ubatch's dedicated communicator, launched
-    # on the comm stream (current stream after the switch above).
-    x = ub_comm.all_reduce(x, stream=comm_stream)
-
-    tbo_switch_to_compute_sync()
-    return x
+    # Out-of-place all_reduce on this ubatch's dedicated communicator, on the
+    # comm stream, overlapping the partner ubatch's compute.
+    return tbo_on_comm(lambda: ub_comm.all_reduce(x, stream=tbo_get_comm_stream()))
 
 
 def _tbo_all_reduce_fake(x: torch.Tensor) -> torch.Tensor: