diff --git a/src/paddlefleet/transformer/cp_utils.py b/src/paddlefleet/transformer/cp_utils.py index 234106d56..2ddbe7ea4 100644 --- a/src/paddlefleet/transformer/cp_utils.py +++ b/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() diff --git a/src/paddlefleet/transformer/mqa_latent_attention.py b/src/paddlefleet/transformer/mqa_latent_attention.py index a74239a0b..3bad6d0fd 100644 --- a/src/paddlefleet/transformer/mqa_latent_attention.py +++ b/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 ) diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index 8e7869002..da518f9ed 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -998,6 +998,36 @@ class TransformerConfig(ModelParallelConfig): layout, which is what makes mixing them safe. """ + mqa_indexer_cp_mode: str | None = None + """Row layout the latent-MQA indexer's forward runs on, under context parallel. + + ``None`` (default) inherits ``cp_balance_mode``: the indexer scores this + rank's own contiguous row slice, so under a causal mask its cost grows with + the rank index. Measured at 256k/cp16: 2.2ms on cp0 vs 66.8ms on cp15 per + layer per pass, i.e. the slowest rank does 1.94x the average and every other + rank waits for it at the next collective. + + ``"dualchunk_p2p"`` splits the global sequence into ``2 * cp_size`` chunks + and has rank ``r`` score chunks ``(2r, 2*cp_size-1-2r)`` instead of its own + ``(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 work is + equal everywhere. Only the indexer's rows move -- attention is already + balanced at a fixed ``index_topk + window`` columns per row, so ``query`` and + the layer output never travel. Rank ``r`` keeps the chunk it already owns and + swaps the other with rank ``cp_size-1-r``, which reduces the exchange to a + single point-to-point sendrecv rather than an all-to-all. + + The global token layout is untouched: this is a layer-local permutation, + undone before the layer returns, so the HCA layers of the same model are + unaffected and ``cp_balance_mode`` must stay contiguous. + + Only the sparse training phase honours this, i.e. ``hybrid_mla_attention= + "mqa_dsa"`` with ``dsa_indexer_use_sparse_loss=True``. That is the only + phase whose indexer runs a top-k over per-rank rows; the warmup phase scores + the whole causal set through a different code path that does not permute + rows, so the other combinations are rejected rather than accepted-and-inert. + """ + v_head_dim: int | None = None """Dimension of the head in the V projection.""" @@ -2470,6 +2500,70 @@ def __post_init__(self): "-2 in csa_compress_ratios." ) + # only support mqa_indexer_cp_mode == dualchunk_p2p if not None + if self.mqa_indexer_cp_mode is not None: + if self.mqa_indexer_cp_mode != "dualchunk_p2p": + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} is " + "invalid. Must be None or 'dualchunk_p2p'." + ) + # The permutation is layer-local: the rows are swapped inside the + # MQA layer and swapped back before it returns, which is only sound + # while the *global* layout is the contiguous one the index tables + # are built against ("build over the global sequence, then take this + # rank's rows"). A dualchunk global layout would double-permute. + # + # Exact value rather than ``startswith("contiguous")``: under a real + # CP group ``MQALatentAttention.__init__`` accepts only + # ``contiguous_allgather``, so admitting ``contiguous_a2a`` here + # would pass config validation and then raise NotImplementedError at + # module construction -- a contract split between two layers. + if self.cp_balance_mode != "contiguous_allgather": + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} needs " + "cp_balance_mode='contiguous_allgather' (the only mode the " + "latent-MQA layer supports under context parallel), got " + f"{self.cp_balance_mode!r}." + ) + # Same membership test as hybrid_mla_cp_mode above: no latent-MQA + # layer means no indexer to rebalance. + if self.experimental_attention_variant != "dsv4_hybrid" or ( + -2 not in (self.csa_compress_ratios or ()) + ): + raise ValueError( + "mqa_indexer_cp_mode can only be set in dsv4_hybrid with " + "-2 in csa_compress_ratios." + ) + # ``MQALatentAttention`` reads the switch in ``_forward_sparse`` + # only. The other two phases reach a different indexer path that + # does not permute rows -- ``hybrid_mla_attention="mha"`` builds no + # latent-MQA layer at all, and ``"mqa_dsa"`` with + # ``dsa_indexer_use_sparse_loss=False`` is the warmup phase, whose + # KL spans the whole per-document causal set with no top-k anywhere. + # Accepting those would start successfully and silently deliver none + # of the rebalance this switch advertises. + if ( + self.hybrid_mla_attention != "mqa_dsa" + or not self.dsa_indexer_use_sparse_loss + ): + raise ValueError( + f"mqa_indexer_cp_mode={self.mqa_indexer_cp_mode!r} only " + "takes effect in the sparse training phase, which is " + "hybrid_mla_attention='mqa_dsa' together with " + "dsa_indexer_use_sparse_loss=True. This config has " + f"hybrid_mla_attention={self.hybrid_mla_attention!r} and " + "dsa_indexer_use_sparse_loss=" + f"{self.dsa_indexer_use_sparse_loss!r}, which runs an " + "indexer path that scores the full causal set and never " + "permutes rows, so the rebalance would be silently inert. " + "Drop mqa_indexer_cp_mode, or move to the sparse phase. " + "Note that train_indexer_only=True is the warmup phase by " + "definition and so cannot carry it either." + ) + # The two-chunks-per-rank split needs an even per-rank row count. + # ``TransformerConfig`` does not carry the sequence length, so that + # is checked at the swap itself (``cp_utils.dualchunk_swap``). + # separate_mtp_headloss validation. if self.separate_mtp_headloss: import warnings as _warnings diff --git a/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py b/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py new file mode 100644 index 000000000..c7a337c15 --- /dev/null +++ b/tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``mqa_indexer_cp_mode="dualchunk_p2p"`` must not change what the layer computes. + +The switch rebalances *which rows* each rank's indexer scores: rank ``r`` takes +global chunks ``(2r, 2*cp_size-1-2r)`` of ``2*cp_size`` instead of its own +``(2r, 2r+1)``, so the causal work is equal across ranks instead of growing 31x +from cp0 to cp15. The permutation is undone before the layer returns, so the +output must be unchanged. + +Two things make that testable as an equality rather than a tolerance: + +* the per-``(row, column)`` indexer score is bit-identical between the layouts. + A CTA covers ``q_stage * q_tokens_per_tile`` rows and both offsets are + multiples of that, so the same global row lands in the same aligned tile and + the kernel visits the same key blocks in the same order; +* the swap itself is pure data movement, and it is an involution, so the + round trip restores the contiguous order exactly. + +What is *not* guaranteed is the top-k tie-break: the two calls see different +``max_k`` on the THD path, so tied scores could in principle resolve +differently. Ties are vanishingly rare (a sum of ``index_n_heads`` fp32 +products), which is why this asserts equality and would flag a systematic +divergence rather than paper over it with a tolerance. +""" + +import types +import unittest + +import numpy as np +import paddle +import paddle.distributed as dist +from paddle.distributed import fleet + +from paddlefleet.transformer.cp_utils import ( + dualchunk_chunk_ids, + dualchunk_partner, + dualchunk_swap, +) +from tests.single_card_tests.transformer import hybrid_mla_utils as U + +CP_SIZE = None +CP_RANK = None +CP_GROUP = None + +S_GLOBAL = 512 + +_FA4_PIN = U._fa4_pin() + + +def setUpModule(): + global CP_SIZE, CP_RANK, CP_GROUP + if dist.get_world_size() < 2: + raise unittest.SkipTest( + "MQA dual-chunk tests require >= 2 GPUs (one peer to swap with)" + ) + _FA4_PIN.__enter__() + world = dist.get_world_size() + strategy = fleet.DistributedStrategy() + strategy.hybrid_configs = { + "dp_degree": 1, + "mp_degree": 1, + "pp_degree": 1, + "sharding_degree": world, + "sep_degree": 1, + "cp_degree": world, + "ep_degree": world, + "moe_sharding_degree": 1, + "order": [ + "sharding", + "moe_sharding", + "pp", + "sep", + "cp", + "dp", + "ep", + "mp", + ], + } + fleet.init(is_collective=True, strategy=strategy) + CP_GROUP = fleet.get_hybrid_communicate_group().get_context_parallel_group() + CP_RANK = CP_GROUP.rank + CP_SIZE = CP_GROUP.nranks + + +def tearDownModule(): + _FA4_PIN.__exit__(None, None, None) + + +def _build(dualchunk, loss_coeff=0.0, seed=7): + """CP layer with the rebalance off/on; identical weights either way.""" + cfg = U._create_mqa_config(mode="mqa_dsa", loss_coeff=loss_coeff) + cfg.dsa_indexer_use_sparse_loss = True + cfg.cp_balance_mode = "contiguous_allgather" + cfg.mqa_indexer_cp_mode = "dualchunk_p2p" if dualchunk else None + cfg.experimental_dataflow = True + cfg.pad_token_id = 0 + cfg.context_parallel_size = CP_SIZE + paddle.seed(seed) + return U._build_module( + cfg, + bf16=True, + pg_collection=types.SimpleNamespace(tp=None, cp=CP_GROUP), + ) + + +def _inputs(s_global, seed=1234): + """Global-length inputs, bitwise identical on every rank.""" + rng = np.random.RandomState(seed) + + def t(shape, scale): + return paddle.to_tensor( + (rng.standard_normal(shape) * scale).astype("float32") + ).cast("bfloat16") + + return { + "query": t([1, s_global, U.H, U.DK], 0.5), + "key": t([1, s_global, 1, U.DK], 0.5), + "w_v": t([U.DV, U.H, U.V_HEAD_DIM], 0.05), + "x": t([1, s_global, U.HIDDEN], 0.5), + "qr": t([1, s_global, U.Q_LORA], 0.5), + } + + +def _leaf(t): + out = t.clone().detach() + out.stop_gradient = False + return out + + +def _max_abs(a, b): + """``max |a - b|`` in fp32, so bf16 tensors compare without overflow.""" + return float((a.astype("float32") - b.astype("float32")).abs().max()) + + +def _rel(a, b): + """``max |a - b| / max |a|``, the shape the dkv tolerance is quoted in.""" + scale = float(a.astype("float32").abs().max()) + return _max_abs(a, b) / max(scale, 1e-12) + + +# The shared cuDNN DSA backward quotes rel ~2e-3 for ``dkv`` between identical +# calls; 5e-3 leaves headroom for the max-abs tail without hiding a real drift. +_GRAD_RTOL = 5e-3 + + +def _run(dualchunk, inp, row_end, loss_coeff): + """One forward+backward of this rank's slice. + + Returns ``(output, input_grads, param_grads)``. The parameter grads are the + only place the indexer loss shows up: ``_indexer_projections`` detaches + ``x``/``qr`` before the indexer touches them, so the KL gradient reaches the + indexer's own weights (``wq_b``, ``wk``, ``weights_proj``) and never the + caller's leaves. Asserting on the output alone would pass even if the + permutation corrupted ``topk_probs``, because the loss scaler forwards + ``output`` unchanged and only injects the gradient in backward. + """ + sl = S_GLOBAL // CP_SIZE + off = CP_RANK * sl + layer = _build(dualchunk, loss_coeff) + if dualchunk: + # Same weights on both sides: rebuilding from the reference state dict is + # unnecessary because ``paddle.seed`` is pinned in ``_build``, but assert + # the switch actually took so a silent no-op cannot pass this test. + assert layer.indexer_dualchunk, "dualchunk_p2p did not take effect" + else: + assert not layer.indexer_dualchunk + + r = { + "query": _leaf(inp["query"][:, off : off + sl]), + "key": _leaf(inp["key"][:, off : off + sl]), + "w_v": _leaf(inp["w_v"]), + "x": _leaf(inp["x"][:, off : off + sl]), + "qr": _leaf(inp["qr"][:, off : off + sl]), + } + out = layer( + r["query"], + r["key"], + None, + None, + attn_mask_startend_row_indices=row_end.clone(), + x=r["x"], + qr=r["qr"], + v_b_proj_weight=r["w_v"], + input_ids=None, + ) + out.sum().backward() + return ( + out, + {k: v.grad for k, v in r.items() if v.grad is not None}, + {n: p.grad for n, p in layer.named_parameters() if p.grad is not None}, + ) + + +class TestMqaIndexerDualChunkCp(unittest.TestCase): + def test_chunk_ids_balance_and_cover(self): + """Ids sum to a constant and every chunk has exactly one owner.""" + owner = {} + for r in range(CP_SIZE): + lo, hi = dualchunk_chunk_ids(r, CP_SIZE) + self.assertEqual(lo + hi, 2 * CP_SIZE - 1) + # The chunk this rank keeps is the one contiguous CP already gave + # it, which is what makes the exchange a single sendrecv. + self.assertEqual(lo // 2, r) + self.assertEqual(hi // 2, dualchunk_partner(r, CP_SIZE)) + for c in (lo, hi): + self.assertNotIn(c, owner) + owner[c] = r + self.assertEqual(sorted(owner), list(range(2 * CP_SIZE))) + + def test_swap_is_an_involution(self): + """Applying the swap twice restores the contiguous rows exactly.""" + sl = S_GLOBAL // CP_SIZE + base = CP_RANK * sl + rows = paddle.arange(base, base + sl, dtype="float32") + x = rows.reshape([1, sl, 1]).expand([1, sl, 4]).clone() + once = dualchunk_swap(x, CP_GROUP, axis=1) + twice = dualchunk_swap(once, CP_GROUP, axis=1) + self.assertEqual(float((twice - x).abs().max()), 0.0) + # And the single swap really did move the expected chunks. + m = sl // 2 + lo, hi = dualchunk_chunk_ids(CP_RANK, CP_SIZE) + want = paddle.concat( + [ + paddle.arange(lo * m, (lo + 1) * m, dtype="float32"), + paddle.arange(hi * m, (hi + 1) * m, dtype="float32"), + ] + ) + self.assertEqual(float((once[0, :, 0] - want).abs().max()), 0.0) + + def test_odd_extent_rejected(self): + """Two chunks per rank needs an even local row count.""" + odd = paddle.zeros([1, 2 * (S_GLOBAL // CP_SIZE) + 1, 2]) + with self.assertRaisesRegex(ValueError, "even extent"): + dualchunk_swap(odd, CP_GROUP, axis=1) + + @U._GPU + def test_output_unchanged(self): + """The rebalance is a pure scheduling change: identical output. + + The output being *exactly* equal is the whole contract, and it is the + strong assertion here: it can only hold if the two layouts selected the + same ``token_indices``, in the same order, from the same scores. + + Gradients then flow from identical inputs through identical code, so they + are only checked against the tolerance the shared cuDNN DSA backward + already carries: it accumulates ``dkv`` with atomics and is not + reproducible even between two identical ``off`` runs (documented as + rel ~2e-3 in ``fusions/mqa_sparse_attn.py``). Both the off-vs-off and the + off-vs-dualchunk pair are held to the same bound, which is what shows the + tolerance belongs to the kernel and not to this switch. + """ + inp = _inputs(S_GLOBAL) + row_end = U._row_end([S_GLOBAL], S_GLOBAL) + ref, gref, pref = _run(False, inp, row_end, loss_coeff=0.0) + ref2, gref2, pref2 = _run(False, inp, row_end, loss_coeff=0.0) + got, ggot, pgot = _run(True, inp, row_end, loss_coeff=0.0) + + self.assertEqual(ref.shape, got.shape) + self.assertEqual( + _max_abs(ref, got), 0.0, "dual-chunk changed the layer output" + ) + self._assert_grads_match( + (gref, gref2, ggot), (pref, pref2, pgot), "no-loss" + ) + + @U._GPU + def test_indexer_loss_gradients_unchanged(self): + """With the KL live, the indexer's own weight grads must also match. + + This is the case the output alone cannot cover: the loss scaler forwards + ``output`` unchanged and only injects the gradient in backward, so a + permutation that corrupted ``topk_probs`` or ``topk_indices`` would show + up nowhere else. ``_indexer_projections`` detaches ``x``/``qr``, so the KL + gradient lands on ``indexer.*`` parameters -- which is what gets compared + here, against the same off-vs-off baseline. + """ + inp = _inputs(S_GLOBAL, seed=99) + row_end = U._row_end([S_GLOBAL // 2, S_GLOBAL // 2], S_GLOBAL) + ref, gref, pref = _run(False, inp, row_end, loss_coeff=0.01) + ref2, gref2, pref2 = _run(False, inp, row_end, loss_coeff=0.01) + got, ggot, pgot = _run(True, inp, row_end, loss_coeff=0.01) + + # The KL has to actually reach the indexer, or every assertion below is + # comparing zeros and the test is vacuous. + indexer_grads = [ + n for n in pref if ".indexer." in n or n.startswith("indexer.") + ] + self.assertTrue(indexer_grads, "no indexer parameter received a grad") + self.assertTrue( + any( + float(pref[n].astype("float32").abs().max()) > 0.0 + for n in indexer_grads + ), + "indexer parameter grads are all zero: the KL never fired, so " + "this test would pass for any permutation", + ) + + self.assertEqual(_max_abs(ref, got), 0.0) + self._assert_grads_match( + (gref, gref2, ggot), (pref, pref2, pgot), "indexer-loss" + ) + + def _assert_grads_match(self, inputs, params, tag): + """Hold off-vs-off and off-vs-dualchunk to the same relative bound.""" + for kind, (a, a2, b) in (("input", inputs), ("param", params)): + self.assertEqual(sorted(a), sorted(b), f"{tag}: {kind} grad keys") + for name in a: + for label, other in (("off-vs-off", a2), ("dualchunk", b)): + rel = _rel(a[name], other[name]) + self.assertLessEqual( + rel, + _GRAD_RTOL, + f"{tag} {kind} grad of {name} ({label}): rel " + f"{rel:.2e} exceeds {_GRAD_RTOL:.0e}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py b/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py index 48e36ffc5..79281ca0a 100644 --- a/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py +++ b/tests/single_card_tests/ai_edited_test/distributed/test_cp_balance_mode_dispatch.py @@ -880,5 +880,157 @@ def test_requires_dsv4_hybrid_with_mla_layer(self): ) +class TestTransformerConfigMqaIndexerCpMode(unittest.TestCase): + """Tests for the mqa_indexer_cp_mode row-rebalance switch.""" + + # Same context hybrid_mla_cp_mode needs: a dsv4_hybrid whose + # csa_compress_ratios declares a latent-MQA layer (the -2), which is the + # only layer carrying an indexer to rebalance. + MQA_HYBRID = { + "experimental_attention_variant": "dsv4_hybrid", + "num_hidden_layers": 2, + "csa_compress_ratios": [-2, 128], + "hybrid_mla_q_lora_rank": 8, + "hybrid_mla_kv_lora_rank": 8, + "hybrid_mla_qk_nope_head_dim": 8, + "hybrid_mla_qk_rope_head_dim": 8, + "hybrid_mla_v_head_dim": 8, + "hybrid_mla_num_attention_heads": 4, + "hybrid_mla_num_key_value_heads": 4, + # The indexer shapes ``hybrid_mla_attention="mqa_dsa"`` requires in both + # DSA phases. ``index_head_dim`` is pinned to 128 by the cuDNN indexer + # and ``index_topk`` must be a multiple of 128. + "dsa_index_n_heads": 4, + "dsa_index_head_dim": 128, + "dsa_index_topk": 128, + } + + # The only phase that reads the switch: ``MQALatentAttention`` consults + # ``indexer_dualchunk`` in ``_forward_sparse`` and nowhere else, so the + # warmup phase would accept the config and rebalance nothing. + SPARSE_PHASE = { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": True, + } + + @staticmethod + def _config(**kwargs): + from paddlefleet.transformer.transformer_config import ( + TransformerConfig, + ) + + return TransformerConfig(**kwargs) + + def test_default_is_none(self): + """Unset means 'score this rank's own contiguous rows'.""" + self.assertIsNone(self._config().mqa_indexer_cp_mode) + + def test_dualchunk_accepted(self): + """The rebalance is layer-local, so it rides on a contiguous global.""" + config = self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + self.assertEqual(config.mqa_indexer_cp_mode, "dualchunk_p2p") + + def test_invalid_value_rejected(self): + """One mode only; a typo must not silently fall back to the default.""" + for mode in ("nonexistent_mode", "dualchunk_allgather"): + with ( + self.subTest(mode=mode), + self.assertRaisesRegex(ValueError, "mqa_indexer_cp_mode"), + ): + self._config( + mqa_indexer_cp_mode=mode, + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_non_contiguous_global_layout_rejected(self): + """A dualchunk global layout would permute the same rows twice.""" + with self.assertRaisesRegex(ValueError, "contiguous_allgather"): + self._config( + cp_balance_mode="dualchunk_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_contiguous_a2a_rejected(self): + """Rejected here rather than later: the latent-MQA layer accepts only + contiguous_allgather under CP, so admitting the sibling contiguous mode + would move the failure to module construction.""" + with self.assertRaisesRegex(ValueError, "contiguous_allgather"): + self._config( + cp_balance_mode="contiguous_a2a", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **self.SPARSE_PHASE, + ) + + def test_requires_dsv4_hybrid_with_mqa_layer(self): + """No latent-MQA layer means no indexer, so the setting is dead.""" + for kwargs in ( + {}, + {**self.MQA_HYBRID, "csa_compress_ratios": [128, 0]}, + ): + with ( + self.subTest(ratios=kwargs.get("csa_compress_ratios")), + self.assertRaisesRegex( + ValueError, "only be set in dsv4_hybrid" + ), + ): + self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **kwargs, + ) + + def test_requires_sparse_phase(self): + """Only ``_forward_sparse`` reads the switch, so nothing else may set it. + + This is the accepted-but-inert combination the layer cannot catch: + ``hybrid_mla_attention`` defaulting to ``"mha"`` builds no latent-MQA + layer, ``"mqa_full_causal"`` has no indexer, and ``"mqa_dsa"`` with + ``dsa_indexer_use_sparse_loss=False`` (which is what + ``train_indexer_only=True`` pins) runs the warmup indexer path, whose KL + spans the whole causal set and never permutes rows. All three would + start successfully and rebalance nothing. + """ + for label, phase in ( + ("default mha", {}), + ("no indexer", {"hybrid_mla_attention": "mqa_full_causal"}), + ( + "warmup", + { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": False, + }, + ), + ( + "train_indexer_only warmup", + { + "hybrid_mla_attention": "mqa_dsa", + "dsa_indexer_use_sparse_loss": False, + "train_indexer_only": True, + }, + ), + ): + with ( + self.subTest(phase=label), + self.assertRaisesRegex( + ValueError, "takes effect in the sparse training phase" + ), + ): + self._config( + cp_balance_mode="contiguous_allgather", + mqa_indexer_cp_mode="dualchunk_p2p", + **self.MQA_HYBRID, + **phase, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py b/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py new file mode 100644 index 000000000..a8c5cd8c9 --- /dev/null +++ b/tests/single_card_tests/ai_edited_test/distributed/test_dualchunk_cp_utils.py @@ -0,0 +1,375 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single-card tests for the dual-chunk CP row helpers in ``cp_utils``. + +``tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py`` covers +the same helpers against real NCCL. This file pins the parts that are pure +arithmetic or pure routing, so the chunk assignment and the peer/op layout are +checked on every single-card run and a regression there does not need two GPUs +to surface. + +The ``MQALatentAttention`` row-plumbing helpers are covered here too, with the +cuDNN indexer stubbed. The CP test exercises them for real but is gated on SM100 +kernels, so on every other box the two ``seq_offset`` values, the ``valid_range`` +slicing and the swap-back would otherwise go unchecked. +""" + +import types +import unittest +from unittest import mock + +import paddle + +from paddlefleet.transformer.cp_utils import ( + dualchunk_chunk_ids, + dualchunk_partner, + dualchunk_swap, +) + + +def _make_mock_group(nranks=4, rank=1): + """Mock process group; ``ranks`` maps group-local index to global rank.""" + group = mock.MagicMock() + group.nranks = nranks + group.rank = rank + # Offset so a test cannot pass by confusing the two rank spaces. + group.ranks = [100 + i for i in range(nranks)] + return group + + +class TestDualChunkChunkIds(unittest.TestCase): + def test_ids_partition_and_balance(self): + """Every chunk has one owner and the per-rank id sum is constant.""" + for cp_size in (1, 2, 4, 8, 16): + with self.subTest(cp_size=cp_size): + owner = {} + for rank in range(cp_size): + lo, hi = dualchunk_chunk_ids(rank, cp_size) + # Constant sum is the whole point: a causal row's candidate + # count grows linearly with its global position, so equal + # id sums mean equal work. + self.assertEqual(lo + hi, 2 * cp_size - 1) + # The kept chunk is the one contiguous CP already placed + # here, which is what makes the exchange one sendrecv. + self.assertEqual(lo, 2 * rank) + for c in (lo, hi): + self.assertNotIn(c, owner) + owner[c] = rank + self.assertEqual(sorted(owner), list(range(2 * cp_size))) + + +class TestDualChunkPartner(unittest.TestCase): + def test_involution_on_even_groups(self): + """``partner(partner(r)) == r``, so one function serves both ways.""" + for cp_size in (2, 4, 16): + for rank in range(cp_size): + p = dualchunk_partner(rank, cp_size) + self.assertEqual(dualchunk_partner(p, cp_size), rank) + self.assertEqual(p, cp_size - 1 - rank) + + def test_no_partner_cases(self): + """``-1`` means 'nothing to swap', not 'rank 0'.""" + self.assertEqual(dualchunk_partner(0, 1), -1) + # Odd group: the middle rank is its own partner, so it keeps both + # chunks rather than sending to itself. + self.assertEqual(dualchunk_partner(1, 3), -1) + self.assertEqual(dualchunk_partner(0, 3), 2) + + +class TestDualChunkSwap(unittest.TestCase): + def test_no_op_without_a_peer(self): + """Degenerate groups return the input untouched, not a copy-with-comm.""" + x = paddle.arange(8).reshape([1, 4, 2]).cast("float32") + self.assertIs(dualchunk_swap(x, None, axis=1), x) + self.assertIs(dualchunk_swap(x, _make_mock_group(1, 0), axis=1), x) + # Odd group's middle rank has no partner. + self.assertIs(dualchunk_swap(x, _make_mock_group(3, 1), axis=1), x) + + def test_odd_extent_rejected(self): + """Two chunks per rank needs an even count on the swapped axis.""" + x = paddle.zeros([1, 5, 2]) + with self.assertRaisesRegex(ValueError, "even extent"): + dualchunk_swap(x, _make_mock_group(4, 1), axis=1) + + def test_routing_and_op_order(self): + """The peer, the halves and the isend/irecv order, without NCCL. + + ``batch_isend_irecv`` is stubbed: this pins the routing contract (send + the odd chunk to ``cp_size-1-rank`` as a *global* rank, keep the even + one, and order the ops by rank) which is what a wrong peer or a + deadlock-prone ordering would break. + """ + cp_size = 4 + x = paddle.arange(16).reshape([1, 8, 2]).cast("float32") + for rank in (1, 2): + with self.subTest(rank=rank): + group = _make_mock_group(cp_size, rank) + captured = [] + + def _p2p_op(op, tensor, peer, grp): + captured.append((op, tensor, peer, grp)) + return (op, tensor, peer) + + with ( + mock.patch("paddle.distributed.P2POp", side_effect=_p2p_op), + mock.patch( + "paddle.distributed.batch_isend_irecv", + return_value=[mock.MagicMock()], + ) as batched, + ): + out = dualchunk_swap(x, group, axis=1) + + batched.assert_called_once() + (ops,) = batched.call_args[0] + self.assertEqual(len(ops), 2) + self.assertEqual(len(captured), 2) + + partner = dualchunk_partner(rank, cp_size) + peers = {c[2] for c in captured} + self.assertEqual( + peers, + {group.ranks[partner]}, + "peer must be the partner's *global* rank", + ) + # Lower rank sends first: harmless for NCCL's grouped p2p, but + # it removes a hang mode if the ops ever degrade to blocking. + # The submitted *list* order is what matters, not the order the + # two P2POp objects happened to be constructed in. + self.assertEqual( + ops[0][0] is paddle.distributed.isend, + rank < partner, + "ops must be submitted send-first only on the lower rank", + ) + + # The sent buffer is the second half; the first half survives + # untouched in the result. + sent = next( + t + for op, t, _, _ in captured + if op is paddle.distributed.isend + ) + self.assertEqual( + float((sent - x[:, 4:]).abs().max()), + 0.0, + "the odd chunk is the one that travels", + ) + self.assertEqual(out.shape, x.shape) + self.assertEqual( + float((out[:, :4] - x[:, :4]).abs().max()), + 0.0, + "the kept chunk must not move", + ) + + +def _half_flip(x, group, axis=1): + """Single-rank stand-in for ``dualchunk_swap``. + + An involution that leaves the first half alone and visibly reorders the + second, which is the only property the row plumbing relies on. A real swap + needs a peer rank; this one makes the swap-out/swap-back symmetry observable + on one card. + """ + m = int(x.shape[axis]) // 2 + keep, give = paddle.split(x, 2, axis=axis) + return paddle.concat([keep, paddle.flip(give, axis)], axis=axis) + + +class TestDualChunkValidRange(unittest.TestCase): + """``MQALatentAttention._dualchunk_valid_range`` reads two chunk offsets.""" + + def test_reads_the_two_chunk_offsets_in_order(self): + """Rows come from ``lo`` then ``hi``, each a length-``m`` global slice. + + The two ``(offset, length)`` pairs are what must line up with the two + ``seq_offset`` values the kernel calls get; asking with the local offset, + or in the wrong order, is silently wrong rather than a crash. + """ + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + asked = [] + + def _chunk( + meta, s_global, doc_start, doc_len, is_valid, offset, length + ): + asked.append((offset, length)) + return paddle.full([1, length, 2], float(offset)) + + fake = types.SimpleNamespace( + cp_rank=1, cp_size=4, _chunk_valid_range=_chunk + ) + out = MQALatentAttention._dualchunk_valid_range( + fake, "meta", 32, "doc_start", "doc_len", "is_valid", 8 + ) + + # cp_rank=1 of 4 owns global chunks (2, 5); m = s // 2 = 4. + self.assertEqual(asked, [(8, 4), (20, 4)]) + self.assertEqual(out.shape, [1, 8, 2]) + self.assertEqual(float(out[0, 0, 0]), 8.0) + self.assertEqual(float(out[0, 4, 0]), 20.0) + + +class TestIndexerTopkDualChunk(unittest.TestCase): + """``MQALatentAttention._indexer_topk_dualchunk`` with the kernel stubbed. + + The real path is covered by ``test_mqa_indexer_dualchunk_cp.py``, but its + value tests need SM100 kernels. What is checked here is the plumbing that + would be wrong on any box: two calls with the two global ``seq_offset`` + values, each fed its own half of the *already dual-chunk-ordered* + ``valid_range``, and the results swapped back to contiguous rows. + """ + + CP_RANK, CP_SIZE, S, TOPK = 1, 4, 8, 4 + + def _run(self, need_loss): + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + calls = [] + + def _kernel(q, k, w, **kw): + calls.append((q, w, kw)) + m = int(q.shape[1]) + # Rows tagged by their global position so the concat-then-swap-back + # is checkable rather than symmetric-by-accident. + rows = ( + kw["seq_offset"] + paddle.arange(m).cast("float32") + ).reshape([1, m, 1]) + out = paddle.expand(rows, [1, m, self.TOPK]) + return (out, None, out) if kw["return_topk_scores"] else (out, None) + + q = paddle.arange(self.S).reshape([1, self.S, 1, 1]).cast("float32") + w = paddle.arange(self.S).reshape([1, self.S, 1]).cast("float32") + vr = paddle.arange(2 * self.S).reshape([1, self.S, 2]).cast("float32") + fake = types.SimpleNamespace( + cp_rank=self.CP_RANK, cp_size=self.CP_SIZE, cp_group="grp" + ) + + with ( + mock.patch( + "paddlefleet.transformer.mqa_latent_attention.dualchunk_swap", + side_effect=_half_flip, + ) as swap, + mock.patch( + "paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn" + ".cudnn_indexer_topk_fwd", + side_effect=_kernel, + ), + ): + selected, scores_out = MQALatentAttention._indexer_topk_dualchunk( + fake, + q, + w, + paddle.zeros([1, self.S, 1]), + self.TOPK, + "doc_lens", + vr, + need_loss, + ) + return calls, swap, selected, scores_out, q, w, vr + + def test_two_calls_carry_the_chunk_offsets_and_their_own_rows(self): + calls, swap, selected, scores_out, q, w, vr = self._run(False) + + self.assertEqual(len(calls), 2) + # chunk ids (2, 5) with m = 4. + self.assertEqual([c[2]["seq_offset"] for c in calls], [8, 20]) + + q_zz, w_zz = _half_flip(q, "grp"), _half_flip(w, "grp") + for i, sl in enumerate((slice(0, 4), slice(4, 8))): + q_seen, w_seen, kw = calls[i] + self.assertEqual(float((q_seen - q_zz[:, sl]).abs().max()), 0.0) + self.assertEqual(float((w_seen - w_zz[:, sl]).abs().max()), 0.0) + # ``vr_zz`` is built in dual-chunk order already, so it is sliced, + # never swapped -- swapping it here would double-permute the rows. + self.assertEqual( + float((kw["valid_range"] - vr[:, sl]).abs().max()), + 0.0, + "valid_range must be sliced, not swapped again", + ) + self.assertFalse(kw["return_topk_scores"]) + + # Swapped back: rows 8..11 stay, 20..23 come back reversed by the stub. + self.assertEqual( + [float(v) for v in selected[0, :, 0]], + [8, 9, 10, 11, 23, 22, 21, 20], + ) + # q, w out; selected back. Nothing else travels without the loss. + self.assertEqual(swap.call_count, 3) + self.assertEqual(scores_out, []) + + def test_scores_are_swapped_back_only_when_the_loss_needs_them(self): + calls, swap, selected, scores_out, *_ = self._run(True) + + self.assertTrue(all(c[2]["return_topk_scores"] for c in calls)) + self.assertEqual(swap.call_count, 4) + (scores,) = scores_out + # Same layout as ``selected``: the KL sees unpermuted contiguous rows. + self.assertEqual( + [float(v) for v in scores[0, :, 0]], + [8, 9, 10, 11, 23, 22, 21, 20], + ) + + +class TestChunkValidRange(unittest.TestCase): + """``MQALatentAttention._chunk_valid_range`` picks one of two row sources. + + Both 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. The + delegation is what can silently break -- passing the caller's ``window`` + instead of ``self.window_size``, or forgetting to unwrap the + ``(valid_range, row_empty)`` pair -- so pin both paths here rather than only + exercising the eager one from the CP test. + """ + + @staticmethod + def _call(fake_self, meta, offset, length): + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + return MQALatentAttention._chunk_valid_range( + fake_self, + meta, + 512, + "doc_start", + "doc_len", + "is_valid", + offset, + length, + ) + + def test_uses_meta_when_present(self): + meta = mock.MagicMock() + meta.indexer_valid_range.return_value = ("vr", "row_empty") + fake = types.SimpleNamespace(window_size=128) + self.assertEqual(self._call(fake, meta, 64, 32), "vr") + meta.indexer_valid_range.assert_called_once_with(128, 64, 32) + + def test_falls_back_to_eager_build(self): + eager = mock.MagicMock(return_value=("vr2", "row_empty")) + fake = types.SimpleNamespace( + window_size=128, _indexer_valid_range=eager + ) + self.assertEqual(self._call(fake, None, 64, 32), "vr2") + eager.assert_called_once_with( + 512, "doc_start", "doc_len", "is_valid", 64, 32 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_configs.yaml b/tests/test_configs.yaml index 0dccc5438..4b77c3bbb 100644 --- a/tests/test_configs.yaml +++ b/tests/test_configs.yaml @@ -154,6 +154,10 @@ tests: - test_case: [tests/multi_card_tests/transformer/test_flash_mask_cp_a2a.py] products: - num_gpus: 2 + # transformer: latent-MQA indexer dual-chunk row rebalance (CP=2) + - test_case: [tests/multi_card_tests/transformer/test_mqa_indexer_dualchunk_cp.py] + products: + - num_gpus: 2 # transformer: CSA context-parallel, including cuDNN indexer docmask path (CP=2) - test_case: [tests/multi_card_tests/transformer/test_csa_attention_cp.py] products: