[Performance] Balance the latent-MQA indexer's CP rows with a dual-chunk swap

src/paddlefleet/transformer/cp_utils.py

@@ -25,6 +25,7 @@
 from __future__ import annotations
 
 import paddle
+import paddle.distributed as dist
 from paddle import Tensor
 
 # ===========================================================================
@@ -164,3 +165,91 @@ def build_causal_mask_cp(
         paddle.zeros([1], dtype="float32"),
     )  # [sq_local, n_comp]
     return mask.unsqueeze(0).expand([batch_size, -1, -1])
+
+
+# ===========================================================================
+# Dual-chunk (zigzag) row swap — CP load balancing for the causal indexer
+# ===========================================================================
+
+
+def dualchunk_chunk_ids(cp_rank: int, cp_size: int) -> tuple[int, int]:
+    """The two ids, out of ``2 * cp_size`` equal chunks, this rank computes.
+
+    Contiguous CP hands rank ``r`` chunks ``(2r, 2r+1)``. This layout keeps
+    ``2r`` and takes ``2*cp_size-1-2r`` instead, so the two ids sum to
+    ``2*cp_size-1`` on **every** rank. A causal row's candidate count grows
+    linearly with its global position, so a constant id sum means constant
+    indexer work per rank: at cp16 the 31x spread between rank 0 and rank 15
+    collapses to 1x.
+
+    Keeping ``2r`` — rather than the more familiar ``(r, 2*cp_size-1-r)``
+    pairing, which balances just as well — is what reduces the exchange to a
+    single pairwise swap; see ``dualchunk_swap``.
+    """
+    lo = 2 * cp_rank
+    return lo, 2 * cp_size - 1 - lo
+
+
+def dualchunk_partner(cp_rank: int, cp_size: int) -> int:
+    """Rank holding the chunk this one wants, or ``-1`` when nothing moves.
+
+    Rank ``r`` wants chunk ``2*cp_size-1-2r``, which contiguous CP placed on
+    rank ``cp_size-1-r``; that rank symmetrically wants ``2r+1`` from here, so
+    the pairing is the involution ``r <-> cp_size-1-r``. Returns ``-1`` for a
+    single-rank group and for the self-paired middle rank of an odd group.
+    """
+    if cp_size <= 1:
+        return -1
+    partner = cp_size - 1 - cp_rank
+    return -1 if partner == cp_rank else partner
+
+
+def dualchunk_swap(x: Tensor, group, axis: int = 1) -> Tensor:
+    """Exchange the second half of ``x`` along ``axis`` with the partner rank.
+
+    In: ``x`` is this rank's contiguous CP shard, i.e. chunks ``(2r, 2r+1)``.
+    Out: chunks ``(2r, 2*cp_size-1-2r)`` — see ``dualchunk_chunk_ids``.
+
+    Both ends of a pair give up their odd chunk and want the other's, so this
+    is one symmetric pairwise exchange rather than a general all-to-all: each
+    rank moves ``x.shape[axis] // 2`` rows in each direction and talks to
+    exactly one peer.
+
+    The map is an **involution**, so calling this again undoes it — which is
+    why one function serves both directions.
+
+    Not differentiable, deliberately. The MQA indexer forward runs wholly
+    under ``paddle.no_grad()`` on detached inputs, and its gradient reaches the
+    weights through ``TileLangCSAIndexerLossAutoScaler`` applied to the
+    *unpermuted* tensors, so there is no gradient to route back through here.
+    """
+    if group is None or group.nranks <= 1:
+        return x
+    partner = dualchunk_partner(group.rank, group.nranks)
+    if partner < 0:
+        return x
+
+    n = int(x.shape[axis])
+    if n % 2 != 0:
+        raise ValueError(
+            "dualchunk_swap needs an even extent on the swapped axis, got "
+            f"{n} on axis {axis} of shape {list(x.shape)}"
+        )
+
+    keep, give = paddle.split(x, 2, axis=axis)
+    give = give.contiguous()
+    take = paddle.empty(give.shape, dtype=give.dtype)
+
+    # ``peer`` is a global rank: ``group.ranks`` maps the group-local index,
+    # matching context_parallel_utils.py:1252-1261.
+    peer = group.ranks[partner]
+    send_op = dist.P2POp(dist.isend, give, peer, group)
+    recv_op = dist.P2POp(dist.irecv, take, peer, group)
+    # Order by rank so a pair never issues two sends before either recv. NCCL's
+    # grouped p2p does not need this, but it costs one branch and removes a
+    # hang mode if the ops ever degrade to blocking.
+    ops = [send_op, recv_op] if group.rank < partner else [recv_op, send_op]
+    for task in dist.batch_isend_irecv(ops):
+        task.wait()
+
+    return paddle.concat([keep, take], axis=axis).contiguous()

src/paddlefleet/transformer/mqa_latent_attention.py

@@ -109,7 +109,11 @@
     preprocess_index,
 )
 from paddlefleet.process_groups_config import ProcessGroupCollection
-from paddlefleet.transformer.cp_utils import all_gather_cp
+from paddlefleet.transformer.cp_utils import (
+    all_gather_cp,
+    dualchunk_chunk_ids,
+    dualchunk_swap,
+)
 from paddlefleet.transformer.csa_attention import (
     TileLangCSAIndexerLossAutoScaler,
     _build_window_topk_idxs_from_doc_bounds,
@@ -755,6 +759,132 @@ def __init__(
         # Fused Triton epilogue for the analytic sink gradient. Only reachable
         # when ``softmax_offset`` exists; a sinkless layer has no sink gradient.
         self.sink_grad_fusion = getattr(config, "dsa_sink_grad_fusion", False)
+        # Row layout the indexer forward runs on; see ``mqa_indexer_cp_mode`` in
+        # ``transformer_config`` and ``_indexer_topk_dualchunk``. Gated on a real
+        # CP group: with ``cp_size == 1`` there is nothing to rebalance.
+        self.indexer_dualchunk = (
+            getattr(config, "mqa_indexer_cp_mode", None) == "dualchunk_p2p"
+            and self.cp_enabled
+            and self.cp_size > 1
+        )
+
+    def _chunk_valid_range(
+        self, meta, s_global, doc_start, doc_len, is_valid, offset, length
+    ):
+        """``valid_range [1, length, 2]`` for the rows at ``[offset, offset+length)``.
+
+        Both the cached ``MQADocMeta`` and the eager fallback expose the same
+        ``(offset, length)`` slice of one global table, which is what lets the
+        dual-chunk layout ask for its two segments by chunk offset instead of
+        needing a second table.
+        """
+        if meta is not None:
+            return meta.indexer_valid_range(self.window_size, offset, length)[0]
+        return self._indexer_valid_range(
+            s_global, doc_start, doc_len, is_valid, offset, length
+        )[0]
+
+    def _dualchunk_valid_range(
+        self, meta, s_global, doc_start, doc_len, is_valid, s
+    ):
+        """``valid_range [1, s, 2]`` for this rank's rows in dual-chunk order.
+
+        The same global table the contiguous path slices once, read instead as
+        the two chunks ``dualchunk_chunk_ids`` assigns, so it lines up
+        row-for-row with what ``dualchunk_swap`` produces and each kernel call
+        can be given its own ``seq_offset``.
+
+        ``row_empty`` deliberately has no counterpart here: it is applied to the
+        results after they come back, i.e. in contiguous order.
+        """
+        lo, hi = dualchunk_chunk_ids(self.cp_rank, self.cp_size)
+        m = s // 2
+        return paddle.concat(
+            [
+                self._chunk_valid_range(
+                    meta, s_global, doc_start, doc_len, is_valid, c * m, m
+                )
+                for c in (lo, hi)
+            ],
+            axis=1,
+        )
+
+    def _indexer_topk_dualchunk(
+        self, q_idx, w_idx, k_idx, topk, doc_lens_arg, vr_zz, need_loss
+    ):
+        """Indexer top-k on the balanced dual-chunk row layout.
+
+        Rank ``r`` scores global chunks ``(2r, 2*cp_size-1-2r)`` out of
+        ``2*cp_size`` instead of its contiguous ``(2r, 2r+1)``. The ids sum to
+        ``2*cp_size-1`` on every rank, and a causal row's candidate count grows
+        linearly with its global position, so the two halves' work sums to a
+        constant: measured 31x between cp0 and cp15 at 256k/cp16 collapses to 1x.
+
+        **Two calls, not one.** The kernel learns where its rows sit globally
+        from ``q_causal_offsets``, one scalar per batch
+        (``csa_indexer_fwd_cudnn.py``), so a single affine map cannot describe two
+        disjoint segments. Each chunk is internally contiguous, so each call
+        carries its own ``seq_offset``. This is the same shape as the query tiling
+        the dense path already does (it passes ``seq_offset + start`` per tile);
+        only the second segment's start jumps.
+
+        ``vr_zz`` must be the ``valid_range`` of exactly these rows, in the same
+        order. A ``seq_offset`` that disagrees with it is silently wrong in one
+        direction: too large only wastes work (the extra columns are masked out of
+        the top-k anyway), too small writes ``-inf`` over legal candidates so they
+        can never be selected.
+
+        The two results are concatenated and swapped back to the contiguous
+        layout before returning, so ``row_empty``, ``window_idxs``, attention, the
+        KL and ``TileLangCSAIndexerLossAutoScaler`` all keep seeing one unpermuted
+        ``[b, s_local, topk]`` tensor. Applying the loss PyLayer per chunk instead
+        would halve ``target.shape[1]``, which is where its backward reads the
+        row-count denominator, and double the gradient.
+
+        The swaps are not differentiable and do not need to be: this whole path
+        runs under ``paddle.no_grad()`` on detached inputs, and the indexer
+        gradient reaches the weights through the loss scaler applied to the
+        *unpermuted* tensors.
+        """
+        from paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn import (
+            cudnn_indexer_topk_fwd,
+        )
+
+        m = int(q_idx.shape[1]) // 2
+        lo, hi = dualchunk_chunk_ids(self.cp_rank, self.cp_size)
+
+        q_zz = dualchunk_swap(q_idx.detach(), self.cp_group, axis=1)
+        w_zz = dualchunk_swap(w_idx.detach(), self.cp_group, axis=1)
+
+        def _chunk(sl, chunk_id):
+            return cudnn_indexer_topk_fwd(
+                q_zz[:, sl].contiguous(),
+                k_idx.detach(),
+                w_zz[:, sl].contiguous(),
+                ratio=1,
+                topk_effective=topk,
+                valid_range=vr_zz[:, sl],
+                doc_lens=doc_lens_arg,
+                seq_offset=chunk_id * m,
+                return_topk_scores=need_loss,
+            )
+
+        r_lo = _chunk(slice(0, m), lo)
+        r_hi = _chunk(slice(m, 2 * m), hi)
+
+        selected = dualchunk_swap(
+            paddle.concat([r_lo[0], r_hi[0]], axis=1), self.cp_group, axis=1
+        )
+        scores_out = []
+        if need_loss:
+            scores_out = [
+                dualchunk_swap(
+                    paddle.concat([r_lo[2], r_hi[2]], axis=1),
+                    self.cp_group,
+                    axis=1,
+                )
+            ]
+        return selected, scores_out
 
     def _needs_indexer_loss(self) -> bool:
         """Whether this forward should build and attach the indexer loss.
@@ -1889,6 +2019,18 @@ def _forward_sparse(
                 valid_range, row_empty = self._indexer_valid_range(
                     s_global, doc_start, doc_len, is_valid, position_offset, s
                 )
+            # ``row_empty`` stays contiguous -- it is applied after the results
+            # come back. ``vr_zz`` is the same table read in dual-chunk row
+            # order, for the two kernel calls: both sources slice a cached
+            # global table by ``(offset, length)``, so asking twice with the
+            # chunk offsets is exactly the row set the swap produces.
+            vr_zz = (
+                self._dualchunk_valid_range(
+                    meta, s_global, doc_start, doc_len, is_valid, s
+                )
+                if self.indexer_dualchunk
+                else None
+            )
 
         q_idx, k_idx, w_idx = self._indexer_projections(
             x, qr, position_offset, grad_enabled=need_loss
@@ -1920,17 +2062,22 @@ def _forward_sparse(
         )
 
         with paddle.no_grad():
-            selected, _, *scores_out = cudnn_indexer_topk_fwd(
-                q_idx.detach(),
-                k_idx.detach(),
-                w_idx.detach(),
-                ratio=1,
-                topk_effective=topk,
-                valid_range=valid_range,
-                doc_lens=doc_lens_arg,
-                seq_offset=position_offset,
-                return_topk_scores=need_loss,
-            )
+            if self.indexer_dualchunk:
+                selected, scores_out = self._indexer_topk_dualchunk(
+                    q_idx, w_idx, k_idx, topk, doc_lens_arg, vr_zz, need_loss
+                )
+            else:
+                selected, _, *scores_out = cudnn_indexer_topk_fwd(
+                    q_idx.detach(),
+                    k_idx.detach(),
+                    w_idx.detach(),
+                    ratio=1,
+                    topk_effective=topk,
+                    valid_range=valid_range,
+                    doc_lens=doc_lens_arg,
+                    seq_offset=position_offset,
+                    return_topk_scores=need_loss,
+                )
             topk_indices = paddle.where(
                 row_empty, paddle.full_like(selected, -1), selected
             )