[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 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. One-byte dtypes travel as uint8: past its custom
+    all-gather size limit `all_gatherv` falls back to PyNCCL, which has no fp8
+    dtype."""
+    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
+    ]
+    # 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):
+        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()
+        for t in inputs
+    ]
+    gathered = get_pcp_group().device_communicator.all_gatherv(sent, dim=0)
+    return [
+        _rerange(g.view(t.dtype), 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/moe.py

@@ -112,6 +112,7 @@
 from atom.utils.tbo.ubatching import (
     tbo_active,
     tbo_enabled,
+    tbo_on_comm,
     tbo_switch_to_compute_sync,
     tbo_yield_and_switch_from_compute_to_comm,
 )
@@ -5372,19 +5373,17 @@ def _run_experts(
                 topk_ids,
                 x_scale,
             )
-            tbo = tbo_active()
-            if tbo:
-                tbo_yield_and_switch_from_compute_to_comm()
             # One grouped collective, rank-major: every PCP shard is one height.
             pcp_group = get_pcp_group()
             heights = [hidden_states.shape[0]] * pcp_group.world_size
-            hidden_states, topk_weights, topk_ids, *scale = gather_rows(
-                lambda ts: all_gatherv(ts, heights, pcp_group), rows
-            )
+
+            def gather():
+                out = gather_rows(lambda ts: all_gatherv(ts, heights, pcp_group), rows)
+                reassign_gathered_shared_ids(self, out[2])
+                return out
+
+            hidden_states, topk_weights, topk_ids, *scale = tbo_on_comm(gather)
             x_scale = scale[0] if scale else None
-            reassign_gathered_shared_ids(self, topk_ids)
-            if tbo:
-                tbo_switch_to_compute_sync()
         gathered = gathered or split
         topk_weights, topk_ids = self.quant_method.place_experts(
             self, topk_weights, topk_ids, gathered
@@ -5407,11 +5406,7 @@ def _run_experts(
             ),
         )
         if split:
-            if tbo:
-                tbo_yield_and_switch_from_compute_to_comm()
-            output = pcp_reduce_scatter(output)
-            if tbo:
-                tbo_switch_to_compute_sync()
+            output = tbo_on_comm(pcp_reduce_scatter, output)
         elif pcp:
             output = pcp_all_reduce(output)
         return output