From e6544ef9ac8077b61b58d36c41235ef3bfdf76b7 Mon Sep 17 00:00:00 2001 From: shenliang03 Date: Wed, 12 Aug 2026 15:39:26 +0800 Subject: [PATCH 1/3] [release/0.4][Improvements] Run the DSA warmup phase on dense MHA instead of latent MQA The hybrid-MLA programme has four phases on a ``csa_compress_ratios == -2`` layer: (1) dense MHA pretrain, (2) MHA -> MQA equivalence, (3) DSA warmup that trains the indexer only, (4) sparse MQA + DSA. Phase 3 is *dense* by construction -- the indexer is still learning, so the main attention must not consume its ranking -- yet it ran on ``MQALatentAttention``, i.e. it drove dense semantics through the block-sparse kernel by handing it a full per-document causal index table. That table is O(s^2) per ``-2`` layer and buys nothing: measured peak memory for one core attention (forward + backward) is 8.2x / 8.5x / 9.0x the dense path at s = 1024 / 2048 / 4096, and the ratio grows with s. The absorbed layout also materialises ``q_absorbed`` (``[b, s, h, 576]``). This routes the warmup phase through a dense per-head attention instead, so its memory profile is exactly phase 1's: * ``MHADSAWarmupAttention`` (new) subclasses ``DotProductAttention`` and delegates the whole attention half to ``super().forward``, adding only the indexer projections and the KL loss. Phase 2 is therefore "phase 1 plus an indexer loss" by construction rather than by assertion, and the q/k it produces are bit-identical to phase 1's (verified with ``apply_rope_fusion`` both off and on). * ``hybrid_mla_indexer.py`` (new) holds the one dispatch predicate, ``latent_mqa_enabled(config)``, plus the indexer plumbing both backends share. ``gpt_layer_specs`` and ``MultiLatentAttention`` are the only callers, so no code path can pick a backend the model would not build. ``dsa_indexer_use_sparse_loss`` now selects the backend class, not just the KL width, and must be passed as a construction kwarg. * ``MQALatentAttention`` keeps only the sparse phases; the full-causal index table construction it needed for the warmup is gone (-434 lines net there). Parameter names, ``state_dict`` keys and saved HF keys are unchanged in every phase (16 keys; phase 1 -> warmup adds the 5 ``indexer.*`` keys and nothing else), so checkpoints stay loadable across the switch. Two fixes fell out of running the new backend: * ``TileLangCSAIndexerLossAutoScaler`` returned its first argument unchanged, and Paddle records a PyLayer returning one of its inputs as an inplace write on it. The dense attention backward saves its own output, so the version bump made the *attention* backward raise ``PermissionDenied: Tensor ... modified by an inplace operation``. The caller now hands the scaler a fresh tensor (``clone`` is a gradient identity). This was invisible before because the scaler already clones when the backbone is frozen. * Three context-parallel cases asserted a bit-identical ``dq`` against phase 1. Dense flashmask accumulates ``dq`` atomically over column blocks, so two runs of the *same* module differ by 3.05e-05 at s = 512 while forward, ``dk`` and ``dv`` are exact. The bound is now self-calibrated from a second reference run measured in the same test run. Rebased on #1679, which adds ``mqa_split_kv_b_proj``. That switch replaces ``kv_b_proj`` with standalone ``k_b_proj`` / ``v_b_proj``, so it only means anything where absorption happens. The warmup phase is now one of the dense phases and keeps ``kv_b_proj``, so its validation was moved onto ``latent_mqa_enabled`` too: the combination is rejected at config time instead of silently changing the parameter set at the warmup -> sparse switch. Tested on SM100 after the rebase: 431 single-card tests over 14 files (config pipeline 24, doc equivalence 35, dsv4 hybrid 87, grad health 22, HF roundtrip 12, latent MQA 63, Muon 67, warmup RoPE/recompute/MTP 14, MLA RoPE CP + VHA 56, warmup doc-mask loss + train_indexer_only 51) and 30 two-card CP tests (test_mqa_dsa_cp 8, test_mqa_dsa_warmup_cp 8, test_indexer_topk_col_mask_cp 4, test_mla_cp_contiguous_allgather 10). The one failure, ``test_documented_bug_gate_proj_is_saved_untransposed``, predates this branch: it pins a defect that has since been fixed outside this repository, and it fails identically on release/0.4. --- src/paddlefleet/models/gpt/gpt_layer_specs.py | 53 +- .../transformer/hybrid_mla_indexer.py | 224 +++ .../transformer/mha_dsa_warmup_attention.py | 475 ++++++ .../transformer/mqa_latent_attention.py | 434 +---- .../transformer/multi_latent_attention.py | 39 +- .../transformer/transformer_config.py | 95 +- .../test_indexer_topk_col_mask_cp.py | 78 +- .../transformer/test_mqa_dsa_cp.py | 75 +- .../transformer/test_mqa_dsa_warmup_cp.py | 960 ++++++++--- .../transformer/hybrid_mla_utils.py | 212 ++- .../transformer/test_dsv4_hybrid_attention.py | 191 ++- .../test_hybrid_mla_config_pipeline.py | 143 +- .../test_hybrid_mla_doc_equivalence.py | 5 +- .../test_hybrid_mla_grad_health.py | 212 ++- .../test_hybrid_mla_hf_roundtrip.py | 94 ++ .../test_hybrid_mla_warmup_doc_mask_loss.py | 1429 ++++++++++++----- ...st_hybrid_mla_warmup_recompute_mtp_rope.py | 751 ++++++--- .../transformer/test_mqa_latent_attention.py | 668 ++++++-- .../test_muon_hybrid_mla_grouping.py | 23 +- .../transformer/test_train_indexer_only.py | 86 +- .../transformer/test_vha_dsv4.py | 15 + 21 files changed, 4521 insertions(+), 1741 deletions(-) create mode 100644 src/paddlefleet/transformer/hybrid_mla_indexer.py create mode 100644 src/paddlefleet/transformer/mha_dsa_warmup_attention.py diff --git a/src/paddlefleet/models/gpt/gpt_layer_specs.py b/src/paddlefleet/models/gpt/gpt_layer_specs.py index 3b9f4dc148..d0506a1be1 100644 --- a/src/paddlefleet/models/gpt/gpt_layer_specs.py +++ b/src/paddlefleet/models/gpt/gpt_layer_specs.py @@ -82,7 +82,12 @@ GatedDeltaNet, GatedDeltaNetSublayersSpec, ) +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled from paddlefleet.transformer.identity_op import IdentityOp +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, + MHADSAWarmupAttentionSublayersSpec, +) from paddlefleet.transformer.mlp import MLP, MLPSublayersSpec from paddlefleet.transformer.mqa_latent_attention import ( MQALatentAttention, @@ -115,6 +120,26 @@ LNImpl = WrappedPaddleNorm +def _hybrid_mla_indexer_spec(backend) -> LayerSpec: + """``DSAIndexer`` spec for the ``csa_compress_ratios == -2`` layers. + + Shared by the two DSA core attentions of a hybrid MLA run -- phase 2's + ``MHADSAWarmupAttention`` and phase 3's ``MQALatentAttention`` -- so that the + indexer parameter names and shapes cannot drift between the phases an HF + checkpoint moves through. + """ + return LayerSpec( + layer=DSAIndexer, + sublayers_spec=DSAIndexerSublayersSpec( + linear_wq_b=backend.linear(), + linear_wk=backend.linear(), + k_norm=paddle.nn.LayerNorm, + linear_weights_proj=backend.linear(), + ), + extra_kwargs={"is_hybrid_mla_indexer": True}, + ) + + def _get_effective_mtp_layers(config: TransformerConfig) -> int: mtp_num_layers = getattr(config, "mtp_num_layers", 0) or 0 nextn_num_layers = getattr(config, "num_nextn_predict_layers", 0) or 0 @@ -340,7 +365,7 @@ def get_attention_spec( and getattr(config, "dsa_index_n_heads", None) is not None ) - if hybrid_mla_attention in ("mqa_dsa", "mqa_full_causal"): + if latent_mqa_enabled(config): # Latent MQA core attention on the KV latent; parameters stay # byte-identical to MHA so an MHA checkpoint loads unchanged. The # DSA indexer is what makes this mode worth running, so it is only @@ -352,21 +377,23 @@ def get_attention_spec( layer=MQALatentAttention, sublayers_spec=MQALatentAttentionSublayersSpec( indexer=( - None - if dense_mqa - else LayerSpec( - layer=DSAIndexer, - sublayers_spec=DSAIndexerSublayersSpec( - linear_wq_b=backend.linear(), - linear_wk=backend.linear(), - k_norm=paddle.nn.LayerNorm, - linear_weights_proj=backend.linear(), - ), - extra_kwargs={"is_hybrid_mla_indexer": True}, - ) + None if dense_mqa else _hybrid_mla_indexer_spec(backend) ), ), ) + elif hybrid_mla_attention == "mqa_dsa": + # Phase 2 (DSA warmup). No top-k on either side, so the attention is + # phase 1's dense MHA verbatim: absorbing into latent MQA here would + # only push a zero-sparsity ``[b, s, s]`` index table through the + # block-sparse kernel. ``latent_mqa_enabled`` above keeps this in + # step with the enclosing ``MLASelfAttention``, which reads the same + # predicate to decide whether to absorb at all. + core_attention = LayerSpec( + layer=MHADSAWarmupAttention, + sublayers_spec=MHADSAWarmupAttentionSublayersSpec( + indexer=_hybrid_mla_indexer_spec(backend), + ), + ) elif use_dsa: # DSA Indexer sublayers spec (duplicated linear, NOT tensor-parallel) dsa_indexer_sublayers = DSAIndexerSublayersSpec( diff --git a/src/paddlefleet/transformer/hybrid_mla_indexer.py b/src/paddlefleet/transformer/hybrid_mla_indexer.py new file mode 100644 index 0000000000..d9257b6ce4 --- /dev/null +++ b/src/paddlefleet/transformer/hybrid_mla_indexer.py @@ -0,0 +1,224 @@ +# 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. + +"""Shared DSA-indexer plumbing for the hybrid MLA (``csa_compress_ratios == -2``) +layers, plus the single predicate that decides whether those layers run the +absorbed latent MQA or the dense MHA of the pretraining phase. + +Two core attentions own a ``DSAIndexer`` on those layers and they run different +*attention* backends, so the indexer-side pieces they do share live here rather +than in either of them: + +* ``MHADSAWarmupAttention`` (phase 2, ``mha_dsa_warmup_attention.py``) -- dense + flashmask attention, exactly the pretraining phase, with a full-candidate KL. +* ``MQALatentAttention`` (phase 3/4, ``mqa_latent_attention.py``) -- block-sparse + attention on the KV latent, with the KL restricted to the selected set. +""" + +from __future__ import annotations + +import paddle +from paddle import Tensor + +from paddlefleet.context_parallel_utils import ContextParallelGatherOp + + +def latent_mqa_enabled(config) -> bool: + """Whether the hybrid MLA layers run *absorbed latent MQA* (not dense MHA). + + The single judgement behind both the spec dispatch + (``gpt_layer_specs.py``) and ``MLASelfAttention.mqa_latent``, which decides + whether ``kv_b_proj`` materialises per-head K/V at all. Splitting it would + let the spec build a dense core attention while the enclosing MLA layer + feeds it absorbed activations. + + * ``"mqa_full_causal"`` -- latent MQA with no indexer (equivalence isolation). + * ``"mqa_dsa"`` -- latent MQA **only in the sparse phase**. The warmup phase + (``dsa_indexer_use_sparse_loss=False``) has no top-k on either side, so + absorption would buy nothing and cost a full ``[b, s, s]`` index table fed + to a block-sparse kernel at zero sparsity; it runs the dense MHA of phase 1 + instead, with the indexer bolted on. + * anything else (``"mha"``, non-``dsv4_hybrid`` models) -- dense MHA. + """ + if getattr(config, "experimental_attention_variant", None) != "dsv4_hybrid": + return False + mode = getattr(config, "hybrid_mla_attention", "mha") + if mode == "mqa_full_causal": + return True + if mode == "mqa_dsa": + return bool(getattr(config, "dsa_indexer_use_sparse_loss", False)) + return False + + +class HybridMLAIndexerMixin: + """Indexer-side helpers shared by the two DSA core attentions. + + Expects the host layer to have set ``self.config``, ``self.indexer`` and the + CP state of :meth:`_init_hybrid_mla_cp_state`. + """ + + # Read by ``MLASelfAttention`` to decide whether to forward ``input_ids`` + # (needed for the indexer-loss row mask, since the packed sequence's trailing + # padding is invisible to ``attn_mask_startend_row_indices``). A capability + # of the core attention rather than a property of the phase, so no core + # attention can be handed a kwarg it does not accept. + accepts_input_ids = True + + def _init_hybrid_mla_cp_state(self, config, pg_collection) -> None: + """Set ``cp_group`` / ``cp_size`` / ``cp_rank`` / ``cp_enabled``. + + Same derivation as ``csa_attention.py:2079-2090``. Deliberately asserts + the *same* ``contiguous_allgather`` constraint the HCA layers of this + model assert (``dsv4_hybrid_attention.py:607-611``), not a weaker one: + the contiguous layout is what makes "build the index tables over the + global sequence, then row-slice this rank's queries" correct, and what + makes the all-gathered KV land in natural global order. + """ + cp_pg = pg_collection.cp if pg_collection is not None else None + if cp_pg is not None and getattr(cp_pg, "nranks", 1) > 1: + self.cp_group = cp_pg + self.cp_size = cp_pg.nranks + self.cp_rank = cp_pg.rank + self.cp_enabled = True + if ( + getattr(config, "cp_balance_mode", None) + != "contiguous_allgather" + ): + raise NotImplementedError( + f"{type(self).__name__} under context parallel requires " + "cp_balance_mode='contiguous_allgather' (the same mode the " + "hybrid model's HCA layers require), got " + f"{getattr(config, 'cp_balance_mode', None)!r}." + ) + else: + self.cp_group = None + self.cp_size = 1 + self.cp_rank = 0 + self.cp_enabled = False + + def _needs_indexer_loss(self) -> bool: + """Whether this forward should build and attach the indexer loss. + + ``paddle.is_grad_enabled()`` is what makes the loss count exactly once + under recompute: the first (no-grad) forward only produces the attention + output, the second one attaches the loss. + """ + return ( + self.training + and paddle.is_grad_enabled() + and self.indexer_loss_coeff > 0 + ) + + def _indexer_projections(self, x, qr, position_offset, grad_enabled): + """``(index_q, index_k, weights)`` from the DSA indexer. + + ``x`` / ``qr`` are always detached first: the indexer loss must never + flow back into the backbone, independently of whether the backbone + parameters are frozen. ``index_k`` comes back all-gathered to + ``s_global`` when CP is on (the indexer gathers the 128-wide key rather + than the hidden states, which is ~32x less traffic). + + ``weights`` is returned exactly as ``DSAIndexer.forward_before_topk`` + produced it, i.e. carrying ``n_heads**-0.5 * head_dim**-0.5``. Every + kernel-backed caller must undo the ``head_dim`` half itself, because both + the cuDNN and the tilelang indexer kernels re-apply ``dim**-0.5`` + internally; only a pure-paddle evaluation of the score (as in + ``dsa_attention.FusedDSAIndexerLoss``) uses it unscaled. + """ + x_det, qr_det = x.detach(), qr.detach() + if grad_enabled: + x_det.stop_gradient = False + qr_det.stop_gradient = False + return self.indexer.forward_before_topk( + x_det, qr_det, position_offset, self.cp_group + ) + with paddle.no_grad(): + return self.indexer.forward_before_topk( + x_det, qr_det, position_offset, self.cp_group + ) + + def _indexer_valid_range( + self, + s_global, + doc_start, + doc_len, + is_valid, + window, + position_offset=0, + s_local=None, + ): + """Candidate range per query, in **global token** space. + + ``window`` is how many trailing causal tokens to exclude, i.e. the + forced local window the sparse phase adds separately: clamping the right + edge to ``doc_start + causal_len - window`` removes every duplicate while + leaving the full top-k budget for distant tokens. Because the clamped end + never exceeds the kernel's own causal limit, no masked ``-inf`` column can + enter the top-k. The warmup phase passes ``0``: it has no forced window, + its candidate set is the whole per-document causal span. + + Built over the global sequence and row-sliced to this CP rank; the two + columns stay global token ids, which is what the kernel's + ``seq_offset``-aware causal bound expects. + + Returns: + ``(valid_range [1, s_local, 2] int32, row_empty [1, s_local, 1])``. + """ + positions = paddle.arange(s_global, dtype="int64") + causal_avail = paddle.minimum(positions - doc_start + 1, doc_len) + n_avail = paddle.clip(causal_avail - window, min=0) + n_avail = paddle.where(is_valid, n_avail, paddle.zeros_like(n_avail)) + valid_range = paddle.stack( + [doc_start, doc_start + n_avail], axis=-1 + ).cast("int32") + if s_local is not None and s_local != s_global: + valid_range = valid_range[ + position_offset : position_offset + s_local + ] + n_avail = n_avail[position_offset : position_offset + s_local] + rows = int(valid_range.shape[0]) + return valid_range.unsqueeze(0), (n_avail == 0).reshape([1, rows, 1]) + + def _indexer_loss_mask(self, input_ids: Tensor | None, b: int, s: int): + """``([b, s] float32 row mask, its row count)`` from ``input_ids``. + + ``(None, None)`` when no ``input_ids`` reached this layer (inference and + the direct-construction unit tests), which keeps the plain row mean. + + Under CP the mask is this rank's row slice but the denominator is the + **global** valid-row count, so summing the per-rank losses reproduces the + single-rank reduction. ``input_ids`` arrives sharded unless + ``experimental_dataflow``, exactly as at ``csa_attention.py:2419-2428``. + """ + if input_ids is None: + return None, None + pad_token_id = getattr(self.config, "pad_token_id", 0) + assert pad_token_id is not None, ( + "pad_token_id must be set in config when input_ids is provided" + ) + if self.cp_enabled: + if not getattr(self.config, "experimental_dataflow", False): + input_ids = ContextParallelGatherOp.apply( + input_ids, axis=1, mode=self.config.cp_balance_mode + ) + loss_mask_global = ( + input_ids.reshape([b, self.cp_size * s]) != pad_token_id + ).astype(paddle.float32) + valid_rows = max(float(loss_mask_global.sum()), 1.0) + offset = self.cp_rank * s + return loss_mask_global[:, offset : offset + s], valid_rows + loss_mask = (input_ids.reshape([b, s]) != pad_token_id).astype( + paddle.float32 + ) + return loss_mask, max(float(loss_mask.sum()), 1.0) diff --git a/src/paddlefleet/transformer/mha_dsa_warmup_attention.py b/src/paddlefleet/transformer/mha_dsa_warmup_attention.py new file mode 100644 index 0000000000..cb8f871dee --- /dev/null +++ b/src/paddlefleet/transformer/mha_dsa_warmup_attention.py @@ -0,0 +1,475 @@ +# 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. + +"""Phase 2 (DSA warmup) core attention: the pretraining dense MHA, plus an +indexer trained over the full causal candidate set. + +Phase 2 freezes the backbone (``train_indexer_only``) and has **no top-k on +either side**: the attention must see exactly the activations phase 1 +pretrained, and the indexer is supervised over every causal column so it cannot +reinforce its own random initial ranking. + +Attention therefore has nothing to gain from the absorbed latent MQA of phase 3 +and a great deal to lose: routing a zero-sparsity candidate set through the +block-sparse kernel means materialising a per-document causal index table +(``[b, s, s]`` int32, 256MB at s=8192, built from an ``[s, s]`` int64 +intermediate twice that size -- ``csa_attention._build_mqa_causal_topk_idxs_ +from_doc_bounds``) and then having the kernel walk all ``s`` columns anyway. +Dense flashmask does the same maths with no ``s x s`` tensor at all. So this +class subclasses :class:`DotProductAttention` and delegates the whole attention +half to ``super().forward``; only the indexer loss is new. + +Consequences of that choice, all deliberate: + +* ``MLASelfAttention.mqa_latent`` is False in this phase + (``hybrid_mla_indexer.latent_mqa_enabled``), so ``kv_b_proj`` materialises + per-head K/V and the layer is bit-for-bit the phase-1 layer. Including its + constraints: a dense MLA attention sink needs + ``FLAGS_flash_attn_version in (3, 4)`` (checked at construction in + ``multi_latent_attention.py``). "Phase 2 runs wherever phase 1 runs" is the + point. +* The sink is *not* in the KL target, exactly as in phase 3: the target + normalises over the indexer's own candidate set, and the sink is outside it by + construction. See ``mqa_latent_attention.MQALatentAttention._attn_target`` + for the measured size of the alternative definition. +* ``softmax_offset`` is built by the shared ``build_softmax_offset`` (inherited + from :class:`DotProductAttention`) and the indexer lives at + ``core_attention.indexer.*``, so the phase-1 -> 2 -> 3 parameter names line up + and an HF checkpoint moves between the phases with no rename mapping. + +What is *not* removed here: ``target`` / ``probs`` / ``columns`` are still +``[1, s, s_global]`` (three 256MB transients at s=8192). That width is the +full-candidate KL objective itself -- ``csa_indexer_topk_fwd`` in its documented +"full-candidate selection" mode returns one slot per candidate -- not an +artefact of the attention backend. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import paddle +import paddle.nn.functional as F +from paddle import Tensor +from paddle.distributed.fleet.meta_parallel import LayerSpec, build_spec_layer + +from paddlefleet.process_groups_config import ProcessGroupCollection +from paddlefleet.transformer.cp_utils import all_gather_cp +from paddlefleet.transformer.csa_attention import ( + TileLangCSAIndexerLossAutoScaler, + _derive_csa_doc_boundaries, + _validate_csa_docmask_shape, +) +from paddlefleet.transformer.dot_product_attention import DotProductAttention +from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.hybrid_mla_indexer import HybridMLAIndexerMixin + +if TYPE_CHECKING: + from paddlefleet.transformer.enums import AttnMaskType + +# Row budget of the KL-target scoring loop, as ``rows x candidate slots``. +# The peak tensor is the fp32 score block ``[chunk, h, s_global]``; at +# s_global=8192 / h=64 this budget gives chunk=16, i.e. 33.5MB. +_TARGET_ROW_SLOTS = 256 * 512 +_NEG_INF = -1e30 +_EPS = 1e-10 + + +@dataclass +class MHADSAWarmupAttentionSublayersSpec: + """Sublayers spec for :class:`MHADSAWarmupAttention`. + + Args: + indexer: ``DSAIndexer`` spec. Always provided in this phase -- a warmup + layer without an indexer would just be the phase-1 dense layer, and + ``gpt_layer_specs`` builds that one instead. + """ + + indexer: LayerSpec | type = None + + +class MHADSAWarmupAttention(HybridMLAIndexerMixin, DotProductAttention): + """Dense MHA attention (phase 1's, unchanged) with the DSA indexer loss.""" + + def __init__( + self, + config, + sublayers_spec: MHADSAWarmupAttentionSublayersSpec, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + pg_collection: ProcessGroupCollection | None = None, + **kwargs, + ): + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + super().__init__( + config=config, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type=attention_type, + pg_collection=pg_collection, + **kwargs, + ) + + DSAIndexerLossLoggingHelper.register_total_num_layers(config) + self._init_hybrid_mla_cp_state(config, pg_collection) + if sublayers_spec.indexer is None: + raise ValueError( + "MHADSAWarmupAttention requires an indexer; a warmup layer " + "without one is just the dense phase-1 layer." + ) + self.indexer = build_spec_layer( + sublayers_spec.indexer, + config=config, + layer_number=layer_number, + pg_collection=pg_collection, + ) + self.indexer_loss_coeff = float( + getattr(config, "dsa_indexer_loss_coeff", 0.0) or 0.0 + ) + + def forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor | None, + attn_mask_startend_row_indices: Tensor | None = None, + attn_mask_type: AttnMaskType | None = None, + attention_bias: Tensor | None = None, + packed_seq_params=None, + use_rr_flash_attention: bool = False, + past_key_values=None, + layer_idx=None, + use_cache: bool = False, + x: Tensor | None = None, + qr: Tensor | None = None, + kv_compressed: Tensor | None = None, + k_pos_emb: Tensor | None = None, + q_absorbed: Tensor | None = None, + v_b_proj_weight: Tensor | None = None, + input_ids: Tensor | None = None, + ) -> Tensor: + """Dense MHA forward, with the indexer loss attached to its output. + + Every attention-side argument is forwarded to + :meth:`DotProductAttention.forward` untouched, so this phase inherits + the whole dense dispatch (flashmask / flashmask-CP / SDPA / eager / + varlen / refined recompute / KV cache) rather than re-implementing any + of it. ``x`` / ``qr`` reach the indexer instead of being ignored, and + ``input_ids`` (which the base class does not accept) only builds the + indexer-loss row mask. + + Returns: + ``[b, s, h * v_head_dim]`` -- this rank's query slice under CP. + """ + output = super().forward( + query, + key, + value, + attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + use_rr_flash_attention=use_rr_flash_attention, + past_key_values=past_key_values, + layer_idx=layer_idx, + use_cache=use_cache, + x=x, + qr=qr, + kv_compressed=kv_compressed, + k_pos_emb=k_pos_emb, + q_absorbed=q_absorbed, + v_b_proj_weight=v_b_proj_weight, + ) + if not self._needs_indexer_loss(): + return output + return self._attach_indexer_loss( + output, + query, + key, + x, + qr, + attn_mask_startend_row_indices, + packed_seq_params, + input_ids, + ) + + def _attach_indexer_loss( + self, + output: Tensor, + query: Tensor, + key: Tensor, + x: Tensor, + qr: Tensor, + row_end: Tensor | None, + packed_seq_params, + input_ids: Tensor | None, + ) -> Tensor: + """Full-candidate indexer KL, attached to ``output``'s gradient. + + One ``csa_indexer_topk_fwd`` call in its documented "full-candidate + selection" mode (``ratio=1``, ``topk_effective=s_global``) gives both the + candidate columns and the indexer's softmax over them; the head dimension + never leaves the kernel. The backward is upstream's ``csa_indexer_bwd`` + via :class:`TileLangCSAIndexerLossAutoScaler`, whose tilelang branch + computes exactly ``(P - Q) * coeff / valid_rows``. + + Recompute: the loss is attached on the grad-enabled forward only + (``_needs_indexer_loss``), so it is counted once, and the no-grad forward + skips the indexer entirely instead of computing and discarding it. + + CP: ``index_k`` is all-gathered to ``s_global`` inside + ``forward_before_topk`` and the per-head ``key`` here, ``valid_range`` is + built over the global sequence and row-sliced, and ``valid_rows`` is the + global valid-row count -- so the per-rank losses sum to the single-rank + one. + """ + from paddlefleet.tilelang_ops import csa_indexer_topk_fwd + + if packed_seq_params is not None: + raise NotImplementedError( + "the DSA warmup indexer loss does not support " + "packed_seq_params; document masking is driven by " + "attn_mask_startend_row_indices." + ) + b, s_local = int(query.shape[0]), int(query.shape[1]) + if b != 1: + raise NotImplementedError( + "the DSA warmup indexer loss requires micro batch size 1 " + f"(documents are packed along the sequence), got b={b}." + ) + s_global = s_local * self.cp_size + position_offset = self.cp_rank * s_local + + self._check_tilelang_indexer_support() + index_q, index_k, weights = self._indexer_projections( + x, qr, position_offset, grad_enabled=True + ) + # ``DSAIndexer`` pre-bakes ``head_dim**-0.5`` into the weights and the + # tilelang indexer kernels apply ``dim**-0.5`` themselves, so undo the + # pre-bake once -- before both the forward call and the weights handed to + # the backward, so the two agree. Getting this wrong is silent: measured + # against a plain-paddle reference the un-baked weights match to max_abs + # 3.0e-8 / cosine 1-1.5e-13, unscaled they give max_abs 7.5e-1 / + # cosine 0.62. + weights = weights * (float(self.indexer.head_dim) ** 0.5) + + with paddle.no_grad(): + if row_end is None: + row_end = paddle.full( + [b, 1, s_global, 1], s_global, dtype="int32" + ) + _validate_csa_docmask_shape(row_end, b, s_global) + doc_start, doc_len, is_valid, _, _ = _derive_csa_doc_boundaries( + row_end, s_global + ) + # No forced window in this phase, so the candidate range is the + # whole per-document causal span. + valid_range, row_empty = self._indexer_valid_range( + s_global, + doc_start, + doc_len, + is_valid, + 0, + position_offset, + s_local, + ) + columns, probs = csa_indexer_topk_fwd( + index_q.detach(), + index_k.detach(), + weights.detach(), + ratio=1, + topk_effective=s_global, + seq_offset=position_offset, + valid_range=valid_range, + ) + columns = paddle.where( + row_empty, paddle.full_like(columns, -1), columns + ) + probs = paddle.where(columns >= 0, probs, paddle.zeros_like(probs)) + target = self._dense_attn_target( + query.detach(), + all_gather_cp(key.detach(), dim=1, group=self.cp_group), + columns, + doc_start, + is_valid, + position_offset, + s_local, + s_global, + ) + loss_mask, valid_rows = self._indexer_loss_mask( + input_ids, b, s_local + ) + # The unmasked branch's ``/cp_size`` has to sit in ``loss_coeff`` + # (and therefore reach the backward), not only in the logged scalar + # the way ``csa_attention`` places it -- see the long comment at + # ``mqa_latent_attention._forward_sparse``. + loss_coeff = ( + self.indexer_loss_coeff + if loss_mask is not None + else self.indexer_loss_coeff / self.cp_size + ) + kl = ( + target * (paddle.log(target + _EPS) - paddle.log(probs + _EPS)) + ).sum(axis=-1) + if loss_mask is None: + loss = kl.mean() * loss_coeff + else: + loss = (kl * loss_mask).sum() / valid_rows * loss_coeff + + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=loss, + layer_number=self.layer_number, + num_layers=DSAIndexerLossLoggingHelper.get_total_num_layers( + self.config + ), + ) + # ``TileLangCSAIndexerLossAutoScaler`` returns its first argument + # unchanged when that argument needs a gradient, and Paddle records a + # PyLayer returning one of its inputs as an inplace write on it + # (version 0 -> 1). The dense attention backward *saves its own output*, + # so the version bump makes the attention backward -- not ours -- raise + # ``PermissionDenied: Tensor ... modified by an inplace operation`` + # (``tensor_wrapper.h:268``). Hand the scaler a fresh tensor to bump + # instead; ``clone`` is a gradient identity. The copy is confined to + # this phase: phase 3 passes a fresh matmul result + # (``mqa_latent_attention._deabsorb``) and the CSA call sites a fresh + # kernel output, and the scaler already clones when the backbone is + # frozen, which is why the production ``train_indexer_only`` run never + # hit this. + if not output.stop_gradient: + output = output.clone() + return TileLangCSAIndexerLossAutoScaler.apply( + output, + target, + index_q, + weights, + index_k, + columns, + probs, + loss_coeff, + "tilelang", + valid_rows, + loss_mask, + ) + + def _check_tilelang_indexer_support(self) -> None: + """Fail loudly on the one tilelang indexer constraint we cannot absorb. + + The candidate *width* needs no check: the wrappers round + ``topk_effective`` up to a power-of-two multiple of their block and crop + the result back (``csa_indexer_fwd.py:430-462``, + ``csa_indexer_bwd.py:617-638``), so any causal span from 1 upwards is + served -- measured at s = 1/2/4/8/16/32/300/384/512/8192. + + The head count is different: ``index_n_heads`` other than 64 trips the + kernel's warp tiling with a bare + ``Check failed: (m_warp * n_warp == num_warps)`` from inside tilelang + (measured with 8). Reject that here rather than at the launch. It is not + checked at config time on purpose -- that would make every + small-geometry unit fixture unrepresentable. + """ + heads = int(self.indexer.n_heads) + if heads != 64: + raise ValueError( + "the tilelang indexer's warp tiling requires " + f"index_n_heads == 64 (measured: 8 fails inside the kernel), " + f"got {heads}." + ) + + def _dense_attn_target( + self, + query: Tensor, + key: Tensor, + columns: Tensor, + doc_start: Tensor, + is_valid: Tensor, + position_offset: int, + s_local: int, + s_global: int, + ) -> Tensor: + """KL target: head-summed attention probs, in the kernel's column order. + + The per-head layout forbids the phase-3 trick of gathering the selected + keys (``[chunk, width, h, dk]`` is 3.2GB at chunk=16 / width=8192 / + h=64 / dk=256), so score in **natural column order** instead -- the full + causal row is the candidate set in this phase anyway -- and permute + afterwards with ``take_along_axis``. That is exact, not an + approximation: ``columns`` holds global token ids and the natural-order + index *is* the global token id. + + Scored in query-row chunks, ``[chunk, h, s_global]`` fp32 at a time. The + matmul runs in the input dtype (bf16) with fp32 accumulation, as the + tilelang kernel does internally for the CSA layers; the softmax and the + L1 normalisation are fp32. + + Args: + query: ``[1, s_local, h, dk]`` detached per-head query (local rows). + key: ``[1, s_global, h, dk]`` per-head key (all-gathered under CP). + columns: ``[1, s_local, s_global]`` int32 candidate ids, ``-1`` for + empty slots. + doc_start / is_valid: ``[s_global]``, from + ``_derive_csa_doc_boundaries``. + + Returns: + ``[1, s_local, s_global]`` float32, rows summing to 1 (0 for empty + rows). + """ + h = int(query.shape[2]) + chunk = max(1, _TARGET_ROW_SLOTS // s_global) + # Head-major once, rather than per chunk: [h, s, dk]. + q_all = query[0].transpose([1, 0, 2]) + k_all = key[0].transpose([1, 0, 2]) + cols = paddle.arange(s_global, dtype="int64") + parts = [] + for start in range(0, s_local, chunk): + end = min(start + chunk, s_local) + lo, hi = position_offset + start, position_offset + end + rows = paddle.arange(lo, hi, dtype="int64").unsqueeze(1) + # Per-document causal, i.e. exactly the range + # ``_indexer_valid_range(window=0)`` handed the kernel. + allowed = ( + (cols.unsqueeze(0) >= doc_start[lo:hi].unsqueeze(1)) + & (cols.unsqueeze(0) <= rows) + & is_valid[lo:hi].unsqueeze(1) + ) + scores = ( + paddle.matmul( + q_all[:, start:end], k_all, transpose_y=True + ).cast("float32") + * self.softmax_scale + ) + scores = paddle.where( + allowed.unsqueeze(0), scores, paddle.full_like(scores, _NEG_INF) + ) + # Each head contributes mass 1; an all-masked (padding) row would + # give a uniform softmax, so zero it explicitly -- a row of zeros + # must stay a row of zeros, because the KL reduction divides by the + # valid-row count, not by the row sum. + probs = F.softmax(scores, axis=-1).sum(axis=0) + parts.append(paddle.where(allowed, probs, paddle.zeros_like(probs))) + target = paddle.concat(parts, axis=0) + target = target / target.sum(axis=-1, keepdim=True).clip(min=_EPS) + + idx = columns[0].cast("int64") + valid = idx >= 0 + target = paddle.take_along_axis( + target, paddle.where(valid, idx, paddle.zeros_like(idx)), axis=-1 + ) + target = paddle.where(valid, target, paddle.zeros_like(target)) + return target.unsqueeze(0) diff --git a/src/paddlefleet/transformer/mqa_latent_attention.py b/src/paddlefleet/transformer/mqa_latent_attention.py index 7b9f6fdb29..ec69ff9f28 100644 --- a/src/paddlefleet/transformer/mqa_latent_attention.py +++ b/src/paddlefleet/transformer/mqa_latent_attention.py @@ -15,28 +15,31 @@ """Latent MQA core attention for hybrid MLA layers, with DSA. ``hybrid_mla_attention`` selects which core attention the -``csa_compress_ratios == -2`` (MLA) layers of a ``dsv4_hybrid`` model run, and -within the DSA mode ``dsa_indexer_use_sparse_loss`` selects the training phase. -``MQALatentAttention._phase()`` is the single place that is decided, and each -phase has its own ``_forward_*`` with no loss code shared between them: +``csa_compress_ratios == -2`` (MLA) layers of a ``dsv4_hybrid`` model run. This +module owns the two modes that attend to the **absorbed KV latent**, i.e. the +ones that consume a sorted candidate set: -* ``"mha"`` -- unchanged dense MLA (MHA); this module is not used. * ``"mqa_full_causal"`` -> ``_forward_full_causal``. Latent MQA with the indexer dropped, attending to the full per-document causal set. Mathematically identical to the dense MHA phase, so it isolates the absorption from the sparsity for equivalence experiments; ``O(s^2)`` in index memory and therefore not a production mode. -* ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=False`` -> ``_forward_warmup`` - (phase 2, DSA warmup, paired with ``train_indexer_only``). The backbone is - frozen and the indexer is random, so **neither side uses top-k**: attention - runs the same full per-document causal set phase 1 did, and the indexer KL - spans every causal column on both sides, via one ``csa_indexer_topk_fwd`` - call in its documented "full-candidate selection" mode. * ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=True`` -> ``_forward_sparse`` - (phase 3). A forced local window plus Lightning-indexer top-k, i.e. DeepSeek + (phase 3/4). A forced local window plus Lightning-indexer top-k, i.e. DeepSeek Sparse Attention on the KV latent, with the KL restricted to that same selected set. This is the only phase that reads ``index_topk``. +The other two modes are *not* this module, and the shared predicate +``hybrid_mla_indexer.latent_mqa_enabled`` is what keeps the spec dispatch and +``MLASelfAttention.mqa_latent`` in step about that: + +* ``"mha"`` -- unchanged dense MLA (MHA). +* ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=False`` (phase 2, DSA warmup, + paired with ``train_indexer_only``) -- also dense MHA, plus the indexer, in + ``mha_dsa_warmup_attention.MHADSAWarmupAttention``. That phase has no top-k on + either side, so absorbing here would buy nothing and cost a zero-sparsity + ``[b, s, s]`` index table pushed through the block-sparse kernel. + The indexer reuses the model-wide ``index_n_heads`` / ``index_head_dim``. Note the ``mqa_*`` modes here are *latent* MQA (this module). A @@ -94,7 +97,6 @@ from paddle import Tensor from paddle.distributed.fleet.meta_parallel import LayerSpec, build_spec_layer -from paddlefleet.context_parallel_utils import ContextParallelGatherOp from paddlefleet.process_groups_config import ProcessGroupCollection from paddlefleet.transformer.cp_utils import all_gather_cp from paddlefleet.transformer.csa_attention import ( @@ -108,6 +110,7 @@ from paddlefleet.transformer.dsa_attention import ( DSAIndexerLossLoggingHelper, ) +from paddlefleet.transformer.hybrid_mla_indexer import HybridMLAIndexerMixin from paddlefleet.transformer.layer import FleetLayer if TYPE_CHECKING: @@ -162,7 +165,7 @@ class MQALatentAttentionSublayersSpec: indexer: LayerSpec | type = None -class MQALatentAttention(FleetLayer): +class MQALatentAttention(HybridMLAIndexerMixin, FleetLayer): """Sparse attention on the absorbed MLA KV latent (``core_attention``). Consumes the pre-absorbed ``query`` / ``key`` produced by @@ -196,33 +199,7 @@ def __init__( pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.pg_collection = pg_collection - # CP state, same derivation as csa_attention.py:2079-2090. - cp_pg = pg_collection.cp if pg_collection is not None else None - if cp_pg is not None and getattr(cp_pg, "nranks", 1) > 1: - self.cp_group = cp_pg - self.cp_size = cp_pg.nranks - self.cp_rank = cp_pg.rank - self.cp_enabled = True - # Deliberately the *same* constraint the HCA layers of this model - # assert (dsv4_hybrid_attention.py:607-611), not a weaker one: the - # contiguous layout is what makes "build the index table over the - # global sequence, then row-slice this rank's queries" correct, and - # it is what makes the all-gathered KV land in natural global order. - if ( - getattr(config, "cp_balance_mode", None) - != "contiguous_allgather" - ): - raise NotImplementedError( - "latent MQA under context parallel requires " - "cp_balance_mode='contiguous_allgather' (the same mode the " - "hybrid model's HCA layers require), got " - f"{getattr(config, 'cp_balance_mode', None)!r}." - ) - else: - self.cp_group = None - self.cp_size = 1 - self.cp_rank = 0 - self.cp_enabled = False + self._init_hybrid_mla_cp_state(config, pg_collection) # ``k_channels`` is the MHA q_head_dim (qk_nope + qk_rope), NOT the 576 # latent width: absorption is exactly score-preserving, so the MHA @@ -286,22 +263,8 @@ def __init__( # when ``softmax_offset`` exists; a sinkless layer has no sink gradient. self.sink_grad_fusion = getattr(config, "dsa_sink_grad_fusion", False) - def _needs_indexer_loss(self) -> bool: - """Whether this forward should build and attach the indexer loss. - - Both DSA phases share this predicate. ``paddle.is_grad_enabled()`` is - what makes the loss count exactly once under recompute: the first - (no-grad) forward only materialises the attention columns, the second one - attaches the loss. - """ - return ( - self.training - and paddle.is_grad_enabled() - and self.indexer_loss_coeff > 0 - ) - def _phase(self) -> str: - """Which of the three training phases this layer runs. + """Which of the two phases this layer runs. The single place the phase is decided, so the attention candidate set and the indexer-loss shape cannot disagree: @@ -310,20 +273,32 @@ def _phase(self) -> str: (``hybrid_mla_attention="mqa_full_causal"``, and the absorption-equivalence unit tests). Per-document full causal attention, no indexer loss. - * ``"warmup"`` -- phase 2 (DSA warmup). The indexer exists but is still - being learned, so attention must not consume its ranking: full causal - attention, and the KL runs over the *full* causal set on both sides. - No top-k anywhere. - * ``"sparse"`` -- phase 3. Attention consumes window + top-k and the + * ``"sparse"`` -- phase 3/4. Attention consumes window + top-k and the KL is restricted to that same selected set. + There is deliberately no warmup state here. Phase 2 has no top-k on + either side, so routing a zero-sparsity candidate set through the + block-sparse kernel would cost a full ``[b, s, s]`` index table for + nothing; it runs the dense MHA of phase 1 instead, in + ``mha_dsa_warmup_attention.MHADSAWarmupAttention``, and + ``hybrid_mla_indexer.latent_mqa_enabled`` is what keeps the spec + dispatch and ``MLASelfAttention.mqa_latent`` from ever building this + class for it. + Read live rather than cached in ``__init__``: a test flipping ``indexer_use_sparse_loss`` on a live module must not be able to desynchronise the two. """ if self.indexer is None: return "full_causal" - return "sparse" if self.indexer_use_sparse_loss else "warmup" + if not self.indexer_use_sparse_loss: + raise ValueError( + "MQALatentAttention has an indexer but " + "dsa_indexer_use_sparse_loss=False; that is the DSA warmup " + "phase, which runs dense MHA in MHADSAWarmupAttention. " + "latent_mqa_enabled() should have dispatched there." + ) + return "sparse" def forward( self, @@ -454,23 +429,6 @@ def forward( s_global, ) - if phase == "warmup": - return self._forward_warmup( - query, - kv, - x, - qr, - v_b_proj_weight, - doc_start, - doc_len, - is_valid, - kv_lora_rank, - input_ids, - position_offset, - s, - s_global, - ) - return self._forward_sparse( query, kv, @@ -508,9 +466,11 @@ def _forward_full_causal( dense MHA phase and is bit-identical across repeated calls (nothing here depends on a top-k tie-break). - Used by two phases -- ``hybrid_mla_attention="mqa_full_causal"``, and - the attention half of the phase-2 warmup, which must not consume the - indexer's ranking while the indexer is still being learned. + Used by ``hybrid_mla_attention="mqa_full_causal"`` (the MHA -> MQA + equivalence isolation) and by the absorption unit tests. The DSA warmup + phase, which also needs a full-causal attention, does *not* come here: + it keeps the dense MHA backend instead, because at zero sparsity the + ``[b, s, s]`` index table below buys nothing. """ b = int(query.shape[0]) token_indices = self._build_full_causal_indices( @@ -521,186 +481,6 @@ def _forward_full_causal( ) return self._deabsorb(core_out, v_b_proj_weight, self.split_kv_b) - # ------------------------------------------------------------------ - # warmup (phase 2) - # ------------------------------------------------------------------ - def _forward_warmup( - self, - query: Tensor, - kv: Tensor, - x: Tensor, - qr: Tensor, - v_b_proj_weight: Tensor, - doc_start: Tensor, - doc_len: Tensor, - is_valid: Tensor, - kv_lora_rank: int, - input_ids: Tensor | None, - position_offset: int, - s_local: int, - s_global: int, - ) -> Tensor: - """Phase 2: frozen backbone, full-causal attention, full-causal KL. - - No top-k on either side. The two halves: - - * **attention** is the same deterministic full-causal set phase 1 uses - (``_forward_full_causal``), so the frozen backbone sees exactly the - activations it was pretrained with while the indexer is still random. - * **the indexer** is supervised over the *whole* per-document causal - span, so it cannot reinforce its own initial ranking. - - Both come from one tilelang call with ``topk_effective = s_global``, which - is the "full-candidate selection" mode ``csa_indexer_topk_fwd`` documents - for exactly this phase -- the CSA layers use it the same way with - ``topk_effective = n_compressed``. The kernel returns the softmax - probabilities over every candidate column plus the column ids, and the - head dimension never leaves the kernel; the backward is upstream's - ``csa_indexer_bwd`` via ``TileLangCSAIndexerLossAutoScaler``, whose - tilelang branch computes exactly ``(P - Q) * coeff / valid_rows``. - - Recompute: the attention column table depends only on the document - boundaries, so the two forwards of a recompute segment are bit-identical - and there is no top-k tie-break to worry about. The loss is attached on - the grad-enabled forward only, so it is counted once; the no-grad forward - skips the indexer entirely rather than computing and discarding it. - - CP: ``index_k`` is all-gathered to ``s_global`` inside - ``forward_before_topk`` and ``kv`` by the caller, ``valid_range`` is built - over the global sequence and row-sliced, and ``valid_rows`` is the global - valid-row count -- so the per-rank losses sum to the single-rank one and - no ``/cp_size`` correction is needed. - """ - from paddlefleet.tilelang_ops import csa_indexer_topk_fwd - - output = self._forward_full_causal( - query, - kv, - v_b_proj_weight, - doc_start, - is_valid, - kv_lora_rank, - position_offset, - s_local, - s_global, - ) - if not self._needs_indexer_loss(): - return output - - b = int(query.shape[0]) - self._check_tilelang_indexer_support() - index_q, index_k, weights = self._indexer_projections( - x, qr, position_offset, grad_enabled=True - ) - # ``DSAIndexer`` pre-bakes ``head_dim**-0.5`` into the weights, and the - # tilelang indexer kernels apply ``dim**-0.5`` themselves -- same - # convention as the cuDNN pair the sparse phase uses, and the opposite of - # the pure-paddle ``FusedDSAIndexerLoss`` reference, which applies none. - # Undo the pre-bake once, before both the forward call and the weights - # handed to the backward, so the two agree. - # Measured (validation_reports/precision_audit_20260809_022929/ops_edge): - # against a plain-paddle reference the un-baked weights match to - # max_abs 3.0e-8 / cosine 1-1.5e-13, while passing them through unscaled - # gives max_abs 7.5e-1 / cosine 0.62. - weights = weights * (float(self.indexer.head_dim) ** 0.5) - # ``topk_effective`` is the causal span itself. The wrapper rounds it up - # to a power-of-two multiple of its block internally and crops the result - # back to the requested width (``csa_indexer_fwd.py:430-462`` / - # ``csa_indexer_bwd.py:617-638``), so there is nothing to round here and - # no surplus ``-1`` slot to carry. Measured at - # s = 1/2/4/8/16/32/300/384/512/8192: the returned width equals - # ``s_global`` exactly, the per-row valid-slot count equals the causal - # length, rows sum to 1 within 9.6e-7 and the backward is finite. - with paddle.no_grad(): - # No forced window in this phase, so the candidate range is the - # whole per-document causal span. - valid_range, row_empty = self._indexer_valid_range( - s_global, - doc_start, - doc_len, - is_valid, - position_offset, - s_local, - window=0, - ) - columns, probs = csa_indexer_topk_fwd( - index_q.detach(), - index_k.detach(), - weights.detach(), - ratio=1, - topk_effective=s_global, - seq_offset=position_offset, - valid_range=valid_range, - ) - columns = paddle.where( - row_empty, paddle.full_like(columns, -1), columns - ) - probs = paddle.where(columns >= 0, probs, paddle.zeros_like(probs)) - target = self._attn_target(query.detach(), kv, columns) - loss_mask, valid_rows = self._indexer_loss_mask( - input_ids, b, s_local - ) - # Same reduction as ``_forward_sparse`` -- see the long comment there - # for why the unmasked branch's ``/cp_size`` has to sit in - # ``loss_coeff`` (and therefore reach the backward) rather than only - # in the logged scalar the way ``csa_attention`` places it. - loss_coeff = ( - self.indexer_loss_coeff - if loss_mask is not None - else self.indexer_loss_coeff / self.cp_size - ) - kl = ( - target * (paddle.log(target + _EPS) - paddle.log(probs + _EPS)) - ).sum(axis=-1) - if loss_mask is None: - loss = kl.mean() * loss_coeff - else: - loss = (kl * loss_mask).sum() / valid_rows * loss_coeff - - DSAIndexerLossLoggingHelper.save_loss_to_tracker( - loss=loss, - layer_number=self.layer_number, - num_layers=DSAIndexerLossLoggingHelper.get_total_num_layers( - self.config - ), - ) - return TileLangCSAIndexerLossAutoScaler.apply( - output, - target, - index_q, - weights, - index_k, - columns, - probs, - loss_coeff, - "tilelang", - valid_rows, - loss_mask, - ) - - def _check_tilelang_indexer_support(self) -> None: - """Fail loudly on the one tilelang indexer constraint we cannot absorb. - - The top-k *width* needs no check: the wrappers round ``topk_effective`` - up to a power-of-two multiple of their block and crop the result back - (``csa_indexer_fwd.py:430-462``, ``csa_indexer_bwd.py:617-638``), so any - causal span from 1 upwards is served -- measured at - s = 1/2/4/8/16/32/300/384/512/8192. - - The head count is different: ``index_n_heads`` other than 64 trips the - kernel's warp tiling with a bare - ``Check failed: (m_warp * n_warp == num_warps)`` from inside tilelang - (measured with 8). Reject that here rather than at the launch. It is not - checked at config time on purpose -- that would make every - small-geometry unit fixture unrepresentable. - """ - heads = int(self.indexer.n_heads) - if heads != 64: - raise ValueError( - "the tilelang indexer's warp tiling requires index_n_heads == 64 " - f"(measured: 8 fails inside the kernel), got {heads}." - ) - # ------------------------------------------------------------------ # index construction / kernel plumbing # ------------------------------------------------------------------ @@ -726,79 +506,6 @@ def _build_full_causal_indices( indices.stop_gradient = True return indices - def _indexer_projections(self, x, qr, position_offset, grad_enabled): - """``(index_q, index_k, weights)`` from the DSA indexer. - - ``x`` / ``qr`` are always detached first: the indexer loss must never - flow back into the backbone, independently of whether the backbone - parameters are frozen. ``index_k`` comes back all-gathered to - ``s_global`` when CP is on (the indexer gathers the 128-wide key rather - than the hidden states, which is ~32x less traffic). - - ``weights`` is returned exactly as ``DSAIndexer.forward_before_topk`` - produced it, i.e. carrying ``n_heads**-0.5 * head_dim**-0.5``. Every - kernel-backed caller must undo the ``head_dim`` half itself, because both - the cuDNN and the tilelang indexer kernels re-apply ``dim**-0.5`` - internally; only a pure-paddle evaluation of the score (as in - ``dsa_attention.FusedDSAIndexerLoss``) uses it unscaled. - """ - x_det, qr_det = x.detach(), qr.detach() - if grad_enabled: - x_det.stop_gradient = False - qr_det.stop_gradient = False - return self.indexer.forward_before_topk( - x_det, qr_det, position_offset, self.cp_group - ) - with paddle.no_grad(): - return self.indexer.forward_before_topk( - x_det, qr_det, position_offset, self.cp_group - ) - - def _indexer_valid_range( - self, - s_global, - doc_start, - doc_len, - is_valid, - position_offset=0, - s_local=None, - window=None, - ): - """Candidate range per query, in **global token** space. - - ``window`` is how many trailing causal tokens to exclude, i.e. the - forced local window the sparse phase adds separately: clamping the right - edge to ``doc_start + causal_len - window`` removes every duplicate while - leaving the full top-k budget for distant tokens. Because the clamped end - never exceeds the kernel's own causal limit, no masked ``-inf`` column can - enter the top-k. Defaults to ``self.csa_window_size``; the warmup phase - passes ``0`` because it has no forced window -- its candidate set is the - whole per-document causal span. - - Built over the global sequence and row-sliced to this CP rank; the two - columns stay global token ids, which is what the kernel's - ``seq_offset``-aware causal bound expects. - - Returns: - ``(valid_range [1, s_local, 2] int32, row_empty [1, s_local, 1])``. - """ - if window is None: - window = self.window_size - positions = paddle.arange(s_global, dtype="int64") - causal_avail = paddle.minimum(positions - doc_start + 1, doc_len) - n_avail = paddle.clip(causal_avail - window, min=0) - n_avail = paddle.where(is_valid, n_avail, paddle.zeros_like(n_avail)) - valid_range = paddle.stack( - [doc_start, doc_start + n_avail], axis=-1 - ).cast("int32") - if s_local is not None and s_local != s_global: - valid_range = valid_range[ - position_offset : position_offset + s_local - ] - n_avail = n_avail[position_offset : position_offset + s_local] - rows = int(valid_range.shape[0]) - return valid_range.unsqueeze(0), (n_avail == 0).reshape([1, rows, 1]) - def _sparse_attn( self, query, kv, token_indices, sm_scale, d_v, indexer_topk=0 ): @@ -869,8 +576,9 @@ def _forward_sparse( The indexer is trained enough to steer attention, so both sides narrow to the selected set. Reached only when ``_phase() == "sparse"``, i.e. - ``dsa_indexer_use_sparse_loss=True``; the warmup phase has its own branch - and shares no loss code with this one. + ``dsa_indexer_use_sparse_loss=True``; the warmup phase is a different + class entirely (``mha_dsa_warmup_attention``) and shares no loss code + with this one. Recompute: the loss is attached on the grad-enabled forward only, so under full recompute it is counted once, on the second pass. Both passes see @@ -897,7 +605,13 @@ def _forward_sparse( :, position_offset : position_offset + s ] valid_range, row_empty = self._indexer_valid_range( - s_global, doc_start, doc_len, is_valid, position_offset, s + s_global, + doc_start, + doc_len, + is_valid, + self.window_size, + position_offset, + s, ) q_idx, k_idx, w_idx = self._indexer_projections( @@ -1071,39 +785,6 @@ def _forward_sparse( loss_mask, ) - def _indexer_loss_mask(self, input_ids, b, s): - """``([b, s] float32 row mask, its row count)`` from ``input_ids``. - - ``(None, None)`` when no ``input_ids`` reached this layer (inference and - the direct-construction unit tests), which keeps the plain row mean. - - Under CP the mask is this rank's row slice but the denominator is the - **global** valid-row count, so summing the per-rank losses reproduces the - single-rank reduction. ``input_ids`` arrives sharded unless - ``experimental_dataflow``, exactly as at ``csa_attention.py:2419-2428``. - """ - if input_ids is None: - return None, None - pad_token_id = getattr(self.config, "pad_token_id", 0) - assert pad_token_id is not None, ( - "pad_token_id must be set in config when input_ids is provided" - ) - if self.cp_enabled: - if not getattr(self.config, "experimental_dataflow", False): - input_ids = ContextParallelGatherOp.apply( - input_ids, axis=1, mode=self.config.cp_balance_mode - ) - loss_mask_global = ( - input_ids.reshape([b, self.cp_size * s]) != pad_token_id - ).astype(paddle.float32) - valid_rows = max(float(loss_mask_global.sum()), 1.0) - offset = self.cp_rank * s - return loss_mask_global[:, offset : offset + s], valid_rows - loss_mask = (input_ids.reshape([b, s]) != pad_token_id).astype( - paddle.float32 - ) - return loss_mask, max(float(loss_mask.sum()), 1.0) - def _attn_target(self, query, kv, kl_columns, lse_indexer=None) -> Tensor: """KL target: head-summed attention probs over the indexer's own columns. @@ -1130,15 +811,14 @@ def _attn_target(self, query, kv, kl_columns, lse_indexer=None) -> Tensor: query: ``[1, s, h, dk]`` detached absorbed query (local rows). kv: ``[1, s_global, dk]`` latent keys (all-gathered under CP). kl_columns: ``[1, s, w]`` int32 global column ids the KL scores, - ``-1`` for empty slots. Column *order* is irrelevant, so the - warmup phase passes the indexer's score-ordered table directly. + ``-1`` for empty slots. Column *order* is irrelevant -- the + indexer's score-ordered table can be passed straight in. lse_indexer: ``[1, s, 64]`` float32 per-head LSE over exactly ``kl_columns`` (``mqa_sparse_attn(indexer_topk=...)``). When present the cuDNN score-recompute kernel does the whole thing in - one launch. ``None`` in the warmup phase -- its candidate set is - not the attention set, so no matching LSE exists -- and when the - budget is not one of ``_LSE_INDEXER_TOPKS``, so the Python path - stays as the reference and the fallback. + one launch. ``None`` when the budget is not one of + ``_LSE_INDEXER_TOPKS``, so the Python path stays as the reference + and the fallback. Returns: ``[1, s, w]`` float32 rows summing to 1 (0 for empty rows). diff --git a/src/paddlefleet/transformer/multi_latent_attention.py b/src/paddlefleet/transformer/multi_latent_attention.py index a6d5c7a31d..53434b0a94 100644 --- a/src/paddlefleet/transformer/multi_latent_attention.py +++ b/src/paddlefleet/transformer/multi_latent_attention.py @@ -60,6 +60,7 @@ ) from paddlefleet.transformer.attention import Attention from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled from paddlefleet.transformer.transformer_config import TransformerConfig from paddlefleet.utils import get_pg_rank, get_pg_size @@ -459,11 +460,13 @@ def __init__( # the per-head K/V produced by ``kv_b_proj`` ("runtime absorption"). This # keeps every parameter byte-identical to the MHA layout -- only the # activations change -- so an MHA checkpoint loads into an MQA run. - self.mqa_latent = getattr( - config, "experimental_attention_variant", None - ) == "dsv4_hybrid" and getattr( - config, "hybrid_mla_attention", "mha" - ) in ("mqa_dsa", "mqa_full_causal") + # + # Only the phases that actually consume a *sorted* candidate set absorb: + # the DSA warmup phase has no top-k anywhere, so it stays on the dense + # MHA path (``MHADSAWarmupAttention``) rather than paying for a + # zero-sparsity block-sparse call. ``gpt_layer_specs`` dispatches the core + # attention on this same predicate, so the two cannot disagree. + self.mqa_latent = latent_mqa_enabled(config) # ``mqa_split_kv_b_proj`` trades that property for speed: # ``kv_b_proj`` is replaced by standalone ``k_b_proj`` / ``v_b_proj`` # absorption parameters, pre-laid-out so each side is one grouped GEMM @@ -580,7 +583,10 @@ def __init__( # deliberately do NOT set the flag ourselves: it is process-global and # would switch every other (HCA / CSA) layer's kernel too. # Latent MQA needs neither check -- its block-sparse kernel - # supports the sink natively and up-casts it internally. Neither do the + # supports the sink natively and up-casts it internally. The DSA warmup + # phase is *not* latent MQA (``latent_mqa_enabled`` is False there), so + # it is covered by this check, deliberately: phase 2 is phase 1's dense + # attention and inherits phase 1's constraints. Neither do the # HySparse absorbed-MQA layers: ``gpt_layer_specs.py:318`` builds every # MLA layer as ``MQASelfAttention`` when ``enable_hy_sparse_attention`` # is on, and for an SWA layer (``is_mqa``, :1919) that subclass runs the @@ -610,9 +616,13 @@ def __init__( f"FLAGS_flash_attn_version={fa_version} (only 3 or 4 reach " "that path). Either export FLAGS_flash_attn_version=4 " "(NOTE: process-global -- it changes the HCA/CSA layers' " - "kernel as well), or run the hybrid MLA layers with " - "hybrid_mla_attention='mqa_dsa' / 'mqa_full_causal', whose " - "block-sparse kernel supports the sink natively." + "kernel as well), or run the hybrid MLA layers on the " + "absorbed latent MQA, whose block-sparse kernel supports " + "the sink natively: hybrid_mla_attention=" + "'mqa_full_causal', or 'mqa_dsa' with " + "dsa_indexer_use_sparse_loss=True. 'mqa_dsa' with " + "dsa_indexer_use_sparse_loss=False (the DSA warmup phase) " + "is *this* dense path by design and does not escape it." ) if "bfloat16" not in str(self.config.params_dtype): raise RuntimeError( @@ -940,10 +950,15 @@ def forward( # The indexer-loss row mask needs ``input_ids``: the packed sequence's # trailing padding is invisible to ``attn_mask_startend_row_indices``. - # Only the non-absorbed-MQA core attention accepts it (and only that one - # owns an indexer), so keep the kwarg off every other core attention. + # Asked of the core attention itself rather than derived from + # ``mqa_latent``, because both hybrid MLA phases with an indexer want it + # while only one of them absorbs -- and no core attention may be handed a + # kwarg it does not accept. core_attn_extra = {} - if self.mqa_latent and kwargs.get("input_ids") is not None: + if ( + getattr(self.core_attention, "accepts_input_ids", False) + and kwargs.get("input_ids") is not None + ): core_attn_extra["input_ids"] = kwargs["input_ids"] if self.mqa_latent: diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index a154fab3de..092e65fc61 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -887,14 +887,25 @@ class TransformerConfig(ModelParallelConfig): - ``"mha"`` (default): ``kv_b_proj`` materialises per-head K/V and dense flash attention runs on them. Leaving the field unset keeps this behaviour. - - ``"mqa_dsa"``: latent MQA -- attention runs on the single shared - ``kv_lora_rank + qk_rope_head_dim`` latent head and a DSA (Lightning) - indexer selects the attended columns (forced local window + top-k). + - ``"mqa_dsa"``: a DSA (Lightning) indexer is trained on these layers. Which + *attention* backend runs is decided by ``dsa_indexer_use_sparse_loss``: + ``True`` (phase 3/4) is latent MQA on the single shared + ``kv_lora_rank + qk_rope_head_dim`` head, attending to the indexer's forced + local window + top-k; ``False`` (phase 2, the warmup) keeps the dense + per-head attention of ``"mha"`` and only bolts the indexer loss on, because + with no top-k on either side latent MQA would route a zero-sparsity + candidate set through the block-sparse kernel and pay a ``[b, s, s]`` index + table for nothing. - ``"mqa_full_causal"``: latent MQA with no indexer -- attend to the full per-document causal set. - Both ``mqa_*`` modes absorb the query against ``kv_b_proj.weight`` at - *runtime* (activation level, not weight level), so the parameter layout stays + The single predicate behind that split is + ``hybrid_mla_indexer.latent_mqa_enabled``, read both by the spec dispatch and + by ``MLASelfAttention.mqa_latent``, so the core attention and the activations + it is fed can never disagree. + + Every mode that absorbs (``mqa_full_causal``, and ``mqa_dsa`` in the sparse + phase) absorbs the query against ``kv_b_proj.weight`` at *runtime* (activation level, not weight level), so the parameter layout stays byte-identical to ``"mha"`` and an MHA checkpoint loads unchanged. The only new weights are the DSA indexer's; ``"mqa_full_causal"`` adds none at all. @@ -911,8 +922,9 @@ class TransformerConfig(ModelParallelConfig): table is ``[b, s, s]`` int32, so memory grows with the square of the sequence length (268 MB per layer at ``s=8192``). - Terminology: the ``mqa_*`` modes above are *latent MQA* (class - ``MQALatentAttention``, ``-2`` layers). A ``csa_compress_ratios`` entry of + Terminology: the modes above that absorb are *latent MQA* (class + ``MQALatentAttention``, ``-2`` layers); the warmup phase is + ``MHADSAWarmupAttention``. A ``csa_compress_ratios`` entry of ``-1`` is *CSA full-causal MQA* (class ``CompressedSparseAttention``) -- a different layer kind that this field does not touch. """ @@ -1194,7 +1206,7 @@ class TransformerConfig(ModelParallelConfig): Covers both indexer flavours of a dsv4-hybrid model: the ``CSAIndexer`` of the CSA layers (``1 < csa_compress_ratios[i] < 128``) and the ``DSAIndexer`` - of the latent MQA layers (``== -2`` with ``hybrid_mla_attention="mqa_dsa"``). + of the hybrid MLA layers (``== -2`` with ``hybrid_mla_attention="mqa_dsa"``). This is about the *checkpoint*, not the training strategy, and is therefore separate from ``train_indexer_only``: @@ -1649,8 +1661,9 @@ def __post_init__(self): raise ValueError( f"hybrid_mla_attention={self.hybrid_mla_attention!r} is invalid. " "It must be one of: 'mha' (per-head K/V materialised, dense flash " - "attention -- the default), 'mqa_dsa' (latent MQA + DSA indexer " - "selecting window + top-k columns), or 'mqa_full_causal' (latent " + "attention -- the default), 'mqa_dsa' (DSA indexer; dense " + "attention while it warms up, then latent MQA on its window + " + "top-k columns), or 'mqa_full_causal' (latent " "MQA with no indexer, full per-document causal set)." ) if self.hybrid_mla_attention != "mha": @@ -1682,13 +1695,14 @@ def __post_init__(self): "exists to isolate absorption from sparsity, not to save memory." ) if self.hybrid_mla_attention == "mqa_dsa": - # On the ``-2`` layers ``dsa_indexer_use_sparse_loss`` decides both - # the indexer-loss candidate set and the attention candidate set, - # because they are one decision: while the indexer is still being - # learned attention must not consume its ranking, so the warmup phase - # runs full-causal attention and a KL over the whole causal set -- - # no top-k on either side. The two training phases are therefore - # fixed pairs, and the two mixed combinations are configuration + # On the ``-2`` layers ``dsa_indexer_use_sparse_loss`` decides the + # indexer-loss candidate set, the attention candidate set *and* the + # attention backend, because they are one decision: while the indexer + # is still being learned attention must not consume its ranking, so + # the warmup phase keeps the dense per-head attention of phase 1 and + # runs a KL over the whole causal set -- no top-k on either side, and + # therefore nothing for latent MQA to absorb. The two training phases + # are fixed pairs, and the two mixed combinations are configuration # mistakes rather than modes. if self.train_indexer_only and self.dsa_indexer_use_sparse_loss: raise ValueError( @@ -1710,8 +1724,8 @@ def __post_init__(self): logger.warning( "hybrid_mla_attention='mqa_dsa' with " "dsa_indexer_use_sparse_loss=False runs the warmup shape " - "(full per-document causal attention plus a KL over the " - "whole causal set, no top-k anywhere) while every backbone " + "(dense per-head attention plus a KL over the whole causal " + "set, no top-k anywhere) while every backbone " "parameter still trains. The " "production warmup phase pairs it with " "train_indexer_only=True; the sparse training phase pairs " @@ -1790,17 +1804,29 @@ def __post_init__(self): "hybrid MLA dimensions must be explicit positive integers; " f"invalid fields: {', '.join(invalid)}" ) - if self.mqa_split_kv_b_proj and ( - self.hybrid_mla_attention - not in ("mqa_dsa", "mqa_full_causal") - ): - raise ValueError( - "mqa_split_kv_b_proj=True only means " - "anything for latent MQA, i.e. " - "hybrid_mla_attention='mqa_dsa' or 'mqa_full_causal'; " - "it splits those modes' kv_b_proj into standalone " - "k_b_proj / v_b_proj absorption parameters." + if self.mqa_split_kv_b_proj: + # Same predicate the layer itself uses, so the error can + # never disagree with what gets built. Note the DSA warmup + # phase is *not* latent MQA: it runs dense MHA and keeps + # ``kv_b_proj``, so allowing the flag there would silently + # change the parameter set at the warmup -> sparse switch. + from paddlefleet.transformer.hybrid_mla_indexer import ( + latent_mqa_enabled, ) + + if not latent_mqa_enabled(self): + raise ValueError( + "mqa_split_kv_b_proj=True only means " + "anything for latent MQA, i.e. " + "hybrid_mla_attention='mqa_full_causal' or " + "'mqa_dsa' with dsa_indexer_use_sparse_loss=True; " + "it splits those modes' kv_b_proj into standalone " + "k_b_proj / v_b_proj absorption parameters, while " + "the dense phases keep kv_b_proj itself. Got " + f"hybrid_mla_attention={self.hybrid_mla_attention!r}," + " dsa_indexer_use_sparse_loss=" + f"{self.dsa_indexer_use_sparse_loss!r}." + ) if self.mqa_split_kv_b_proj and getattr( self, "enable_hy_sparse_attention", False ): @@ -1847,9 +1873,10 @@ def __post_init__(self): ) if self.dsa_index_head_dim != 128: raise ValueError( - "hybrid_mla_attention='mqa_dsa' uses the cuDNN " - "indexer, which requires index_head_dim=128, got " - f"{self.dsa_index_head_dim}." + "hybrid_mla_attention='mqa_dsa' runs the indexer " + "through the cuDNN (sparse phase) / tilelang (warmup " + "phase) kernels, which require index_head_dim=128, " + f"got {self.dsa_index_head_dim}." ) # ``index_n_heads`` is deliberately *not* pinned to 64 here. # The warmup phase's tilelang indexer does need exactly 64 @@ -1858,7 +1885,7 @@ def __post_init__(self): # enforcing it at config time would make every small-geometry # unit fixture unrepresentable. The check lives at the first # use instead -- - # ``MQALatentAttention._check_tilelang_full_candidate_support`` + # ``MHADSAWarmupAttention._check_tilelang_indexer_support`` # -- which still raises before any kernel launch. if self.dsa_indexer_use_sparse_loss and ( self.dsa_index_topk % 128 != 0 @@ -1953,7 +1980,7 @@ def __post_init__(self): "least one Indexer, and this config builds none, so there " "would be no trainable parameter left. Either a CSA layer " "(csa_dense_mode=False plus some 1 < csa_compress_ratios[i] " - "< 128) or a latent MQA layer with a DSA indexer " + "< 128) or a hybrid MLA layer with a DSA indexer " "(hybrid_mla_attention='mqa_dsa' plus some " "csa_compress_ratios[i] == -2). Got " f"csa_dense_mode={self.csa_dense_mode}, " diff --git a/tests/multi_card_tests/transformer/test_indexer_topk_col_mask_cp.py b/tests/multi_card_tests/transformer/test_indexer_topk_col_mask_cp.py index eba8c79d61..b77bb6fd89 100644 --- a/tests/multi_card_tests/transformer/test_indexer_topk_col_mask_cp.py +++ b/tests/multi_card_tests/transformer/test_indexer_topk_col_mask_cp.py @@ -24,12 +24,14 @@ *is* the ratio-causal limit (``csa_attention.get_valid_range``), and on the dense path ``shift_scores_to_local_window`` fills the tail with ``-inf`` (``docmask_utils.py:137-182``), so nothing finite sits past ``seq_lens``. -* hybrid MLA latent MQA + DSA: ``_indexer_valid_range`` deliberately clamps the - end ``csa_window_size`` *before* the diagonal, because the forced local window - already covers those tokens (``mqa_latent_attention.py:398-405``). On the THD - path the scores are the kernel's own document-causal ones, so columns in - ``[seq_lens, causal_len)`` are finite and ``topk`` will return them -- - duplicating the window and, at the diagonal, leaking the query's own column. +* hybrid MLA latent MQA + DSA (phase 3): ``_indexer_valid_range`` deliberately + clamps the end ``csa_window_size`` *before* the diagonal, because the forced + local window already covers those tokens + (``hybrid_mla_indexer.py:178-181``, called with ``window=self.window_size`` + at ``mqa_latent_attention.py:550-558``). On the THD path the scores are the + kernel's own document-causal ones, so columns in ``[seq_lens, causal_len)`` + are finite and ``topk`` will return them -- duplicating the window and, at + the diagonal, leaking the query's own column. Masking the columns before ``topk`` fixes the second caller. This file is the multi-card evidence for both halves of the claim, on the real kernels rather @@ -40,7 +42,10 @@ bitwise, because nothing finite lives past ``seq_lens``. * ``TestHybridMLAColumnMaskLoadBearingCP``: the latent MQA + DSA layer under CP -- without the mask the top-k *does* return out-of-window columns, so the - no-op above is not a vacuous property of the helper. + no-op above is not a vacuous property of the helper. Its phase-2 subtest is a + different core attention entirely (dense MHA, ``mha_dsa_warmup_attention``), + driven through ``test_mqa_dsa_warmup_cp.run_warmup_cp``, and asserts the + helper is never reached at all. No CSA CP test reached this helper before: ``test_csa_attention_cp.py``'s real-kernel classes all run ``csa_indexer_backend="unfused"``, and its one cuDNN @@ -68,6 +73,7 @@ class monkeypatches ``cudnn_indexer_topk_fwd`` away entirely. # same import style as ``test_mqa_dsa_warmup_cp`` reusing ``test_mqa_dsa_cp``. import test_csa_attention_cp as C import test_mqa_dsa_cp as H +import test_mqa_dsa_warmup_cp as W import paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn as M from paddlefleet.transformer.csa_attention import CSADocMaskMetadata @@ -88,9 +94,13 @@ class monkeypatches ``cudnn_indexer_topk_fwd`` away entirely. def setUpModule(): - # One fleet init for both harnesses; ``test_csa_attention_cp``'s builders - # only read its module globals, so mirroring them is enough. - H.setUpModule() + # One fleet init for all three harnesses; ``test_csa_attention_cp``'s + # builders only read its module globals, so mirroring them is enough. + # Going through ``W.setUpModule`` (which calls ``H.setUpModule`` and then + # clears ``parallel_state._CONTEXT_PARALLEL_GROUP``) is what keeps + # ``run_warmup_cp``'s CP=1 reference module out of the CP path -- the dense + # attention half reads that process-global group, not ``pg_collection.cp``. + W.setUpModule() C.CP_SIZE, C.CP_RANK, C.CP_GROUP = H.CP_SIZE, H.CP_RANK, H.CP_GROUP @@ -312,8 +322,8 @@ class TestHybridMLAColumnMaskLoadBearingCP(unittest.TestCase): """The same mask is what keeps latent MQA + DSA inside its window. ``H._STRADDLE`` tiles ``s_global`` exactly, so ``doc_lens`` is handed to the - kernel and the THD path runs (``mqa_latent_attention.py:574-589``) -- the one - layout where the scores past ``seq_lens`` are still finite. + kernel and the THD path runs (``mqa_latent_attention.py:574-596``) -- the + one layout where the scores past ``seq_lens`` are still finite. """ def _run(self, sparse_loss, loss_coeff): @@ -327,6 +337,24 @@ def _run(self, sparse_loss, loss_coeff): ) return spy + def _run_warmup(self, loss_coeff): + """Phase 2 on its own harness: a dense MHA layer, not latent MQA. + + ``H.run_core_cp`` can no longer drive this phase -- it builds an + ``MQALatentAttention`` and feeds the absorbed latent layout, while + ``dsa_indexer_use_sparse_loss=False`` now selects + ``MHADSAWarmupAttention`` (``hybrid_mla_indexer.py:59-60``). The + warmup suite's runner builds the right module and the per-head dense + inputs it wants. + """ + with _spying() as spy: + W.run_warmup_cp( + H._STRADDLE, + loss_coeff=loss_coeff, + with_input_ids=loss_coeff > 0, + ) + return spy + def _report(self, spy, tag): totals = { k: spy.group_total(k) @@ -385,19 +413,21 @@ def test_1_phase3_mask_is_load_bearing(self): def test_2_warmup_never_reaches_the_topk_helper(self): """Phase 2 never reaches this helper, so the column mask cannot apply. - ``_forward_warmup`` (``mqa_latent_attention.py:453-585``) supervises the - indexer over the **whole** per-document causal span: one *tilelang* - ``csa_indexer_topk_fwd`` with ``topk_effective=s_global`` and - ``window=0`` (``mqa_latent_attention.py:524-541``), imported directly - rather than through ``csa_indexer_backend``. So the cuDNN helper this - file spies on -- the one that carries the column mask -- is not called - on either the attention or the loss side, and with no forced window - there is no window duplication for a mask to prevent. Asserted rather - than dropped because it is the sharpest statement of the phase-2 - contract, and because it bounds ``test_1``'s claim: the mask is - load-bearing for phase 3 only. + The warmup phase is dense MHA plus an indexer + (``MHADSAWarmupAttention``): its attention half is + ``DotProductAttention.forward`` + (``mha_dsa_warmup_attention.py:179-198``), which selects no columns at + all, and its KL runs one *tilelang* ``csa_indexer_topk_fwd`` with + ``topk_effective=s_global`` and ``window=0`` + (``mha_dsa_warmup_attention.py:283-300``), imported directly rather + than through ``csa_indexer_backend``. So the cuDNN helper this file + spies on -- the one that carries the column mask -- is not called on + either side, and with no forced window there is no window duplication + for a mask to prevent. Asserted rather than dropped because it is the + sharpest statement of the phase-2 contract, and because it bounds + ``test_1``'s claim: the mask is load-bearing for phase 3 only. """ - spy = self._run(sparse_loss=False, loss_coeff=0.1) + spy = self._run_warmup(loss_coeff=0.1) totals = self._report(spy, "mqa_dsa/warmup") self.assertEqual( totals["rows"], diff --git a/tests/multi_card_tests/transformer/test_mqa_dsa_cp.py b/tests/multi_card_tests/transformer/test_mqa_dsa_cp.py index 4e955746f5..878d6e4c81 100644 --- a/tests/multi_card_tests/transformer/test_mqa_dsa_cp.py +++ b/tests/multi_card_tests/transformer/test_mqa_dsa_cp.py @@ -38,8 +38,9 @@ 4. The selected column set: the sparse kernel's ``token_indices`` under CP must equal the reference's rows for this rank. 5. Indexer-loss normalisation, masked (global denominator) and unmasked - (``/cp_size`` on the phase-3 path), and the phase-2 - ``dsa_indexer_use_sparse_loss=False`` full-causal KL. + (``/cp_size``), on the phase-3 sparse path. The phase-2 warmup path is a + different core attention (dense MHA, ``mha_dsa_warmup_attention.py``) and is + covered by ``test_mqa_dsa_warmup_cp.py``. 6. The attention sink under CP. 7. The ``cp_balance_mode`` guard. @@ -61,6 +62,7 @@ from paddle.distributed import fleet from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled _HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert( @@ -123,8 +125,9 @@ def setUpModule(): def _build(mode, cp_group, loss_coeff=0.0, sink=None, sparse_loss=True, seed=7): """CP=1 reference (``cp_group is None``) or CP layer, identical weights.""" - cfg = U._create_mqa_config(mode=mode, loss_coeff=loss_coeff) - cfg.dsa_indexer_use_sparse_loss = sparse_loss + cfg = U._create_mqa_config( + mode=mode, loss_coeff=loss_coeff, sparse_loss=sparse_loss + ) cfg.cp_balance_mode = "contiguous_allgather" # Production EB dataflow hands every rank the *global* input_ids, which is # what ``_indexer_loss_mask`` assumes when this flag is set @@ -132,6 +135,17 @@ def _build(mode, cp_group, loss_coeff=0.0, sink=None, sparse_loss=True, seed=7): cfg.experimental_dataflow = True cfg.pad_token_id = 0 cfg.context_parallel_size = 1 if cp_group is None else cp_group.nranks + # This harness feeds the *latent* layout (one 576-wide key head plus + # ``v_b_proj_weight``, ``value=None``) and reads the block-sparse kernel's + # index table out of ``U._CAPTURED``, so it can only drive latent MQA. The + # phase-2 warmup backend is dense per-head MHA and has its own runner, + # ``test_mqa_dsa_warmup_cp.run_warmup_cp``. + assert latent_mqa_enabled(cfg), ( + f"this harness cannot drive hybrid_mla_attention=" + f"{cfg.hybrid_mla_attention!r} with dsa_indexer_use_sparse_loss=" + f"{sparse_loss!r}: that selects the dense MHADSAWarmupAttention " + "(hybrid_mla_indexer.py:59-60) -- use test_mqa_dsa_warmup_cp instead" + ) paddle.seed(seed) return U._build_module( cfg, @@ -190,12 +204,12 @@ def _rel(a, e): def _logged_indexer_loss(layer_number=1): """The indexer loss this layer just pushed into the logging tracker. - ``_forward_warmup`` / ``_forward_sparse`` reduce the KL with the coefficient - and denominator their phase picked (phase 2 always the global row count; - phase 3 the global valid-row count when masked, ``/cp_size`` when not) and - hand exactly that scalar to ``DSAIndexerLossLoggingHelper``, so the tracker - is a direct read of the normalisation -- no gradient indirection. ``0.0`` - when the step attached no loss. + ``MQALatentAttention._forward_sparse`` reduces the KL with the coefficient + and denominator its phase picked (the global valid-row count when masked, + ``/cp_size`` when not) and hands exactly that scalar to + ``DSAIndexerLossLoggingHelper``, so the tracker is a direct read of the + normalisation -- no gradient indirection. ``0.0`` when the step attached no + loss. """ values = DSAIndexerLossLoggingHelper.tracker.get("values") if values is None: @@ -536,8 +550,13 @@ def test_7_indexer_loss_normalisation(self): to land on the CP=1 reference. A per-rank denominator is off by ``cp_size`` (or by the pad imbalance, which ``_input_ids`` puts on the last rank only) and fails here. + + Phase 2 (``sparse_loss=False``) is a different core attention entirely + and its own normalisation sweep lives in + ``test_mqa_dsa_warmup_cp.py::TestWarmupCP::test_4``, which additionally + reads the coefficient handed to the backward. """ - for masked, sparse in ((True, True), (False, True), (True, False)): + for masked, sparse in ((True, True), (False, True)): with self.subTest(masked=masked, sparse_loss=sparse): res = run_core_cp( "mqa_dsa", @@ -560,21 +579,27 @@ def test_8_cp_balance_mode_guard(self): The hybrid model's HCA layers assert the same mode (``dsv4_hybrid_attention.py:607-611``); the MQA layer must refuse the - others rather than silently attend against a permuted KV. + others rather than silently attend against a permuted KV. Swept over + ``dsa_indexer_use_sparse_loss`` because the guard now lives in the mixin + both phases share (``hybrid_mla_indexer.py:94-103``), so both core + attentions have to raise -- phase 2 from + ``MHADSAWarmupAttention.__init__``'s + ``_init_hybrid_mla_cp_state`` call (``mha_dsa_warmup_attention.py:128``). """ - cfg = U._create_mqa_config(mode="mqa_dsa") - cfg.context_parallel_size = CP_SIZE - for mode in ("p2p", "zigzag", None): - with self.subTest(cp_balance_mode=mode): - cfg.cp_balance_mode = mode - with self.assertRaises(NotImplementedError): - U._build_module( - cfg, - bf16=True, - pg_collection=types.SimpleNamespace( - tp=None, cp=CP_GROUP - ), - ) + for sparse_loss in (True, False): + cfg = U._create_mqa_config(mode="mqa_dsa", sparse_loss=sparse_loss) + cfg.context_parallel_size = CP_SIZE + for mode in ("p2p", "zigzag", None): + with self.subTest(cp_balance_mode=mode, sparse=sparse_loss): + cfg.cp_balance_mode = mode + with self.assertRaises(NotImplementedError): + U._build_module( + cfg, + bf16=True, + pg_collection=types.SimpleNamespace( + tp=None, cp=CP_GROUP + ), + ) if __name__ == "__main__": diff --git a/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py b/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py index ec0bcd4d8d..752bf3c29a 100644 --- a/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py +++ b/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py @@ -12,37 +12,52 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Context parallel for the DSA **warmup** phase, and for padded layouts. - -Two gaps in ``test_mqa_dsa_cp.py`` / ``test_mla_cp_contiguous_allgather.py``: - -1. ``hybrid_mla_attention="mqa_dsa"`` with ``dsa_indexer_use_sparse_loss=False`` - -- the phase-2 (DSA warmup) pairing, where attention consumes the full - per-document causal set while the indexer is trained by a KL over that same - full causal set (no top-k anywhere). The existing CP suites run ``True`` - everywhere except the one ``(masked=True, sparse=False)`` subtest of - ``test_mqa_dsa_cp.py::test_7``, which only observes parameter gradients. The - warmup mode takes its own branch (``mqa_latent_attention._forward_warmup``) - whose index table and causal mask are both built at ``s_global`` and - row-sliced, so it needs its own CP evidence: that the attention output is the - CP=1 reference, that the table really is the global one sliced, that the mode - is bit-identical to ``"mqa_full_causal"`` under CP, and that the full-causal - KL normalises across the CP group on both the masked and the unmasked branch - (``_indexer_loss_mask`` / the all-ones fallback, both of which divide by the - **global** row count so the per-rank losses simply add up). - -2. A layout with genuine **row-validity pad rows**. ``_STRADDLE`` sums to - exactly ``S_GLOBAL`` and ``U._row_end`` folds any trailing gap into one final - document, so ``is_valid`` has been all-``True`` in every CP test so far; the - pad-row path (all-``-1`` index row -> zero output, zero ``dq``) was only ever - audited on one card. ``[475] @ s=512`` puts 37 pad rows on the last rank - only, which is also the pad-imbalance the loss denominator has to survive. - -Everything reuses ``test_mqa_dsa_cp``'s harness (fleet init, CP globals, -``run_core_cp``, ``_check``, ``_check_index_sets``). No ``if rank == X`` -short-circuit exists in this file: every collective (``run_core_cp``'s -all-reduces, the all-gather inside the layer) is issued on all ranks and only -the assertions are rank-conditional. +"""Context parallel for the DSA **warmup** phase (phase 2), and for padded +layouts. + +Phase 2 is ``hybrid_mla_attention="mqa_dsa"`` with +``dsa_indexer_use_sparse_loss=False``: no top-k on either side. It is no longer +latent MQA -- ``hybrid_mla_indexer.latent_mqa_enabled`` returns False for it +(``hybrid_mla_indexer.py:59-60``), so the spec builds +:class:`MHADSAWarmupAttention`, which delegates its whole attention half to +``DotProductAttention.forward`` (``mha_dsa_warmup_attention.py:103``, +``:179-198``) and only adds the full-candidate indexer KL. There is no +``[b, s, s]`` index table and no block-sparse kernel call in this phase at all. + +That moves, but does not remove, the CP evidence this file owes: + +1. Two independent CP states have to be right at once. The attention half reads + the *process-global* ``paddlefleet.parallel_state`` group, at construction + (``dot_product_attention.py:155``) and at dispatch (``:397``, ``:565-570``); + the indexer half reads ``pg_collection.cp`` + (``hybrid_mla_indexer.py:88-108``). ``fleet.init`` sets only the former's + fleet counterpart, so the runner below toggles ``parallel_state`` explicitly + and the CP=1 reference is built *and* run with it off. +2. "The full-causal ``[b, s, s]`` table is the global build, row-sliced" + inverts into two claims: the per-rank output is **bit-identical** to the + phase-1 dense attention fed the same local slice + (``hybrid_mla_utils._build_phase1_dense_module``), and the *indexer's* + candidate set -- ``csa_indexer_topk_fwd``'s ``columns``, which are global + token ids -- is the whole per-document causal set built over the global + sequence and row-sliced (``hybrid_mla_indexer._indexer_valid_range``). +3. The full-candidate KL still has to normalise across the CP group: the masked + branch divides by the **global** valid-row count + (``hybrid_mla_indexer.py:215-220``) and the unmasked one folds ``/cp_size`` + into the coefficient handed to the *backward* + (``mha_dsa_warmup_attention.py:322-326``). +4. A layout with genuine row-validity pad rows. ``_STRADDLE`` sums to exactly + ``S_GLOBAL`` and ``U._row_end`` folds any trailing gap into one final + document, so ``is_valid`` is all-``True`` in every other CP test. + ``[475] @ s=512`` puts 37 pad rows on the last rank only, which is also the + pad imbalance the loss denominator has to survive. The phase-3 (latent, + block-sparse) pad-row cases stay here as well, on ``test_mqa_dsa_cp``'s + harness, so a pad-row failure can be attributed to a phase. + +Phase 3/4's own loss normalisation is swept in +``test_mqa_dsa_cp.py::TestMQADSACP::test_7``; this file owns phase 2. + +No ``if rank == X`` short-circuit exists in this file: every collective is +issued on all ranks and only the assertions are rank-conditional. Run (2 or 4 GPUs):: @@ -53,8 +68,11 @@ test_mqa_dsa_warmup_cp.py """ +import contextlib +import types import unittest +import numpy as np import paddle import paddle.distributed as dist @@ -63,8 +81,17 @@ # ``test_mla_cp_recompute`` importing ``test_mla_cp_contiguous_allgather``). import test_mqa_dsa_cp as H +from paddlefleet import parallel_state as ps +from paddlefleet.transformer import mha_dsa_warmup_attention as warmup_mod from paddlefleet.transformer.csa_attention import _derive_csa_doc_boundaries -from paddlefleet.transformer.mqa_latent_attention import MQALatentAttention +from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, +) + +U = H.U S_GLOBAL = H.S_GLOBAL _STRADDLE = H._STRADDLE @@ -77,15 +104,73 @@ _PAD_DOC_LEN = 475 _N_PAD = S_GLOBAL - _PAD_DOC_LEN +FWD_RTOL = H.FWD_RTOL +GRAD_RTOL = H.GRAD_RTOL + +# Floor for the warmup-vs-phase-1 ``dq`` comparison, used only when the phase-1 +# module's own two runs happen to agree exactly on this batch: one bf16 ulp at +# the scale of these gradients (2**-8 relative, |dq| ~ 1e-2 in this fixture). +# The measured spread between two runs of the same module is 3.052e-05 at +# s=512, so this floor is not what carries the assertion -- see +# ``_assert_phase1_identical``. +_DQ_ATOMIC_FLOOR = 1e-4 + # The logged indexer loss is a bf16-fed fp32 KL reduction, so the per-rank sum # lands a few 1e-4 off the single-rank value. A wrong denominator is off by # ``cp_size`` (100% at CP=2), so this bound is three orders of magnitude away # from the failure it has to catch. LOSS_RTOL = 5e-3 +# Positional arguments of ``TileLangCSAIndexerLossAutoScaler.apply`` as phase 2 +# calls it (``mha_dsa_warmup_attention.py:342-354``). +_ARG_LOSS_COEFF = 7 +_ARG_LOSS_MASK = 10 + def setUpModule(): H.setUpModule() + # The attention half is ``DotProductAttention``'s and takes its CP state + # from the *process-global* ``parallel_state`` group, at construction + # (``dot_product_attention.py:155``) and at dispatch (``:397``, ``:565``), + # not from ``pg_collection.cp`` the way the indexer half does + # (``hybrid_mla_indexer.py:88-108``). ``fleet.init`` does not set it, so + # keep it off by default -- the CP=1 reference must build *and* run non-CP + # -- and turn it on only around the CP modules. + ps._CONTEXT_PARALLEL_GROUP = None + + +@contextlib.contextmanager +def _cp_enabled(): + """Make ``paddlefleet.parallel_state`` report this test's CP group.""" + ps._CONTEXT_PARALLEL_GROUP = H.CP_GROUP + try: + yield + finally: + ps._CONTEXT_PARALLEL_GROUP = None + + +@contextlib.contextmanager +def _capture_loss_args(): + """Record every ``TileLangCSAIndexerLossAutoScaler.apply`` argument list. + + Phase 2 attaches its loss through ``mha_dsa_warmup_attention``'s own symbol + (``mha_dsa_warmup_attention.py:342``), so that is the one to patch; + ``mqa_latent_attention``'s is phase 3's. + """ + real = warmup_mod.TileLangCSAIndexerLossAutoScaler + calls = [] + + class _Spy: + @staticmethod + def apply(*args, **kwargs): + calls.append(args) + return real.apply(*args, **kwargs) + + warmup_mod.TileLangCSAIndexerLossAutoScaler = _Spy + try: + yield calls + finally: + warmup_mod.TileLangCSAIndexerLossAutoScaler = real def _pad_row_end(doc_len, s_global): @@ -110,19 +195,225 @@ def _maxabs(a, b): return float((a.cast("float32") - b.cast("float32")).abs().max()) +def _bit_equal(a, b): + """Bit equality of two bf16 tensors, tolerating matching ``NaN``s. + + A fully masked (pad) row is left to the flashmask kernel, which may return + ``NaN`` there; ``_maxabs`` would then read ``nan`` and compare false against + every bound, including ``== 0.0``. What this file asserts is that phase 2 + and phase 1 produce *the same bits*, ``NaN`` included. + """ + x = a.cast("float32").numpy() + y = b.cast("float32").numpy() + return bool(np.array_equal(x, y, equal_nan=True)) + + def _local_slice(): """``(row_offset, rows)`` of this CP rank's query rows.""" rows = S_GLOBAL // H.CP_SIZE return H.CP_RANK * rows, rows +def _warmup_cfg(cp_size, loss_coeff=0.0): + """Phase-2 config: ``mqa_dsa`` with the sparse-loss switch off.""" + cfg = U._create_mqa_config( + mode="mqa_dsa", loss_coeff=loss_coeff, sparse_loss=False + ) + cfg.cp_balance_mode = "contiguous_allgather" + # Production EB dataflow hands every rank the *global* ``input_ids``, which + # is the branch ``_indexer_loss_mask`` takes when this flag is set + # (``hybrid_mla_indexer.py:210-214``). + cfg.experimental_dataflow = True + cfg.pad_token_id = 0 + cfg.context_parallel_size = cp_size + assert not latent_mqa_enabled(cfg), ( + "mqa_dsa + dsa_indexer_use_sparse_loss=False must not select latent " + "MQA (hybrid_mla_indexer.py:59-60); this suite would silently be " + "testing phase 3 instead of the warmup phase" + ) + return cfg + + +def _build_warmup(cp_group, loss_coeff=0.0, seed=7): + """CP=1 reference (``cp_group is None``) or CP layer, identical weights. + + Must be called under :func:`_cp_enabled` for the CP layer and outside it for + the reference: ``pg_collection.cp`` drives the indexer half only. + """ + cfg = _warmup_cfg(1 if cp_group is None else cp_group.nranks, loss_coeff) + paddle.seed(seed) + module = U._build_module( + cfg, + bf16=True, + pg_collection=types.SimpleNamespace(tp=None, cp=cp_group), + ) + assert isinstance(module, MHADSAWarmupAttention), ( + f"the phase-2 fixture built {type(module).__name__}, not the dense " + "MHADSAWarmupAttention" + ) + return module + + +def _capture_columns(module, store): + """Record ``_dense_attn_target``'s ``columns`` (global token ids). + + ``RecordingWarmupMHA`` keeps the KL *target*, which is already permuted into + the kernel's column order, so the candidate ids themselves have to be taken + from the call's third positional argument + (``mha_dsa_warmup_attention.py:305-314``). + """ + inner = module._dense_attn_target + + def wrapper(*args, **kwargs): + store.append(args[2].numpy().copy()) + return inner(*args, **kwargs) + + module._dense_attn_target = wrapper + + +def _forward(module, q, k, v, row_end, x=None, qr=None, ids=None): + """One dense forward, per-document causal, ``row_end`` at global length. + + ``attn_mask_type`` has to be passed explicitly: on a ``None`` argument + ``DotProductAttention`` derives ``is_causal = attn_mask_type == + AttnMaskType.causal`` (``dot_product_attention.py:562``), i.e. non-causal. + The mask is global on both sides because + ``expand_attn_mask_startend_row_indices_for_cp`` builds its second column + from ``arange(s_local * cp_size)`` and expands it onto the mask's own rows + (``:319-342``), so a local-length table would not even broadcast; the CP + kernel then slices it per rank (``context_parallel_utils.py:1979-1992``). + + ``x`` left ``None`` selects the phase-1 module's narrower signature (it + takes no ``input_ids`` -- ``accepts_input_ids`` is the phase-2 mixin's, + ``hybrid_mla_indexer.py:76``). + """ + extra = {} if x is None else {"x": x, "qr": qr, "input_ids": ids} + return module( + q, + k, + v, + None, + attn_mask_startend_row_indices=row_end.clone(), + attn_mask_type=AttnMaskType.causal, + **extra, + ) + + +def run_warmup_cp( + doc_lens, + loss_coeff=0.0, + with_input_ids=False, + row_end=None, + s_global=S_GLOBAL, +): + """CP=1 reference, the CP layer, and the phase-1 dense module on one batch. + + Three modules rather than the harness' two: the phase-2 attention half is + ``DotProductAttention.forward`` verbatim + (``mha_dsa_warmup_attention.py:179-198``), so the phase-1 module fed this + rank's slice is the reference for *bit* equality, while the CP=1 phase-2 + module is the reference for CP equivalence. + """ + sl = s_global // H.CP_SIZE + off = H.CP_RANK * sl + if row_end is None: + row_end = U._row_end(doc_lens, s_global) + q, k, v, x, qr = U._make_dense_inputs(s_global, seed=3) + ids = H._input_ids(s_global) if with_input_ids else None + + ref = _build_warmup(None, loss_coeff) + with _cp_enabled(): + cpl = _build_warmup(H.CP_GROUP, loss_coeff) + phase1 = U._build_phase1_dense_module(_warmup_cfg(H.CP_SIZE), bf16=True) + cpl.set_state_dict(ref.state_dict()) + + U._CAPTURED.clear() + U._WARMUP_TARGETS.clear() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + cols_ref = [] + _capture_columns(ref, cols_ref) + ra = [H._leaf(t) for t in (q, k, v, x, qr)] + out_ref = _forward(ref, ra[0], ra[1], ra[2], row_end, ra[3], ra[4], ids) + out_ref.sum().backward() + logged_ref = H._logged_indexer_loss() + targets_ref = list(U._WARMUP_TARGETS) + + U._CAPTURED.clear() + U._WARMUP_TARGETS.clear() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + cols_cp = [] + _capture_columns(cpl, cols_cp) + cb = [H._leaf(t[:, off : off + sl]) for t in (q, k, v, x, qr)] + p1 = [H._leaf(t[:, off : off + sl]) for t in (q, k, v)] + # Second phase-1 run on fresh leaves, same module, same inputs: the + # kernel's own run-to-run spread, used as the bound for the warmup-vs-phase-1 + # ``dq`` comparison. Measured: the forward, ``dk`` and ``dv`` are bitwise + # reproducible, ``dq`` is not (3.052e-05 at s=512 between two runs of the + # *same* module), because the flashmask backward accumulates ``dq`` across + # column blocks with atomics. + p1c = [H._leaf(t[:, off : off + sl]) for t in (q, k, v)] + with _cp_enabled(), _capture_loss_args() as loss_args: + out = _forward(cpl, cb[0], cb[1], cb[2], row_end, cb[3], cb[4], ids) + out.sum().backward() + out_p1 = _forward(phase1, p1[0], p1[1], p1[2], row_end) + out_p1.sum().backward() + out_p1c = _forward(phase1, p1c[0], p1c[1], p1c[2], row_end) + out_p1c.sum().backward() + logged_cp = H._logged_indexer_loss() + + # Parameter grads: this rank only saw its own query rows, so the CP group's + # SUM is the reference. The attention half owns no parameter in this fixture + # (no sink is configured, so ``build_softmax_offset`` returns ``None``), + # which is also why the phase-1 module needs no ``set_state_dict``. + ref_named = dict(ref.named_parameters()) + param_err = {} + for name, p in cpl.named_parameters(): + if p.grad is None: + continue + g = p.grad.contiguous() + dist.all_reduce(g, group=H.CP_GROUP) + rg = ref_named[name].grad + param_err[name] = None if rg is None else H._rel(g, rg) + + return { + "fwd": H._rel(out, out_ref[:, off : off + sl]), + "dq": H._rel(cb[0].grad, ra[0].grad[:, off : off + sl]), + "dk": H._rel(cb[1].grad, ra[1].grad[:, off : off + sl]), + "dv": H._rel(cb[2].grad, ra[2].grad[:, off : off + sl]), + "param_err": param_err, + "out": out.detach(), + "ref_out": out_ref.detach(), + "dq_local": cb[0].grad.detach(), + # The phase-1 dense module on this rank's slice: the reference for the + # output (bitwise) and for ``dq`` (within the control spread below). + "phase1_out": out_p1.detach(), + "phase1_dq": p1[0].grad.detach(), + "phase1_out_control": out_p1c.detach(), + "phase1_dq_control": p1c[0].grad.detach(), + # ``columns`` are global token ids on both sides, so the reference's + # rows compare to this rank's directly, with no offset arithmetic. + "cols_ref_slice": cols_ref[-1][:, off : off + sl] if cols_ref else None, + "cols_cp": cols_cp[-1] if cols_cp else None, + "targets_ref": targets_ref, + "targets_cp": list(U._WARMUP_TARGETS), + # ``RecordingMQA``'s hook: any entry means a block-sparse kernel call + # happened, which phase 2 must never make. + "captured": len(U._CAPTURED), + "loss_args": list(loss_args), + "logged_ref": logged_ref, + "logged_cp": logged_cp, + "row_end": row_end, + "off": off, + "rows": sl, + } + + class _CPChecks(unittest.TestCase): - """Shared assertions, borrowed from the harness class. + """Shared assertions. - ``_check`` (forward + every gradient against the CP=1 reference) and - ``_check_index_sets`` (the sparse kernel is handed the reference's own - columns) are the same contract here, so call the harness' own - implementations rather than restating them. + ``_check`` / ``_check_index_sets`` are the harness' own implementations, + used by the latent (phase 3/4) pad-row cases; everything below them is the + dense phase-2 contract. """ def _check(self, res, tag): @@ -131,23 +422,207 @@ def _check(self, res, tag): def _check_index_sets(self, res, tag): H.TestMQADSACP._check_index_sets(self, res, tag) - def _assert_full_causal(self, idx, row_end, tag): - """Every local row's selected set == its whole per-document causal set. + def _check_dense(self, res, tag): + """CP=N == CP=1 on this rank's slice: output and every gradient.""" + for key, bound in ( + ("fwd", FWD_RTOL), + ("dq", GRAD_RTOL), + ("dk", GRAD_RTOL), + ("dv", GRAD_RTOL), + ): + self.assertLess(res[key], bound, f"{tag}: {key} {res[key]:.3e}") + for name, err in res["param_err"].items(): + self.assertIsNotNone( + err, f"{tag}: reference has no grad for {name}" + ) + self.assertLess(err, GRAD_RTOL, f"{tag}: param {name} {err:.3e}") + worst = max(res["param_err"].values(), default=0.0) + print( + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: " + f"fwd={res['fwd']:.2e} dq={res['dq']:.2e} dk={res['dk']:.2e} " + f"dv={res['dv']:.2e} param_max={worst:.2e}", + flush=True, + ) + + def _assert_no_block_sparse(self, res, tag): + """No ``[b, s, s]`` table, no block-sparse call -- the inverted claim. - This is what separates the warmup mode from phase 3: under - ``window + top-k`` a row longer than ``window + index_topk`` selects a - strict subset, so this assertion would fail there. + The old suite asserted phase 2 built the *same* full-causal index table + as ``mqa_full_causal``. Phase 2 no longer builds one at all, so the + evidence is that ``RecordingMQA``'s ``_sparse_attn`` hook never fired. + """ + self.assertEqual( + res["captured"], + 0, + f"{tag}: {res['captured']} block-sparse kernel call(s) were made " + "in the warmup phase, which must run dense flashmask only", + ) + + def _assert_phase1_identical(self, res, tag): + """Per-rank output == the phase-1 dense module's bitwise; ``dq`` == it to + within the kernel's own run-to-run spread. + + This is the other half of the inversion: "phase 2 is phase 1 plus an + indexer loss" is only true if the attention half is untouched, and the + indexer loss reaches the output solely through + ``TileLangCSAIndexerLossAutoScaler``'s *backward* + (``mha_dsa_warmup_attention.py:342-354``), so even a live loss must not + move the forward. Asserted with the CP row-slicing in the path. + + ``dq`` cannot be asserted bitwise, and the bound is measured rather than + chosen: two backward passes of the *same* phase-1 module on the same + inputs already disagree (3.052e-05 at s=512), because the flashmask + backward accumulates ``dq`` over column blocks with atomics. The + forward, ``dk`` and ``dv`` are bitwise reproducible, so those stay exact. + The control is computed in this same run + (``run_warmup_cp``: ``phase1_dq_control``), so a real divergence -- which + would be orders of magnitude larger, as the un-sliced-mask and + wrong-offset controls in this suite show -- still fails. + """ + same = _bit_equal(res["out"], res["phase1_out"]) + delta = _maxabs(res["out"], res["phase1_out"]) + print( + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: output vs " + f"phase-1 dense maxabs={delta:.3e} bit_equal={same}", + flush=True, + ) + self.assertTrue( + same, + f"{tag}: the warmup output is not bit-identical to the phase-1 " + f"dense module's (maxabs={delta:.3e}), so phase 2 is no longer " + "'phase 1 plus an indexer loss'", + ) + dq_delta = _maxabs(res["dq_local"], res["phase1_dq"]) + control = _maxabs(res["phase1_dq"], res["phase1_dq_control"]) + print( + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: dq vs phase-1 " + f"dense maxabs={dq_delta:.3e} (phase-1 self-control " + f"{control:.3e})", + flush=True, + ) + self.assertLessEqual( + dq_delta, + max(control, _DQ_ATOMIC_FLOOR), + f"{tag}: the warmup dq differs from the phase-1 dense module's by " + f"{dq_delta:.3e}, more than that module's own run-to-run spread " + f"({control:.3e}), so phase 2's attention backward is not phase " + "1's", + ) + + def _assert_full_causal_columns(self, res, tag): + """Every local row's candidate set == its whole per-document causal set. + + ``csa_indexer_topk_fwd`` runs in full-candidate mode here (``ratio=1``, + ``topk_effective=s_global``), fed a ``valid_range`` built over the + global sequence and row-sliced with ``window=0`` + (``mha_dsa_warmup_attention.py:283-300``, + ``hybrid_mla_indexer.py:151-191``). Under phase 3's window + top-k a row + longer than ``window + index_topk`` would select a strict subset, so + this assertion separates the two phases without needing phase 3's + columns for comparison. """ off, rows = _local_slice() - doc_start, is_valid = _doc_bounds(row_end, S_GLOBAL) + doc_start, is_valid = _doc_bounds(res["row_end"], S_GLOBAL) + cols = res["cols_cp"] + self.assertIsNotNone(cols, f"{tag}: the indexer produced no columns") + self.assertEqual(int(cols.shape[1]), rows, f"{tag}: row count") for r in range(rows): q = off + r - got = {int(c) for c in idx[0][r] if c >= 0} + got = {int(c) for c in cols[0][r] if c >= 0} want = set(range(doc_start[q], q + 1)) if is_valid[q] else set() self.assertEqual( got, want, f"{tag}: row {r} (global {q}) is not full-causal" ) + def _assert_loss_cp_sum(self, res, tag): + """Sum the per-rank logged loss and compare to the CP=1 value.""" + total = paddle.to_tensor([res["logged_cp"]], dtype="float64") + dist.all_reduce(total, group=H.CP_GROUP) + got, want = float(total[0]), res["logged_ref"] + self.assertGreater(abs(want), 0.0, f"{tag}: reference logged no loss") + rel = abs(got - want) / abs(want) + print( + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: " + f"this_rank={res['logged_cp']:.6e} sum(loss)={got:.6e} " + f"ref={want:.6e} rel={rel:.3e}", + flush=True, + ) + self.assertLess( + rel, + LOSS_RTOL, + f"{tag}: CP loss sum {got:.6e} != CP=1 {want:.6e} " + f"(rel {rel:.3e}); a per-rank denominator would be off by " + f"~{H.CP_SIZE}x", + ) + return want + + def _assert_loss_coeff(self, res, masked, loss_coeff, tag): + """The unmasked ``/cp_size`` must reach the *backward*, not just the log. + + ``csa_attention`` folds it into the logged scalar only; phase 2 puts it + in the coefficient it hands the autoscaler + (``mha_dsa_warmup_attention.py:322-326``), which is the value the + tilelang backward multiplies ``(P - Q)`` by. The logged loss cannot see + the difference -- both placements produce the same scalar -- so read the + argument itself. + """ + self.assertEqual(len(res["loss_args"]), 1, f"{tag}: one loss per layer") + args = res["loss_args"][0] + got = float(args[_ARG_LOSS_COEFF]) + want = loss_coeff if masked else loss_coeff / H.CP_SIZE + self.assertAlmostEqual( + got, + want, + places=9, + msg=f"{tag}: the backward got loss_coeff={got:.6e}, expected " + f"{want:.6e} (masked={masked}, cp_size={H.CP_SIZE})", + ) + mask = args[_ARG_LOSS_MASK] + if masked: + self.assertIsNotNone(mask, f"{tag}: input_ids built no row mask") + self.assertEqual( + list(mask.shape), [1, res["rows"]], f"{tag}: mask is not local" + ) + else: + self.assertIsNone(mask, f"{tag}: a mask appeared without input_ids") + + def _assert_columns_match_reference(self, res, tag): + """The CP layer's candidate set == the CP=1 layer's, for these rows. + + Set equality per row rather than sequence equality: the kernel's slot + order is its own business, and ``columns`` is what the KL target is + permuted into (``mha_dsa_warmup_attention.py:454-459``). + """ + want, got = res["cols_ref_slice"], res["cols_cp"] + self.assertIsNotNone(got, f"{tag}: the CP layer produced no columns") + self.assertIsNotNone(want, f"{tag}: the reference produced no columns") + self.assertEqual(list(got.shape), list(want.shape), f"{tag}: shape") + for r in range(int(got.shape[1])): + a = {int(c) for c in want[0][r] if c >= 0} + b = {int(c) for c in got[0][r] if c >= 0} + self.assertEqual( + b, + a, + f"{tag}: row {r} (global {res['off'] + r}) differs: " + f"missing={sorted(a - b)[:8]} extra={sorted(b - a)[:8]}", + ) + + def _rel_on_valid_rows(self, res): + """Relative error against the CP=1 reference, valid rows only. + + A fully masked pad row is the flashmask kernel's business and may come + back as ``NaN``, which would poison a whole-slice norm. Bit equality to + the phase-1 module (``_assert_phase1_identical``) is what covers those + rows here. + """ + off, rows = _local_slice() + _, is_valid = _doc_bounds(res["row_end"], S_GLOBAL) + keep = [r for r in range(rows) if is_valid[off + r]] + got = res["out"].cast("float32").numpy()[0][keep] + want = res["ref_out"].cast("float32").numpy()[0][off : off + rows][keep] + denom = max(float(np.linalg.norm(want)), 1e-12) + return float(np.linalg.norm(got - want)) / denom + class TestWarmupCP(_CPChecks): """``mqa_dsa`` + ``dsa_indexer_use_sparse_loss=False`` under CP.""" @@ -156,239 +631,196 @@ class TestWarmupCP(_CPChecks): def test_1_warmup_forward_equivalence(self): """CP=N == CP=1 on the warmup path, with and without a live loss. - ``dsa_indexer_loss_coeff == 0`` returns straight after the full-causal - attention and never touches the indexer projections; - ``> 0`` additionally runs the full-candidate indexer KL - (``csa_indexer_topk_fwd`` with ``topk_effective=s_global``) over the - whole causal set. Neither may perturb the attention output, so both are - checked against the same reference. + ``dsa_indexer_loss_coeff == 0`` returns straight after the dense + attention and never touches the indexer projections + (``mha_dsa_warmup_attention.py:199-200``); ``> 0`` additionally runs the + full-candidate indexer KL. Neither may perturb the attention output, so + both are checked against the same reference *and* against the phase-1 + dense module. """ for coeff in (0.0, 0.1): with self.subTest(loss_coeff=coeff): tag = f"warmup/coeff={coeff}" - res = H.run_core_cp( - "mqa_dsa", - _STRADDLE, - loss_coeff=coeff, - with_input_ids=coeff > 0, - sparse_loss=False, + res = run_warmup_cp( + _STRADDLE, loss_coeff=coeff, with_input_ids=coeff > 0 ) - self._check(res, tag) - self._check_index_sets(res, tag) - self._assert_full_causal(res["idx_cp"], res["row_end"], tag) + self._check_dense(res, tag) + self._assert_no_block_sparse(res, tag) + self._assert_phase1_identical(res, tag) + if coeff > 0: + self._assert_full_causal_columns(res, tag) @H.U._GPU - def test_2_warmup_token_indices_are_the_global_table_row_sliced(self): - """The kernel's table == the global build, sliced -- bitwise. - - A per-rank build (``s_local`` rows, ``doc_start`` sliced first) clips - each row at ``q - position_offset`` and drops the prefix owned by lower - ranks. The control at the end constructs exactly that and asserts it - differs, so this comparison cannot be vacuous on rank > 0. + def test_2_candidate_columns_are_the_global_set_row_sliced(self): + """The indexer's candidates == the global build, sliced. + + ``columns`` are global token ids, so this rank's rows must equal the + CP=1 reference's same rows with no rebasing. The control at the end + constructs the set a CP-unaware build would produce -- ``valid_range`` + derived from *local* coordinates, i.e. rows clipped at + ``q - position_offset`` with the prefix owned by lower ranks dropped -- + and asserts it differs, so the comparison cannot be vacuous on rank > 0. """ - off, rows = _local_slice() - res = H.run_core_cp("mqa_dsa", _STRADDLE, sparse_loss=False) - row_end = res["row_end"] - doc_start, _, is_valid, _, _ = _derive_csa_doc_boundaries( - row_end, S_GLOBAL - ) - want = MQALatentAttention._build_full_causal_indices( - 1, S_GLOBAL, doc_start, is_valid - )[:, off : off + rows] - got = paddle.to_tensor(res["idx_cp"]) - self.assertEqual(list(got.shape), list(want.shape), "table shape") - drift = int((got.cast("int64") != want.cast("int64")).sum()) - self.assertEqual( - drift, - 0, - f"{drift} of {int(got.numel())} slots differ from the global table", - ) + tag = "columns" + res = run_warmup_cp(_STRADDLE, loss_coeff=0.1, with_input_ids=True) + self._assert_columns_match_reference(res, tag) + self._assert_full_causal_columns(res, tag) - # Control: the same builder driven at the local length, with the - # document starts rebased (and clipped) into local coordinates -- the - # shape a CP-unaware implementation would produce. - local = MQALatentAttention._build_full_causal_indices( - 1, - rows, - paddle.clip(doc_start[off : off + rows] - off, min=0), - is_valid[off : off + rows], - ) + off, rows = _local_slice() + doc_start, is_valid = _doc_bounds(res["row_end"], S_GLOBAL) + widths = [ + len({int(c) for c in res["cols_cp"][0][r] if c >= 0}) + for r in range(rows) + ] print( - f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] token_indices drift vs " - f"global table = {drift}", + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] candidate widths: " + f"min={min(widths)} max={max(widths)}", flush=True, ) if H.CP_RANK == 0: return - self.assertGreater( - int((local.cast("int64") != want[:, :, :rows].cast("int64")).sum()), - 0, - "a per-rank index build was indistinguishable from the global " - "one, so this test is vacuous", + local = [ + len(range(max(doc_start[off + r] - off, 0), r + 1)) + if is_valid[off + r] + else 0 + for r in range(rows) + ] + self.assertNotEqual( + widths, + local, + "a local-coordinate candidate build was indistinguishable from " + "the global one, so this test is vacuous", ) @H.U._GPU - def test_3_warmup_equals_mqa_full_causal_under_cp(self): - """Warmup output == ``hybrid_mla_attention="mqa_full_causal"`` output. - - Both modes call ``_build_full_causal_indices`` and then the same sparse - kernel with the same inputs -- ``_forward_warmup`` reaches it *through* - ``_forward_full_causal`` -- so on each rank this must be *bitwise* equal, - not merely close. The single-card claim (maxabs 0.0) is asserted here - with the CP row-slicing in the path, and it is also what pins the - phase-2 attention set to the frozen backbone's phase-1 one. + def test_3_warmup_output_is_real_per_document_causal_attention(self): + """The non-vacuity control for the phase-1 bit-identity claim. + + ``_assert_phase1_identical`` compares two runs of the same code path, so + on its own it would also pass if both returned garbage (or zeros). Pin + this rank's rows to an independent fp32 per-document causal MHA + (``hybrid_mla_utils._dense_mha_reference``) over the *global* batch, at + the module's own scale: ``softmax_scale = 1/sqrt(k_channels)`` + (``dot_product_attention.py:210-215``, with ``k_channels=K_CHANNELS`` + from ``_build_phase1_dense_module``), which is what + ``_dense_attn_target`` uses for the KL target too + (``mha_dsa_warmup_attention.py:440``). """ - ref = H.run_core_cp("mqa", _STRADDLE) - for coeff in (0.0, 0.1): - with self.subTest(loss_coeff=coeff): - warm = H.run_core_cp( - "mqa_dsa", - _STRADDLE, - loss_coeff=coeff, - with_input_ids=coeff > 0, - sparse_loss=False, - ) - delta = _maxabs(warm["out"], ref["out"]) - print( - f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] warmup vs " - f"mqa_full_causal (coeff={coeff}): maxabs={delta:.3e}", - flush=True, - ) - self.assertEqual( - delta, - 0.0, - f"warmup output is not bit-identical to mqa_full_causal " - f"(coeff={coeff}, maxabs={delta:.3e})", - ) + off, rows = _local_slice() + res = run_warmup_cp(_STRADDLE, loss_coeff=0.1, with_input_ids=True) + q, k, v, _, _ = U._make_dense_inputs(S_GLOBAL, seed=3) + want = U._dense_mha_reference( + q, k, v, res["row_end"], U.K_CHANNELS**-0.5 + ) + err = H._rel(res["out"], want[:, off : off + rows]) + print( + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] vs fp32 dense MHA: " + f"rel={err:.3e}", + flush=True, + ) + self.assertLess( + err, + FWD_RTOL, + f"the warmup output is not per-document causal MHA (rel {err:.3e})", + ) @H.U._GPU def test_4_indexer_loss_cp_normalisation(self): - """The logged indexer loss must sum to the CP=1 value across the group. + """The logged loss must sum to the CP=1 value, on both mask branches. Read straight out of ``DSAIndexerLossLoggingHelper``, so it observes the - denominator itself rather than its shadow in the gradients. Phase 2 - (``sparse_loss=False``) has one formula on both branches: the row mask is - always a real tensor -- ``_indexer_loss_mask``'s global-count mask when - ``input_ids`` reached the layer, an all-ones mask over ``b * s_global`` - when it did not -- so every rank contributes its own rows to one global - mean and the per-rank losses are partial sums. Phase 3 - (``sparse_loss=True``) still keeps the two-branch form (global count when - masked, a local mean times ``1 / cp_size`` when not); both must land on - the CP=1 value. ``sparse_loss`` is swept so a failure can be attributed: - warmup-only means the full-causal KL broke it, both means the - pre-existing normalisation is wrong. - - Note this layout does not exercise the *width* difference between the two - ``sparse_loss`` values as sharply as it could: with documents of - 200/150/162 and a 128-wide forced window, most rows have few non-local - candidates. ``test_5`` covers the width itself. + denominator rather than its shadow in the gradients. Masked divides by + the **global** valid-row count (``hybrid_mla_indexer.py:215-220``); + unmasked takes the plain local mean with ``/cp_size`` folded into the + coefficient, which ``_assert_loss_coeff`` checks reaches the backward. """ - for sparse in (False, True): - for masked in (True, False): - with self.subTest(sparse_loss=sparse, masked=masked): - tag = f"loss/sparse={sparse}/masked={masked}" - res = H.run_core_cp( - "mqa_dsa", - _STRADDLE, - loss_coeff=0.1, - with_input_ids=masked, - sparse_loss=sparse, - ) - self.assertTrue( - any(n.startswith("indexer.") for n in res["param_err"]), - f"{tag}: the indexer received no gradient", - ) - self._check(res, tag) - self._assert_loss_cp_sum(res, tag) + for masked in (True, False): + with self.subTest(masked=masked): + tag = f"loss/masked={masked}" + res = run_warmup_cp( + _STRADDLE, loss_coeff=0.1, with_input_ids=masked + ) + self.assertTrue( + any(n.startswith("indexer.") for n in res["param_err"]), + f"{tag}: the indexer received no gradient", + ) + self._check_dense(res, tag) + self._assert_no_block_sparse(res, tag) + self._assert_loss_cp_sum(res, tag) + self._assert_loss_coeff(res, masked, 0.1, tag) @H.U._GPU - def test_5_widened_warmup_loss_table_under_cp(self): - """One 512-long document, so the two phases' KL supports differ widely. - - Phase 2's KL spans the whole per-document causal set (up to 512 columns - on this layout) while phase 3's spans window(128) + top-k(128), so the - two logged losses must differ -- that is the non-vacuity control for - ``test_4`` -- while each still normalises across the CP group. + def test_5_full_candidate_kl_is_wider_than_a_windowed_one(self): + """One 512-long document: the KL support is the whole causal row. + + The width itself is the phase difference, so assert it directly instead + of comparing against a phase-3 run: the widest local candidate row must + exceed ``window + index_topk``, which is all phase 3 can ever supervise, + and the target the KL is taken over must be ``s_global`` wide + (``mha_dsa_warmup_attention.py:391-417``). Taken as a MAX over the CP + group: on rank 0 at CP=2 the widest row is exactly ``s_local``, so a + per-rank assertion would be a false failure there. """ - logged = {} - for sparse in (False, True): - with self.subTest(sparse_loss=sparse): - tag = f"wide/sparse={sparse}" - res = H.run_core_cp( - "mqa_dsa", - [S_GLOBAL], - loss_coeff=0.1, - with_input_ids=True, - sparse_loss=sparse, - ) - self._check(res, tag) - logged[sparse] = self._assert_loss_cp_sum(res, tag) - spread = abs(logged[False] - logged[True]) / abs(logged[True]) - print( - f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] wide vs narrow KL: " - f"{logged[False]:.6e} vs {logged[True]:.6e} rel={spread:.3e}", - flush=True, + tag = "wide" + res = run_warmup_cp([S_GLOBAL], loss_coeff=0.1, with_input_ids=True) + self._check_dense(res, tag) + self._assert_loss_cp_sum(res, tag) + self._assert_full_causal_columns(res, tag) + self.assertEqual( + list(res["targets_cp"][-1].shape), + [1, res["rows"], S_GLOBAL], + f"{tag}: the KL target is not the full candidate width", ) - self.assertGreater( - spread, - 1e-3, - "the widened warmup KL table produced the same loss as the narrow " - f"one ({logged[False]:.6e} vs {logged[True]:.6e}), so " - "dsa_indexer_use_sparse_loss is not changing the table here and " - "test_4 is width-blind", + widest = max( + len({int(c) for c in res["cols_cp"][0][r] if c >= 0}) + for r in range(res["rows"]) ) - - def _assert_loss_cp_sum(self, res, tag): - """Sum the per-rank logged loss and compare to the CP=1 value.""" - total = paddle.to_tensor([res["logged_cp"]], dtype="float64") - dist.all_reduce(total, group=H.CP_GROUP) - got, want = float(total[0]), res["logged_ref"] - self.assertGreater(abs(want), 0.0, f"{tag}: reference logged no loss") - rel = abs(got - want) / abs(want) + group_max = paddle.to_tensor([widest], dtype="int64") + dist.all_reduce(group_max, group=H.CP_GROUP, op=dist.ReduceOp.MAX) + narrow = U.WINDOW + U.INDEX_TOPK print( - f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: " - f"this_rank={res['logged_cp']:.6e} sum(loss)={got:.6e} " - f"ref={want:.6e} rel={rel:.3e}", + f"[warmup-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: widest local " + f"row={widest} group_max={int(group_max[0])} narrow={narrow}", flush=True, ) - self.assertLess( - rel, - LOSS_RTOL, - f"{tag}: CP loss sum {got:.6e} != CP=1 {want:.6e} " - f"(rel {rel:.3e}); a per-rank denominator would be off by " - f"~{H.CP_SIZE}x", + self.assertGreater( + int(group_max[0]), + narrow, + f"{tag}: the widest candidate row across the CP group is " + f"{int(group_max[0])} <= window+index_topk ({narrow}), so this " + "layout cannot tell the full-candidate KL from a windowed one", ) - return want class TestPadRowsCP(_CPChecks): """``[475] @ s=512``: 37 real pad rows, all on the last CP rank. - Every CP suite so far used layouts whose documents tile the sequence, so - ``is_valid`` was all-``True`` and the pad-row path -- an all-``-1`` index - row, which the kernel must turn into a zero output row with a zero ``dq`` - -- was never reached under CP. The rows also sit entirely on the last rank, - which is the pad imbalance a per-rank loss denominator cannot survive. + Every other CP suite uses layouts whose documents tile the sequence, so + ``is_valid`` is all-``True`` and the pad-row path is never reached under CP. + The rows also sit entirely on the last rank, which is the pad imbalance a + per-rank loss denominator cannot survive. Both phases are covered so a + pad-row failure can be attributed to one: the latent (phase 3/4) cases run + on the harness' own runner, the dense warmup case on this file's. """ - def _run(self, mode, sparse_loss): - row_end = _pad_row_end(_PAD_DOC_LEN, S_GLOBAL) + def _run_latent(self, mode, sparse_loss): return H.run_core_cp( mode, None, loss_coeff=0.1, with_input_ids=True, sparse_loss=sparse_loss, - row_end=row_end, + row_end=_pad_row_end(_PAD_DOC_LEN, S_GLOBAL), ) - def _check_pad(self, res, tag): + def _local_pad_rows(self, row_end, tag): + """This rank's pad rows, after asserting the group really has ``_N_PAD``. + + Only the last rank owns them, so the count is checked on the group. + """ off, rows = _local_slice() - _, is_valid = _doc_bounds(res["row_end"], S_GLOBAL) + _, is_valid = _doc_bounds(row_end, S_GLOBAL) local_pad = [r for r in range(rows) if not is_valid[off + r]] - - # The layout must actually produce pad rows *somewhere*: assert it on - # the group, not on this rank, since only the last rank owns them. n_pad = paddle.to_tensor([len(local_pad)], dtype="int64") dist.all_reduce(n_pad, group=H.CP_GROUP) self.assertEqual( @@ -397,7 +829,11 @@ def _check_pad(self, res, tag): f"{tag}: the layout produced {int(n_pad[0])} pad rows, expected " f"{_N_PAD} -- the fixture no longer tests what it claims", ) + return off, local_pad + def _check_pad_latent(self, res, tag): + """Phase 3/4: an all-``-1`` index row must give a zero output and dq.""" + off, local_pad = self._local_pad_rows(res["row_end"], tag) out = res["out"].cast("float32") dq = res["dq_local"].cast("float32") for r in local_pad: @@ -419,37 +855,73 @@ def _check_pad(self, res, tag): ) print( f"[padrows-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: " - f"local_pad_rows={len(local_pad)} fwd={res['fwd']:.2e} " - f"per_pos_max={max(res['per_pos']):.3e}", + f"local_pad_rows={len(local_pad)} fwd={res['fwd']:.2e}", flush=True, ) @H.U._GPU def test_1_pad_rows_mqa_full_causal(self): - res = self._run("mqa", True) + res = self._run_latent("mqa", True) self._check(res, "pad/mqa_full_causal") self._check_index_sets(res, "pad/mqa_full_causal") - self._check_pad(res, "pad/mqa_full_causal") + self._check_pad_latent(res, "pad/mqa_full_causal") @H.U._GPU - def test_2_pad_rows_warmup(self): - res = self._run("mqa_dsa", False) - self._check(res, "pad/warmup") - self._check_index_sets(res, "pad/warmup") - self._assert_full_causal(res["idx_cp"], res["row_end"], "pad/warmup") - self._check_pad(res, "pad/warmup") + def test_2_pad_rows_warmup_dense(self): + """Phase 2's pad rows: whatever phase 1 does, plus an empty KL row. + + The output of a fully masked row is the dense flashmask kernel's own + behaviour, not this phase's, so it is pinned by bit equality to the + phase-1 module rather than by a value; the ``NaN``-tolerant comparison + exists for exactly that reason. What phase 2 *does* own is the indexer + half: a pad row must carry no candidate and no KL mass + (``mha_dsa_warmup_attention.py:301-304``, ``:450``), and the loss + denominator must still be the global valid-row count with 37 of the pad + rows on one rank. + """ + tag = "pad/warmup" + row_end = _pad_row_end(_PAD_DOC_LEN, S_GLOBAL) + res = run_warmup_cp( + None, loss_coeff=0.1, with_input_ids=True, row_end=row_end + ) + off, local_pad = self._local_pad_rows(row_end, tag) + self._assert_no_block_sparse(res, tag) + self._assert_phase1_identical(res, tag) + self._assert_full_causal_columns(res, tag) + + target = res["targets_cp"][-1] + for r in local_pad: + self.assertEqual( + int((res["cols_cp"][0][r] >= 0).sum()), + 0, + f"{tag}: pad row {r} (global {off + r}) has candidates", + ) + self.assertEqual( + float(np.abs(target[0][r]).max()), + 0.0, + f"{tag}: pad row {r} (global {off + r}) carries KL mass", + ) + + err = self._rel_on_valid_rows(res) + self.assertLess(err, FWD_RTOL, f"{tag}: valid-row forward {err:.3e}") + for name, perr in res["param_err"].items(): + self.assertIsNotNone(perr, f"{tag}: no reference grad for {name}") + self.assertLess(perr, GRAD_RTOL, f"{tag}: param {name} {perr:.3e}") + self._assert_loss_cp_sum(res, tag) + self._assert_loss_coeff(res, True, 0.1, tag) + print( + f"[padrows-cp{H.CP_SIZE} rank{H.CP_RANK}] {tag}: " + f"local_pad_rows={len(local_pad)} valid_row_fwd={err:.2e}", + flush=True, + ) @H.U._GPU def test_3_pad_rows_sparse(self): - """Same layout on the phase-3 (``window + top-k``) path. - - Included so a pad-row failure can be attributed: if it reproduces here - it is not specific to the warmup branch this change introduced. - """ - res = self._run("mqa_dsa", True) + """Same layout on the phase-3 (``window + top-k``) path.""" + res = self._run_latent("mqa_dsa", True) self._check(res, "pad/sparse") self._check_index_sets(res, "pad/sparse") - self._check_pad(res, "pad/sparse") + self._check_pad_latent(res, "pad/sparse") if __name__ == "__main__": diff --git a/tests/single_card_tests/transformer/hybrid_mla_utils.py b/tests/single_card_tests/transformer/hybrid_mla_utils.py index 3dbeb58a7b..8f7faeb185 100644 --- a/tests/single_card_tests/transformer/hybrid_mla_utils.py +++ b/tests/single_card_tests/transformer/hybrid_mla_utils.py @@ -41,11 +41,17 @@ from paddle.distributed.fleet.meta_parallel import LayerSpec from paddlefleet.transformer.csa_attention import _derive_csa_doc_boundaries +from paddlefleet.transformer.dot_product_attention import DotProductAttention from paddlefleet.transformer.dsa_attention import ( DSAIndexer, DSAIndexerSublayersSpec, ) from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, + MHADSAWarmupAttentionSublayersSpec, +) from paddlefleet.transformer.mqa_latent_attention import ( MQALatentAttention, MQALatentAttentionSublayersSpec, @@ -143,17 +149,26 @@ def _dsa_kernels_available(): } -def _create_mqa_config(mode="mqa", loss_coeff=0.0, num_hidden_layers=2): +def _create_mqa_config( + mode="mqa", loss_coeff=0.0, num_hidden_layers=2, sparse_loss=True +): """dsv4_hybrid config for a ``csa_compress_ratios == -2`` layer. ``mode`` is a test-only fixture label mapped onto the production ``hybrid_mla_attention`` enum by ``_HYBRID_MLA_ATTENTION``: ``"mha"`` keeps the dense per-head path, ``"mqa"`` selects ``"mqa_full_causal"`` (latent MQA over the full per-document causal set, no indexer) and ``"mqa_dsa"`` selects - ``"mqa_dsa"`` (latent MQA + DSA indexer). Whether an indexer actually exists - is expressed by ``_build_module`` attaching one to the sublayers spec, - mirroring the production source which reads the layer path from the spec, - not from a config string. + ``"mqa_dsa"`` (DSA indexer). Whether an indexer actually exists is expressed + by ``_build_module`` attaching one to the sublayers spec, mirroring the + production source which reads the layer path from the spec, not from a + config string. + + ``sparse_loss`` is ``dsa_indexer_use_sparse_loss``, i.e. the phase switch of + ``"mqa_dsa"``. It picks the *attention backend* as well as the loss width, so + ``_build_module`` dispatches on it through the production predicate + ``hybrid_mla_indexer.latent_mqa_enabled``: ``True`` (phase 3/4) builds latent + MQA, ``False`` (phase 2, DSA warmup) builds the dense per-head + ``MHADSAWarmupAttention``. Attributes are assigned after construction so that ``__post_init__`` validation (exercised by the production model config, not by this unit) is @@ -188,7 +203,7 @@ def _create_mqa_config(mode="mqa", loss_coeff=0.0, num_hidden_layers=2): config.dsa_index_topk = INDEX_TOPK config.csa_window_size = WINDOW config.dsa_indexer_loss_coeff = loss_coeff - config.dsa_indexer_use_sparse_loss = True + config.dsa_indexer_use_sparse_loss = sparse_loss config.dsa_indexer_rotary_interleaved = False # The -2 layers are uncompressed, hence plain RoPE (base 10000); YaRN only # applies to the compressed HCA layers. @@ -222,6 +237,39 @@ def _sparse_attn( ) +class RecordingWarmupMHA(MHADSAWarmupAttention): + """Phase-2 counterpart of :class:`RecordingMQA`. + + There is no index table to capture: the whole point of the phase-2 backend + is that no ``[b, s, s]`` table and no block-sparse call exist. So the hook + records the KL target instead (into ``_WARMUP_TARGETS``), and ``_CAPTURED`` + staying empty across a phase-2 forward *is* the assertion that the + block-sparse kernel was never reached. + """ + + def _dense_attn_target(self, *args, **kwargs): + target = super()._dense_attn_target(*args, **kwargs) + _WARMUP_TARGETS.append(target.astype("float32").numpy().copy()) + return target + + +_WARMUP_TARGETS = [] + + +def _indexer_layer_spec(): + """``LayerSpec`` for the hybrid-MLA ``DSAIndexer``, shared by both phases.""" + return LayerSpec( + layer=DSAIndexer, + sublayers_spec=DSAIndexerSublayersSpec( + linear_wq_b=BiasedLinear, + linear_wk=BiasedLinear, + k_norm=LayerNormStub, + linear_weights_proj=BiasedLinear, + ), + extra_kwargs={"is_hybrid_mla_indexer": True}, + ) + + def _build_module( config, layer_number=1, @@ -230,7 +278,12 @@ def _build_module( is_mtp=False, pg_collection=None, ): - """Build a ``RecordingMQA``. + """Build the core attention this config selects. + + ``RecordingWarmupMHA`` (dense per-head, phase 2) when the config has an + indexer but ``latent_mqa_enabled`` is False, ``RecordingMQA`` (latent MQA) + otherwise. The dispatch is the production predicate itself, so a fixture can + never build a backend the model would not. ``pg_collection`` must be passed explicitly by the context-parallel tests (``tests/multi_card_tests/transformer/test_mqa_dsa_cp.py``): left ``None`` @@ -238,43 +291,82 @@ def _build_module( which inside a ``fleet.init``-ed process would hand the CP=1 reference the real CP group. """ - indexer = None - if getattr(config, "_build_dsa_indexer", False): - indexer = LayerSpec( - layer=DSAIndexer, - sublayers_spec=DSAIndexerSublayersSpec( - linear_wq_b=BiasedLinear, - linear_wk=BiasedLinear, - k_norm=LayerNormStub, - linear_weights_proj=BiasedLinear, + has_indexer = bool(getattr(config, "_build_dsa_indexer", False)) + indexer = _indexer_layer_spec() if has_indexer else None + common = { + "config": config, + "layer_number": layer_number, + "attn_mask_type": AttnMaskType.causal, + "attention_type": "self", + "k_channels": K_CHANNELS, + "is_mtp_layer": is_mtp, + "pg_collection": pg_collection, + } + if has_indexer and not latent_mqa_enabled(config): + module = RecordingWarmupMHA( + sublayers_spec=MHADSAWarmupAttentionSublayersSpec(indexer=indexer), + # ``MLASelfAttention`` passes these three to every core attention + # (multi_latent_attention.py:551-554); the latent-MQA class derives + # them from the config instead, which is why only this branch needs + # them. + v_channels=V_HEAD_DIM, + num_attention_heads=H, + num_key_value_heads=1, + **common, + ) + else: + module = RecordingMQA( + sublayers_spec=MQALatentAttentionSublayersSpec(indexer=indexer), + **common, + ) + if bf16: + # ``rotate_activation`` asserts bf16 inputs, so the indexer projections + # must hold bf16 weights. + module.to(dtype="bfloat16") + if sink is not None: + # In production both core attentions build this parameter via the shared + # ``build_softmax_offset`` helper (name ``core_attention.softmax_offset`` + # in every phase, so an MHA checkpoint stays loadable). This unit uses a + # default config with no sink configured, so the helper returns ``None``; + # inject the sink here instead. Created *after* ``to(dtype=...)`` and in + # the module dtype: production uses ``params_dtype`` (bf16), which is + # what the FA4 cute kernel of the dense path requires, and the DSA path + # returns the sink gradient in the parameter's own dtype. + module.softmax_offset = module.create_parameter( + shape=[H], + dtype="bfloat16" if bf16 else "float32", + default_initializer=paddle.nn.initializer.Assign( + np.asarray(sink, dtype="float32") ), - extra_kwargs={"is_hybrid_mla_indexer": True}, ) - module = RecordingMQA( + return module + + +def _build_phase1_dense_module( + config, layer_number=1, bf16=False, sink=None, is_mtp=False +): + """A plain :class:`DotProductAttention`, i.e. the phase-1 attention. + + The reference the phase-2 backend is asserted *bit-identical* to: phase 2 + delegates its whole attention half to ``super().forward``, so anything but + bit equality here means the warmup phase is no longer "phase 1 plus an + indexer loss". Built with exactly the kwargs ``MLASelfAttention`` passes + (multi_latent_attention.py:542-557). + """ + module = DotProductAttention( config=config, - sublayers_spec=MQALatentAttentionSublayersSpec(indexer=indexer), layer_number=layer_number, attn_mask_type=AttnMaskType.causal, attention_type="self", - k_channels=K_CHANNELS, is_mtp_layer=is_mtp, - pg_collection=pg_collection, + k_channels=K_CHANNELS, + v_channels=V_HEAD_DIM, + num_attention_heads=H, + num_key_value_heads=1, ) if bf16: - # ``rotate_activation`` asserts bf16 inputs, so the indexer projections - # must hold bf16 weights. module.to(dtype="bfloat16") if sink is not None: - # In production ``MQALatentAttention.__init__`` builds this parameter - # via the shared ``build_softmax_offset`` helper (name - # ``core_attention.softmax_offset``, identical to the dense - # ``DotProductAttention`` phase, so an MHA checkpoint stays loadable). - # This unit uses a default config with no sink configured, so the - # helper returns ``None``; inject the sink here instead. Created *after* - # ``to(dtype=...)`` and in the module dtype: production uses - # ``params_dtype`` (bf16), which is what the FA4 cute kernel of the - # dense path requires, and the DSA path returns the sink gradient in - # the parameter's own dtype. module.softmax_offset = module.create_parameter( shape=[H], dtype="bfloat16" if bf16 else "float32", @@ -328,6 +420,60 @@ def _rel(actual, expected): return float((a - e).norm() / e.norm().clip(min=1e-12)) +def _make_dense_inputs(seqlen, seed=0): + """``(query, key, value, x, qr)`` in the **per-head** phase-1/2 layout. + + ``_make_inputs`` above produces the absorbed latent layout (one 576-wide + key head plus the de-absorption weight), which the dense backend cannot + consume: ``kv_b_proj`` has already materialised per-head K/V by the time it + is called, so it sees ``[b, s, H, K_CHANNELS]`` queries and keys and a + ``[b, s, H, V_HEAD_DIM]`` value. ``x`` / ``qr`` are the indexer's inputs and + are identical in both layouts. + """ + paddle.seed(seed) + query = (paddle.randn([1, seqlen, H, K_CHANNELS]) * 0.5).cast("bfloat16") + key = (paddle.randn([1, seqlen, H, K_CHANNELS]) * 0.5).cast("bfloat16") + value = (paddle.randn([1, seqlen, H, V_HEAD_DIM]) * 0.5).cast("bfloat16") + x = (paddle.randn([1, seqlen, HIDDEN]) * 0.5).cast("bfloat16") + qr = (paddle.randn([1, seqlen, Q_LORA]) * 0.5).cast("bfloat16") + return query, key, value, x, qr + + +def _dense_mha_reference(query, key, value, row_end, scale, sink=None): + """fp32 per-document causal MHA over the per-head layout. + + The dense counterpart of ``_dense_reference``: same masking and same sink + semantics (one value-less softmax column per head), but on real per-head K/V + instead of the shared latent. Returns ``[1, s, H * V_HEAD_DIM]``. + """ + seqlen = int(query.shape[1]) + doc_start, is_valid = _doc_meta(row_end, seqlen) + pos = np.arange(seqlen) + allowed = ( + (pos[None, :] <= pos[:, None]) + & (pos[None, :] >= doc_start[:, None]) + & is_valid[:, None] + ) + q = query[0].cast("float32") + k = key[0].cast("float32") + v = value[0].cast("float32") + scores = paddle.einsum("shd,thd->sht", q, k) * scale + keep = paddle.to_tensor(allowed).unsqueeze(1) + scores = paddle.where(keep, scores, paddle.full_like(scores, -1e30)) + if sink is None: + probs = F.softmax(scores, axis=-1) + else: + sink_col = paddle.to_tensor(np.asarray(sink, dtype="float32")).reshape( + [1, H, 1] + ) + sink_col = paddle.expand(sink_col, [seqlen, H, 1]) + probs = F.softmax(paddle.concat([scores, sink_col], axis=-1), axis=-1) + probs = probs[:, :, :seqlen] + out = paddle.einsum("sht,thv->shv", probs, v) + row_ok = paddle.to_tensor(is_valid).cast("float32").reshape([seqlen, 1, 1]) + return (out * row_ok).reshape([1, seqlen, H * V_HEAD_DIM]) + + def _dense_reference(query, key, w_v, row_end, scale, sink=None): """Per-document full-causal attention on the latent, computed in fp32. diff --git a/tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py b/tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py index 0e9bb0a359..0c4d7847ac 100644 --- a/tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py +++ b/tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py @@ -14,6 +14,7 @@ import unittest from functools import wraps +from types import SimpleNamespace from unittest.mock import MagicMock, patch import numpy as np @@ -97,6 +98,7 @@ def wrapper(*args, **kwargs): build_document_rope_freqs, ) from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled from paddlefleet.transformer.mqa_latent_attention import ( MQALatentAttention, ) @@ -3166,42 +3168,75 @@ def test_rejects_unfoldable_dims_hidden_behind_a_singleton(self): class TestHybridMLAAttentionSinkParameter(unittest.TestCase): - """``add_full_attention_sink_bias`` must give the same parameter in both - hybrid MLA phases. + """``add_full_attention_sink_bias`` must give the same parameter in all + three hybrid MLA phases. The sink is created by ``build_softmax_offset`` *on the core attention*, so - its state_dict name is ``core_attention.softmax_offset`` for the dense MHA - phase (``hybrid_mla_attention="mha"``, where ``DotProductAttention`` consumes - it as the FA4 ``learnable_sink``) as well as for the latent MQA + DSA - indexer phase (``hybrid_mla_attention="mqa_dsa"``, where - ``MQALatentAttention`` feeds it to the block-sparse kernel). Keeping one - name, one shape and one dtype is what lets an MHA checkpoint load into a - latent-MQA run unchanged. + its state_dict name is ``core_attention.softmax_offset`` for all of them: + phase 1 dense MHA (``"mha"`` -> ``DotProductAttention``, which consumes it as + the FA4 ``learnable_sink``), phase 2 DSA warmup (``"mqa_dsa"`` + + ``dsa_indexer_use_sparse_loss=False`` -> ``MHADSAWarmupAttention``, i.e. the + same dense consumer plus an indexer) and phase 3 (``"mqa_dsa"`` + + ``dsa_indexer_use_sparse_loss=True`` -> ``MQALatentAttention``, which feeds + it to the block-sparse kernel). Keeping one name, one shape and one dtype is + what lets an MHA checkpoint load into a latent-MQA run unchanged. + + WAS: two phases, ``_PHASES = ("mha", "mqa_dsa")``, and the dense-only FA4 / + bf16 construction guards were asserted to NOT apply to ``"mqa_dsa"``. NO + LONGER TRUE for phase 2: it is dense now, so it inherits both guards on + purpose ("phase 2 runs wherever phase 1 runs"). The guards are pinned on the + latent phase (``sparse_loss=True``) instead, and phase 2 is asserted to be + gated exactly like phase 1. """ - def _build(self, hybrid_mla_attention, sink): + # (hybrid_mla_attention, dsa_indexer_use_sparse_loss) per phase, in + # training order. ``"mqa_full_causal"`` (indexer-less latent MQA, equivalence + # experiments only) shares phase 3's sink layout and is not repeated here. + _MHA = ("mha", False) + _WARMUP = ("mqa_dsa", False) + _SPARSE = ("mqa_dsa", True) + _PHASES = (_MHA, _WARMUP, _SPARSE) + + @staticmethod + def _is_latent(hybrid_mla_attention, sparse_loss): + """The production dispatch predicate, not a copy of its rule.""" + return latent_mqa_enabled( + SimpleNamespace( + experimental_attention_variant="dsv4_hybrid", + hybrid_mla_attention=hybrid_mla_attention, + dsa_indexer_use_sparse_loss=sparse_loss, + ) + ) + + def _build(self, hybrid_mla_attention, sparse_loss, sink): # ``MultiLatentAttention.__init__`` refuses to create the sink for the - # dense MHA phase unless ``FLAGS_flash_attn_version in (3, 4)``: that - # phase consumes it as ``flashmask_attention_func(learnable_sink=...)``, - # which only exists on the cute path. The image default is 2, so flip - # the flag for the construction and restore it -- these tests only - # inspect the parameter, and the process-global default must stay - # untouched for the other suites in this file. The absorbed phase owns - # its sink inside the block-sparse kernel and needs no flag. - needs_fa4 = sink and hybrid_mla_attention == "mha" + # dense phases unless ``FLAGS_flash_attn_version in (3, 4)``: they + # consume it as ``flashmask_attention_func(learnable_sink=...)``, which + # only exists on the cute path. The image default is 2, so flip the flag + # for the construction and restore it -- these tests only inspect the + # parameter, and the process-global default must stay untouched for the + # other suites in this file. The latent phase owns its sink inside the + # block-sparse kernel and needs no flag. + needs_fa4 = sink and not self._is_latent( + hybrid_mla_attention, sparse_loss + ) previous = paddle.get_flags(["FLAGS_flash_attn_version"])[ "FLAGS_flash_attn_version" ] if needs_fa4: paddle.set_flags({"FLAGS_flash_attn_version": 4}) try: - return self._build_raw(hybrid_mla_attention, sink) + return self._build_raw(hybrid_mla_attention, sparse_loss, sink) finally: if needs_fa4: paddle.set_flags({"FLAGS_flash_attn_version": previous}) def _build_raw( - self, hybrid_mla_attention, sink, params_dtype=paddle.bfloat16 + self, + hybrid_mla_attention, + sparse_loss, + sink, + params_dtype=paddle.bfloat16, ): """Build without touching ``FLAGS_flash_attn_version``.""" model_parallel_cuda_manual_seed(_SEED) @@ -3211,6 +3246,7 @@ def _build_raw( params_dtype=params_dtype, csa_compress_ratios=[-2, 128], hybrid_mla_attention=hybrid_mla_attention, + dsa_indexer_use_sparse_loss=sparse_loss, add_full_attention_sink_bias=sink, # The -2 layer's DSA indexer (built only for # hybrid_mla_attention="mqa_dsa") reads these model-wide fields; the @@ -3241,29 +3277,40 @@ def _sink_keys(module): if name.endswith("softmax_offset") ) - # The two hybrid MLA phases: dense MHA (``"mha"``) and latent MQA + DSA - # indexer (``"mqa_dsa"``). The indexer-less latent MQA mode - # (``"mqa_full_causal"``) exists only for equivalence experiments and shares - # the ``"mqa_dsa"`` sink layout, so it is not a separate phase here. - _PHASES = ("mha", "mqa_dsa") + def test_core_attention_class_per_phase(self): + """The three phases really are three different core-attention classes. + + Without this the rest of the class could pass while every phase built + the same module. + """ + expected = { + self._MHA: "DotProductAttention", + self._WARMUP: "MHADSAWarmupAttention", + self._SPARSE: "MQALatentAttention", + } + for phase, core in expected.items(): + with self.subTest(phase=phase): + mla = self._build(*phase, sink=True) + self.assertEqual(type(mla.core_attention).__name__, core) + self.assertIs(mla.mqa_latent, self._is_latent(*phase)) def test_disabled_creates_no_parameter(self): - for hybrid_mla_attention in self._PHASES: - with self.subTest(hybrid_mla_attention=hybrid_mla_attention): - mla = self._build(hybrid_mla_attention, sink=False) + for phase in self._PHASES: + with self.subTest(phase=phase): + mla = self._build(*phase, sink=False) self.assertIsNone(mla.core_attention.softmax_offset) self.assertEqual(self._sink_keys(mla), []) def test_enabled_creates_one_zero_initialised_per_head_parameter(self): - for hybrid_mla_attention in self._PHASES: - with self.subTest(hybrid_mla_attention=hybrid_mla_attention): - mla = self._build(hybrid_mla_attention, sink=True) + for phase in self._PHASES: + with self.subTest(phase=phase): + mla = self._build(*phase, sink=True) sink = mla.core_attention.softmax_offset self.assertIsNotNone(sink) # One logit per local head of the hybrid MLA layer. self.assertEqual(list(sink.shape), [64]) - # bf16 == params_dtype: the FA4 cute kernel of the dense MHA - # phase asserts learnable_sink.dtype == bfloat16. + # bf16 == params_dtype: the FA4 cute kernel of the dense phases + # asserts learnable_sink.dtype == bfloat16. self.assertEqual(sink.dtype, paddle.bfloat16) # The shared ``build_softmax_offset`` helper initialises the # learnable sink with ``config.init_method`` (normal, std=0.02) @@ -3278,67 +3325,77 @@ def test_enabled_creates_one_zero_initialised_per_head_parameter(self): ) def test_state_dict_name_is_identical_across_phases(self): - # The valuable "an MHA checkpoint stays loadable into a latent-MQA - # run" guarantee: same name, shape and dtype in both phases. - mha = self._build(hybrid_mla_attention="mha", sink=True) - mqa = self._build(hybrid_mla_attention="mqa_dsa", sink=True) - self.assertEqual(self._sink_keys(mha), self._sink_keys(mqa)) - self.assertEqual( - mha.core_attention.softmax_offset.shape, - mqa.core_attention.softmax_offset.shape, - ) - self.assertEqual( - mha.core_attention.softmax_offset.dtype, - mqa.core_attention.softmax_offset.dtype, - ) - - def test_mha_sink_without_fa4_is_rejected_at_construction(self): - # The dense MHA phase reaches the sink only through the flashmask cute + # The valuable "an MHA checkpoint stays loadable by the later phases" + # guarantee: same name, shape and dtype in all three. + seen = {} + for phase in self._PHASES: + mla = self._build(*phase, sink=True) + sink = mla.core_attention.softmax_offset + seen[phase] = ( + tuple(self._sink_keys(mla)), + tuple(sink.shape), + str(sink.dtype), + ) + self.assertEqual(len(set(seen.values())), 1, seen) + self.assertEqual(seen[self._MHA][0], ("core_attention.softmax_offset",)) + + def test_dense_sink_without_fa4_is_rejected_at_construction(self): + # The dense phases reach the sink only through the flashmask cute # kernel, which is gated on FLAGS_flash_attn_version in (3, 4). With the # image default of 2 the run used to die at the *first forward* on an # opaque ``learnable_sink is only supported on the flashmask v4 (cute) # path`` assertion; ``MultiLatentAttention.__init__`` now refuses up - # front. + # front. Phase 2 is dense, so it is gated exactly like phase 1 -- that + # inheritance is the point, not an oversight. previous = paddle.get_flags(["FLAGS_flash_attn_version"])[ "FLAGS_flash_attn_version" ] paddle.set_flags({"FLAGS_flash_attn_version": 2}) try: - with self.assertRaisesRegex( - RuntimeError, "FLAGS_flash_attn_version" - ): - self._build_raw(hybrid_mla_attention="mha", sink=True) - # The absorbed phase owns its sink inside the block-sparse kernel, - # so the flag must NOT gate it. + for phase in (self._MHA, self._WARMUP): + with ( + self.subTest(phase=phase), + self.assertRaisesRegex( + RuntimeError, "FLAGS_flash_attn_version" + ), + ): + self._build_raw(*phase, sink=True) + # The latent phase owns its sink inside the block-sparse kernel, so + # the flag must NOT gate it. self.assertIsNotNone( self._build_raw( - hybrid_mla_attention="mqa_dsa", sink=True + *self._SPARSE, sink=True ).core_attention.softmax_offset ) finally: paddle.set_flags({"FLAGS_flash_attn_version": previous}) - def test_mha_sink_with_non_bf16_params_dtype_is_rejected(self): + def test_dense_sink_with_non_bf16_params_dtype_is_rejected(self): # Second requirement of the same cute kernel: it asserts the learnable # sink is bf16, and the sink is created with ``params_dtype``. An fp32 # run therefore used to pass construction and die on the first forward # with a terse ``learnable_sink must be bfloat16`` that names no config - # knob. The guard now names ``params_dtype`` at construction time. + # knob. The guard now names ``params_dtype`` at construction time, for + # both dense phases. previous = paddle.get_flags(["FLAGS_flash_attn_version"])[ "FLAGS_flash_attn_version" ] paddle.set_flags({"FLAGS_flash_attn_version": 4}) try: - with self.assertRaisesRegex(RuntimeError, "params_dtype"): - self._build_raw( - hybrid_mla_attention="mha", - sink=True, - params_dtype=paddle.float32, - ) - # The block-sparse kernel up-casts the sink itself, so the absorbed + for phase in (self._MHA, self._WARMUP): + with ( + self.subTest(phase=phase), + self.assertRaisesRegex(RuntimeError, "params_dtype"), + ): + self._build_raw( + *phase, + sink=True, + params_dtype=paddle.float32, + ) + # The block-sparse kernel up-casts the sink itself, so the latent # phase must stay dtype agnostic. mqa = self._build_raw( - hybrid_mla_attention="mqa_dsa", + *self._SPARSE, sink=True, params_dtype=paddle.float32, ) diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py b/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py index e4e72a4ee9..3441a67a4e 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py @@ -32,8 +32,15 @@ experiment). * ``..._non_absorbed_mqa_hca_dsa`` -- phase 2, ``"mqa_dsa"`` (+ YAML ``train_indexer_only``). + No top-k on either side, so the + attention half is phase 1's dense + MHA (``MHADSAWarmupAttention``) + with the indexer bolted on. * ``..._non_absorbed_mqa_hca_dsa_sparse_loss`` -- phase 3, ``"mqa_dsa"`` with ``dsa_indexer_use_sparse_loss``. + The only config-level difference + from phase 2, and the flag that + selects latent MQA. * ``ernielite_layer43_mqa_hca`` -- CSA full-causal MQA (``csa_compress_ratios == -1``). NOT a hybrid-MLA config: a @@ -199,8 +206,8 @@ def _assert_table(self, name, expected_core, expected_indexer): def test_mha_baseline_uses_dot_product_no_indexer(self): self._assert_table(_MHA, "DotProductAttention", None) - def test_mqa_uses_mqa_latent_with_dsa_indexer(self): - """The *other* production ``"mqa_dsa"`` config dispatches identically. + def test_sparse_loss_uses_mqa_latent_with_dsa_indexer(self): + """Phase 3: ``"mqa_dsa"`` + sparse loss dispatches to latent MQA. NAME/FIXTURE HISTORY: this used to run a config called ``ernielite_layer43_mla_mqa_hca`` -- an indexer-less latent-MQA variant @@ -213,12 +220,14 @@ def test_mqa_uses_mqa_latent_with_dsa_indexer(self): for that config now lives in ``test_csa_full_causal_mqa_config_never_builds_latent_mqa``. - The property kept here is unchanged in strength: *every* production - config that sets ``hybrid_mla_attention="mqa_dsa"`` must dispatch to - ``MQALatentAttention`` + ``DSAIndexer`` on all seven ``-2`` layers. Since - the enum rename the second such config is the phase-3 sparse-loss one, - so this pins that ``dsa_indexer_use_sparse_loss`` does NOT leak into the - class dispatch. + WAS: "``dsa_indexer_use_sparse_loss`` does NOT leak into the class + dispatch", i.e. both ``"mqa_dsa"`` configs built ``MQALatentAttention``. + NO LONGER TRUE: it is now the *only* thing that selects the class on a + ``"mqa_dsa"`` config (``hybrid_mla_indexer.latent_mqa_enabled``), because + the zero-sparsity phase-2 candidate set is served by dense MHA instead. + The flip itself is pinned in + ``test_sparse_loss_flag_alone_flips_the_core_attention_class``; this test + keeps the phase-3 half of the table. """ self._assert_table(_SPARSE_LOSS, "MQALatentAttention", "DSAIndexer") @@ -256,8 +265,54 @@ def test_csa_full_causal_mqa_config_never_builds_latent_mqa(self): self.assertEqual(ratio, 128, f"L{li}") self.assertEqual(seen_minus1, list(_MINUS2_LAYERS)) - def test_mqa_dsa_uses_mqa_latent_with_dsa_indexer(self): - self._assert_table(_DSA, "MQALatentAttention", "DSAIndexer") + def test_mqa_dsa_uses_dense_warmup_mha_with_dsa_indexer(self): + """Phase 2 dispatches to dense MHA + the DSA indexer, not latent MQA. + + WAS ``self._assert_table(_DSA, "MQALatentAttention", "DSAIndexer")``. + Phase 2 has no top-k on either side, so routing it through the + block-sparse latent-MQA kernel only bought a per-document causal + ``[b, s, s]`` index table (256MB/layer at s=8192) that the kernel then + walked in full; it now runs phase 1's dense attention with the indexer + bolted on (``mha_dsa_warmup_attention.py:103``, dispatched at + ``gpt_layer_specs.py`` via ``latent_mqa_enabled``). + """ + self._assert_table(_DSA, "MHADSAWarmupAttention", "DSAIndexer") + + def test_sparse_loss_flag_alone_flips_the_core_attention_class(self): + """``dsa_indexer_use_sparse_loss`` is the whole phase-2/3 dispatch. + + Both directions from one provider, so the assertion cannot pass by two + configs differing in some other field: ``False`` (phase 2, no top-k + anywhere) must give ``MHADSAWarmupAttention``, ``True`` (phase 3) must + give ``MQALatentAttention``, and the indexer spec and the enclosing + ``MLASelfAttention`` must be the same on both sides -- that is what makes + the phase switch a pure backend swap with no parameter rename. + + This is the test that catches a dispatch regression, and it pins + ``hybrid_mla_indexer.latent_mqa_enabled`` (the single predicate shared by + the spec dispatch and ``MLASelfAttention.mqa_latent``) rather than + re-deriving the rule. + """ + from paddlefleet.transformer.hybrid_mla_indexer import ( + latent_mqa_enabled, + ) + + _, provider = _load_provider(_DSA) + self.assertEqual(provider.hybrid_mla_attention, "mqa_dsa") + self.assertFalse(provider.dsa_indexer_use_sparse_loss) + expected = {False: "MHADSAWarmupAttention", True: "MQALatentAttention"} + for sparse_loss, core in expected.items(): + provider.dsa_indexer_use_sparse_loss = sparse_loss + self.assertIs(latent_mqa_enabled(provider), sparse_loss) + for li in _MINUS2_LAYERS: + with self.subTest(sparse_loss=sparse_loss, layer=li): + ratio, attn_cls, core_cls, indexer_cls = _dispatch( + provider, li + ) + self.assertEqual(ratio, -2) + self.assertEqual(attn_cls, "MLASelfAttention") + self.assertEqual(core_cls, core) + self.assertEqual(indexer_cls, "DSAIndexer") def test_mqa_full_causal_drops_the_indexer(self): """``"mqa_full_causal"`` keeps the latent MQA core, removes the indexer. @@ -361,8 +416,10 @@ class TestSinkParameterOnRealModules(unittest.TestCase): """Build the real ``-2`` layer and check the learnable sink. Consumer: ``build_softmax_offset`` (dot_product_attention.py:87), called by - BOTH ``DotProductAttention.__init__`` (MHA phase) and - ``MQALatentAttention.__init__`` (the latent MQA modes). Proves the + ALL THREE ``-2`` core-attention classes: ``DotProductAttention.__init__`` + (phase 1), ``MHADSAWarmupAttention`` (phase 2, via the same + ``DotProductAttention.__init__``) and ``MQALatentAttention.__init__`` + (``"mqa_full_causal"`` and phase 3). Proves the model-wide ``add_full_attention_sink_bias`` JSON flag reaches a real bf16 [num_heads] param at the SAME state_dict key (``core_attention.softmax_offset``) in both phases -- which is what keeps an @@ -417,12 +474,16 @@ def test_mha_baseline_has_no_sink(self): ) def test_mqa_and_dsa_have_trainable_bf16_per_head_sink(self): + # Phase 2 is dense MHA now, phase 3 latent MQA; the sink is built by the + # shared ``build_softmax_offset`` either way, which is the property here. + cores = { + _DSA: "MHADSAWarmupAttention", + _SPARSE_LOSS: "MQALatentAttention", + } for name in _MQA_DSA_CFGS: with self.subTest(config=name): mod = self._build(name) - self.assertEqual( - type(mod.core_attention).__name__, "MQALatentAttention" - ) + self.assertEqual(type(mod.core_attention).__name__, cores[name]) sink = mod.core_attention.softmax_offset self.assertIsNotNone(sink) self.assertEqual(list(sink.shape), [64]) @@ -498,7 +559,12 @@ def test_dsa_indexer_reflects_model_wide_index_fields(self): # block size). provider.dsa_index_n_heads = 32 provider.dsa_index_topk = 256 - mod = _build_real_attn(provider, _MINUS2_LAYERS[0]) + # ``_DSA`` is phase 2 = dense MHA + sink, so the build now goes through + # the ``FLAGS_flash_attn_version in (3, 4)`` guard that used to apply to + # the baseline only (multi_latent_attention.py:587-614). Pin the + # production value, as ``TestSinkParameterOnRealModules._build`` does. + with _flash_attn_version(_production_fa_version()): + mod = _build_real_attn(provider, _MINUS2_LAYERS[0]) idx = mod.core_attention.indexer self.assertEqual(type(idx).__name__, "DSAIndexer") self.assertEqual(idx.n_heads, 32) @@ -886,9 +952,7 @@ class TestConfigDeltas(unittest.TestCase): # deliberate and identical in all four, so they are pinned as a group -- if # one variant drifts out of the group this fails. _YAML_COMMON_DELTA = { - # The MLA layers use plain RoPE; the fused path is the YaRN one. - "apply_rope_fusion": (True, False), - "dsv4_yarn_rope_fusion": (True, _MISSING), + "dsv4_yarn_rope_fusion": (False, _MISSING), # full/uniform/1 recompute instead of the selective module list. "recompute_granularity": ("selective", "full"), "recompute_method": (_MISSING, "uniform"), @@ -897,6 +961,26 @@ class TestConfigDeltas(unittest.TestCase): ["full_attn", "moe_gate_up", "moe_premute", "mhc_forward"], _MISSING, ), + # Checkpointing / logging / accumulation knobs the *baseline* moved on + # after the four experiment yamls forked from it. Not this feature: the + # delta is identical across all four variants, including + # ``ernielite_layer43_mqa_hca``, which carries no ``hybrid_mla_*`` key + # at all. Pinned here rather than resynced in the yamls, because those + # files drive submitted jobs -- editing them would change run behaviour + # to make a test pass. + "enable_zero_cost_checkpoint": (True, False), + "flash_device_save_steps": (50, _MISSING), + "zcc_workers_num": (1, _MISSING), + "save_steps": (400, 200), + "save_hf_steps": (800, 200), + "gradient_accumulation_steps": (2, 1), + "global_logging_interval": (20, 10), + "load_process_num": (3, 8), + "sharding_comm_buffer_size_MB": (4096, 2048), + "output_dir": ( + "./output/ernielite_pretrain8k_fleet_exp2-32", + "./output/ernielite_pretrain8k_fleet_exp2-13", + ), } def _yaml_allowlist(self, name): @@ -905,9 +989,17 @@ def _yaml_allowlist(self, name): f"./model_config_separated/conf/fleet_align/{_MHA}", f"./model_config_separated/conf/fleet_align/{name}", ) + # ``apply_rope_fusion`` is NOT a common delta: the fused MLA RoPE kernel + # needs the per-head K/V that absorption never materialises, so only the + # *latent* variants turn it off. The downgrade lives inside + # ``if self.mqa_latent:`` (``multi_latent_attention.py:485-496``) and is a + # warning plus an eager-RoPE fallback, not a construction error -- so the + # DSA warmup phase, whose ``mqa_latent`` is False, keeps the baseline's + # fused kernel exactly like phase 1. Its own RoPE is on the independent + # ``dsa_indexer_rope_fusion`` (``dsa_attention.py:458-461``). + if name != _DSA: + allowed["apply_rope_fusion"] = (True, False) if name in _MQA_DSA_CFGS: - # The DSA indexer on the -2 layers only has a cuDNN backward. - allowed["csa_indexer_backend"] = ("tilelang", "cudnn") # ``indexer_init_from_scratch`` is mandatory once an indexer exists # (``modeling.py`` hard-errors on ``None``), so both mqa_dsa phases # set it -- with opposite values, which is the point: phase 2 @@ -940,6 +1032,15 @@ def _yaml_allowlist(self, name): if name == _DSA: # Phase 2 = DSA warmup: only the indexer trains. allowed["train_indexer_only"] = (self._MISSING, True) + if name == _SPARSE_LOSS: + # Documentary only, and only on this config: neither phase reads it + # for a -2 layer (phase 2 hardcodes ``csa_indexer_topk_fwd``, + # ``mha_dsa_warmup_attention.py``; phase 3 hardcodes + # ``cudnn_indexer_topk_fwd``, ``mqa_latent_attention.py:535,586``), + # and the ratio-128 HCA layers set ``indexer = None`` + # (``csa_attention.py:2038-2049``). Phase 2 therefore keeps the + # baseline's ``tilelang`` rather than diverging for no reason. + allowed["csa_indexer_backend"] = ("tilelang", "cudnn") return allowed @classmethod diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_doc_equivalence.py b/tests/single_card_tests/transformer/test_hybrid_mla_doc_equivalence.py index f6530afd89..650bb815a8 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_doc_equivalence.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_doc_equivalence.py @@ -350,8 +350,11 @@ def test_window_indexer_partition_all_layouts(self): window = _build_window_topk_idxs_from_doc_bounds( 1, seqlen, WINDOW, doc_start, is_valid ).numpy() + # ``window`` is a required positional arg on + # ``HybridMLAIndexerMixin._indexer_valid_range``; it used to + # default to ``self.window_size``, which is ``WINDOW`` here. vr, row_empty = self.module._indexer_valid_range( - seqlen, doc_start, doc_len, is_valid + seqlen, doc_start, doc_len, is_valid, WINDOW ) vr = vr.numpy()[0] row_empty = row_empty.numpy().reshape([seqlen]) diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_grad_health.py b/tests/single_card_tests/transformer/test_hybrid_mla_grad_health.py index 9dc35cdeb1..f8eaf5571d 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_grad_health.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_grad_health.py @@ -18,13 +18,24 @@ layer (csa ratio == -2, experimental_attention_variant == "dsv4_hybrid"). * ``"mha"`` -- dense MLA (MLASelfAttention + DotProductAttention). - * ``"mqa_dsa"`` -- latent MQA: runtime activation-level absorption on the - shared KV latent PLUS a cuDNN block-sparse DSA indexer (gpt_layer_specs - always builds the indexer for such a layer). With the forced local window - (``csa_window_size``) >= sequence length the indexer selects no extra - tokens, so the absorbed path degenerates to full per-document causal - attention -- mathematically equal to the dense MHA phase, which is what - the equivalence items below rely on. + * ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=True`` (phase 3/4) -- latent + MQA: runtime activation-level absorption on the shared KV latent PLUS a + cuDNN block-sparse DSA indexer (gpt_layer_specs always builds the indexer + for such a layer). With the forced local window (``csa_window_size``) >= + sequence length the indexer selects no extra tokens, so the absorbed path + degenerates to full per-document causal attention -- mathematically equal + to the dense MHA phase, which is what the equivalence items below rely on. + * ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=False`` (phase 2, DSA warmup) + -- ``MHADSAWarmupAttention``: phase 1's dense MHA with the indexer's + full-candidate KL bolted on. The warmup phase has no top-k on either side, + so there is nothing for latent MQA to absorb and no block-sparse kernel is + involved. The predicate is ``hybrid_mla_indexer.latent_mqa_enabled``. + +``sparse_loss`` therefore selects the attention *backend*, not merely the loss +width. ``_build`` defaults it to ``True`` so the equivalence / magnitude / +sink / stress / stability items keep measuring the latent path they were +written for; item 10 (the indexer KL at its production phase-2 setting) passes +``False``, and item 4 walks all three backends. The old 3-state ``hybrid_mla_attn_mode`` {mha, mqa, mqa_dsa} became the ``hybrid_mla_attention`` enum. Only two of its values are exercised here: "mha" @@ -85,6 +96,10 @@ def _try_use_cuda_device(): from paddlefleet.tensor_parallel.random import ( model_parallel_cuda_manual_seed, ) +from paddlefleet.transformer.hybrid_mla_indexer import latent_mqa_enabled +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, +) from paddlefleet.transformer.multi_latent_attention import ( MLASelfAttention, # noqa: F401 (import forces hybrid_mla_* field registration) ) @@ -137,11 +152,21 @@ def __init__(self, tp_nranks=1, cp_nranks=1): # == 0, kv_lora_rank 512 + qk_rope 64 == 576 latent key width, v == 512. # --------------------------------------------------------------------------- # def _make_config( - mode, sink, indexer=False, dtype=paddle.bfloat16, split_kv_b=False + mode, + sink, + indexer=False, + dtype=paddle.bfloat16, + split_kv_b=False, + sparse_loss=True, ): # ``indexer`` is vestigial: the hybrid-MLA layer's DSA indexer is now built # unconditionally for ``hybrid_mla_attention="mqa_dsa"`` (see # gpt_layer_specs), so it is retained only for call-site compatibility. + # + # ``sparse_loss`` (``dsa_indexer_use_sparse_loss``) is the phase switch of + # ``"mqa_dsa"`` and picks the attention backend as well as the KL width: + # ``True`` -> latent MQA (phase 3/4), ``False`` -> the dense + # ``MHADSAWarmupAttention`` of phase 2. Ignored by ``mode="mha"``. del indexer hybrid_mla_attention = "mha" if mode == "mha" else "mqa_dsa" cfg = TransformerConfig( @@ -187,7 +212,7 @@ def _make_config( dsa_index_head_dim=128, dsa_index_topk=128, dsa_indexer_loss_coeff=1.0, - dsa_indexer_use_sparse_loss=False, + dsa_indexer_use_sparse_loss=sparse_loss, dsa_indexer_rotary_interleaved=False, apply_rope_fusion=False, attention_dropout=0.0, @@ -209,35 +234,56 @@ def _build( indexer=False, dtype=paddle.bfloat16, split_kv_b=False, + sparse_loss=True, ): """Build the layer-0 (hybrid-MLA) self-attention for ``mode``.""" model_parallel_cuda_manual_seed(_SEED) - cfg = _make_config(mode, sink, indexer, dtype=dtype, split_kv_b=split_kv_b) + cfg = _make_config( + mode, + sink, + indexer, + dtype=dtype, + split_kv_b=split_kv_b, + sparse_loss=sparse_loss, + ) spec = get_gpt_layer_local_spec( config=cfg, normalization=cfg.normalization, layer_number=0, ).sublayers_spec.self_attn - with _fa4_for_mha_sink(mode, sink): + with _fa4_for_dense_sink(cfg): module = build_spec_layer( spec, config=cfg, layer_number=0, pg_collection=_FakePGCollection() ) return module -@contextlib.contextmanager -def _fa4_for_mha_sink(mode, sink): - """Satisfy the ``mha`` + sink FA4 requirement for the duration of a build. +def _uses_latent_mqa(mode, sparse_loss=True, **_ignored): + """Whether ``_build`` yields the absorbed latent-MQA backend. - ``MultiLatentAttention.__init__`` refuses to create the sink for ``mha`` - unless ``FLAGS_flash_attn_version in (3, 4)``, because ``mha`` consumes it as - ``flashmask_attention_func(learnable_sink=...)`` which only exists on the - cute path (multi_latent_attention.py, guard next to the parameter creation). + Routed through the production predicate instead of re-deriving the rule, so + a fixture can never disagree with ``gpt_layer_specs``' dispatch. + """ + return latent_mqa_enabled( + _make_config(mode, sink=False, sparse_loss=sparse_loss) + ) + + +@contextlib.contextmanager +def _fa4_for_dense_sink(cfg): + """Satisfy the dense-MLA + sink FA4 requirement for the duration of a build. + + ``MultiLatentAttention.__init__`` refuses to create the sink unless + ``FLAGS_flash_attn_version in (3, 4)`` whenever the layer is *not* latent + MQA (multi_latent_attention.py:584-607), because the dense path consumes it + as ``flashmask_attention_func(learnable_sink=...)`` which only exists on the + cute path. That now covers two phases, not one: phase 1 (``"mha"``) and the + phase-2 DSA warmup, which deliberately inherits every phase-1 constraint. The default flag value in this image is 2. These tests only inspect the parameter, so we flip the flag around construction and restore it -- the process-global default must stay untouched for the other suites. """ - if not (sink and mode == "mha"): + if not (cfg.add_full_attention_sink_bias and not latent_mqa_enabled(cfg)): yield return previous = paddle.get_flags(["FLAGS_flash_attn_version"])[ @@ -250,12 +296,17 @@ def _fa4_for_mha_sink(mode, sink): paddle.set_flags({"FLAGS_flash_attn_version": previous}) -# Modes exercised here, as (label, vestigial ``indexer`` argument). Only two of -# the three ``hybrid_mla_attention`` values are relevant to gradient health: -# ``"mha"`` and ``"mqa_dsa"`` (which always builds the indexer). The -# indexer-less ``"mqa_full_causal"`` is a non-production equivalence experiment, -# so its redundant entry was dropped. -_MODES = (("mha", False), ("mqa_dsa", True)) +# The three hybrid-MLA backends reachable from this fixture, as +# ``(label, _build kwargs)``. ``"mqa_dsa"`` is two backends now rather than one: +# ``dsa_indexer_use_sparse_loss`` picks the attention as well as the loss width, +# so the warmup phase runs phase 1's dense MHA (``MHADSAWarmupAttention``) and +# only the sparse phase runs latent MQA. The indexer-less ``"mqa_full_causal"`` +# is a non-production equivalence experiment covered elsewhere. +_MODES = ( + ("mha", {"mode": "mha"}), + ("mqa_dsa_warmup", {"mode": "mqa_dsa", "sparse_loss": False}), + ("mqa_dsa_sparse", {"mode": "mqa_dsa", "sparse_loss": True}), +) # Parameters we expect to carry gradient in the hybrid-MLA layer (weights only; # the sink is handled separately because mha's is inert here, see item 3). @@ -786,12 +837,19 @@ def test_mqa_sink_grad_finite_nonzero_and_matches_finite_diff(self): @_skip_if_no_cuda class TestItem4GradientFlowCompleteness(unittest.TestCase): """Item 4: every parameter that should receive a gradient does -- no None, - no all-zero -- in all reachable modes, sink on and off.""" + no all-zero -- in all reachable backends, sink on and off. + + "All reachable backends" is three, not two, since the phase split: dense + phase 1 (``"mha"``), the dense phase-2 DSA warmup + (``MHADSAWarmupAttention``) and latent MQA (phase 3/4). Only the last one + needs the SM100+ block-sparse kernel; the warmup phase runs wherever phase 1 + runs, which is the point of it being phase 1's attention. + """ SEQ = 64 - def _check(self, mode, sink, indexer): - m = _build(mode, sink=sink, indexer=indexer) + def _check(self, sink, **build_kwargs): + m = _build(sink=sink, **build_kwargs) _, g = _grads( m, _hidden(self.SEQ, seed=2), _row_end(self.SEQ), weighted=True ) @@ -814,23 +872,26 @@ def _check(self, mode, sink, indexer): return missing, dead, sink_state def test_all_weight_params_receive_gradient(self): - for mode, idx in _MODES: - if mode == "mqa_dsa" and not _HAS_DSA: + for label, kwargs in _MODES: + latent = _uses_latent_mqa(**kwargs) + if latent and not _HAS_DSA: continue for sink in (False, True): - if mode == "mha" and sink: - # mha + sink cannot run here (FA4 cute kernel absent); the - # forward raises. Documented under item 3; skip flow check. + if sink and not latent: + # A dense MLA sink cannot run here (FA4 cute kernel + # absent); the forward raises. That covers ``"mha"`` and the + # phase-2 warmup alike -- the warmup phase inherits phase + # 1's constraints deliberately. Documented under item 3. continue - with self.subTest(mode=mode, sink=sink): - missing, dead, sink_state = self._check(mode, sink, idx) - self.assertEqual(missing, [], f"{mode}: missing grads") - self.assertEqual(dead, [], f"{mode}: all-zero grads") - if sink and mode != "mha": + with self.subTest(backend=label, sink=sink): + missing, dead, sink_state = self._check(sink, **kwargs) + self.assertEqual(missing, [], f"{label}: missing grads") + self.assertEqual(dead, [], f"{label}: all-zero grads") + if sink: self.assertEqual( sink_state, "live", - f"{mode}: sink got no usable gradient", + f"{label}: sink got no usable gradient", ) @@ -1032,22 +1093,36 @@ def test_weight_decay_pull_on_sink(self): @_skip_if_no_cuda class TestItem9InputIdsReachTheIndexerLossMask(unittest.TestCase): - """``input_ids`` must reach the MQA core attention, and only that one. + """``input_ids`` must reach an indexer-owning core attention, and no other. The indexer-loss row mask has to come from ``input_ids != pad_token_id``: ``attn_mask_startend_row_indices`` folds a packed sequence's trailing padding into the last document, so the document metadata alone reports those rows as valid. ``MultiLatentAttention.forward`` therefore forwards - ``input_ids`` to ``core_attention`` -- but only under ``mqa_latent``, since + ``input_ids`` to ``core_attention`` -- but only when the core attention + advertises ``accepts_input_ids`` (multi_latent_attention.py:945-950), since ``DotProductAttention.forward`` has no such parameter and no ``**kwargs``. + + That capability flag, not ``mqa_latent``, is the gate: BOTH indexer-owning + phases need the mask (``HybridMLAIndexerMixin.accepts_input_ids = True``, + hybrid_mla_indexer.py:76) while only phase 3/4 absorbs, so the phase-2 dense + warmup backend must be covered here too -- it is the phase that actually + runs packed sequences with a padded tail in production. """ SEQ = 64 - def test_mqa_core_attention_receives_input_ids(self): - if not _HAS_DSA: - self.skipTest("absorbed MQA forward requires SM100+ DSA kernel") - module = _build("mqa_dsa") + def test_indexer_core_attentions_receive_input_ids(self): + for label, kwargs in _MODES: + if kwargs["mode"] == "mha": + continue + if _uses_latent_mqa(**kwargs) and not _HAS_DSA: + continue + with self.subTest(backend=label): + self._assert_forwarded(**kwargs) + + def _assert_forwarded(self, **build_kwargs): + module = _build(**build_kwargs) seen = {} inner = module.core_attention.forward @@ -1064,6 +1139,10 @@ def recording(*args, **kwargs): attn_mask_startend_row_indices=_row_end(self.SEQ), input_ids=input_ids, ) + self.assertTrue( + getattr(module.core_attention, "accepts_input_ids", False), + "an indexer-owning core attention must accept input_ids", + ) self.assertIsNotNone(seen["input_ids"]) self.assertEqual(list(seen["input_ids"].shape), [1, self.SEQ]) @@ -1095,13 +1174,19 @@ class TestItem10IndexerKLValueAndDenominator(unittest.TestCase): divides by too) equals the non-pad token count, so forward and backward agree; * the coefficient is applied exactly once (linear in coeff). + + ``dsa_indexer_use_sparse_loss=False`` now also selects the *backend*, so the + module under test is ``MHADSAWarmupAttention`` (dense MHA + full-candidate + KL) rather than latent MQA. Only the construction and the spied-on module + change; every observable above is a property of the KL, which is why this is + still one test. """ SEQ = 256 NPAD = 48 # trailing padding folded into the last document's row range def _run_and_capture(self, coeff): - import paddlefleet.transformer.mqa_latent_attention as mqamod + import paddlefleet.transformer.mha_dsa_warmup_attention as warmupmod from paddlefleet.transformer.dsa_attention import ( DSAIndexerLossLoggingHelper as LOG, ) @@ -1109,13 +1194,14 @@ def _run_and_capture(self, coeff): cap = {} # Phase 2 attaches its indexer loss through the shared # ``TileLangCSAIndexerLossAutoScaler`` PyLayer (imported into - # ``mqa_latent_attention`` from ``csa_attention``), with the - # ``"tilelang"`` backend tag. One spy on that boundary carries every - # observable this test needs -- ``P`` (``topk_probs``, already softmaxed - # by ``csa_indexer_topk_fwd``), ``Q`` (``target``), the row mask, the + # ``mha_dsa_warmup_attention`` from ``csa_attention``; phase 3 imports + # the same symbol into ``mqa_latent_attention``), with the ``tilelang`` + # backend tag. One spy on that boundary carries every observable this + # test needs -- ``P`` (``topk_probs``, already softmaxed by + # ``csa_indexer_topk_fwd``), ``Q`` (``target``), the row mask, the # denominator and the coefficient -- and they are exactly the tensors the # tilelang ``csa_indexer_bwd`` differentiates. - real = mqamod.TileLangCSAIndexerLossAutoScaler + real = warmupmod.TileLangCSAIndexerLossAutoScaler # Bound by position, so pin the order: a reordered signature would hand # the spy the wrong tensors instead of failing. expected_args = [ @@ -1179,14 +1265,25 @@ def apply( loss_mask, ) - module = _build("mqa_dsa") + module = _build("mqa_dsa", sparse_loss=False) + # ``dsa_indexer_use_sparse_loss=False`` is the phase-2 pair, so the + # backend must be the dense warmup one. Asserted on the class (and the + # config below), not on a ``indexer_use_sparse_loss`` attribute: + # ``MHADSAWarmupAttention`` has none -- the phase is a property of the + # config, and the backend the config selected is what the class name + # reports. + self.assertIsInstance( + module.core_attention, + MHADSAWarmupAttention, + "not the phase-2 backend", + ) + self.assertFalse(module.config.dsa_indexer_use_sparse_loss) module.core_attention.indexer_loss_coeff = float(coeff) - module.core_attention.indexer_use_sparse_loss = False module.train() ids = np.ones((1, self.SEQ), dtype="int64") ids[0, self.SEQ - self.NPAD :] = 0 # pad_token_id == 0 input_ids = paddle.to_tensor(ids) - mqamod.TileLangCSAIndexerLossAutoScaler = Spy + warmupmod.TileLangCSAIndexerLossAutoScaler = Spy LOG.clean_loss_in_tracker() try: module( @@ -1198,13 +1295,16 @@ def apply( input_ids=input_ids, ) finally: - mqamod.TileLangCSAIndexerLossAutoScaler = real + warmupmod.TileLangCSAIndexerLossAutoScaler = real logged = float(LOG.tracker["values"].astype("float32").sum()) return logged, cap def test_kl_value_denominator_and_coeff(self): if not _HAS_DSA: - self.skipTest("absorbed MQA forward requires SM100+ DSA kernel") + # The phase-2 attention itself is plain dense flashmask, but its KL + # runs the tilelang indexer kernels. ``is_dsa_available`` is the + # available proxy for "an SM100+ box with those kernels built". + self.skipTest("the tilelang indexer KL requires an SM100+ box") logged, cap = self._run_and_capture(coeff=0.01) nonpad = float(self.SEQ - self.NPAD) diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_hf_roundtrip.py b/tests/single_card_tests/transformer/test_hybrid_mla_hf_roundtrip.py index facf228377..ff0a293209 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_hf_roundtrip.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_hf_roundtrip.py @@ -52,6 +52,12 @@ (``load_state_dict.py:325-327``) whose text is wrong -- the keys *are* in the model state dict. The current behaviour is pinned as a documented hazard, and the behaviour we would want is a companion ``expectedFailure``. +5. **Phase 2 and phase 3 expose the same HF surface.** They now run *different* + core-attention classes (``MHADSAWarmupAttention`` vs ``MQALatentAttention``, + picked by ``latent_mqa_enabled``), so the key set / shapes / dtypes on disk + are pinned equal and a phase-2 checkpoint is loaded into a phase-3 module + bit-exactly. HF safetensors is the only supported phase switch, so a name + drift between the two classes would be an unguarded silent re-init. """ import unittest @@ -65,6 +71,7 @@ from .hybrid_mla_utils import ( _CONFIG_DIR, _DSA_CFG as _DSA, + _DSA_SPARSE_LOSS_CFG as _DSA3, _MHA_CFG as _MHA, _PARENT_REPO_AVAILABLE, _add_repo_root_to_sys_path, @@ -436,6 +443,93 @@ def test_mqa_dsa_round_trip_is_bitwise_exact(self): self.assertIn(key, differ) +@_requires_cuda +class TestPhase2Phase3HFIdentity(unittest.TestCase): + """Q5: the phase-2 -> phase-3 hand-over, which is the *only* supported way + to switch phases (HF safetensors; DCP is out of scope by ruling). + + Phase 2 no longer runs the same class as phase 3: ``latent_mqa_enabled`` + (``hybrid_mla_indexer.py``) sends ``mqa_dsa`` + + ``dsa_indexer_use_sparse_loss=False`` to ``MHADSAWarmupAttention`` and + ``+ True`` to ``MQALatentAttention`` (``gpt_layer_specs.py`` dispatch). Two + different core-attention classes could easily disagree on parameter names, + which would break the hand-over silently -- ``_ -> key`` would re-init and + the rest would be an "Unexpected keys" warning. So pin the HF surface as + *identical*: same key set, same shapes, same dtypes, and a bit-exact + phase-2 -> phase-3 load. + """ + + def _emit(self, cfg_name, seed=11): + attn = _attn(cfg_name, seed=seed) + keys = {PREFIX + k for k in attn.state_dict()} + _, inverse = _aoa(cfg_name, indexer_init_from_scratch=True) + with TemporaryDirectory() as tmp: + path = Path(tmp) / "hf" + emitted = _save_hf( + attn, _filter(inverse, keys, fleet_on_left=True), path + ) + disk = set(_on_disk(path)) + return attn, {e[0]: (e[1], e[2]) for e in emitted}, disk + + def test_phase2_and_phase3_really_are_different_classes(self): + """ANTI-VACUITY: without this, the identity below could just be two + builds of the same class. + """ + from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, + ) + from paddlefleet.transformer.mqa_latent_attention import ( + MQALatentAttention, + ) + + self.assertIsInstance( + _attn(_DSA, seed=11).core_attention, MHADSAWarmupAttention + ) + self.assertIsInstance( + _attn(_DSA3, seed=11).core_attention, MQALatentAttention + ) + + def test_phase2_and_phase3_hf_surfaces_are_identical(self): + _, t2, disk2 = self._emit(_DSA) + _, t3, disk3 = self._emit(_DSA3) + self.assertEqual( + sorted(disk2), sorted(disk3), "phase-2/3 HF key sets diverged" + ) + self.assertEqual(sorted(t2), sorted(t3)) + for key in sorted(t2): + with self.subTest(key=key): + self.assertEqual(t2[key], t3[key], "shape/dtype diverged") + for key in _INDEXER_KEYS: + self.assertIn(PREFIX + key, disk2) + self.assertIn(PREFIX + key, disk3) + self.assertIn(PREFIX + "core_attention.softmax_offset", disk2) + + def test_phase2_checkpoint_loads_into_phase3_bitwise(self): + """The hand-over itself: a warmup checkpoint must restore into the + sparse module exactly, indexer included (so + ``indexer_init_from_scratch`` has to be off, cf. + ``TestAddPrimitiveDiscardsTrainedIndexer``). + """ + src = _attn(_DSA, seed=11) + dst = _attn(_DSA3, seed=99) + keys = {PREFIX + k for k in src.state_dict()} + self.assertEqual(keys, {PREFIX + k for k in dst.state_dict()}) + _, inverse = _aoa(_DSA, indexer_init_from_scratch=True) + forward, _ = _aoa(_DSA3, indexer_init_from_scratch=False) + own = {k: v.clone() for k, v in dst.state_dict().items()} + differ = _diff_keys(own, src.state_dict()) + self.assertGreaterEqual(len(differ), 8, "too few moving tensors") + for key in _INDEXER_RANDOM_KEYS: + self.assertIn(key, differ, "the indexer must actually move") + with TemporaryDirectory() as tmp: + path = Path(tmp) / "phase2" + _save_hf(src, _filter(inverse, keys, fleet_on_left=True), path) + _load_hf(dst, _filter(forward, keys, fleet_on_left=False), path) + self.assertEqual(_diff_keys(dst.state_dict(), src.state_dict()), []) + for key, ref in src.state_dict().items(): + self.assertEqual(dst.state_dict()[key].dtype, ref.dtype) + + @_requires_cuda class TestCrossPhaseLoad(unittest.TestCase): """Q2: a phase-1 ``"mha"`` checkpoint into a ``"mqa_dsa"`` module.""" diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py index c1ca3ed013..f3d1a981d0 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py @@ -15,44 +15,56 @@ """Document masking and indexer-loss arithmetic of the hybrid-MLA WARMUP phase. The warmup phase is ``hybrid_mla_attention="mqa_dsa"`` with -``dsa_indexer_use_sparse_loss=False``: the indexer is still being learned, so -**attention** consumes the full per-document causal table -(``MQALatentAttention._build_full_causal_indices`` -> -``_build_mqa_causal_topk_idxs_from_doc_bounds``) and the KL is scored over that -same full causal set (``MQALatentAttention._forward_warmup`` -> one -``paddlefleet.tilelang_ops.csa_indexer_topk_fwd`` call at -``topk_effective = s_global``). The phase-3 cuDNN top-k kernel runs zero times -on this path, which is exactly why the mask semantics and the loss reduction -need pinning here and not only on the phase-3 path -(``dsa_indexer_use_sparse_loss=True``), where -``test_hybrid_mla_doc_equivalence.py`` and ``test_mqa_latent_attention.py`` -already cover them. +``dsa_indexer_use_sparse_loss=False``. It no longer runs the absorbed latent +MQA: with no top-k on either side the block-sparse kernel was being handed a +zero-sparsity ``[b, s, s]`` int32 index table and made to walk all ``s`` +columns anyway, so the phase now runs *exactly* phase 1's dense attention -- +``MHADSAWarmupAttention(HybridMLAIndexerMixin, DotProductAttention)``, which +delegates its whole attention half to ``super().forward`` -- with the indexer +bolted on: one ``paddlefleet.tilelang_ops.csa_indexer_topk_fwd`` call in +full-candidate mode (``ratio=1``, ``topk_effective=s_global``) and a KL against +the head-summed dense attention distribution +(``MHADSAWarmupAttention._dense_attn_target``). + +That inverts the two central claims of the previous revision of this file. +Where it asserted "the attention index table is exactly the full per-document +causal set", the assertion is now "no index table is built at all and the +block-sparse kernel is never reached"; where it asserted "warmup output == +``mqa_full_causal`` output bitwise", it is now "warmup output == phase-1 dense +output bitwise". Everything else -- mask semantics, the KL column set, the +reduction denominator, gradient health -- survives unchanged in intent, on the +dense per-head call shape. What is proven here, and nowhere else: -* ``TestWarmupFullCausalTable`` -- row ``i`` of the attention table is - *exactly* the column set ``[doc_start[i], i]``, on the layouts below plus a - layout with genuine pad rows. -* ``TestWarmupCrossDocumentIsolation`` -- zero cross-document leakage, in the - strong form the builder's docstring claims ("bit-identical to running each - document on its own"): measured ``maxabs == 0.0`` exactly, in both eval (the - ``:495`` early exit) and train (the ``:600`` branch). -* ``TestWarmupPadRows`` -- rows with ``is_valid == False``. ``_row_end`` cannot - produce them (it turns the trailing gap into one final valid document), so - ``_pad_row_end`` below keeps the gap folded into the last document instead, - which is what a real packed batch's padding tail looks like. -* ``TestWarmupIndexerLossPrecision`` -- the KL column set (the whole causal set, - with no cuDNN top-k call anywhere), the row mask coming from ``input_ids`` (not - from the document metadata), the reduction denominator, and the claim that the - KL target is the head-summed *full-causal* attention distribution. -* ``TestWarmupGradHealth`` -- the five indexer parameters, the detached indexer - inputs, and the attention-side ``dq`` / ``dkv`` / ``d_sink``. - -Shape caveat: ``WINDOW + INDEX_TOPK == 256`` in ``hybrid_mla_utils``, so a -``seqlen=256`` fixture has a *saturated* sparse budget -- the phase-3 top-k -table would already equal the full causal set there, and no assertion at that -shape can tell the two apart. ``_LAYOUTS`` therefore also carries ``seqlen=512`` -layouts, which is the discriminating shape. +* ``TestWarmupCandidateRange`` -- ``_indexer_valid_range(window=0)`` is exactly + the per-document causal span ``[doc_start[i], i]``, on every layout below + plus layouts with genuine pad rows and with documents shorter than the sparse + phase's forced window. Pure integer arithmetic, no kernel, hence not GPU + gated. +* ``TestWarmupIsPhase1Dense`` -- the output is **bit-identical** to a plain + ``DotProductAttention`` built from the same config, ``_CAPTURED`` stays empty + (no block-sparse call), and exactly one KL target of shape ``[1, s, s]`` is + built per grad-enabled forward. +* ``TestWarmupCrossDocumentIsolation`` -- zero cross-document leakage, measured + in the only form that stays exact on a dense kernel: replace every *other* + document's K/V with noise and this document's output does not move one bit. +* ``TestWarmupPadRows`` -- rows with ``is_valid == False`` contribute nothing + and receive nothing: output, ``dq`` and their KL rows are exactly zero. +* ``TestWarmupIndexerLossPrecision`` -- the KL column set (the whole causal + set, with no cuDNN top-k call anywhere), the row mask coming from + ``input_ids`` (not from the document metadata), the reduction denominator, + and the claim that the target is the head-summed dense attention + distribution. +* ``TestDenseAttnTarget`` -- ``_dense_attn_target`` against a naive + plain-paddle/numpy reference, including an all-padding row and a non-identity + column permutation. This is the one piece of phase-2 maths with no upstream + counterpart, so it gets a direct reference check rather than only the + end-to-end one above. +* ``TestWarmupGradHealth`` -- the indexer parameters, the detached indexer + inputs, and the attention-side ``dq`` / ``dk`` / ``dv``. +* ``TestShortDocumentLayouts`` -- the layouts that starve the *sparse* phase's + indexer cannot starve this one, and phase 2 still equals phase 1 on them. Shared fixtures come from ``hybrid_mla_utils``. @@ -72,45 +84,67 @@ import paddle import paddle.nn.functional as F -import paddlefleet.transformer.mqa_latent_attention as mqa_mod +import paddlefleet.transformer.mha_dsa_warmup_attention as warmup_mod from paddlefleet.transformer.csa_attention import _derive_csa_doc_boundaries from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import ( + HybridMLAIndexerMixin, + latent_mqa_enabled, +) +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, +) from .hybrid_mla_utils import ( _CAPTURED, _GPU, + _WARMUP_TARGETS, V_HEAD_DIM, WINDOW, H, _build_module, - _check_index_invariants, + _build_phase1_dense_module, _create_mqa_config, + _dense_mha_reference, _doc_meta, - _make_inputs, + _flash_attn_version, + _make_dense_inputs, + _production_fa_version, + _rel, _row_end, ) -_EPS = 1e-10 # mqa_latent_attention._EPS, the KL/renormalisation epsilon +_EPS = 1e-10 # mha_dsa_warmup_attention._EPS, the KL/renormalisation epsilon -# The four required layouts, as ``(doc_lens, seqlen)``. ``_row_end`` turns any +# The required layouts, as ``(doc_lens, seqlen)``. ``_row_end`` turns any # trailing gap into one more *valid* document, so "with a trailing gap" is still # a well-formed multi-document layout here; genuine pad rows need # ``_pad_row_end`` and live in ``_PAD_LAYOUTS``. # -# ``seqlen=256`` is a *saturated* budget: ``WINDOW + INDEX_TOPK == 128 + 128 == -# 256``, so at that shape the phase-3 sparse table already covers every row's -# causal length and "attention takes the full causal set" is indistinguishable -# from "attention takes the indexer's top-k". The last two layouts are therefore -# at ``seqlen=512``, where the sparse budget covers at most 256 of a row's up to -# 512 causal columns -- only there does an exact-column-set assertion actually -# discriminate the warmup table from the top-k table. +# The two ``seqlen=512`` entries are kept, with new justifications. They used to +# be the only shapes where an exact-column-set assertion could tell the warmup +# index table apart from the sparse phase's ``WINDOW + INDEX_TOPK == 128 + 128 +# == 256`` budget. Phase 2 builds no index table at all now, so that +# discrimination is gone -- but 512 is still the only shape here that reaches +# two properties: +# +# * a row's causal length exceeds the sparse phase's forced ``WINDOW``, so +# ``_indexer_valid_range(window=0)`` and ``window=WINDOW`` disagree on most +# rows instead of only on the tail -- what ``TestWarmupCandidateRange`` +# measures; +# * ``_dense_attn_target`` runs **more than one row chunk**: ``chunk = +# _TARGET_ROW_SLOTS // s_global == 131072 // 512 == 256 < 512`` +# (``mha_dsa_warmup_attention.py:399-400``), whereas at ``s = 256`` the chunk +# is 512 and the loop body runs once. The ``paddle.concat`` of the parts and +# the per-chunk row offsets are therefore only exercised at 512. _LAYOUTS = [ ([256], 256), # one document spanning the whole buffer ([40, 216], 256), # two documents, tiles the buffer ([100, 50, 106], 256), # three documents, none a multiple of the window ([127, 65], 256), # trailing gap -> a third (valid) document of 64 - ([512], 512), # discriminating: causal length 512 > window+topk = 256 - ([200, 312], 512), # discriminating, multi-document + ([512], 512), # two target row chunks; causal length > WINDOW + ([200, 312], 512), # two target row chunks, multi-document ] # Layouts whose trailing gap stays *outside* every document, i.e. real pad rows. @@ -120,6 +154,24 @@ ([100, 50, 60], 256), # 46 pad rows after three documents ] +# Documents no longer than the sparse phase's forced window. ``cand_i = +# max(causal_len_i - WINDOW, 0)``, so with ``window=WINDOW`` every row of such a +# document has an *empty* candidate range -- the starvation that made the old +# warmup log a KL of exactly 0.0. Phase 2 passes ``window=0``, so the same +# layouts keep their full causal span; and since its attention half is phase 1's, +# they are no longer special for the attention either. Both claims are what +# ``TestShortDocumentLayouts`` keeps regressing. +_SHORT_DOC_LAYOUTS = [ + ([64, 64, 64, 64], 256, "every doc half the window: 256/256 starved"), + ([128, 128], 256, "every doc exactly the window: 256/256 starved"), + ([1] * 8 + [120, 128], 256, "single-token docs mixed with window-sized"), + ([2, 3, 5, 7, 11, 100], 128, "prime tiny docs, all far below the window"), + ([129, 127], 256, "one row past the window: 255/256 starved"), + ([1] * 8 + [248], 256, "single-token docs then one long document"), + ([255, 1], 256, "long document plus a single-token tail"), + ([WINDOW], WINDOW, "s == csa_window_size: the whole buffer is starved"), +] + def _pad_row_end(doc_lens, seqlen): """``[1, 1, s, 1]`` int32 ``row_end`` that produces real pad rows. @@ -145,32 +197,80 @@ def _pad_row_end(doc_lens, seqlen): def _segments(row_end, seqlen): """``[(start, length), ...]`` for every document, from the production - deriver, so the "run this document alone" reference isolates exactly what - the packed kernel is supposed to.""" + deriver, so a per-document perturbation hits exactly the span the kernel is + supposed to keep isolated.""" _, _, _, doc_lens, doc_starts = _derive_csa_doc_boundaries(row_end, seqlen) return list(zip(doc_starts.numpy().tolist(), doc_lens.numpy().tolist())) -def _warmup_module(loss_coeff=0.01, sink=None, sparse_loss=False): - """A ``"mqa_dsa"`` module with the phase switch off *from construction*. +def _valid_range(row_end, seqlen, window=0, position_offset=0, s_local=None): + """``(valid_range, row_empty)`` as numpy, from the production mixin. + + ``HybridMLAIndexerMixin._indexer_valid_range`` touches no instance state -- + it is pure integer arithmetic on the document bounds -- so it is called + unbound with ``None`` for ``self``. That keeps every candidate-range + assertion kernel-free, which is what the old + ``_build_full_causal_indices``-based assertions were: the phase-2 candidate + *range* is the successor of the phase-2 index *table*, the table being what + this change removed. + """ + doc_start, doc_len, is_valid, _, _ = _derive_csa_doc_boundaries( + row_end, seqlen + ) + valid_range, row_empty = HybridMLAIndexerMixin._indexer_valid_range( + None, + seqlen, + doc_start, + doc_len, + is_valid, + window, + position_offset, + s_local, + ) + return valid_range.numpy(), row_empty.numpy() + + +def _warmup_config(loss_coeff=0.01): + """The phase-2 config: ``"mqa_dsa"`` with the phase switch off. - ``sparse_loss=True`` builds the phase-3 module instead, used only as the - control that decides whether a finding is specific to this change. + ``sparse_loss=False`` is not just the loss width any more -- it is what + ``hybrid_mla_indexer.latent_mqa_enabled`` reads to pick the *backend* + (``hybrid_mla_indexer.py:56-59``), so this one kwarg is what makes the module + below a dense ``MHADSAWarmupAttention`` rather than a latent + ``MQALatentAttention``. """ - config = _create_mqa_config("mqa_dsa", loss_coeff=loss_coeff) - config.dsa_indexer_use_sparse_loss = sparse_loss + config = _create_mqa_config( + "mqa_dsa", loss_coeff=loss_coeff, sparse_loss=False + ) config.pad_token_id = 0 + return config + + +def _warmup_module(loss_coeff=0.01, sink=None): + """A phase-2 module, asserted to *be* the dense warmup backend. + + The old form of this assertion was ``module.indexer_use_sparse_loss is + False``. That attribute belonged to ``MQALatentAttention``'s phase dispatch + and ``MHADSAWarmupAttention`` has none -- the phase is now expressed by the + class itself, so assert the class. ``latent_mqa_enabled`` is asserted too, + because it is the production predicate that both this fixture's builder and + ``gpt_layer_specs`` dispatch on (``gpt_layer_specs.py``, + ``hybrid_mla_utils.py:305``). + """ + config = _warmup_config(loss_coeff) + assert not latent_mqa_enabled(config) module = _build_module(config, bf16=True, sink=sink) assert module.indexer is not None - assert module.indexer_use_sparse_loss is sparse_loss + assert isinstance(module, MHADSAWarmupAttention), type(module) return module # Positional signature of ``TileLangCSAIndexerLossAutoScaler.forward`` (minus -# ``ctx``) -- the PyLayer phase 2 now attaches its loss through, imported into -# ``mqa_latent_attention`` from ``csa_attention``. The spy below binds by -# position, so a reordered signature would silently hand it the wrong tensors -# instead of failing. Assert the order rather than trust it. +# ``ctx``) -- the PyLayer phase 2 attaches its loss through, imported into +# ``mha_dsa_warmup_attention`` from ``csa_attention`` +# (``csa_attention.py:1204-1218``). The spy below binds by position, so a +# reordered signature would silently hand it the wrong tensors instead of +# failing. Assert the order rather than trust it. _LOSS_ARGS = [ "output", "target", @@ -194,7 +294,7 @@ def _positional(columns, values=None): descending score, with ``-1`` in the unused slots. That is not position order, so anything compared against a positional reference has to be scattered first. ``values is None`` returns the boolean "row ``i`` scored - column ``c``" table, which is what the old dense ``causal_mask > -1`` was. + column ``c``" table. """ b, s, width = columns.shape dtype = bool if values is None else values.dtype @@ -218,27 +318,28 @@ def _positional(columns, values=None): def _capture_loss_args(): """Capture what the warmup KL is actually reduced over. - One spy, on the ``TileLangCSAIndexerLossAutoScaler`` boundary phase 2 now + One spy, on the ``TileLangCSAIndexerLossAutoScaler`` boundary phase 2 attaches its loss at, because every observable is an argument of that single call: ``P`` (``topk_probs``, already softmaxed by the kernel), ``Q`` - (``target``, from ``_attn_target``), the column ids, the row mask, the - denominator and the coefficient. These are literally the tensors the - backward (upstream's tilelang ``csa_indexer_bwd``) differentiates, so - anything asserted here is what the gradient sees. - - ``cap["probs"]`` / ``cap["target"]`` / ``cap["columns"]`` are in the - kernel's column layout; ``cap["live"]`` and ``cap["dense_target"]`` are the - position-space scatters for tests that compare against a positional - reference. A sum over the last axis -- which is all the KL does -- is - permutation invariant and may be taken on the column layout directly. - - ``cap`` stays empty if the module was not in the warmup phase -- which is - itself the assertion that phase 3 does not reach this code. Phase 3 - attaches through the *same* PyLayer now, so the discriminator is the - ``indexer_backend`` tag: ``"tilelang"`` is phase 2's full-candidate kernel, - ``"cudnn"`` is phase 3's top-k one, and only the former is recorded. + (``target``, from ``_dense_attn_target``), the column ids, the row mask, the + denominator and the coefficient. These are literally the tensors the backward + (upstream's tilelang ``csa_indexer_bwd``) differentiates, so anything + asserted here is what the gradient sees. + + Patched in ``mha_dsa_warmup_attention``'s namespace: phase 3 imports the + *same* PyLayer into ``mqa_latent_attention``, so patching the shared + definition would spy on both phases at once. Rebinding this module's name + makes "the spy fired" mean "phase 2's code path ran", with no need for the + old ``indexer_backend`` discriminator -- which is instead asserted, since + ``"tilelang"`` is the only backend this phase may select + (``mha_dsa_warmup_attention.py:350``). + + ``cap["probs"]`` / ``cap["target"]`` / ``cap["columns"]`` are in the kernel's + column layout; ``cap["live"]`` and ``cap["dense_target"]`` are the + position-space scatters. A sum over the last axis -- which is all the KL does + -- is permutation invariant and may be taken on the column layout directly. """ - real = mqa_mod.TileLangCSAIndexerLossAutoScaler + real = warmup_mod.TileLangCSAIndexerLossAutoScaler actual = [ name for name in inspect.signature(real.forward).parameters @@ -266,26 +367,24 @@ def apply( num_rows_override=None, loss_mask=None, ): - if indexer_backend == "tilelang": - cap["columns"] = topk_indices.numpy().copy() - cap["probs"] = topk_probs.astype("float32").numpy().copy() - cap["target"] = target.astype("float32").numpy().copy() - cap["width"] = int(topk_indices.shape[-1]) - # ``loss_mask`` / ``num_rows_override`` are ``None`` when no - # ``input_ids`` reached the layer -- the same unmasked branch - # ``csa_attention`` takes -- so record that rather than assuming - # a synthesised all-ones mask. - cap["mask"] = ( - None - if loss_mask is None - else loss_mask.astype("float32").numpy().copy() - ) - cap["num_rows"] = ( - None - if num_rows_override is None - else float(num_rows_override) - ) - cap["coeff"] = float(loss_coeff) + cap["backend"] = indexer_backend + cap["columns"] = topk_indices.numpy().copy() + cap["probs"] = topk_probs.astype("float32").numpy().copy() + cap["target"] = target.astype("float32").numpy().copy() + cap["width"] = int(topk_indices.shape[-1]) + # ``loss_mask`` / ``num_rows_override`` are ``None`` when no + # ``input_ids`` reached the layer -- the same unmasked branch + # ``csa_attention`` takes -- so record that rather than assuming a + # synthesised all-ones mask. + cap["mask"] = ( + None + if loss_mask is None + else loss_mask.astype("float32").numpy().copy() + ) + cap["num_rows"] = ( + None if num_rows_override is None else float(num_rows_override) + ) + cap["coeff"] = float(loss_coeff) return real.apply( output, target, @@ -300,40 +399,67 @@ def apply( loss_mask, ) - mqa_mod.TileLangCSAIndexerLossAutoScaler = _Spy + warmup_mod.TileLangCSAIndexerLossAutoScaler = _Spy try: yield cap finally: - mqa_mod.TileLangCSAIndexerLossAutoScaler = real + warmup_mod.TileLangCSAIndexerLossAutoScaler = real if cap: cap["live"] = _positional(cap["columns"]) cap["dense_target"] = _positional(cap["columns"], cap["target"]) -def _forward(module, tensors, row_end, w_v, training, input_ids=None): - """One forward. ``tensors`` is ``(query, key, x, qr)``.""" +def _forward(module, tensors, row_end, training, input_ids=None): + """One forward in the **dense per-head** call shape. + + ``tensors`` is ``(query, key, value, x, qr)`` from ``_make_dense_inputs``. + Three things changed with the backend and all three are load-bearing: + + * a real per-head ``value`` replaces ``v_b_proj_weight``, which the dense + path has no use for -- ``kv_b_proj`` ran before the core attention; + * ``attn_mask_type`` must be passed explicitly, because + ``DotProductAttention`` derives ``is_causal`` from it + (``dot_product_attention.py:562``) and defaults it to ``None``, i.e. *not* + causal. Omitting it silently tests a bidirectional attention; + * ``input_ids`` is accepted only because ``MHADSAWarmupAttention`` declares + ``accepts_input_ids`` and strips it before ``super().forward`` + (``mha_dsa_warmup_attention.py:175-210``). + """ module.train() if training else module.eval() - query, key, x, qr = tensors + query, key, value, x, qr = tensors return module( query, key, - None, + value, None, row_end, - v_b_proj_weight=w_v, + attn_mask_type=AttnMaskType.causal, x=x, qr=qr, input_ids=input_ids, ) +def _phase1_forward(module, tensors, row_end, training): + """The same forward on a plain ``DotProductAttention`` (phase 1).""" + module.train() if training else module.eval() + query, key, value, _, _ = tensors + return module( + query, + key, + value, + None, + row_end, + attn_mask_type=AttnMaskType.causal, + ) + + def _leaves(seqlen, seed=1): - """``(query, key, x, qr)`` as differentiable leaves, plus ``w_v``.""" - query, key, w_v, x, qr = _make_inputs(seqlen, seed=seed, with_hidden=True) - tensors = [query, key, x, qr] + """``[query, key, value, x, qr]`` as differentiable leaves.""" + tensors = list(_make_dense_inputs(seqlen, seed=seed)) for tensor in tensors: tensor.stop_gradient = False - return tensors, w_v + return tensors def _fp32(tensor): @@ -341,51 +467,158 @@ def _fp32(tensor): return tensor.cast("float32").numpy() -class TestWarmupFullCausalTable(unittest.TestCase): - """Row ``i`` of the warmup attention table is exactly ``[doc_start[i], i]``. +class _TargetHost: + """Minimal host for an unbound ``_dense_attn_target`` call. - ``_build_full_causal_indices`` is ``indices = doc_start + offsets`` masked by - ``(indices > positions) | ~is_valid``, a pure integer function of the - document bounds -- no kernel, no float -- so it is assertable as an exact set - equality rather than a bound. Kernel-free, hence not GPU gated. + The method reads exactly one attribute of its host, ``self.softmax_scale`` + (``mha_dsa_warmup_attention.py:427``), and calls no other method, so it can + be evaluated without building an attention -- which keeps + ``TestDenseAttnTarget`` kernel-free and lets it use fp32 inputs instead of + the bf16 the flashmask path requires. """ - def _assert_exact(self, row_end, seqlen): - doc_start, _, is_valid, _, _ = _derive_csa_doc_boundaries( - row_end, seqlen - ) - table = mqa_mod.MQALatentAttention._build_full_causal_indices( - 1, seqlen, doc_start, is_valid - ).numpy() - self.assertEqual(list(table.shape), [1, seqlen, seqlen]) - starts = doc_start.numpy() - valid = is_valid.numpy().astype(bool) - for row in range(seqlen): - cols = table[0, row] - got = set(cols[cols >= 0].tolist()) - want = ( - set(range(int(starts[row]), row + 1)) if valid[row] else set() - ) - self.assertEqual(got, want, f"row {row}: column set is not exact") - # The padding must be right-aligned ``-1``, never interleaved: the - # kernel walks the row until its per-query length runs out. - self.assertEqual( - cols[: len(got)].tolist(), - sorted(got), - f"row {row}: selected columns are not a left-packed run", - ) - self.assertTrue( - bool((cols[len(got) :] == -1).all()), - f"row {row}: non ``-1`` padding after the causal run", + def __init__(self, softmax_scale): + self.softmax_scale = softmax_scale + + +def _dense_target(query, key, columns, doc_start, is_valid, scale, offset=0): + """``_dense_attn_target`` under test, with the shapes it gets in production.""" + s_local, s_global = int(query.shape[1]), int(key.shape[1]) + return MHADSAWarmupAttention._dense_attn_target( + _TargetHost(scale), + query, + key, + columns, + doc_start, + is_valid, + offset, + s_local, + s_global, + ) + + +def _dyadic(shape, seed): + """fp32 values ``k / 8``, ``|k| <= 8`` -- exact in fp32 *and* in TF32. + + Both need 4 mantissa bits, their products 8, and a 256-term dot product of + them stays a multiple of ``1/64`` below 256, i.e. under 15 bits. So the + matmul inside ``_dense_attn_target`` is exact whatever the matmul precision + flag is, and the only residual against a float64 reference is the fp32 + softmax. Random bf16 inputs would instead bury the comparison under a 1e-3 + rounding floor. + """ + rng = np.random.default_rng(seed) + return rng.integers(-8, 9, size=shape).astype("float32") / 8.0 + + +def _fake_columns( + doc_start, is_valid, s_local, s_global, offset=0, reverse=False +): + """A stand-in for the indexer's ``[1, s_local, s_global]`` column table. + + Same contract as the kernel's output: the live slots of a row are its causal + candidate set, left-packed, the rest ``-1``, and an invalid row is all + ``-1``. ``reverse=True`` reverses the live run, which is the only way to + catch a ``_dense_attn_target`` that forgot the ``take_along_axis`` + permutation -- with the identity order, "permuted" and "not permuted" are the + same array. + """ + cols = np.full([s_local, s_global], -1, dtype="int32") + for local_row in range(s_local): + row = offset + local_row + if not bool(is_valid[row]): + continue + live = list(range(int(doc_start[row]), row + 1)) + if reverse: + live.reverse() + cols[local_row, : len(live)] = live + return paddle.to_tensor(cols).unsqueeze(0) + + +def _naive_target(query, key, columns, doc_start, is_valid, scale, offset=0): + """float64 loop reference for ``_dense_attn_target``. + + Deliberately the slowest possible spelling of the docstring's claim: for each + query row, score its causal columns one head at a time, softmax that head in + float64, sum the heads, L1-normalise, then place the result at the slot + ``columns`` names. No chunking, no masking-by-``-inf``, no vectorisation, so + it shares no structure with the implementation. + """ + q = query.numpy()[0].astype("float64") + k = key.numpy()[0].astype("float64") + cols_np = columns.numpy()[0] + s_local, heads = q.shape[0], q.shape[1] + s_global = k.shape[0] + out = np.zeros([s_local, s_global], dtype="float64") + for local_row in range(s_local): + row = offset + local_row + if not bool(is_valid[row]): + continue + live = list(range(int(doc_start[row]), row + 1)) + acc = np.zeros([s_global], dtype="float64") + for head in range(heads): + logits = np.array( + [ + float(q[local_row, head] @ k[col, head]) * scale + for col in live + ] ) - _check_index_invariants(self, table, row_end, seqlen, expect_full=True) + probs = np.exp(logits - logits.max()) + acc[live] += probs / probs.sum() + out[local_row] = acc / max(acc.sum(), _EPS) + permuted = np.zeros_like(out) + for local_row in range(s_local): + slots = cols_np[local_row] + keep = slots >= 0 + permuted[local_row, keep] = out[local_row, slots[keep]] + return permuted + + +class TestWarmupCandidateRange(unittest.TestCase): + """The warmup candidate range is exactly the per-document causal span. + + This is the successor of ``TestWarmupFullCausalTable``. That class asserted + that row ``i`` of the phase-2 ``[b, s, s]`` *index table* held exactly + ``[doc_start[i], i]``; there is no index table any more, and the surviving + carrier of the same claim is the ``valid_range`` phase 2 hands the indexer + kernel -- ``_indexer_valid_range(..., window=0)``, whose two columns are the + inclusive-exclusive bounds of that identical set + (``hybrid_mla_indexer.py:174-183``). Still pure integer arithmetic on the + document bounds, no kernel and no float, hence not GPU gated. + """ + + def _assert_range(self, row_end, seqlen): + doc_start, is_valid = _doc_meta(row_end, seqlen) + valid_range, row_empty = _valid_range(row_end, seqlen, window=0) + self.assertEqual(list(valid_range.shape), [1, seqlen, 2]) + self.assertEqual(list(row_empty.shape), [1, seqlen, 1]) + valid = is_valid.astype(bool) + for row in range(seqlen): + low, high = valid_range[0, row].tolist() + if valid[row]: + self.assertEqual( + [low, high], + [int(doc_start[row]), row + 1], + f"row {row}: candidate range is not [doc_start, i]", + ) + else: + self.assertEqual( + high - low, 0, f"row {row}: a pad row has candidates" + ) + # ``row_empty`` is what zeroes a row's columns and probs + # (``mha_dsa_warmup_attention.py:301-306``), so it must be exactly the + # pad-row predicate: with ``window=0`` every valid row keeps at least its + # own diagonal, so no valid row may be reported empty. + np.testing.assert_array_equal( + row_empty.reshape(-1).astype(bool), ~valid + ) - def test_exact_column_set_all_layouts(self): + def test_candidate_range_all_layouts(self): for layout, seqlen in _LAYOUTS: with self.subTest(layout=layout): - self._assert_exact(_row_end(layout, seqlen), seqlen) + self._assert_range(_row_end(layout, seqlen), seqlen) - def test_exact_column_set_with_pad_rows(self): + def test_candidate_range_with_pad_rows(self): for layout, seqlen in _PAD_LAYOUTS: with self.subTest(layout=layout, pad=True): row_end = _pad_row_end(layout, seqlen) @@ -395,21 +628,199 @@ def test_exact_column_set_with_pad_rows(self): seqlen - sum(layout), "``_pad_row_end`` did not produce the pad rows", ) - self._assert_exact(row_end, seqlen) + self._assert_range(row_end, seqlen) + + def test_window_zero_is_what_saves_the_short_documents(self): + """``window=0`` vs ``window=WINDOW`` on documents shorter than the window. + + The kernel-free half of the old + ``test_window_length_sequence_still_trains_the_indexer_in_warmup``: the + sparse phase clamps the candidate end a full ``csa_window_size`` before + the diagonal, so on these layouts *every* row's range is empty and the KL + is identically 0.0 -- a floor in the very phase where the indexer does all + of its learning. Phase 2 passes ``0`` + (``mha_dsa_warmup_attention.py:283-291``) and keeps every row. Asserting + both sides of the comparison here, rather than building a phase-3 module, + keeps this file free of latent-MQA construction: ``MQALatentAttention`` + with an indexer and ``sparse_loss=False`` is now a hard error state + (``mqa_latent_attention.py:253-288``), so a phase-3 control would only be + exercising another suite's backend. + """ + for layout, seqlen, note in _SHORT_DOC_LAYOUTS: + with self.subTest(layout=layout, note=note): + row_end = _row_end(layout, seqlen) + _, empty_sparse = _valid_range(row_end, seqlen, window=WINDOW) + _, empty_warmup = _valid_range(row_end, seqlen, window=0) + self.assertGreater( + int(empty_sparse.sum()), + 0, + f"{note}: the sparse window starves no row here", + ) + self.assertEqual( + int(empty_warmup.sum()), + 0, + f"{note}: warmup starved a row although window == 0", + ) + self._assert_range(row_end, seqlen) + + def test_range_row_slice_matches_the_global_rows(self): + """The ``s_local`` branch is a pure row slice of the global range. + + ``_indexer_valid_range`` builds over ``s_global`` then slices + ``[position_offset : position_offset + s_local]`` + (``hybrid_mla_indexer.py:184-190``), which is what makes the returned + bounds *global* token ids that the kernel's ``seq_offset`` causal bound + can consume. Cheap to pin here and integer-exact. + """ + seqlen, layout = 256, [40, 216] + row_end = _row_end(layout, seqlen) + full, full_empty = _valid_range(row_end, seqlen, window=0) + half = seqlen // 2 + for offset in (0, half): + with self.subTest(offset=offset): + got, got_empty = _valid_range( + row_end, seqlen, 0, position_offset=offset, s_local=half + ) + self.assertEqual(list(got.shape), [1, half, 2]) + np.testing.assert_array_equal( + got[0], full[0, offset : offset + half] + ) + np.testing.assert_array_equal( + got_empty[0], full_empty[0, offset : offset + half] + ) + + +@_GPU +class TestWarmupIsPhase1Dense(unittest.TestCase): + """Phase 2's attention half **is** phase 1's, and nothing else runs. + + The inversion of the old ``warmup == mqa_full_causal`` bit-identity claim. + ``MHADSAWarmupAttention.forward`` forwards every attention-side argument to + ``DotProductAttention.forward`` untouched and returns its output through a + PyLayer that passes it through unchanged (``mha_dsa_warmup_attention.py: + 175-210``, ``csa_attention.py:1230-1239``), so anything but **bit equality** + against a plain ``DotProductAttention`` means the warmup phase is no longer + "phase 1 plus an indexer loss". + + Three things are asserted together, because each alone is satisfiable for the + wrong reason: + + * bit equality with ``_build_phase1_dense_module``; + * ``_CAPTURED`` empty -- no index table was built and the block-sparse kernel + was never entered. This is the direct successor of the old + ``[b, s, s]``-table assertions: there is nothing left to inspect, and the + absence *is* the property (``hybrid_mla_utils.py:240-253``); + * agreement with the fp32 per-document causal reference, which is what rules + out "both modules are wrong in the same way" -- bit equality against a + broken phase 1 would hold happily. + """ + + @classmethod + def setUpClass(cls): + try: + paddle.set_flags({"FLAGS_cudnn_deterministic": True}) + except Exception: + pass + + def setUp(self): + _CAPTURED.clear() + _WARMUP_TARGETS.clear() + DSAIndexerLossLoggingHelper.tracker.clear() + + def _compare(self, layout, seqlen, training): + config = _warmup_config() + warmup = _build_module(config, bf16=True) + self.assertIsInstance(warmup, MHADSAWarmupAttention) + phase1 = _build_phase1_dense_module(config, bf16=True) + # Both are parameter-free on the attention side (q/k/v arrive already + # projected and no sink is configured), so bit equality needs no weight + # copy -- only the same ``softmax_scale``, which both derive from the same + # config (``dot_product_attention.py:211-215``). + self.assertEqual(warmup.softmax_scale, phase1.softmax_scale) + + tensors = _leaves(seqlen) + row_end = _row_end(layout, seqlen) + _CAPTURED.clear() + _WARMUP_TARGETS.clear() + got = _forward( + warmup, + tensors, + row_end, + training=training, + input_ids=paddle.ones([1, seqlen], dtype="int64"), + ) + want = _phase1_forward(phase1, tensors, row_end, training=training) + self.assertEqual(list(got.shape), [1, seqlen, H * V_HEAD_DIM]) + self.assertEqual( + float(np.abs(_fp32(got) - _fp32(want)).max()), + 0.0, + f"{layout}: warmup output is not bit-identical to phase 1", + ) + self.assertEqual( + _CAPTURED, [], f"{layout}: the block-sparse kernel was reached" + ) + return got, tensors, row_end + + def test_bit_identical_to_phase1_eval(self): + for layout, seqlen in _LAYOUTS: + with self.subTest(layout=layout): + self._compare(layout, seqlen, training=False) + self.assertEqual( + _WARMUP_TARGETS, + [], + "eval built a KL target: the loss is not grad-gated", + ) + + def test_bit_identical_to_phase1_train(self): + for layout, seqlen in _LAYOUTS: + with self.subTest(layout=layout): + got, _, _ = self._compare(layout, seqlen, training=True) + got.cast("float32").sum().backward() + # Exactly one KL target per grad-enabled forward, spanning every + # column: the loss is attached once, not per chunk and not twice + # under the two forwards recompute would run. + self.assertEqual(len(_WARMUP_TARGETS), 1) + self.assertEqual(_WARMUP_TARGETS[0].shape, (1, seqlen, seqlen)) + + def test_output_matches_the_fp32_per_document_causal_reference(self): + """Both modules compute per-document causal MHA, not just the same thing. + + The residual is the bf16 flashmask kernel against an fp32 einsum + reference; measured 1.969e-3 relative at ``s=256``, docs ``[100, 156]``. + A leaked cross-document column or a lost row mask moves this by orders of + magnitude, which is what makes a loose bound sufficient here -- the tight + statement is the bit equality above. + """ + seqlen, layout = 256, [100, 156] + warmup = _warmup_module() + tensors = _leaves(seqlen) + row_end = _row_end(layout, seqlen) + got = _forward(warmup, tensors, row_end, training=False) + want = _dense_mha_reference( + tensors[0], tensors[1], tensors[2], row_end, warmup.softmax_scale + ) + rel = _rel(got, want) + print(f"\n[warmup] dense output vs fp32 reference: rel={rel:.3e}") + self.assertLess(rel, 1e-2) @_GPU class TestWarmupCrossDocumentIsolation(unittest.TestCase): - """Zero cross-document leakage, in the builder's own strong form. - - ``_build_mqa_causal_topk_idxs_from_doc_bounds`` claims a packed batch is - "bit-identical to running each document on its own". Measured on the warmup - path: ``maxabs == 0.0`` exactly for every layout, in both modes -- eval takes - the ``:495`` early exit (indexer projections skipped entirely) and train - takes the ``:600`` branch (indexer runs, for the loss only). Exactness is - expected rather than merely hoped for, because the sparse kernel reduces each - query row over its own listed columns only, and those columns are identical - in the packed and the single-document run once the column ids are shifted. + """Zero cross-document leakage, in the only exact form a dense kernel allows. + + The old mechanism -- run each document alone and demand the packed output + match bitwise -- was sound on the block-sparse kernel, where each query row + reduces over its own listed columns and nothing else. It cannot survive here: + a single-document rerun changes the sequence length, so flashmask picks + different tiles and a different accumulation order, and bf16 reassociation + makes the comparison approximate for reasons that have nothing to do with + leakage. + + The intent survives at full strength with the shapes held fixed instead: + replace every *other* document's K/V with noise and require this document's + rows not to move one bit. Same claim ("this document's output is a function of + this document only"), same exactness, and it is now the *masking* that is + under test rather than the kernel's tiling. """ @classmethod @@ -421,66 +832,73 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() self.module = _warmup_module() def _worst_leak(self, layout, seqlen, training): - tensors, w_v = _leaves(seqlen) - query, key, x, qr = tensors + """Worst per-document deviation, or ``None`` for a one-document layout.""" row_end = _row_end(layout, seqlen) + segments = _segments(row_end, seqlen) + if len(segments) < 2: + return None # nothing to leak *from* + tensors = _leaves(seqlen) DSAIndexerLossLoggingHelper.tracker.clear() - packed = _fp32( - _forward(self.module, tensors, row_end, w_v, training=training) - ) + base = _fp32(_forward(self.module, tensors, row_end, training=training)) worst = 0.0 - for start, length in _segments(row_end, seqlen): + for start, length in segments: sl = slice(start, start + length) + noisy = list(tensors) + paddle.seed(23) + for idx in (1, 2): # key, value + other = paddle.randn(tensors[idx].shape).cast("bfloat16") * 4.0 + other[:, sl] = tensors[idx][:, sl] + noisy[idx] = other DSAIndexerLossLoggingHelper.tracker.clear() - piece = _fp32( - _forward( - self.module, - ( - query[:, sl].contiguous(), - key[:, sl].contiguous(), - x[:, sl].contiguous(), - qr[:, sl].contiguous(), - ), - _row_end([length], length), - w_v, - training=training, - ) + perturbed = _fp32( + _forward(self.module, noisy, row_end, training=training) + ) + worst = max( + worst, float(np.abs(base[:, sl] - perturbed[:, sl]).max()) ) - worst = max(worst, float(np.abs(packed[:, sl] - piece).max())) return worst - def test_packed_equals_single_bitwise_eval(self): + def _check_all(self, training): + measured = 0 for layout, seqlen in _LAYOUTS: with self.subTest(layout=layout): - worst = self._worst_leak(layout, seqlen, training=False) + worst = self._worst_leak(layout, seqlen, training) + if worst is None: + continue + measured += 1 self.assertEqual( - worst, 0.0, f"{layout}: packed != single (eval)" + worst, 0.0, f"{layout}: cross-document leakage" ) + # The single-document layouts of ``_LAYOUTS`` have no other document to + # leak from; the rest must all have been measured. + self.assertEqual(measured, 4) - def test_packed_equals_single_bitwise_train(self): - for layout, seqlen in _LAYOUTS: - with self.subTest(layout=layout): - worst = self._worst_leak(layout, seqlen, training=True) - self.assertEqual( - worst, 0.0, f"{layout}: packed != single (train)" - ) + def test_other_documents_cannot_move_this_one_eval(self): + self._check_all(training=False) + + def test_other_documents_cannot_move_this_one_train(self): + self._check_all(training=True) @_GPU class TestWarmupPadRows(unittest.TestCase): - """Rows outside every document must produce nothing and receive nothing. - - A pad row's table entry is all ``-1`` with a forced per-query length of 1 - (``_build_mqa_causal_topk_idxs_from_doc_bounds`` line 197-198), so the kernel - softmaxes over an empty set. Both the output and ``dq`` must be exactly zero: - a non-zero output would inject padding into the residual stream, and a - non-zero ``dq`` would train on it. Checked with the sink OFF and ON, because - the sink is the one column that survives the masking -- it is value-less, so - it must still contribute nothing. + """Rows outside every document produce nothing and receive nothing. + + Same intent as before, on the new backend and with one more observable. The + old carrier of "a pad row selected no column" was the index table's all-``-1`` + row; there is no table, so the KL's own column table takes over -- phase 2 + forces ``columns`` to ``-1`` and ``probs`` to 0 wherever ``row_empty`` + (``mha_dsa_warmup_attention.py:301-306``), and ``_dense_attn_target`` leaves + those rows all-zero, so the pad rows contribute exactly 0 to the KL as well as + to the attention output. + + Output and ``dq`` must be exactly zero: a non-zero output would inject + padding into the residual stream and a non-zero ``dq`` would train on it. """ @classmethod @@ -492,6 +910,7 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() def _check(self, layout, seqlen, sink): @@ -501,17 +920,32 @@ def _check(self, layout, seqlen, sink): pad = ~is_valid.astype(bool) self.assertTrue(pad.any(), f"{layout}: no pad row was produced") - tensors, w_v = _leaves(seqlen) - _CAPTURED.clear() - out = _forward(module, tensors, row_end, w_v, training=True) + tensors = _leaves(seqlen) + with _capture_loss_args() as cap: + out = _forward( + module, + tensors, + row_end, + training=True, + input_ids=paddle.ones([1, seqlen], dtype="int64"), + ) paddle.seed(7) upstream = paddle.randn([1, seqlen, H * V_HEAD_DIM]).cast("float32") (out.cast("float32") * upstream).sum().backward() - table = _CAPTURED[-1] self.assertTrue( - bool((table[0][pad] == -1).all()), - f"{layout}: a pad row selected a column", + bool((cap["columns"][0][pad] == -1).all()), + f"{layout}: a pad row kept an indexer candidate", + ) + self.assertEqual( + float(np.abs(cap["target"][0][pad]).max()), + 0.0, + f"{layout}: a pad row carries KL target mass", + ) + self.assertEqual( + float(np.abs(cap["probs"][0][pad]).max()), + 0.0, + f"{layout}: a pad row carries indexer probability mass", ) out_np = _fp32(out)[0] self.assertEqual( @@ -529,6 +963,7 @@ def _check(self, layout, seqlen, sink): # the assertions above for the wrong reason. self.assertGreater(float(np.abs(out_np[~pad]).max()), 0.0) self.assertGreater(float(np.abs(dq[~pad]).max()), 0.0) + self.assertGreater(float(np.abs(cap["target"][0][~pad]).max()), 0.0) module.clear_gradients() def test_pad_rows_are_inert_sinkless(self): @@ -537,10 +972,30 @@ def test_pad_rows_are_inert_sinkless(self): self._check(layout, seqlen, sink=None) def test_pad_rows_are_inert_with_sink(self): + """The sink is the one column masking cannot remove, so it keeps its own + case -- under the production ``FLAGS_flash_attn_version``, and skipped + where the installed flashmask cannot take a learnable sink. + + Phase 2 inherits phase 1's sink constraint deliberately ("phase 2 runs + wherever phase 1 runs"): dense MLA with ``add_full_attention_sink_bias`` + requires ``FLAGS_flash_attn_version in (3, 4)`` + (``multi_latent_attention.py:584-614``), which a bare pytest process does + not set, hence the ``_flash_attn_version`` pin. Beyond that, this + fixture's head dims (``K_CHANNELS=256`` / ``V_HEAD_DIM=64``) are outside + FA4's supported set, so ``get_fa_version`` downgrades to FA2 + (``flash_mask_facade.py:39-101``) where ``learnable_sink`` is only + available if the installed kernel takes the kwarg + (``flash_mask_facade.py:104-199``). That is a property of the environment + rather than of this change, so it skips instead of failing. + """ sink = np.linspace(1.0, 3.0, H) - for layout, seqlen in _PAD_LAYOUTS: - with self.subTest(layout=layout): - self._check(layout, seqlen, sink=sink) + with _flash_attn_version(_production_fa_version()): + for layout, seqlen in _PAD_LAYOUTS: + with self.subTest(layout=layout): + try: + self._check(layout, seqlen, sink=sink) + except NotImplementedError as exc: + self.skipTest(f"no learnable_sink support here: {exc}") @_GPU @@ -549,7 +1004,10 @@ class TestWarmupIndexerLossPrecision(unittest.TestCase): Every number below is read at the ``TileLangCSAIndexerLossAutoScaler`` boundary, i.e. exactly what the backward differentiates, and cross-checked - against an independent fp32 recomputation of the logged scalar. + against an independent fp32 recomputation of the logged scalar. Unchanged in + intent from the previous revision; what changed is that there is no attention + index table to compare the KL's column set against, so the reference is the + per-document causal predicate itself. """ @classmethod @@ -561,19 +1019,20 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() def _step(self, module, seqlen, layout, input_ids=None, seed=1): - """One training step; returns ``(logged_kl, captured, attn_table)``.""" - tensors, w_v = _leaves(seqlen, seed=seed) + """One training step; returns ``(logged_kl, captured)``.""" + tensors = _leaves(seqlen, seed=seed) _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() with _capture_loss_args() as cap: out = _forward( module, tensors, _row_end(layout, seqlen), - w_v, training=True, input_ids=input_ids, ) @@ -584,7 +1043,9 @@ def _step(self, module, seqlen, layout, input_ids=None, seed=1): .sum() ) module.clear_gradients() - return logged, cap, _CAPTURED[-1].copy() + # No index table exists in this phase, on any step. + self.assertEqual(_CAPTURED, []) + return logged, cap @staticmethod def _kl_per_row(cap): @@ -594,24 +1055,20 @@ def _kl_per_row(cap): axis=-1 ) - def test_kl_column_set_is_the_attention_column_set_and_no_topk_runs(self): - """One column set serves both consumers, and no cuDNN top-k is called. + def test_kl_spans_the_whole_causal_set_and_no_topk_runs(self): + """The KL's live columns are the per-document causal set, from one + full-candidate tilelang call and no cuDNN top-k. - This replaces the old "two distinct tables" expectation: phase 2 used to - ask the indexer for a *wider* top-k table - (``max(index_topk, min(2048, s//128*128))``) to score the KL over, which - at the production ``index_topk=2048`` silently degenerated into the very - phase-3 table it was supposed to widen. Phase 2 now scores every causal - column instead, so the assertion is an equality between the KL's live - columns and the attention table -- checked element-wise per row, - including the diagonal the old indexer candidate range structurally - excluded (``_indexer_valid_range`` clamped at - ``causal_len - window_size``). + The previous revision asserted an equality between the KL's live columns + and the *attention* index table, both being the full causal set. Phase 2 + has no attention table, so the reference is now the causal predicate + itself -- which is the stronger of the two statements anyway: it does not + depend on a second table being right. Two call counts pin *which* selector produced those columns: the cuDNN top-k kernel (phase 3's) is called zero times, and the tilelang indexer exactly once per step at ``topk_effective == s``, its documented - full-candidate mode. + full-candidate mode (``mha_dsa_warmup_attention.py:292-300``). """ import paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn as fwd_mod import paddlefleet.tilelang_ops as tl_mod @@ -636,32 +1093,23 @@ def recording_tl(*args, **kwargs): for seqlen in (16, 128, 256, 300, 384, 512): with self.subTest(seqlen=seqlen): before = len(tl_widths) - _, cap, attn_table = self._step(module, seqlen, [seqlen]) + _, cap = self._step(module, seqlen, [seqlen]) # The KL table is exactly the causal span -- no rounding. The # tilelang wrapper pads ``topk_effective`` up to its block # internally and crops the result back - # (``csa_indexer_fwd.py:430-462``), so short and non - # -power-of-two lengths are served too; ``seqlen=16`` is below - # the block size on purpose. + # (``csa_indexer_fwd.py:430-462``), so short and + # non-power-of-two lengths are served too; ``seqlen=16`` is + # below the block size on purpose. + self.assertEqual(cap["backend"], "tilelang") self.assertEqual(cap["target"].shape[-1], seqlen) self.assertEqual(cap["probs"].shape[-1], seqlen) self.assertEqual(cap["width"], seqlen) - self.assertEqual(attn_table.shape[-1], seqlen) - # One tilelang call, over every candidate column. self.assertEqual(tl_widths[before:], [seqlen]) - # The KL's live columns are exactly attention's columns. + # One document over the whole buffer, so the causal set of + # row ``i`` is ``[0, i]`` -- a lower-triangular live table. live = cap["live"][0] - for row in range(seqlen): - cols = attn_table[0, row] - self.assertEqual( - set(cols[cols >= 0].tolist()), - set(np.flatnonzero(live[row]).tolist()), - f"s={seqlen} row {row}: KL and attention disagree", - ) - # ... the diagonal included, on every row. - self.assertTrue( - bool(live[np.arange(seqlen), np.arange(seqlen)].all()) - ) + expected = np.tril(np.ones([seqlen, seqlen], dtype=bool)) + np.testing.assert_array_equal(live, expected) finally: fwd_mod.cudnn_indexer_topk_fwd = inner tl_mod.csa_indexer_topk_fwd = inner_tl @@ -681,7 +1129,7 @@ def test_pad_tail_excluded_from_indexer_loss_warmup(self): module = _warmup_module() ids = np.zeros([1, seqlen], dtype="int64") ids[0, :real_tokens] = np.arange(1, real_tokens + 1) - logged, cap, _ = self._step( + logged, cap = self._step( module, seqlen, [seqlen], input_ids=paddle.to_tensor(ids) ) @@ -698,12 +1146,10 @@ def test_pad_tail_excluded_from_indexer_loss_warmup(self): self.assertLess(abs(logged - float(ref)) / abs(float(ref)), 1e-5) # Two discriminators, so this is not a tautology. # (1) the denominator: had ``B*Sq`` driven it, the scalar would be - # 256/200 = 1.28x smaller. (The old form of this check compared the - # masked mean against the *plain* mean of the same rows, which - # discriminated only because the top-k KL was near-zero on the padding - # tail. The full-causal KL is the same order of magnitude on every row -- - # measured 0.7% apart here -- so that comparison no longer separates - # anything and the denominator has to be pinned directly.) + # 256/200 = 1.28x smaller. The full-candidate KL is the same order of + # magnitude on every row, so comparing against the plain mean of the same + # rows would separate nothing -- the denominator has to be pinned + # directly. wrong_denominator = float( (kl_per_row * cap["mask"]).sum() / seqlen * cap["coeff"] ) @@ -721,8 +1167,9 @@ def test_pad_tail_excluded_from_indexer_loss_warmup(self): def test_no_input_ids_uses_the_plain_row_mean(self): """Without ``input_ids`` the reduction is ``kl.mean() * coeff``. - ``_indexer_loss_mask`` returns ``(None, None)`` and ``_forward_warmup`` - passes that straight down, which is the same unmasked branch + ``_indexer_loss_mask`` returns ``(None, None)`` + (``hybrid_mla_indexer.py:203-205``) and ``_attach_indexer_loss`` passes + that straight down, which is the same unmasked branch ``csa_attention._compute_fused_indexer_target`` takes: the backward then falls back to the kernel's own ``1/(B*Sq)``, and only the *logged* scalar carries the ``/cp_size`` CP correction. So what to assert here is that @@ -731,7 +1178,7 @@ def test_no_input_ids_uses_the_plain_row_mean(self): """ seqlen = 256 module = _warmup_module() - logged, cap, _ = self._step(module, seqlen, [seqlen], input_ids=None) + logged, cap = self._step(module, seqlen, [seqlen], input_ids=None) self.assertIsNone(cap["mask"]) self.assertIsNone(cap["num_rows"]) ref = float(self._kl_per_row(cap).mean() * cap["coeff"]) @@ -745,16 +1192,12 @@ def test_kl_target_is_normalised_and_zero_off_the_causal_set(self): """The KL target is L1-normalised per row and exactly zero on columns the per-document causal mask excludes. - The old form of this test asserted zero on ``-1`` top-k slots and - allowed whole rows with *no* candidate at all (the window clamp could - empty a row). Neither exists now: the column set is the causal set, so - every row has at least its own diagonal. A row shorter than the ``s``-wide - table still comes back ``-1``-padded, and those dead slots are where the - "excluded column" assertion now lands. + A row shorter than the ``s``-wide table comes back ``-1``-padded, and + those dead slots are where the "excluded column" assertion lands. """ seqlen, layout = 256, [40, 216] module = _warmup_module() - _, cap, _ = self._step(module, seqlen, layout) + _, cap = self._step(module, seqlen, layout) target = cap["target"][0] masked = cap["columns"][0] < 0 self.assertTrue(masked.any(), "layout masks nothing") @@ -772,27 +1215,25 @@ def test_kl_target_is_normalised_and_zero_off_the_causal_set(self): ) def test_kl_target_is_the_full_causal_attention_distribution(self): - """The intended semantics, measured. - - In warmup both sides of the KL span the whole per-document causal set, so - the target is the head-summed attention distribution over *all* causal - columns, L1-normalised -- no restriction to an indexer candidate subset - and no renormalisation onto it, which is what the previous revision did. - - The reference below is an independent fp32 numpy recomputation (full - per-document causal softmax over all ``s`` columns, then head-summed and - L1-normalised), so agreement is evidence about the semantics, not a - restatement of the code. The residual is the bf16 rounding of the inputs. + """The intended semantics, measured end to end. + + Both sides of the warmup KL span the whole per-document causal set, so the + target is the head-summed *dense* attention distribution over all causal + columns, L1-normalised. The reference below is an independent fp32 + recomputation on the per-head layout -- ``"shd,thd->sht"`` now that K is + per-head rather than the shared latent -- so agreement is evidence about + the semantics rather than a restatement of the code. The residual is the + bf16 rounding of the inputs. """ seqlen, layout = 256, [40, 216] module = _warmup_module() - tensors, w_v = _leaves(seqlen) + tensors = _leaves(seqlen) query, key = tensors[0], tensors[1] row_end = _row_end(layout, seqlen) - _CAPTURED.clear() - DSAIndexerLossLoggingHelper.tracker.clear() with _capture_loss_args() as cap: - out = _forward(module, tensors, row_end, w_v, training=True) + out = _forward( + module, tensors, row_end, training=True, input_ids=None + ) out.cast("float32").sum().backward() module.clear_gradients() @@ -805,9 +1246,9 @@ def test_kl_target_is_the_full_causal_attention_distribution(self): ) scores = ( paddle.einsum( - "shd,td->sht", + "shd,thd->sht", query.detach()[0].cast("float32"), - key.detach().squeeze(2)[0].cast("float32"), + key.detach()[0].cast("float32"), ) * module.softmax_scale ) @@ -820,8 +1261,8 @@ def test_kl_target_is_the_full_causal_attention_distribution(self): reference = head_sum / np.maximum( head_sum.sum(axis=-1, keepdims=True), _EPS ) - # The mask the implementation used must be the same predicate. The - # kernel emits columns score-descending, so the comparison is on the + # The mask the implementation used must be the same predicate. The kernel + # emits columns score-descending, so the comparison is on the # position-space scatter, not on the raw column order. np.testing.assert_array_equal(cap["live"][0], allowed) @@ -837,49 +1278,194 @@ def test_kl_target_is_the_full_causal_attention_distribution(self): self.assertLess(norm_rel, 3e-2) self.assertLess(max_abs, 3e-2) - def test_window_length_sequence_still_trains_the_indexer_in_warmup(self): - """At ``s == csa_window_size`` phase 2 now learns, phase 3 still cannot. - - Pre-existing and unchanged for phase 3: ``_indexer_valid_range`` clamps - the candidate range at ``causal_len - window_size``, so at ``s == window`` - every row's range is empty, the KL is exactly 0 and the indexer learns - nothing that step. The old warmup shared that clamp and logged the same - 0.0 -- a floor in the very phase where the indexer does all of its - learning. ``_forward_warmup`` passes ``window=0`` to - ``_indexer_valid_range`` now, so the candidate range is the whole causal - span and the KL is strictly positive at that shape (measured 9.10e-05); - the phase-3 control still logs 0.0, which is what makes this a property of - the change rather than of the fixture. + def test_window_length_sequence_still_trains_the_indexer(self): + """At ``s == csa_window_size`` the warmup KL is strictly positive. + + The kernel half of the old + ``test_window_length_sequence_still_trains_the_indexer_in_warmup``: with + ``window=0`` every row keeps its full causal span, so the KL is nonzero + (measured 9.10e-05) at the very shape where the sparse phase's clamp + empties every candidate range and logs exactly 0.0. The phase-3 control + that used to sit here is now the kernel-free comparison in + ``TestWarmupCandidateRange + .test_window_zero_is_what_saves_the_short_documents``, because building an + ``MQALatentAttention`` with ``sparse_loss=False`` is a hard error now + (``mqa_latent_attention.py:253-288``). """ - seqlen = WINDOW - sparse = _warmup_module(sparse_loss=True) - logged_sparse, cap_sparse, _ = self._step(sparse, seqlen, [seqlen]) - self.assertEqual(logged_sparse, 0.0) - self.assertEqual(cap_sparse, {}, "phase 3 reached the warmup KL") - module = _warmup_module() - logged, cap, _ = self._step(module, seqlen, [seqlen]) + logged, cap = self._step(module, WINDOW, [WINDOW]) self.assertGreater(logged, 0.0) - self.assertEqual(cap["target"].shape[-1], seqlen) + self.assertEqual(cap["target"].shape[-1], WINDOW) self.assertGreater(float(np.abs(cap["target"]).max()), 0.0) # ... and it keeps scaling with the sequence, as before. - logged_2w, cap_2w, _ = self._step(module, 2 * WINDOW, [2 * WINDOW]) + logged_2w, cap_2w = self._step(module, 2 * WINDOW, [2 * WINDOW]) self.assertGreater(logged_2w, 0.0) self.assertGreater(float(np.abs(cap_2w["target"]).max()), 0.0) +class TestDenseAttnTarget(unittest.TestCase): + """``_dense_attn_target`` against a naive reference. NEW in this revision. + + The KL target builder is the one piece of phase-2 arithmetic with no upstream + counterpart: phase 3 gathers the selected keys and scores them in the kernel's + own column order, which the per-head layout makes impossible here (``[chunk, + width, h, dk]`` is 3.2GB at production shapes), so this phase scores in + *natural* column order and permutes afterwards + (``mha_dsa_warmup_attention.py:380-460``). Everywhere else in this file the + target is only checked end to end, where a wrong permutation and a wrong mask + are hard to tell apart. + + The reference is a float64 triple loop -- one row, one head, one column at a + time -- sharing no structure with the implementation: no chunking, no + ``-inf`` masking, no vectorisation. + + Inputs are ``_dyadic`` fp32 rather than the production bf16, so the matmul is + exact and the only residual is the fp32 softmax; a bf16 fixture would bury the + comparison under a ~1e-3 rounding floor. Kernel-free, hence not GPU gated. + """ + + SEQ, HEADS, DK, SCALE = 24, 4, 8, 0.125 + DOCS, PAD_DOCS = [10, 14], [10, 8] + + def _case(self, row_end, seed=5, offset=0, s_local=None, reverse=False): + """``(got, want, is_valid)`` for one configuration.""" + seqlen = self.SEQ + rows = seqlen if s_local is None else s_local + doc_start, _, is_valid, _, _ = _derive_csa_doc_boundaries( + row_end, seqlen + ) + starts, valid = doc_start.numpy(), is_valid.numpy() + query = paddle.to_tensor(_dyadic([1, rows, self.HEADS, self.DK], seed)) + key = paddle.to_tensor( + _dyadic([1, seqlen, self.HEADS, self.DK], seed + 1) + ) + columns = _fake_columns( + starts, valid, rows, seqlen, offset=offset, reverse=reverse + ) + got = _dense_target( + query, key, columns, doc_start, is_valid, self.SCALE, offset + ) + want = _naive_target( + query, key, columns, starts, valid, self.SCALE, offset + ) + self.assertEqual(list(got.shape), [1, rows, seqlen]) + return got.numpy()[0].astype("float64"), want, valid + + def _assert_close(self, got, want, note): + max_abs = float(np.abs(got - want).max()) + self.assertLess(max_abs, 5e-6, f"{note}: max_abs={max_abs:.3e}") + + def test_matches_the_naive_reference(self): + for label, row_end in ( + ("all rows valid", _row_end(self.DOCS, self.SEQ)), + ("with pad rows", _pad_row_end(self.PAD_DOCS, self.SEQ)), + ): + for reverse in (False, True): + with self.subTest(layout=label, reverse=reverse): + got, want, _ = self._case(row_end, reverse=reverse) + self._assert_close(got, want, f"{label}/{reverse}") + + def test_the_permutation_is_actually_applied(self): + """A reversed column table must produce a reversed target row. + + Without this the identity-ordered comparison above would pass on an + implementation that dropped the ``take_along_axis``: the two orders are + the same array. The kernel really does emit score-descending order, so the + permutation is the normal case, not an edge one. + """ + row_end = _row_end(self.DOCS, self.SEQ) + natural, _, _ = self._case(row_end, reverse=False) + reversed_, _, _ = self._case(row_end, reverse=True) + moved = 0 + for row in range(self.SEQ): + live = int(np.count_nonzero(natural[row])) + if live < 2: + continue + moved += 1 + np.testing.assert_allclose( + reversed_[row, :live], natural[row, :live][::-1], atol=5e-6 + ) + self.assertGreater(moved, 0, "no row had two live columns") + + def test_empty_rows_stay_zero_and_the_rest_sum_to_one(self): + """The all-padding row, and the row-sum contract around it. + + An empty row is the only place the implementation cannot rely on the + softmax: every column is ``-inf``-masked, so ``F.softmax`` returns a + *uniform* row, which is why the mask is re-applied after it and the + normalisation clips its denominator + (``mha_dsa_warmup_attention.py:440-449``). A pad row leaking ``1/s`` per + head would still look normalised, so it has to be pinned at exactly zero: + the KL reduction divides by the valid-row count, not by the row sum. + """ + row_end = _pad_row_end(self.PAD_DOCS, self.SEQ) + got, _, valid = self._case(row_end) + pad = ~valid.astype(bool) + self.assertTrue(pad.any(), "the fixture produced no pad row") + self.assertEqual( + float(np.abs(got[pad]).max()), 0.0, "a pad row is not exactly zero" + ) + sums = got[~pad].sum(axis=-1) + self.assertLess(float(np.abs(sums - 1.0).max()), 5e-6) + + def test_chunking_does_not_change_the_result(self): + """Row chunking is a memory device, not part of the maths. + + ``chunk = max(1, _TARGET_ROW_SLOTS // s_global)`` + (``mha_dsa_warmup_attention.py:399-400``) is 512 rows at the production + ``s = 256`` and 1 row at ``s = 131072``, so in most unit fixtures the loop + body runs exactly once and its row offsets are never exercised. Shrinking + the budget to 2, 5 and 1 rows per chunk is the cheap way to cover them; + with ``_dyadic`` inputs the matmul is exact, so every run must agree + **bitwise** with the single-chunk one, and a mis-indexed chunk shows up as + a shifted row rather than a rounding difference. + """ + row_end = _row_end(self.DOCS, self.SEQ) + whole, _, _ = self._case(row_end) + original = warmup_mod._TARGET_ROW_SLOTS + try: + # 24 columns, so these give chunk = 2, 5 and 1 rows. + for slots in (48, 120, 24): + with self.subTest(row_slots=slots): + warmup_mod._TARGET_ROW_SLOTS = slots + chunked, want, _ = self._case(row_end) + self.assertEqual( + float(np.abs(chunked - whole).max()), + 0.0, + "chunking changed the target", + ) + self._assert_close(chunked, want, f"slots={slots}") + finally: + warmup_mod._TARGET_ROW_SLOTS = original + + def test_context_parallel_row_slice(self): + """``s_local < s_global``: this rank's rows, global columns. + + The CP shape, which the single-card end-to-end tests never reach: rows are + ``[position_offset, position_offset + s_local)`` while the columns and the + document metadata stay global. Cheap to cover here because the method is + pure paddle; the actual CP wiring is asserted in the multi-card suite. + """ + row_end = _row_end(self.DOCS, self.SEQ) + half = self.SEQ // 2 + for offset in (0, half): + with self.subTest(offset=offset): + got, want, _ = self._case(row_end, offset=offset, s_local=half) + self._assert_close(got, want, f"offset={offset}") + self.assertEqual(got.shape, (half, self.SEQ)) + + @_GPU class TestWarmupGradHealth(unittest.TestCase): """Warmup is where the indexer does all of its learning, so a silently gradient-free indexer parameter would waste the whole phase. - Extends ``test_mqa_latent_attention.TestMQADSAWarmupPhase - .test_indexer_gradients_flow_in_the_warmup_phase`` with the attention side - (``dq`` / ``dkv`` / ``d_sink``) in the same step, on a multi-document layout - and with the sink both OFF and ON. ``dkv`` is only checked finite and - non-zero: it carries a known ~2e-3 run-to-run jitter from the atomic - accumulation in ``csa_sparse_attn_bwd_cudnn``, shared with the ordinary - CSA/HCA layouts, so its exact value is not a warmup property. + Same intent as before, with the attention side moved to the dense layout: + ``dq`` / ``dk`` / ``dv`` instead of ``dq`` / ``dkv`` / ``dw_v``, since + ``v_b_proj_weight`` never reaches this backend -- ``kv_b_proj`` has already + materialised per-head K/V. They are only checked finite and non-zero: bf16 + flashmask reduction order gives them a small run-to-run jitter shared with + every other dense layer, so their exact values are not a warmup property. """ @classmethod @@ -891,6 +1477,7 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() def _assert_live(self, name, grad): @@ -907,10 +1494,13 @@ def _assert_live(self, name, grad): def _check(self, sink): seqlen, layout = 256, [40, 216] module = _warmup_module(sink=sink) - tensors, w_v = _leaves(seqlen) - w_v.stop_gradient = False + tensors = _leaves(seqlen) out = _forward( - module, tensors, _row_end(layout, seqlen), w_v, training=True + module, + tensors, + _row_end(layout, seqlen), + training=True, + input_ids=paddle.ones([1, seqlen], dtype="int64"), ) paddle.seed(11) upstream = paddle.randn([1, seqlen, H * V_HEAD_DIM]).cast("float32") @@ -925,17 +1515,19 @@ def _check(self, sink): ("indexer.weights_proj.weight", indexer.weights_proj.linear.weight), ): self._assert_live(name, param.grad) - query, key, x, qr = tensors + query, key, value, x, qr = tensors self._assert_live("dq", query.grad) - self._assert_live("dkv", key.grad) - self._assert_live("dw_v", w_v.grad) + self._assert_live("dk", key.grad) + self._assert_live("dv", value.grad) if sink is not None: self._assert_live("d_sink", module.softmax_offset.grad) self.assertEqual( module.softmax_offset.grad.dtype, module.softmax_offset.dtype ) - # The indexer learns from its own KL only: its inputs stay detached, so - # no indexer gradient may leak into the backbone through ``x`` / ``qr``. + # The indexer learns from its own KL only: its inputs stay detached + # (``hybrid_mla_indexer.py:140-141``), so no indexer gradient may leak + # into the backbone through ``x`` / ``qr`` -- and since the dense + # attention half ignores both, ``None`` is the only correct answer. self.assertIsNone( x.grad, "x.grad is not None: indexer input not detached" ) @@ -949,43 +1541,33 @@ def test_grad_health_sinkless(self): self._check(sink=None) def test_grad_health_with_sink(self): - self._check(sink=np.linspace(1.0, 3.0, H)) - - -# Layouts where the indexer has little or nothing left to choose from, because -# the forced window already covers each document. ``cand_i = max(causal_len_i - -# WINDOW, 0)``, so a document of length <= WINDOW gives *every* one of its rows -# an empty candidate range. -_STARVED_LAYOUTS = [ - ([64, 64, 64, 64], 256, "every doc half the window: 256/256 rows starved"), - ([128, 128], 256, "every doc exactly the window: 256/256 rows starved"), - ([1] * 8 + [120, 128], 256, "single-token docs mixed with window-sized"), - ([2, 3, 5, 7, 11, 100], 128, "prime tiny docs, all far below the window"), - ([129, 127], 256, "one row past the window: 255/256 starved"), - ([1] * 8 + [248], 256, "single-token docs then one long document"), - ([255, 1], 256, "long document plus a single-token tail"), -] + with _flash_attn_version(_production_fa_version()): + try: + self._check(sink=np.linspace(1.0, 3.0, H)) + except NotImplementedError as exc: + self.skipTest(f"no learnable_sink support here: {exc}") @_GPU -class TestStarvedIndexerCandidates(unittest.TestCase): - """Packed documents too short for the indexer to pick anything. - - Real packing is dominated by short documents, so this is not a synthetic - edge: ``_indexer_valid_range`` clamps the candidate end a full - ``csa_window_size`` before the diagonal, so any document no longer than the - window leaves *every* one of its rows with zero candidates. The interesting - question is what that costs, and the answer must be "nothing": - - * the forced window already spans the whole document, so the phase-3 - attention table still contains the complete per-document causal set -- - asserted as ``want - got == empty set``, not as a count; - * consequently all three ``hybrid_mla_attention`` shapes (phase 3, warmup, - ``mqa_full_causal``) must produce **bit-identical** output on these - layouts, which is a much stronger statement than "no crash"; - * an all-``-1`` candidate row means a softmax over an empty set, the classic - NaN source, so output and every gradient are checked finite; - * and no starved row may borrow a column from a neighbouring document. +class TestShortDocumentLayouts(unittest.TestCase): + """Packed documents no longer than the sparse phase's forced window. + + Successor of ``TestStarvedIndexerCandidates``. Real packing is dominated by + short documents, so this is not a synthetic edge; what changed is that phase 2 + cannot be starved by them at all. The three old claims map as follows: + + * "the phase-3 attention table still contains the whole causal set" -> there + is no table; the candidate *range* claim is asserted kernel-free in + ``TestWarmupCandidateRange + .test_window_zero_is_what_saves_the_short_documents``; + * "all three ``hybrid_mla_attention`` shapes agree bitwise" -> the comparison + that matters now is phase 2 against **phase 1**, and it must hold on every + layout rather than only on the fully starved ones, because the two run the + same kernel; + * "an empty candidate row is the classic NaN source, so check finiteness" -> + survives unchanged. It is now the *indexer* rows that can be empty (pad + rows) rather than whole documents, and the KL still divides by a clipped + row sum, so the check is still worth its cost. """ @classmethod @@ -995,111 +1577,70 @@ def setUpClass(cls): except Exception: pass - def _module(self, mode, sparse_loss): - config = _create_mqa_config(mode, loss_coeff=0.01) - if sparse_loss is not None: - config.dsa_indexer_use_sparse_loss = sparse_loss - config.pad_token_id = 0 - return _build_module(config, bf16=True) - - def _run(self, mode, sparse_loss, seqlen, layout, backward): - tensors, w_v = _leaves(seqlen) - row_end = _row_end(layout, seqlen) - module = self._module(mode, sparse_loss) + def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() - out = _forward( - module, - tensors, - row_end, - w_v, - training=backward, - input_ids=paddle.ones([1, seqlen], dtype="int64"), - ) - grads = {} - if backward: - out.cast("float32").sum().backward() - for name, tensor in zip(("query", "key", "x", "qr"), tensors): - grads[name] = tensor.grad - for name, param in module.named_parameters(): - grads[name] = param.grad - return out, grads - @staticmethod - def _starved_rows(row_end, seqlen): - doc_start, _, _, _, _ = _derive_csa_doc_boundaries(row_end, seqlen) - starts = doc_start.numpy() - causal_len = np.arange(seqlen) + 1 - starts - return starts, int((np.maximum(causal_len - WINDOW, 0) == 0).sum()) - - def test_starved_rows_keep_the_whole_causal_set(self): - for layout, seqlen, note in _STARVED_LAYOUTS: - with self.subTest(layout=layout, note=note): + def test_short_documents_still_equal_phase1_bitwise(self): + for layout, seqlen, note in _SHORT_DOC_LAYOUTS: + with self.subTest(note=note): + config = _warmup_config() + warmup = _build_module(config, bf16=True) + phase1 = _build_phase1_dense_module(config, bf16=True) + tensors = _leaves(seqlen) row_end = _row_end(layout, seqlen) - starts, starved = self._starved_rows(row_end, seqlen) - self.assertGreater( - starved, 0, "layout starves no row: the case is not covered" + _CAPTURED.clear() + got = _forward( + warmup, + tensors, + row_end, + training=True, + input_ids=paddle.ones([1, seqlen], dtype="int64"), ) - self._run("mqa_dsa", True, seqlen, layout, backward=False) - table = _CAPTURED[-1][0] - for row in range(seqlen): - cols = table[row] - got = set(cols[cols >= 0].tolist()) - want = set(range(int(starts[row]), row + 1)) - self.assertEqual( - want - got, - set(), - f"row {row}: causal columns missing from the table", - ) - self.assertEqual( - got - want, - set(), - f"row {row}: table has columns outside its document", - ) - - def test_all_three_modes_agree_bitwise_when_starved(self): - for layout, seqlen, note in _STARVED_LAYOUTS: - _, starved = self._starved_rows(_row_end(layout, seqlen), seqlen) - if starved < seqlen: - continue # only fully starved layouts must collapse to equality - with self.subTest(layout=layout, note=note): - phase3, _ = self._run("mqa_dsa", True, seqlen, layout, False) - warmup, _ = self._run("mqa_dsa", False, seqlen, layout, False) - causal, _ = self._run("mqa", None, seqlen, layout, False) + want = _phase1_forward(phase1, tensors, row_end, training=True) self.assertEqual( - float(np.abs(_fp32(phase3) - _fp32(warmup)).max()), + float(np.abs(_fp32(got) - _fp32(want)).max()), 0.0, - "phase 3 != warmup although the window covers everything", + f"{note}: warmup != phase 1", ) - self.assertEqual( - float(np.abs(_fp32(warmup) - _fp32(causal)).max()), - 0.0, - "warmup != mqa_full_causal although the table is identical", + self.assertEqual(_CAPTURED, []) + + def test_short_documents_stay_finite(self): + for layout, seqlen, note in _SHORT_DOC_LAYOUTS: + with self.subTest(note=note): + module = _warmup_module() + tensors = _leaves(seqlen) + out = _forward( + module, + tensors, + _row_end(layout, seqlen), + training=True, + input_ids=paddle.ones([1, seqlen], dtype="int64"), ) - - def test_starved_rows_stay_finite(self): - for mode, sparse_loss in ( - ("mqa_dsa", False), - ("mqa_dsa", True), - ("mqa", None), - ): - for layout, seqlen, note in _STARVED_LAYOUTS: - with self.subTest(mode=mode, sparse=sparse_loss, note=note): - out, grads = self._run( - mode, sparse_loss, seqlen, layout, backward=True + out.cast("float32").sum().backward() + self.assertTrue( + np.isfinite(_fp32(out)).all(), + f"{note}: non-finite output", + ) + grads = [ + (name, tensor.grad) + for name, tensor in zip( + ("query", "key", "value", "x", "qr"), tensors ) - array = _fp32(out) + ] + grads += [ + (name, param.grad) + for name, param in module.named_parameters() + ] + for name, grad in grads: + if grad is None: + continue self.assertTrue( - np.isfinite(array).all(), - "an empty candidate row produced NaN/Inf output", + np.isfinite(grad.cast("float32").numpy()).all(), + f"{note}: gradient {name} is not finite", ) - for name, grad in grads.items(): - if grad is None: - continue - self.assertTrue( - np.isfinite(grad.cast("float32").numpy()).all(), - f"gradient {name} is not finite", - ) + module.clear_gradients() if __name__ == "__main__": diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_recompute_mtp_rope.py b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_recompute_mtp_rope.py index 05ec38fcf0..e820c5a9ca 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_recompute_mtp_rope.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_recompute_mtp_rope.py @@ -15,38 +15,43 @@ """Recompute / MTP / RoPE regression for the phase-2 (warmup) shape of ``hybrid_mla_attention="mqa_dsa"``. -The warmup shape is selected by ``dsa_indexer_use_sparse_loss=False``: attention -consumes the full per-document causal table while the indexer's top-k feeds the -wide KL loss only. ``test_mqa_latent_attention.TestMQADSAWarmupPhase`` covers the -single-forward behaviour of that path; this module covers the three axes that -compose *around* it and were previously only exercised in phase 3 -(``dsa_indexer_use_sparse_loss=True``): +The warmup shape is selected by ``dsa_indexer_use_sparse_loss=False``, and since +that phase has no top-k on *either* side it runs phase 1's **dense MHA** +(``MHADSAWarmupAttention``, a ``DotProductAttention`` subclass) with the indexer +bolted on; the absorbed latent MQA of ``MQALatentAttention`` is phase 3/4 only +(``hybrid_mla_indexer.latent_mqa_enabled``). So there is no ``[b, s, s]`` index +table and no block-sparse call anywhere below -- ``_CAPTURED`` staying empty is +itself an assertion -- and the object the backward consumes is the KL target +(``_WARMUP_TARGETS``) instead. This module covers the three axes that compose +*around* the single forward: 1. **Recompute** (``TestWarmupRecompute``) -- the real production wrapping, ``paddle.distributed.fleet.utils.recompute`` around - ``MQALatentAttention.forward``, which is how ``full_recompute`` wraps - ``_forward_impl``. ON must equal OFF, the ``token_indices`` must be - re-derived bit-identically, and the indexer loss must be attached exactly - once, on the grad-enabled pass. Warmup is strictly stronger than phase 3 - here: the table is ``_build_full_causal_indices``, a pure integer function of - the document bounds, so it is bit-identical even on the single-document - layout whose top-k order phase 3 cannot reproduce. It also *skips the indexer - entirely* on the ``no_grad`` pass (the ``mqa_latent_attention.py:495`` early - exit), which phase 3 does not. + ``MHADSAWarmupAttention.forward``, which is how ``full_recompute`` wraps + ``_forward_impl``. ON must equal OFF, the KL target must be re-derived + bit-identically, and the indexer loss must be attached exactly once, on the + grad-enabled pass: ``_needs_indexer_loss`` gates on + ``paddle.is_grad_enabled()`` (``hybrid_mla_indexer.py:110-121``), so the + ``no_grad`` pass skips the indexer entirely -- which phase 3, whose attention + consumes the ranking, cannot do. That contrast is the discriminator of + ``test_no_grad_pass_skips_the_indexer_only_in_warmup``. 2. **MTP** (``TestWarmupMTP``) -- ``MultiTokenPredictionLayer`` builds its ``transformer_layer`` without passing ``pg_collection``, so the MTP ``-2`` - layer is the same ``MQALatentAttention`` class with the same config; the - warmup shape must therefore be live there too. Plus the tracker denominator: - ``track_indexer_metrics`` now takes the enum string and must count the MTP - ``-2`` entry of ``csa_compress_ratios``. + layer is the same core-attention class reading the same config; the warmup + shape must therefore be live there too, and bit-identical to the phase-1 + dense MTP layer. Plus the tracker denominator: ``track_indexer_metrics`` + takes the enum string and must count the MTP ``-2`` entry of + ``csa_compress_ratios``. 3. **RoPE** (``TestWarmupRope``) -- the switch must not touch RoPE. The main attention's rotary application happens in ``MLASelfAttention`` before the ``mqa_latent`` branch, and the DSA indexer keeps its own plain-RoPE - (``dsa_indexer_rotary_interleaved``) which is still evaluated in warmup even - though attention does not consume its ranking. Also covers the + (``dsa_indexer_rotary_interleaved``), evaluated in warmup even though + attention does not consume its ranking -- both phases reach it through the + shared ``HybridMLAIndexerMixin._indexer_projections``. Also covers the construction-time ``apply_rope_fusion`` x latent-MQA behaviour: the latent - layer downgrades to eager RoPE and warns (non-latent layers keep the global - fusion), plus the ``mqa_latent_rope_fusion=True`` opt-in that fuses it instead. + layer downgrades to eager RoPE and warns while the non-latent layers -- + phase 1, ``mha``, and now phase 2 -- keep the global fusion, plus the + ``mqa_latent_rope_fusion=True`` opt-in that fuses it instead. Every RoPE assertion is against the independent fp64 reference of ``test_hybrid_mla_rope_audit``, never against the implementation itself. @@ -55,10 +60,12 @@ recompute-ON vs recompute-OFF says nothing about whether the *two* forwards of a recomputed step agree with each other, since paddle discards the first one's output. That class keeps both and requires ``maxabs == 0.0``, on all three -``hybrid_mla_attention`` shapes and on a ``seqlen`` where the sparse budget does -not already cover the whole causal range. +``hybrid_mla_attention`` shapes -- the two latent ones on their own call shape, +untouched -- and on a ``seqlen`` where the sparse budget does not already cover +the whole causal range. """ +import types import unittest from unittest import mock @@ -67,6 +74,10 @@ from paddlefleet.transformer.dot_product_attention import DotProductAttention from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, +) from paddlefleet.transformer.multi_latent_attention import ( MLASelfAttention, MLASelfAttentionSublayersSpec, @@ -75,14 +86,16 @@ from .hybrid_mla_utils import ( _CAPTURED, _GPU, + _WARMUP_TARGETS, HIDDEN, INDEX_TOPK, Q_LORA, BiasedLinear, LayerNormStub, _build_module, - _check_index_invariants, + _build_phase1_dense_module, _create_mqa_config, + _make_dense_inputs, _make_inputs, _rel, _row_end, @@ -93,7 +106,6 @@ ref_inv_freq, ref_rope_halfsplit, ) -from .test_mqa_latent_attention import _fp32, _full_causal_table SEQLEN = 256 # Two documents, the second longer than the forced window (so the indexer's @@ -115,10 +127,21 @@ def _prod_csa_ratios(): def _warmup_config(loss_coeff=0.01, **overrides): - """A ``"mqa_dsa"`` config with the phase-2 switch off from construction.""" - config = _create_mqa_config("mqa_dsa", loss_coeff=loss_coeff, **overrides) - config.dsa_indexer_use_sparse_loss = False - return config + """A ``"mqa_dsa"`` config with the phase-2 switch off from construction. + + ``sparse_loss=False`` is what makes ``_build_module`` pick the dense + ``MHADSAWarmupAttention`` backend, through the production predicate + ``hybrid_mla_indexer.latent_mqa_enabled``, so it has to be set *at* + construction rather than assigned afterwards. + """ + return _create_mqa_config( + "mqa_dsa", loss_coeff=loss_coeff, sparse_loss=False, **overrides + ) + + +def _fp32(tensor): + """bf16 -> fp32 numpy; the widening is exact, so bit equality survives.""" + return tensor.cast("float32").numpy() def _leaf(tensor): @@ -127,6 +150,62 @@ def _leaf(tensor): return out +def _dense_call(module, query, key, value, row_end, x, qr, input_ids=None): + """The phase-2 forward call shape. + + Per-head q/k/v (``_make_dense_inputs``), no ``v_b_proj_weight`` (nothing is + absorbed), and an **explicit** ``attn_mask_type``: the inherited + ``DotProductAttention.forward`` resolves ``is_causal`` from this arg and + leaves it False when it is omitted, so dropping it would silently compare a + bidirectional attention against a causal reference. + """ + return module( + query, + key, + value, + None, + row_end, + attn_mask_type=AttnMaskType.causal, + x=x, + qr=qr, + input_ids=input_ids, + ) + + +class _ForwardSpy: + """Count and keep the output of every ``module.forward`` call. + + ``_CAPTURED`` used to double as the forward counter (one entry per + block-sparse call), but the phase-2 backend never reaches that kernel, so + "did recompute really re-forward the layer?" has to be answered by the layer + itself now. Patching the bound method rather than the class keeps the two + concurrent module instances of ``_indexer_call_count`` independent. + """ + + def __init__(self, module): + self.module = module + self.outputs = [] + + def __enter__(self): + real = type(self.module).forward + + def spy(zelf, *args, **kwargs): + result = real(zelf, *args, **kwargs) + tensor = result[0] if isinstance(result, tuple) else result + self.outputs.append(_fp32(tensor.detach()).copy()) + return result + + self.module.forward = types.MethodType(spy, self.module) + return self + + def __exit__(self, *exc_info): + del self.module.forward + return False + + def __len__(self): + return len(self.outputs) + + def _grad_rel(g_on, g_off): if g_on is None and g_off is None: return 0.0 @@ -141,14 +220,14 @@ def _tracker_value(slot): @_GPU class TestWarmupRecompute(unittest.TestCase): - """``recompute(MQALatentAttention.forward)`` in the warmup phase. + """``recompute(MHADSAWarmupAttention.forward)`` in the warmup phase. This is the production wrapping: ``full_recompute`` hands ``TransformerLayer._forward_impl`` to ``paddle.distributed.fleet.utils. recompute``, whose reentrant implementation runs the wrapped callable once under ``no_grad`` (to produce the output) and once more with grad enabled - during backward. Both forwards must agree on the sparsity pattern, or the - backward differentiates a set of columns the forward never used. + during backward. Both forwards must agree, or the backward differentiates + something the forward never produced. """ @classmethod @@ -160,29 +239,35 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() self.module = _build_module(_warmup_config(), bf16=True) + # ``indexer_use_sparse_loss`` is a latent-MQA attribute; the phase-2 + # backend has no such field, the config is the only place the phase + # lives. + self.assertIsInstance(self.module, MHADSAWarmupAttention) self.assertIsNotNone(self.module.indexer) - self.assertFalse(self.module.indexer_use_sparse_loss) + self.assertFalse(self.module.config.dsa_indexer_use_sparse_loss) def _run(self, module, row_end, use_recompute, seed=7): - """One train step, with or without recompute. Returns output + grads.""" + """One train step, with or without recompute. + + Returns ``(output, grads, n_forwards)``. + """ from paddle.distributed.fleet.utils import recompute - query, key, w_v, x, qr = _make_inputs( - SEQLEN, seed=seed, with_hidden=True - ) + query, key, value, x, qr = _make_dense_inputs(SEQLEN, seed=seed) module.train() module.clear_gradients() q = _leaf(query) def fn(qin): - return module( - qin, key, None, None, row_end, v_b_proj_weight=w_v, x=x, qr=qr - ) + return _dense_call(module, qin, key, value, row_end, x, qr) - out = recompute(fn, q) if use_recompute else fn(q) - out.cast("float32").sum().backward() + with _ForwardSpy(module) as spy: + out = recompute(fn, q) if use_recompute else fn(q) + out.cast("float32").sum().backward() + n_forwards = len(spy) grads = { name: (None if p.grad is None else p.grad.detach().cast("float32")) for name, p in module.named_parameters() @@ -190,38 +275,68 @@ def fn(qin): grads["__query__"] = ( None if q.grad is None else q.grad.detach().cast("float32") ) - return out.detach().cast("float32"), grads + return out.detach().cast("float32"), grads, n_forwards + + def _check_target_rows(self, target): + """The KL target is a per-row distribution over the candidate set. + + The phase-2 replacement for ``_check_index_invariants``: there is no + column table left to audit, so audit what the backward actually consumes + instead. ``mha_dsa_warmup_attention.py:414-417`` documents the shape and + the contract -- ``[1, s_local, s_global]``, rows summing to 1, empty + rows all-zero -- and ``:445-450`` is where the explicit zeroing of an + all-masked row happens (a uniform softmax there would poison a KL + reduction that divides by the valid-row count, not by the row sum). + """ + self.assertEqual(list(target.shape), [1, SEQLEN, SEQLEN]) + self.assertTrue(bool((target >= 0.0).all()), "negative target mass") + sums = target[0].sum(axis=-1) + for row in range(SEQLEN): + # ``_row_end`` makes every row valid, so every row is a distr. + self.assertAlmostEqual( + float(sums[row]), + 1.0, + delta=1e-4, + msg=f"target row {row} does not sum to 1", + ) def _equivalence(self, layout): row_end = _row_end(layout, SEQLEN) - expected = _full_causal_table(layout, SEQLEN) _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() - out_off, g_off = self._run(self.module, row_end, use_recompute=False) - idx_off = _CAPTURED[-1] + out_off, g_off, n_fwd_off = self._run( + self.module, row_end, use_recompute=False + ) + target_off = _WARMUP_TARGETS[-1] loss_off = _tracker_value(self.module.layer_number - 1) - n_calls_off = len(_CAPTURED) _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() - out_on, g_on = self._run(self.module, row_end, use_recompute=True) + out_on, g_on, n_fwd_on = self._run( + self.module, row_end, use_recompute=True + ) loss_on = _tracker_value(self.module.layer_number - 1) # The recompute forward really ran (otherwise the rest proves nothing). - self.assertEqual(n_calls_off, 1) + self.assertEqual(n_fwd_off, 1) self.assertGreaterEqual( - len(_CAPTURED), 2, "recompute did not re-forward the layer" + n_fwd_on, 2, "recompute did not re-forward the layer" ) - # Bit-identical index tables: across the two recompute passes, against - # the no-recompute run, and against the analytic table. No tolerance -- - # the warmup table has no floating-point scoring in it at all. - for cap in _CAPTURED: - np.testing.assert_array_equal(cap, idx_off) - np.testing.assert_array_equal(cap, expected) - _check_index_invariants( - self, idx_off, row_end, SEQLEN, expect_full=True + # Phase 2 is dense: no ``[b, s, s]`` table is built and the block-sparse + # kernel is never entered, on either pass. + self.assertEqual( + len(_CAPTURED), 0, "phase 2 reached the block-sparse kernel" ) + # The KL target is built on the grad-enabled pass only, so two forwards + # produce exactly one -- and it is bit-identical to the single-forward + # run's. No tolerance: the recomputed forward re-executes the same + # kernels on the same saved inputs. + self.assertEqual(len(_WARMUP_TARGETS), 1) + np.testing.assert_array_equal(_WARMUP_TARGETS[-1], target_off) + self._check_target_rows(target_off) out_rel = _rel(out_on, out_off) self.assertEqual(set(g_on), set(g_off)) @@ -254,25 +369,65 @@ def test_recompute_equivalence_two_documents(self): self._equivalence(TWO_DOCS) def test_recompute_equivalence_single_document(self): - """The layout phase 3 cannot assert on: warmup is exact here.""" + """The layout phase 3 cannot assert on. + + Phase 3's top-k emission order drifts between two forwards of the same + single full-length document, so it can only compare the selected *set*. + Phase 2 selects nothing, so its target is reproduced exactly. + """ self._equivalence(ONE_DOC) + def _run_latent(self, module, row_end, seed=7): + """One recomputed train step on the **phase-3** latent module. + + Deliberately kept on the latent call shape (``_make_inputs`` plus + ``v_b_proj_weight``): phase 3 is unchanged by the dense-warmup rework + and appears here only as the contrast of + ``test_no_grad_pass_skips_the_indexer_only_in_warmup``. + """ + from paddle.distributed.fleet.utils import recompute + + query, key, w_v, x, qr = _make_inputs( + SEQLEN, seed=seed, with_hidden=True + ) + module.train() + module.clear_gradients() + q = _leaf(query) + + def fn(qin): + return module( + qin, key, None, None, row_end, v_b_proj_weight=w_v, x=x, qr=qr + ) + + with _ForwardSpy(module) as spy: + out = recompute(fn, q) + out.cast("float32").sum().backward() + return len(spy) + def _indexer_call_count(self, use_sparse_loss): """Indexer selector calls across both passes, per backend. - Two selectors have to be counted separately now: phase 3 selects with - the **cuDNN** top-k kernel, phase 2 with the **tilelang** one at - ``topk_effective = s_global`` (its full-candidate mode). Counting only - cuDNN would report "zero top-k calls" for warmup and hide the one call - it does make. + Two selectors have to be counted separately: phase 3 selects with the + **cuDNN** top-k kernel, phase 2 with the **tilelang** one at + ``topk_effective = s_global`` (its full-candidate mode, + ``mha_dsa_warmup_attention.py:292-300``). Counting only cuDNN would + report "zero top-k calls" for warmup and hide the one call it does make. + + Returns ``(n_before_topk, cudnn_topks, tilelang_topks, n_forwards, + n_sparse_kernel_calls)``. """ import paddlefleet.cudnn_ops.indexer.csa_indexer_fwd_cudnn as fwd_mod import paddlefleet.tilelang_ops as tl_mod - config = _create_mqa_config("mqa_dsa", loss_coeff=0.01) - config.dsa_indexer_use_sparse_loss = use_sparse_loss + config = _create_mqa_config( + "mqa_dsa", loss_coeff=0.01, sparse_loss=use_sparse_loss + ) module = _build_module(config, bf16=True) - self.assertEqual(module.indexer_use_sparse_loss, use_sparse_loss) + # The phase switch picks the backend, so the fixture cannot silently + # test the same class twice. + self.assertEqual( + isinstance(module, MHADSAWarmupAttention), not use_sparse_loss + ) cudnn_calls = [] tl_calls = [] @@ -297,55 +452,78 @@ def rec_before(*args, **kwargs): tl_mod.csa_indexer_topk_fwd = rec_tl module.indexer.forward_before_topk = rec_before _CAPTURED.clear() + _WARMUP_TARGETS.clear() + row_end = _row_end(TWO_DOCS, SEQLEN) try: - self._run(module, _row_end(TWO_DOCS, SEQLEN), use_recompute=True) + if use_sparse_loss: + n_forwards = self._run_latent(module, row_end) + else: + n_forwards = self._run(module, row_end, use_recompute=True)[2] finally: fwd_mod.cudnn_indexer_topk_fwd = inner_topk tl_mod.csa_indexer_topk_fwd = inner_tl module.indexer.forward_before_topk = inner_before - return len(before_calls), cudnn_calls, tl_calls, len(_CAPTURED) + return ( + len(before_calls), + cudnn_calls, + tl_calls, + n_forwards, + len(_CAPTURED), + ) def test_no_grad_pass_skips_the_indexer_only_in_warmup(self): """The warmup early exit, observed through the recompute double pass. - Under recompute the layer is forwarded twice: once under ``no_grad`` - (no loss needed) and once grad-enabled. In warmup the first pass takes - the early exit, so the indexer projections run exactly *once* even though - attention was built twice, and the **cuDNN** top-k kernel -- phase 3's - selector -- runs **zero** times: warmup reads no ``index_topk``. What it - does run is one **tilelang** call at ``topk_effective == SEQLEN``, its - full-candidate mode, and exactly one, on the grad-enabled pass only. - Phase 3 has no such exit: attention consumes the ranking, so the - projections and the cuDNN kernel both run on both passes. The contrast is - the discriminator. + Under recompute the layer is forwarded twice: once under ``no_grad`` (no + loss needed) and once grad-enabled. ``_needs_indexer_loss`` is False + on the first one (``hybrid_mla_indexer.py:117-121``), so in warmup the + indexer projections run exactly *once* even though attention was built + twice, and the **cuDNN** top-k kernel -- phase 3's selector -- runs + **zero** times: warmup reads no ``index_topk``. What it does run is one + **tilelang** call at ``topk_effective == SEQLEN``, its full-candidate + mode, and exactly one, on the grad-enabled pass only. Phase 3 has no + such exit: attention consumes the ranking, so the projections and the + cuDNN kernel both run on both passes. The contrast is the discriminator. + + The sparse-kernel counts are the other half of the contrast: phase 3 + enters it once per forward, phase 2 never. """ - n_before_w, cudnn_w, tl_w, n_attn_w = self._indexer_call_count(False) - n_before_s, cudnn_s, tl_s, n_attn_s = self._indexer_call_count(True) + n_before_w, cudnn_w, tl_w, n_fwd_w, n_sparse_w = ( + self._indexer_call_count(False) + ) + n_before_s, cudnn_s, tl_s, n_fwd_s, n_sparse_s = ( + self._indexer_call_count(True) + ) print( f"[warmup indexer calls] warmup before_topk={n_before_w} " - f"cudnn={cudnn_w} tilelang={tl_w} attn_forwards={n_attn_w} || " + f"cudnn={cudnn_w} tilelang={tl_w} forwards={n_fwd_w} " + f"sparse_kernel={n_sparse_w} || " f"phase3 before_topk={n_before_s} cudnn={cudnn_s} " - f"tilelang={tl_s} attn_forwards={n_attn_s}" + f"tilelang={tl_s} forwards={n_fwd_s} sparse_kernel={n_sparse_s}" ) - self.assertGreaterEqual(n_attn_w, 2, "recompute did not re-forward") - self.assertGreaterEqual(n_attn_s, 2, "recompute did not re-forward") + self.assertGreaterEqual(n_fwd_w, 2, "recompute did not re-forward") + self.assertGreaterEqual(n_fwd_s, 2, "recompute did not re-forward") self.assertEqual(n_before_w, 1) self.assertEqual(cudnn_w, [], "warmup called the cuDNN top-k kernel") # One tilelang call, on the grad-enabled pass only, over every column. self.assertEqual(tl_w, [SEQLEN]) + self.assertEqual(n_sparse_w, 0, "warmup ran the block-sparse kernel") # Same wrapping, sparse phase: no early exit, both passes pay for it. self.assertEqual(n_before_s, 2) self.assertEqual(len(cudnn_s), 2) self.assertEqual(cudnn_s, [INDEX_TOPK, INDEX_TOPK]) self.assertEqual(tl_s, [], "phase 3 selected with the tilelang kernel") + self.assertGreaterEqual(n_sparse_s, 2) def _mtp_config(use_sparse_loss, loss_coeff=0.01): """Production MTP shape: 43 backbone layers + 1 next-n predict layer.""" config = _create_mqa_config( - "mqa_dsa", loss_coeff=loss_coeff, num_hidden_layers=43 + "mqa_dsa", + loss_coeff=loss_coeff, + num_hidden_layers=43, + sparse_loss=use_sparse_loss, ) - config.dsa_indexer_use_sparse_loss = use_sparse_loss config.num_nextn_predict_layers = 1 config.pad_token_id = 0 return config @@ -357,9 +535,10 @@ class TestWarmupMTP(unittest.TestCase): ``MultiTokenPredictionLayer.__init__`` builds its ``transformer_layer`` without passing ``pg_collection`` (``multi_token_prediction.py:419-423``), so - the MTP ``-2`` layer is an ordinary ``MQALatentAttention`` reading the same - config -- there is no MTP-specific branch that could keep it on the phase-3 - shape. These tests pin that: same construction as + the MTP ``-2`` layer reads the same config as the backbone ones and its core + attention comes from the same ``latent_mqa_enabled`` dispatch -- there is no + MTP-specific branch that could keep it on the phase-3 latent shape. These + tests pin that: same construction as ``test_hybrid_mla_mtp_layer43_w6._build_mtp_module`` but with the phase-2 switch off. """ @@ -373,6 +552,7 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() DSAIndexerLossLoggingHelper.num_layers = None @@ -387,44 +567,74 @@ def _build(use_sparse_loss=False, loss_coeff=0.01): assert module.layer_number == 0 return module - def _call(self, module, row_end, training=True): - query, key, w_v, x, qr = _make_inputs(SEQLEN, seed=5, with_hidden=True) + @staticmethod + def _call(module, row_end, training=True, input_ids=None): + query, key, value, x, qr = _make_dense_inputs(SEQLEN, seed=5) module.train() if training else module.eval() + return _dense_call(module, query, key, value, row_end, x, qr, input_ids) + + @staticmethod + def _call_phase1(module, row_end): + """Same inputs, on a bare ``DotProductAttention``. + + It does not accept ``input_ids`` (``accepts_input_ids`` is a capability + of the two indexer-owning core attentions only), so the phase-2 call + helper cannot be reused verbatim. + """ + query, key, value, x, qr = _make_dense_inputs(SEQLEN, seed=5) + module.eval() return module( - query, key, None, None, row_end, v_b_proj_weight=w_v, x=x, qr=qr + query, + key, + value, + None, + row_end, + attn_mask_type=AttnMaskType.causal, + x=x, + qr=qr, ) - def test_mtp_attention_is_full_causal_in_warmup(self): - """The MTP ``-2`` layer's attention takes the full causal table, and its - output is bit-identical to the indexer-less ``mqa_full_causal`` MTP - layer -- the same equality the backbone layer has.""" - row_end = _row_end(TWO_DOCS, SEQLEN) - module = self._build() - self.assertFalse(module.indexer_use_sparse_loss) + def test_mtp_warmup_attention_is_the_phase1_dense_layer(self): + """The MTP ``-2`` layer's attention is bit-identical to phase 1's. - reference_cfg = _create_mqa_config("mqa", num_hidden_layers=43) - reference_cfg.num_nextn_predict_layers = 1 - reference = _build_module( - reference_cfg, layer_number=0, bf16=True, is_mtp=True + The inverse of the assertion this test used to make (that the MTP layer + built the same full-causal ``[b, s, s]`` table as an indexer-less + ``mqa_full_causal`` MTP layer): phase 2 delegates its whole attention + half to ``DotProductAttention.forward`` + (``mha_dsa_warmup_attention.py:179-198``) and the indexer only rides on + the output's gradient, so the forward must equal the plain dense layer + exactly -- and no block-sparse call may appear. + """ + row_end = _row_end(TWO_DOCS, SEQLEN) + config = _mtp_config(use_sparse_loss=False) + module = _build_module(config, layer_number=0, bf16=True, is_mtp=True) + self.assertIsInstance(module, MHADSAWarmupAttention) + reference = _build_phase1_dense_module( + config, layer_number=0, bf16=True, is_mtp=True + ) + self.assertNotIsInstance(reference, MHADSAWarmupAttention) + # Phase 2 adds indexer parameters and nothing else, which is why no + # state_dict has to be copied for the comparison below to be meaningful: + # the attention half is parameter-free in this fixture (no sink + # configured, so ``build_softmax_offset`` returns None). + self.assertEqual( + {k for k in module.state_dict() if not k.startswith("indexer.")}, + set(reference.state_dict()), ) - self.assertIsNone(reference.indexer) - - _CAPTURED.clear() - out_ref = _fp32(self._call(reference, row_end, training=False)) - table_ref = _CAPTURED[-1] _CAPTURED.clear() + _WARMUP_TARGETS.clear() out_warm = _fp32(self._call(module, row_end, training=True)) - table_warm = _CAPTURED[-1] + out_ref = _fp32(self._call_phase1(reference, row_end)) - expected = _full_causal_table(TWO_DOCS, SEQLEN) - np.testing.assert_array_equal(table_warm, expected) - np.testing.assert_array_equal(table_warm, table_ref) - _check_index_invariants( - self, table_warm, row_end, SEQLEN, expect_full=True + self.assertEqual( + len(_CAPTURED), 0, "the MTP warmup layer built an index table" + ) + self.assertEqual( + len(_WARMUP_TARGETS), 1, "the MTP warmup layer skipped the indexer" ) maxabs = float(np.max(np.abs(out_warm - out_ref))) - print(f"[warmup mtp] out maxabs vs mqa_full_causal MTP = {maxabs!r}") + print(f"[warmup mtp] out maxabs vs the phase-1 dense MTP = {maxabs!r}") np.testing.assert_array_equal(out_warm, out_ref) def test_mtp_indexer_loss_denominator_counts_the_mtp_minus2_layer(self): @@ -533,28 +743,38 @@ class TestWarmupRope(unittest.TestCase): """The phase-2 switch must not reach RoPE, anywhere. Two independent RoPE users live on a ``-2`` layer: ``MLASelfAttention`` - rotates q/k *before* dispatching to the ``mqa_latent`` branch, and the DSA - indexer keeps its own plain RoPE. The switch changes neither; every - assertion below is against the fp64 reference of - ``test_hybrid_mla_rope_audit``, never against the implementation. + rotates q/k *before* dispatching on ``mqa_latent``, and the DSA indexer + keeps its own plain RoPE. The switch changes neither; every assertion below + is against the fp64 reference of ``test_hybrid_mla_rope_audit``, never + against + the implementation. """ - def _mla(self, mode, use_sparse_loss=True): - config = _create_mqa_config(mode) - config.dsa_indexer_use_sparse_loss = use_sparse_loss + def _mla(self, mode, use_sparse_loss=True, rope_fusion=False): + config = _create_mqa_config(mode, sparse_loss=use_sparse_loss) + # The shared fixture pins fusion off (``hybrid_mla_utils.py:214``); + # the production phase-1/phase-2 yamls turn it on, so one test needs + # to opt back in. + config.apply_rope_fusion = rope_fusion paddle.seed(123) return MLASelfAttention( config=config, sublayers_spec=_MLA_SPEC, layer_number=1 ) def test_main_attention_rope_is_untouched_by_the_phase_switch(self): - """warmup q/k == phase-3 q/k == the ``mha`` rope sub-blocks, exactly. + """warmup q/k == the ``mha`` q/k exactly, and == phase 3's rope block. ``get_query_key_value_tensors`` is where the rotation happens. Sharing one ``state_dict`` across the three modes (the key sets are identical -- that is the whole point of activation-level absorption) makes the comparison meaningful, and the result must be bit-equality, since the switch is not read on this code path at all. + + Phase 2 is now *non*-latent (``latent_mqa_enabled`` is False for + ``mqa_dsa`` + ``sparse_loss=False``), which strengthens the first half + of this test from "the rope sub-block matches" to "the whole q and k + match phase 1's"; the cross-layout rope-sub-block comparison moves to + warmup-vs-phase-3, where the shapes genuinely differ. """ mha = self._mla("mha") warm = self._mla("mqa_dsa", use_sparse_loss=False) @@ -565,7 +785,8 @@ def test_main_attention_rope_is_untouched_by_the_phase_switch(self): warm.set_state_dict(state) ph3.set_state_dict(state) self.assertFalse(mha.mqa_latent) - self.assertTrue(warm.mqa_latent) + self.assertFalse(warm.mqa_latent) + self.assertTrue(ph3.mqa_latent) paddle.seed(7) hidden = paddle.randn([1, 64, HIDDEN]) * 0.5 @@ -577,19 +798,21 @@ def test_main_attention_rope_is_untouched_by_the_phase_switch(self): rope_dim = mha.config.hybrid_mla_qk_rope_head_dim pairs = { - "warm_vs_ph3_q": (out["warm"][0], out["ph3"][0]), - "warm_vs_ph3_k": (out["warm"][1], out["ph3"][1]), + # Phase 2 runs phase 1's dense path, so this is full-tensor + # equality, not a sub-block one. + "mha_vs_warm_q": (out["mha"][0], out["warm"][0]), + "mha_vs_warm_k": (out["mha"][1], out["warm"][1]), # The latent q/k carry the rope block in their trailing dims; the - # nope halves differ in shape between mha and latent, the rope - # sub-block does not. ``mha`` keeps K per head, the latent path - # keeps the single shared head, so take head 0 on both. - "mha_vs_warm_q_pe": ( - out["mha"][0][..., -rope_dim:], + # nope halves differ in shape between dense and latent, the rope + # sub-block does not. The dense path keeps K per head, the latent + # path keeps the single shared head, so take head 0 on both. + "warm_vs_ph3_q_pe": ( out["warm"][0][..., -rope_dim:], + out["ph3"][0][..., -rope_dim:], ), - "mha_vs_warm_k_pe": ( - out["mha"][1][:, :, :1, -rope_dim:], + "warm_vs_ph3_k_pe": ( out["warm"][1][:, :, :1, -rope_dim:], + out["ph3"][1][:, :, :1, -rope_dim:], ), } measured = {} @@ -601,9 +824,61 @@ def test_main_attention_rope_is_untouched_by_the_phase_switch(self): a.numpy(), b.numpy(), err_msg=f"{name} moved" ) + def test_warmup_matches_phase1_under_fused_rope(self): + """The same q/k equality, but with ``apply_rope_fusion`` actually on. + + The fixture family runs with fusion off + (``hybrid_mla_utils.py:214``), so without this the "phase 2 == phase 1" + claim was only ever measured on the eager kernel -- while the production + phase-2 yaml sets ``apply_rope_fusion: true`` to match the baseline. The + two are mathematically equal and *not* bit-identical, so a fixture that + silently kept eager could not have caught a production config where one + phase fuses and the other does not (which is exactly what the phase-2 + yaml carried until this rework: ``false`` against the baseline's + ``true``). + + Phase 3 is deliberately absent here: it is latent, so it downgrades + itself to eager (``test_apply_rope_fusion_downgrades_on_latent_mqa``) + and a bitwise comparison against a fused tensor would be meaningless. + """ + mha = self._mla("mha", rope_fusion=True) + warm = self._mla("mqa_dsa", use_sparse_loss=False, rope_fusion=True) + self.assertFalse(mha.mqa_latent) + self.assertFalse(warm.mqa_latent) + self.assertTrue(mha.config.apply_rope_fusion) + self.assertTrue(warm.config.apply_rope_fusion) + warm.set_state_dict(mha.state_dict()) + + paddle.seed(7) + hidden = paddle.randn([1, 64, HIDDEN]) * 0.5 + # Training mode is mandatory here: the fused MLA RoPE path raises + # ``NotImplementedError: apply_rope_fusion does not support dynamic + # inference yet`` under ``eval()`` + # (``multi_latent_attention.py:1808-1812``). Training is the only mode + # the three phases are ever run in, so that gap is out of scope. + self.assertTrue(mha.training and warm.training) + q_ref, k_ref = mha.get_query_key_value_tensors(hidden)[:2] + q_got, k_got = warm.get_query_key_value_tensors(hidden)[:2] + print( + "[warmup rope fused] q maxabs=" + f"{float((q_got - q_ref).abs().max()):.3e} k maxabs=" + f"{float((k_got - k_ref).abs().max()):.3e}" + ) + np.testing.assert_array_equal( + q_got.numpy(), q_ref.numpy(), err_msg="fused q moved" + ) + np.testing.assert_array_equal( + k_got.numpy(), k_ref.numpy(), err_msg="fused k moved" + ) + def test_indexer_rope_is_plain_and_correct_in_warmup(self): """The indexer's own RoPE, in warmup, against the fp64 reference. + Reached off the dense phase-2 backend now (``_build_module`` returns an + ``MHADSAWarmupAttention``), but it is the very same ``DSAIndexer`` + instance: both phases build it from ``_indexer_layer_spec`` and call it + through ``HybridMLAIndexerMixin._indexer_projections``. + Three things at once: the frequency table is plain RoPE (base 10000, not the compressed layers' YaRN / ``csa_compress_rotary_base``), the ``dsa_indexer_rotary_interleaved`` layout switch is live (each setting @@ -695,15 +970,20 @@ def test_indexer_rope_output_matches_phase_three_bitwise(self): """``forward_before_topk`` q/k are identical across the two phases. The indexer still runs in warmup; only the *consumption* of its ranking - changes. Copying one ``state_dict`` across, the pre-top-k activations -- - which is where all the RoPE is -- must be bit-equal. + changes -- and, since the rework, the class that owns it: phase 2's + ``MHADSAWarmupAttention`` and phase 3's ``MQALatentAttention`` build the + same ``DSAIndexer`` under the same ``indexer.*`` names, which is what + makes copying one ``state_dict`` across legal. The pre-top-k activations + -- which is where all the RoPE is -- must then be bit-equal. """ seqlen = 64 warm = _build_module(_warmup_config(), bf16=True) ph3 = _build_module( _create_mqa_config("mqa_dsa", loss_coeff=0.01), bf16=True ) + self.assertIsInstance(warm, MHADSAWarmupAttention) self.assertTrue(ph3.indexer_use_sparse_loss) + self.assertEqual(set(warm.state_dict()), set(ph3.state_dict())) ph3.set_state_dict(warm.state_dict()) paddle.seed(11) @@ -732,16 +1012,18 @@ def test_apply_rope_fusion_downgrades_on_latent_mqa(self): downgrades *itself* to eager RoPE and warns, so the non-latent HCA/CSA layers of the same model keep the global fusion. This checks: - - both latent modes (``mqa_full_causal``, ``mqa_dsa``, incl. the warmup - pairing) construct, stay ``mqa_latent=True`` and resolve the per-layer - decision (``apply_rope_fusion and not mqa_latent``) to False (eager), - emitting the downgrade warning; - - the ``mha`` positive control stays non-latent and keeps fusion enabled - -- proving the downgrade is scoped to latent MQA, not "fusion is broken - here"; + - both latent modes (``mqa_full_causal`` and the *sparse* pairing of + ``mqa_dsa``) construct, stay ``mqa_latent=True`` and resolve the + per-layer decision (``apply_rope_fusion and not mqa_latent``) to + False (eager), emitting the downgrade warning; + - the non-latent controls -- ``mha`` and, since the rework, phase 2 + (``mqa_dsa`` + ``sparse_loss=False``) -- keep fusion on and emit + no warning, proving the downgrade is scoped to *absorption*, not to + "``mqa_dsa`` breaks fusion". Phase 2 running on the dense path is + exactly why it keeps it; - ``mqa_latent_rope_fusion=True`` is the opt-in alternate path: the - latent layer constructs without the downgrade warning because it takes - the fused rotate_half branch instead. + latent layer constructs without the downgrade warning because it + takes the fused rotate_half branch instead. """ import paddlefleet.transformer.multi_latent_attention as _mla @@ -751,48 +1033,49 @@ def _effective(module, config): _DOWNGRADE = "has no effect on the RoPE" + def _build(mode, sparse): + config = _create_mqa_config(mode, sparse_loss=sparse) + config.apply_rope_fusion = True + with mock.patch.object(_mla.logger, "warning") as warn: + module = MLASelfAttention( + config=config, + sublayers_spec=_MLA_SPEC, + layer_number=1, + ) + warned = any( + _DOWNGRADE in str(c.args[0]) for c in warn.call_args_list + ) + return config, module, warned + # ---- latent modes downgrade to eager and warn ---- - for mode, sparse in ( - ("mqa", True), - ("mqa_dsa", True), - ("mqa_dsa", False), - ): + for mode, sparse in (("mqa", True), ("mqa_dsa", True)): with self.subTest( mode=mode, use_sparse_loss=sparse, path="downgrade" ): - config = _create_mqa_config(mode) - config.dsa_indexer_use_sparse_loss = sparse - config.apply_rope_fusion = True - with mock.patch.object(_mla.logger, "warning") as warn: - module = MLASelfAttention( - config=config, - sublayers_spec=_MLA_SPEC, - layer_number=1, - ) + config, module, warned = _build(mode, sparse) self.assertTrue(module.mqa_latent) self.assertFalse(_effective(module, config)) self.assertTrue( - any( - _DOWNGRADE in str(c.args[0]) - for c in warn.call_args_list - ), - f"{mode}: expected the eager-downgrade warning", + warned, f"{mode}: expected the eager-downgrade warning" ) - # ---- mha control: non-latent keeps the global fusion ---- - with self.subTest(mode="mha", path="enabled"): - config = _create_mqa_config("mha") - config.apply_rope_fusion = True - module = MLASelfAttention( - config=config, sublayers_spec=_MLA_SPEC, layer_number=1 - ) - self.assertFalse(module.mqa_latent) - self.assertTrue(_effective(module, config)) + # ---- non-latent controls keep the global fusion, silently ---- + for mode, sparse in (("mha", True), ("mqa_dsa", False)): + with self.subTest( + mode=mode, use_sparse_loss=sparse, path="enabled" + ): + config, module, warned = _build(mode, sparse) + self.assertFalse(module.mqa_latent) + self.assertTrue(_effective(module, config)) + self.assertFalse( + warned, + f"{mode}: a non-latent layer must not be downgraded", + ) # ---- opt-in fused rotate_half path: no downgrade warning ---- for mode in ("mqa", "mqa_dsa"): with self.subTest(mode=mode, path="mqa_latent_rope_fusion"): - config = _create_mqa_config(mode) + config = _create_mqa_config(mode, sparse_loss=True) config.apply_rope_fusion = True config.mqa_latent_rope_fusion = True with mock.patch.object(_mla.logger, "warning") as warn: @@ -827,6 +1110,12 @@ class TestRecomputeInnerForwardBitIdentical(unittest.TestCase): on the same saved inputs, and the index table is an integer function of the document bounds. The captured-call count is asserted first, otherwise a single-forward implementation would make the comparison vacuous. + + The output comparison covers all three ``hybrid_mla_attention`` shapes; the + *index table* comparison only applies to the two latent ones, because the + phase-2 backend has no index table. Its analogue is the KL target, which + exists once per recomputed step (the no-grad forward attaches no loss), so + for phase 2 the invariant is a count, not a diff. """ @classmethod @@ -836,45 +1125,56 @@ def setUpClass(cls): except Exception: pass - # ``mqa`` -> "mqa_full_causal"; the two ``mqa_dsa`` rows are warmup (wide - # loss, full-causal attention) and phase 3 (narrow loss, top-k attention). - _MODES = (("mqa_dsa", False), ("mqa_dsa", True), ("mqa", None)) + # ``mqa`` -> "mqa_full_causal"; ``("mqa_dsa", True)`` is phase 3 (narrow + # loss, top-k attention). Both are latent MQA and keep the latent call + # shape. ``("mqa_dsa", False)`` -- phase 2 -- is dense MHA and is listed + # apart because only the output half of this class applies to it. + _LATENT_MODES = (("mqa_dsa", True), ("mqa", None)) + _WARMUP_MODE = ("mqa_dsa", False) + _MODES = (_WARMUP_MODE, *_LATENT_MODES) # 256 saturates window+topk, 512 does not -- see the module docstring of # ``test_hybrid_mla_warmup_doc_mask_loss``. _SHAPES = ((SEQLEN, TWO_DOCS), (512, [200, 312])) def _module(self, mode, sparse_loss): - config = _create_mqa_config(mode, loss_coeff=0.01) - if sparse_loss is not None: - config.dsa_indexer_use_sparse_loss = sparse_loss + """``sparse_loss`` has to reach ``_create_mqa_config`` as a kwarg. + + It selects the attention backend through ``latent_mqa_enabled``, so + assigning ``config.dsa_indexer_use_sparse_loss`` after construction + would build the wrong class. + """ + kwargs = {} if sparse_loss is None else {"sparse_loss": sparse_loss} + config = _create_mqa_config(mode, loss_coeff=0.01, **kwargs) config.pad_token_id = 0 return _build_module(config, bf16=True) def _capture_two_forwards(self, module, seqlen, layout): - """Run one recomputed train step; return both forwards' outputs.""" - import types + """Run one recomputed train step; return both forwards' outputs. + The call shape follows the backend the production dispatch picked + (``latent_mqa_enabled`` inside ``_build_module``): the latent rows keep + the absorbed signature, phase 2 gets the per-head dense one. + """ from paddle.distributed.fleet.utils import recompute - query, key, w_v, x, qr = _make_inputs(seqlen, seed=7, with_hidden=True) + dense = isinstance(module, MHADSAWarmupAttention) + input_ids = paddle.ones([1, seqlen], dtype="int64") row_end = _row_end(layout, seqlen) - module.train() - module.clear_gradients() - q = _leaf(query) + if dense: + query, key, value, x, qr = _make_dense_inputs(seqlen, seed=7) - outs = [] - real = type(module).forward + def _call(qin): + return _dense_call( + module, qin, key, value, row_end, x, qr, input_ids + ) - def spy(zelf, *args, **kwargs): - result = real(zelf, *args, **kwargs) - tensor = result[0] if isinstance(result, tuple) else result - outs.append(tensor.detach().cast("float32").numpy().copy()) - return result + else: + query, key, w_v, x, qr = _make_inputs( + seqlen, seed=7, with_hidden=True + ) - module.forward = types.MethodType(spy, module) - try: - out = recompute( - lambda qin: module( + def _call(qin): + return module( qin, key, None, @@ -883,20 +1183,24 @@ def spy(zelf, *args, **kwargs): v_b_proj_weight=w_v, x=x, qr=qr, - input_ids=paddle.ones([1, seqlen], dtype="int64"), - ), - q, - ) + input_ids=input_ids, + ) + + module.train() + module.clear_gradients() + q = _leaf(query) + + with _ForwardSpy(module) as spy: + out = recompute(_call, q) out.cast("float32").sum().backward() - finally: - del module.forward - return outs + return spy.outputs def test_inner_forward_output_is_bit_identical(self): for mode, sparse_loss in self._MODES: for seqlen, layout in self._SHAPES: with self.subTest(mode=mode, sparse=sparse_loss, s=seqlen): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() module = self._module(mode, sparse_loss) outs = self._capture_two_forwards(module, seqlen, layout) @@ -914,7 +1218,8 @@ def test_inner_forward_output_is_bit_identical(self): ) def test_inner_forward_index_table_is_bit_identical(self): - for mode, sparse_loss in self._MODES: + """Latent MQA only -- phase 2 has no index table to compare.""" + for mode, sparse_loss in self._LATENT_MODES: for seqlen, layout in self._SHAPES: with self.subTest(mode=mode, sparse=sparse_loss, s=seqlen): _CAPTURED.clear() @@ -931,6 +1236,38 @@ def test_inner_forward_index_table_is_bit_identical(self): "the recomputed forward selected different columns", ) + def test_warmup_recompute_has_no_table_and_one_kl_target(self): + """The phase-2 counterpart of the index-table comparison above. + + Two forwards happen, yet the block-sparse kernel is never reached and + exactly one KL target is built: the no-grad forward is gated out by + ``_needs_indexer_loss`` (``hybrid_mla_indexer.py:110-121``), so the + recomputed step pays for the indexer once rather than twice. + """ + mode, sparse_loss = self._WARMUP_MODE + for seqlen, layout in self._SHAPES: + with self.subTest(s=seqlen): + _CAPTURED.clear() + _WARMUP_TARGETS.clear() + DSAIndexerLossLoggingHelper.tracker.clear() + module = self._module(mode, sparse_loss) + self.assertIsInstance(module, MHADSAWarmupAttention) + outs = self._capture_two_forwards(module, seqlen, layout) + self.assertGreaterEqual(len(outs), 2, "no second forward") + self.assertEqual( + len(_CAPTURED), + 0, + "the warmup phase reached the block-sparse kernel", + ) + self.assertEqual( + len(_WARMUP_TARGETS), + 1, + "the KL target was built on the no-grad forward too", + ) + self.assertEqual( + tuple(_WARMUP_TARGETS[-1].shape), (1, seqlen, seqlen) + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/single_card_tests/transformer/test_mqa_latent_attention.py b/tests/single_card_tests/transformer/test_mqa_latent_attention.py index ba9857dc0f..c432fa49f3 100644 --- a/tests/single_card_tests/transformer/test_mqa_latent_attention.py +++ b/tests/single_card_tests/transformer/test_mqa_latent_attention.py @@ -14,18 +14,28 @@ """Unit tests for :mod:`paddlefleet.transformer.mqa_latent_attention`. -``hybrid_mla_attention`` set to ``"mqa_dsa"`` or ``"mqa_full_causal"`` turns the -hybrid MLA (``csa_compress_ratios == -2``) layers of a ``dsv4_hybrid`` model into -:class:`MQALatentAttention` (latent MQA). The module picks its path from the -sublayers spec, not from any config string: +``hybrid_mla_attention`` decides which core attention the hybrid MLA +(``csa_compress_ratios == -2``) layers of a ``dsv4_hybrid`` model run. +:class:`MQALatentAttention` (latent MQA) owns exactly the two modes that attend +to the **absorbed KV latent**, i.e. the ones that consume a sorted candidate +set, and it picks between them from the sublayers spec rather than from any +config string: * ``MQALatentAttentionSublayersSpec(indexer=None)`` -- per-document full-causal attention on the latent, mathematically equal to MHA. This is what production - builds for ``"mqa_full_causal"``; ``gpt_layer_specs`` always attaches an - indexer for ``"mqa_dsa"``. The absorption-equivalence tests here drive it by - constructing the layer directly with ``indexer=None``. -* an indexer spec -- forced local window + Lightning-indexer top-k, i.e. DSA on - the KV latent. + builds for ``"mqa_full_causal"``. The absorption-equivalence tests here drive + it by constructing the layer directly with ``indexer=None``. +* an indexer spec **plus** ``dsa_indexer_use_sparse_loss=True`` (phase 3/4) -- + forced local window + Lightning-indexer top-k, i.e. DSA on the KV latent. + +The other two modes are not this class, and +``hybrid_mla_indexer.latent_mqa_enabled`` is the single predicate that keeps the +spec dispatch and ``MLASelfAttention.mqa_latent`` in step about it: ``"mha"`` +and ``"mqa_dsa"`` + ``dsa_indexer_use_sparse_loss=False`` (phase 2, the DSA +warmup) both run dense per-head attention, the latter in +``mha_dsa_warmup_attention.MHADSAWarmupAttention``. An indexer on this class +with the sparse loss off is therefore an **error state**, not a phase, and +``_phase()`` raises for it. Coverage: 1. Guards -- unsupported configurations fail loudly (no GPU needed). @@ -45,9 +55,12 @@ path) and takes a finite non-zero fp32 gradient. There is a single sink switch now, so the config no longer rejects any combination. 7. The phase-2 (warmup) shape of ``"mqa_dsa"``, selected by - ``dsa_indexer_use_sparse_loss=False``: attention consumes the full - per-document causal table (bit-identical to ``"mqa_full_causal"``) while - the indexer's top-k serves the wide KL loss only. + ``dsa_indexer_use_sparse_loss=False``: it is the **dense** backend + ``MHADSAWarmupAttention``, bit-identical to phase 1's attention, building + no index table at all while the indexer's KL spans every causal column. + Plus the two guards that keep the split honest: ``latent_mqa_enabled`` over + all four config combinations, and ``MQALatentAttention`` refusing to run + that phase itself. 8. Migration: the renamed config keys (``non_absorbed_mqa*``, ``csa_train_indexer_only``, ``csa_indexer_init_from_scratch``) ship without an alias, so a stale config must raise rather than be absorbed into a @@ -74,9 +87,18 @@ _derive_csa_doc_boundaries, ) from paddlefleet.transformer.dsa_attention import DSAIndexerLossLoggingHelper +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.hybrid_mla_indexer import ( + HybridMLAIndexerMixin, + latent_mqa_enabled, +) +from paddlefleet.transformer.mha_dsa_warmup_attention import ( + MHADSAWarmupAttention, +) from paddlefleet.transformer.mqa_latent_attention import ( _LSE_INDEXER_TOPKS, MQALatentAttention, + MQALatentAttentionSublayersSpec, _HashableTensor, ) from paddlefleet.transformer.transformer_config import TransformerConfig @@ -84,6 +106,7 @@ from .hybrid_mla_utils import ( _CAPTURED, _GPU, + _WARMUP_TARGETS, DK, DV, HIDDEN, @@ -96,9 +119,13 @@ WINDOW, H, _build_module, + _build_phase1_dense_module, _check_index_invariants, _create_mqa_config, + _dense_mha_reference, _dense_reference, + _indexer_layer_spec, + _make_dense_inputs, _make_inputs, _rel, _row_end, @@ -123,6 +150,10 @@ def _full_causal_table(layout, seqlen): """The per-document full-causal ``[1, s, s]`` table, from the production builder itself -- it is a pure integer function of the document bounds. + + Still the attention table of the indexer-less latent path; for phase 2 it is + now only the *analytic* per-document causal set (that phase materialises no + table at all), which is what its indexer KL must span. """ row_end = _row_end(layout, seqlen) doc_start, _, is_valid, _, _ = _derive_csa_doc_boundaries(row_end, seqlen) @@ -137,6 +168,34 @@ def _fp32(tensor): return tensor.cast("float32").numpy() +def _build_latent_warmup_module(loss_coeff=0.01): + """A ``MQALatentAttention`` in the (now illegal) phase-2 combination. + + ``_build_module`` cannot produce this: it dispatches on the production + predicate ``latent_mqa_enabled``, which sends ``"mqa_dsa"`` + + ``dsa_indexer_use_sparse_loss=False`` to ``MHADSAWarmupAttention``. Reaching + the guard therefore means building the latent class by hand -- which is + exactly the situation the guard is for: a spec change that forgets the + predicate must fail loudly instead of quietly running a zero-sparsity + block-sparse kernel. + + Local helper on purpose: ``hybrid_mla_utils`` is shared with other suites. + """ + config = _create_mqa_config( + "mqa_dsa", loss_coeff=loss_coeff, sparse_loss=False + ) + return MQALatentAttention( + config=config, + sublayers_spec=MQALatentAttentionSublayersSpec( + indexer=_indexer_layer_spec() + ), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + k_channels=K_CHANNELS, + ) + + class TestMQAGuards(unittest.TestCase): """Unsupported configurations must fail loudly, not silently mis-compute.""" @@ -210,11 +269,60 @@ def test_softmax_scale_is_the_mha_scale(self): class TestMQAIndexRanges(unittest.TestCase): """The forced window and the indexer candidate range partition the causal - set: no overlap (would double-count) and no gap (would waste budget).""" + set: no overlap (would double-count) and no gap (would waste budget). + + ``_indexer_valid_range`` now lives on the shared + ``HybridMLAIndexerMixin``, because both DSA phases build a candidate range + from it, and its ``window`` argument became a **required positional** one + placed before ``position_offset`` (the warmup phase passes ``0``). It is + reached here through the latent class, which is one of its two callers. + """ def setUp(self): self.module = _build_module(_create_mqa_config("mqa")) + def test_the_range_builder_is_the_shared_mixins(self): + """Retargeted, not deleted: the method moved out of the latent class. + + Both DSA phases must clamp the candidate range identically, so a copy + per class would be a silent divergence risk. Pin the ownership, and pin + that ``window`` is mandatory -- it used to default to + ``self.window_size``, so an updated caller that forgets it would + otherwise keep working while a phase-2 caller silently subtracted a + 128-wide window it does not have. + """ + self.assertIs( + MQALatentAttention._indexer_valid_range, + HybridMLAIndexerMixin._indexer_valid_range, + ) + self.assertIs( + MHADSAWarmupAttention._indexer_valid_range, + HybridMLAIndexerMixin._indexer_valid_range, + ) + seqlen = 32 + args = self._doc_bounds([seqlen], seqlen) + with self.assertRaises(TypeError): + self.module._indexer_valid_range(seqlen, *args) + # ``window=0`` (what the warmup phase passes) is the whole per-document + # causal span, diagonal included; the sparse phase's ``WINDOW`` cuts the + # trailing window off the same range. + valid_range, row_empty = self.module._indexer_valid_range( + seqlen, *args, 0 + ) + vr = valid_range.numpy()[0] + for q in range(seqlen): + self.assertEqual((int(vr[q, 0]), int(vr[q, 1])), (0, q + 1)) + self.assertFalse(bool(row_empty.numpy().any())) + + @staticmethod + def _doc_bounds(layout, seqlen): + """``(doc_start, doc_len, is_valid)`` for one layout.""" + row_end = _row_end(layout, seqlen) + doc_start, doc_len, is_valid, _, _ = _derive_csa_doc_boundaries( + row_end, seqlen + ) + return doc_start, doc_len, is_valid + def test_window_and_indexer_range_partition_causal_set(self): seqlen = 256 for layout in _LAYOUTS: @@ -227,7 +335,7 @@ def test_window_and_indexer_range_partition_causal_set(self): 1, seqlen, WINDOW, doc_start, is_valid ).numpy() valid_range, row_empty = self.module._indexer_valid_range( - seqlen, doc_start, doc_len, is_valid + seqlen, doc_start, doc_len, is_valid, WINDOW ) self._assert_partition( window, @@ -357,8 +465,9 @@ def test_index_dims_are_validated(self): def test_warmup_phase_needs_no_index_topk(self): """Phase 2 must not be forced to carry a top-k budget. - ``_forward_warmup`` never selects a top-k, so a kernel-illegal (or - simply absent, hence default) ``index_topk`` must not block startup -- + ``MHADSAWarmupAttention`` never selects a top-k on either side, so a + kernel-illegal (or simply absent, hence default) ``index_topk`` must not + block startup -- while the sparse phase still rejects it. The production phase-2 ``model_config.json`` relies on this: it ships no ``index_topk`` at all. """ @@ -428,12 +537,26 @@ def test_split_kv_b_proj_only_means_anything_for_latent_mqa(self): # to split, so silently accepting it would hide a mis-set config. with self.assertRaisesRegex(ValueError, "only means"): TransformerConfig(**self._kwargs(mqa_split_kv_b_proj=True)) - for mode in ("mqa_dsa", "mqa_full_causal"): + # The DSA warmup phase is one of those dense paths: it keeps kv_b_proj, + # so accepting the flag there would silently change the parameter set + # at the warmup -> sparse switch instead of at a config edit. + with self.assertRaisesRegex(ValueError, "only means"): + TransformerConfig( + **self._mqa_dsa_kwargs( + mqa_split_kv_b_proj=True, + dsa_indexer_use_sparse_loss=False, + ) + ) + for mode, extra in ( + ("mqa_dsa", {"dsa_indexer_use_sparse_loss": True}), + ("mqa_full_causal", {}), + ): with self.subTest(mode=mode): config = TransformerConfig( **self._mqa_dsa_kwargs( hybrid_mla_attention=mode, mqa_split_kv_b_proj=True, + **extra, ) ) self.assertTrue(config.mqa_split_kv_b_proj) @@ -447,6 +570,7 @@ def test_split_kv_b_proj_rejects_hy_sparse_attention(self): TransformerConfig( **self._mqa_dsa_kwargs( hybrid_mla_attention="mqa_dsa", + dsa_indexer_use_sparse_loss=True, mqa_split_kv_b_proj=True, enable_hy_sparse_attention=True, ) @@ -587,6 +711,129 @@ def test_from_config_accepts_the_current_key_names(self): self.assertTrue(config.indexer_init_from_scratch) +class TestLatentMqaEnabledPredicate(unittest.TestCase): + """``latent_mqa_enabled`` (``hybrid_mla_indexer.py:37-61``) decides alone. + + Both the spec dispatch (``gpt_layer_specs.py``) and + ``MLASelfAttention.mqa_latent`` read this one predicate, so if it drifts the + spec builds one core attention while the enclosing layer feeds it the other + one's activations -- absorbed latents into dense MHA, or per-head K/V into + the block-sparse kernel. Pinned over every combination that reaches it. + """ + + @staticmethod + def _dsv4(**overrides): + return TransformerConfig(**TestHybridMLAConfig._kwargs(**overrides)) + + def test_non_dsv4_models_never_run_latent_mqa(self): + config = TransformerConfig( + num_hidden_layers=2, hidden_size=HIDDEN, num_attention_heads=H + ) + self.assertEqual(config.hybrid_mla_attention, "mha") + self.assertIs(latent_mqa_enabled(config), False) + # The variant gate is checked before the mode, so it holds even for a + # mode that would otherwise say True. ``__post_init__`` rejects that + # pair outright (transformer_config.py:1628-1649), hence the + # post-construction assignment: defence in depth, not a live config. + config.hybrid_mla_attention = "mqa_full_causal" + self.assertIs(latent_mqa_enabled(config), False) + + def test_mha_mode_is_dense_even_on_dsv4(self): + self.assertIs(latent_mqa_enabled(self._dsv4()), False) + + def test_mqa_full_causal_is_latent(self): + config = self._dsv4(hybrid_mla_attention="mqa_full_causal") + self.assertIs(latent_mqa_enabled(config), True) + # No indexer in this mode, so the sparse-loss switch is irrelevant. + config.dsa_indexer_use_sparse_loss = False + self.assertIs(latent_mqa_enabled(config), True) + + def test_mqa_dsa_is_latent_only_in_the_sparse_phase(self): + for sparse_loss in (False, True): + with self.subTest(dsa_indexer_use_sparse_loss=sparse_loss): + config = TransformerConfig( + **TestHybridMLAConfig._mqa_dsa_kwargs( + dsa_indexer_use_sparse_loss=sparse_loss + ) + ) + self.assertIs(latent_mqa_enabled(config), sparse_loss) + # The unit fixture must agree with the production config or + # every phase-2 test below would exercise the wrong backend. + self.assertIs( + latent_mqa_enabled( + _create_mqa_config("mqa_dsa", sparse_loss=sparse_loss) + ), + sparse_loss, + ) + + +class TestLatentMqaRefusesTheWarmupPhase(unittest.TestCase): + """An indexer with the sparse loss off is an error state, not a phase. + + ``MQALatentAttention`` used to *implement* phase 2 (``_forward_warmup``). + Now that phase runs dense MHA in ``MHADSAWarmupAttention``, so this class + seeing that combination means the dispatch predicate was bypassed -- which + must fail loudly rather than quietly build a ``[b, s, s]`` index table and + feed it to the block-sparse kernel at zero sparsity + (``mqa_latent_attention.py:279-288``). + """ + + S = 64 + + def _assert_message(self, message): + for fragment in ( + "dsa_indexer_use_sparse_loss=False", + "DSA warmup phase", + "dense MHA", + "MHADSAWarmupAttention", + ): + self.assertIn(fragment, message) + + def test_phase_raises_and_names_the_dense_backend(self): + module = _build_latent_warmup_module() + self.assertIsNotNone(module.indexer) + self.assertFalse(module.indexer_use_sparse_loss) + with self.assertRaises(ValueError) as raised: + module._phase() + self._assert_message(str(raised.exception)) + + def test_the_forward_refuses_before_touching_a_kernel(self): + """The guard sits in front of the whole sparse path, not inside it. + + ``_phase`` is consulted after the document bounds and before any index + table or kernel launch (``mqa_latent_attention.py:372``), so this needs + no GPU: a well-formed call that would previously have run the warmup + forward now raises instead. + """ + module = _build_latent_warmup_module() + module.eval() + query = paddle.zeros([1, self.S, H, DK], dtype="float32") + key = paddle.zeros([1, self.S, 1, DK], dtype="float32") + w_v = paddle.zeros([DV, H, V_HEAD_DIM], dtype="float32") + with self.assertRaises(ValueError) as raised: + module( + query, + key, + None, + None, + _row_end([self.S], self.S), + v_b_proj_weight=w_v, + x=paddle.zeros([1, self.S, HIDDEN], dtype="float32"), + qr=paddle.zeros([1, self.S, Q_LORA], dtype="float32"), + ) + self._assert_message(str(raised.exception)) + + def test_the_two_surviving_phases_still_resolve(self): + # The guard must not have swallowed the legal states: no indexer is + # full causal, sparse loss on is the sparse phase. + latent = _build_module(_create_mqa_config("mqa")) + self.assertIsNone(latent.indexer) + self.assertEqual(latent._phase(), "full_causal") + sparse = _build_module(_create_mqa_config("mqa_dsa", loss_coeff=0.01)) + self.assertIsNotNone(sparse.indexer) + self.assertEqual(sparse._phase(), "sparse") + + @_GPU class TestMQAEquivalence(unittest.TestCase): """The indexer-less full-causal path is mathematically identical to MHA.""" @@ -842,26 +1089,20 @@ def test_use_sparse_loss_switches_both_attention_and_loss_width(self): """``dsa_indexer_use_sparse_loss`` picks the whole training phase. On these uncompressed ``-2`` layers the switch is one decision with two - effects (``MQALatentAttention._phase``), not just the KL width it selects + effects (``MQALatentAttention._phase``), not just the KL width it picks for the CSA layers of the same model (``_resolve_csa_indexer_loss_topk_effective``): - * ``False`` -- phase 2 (warmup, ``_forward_warmup``). Attention consumes - the **full per-document causal** table, because a freshly initialised - indexer's ranking must not steer attention yet, and the KL is scored - over that same full causal set -- so ``_attn_target``, the top-k KL - target builder, is never called at all. (It used to be called with a - *widened* top-k table; at the production ``index_topk=2048`` that - widening degenerated back into the phase-3 table, which is why the - phase now shares no loss code with phase 3.) * ``True`` -- phase 3 (``_forward_sparse``). Attention consumes ``window + index_topk`` and the KL is restricted to that same set, so ``_attn_target`` is called once per step at exactly ``index_topk``. - - Both column sets are asserted. The ``False`` attention table is *exactly* - assertable: it is ``_build_full_causal_indices``, a pure integer function - of the document bounds with no floating-point scoring in it, so it is - reproducible and equal to the builder's own output element for element. + * ``False`` -- phase 2, which is no longer this class's phase: it runs + dense MHA in ``MHADSAWarmupAttention`` (see + ``TestMQADSAWarmupPhase``). Flipping the switch on a live latent + module therefore *raises* instead of widening the attention table to + the full per-document causal set. That inversion is the point: the + ``[b, s, s]`` table the old ``False`` branch built for a zero-sparsity + kernel does not exist any more. The ``True`` path stays statistical, which is the pre-existing measured fact this test still records: on a single full-length document neither @@ -878,21 +1119,19 @@ def test_use_sparse_loss_switches_both_attention_and_loss_width(self): inner_target = self.module._attn_target def recording_target(query_, kv_, kl_columns, lse_indexer=None): - # The KL's column set is the indexer's candidate set: the top-k in - # the sparse phase, every causal column in warmup. The forced window - # is never in it. + # The KL's column set is the indexer's candidate set, i.e. the + # top-k. The forced window is never in it. loss_widths.append(int(kl_columns.shape[-1])) return inner_target(query_, kv_, kl_columns, lse_indexer) self.module._attn_target = recording_target - def run(sparse): + def run(): _CAPTURED.clear() DSAIndexerLossLoggingHelper.tracker.clear() tensors = [t.clone() for t in (query, key, x, qr)] for tensor in tensors: tensor.stop_gradient = False - self.module.indexer_use_sparse_loss = sparse self.module.train() out = self.module( tensors[0], @@ -910,37 +1149,34 @@ def run(sparse): float(DSAIndexerLossLoggingHelper.tracker["values"][0]), ) - idx_a, loss_sparse = run(True) - idx_b, _ = run(True) - idx_full, loss_full = run(False) + idx_a, loss_sparse = run() + idx_b, _ = run() - # The KL column set: exactly ``index_topk`` under ``True``; under - # ``False`` the same builder is reached but over the whole causal span, - # so the width is ``s``. - self.assertEqual(loss_widths, [INDEX_TOPK, INDEX_TOPK, seqlen]) - # The KL never scores the forced window: its width is exactly the - # indexer's candidate budget, not ``WINDOW + INDEX_TOPK``. + # The KL column set is exactly ``index_topk`` wide... + self.assertEqual(loss_widths, [INDEX_TOPK, INDEX_TOPK]) + # ...and never covers the forced window, i.e. it is the indexer's own + # candidate budget, not ``WINDOW + INDEX_TOPK``. self.assertNotIn(WINDOW + INDEX_TOPK, loss_widths) - # The attention table: window + top-k under ``True``, the full causal - # table (width ``s``) under ``False``. + # The attention table is window + top-k. for table in (idx_a, idx_b): self.assertEqual(int(table.shape[-1]), WINDOW + INDEX_TOPK) - self.assertEqual(list(idx_full.shape), [1, seqlen, seqlen]) - np.testing.assert_array_equal( - idx_full, _full_causal_table([seqlen], seqlen) - ) - # The measured identical-call drift of the ``True`` table, kept as the - # reason its width -- not its contents -- is what gets asserted. + # The measured identical-call drift of the table, kept as the reason its + # width -- not its contents -- is what gets asserted. drift = float((idx_a != idx_b).mean()) self.assertLess(drift, 0.05) - self.assertGreater(loss_sparse, 0.0) - self.assertGreater(loss_full, 0.0) - # A wider renormalisation set means a different KL; equal values would - # mean the switch never reached the loss. - self.assertGreater(abs(loss_full - loss_sparse), 1e-6) + + # The other half of the switch: no wider table, a refusal. ``_phase`` is + # read live, so the flip takes effect on this same module. + self.module.indexer_use_sparse_loss = False + _CAPTURED.clear() + with self.assertRaisesRegex(ValueError, "MHADSAWarmupAttention"): + run() + self.assertEqual(_CAPTURED, [], "an index table was built anyway") + self.assertEqual(loss_widths, [INDEX_TOPK, INDEX_TOPK]) + self.module.indexer_use_sparse_loss = True def test_recompute_double_forward_is_consistent(self): """Reentrant recompute runs the layer twice: pass 1 under ``no_grad``. @@ -1061,15 +1297,24 @@ class TestMQADSAWarmupPhase(unittest.TestCase): """Phase 2 of ``"mqa_dsa"``: ``dsa_indexer_use_sparse_loss=False``. The indexer is still being learned, so attention must not consume its - ranking: it attends to the full per-document causal set (bit-identical to - ``hybrid_mla_attention="mqa_full_causal"``) while the indexer's top-k feeds - the wide KL loss only. ``TestMQADSA`` covers the phase-3 shape - (``True``), where attention consumes ``window + index_topk``. - - Kept as its own class rather than folded into ``TestMQADSA``: the module - fixture differs (the switch is off from construction, not flipped mid-test), - and everything here is an exact assertion, because the full-causal table is - integer-only. + ranking -- and with no top-k on either side there is nothing for absorbed + latent MQA to save. So this phase runs **phase 1's dense MHA** with the + indexer bolted on (``mha_dsa_warmup_attention.MHADSAWarmupAttention``), + which ``latent_mqa_enabled`` selects (``TestLatentMqaEnabledPredicate``) + and which ``MQALatentAttention`` now refuses to impersonate + (``TestLatentMqaRefusesTheWarmupPhase``). ``TestMQADSA`` covers phase 3, + where attention does consume ``window + index_topk``. + + Kept in this file although the backend moved: what these tests are about is + the phase boundary, and the phase-3 fixtures they contrast with live here. + Two class-wide inversions of the old assertions, both consequences of the + dense backend: + + * ``_CAPTURED`` (the block-sparse kernel's ``token_indices``) must stay + **empty** -- no ``[b, s, s]`` index table is built at all, where the old + warmup built one and then walked all ``s`` columns at zero sparsity; + * the reference is ``_build_phase1_dense_module``, a plain + ``DotProductAttention``, and agreement with it is **bit** equality. """ SEQLEN = 256 @@ -1086,97 +1331,121 @@ def setUpClass(cls): def setUp(self): _CAPTURED.clear() + _WARMUP_TARGETS.clear() DSAIndexerLossLoggingHelper.tracker.clear() - self.module = self._build_warmup() + self.config = _create_mqa_config( + "mqa_dsa", loss_coeff=0.01, sparse_loss=False + ) + self.module = self._build_warmup(self.config) self.row_end = _row_end(self.LAYOUT, self.SEQLEN) @staticmethod - def _build_warmup(): - """A ``"mqa_dsa"`` module with the switch off from construction.""" - config = _create_mqa_config("mqa_dsa", loss_coeff=0.01) - config.dsa_indexer_use_sparse_loss = False + def _build_warmup(config): + """The dense phase-2 backend, picked by the production predicate.""" module = _build_module(config, bf16=True) + assert isinstance(module, MHADSAWarmupAttention), type(module) + assert not isinstance(module, MQALatentAttention) assert module.indexer is not None - assert module.indexer_use_sparse_loss is False return module def _inputs(self, seed=0): - return _make_inputs(self.SEQLEN, seed=seed, with_hidden=True) + return _make_dense_inputs(self.SEQLEN, seed=seed) + + def _ids(self): + """All-valid ``input_ids``; ``pad_token_id`` defaults to 0.""" + return paddle.ones([1, self.SEQLEN], dtype="int64") - def _call(self, module, tensors, w_v, training, differentiable=False): + def _call( + self, module, tensors, training, differentiable=False, input_ids=None + ): + """The phase-2 call shape: per-head q/k/v plus the indexer's inputs. + + ``attn_mask_type`` is explicit because ``DotProductAttention`` leaves + ``is_causal`` False when it is omitted, and the document mask alone + would then leave the upper triangle unmasked. ``input_ids`` is the only + kwarg the phase-1 reference does not take, so it is passed only when + given -- which is how ``MLASelfAttention`` forwards it too, gated on + ``accepts_input_ids``. + """ module.train() if training else module.eval() - query, key, x, qr = tensors + query, key, value, x, qr = tensors if differentiable: for tensor in tensors: tensor.stop_gradient = False + extra = {} if input_ids is None else {"input_ids": input_ids} return module( query, key, - None, + value, None, self.row_end, - v_b_proj_weight=w_v, + attn_mask_type=AttnMaskType.causal, x=x, qr=qr, + **extra, ) - def test_attention_output_equals_the_indexer_less_full_causal_path(self): - """The core invariant: the indexer's *existence* must not move a bit. + def test_attention_output_is_bit_identical_to_phase_1(self): + """The core invariant, inverted: phase 2 *is* phase 1's attention. + + The old assertion compared against the indexer-less **latent** path. + The reference is now the dense ``DotProductAttention`` of phase 1 + itself, and the claim is stronger: the whole attention half is + ``super().forward`` (``mha_dsa_warmup_attention.py:179-198``), so the + indexer loss is the only new thing. Nothing needs weight copying -- + attention consumes no module parameter here (q/k/v are inputs and + ``softmax_offset`` is ``None`` in both) -- so a difference could only + come from the backend. Measured maxabs 0.0 on SM103 / FA4. - Both paths call the same ``_build_full_causal_indices`` and then the - same sparse kernel, so the outputs must be bit-identical, not merely - close. Nothing needs weight copying: on this path attention consumes no - module parameter at all -- the query/key/``v_b_proj_weight`` are inputs - and ``softmax_offset`` is ``None`` in both -- so the only thing that - could differ is the index table. Asserted in both modes, so the - ``:495`` early exit and the full ``:600`` branch are each covered. + Asserted in eval and in train mode, i.e. with the indexer branch both + skipped (``mha_dsa_warmup_attention.py:199-200``) and taken. """ - query, key, w_v, x, qr = self._inputs() - reference = _build_module(_create_mqa_config("mqa"), bf16=True) - self.assertIsNone(reference.indexer) + tensors = self._inputs() + reference = _build_phase1_dense_module(self.config, bf16=True) self.assertIsNone(reference.softmax_offset) self.assertIsNone(self.module.softmax_offset) self.assertEqual(self.module.softmax_scale, reference.softmax_scale) - reference.eval() - out_ref = _fp32( - self._call(reference, (query, key, x, qr), w_v, training=False) - ) + out_ref = _fp32(self._call(reference, tensors, training=False)) for training in (False, True): with self.subTest(training=training): DSAIndexerLossLoggingHelper.tracker.clear() - tensors = [t.clone() for t in (query, key, x, qr)] + clones = [t.clone() for t in tensors] out = self._call( self.module, - tensors, - w_v, + clones, training=training, differentiable=training, + input_ids=self._ids(), ) np.testing.assert_array_equal(_fp32(out), out_ref) + self.assertEqual(_CAPTURED, [], "the block-sparse kernel was reached") - def test_token_indices_are_the_full_causal_table(self): - """The captured table is ``[b, s, s]`` and element-wise equal to the - builder's own output, over several document layouts.""" - query, key, w_v, x, qr = self._inputs() + def test_output_matches_the_fp32_dense_mha_reference(self): + """Independent check that the delegated half really is per-document + causal MHA, rather than merely equal to another copy of itself. + + The document layouts are varied here rather than in the bit-identity + test because masking is what this one is about -- and no layout may + produce an index table. + """ + tensors = self._inputs(seed=3) + query, key, value = tensors[0], tensors[1], tensors[2] for layout in ([self.SEQLEN], self.LAYOUT, [3, WINDOW, WINDOW + 1, 1]): with self.subTest(layout=layout): self.row_end = _row_end(layout, self.SEQLEN) _CAPTURED.clear() - self._call( - self.module, (query, key, x, qr), w_v, training=False - ) - table = _CAPTURED[-1] - self.assertEqual( - list(table.shape), [1, self.SEQLEN, self.SEQLEN] - ) - np.testing.assert_array_equal( - table, _full_causal_table(layout, self.SEQLEN) - ) - # ... and the table is sound in its own right: the whole causal - # set, no duplicate, nothing cross-document. - _check_index_invariants( - self, table, self.row_end, self.SEQLEN, expect_full=True + out = self._call(self.module, tensors, training=False) + ref = _dense_mha_reference( + query, + key, + value, + self.row_end, + self.module.softmax_scale, ) + # bf16 flashmask vs the fp32 reference: measured rel 1.969e-3 + # at [100, 156]; 3.5e-3 keeps the phase-3 tests' margin. + self.assertLess(_rel(out, ref), 3.5e-3) + self.assertEqual(_CAPTURED, [], "an index table was built") def test_warmup_undoes_the_indexer_weight_prebake_for_tilelang(self): """The tilelang indexer re-applies ``head_dim**-0.5``, so the pre-bake @@ -1217,13 +1486,16 @@ def recording_tl(*args, **kwargs): seen["probs"] = probs.cast("float32").numpy().copy() return columns, probs - query, key, w_v, x, qr = self._inputs() - tensors = [t.clone() for t in (query, key, x, qr)] + tensors = [t.clone() for t in self._inputs()] self.module._indexer_projections = recording_proj tl_mod.csa_indexer_topk_fwd = recording_tl try: self._call( - self.module, tensors, w_v, training=True, differentiable=True + self.module, + tensors, + training=True, + differentiable=True, + input_ids=self._ids(), ) finally: self.module._indexer_projections = inner_proj @@ -1285,13 +1557,19 @@ def recording_tl(*args, **kwargs): def test_warmup_scores_every_causal_column_via_tilelang(self): """Phase 2 scores the whole causal span, in one tilelang call. - Two things are pinned. First, the **cuDNN** top-k kernel -- phase 3's + Three things are pinned. First, the **cuDNN** top-k kernel -- phase 3's selector -- is called zero times: this phase reads no ``index_topk``, no window and no clamped candidate range. Second, the tilelang indexer is called exactly once at ``topk_effective == s``, its documented - "full-candidate selection" mode, and the columns it comes back with are - exactly the attention table's, diagonal included -- the very column the - old clamped candidate range could never return. + "full-candidate selection" mode. Third, the columns it comes back with + are exactly the per-document causal set, diagonal included -- the very + column the old clamped candidate range could never return. + + The causal set is now only *analytic* (``_full_causal_table``): with the + dense backend there is no attention table to compare against, which is + itself asserted -- ``_CAPTURED`` stays empty while the KL still spans + every causal column, i.e. the width came without the ``[b, s, s]`` + transient. Before the phase split this test demanded one *cuDNN* call for a widened loss table. That widening was the bug: at the production @@ -1317,13 +1595,16 @@ def recording_tl(*args, **kwargs): tl_columns.append(columns.numpy().copy()) return columns, probs - query, key, w_v, x, qr = self._inputs() - tensors = [t.clone() for t in (query, key, x, qr)] + tensors = [t.clone() for t in self._inputs()] fwd_mod.cudnn_indexer_topk_fwd = recording_cudnn tl_mod.csa_indexer_topk_fwd = recording_tl try: out = self._call( - self.module, tensors, w_v, training=True, differentiable=True + self.module, + tensors, + training=True, + differentiable=True, + input_ids=self._ids(), ) out.cast("float32").sum().backward() finally: @@ -1335,32 +1616,36 @@ def recording_tl(*args, **kwargs): ) self.assertEqual(tl_widths, [self.SEQLEN]) self.assertIn("values", DSAIndexerLossLoggingHelper.tracker) + self.assertEqual(_CAPTURED, [], "the block-sparse kernel was reached") - attn_table = _CAPTURED[-1] - np.testing.assert_array_equal( - attn_table, _full_causal_table(self.LAYOUT, self.SEQLEN) - ) + causal = _full_causal_table(self.LAYOUT, self.SEQLEN) kl_columns = tl_columns[0] for row in range(self.SEQLEN): - attn_cols = attn_table[0, row] + causal_cols = causal[0, row] kl_cols = kl_columns[0, row] self.assertEqual( set(kl_cols[kl_cols >= 0].tolist()), - set(attn_cols[attn_cols >= 0].tolist()), - f"row {row}: KL and attention column sets differ", + set(causal_cols[causal_cols >= 0].tolist()), + f"row {row}: KL and causal column sets differ", ) last = self.SEQLEN - 1 - self.assertIn(last, set(attn_table[0, last].tolist())) self.assertIn(last, set(kl_columns[0, last].tolist())) + # The KL target is built over that same width, once. + self.assertEqual( + [t.shape for t in _WARMUP_TARGETS], + [(1, self.SEQLEN, self.SEQLEN)], + ) def test_eval_early_exit_matches_the_training_forward(self): - """``:495`` skips the indexer projections entirely under ``eval()``. + """The no-loss forward skips the indexer entirely. - Attention does not consume the indexer in this phase, so with nothing to - learn this step there is nothing to compute -- and the attention output - must be bit-identical to the training forward, which does run them. + ``_needs_indexer_loss`` gates the whole second half + (``mha_dsa_warmup_attention.py:199-200``), so with nothing to learn this + step there is nothing to compute -- and because the attention half is + the same ``super().forward`` either way, the output must be + bit-identical to the training forward, which does run the indexer. """ - query, key, w_v, x, qr = self._inputs(seed=2) + tensors = self._inputs(seed=2) calls = [] inner = self.module.indexer.forward_before_topk @@ -1370,48 +1655,65 @@ def recording(*args, **kwargs): self.module.indexer.forward_before_topk = recording - tensors = [t.clone() for t in (query, key, x, qr)] out_train = _fp32( self._call( - self.module, tensors, w_v, training=True, differentiable=True + self.module, + [t.clone() for t in tensors], + training=True, + differentiable=True, + input_ids=self._ids(), ) ) self.assertEqual(len(calls), 1) self.assertIn("values", DSAIndexerLossLoggingHelper.tracker) + self.assertEqual(len(_WARMUP_TARGETS), 1) DSAIndexerLossLoggingHelper.tracker.clear() out_eval = _fp32( - self._call(self.module, (query, key, x, qr), w_v, training=False) + self._call( + self.module, tensors, training=False, input_ids=self._ids() + ) ) self.assertEqual(len(calls), 1, "eval must not run the indexer at all") + self.assertEqual(len(_WARMUP_TARGETS), 1, "eval built a KL target") self.assertNotIn("values", DSAIndexerLossLoggingHelper.tracker) np.testing.assert_array_equal(out_train, out_eval) def test_indexer_gradients_flow_in_the_warmup_phase(self): - """All five indexer parameters keep a finite non-zero gradient. + """Every indexer parameter keeps a finite non-zero gradient. Phase 2 is where the indexer does all of its learning (the backbone is frozen by the trainer), so a silently gradient-free indexer parameter would waste the entire phase. Same contract as ``TestMQADSA.test_backward_produces_finite_grads_and_reports_loss``, - extended to ``k_norm`` and driven with attention detached from the - indexer's output. + widened to the whole parameter set and driven through the dense backend, + whose attention half does not touch the indexer at all -- the gradients + can only come from the KL attached to the output + (``mha_dsa_warmup_attention.py:342-354``). Measured range on this + fixture: 9.3e-8 .. 2.2e-7 over the 8 parameters. """ - query, key, w_v, x, qr = self._inputs() - tensors = [query, key, x, qr] + tensors = self._inputs() + query, x, qr = tensors[0], tensors[3], tensors[4] out = self._call( - self.module, tensors, w_v, training=True, differentiable=True + self.module, + tensors, + training=True, + differentiable=True, + input_ids=self._ids(), ) out.cast("float32").sum().backward() indexer = self.module.indexer - params = { - "wq_b.weight": indexer.wq_b.linear.weight, - "wk.weight": indexer.wk.linear.weight, - "k_norm.weight": indexer.k_norm.weight, - "k_norm.bias": indexer.k_norm.bias, - "weights_proj.weight": indexer.weights_proj.linear.weight, - } - for name, param in params.items(): + # Every parameter, not a hand-listed subset: the measured set is 8 + # (weight+bias of wq_b / wk / k_norm / weights_proj), and a new one + # appearing gradient-free would otherwise go unnoticed. + named = dict(indexer.named_parameters()) + self.assertGreaterEqual(len(named), 5) + for expected in ("wq_b", "wk", "k_norm", "weights_proj"): + self.assertTrue( + any(expected in name for name in named), + f"indexer has no {expected} parameter any more", + ) + for name, param in named.items(): self.assertIsNotNone(param.grad, f"indexer.{name} has no gradient") self.assertTrue( bool(paddle.isfinite(param.grad.cast("float32")).all()), @@ -1430,37 +1732,50 @@ def test_indexer_gradients_flow_in_the_warmup_phase(self): self.assertTrue(bool(paddle.isfinite(query.grad.cast("float32")).all())) self.assertIn("values", DSAIndexerLossLoggingHelper.tracker) - def test_recompute_double_forward_table_is_bit_identical(self): - """Stronger than the phase-3 equivalent, and on the harder layout. + def test_recompute_double_forward_attaches_the_loss_once(self): + """Reentrant recompute runs the layer twice: pass 1 under ``no_grad``. - ``TestMQADSA.test_recompute_double_forward_is_consistent`` has to pick a - two-document layout because the top-k kernel's emitted order drifts on a - single full-length document. Phase 2's table contains no floating-point - scoring at all, so it is bit-identical across the two passes *and* equal - to the analytic table -- assert both, on the single-document layout that - the phase-3 path cannot use. + The loss must be attached on the grad-enabled pass only -- otherwise it + would be counted twice -- and, since the attention half is dense + flashmask on a fixed mask, the two passes must produce the same output + bit for bit. The old form of this test asserted the same thing about the + index table; there is none any more, so the KL target takes its place: + it is built once, on the differentiable pass. + + Single document on purpose: the phase-3 equivalent + (``TestMQADSA.test_recompute_double_forward_is_consistent``) has to + avoid that layout because the top-k kernel's emitted order drifts on it. + Phase 2 has no top-k, so the hard layout is available. """ - seqlen = self.SEQLEN - self.row_end = _row_end([seqlen], seqlen) - query, key, w_v, x, qr = self._inputs() - query.stop_gradient = False - expected = _full_causal_table([seqlen], seqlen) + self.row_end = _row_end([self.SEQLEN], self.SEQLEN) + tensors = self._inputs() + tensors[0].stop_gradient = False - _CAPTURED.clear() with paddle.no_grad(): - self._call(self.module, (query, key, x, qr), w_v, training=True) - first = _CAPTURED[-1] + first = _fp32( + self._call( + self.module, + tensors, + training=True, + input_ids=self._ids(), + ) + ) + self.assertEqual(_WARMUP_TARGETS, []) self.assertNotIn( "values", DSAIndexerLossLoggingHelper.tracker, "indexer loss must not be attached on the no_grad pass", ) - self._call(self.module, (query, key, x, qr), w_v, training=True) - second = _CAPTURED[-1] + second = _fp32( + self._call( + self.module, tensors, training=True, input_ids=self._ids() + ) + ) np.testing.assert_array_equal(first, second) - np.testing.assert_array_equal(first, expected) + self.assertEqual(len(_WARMUP_TARGETS), 1) self.assertIn("values", DSAIndexerLossLoggingHelper.tracker) + self.assertEqual(_CAPTURED, [], "the block-sparse kernel was reached") class TestHashableTensor(unittest.TestCase): @@ -1680,7 +1995,10 @@ def test_lse_present_selects_the_kernel(self): def test_lse_absent_selects_the_reference(self): got = MQALatentAttention._attn_target(self.stub, "q", "kv", "idx", None) self.assertEqual((got, self.calls), ("python", ["python"])) - # The default is the fallback too: phase 2 calls it with three args. + # The default is the fallback too. No production caller omits the LSE + # any more (phase 2 used to), but the parameter's default is what makes + # the reference reachable from a bare three-argument call, which is how + # every reference-vs-kernel comparison in this file invokes it. self.assertEqual( MQALatentAttention._attn_target(self.stub, "q", "kv", "idx"), "python", diff --git a/tests/single_card_tests/transformer/test_muon_hybrid_mla_grouping.py b/tests/single_card_tests/transformer/test_muon_hybrid_mla_grouping.py index e82711ab1b..cc9c3ca7c0 100644 --- a/tests/single_card_tests/transformer/test_muon_hybrid_mla_grouping.py +++ b/tests/single_card_tests/transformer/test_muon_hybrid_mla_grouping.py @@ -238,7 +238,15 @@ def slice_fn(self, rel): class TestMqaDsaRouting(unittest.TestCase): - """The phase-2 config: latent MQA + DSA indexer + sink.""" + """The phase-2 config: dense warmup MHA + DSA indexer + sink. + + (WAS "latent MQA + DSA indexer + sink". Phase 2 has no top-k on either side + and now runs phase 1's dense attention -- ``MHADSAWarmupAttention`` -- + instead of the block-sparse latent MQA. Muon routing keys off parameter + names only, and the two phases own byte-identical parameter sets, so nothing + below changes; see ``test_mqa_dsa_adds_exactly_the_indexer`` and + ``test_sparse_loss_phase_matches_warmup_phase``.) + """ @classmethod def setUpClass(cls): @@ -413,8 +421,21 @@ def test_mqa_dsa_adds_exactly_the_indexer(self): def test_sparse_loss_phase_matches_warmup_phase(self): """dsa_indexer_use_sparse_loss changes the attention/loss shape, never the optimizer routing. + + Since the flag also selects the core-attention *class* (phase 2 dense + ``MHADSAWarmupAttention`` vs phase 3 ``MQALatentAttention``, + ``hybrid_mla_indexer.latent_mqa_enabled``), assert the parameter names + themselves are equal first: that is the property the HF-checkpoint stage + switch rests on, and Muon routing is a function of those names. """ warm, sparse = self.envs[_DSA_CFG], self.envs[_DSA_SPARSE_LOSS_CFG] + self.assertEqual(set(sparse.params), set(warm.params)) + self.assertEqual( + type(warm.attn.core_attention).__name__, "MHADSAWarmupAttention" + ) + self.assertEqual( + type(sparse.attn.core_attention).__name__, "MQALatentAttention" + ) self.assertEqual(sparse.muon, warm.muon) self.assertEqual(sparse.adamw, warm.adamw) self.assertEqual( diff --git a/tests/single_card_tests/transformer/test_train_indexer_only.py b/tests/single_card_tests/transformer/test_train_indexer_only.py index 30f779a6d6..be5e039b08 100644 --- a/tests/single_card_tests/transformer/test_train_indexer_only.py +++ b/tests/single_card_tests/transformer/test_train_indexer_only.py @@ -28,6 +28,13 @@ never runs backward - no error, normal loss curve, zero indexer updates. ``keep_indexer_grad_path`` re-enters the graph to keep that path alive. +Neither depends on which attention the ``-2`` hybrid-MLA layers run: phase 2 now +dispatches them to the dense ``MHADSAWarmupAttention`` rather than latent MQA, +but that class attaches its KL through the very same +``TileLangCSAIndexerLossAutoScaler`` and sits inside the very same recompute +segments, so both failure modes are reached identically. Accordingly every guard +below is expressed on ``config`` / the PyLayers, never on the attention class. + The tests below run on CPU and do not need the TileLang/cuDNN kernels. """ @@ -68,6 +75,13 @@ def _make_dsv4_config(**overrides): "dsa_index_head_dim": 64, "dsa_index_topk": 16, "dsa_indexer_loss_coeff": 0.01, + # Phase 2. For a CSA layer this widens the KL and makes the main + # attention walk the whole compressed range; for a ``-2`` hybrid-MLA + # layer it additionally selects the attention *backend* -- the dense + # ``MHADSAWarmupAttention`` rather than latent MQA + # (hybrid_mla_indexer.py::latent_mqa_enabled). Either way it is the + # legal companion of ``train_indexer_only=True``; ``True`` here is the + # phase-3 value and is rejected with it. "dsa_indexer_use_sparse_loss": False, "csa_dense_mode": False, } @@ -76,13 +90,22 @@ def _make_dsv4_config(**overrides): def _make_mqa_config(**overrides): - """dsv4-hybrid config whose ``-2`` layers run latent MQA + DSA indexer. + """dsv4-hybrid config whose ``-2`` layers carry a DSA indexer. A ``-2`` layer is a hybrid MLA layer, so ``__post_init__`` demands the whole ``hybrid_mla_*`` block before it ever gets to the ``train_indexer_only`` checks. ``csa_dense_mode=True`` matches the real new-attention configs: the 128/HCA layers have no CSAIndexer, so the DSAIndexer of the ``-2`` layers is the only Indexer in the model. + + Which attention those ``-2`` layers run is a separate question that + ``train_indexer_only`` deliberately does not ask: with the inherited + ``dsa_indexer_use_sparse_loss=False`` it is the dense + ``MHADSAWarmupAttention`` (phase 2), and with ``True`` the absorbed latent + ``MQALatentAttention`` (phase 3). Only the *existence* of the indexer + matters here, which is exactly what ``has_mqa_indexer`` tests + (transformer_config.py:1892-1897 -- ``"mqa_dsa"`` plus a ``-2`` ratio, with + no reference to the sparse-loss flag). """ kwargs = { "csa_dense_mode": True, @@ -146,14 +169,17 @@ def test_no_indexer_layer_rejected(self): def test_mqa_dsa_indexer_satisfies_the_check_without_any_csa_layer(self): # The new-attention phase 2: every CSA-ratio layer is HCA (128) and # csa_dense_mode is on, so there is no CSAIndexer anywhere. The only - # Indexer is the DSAIndexer of the ``-2`` latent MQA layers. + # Indexer is the DSAIndexer of the ``-2`` hybrid-MLA layers, which in + # this phase run dense MHA (``MHADSAWarmupAttention``) with that indexer + # bolted on. config = _make_mqa_config(train_indexer_only=True) self.assertTrue(config.train_indexer_only) self.assertTrue(config.csa_dense_mode) def test_mqa_full_causal_rejected(self): - # "mqa_full_causal" drops the DSAIndexer (gpt_layer_specs.py passes - # indexer=None), which is exactly the phase-1 shape: nothing to train. + # "mqa_full_causal" drops the DSAIndexer (gpt_layer_specs.py:379-384 + # passes indexer=None), which is exactly the phase-1 shape: nothing to + # train. with self.assertRaisesRegex(ValueError, "build at least one Indexer"): _make_mqa_config( train_indexer_only=True, @@ -162,8 +188,9 @@ def test_mqa_full_causal_rejected(self): def test_mqa_without_hybrid_layer_rejected(self): # ``hybrid_mla_attention`` only does anything to ``-2`` layers; without - # one there is no MQALatentAttention and therefore no DSAIndexer. That - # is now caught earlier and more precisely, by the unconditional + # one nothing dispatches to MHADSAWarmupAttention / MQALatentAttention + # and therefore no DSAIndexer is built (gpt_layer_specs.py:368-398). + # That is now caught earlier and more precisely, by the unconditional # ``hybrid_mla_attention`` guard, rather than by the # ``train_indexer_only`` "build at least one Indexer" check. with self.assertRaisesRegex(ValueError, "only applies to MLA layers"): @@ -230,6 +257,43 @@ def test_non_tensor_passthrough(self): config = _make_dsv4_config(train_indexer_only=True) self.assertIsNone(keep_indexer_grad_path(None, config)) + def test_gating_is_train_indexer_only_and_nothing_else(self): + """``train_indexer_only`` alone decides; the phase must not leak in. + + ``recompute_utils.py:49`` is the only config read in the whole function, + so the anchor must appear for every config that sets the flag, whatever + Indexer kind or attention backend the rest of the config selects. Pinned + explicitly because phase 2 moved the ``-2`` layers from latent MQA to + ``MHADSAWarmupAttention``, changing which class sits inside the + recompute segment; a future "only for backend X" shortcut here would + reintroduce the silent zero-gradient bug for the other one. + + ``dsa_indexer_use_sparse_loss=True`` is only illegal *together with* + ``train_indexer_only`` on a ``-2`` layer (transformer_config.py:1656, + :1666), so the CSA-shaped ``_make_dsv4_config`` is what makes the + sparse-flag half of this assertion constructible at all. + """ + cases = ( + ("csa_indexer", lambda **kw: _make_dsv4_config(**kw)), + ( + "csa_indexer_sparse_loss", + lambda **kw: _make_dsv4_config( + dsa_indexer_use_sparse_loss=True, **kw + ), + ), + ("mqa_dsa_warmup_indexer", lambda **kw: _make_mqa_config(**kw)), + ) + for label, make in cases: + with self.subTest(config=label): + on = keep_indexer_grad_path( + self.hidden, make(train_indexer_only=True) + ) + self.assertIsNot(on, self.hidden) + self.assertFalse(on.stop_gradient) + self.assertIs( + keep_indexer_grad_path(self.hidden, make()), self.hidden + ) + class _IndexerBranch(nn.Layer): """Minimal stand-in for one CSA layer: frozen main path + trainable indexer. @@ -299,6 +363,10 @@ def test_dsa_autoscaler_still_passes_grad_when_backbone_trains(self): def test_tilelang_autoscaler_forward_returns_fresh_tensor_when_frozen(self): # Only the forward contract is checked here: the backward needs the # TileLang/cuDNN indexer kernels, which single-card CPU tests skip. + # This PyLayer is shared by both indexer-bearing backends -- phase 3's + # ``MQALatentAttention`` and phase 2's ``MHADSAWarmupAttention``, which + # imports the same symbol from ``csa_attention`` -- so one forward-shape + # guard covers the frozen-backbone contract for the whole family. output = paddle.randn([2, 4, 8]) output.stop_gradient = True index_q = paddle.randn([2, 4, 8]) @@ -311,14 +379,18 @@ def test_tilelang_autoscaler_forward_returns_fresh_tensor_when_frozen(self): topk_probs = paddle.rand([2, 4, 2]) target = paddle.rand([2, 4, 2]) + # Positional order is the real one (csa_attention.py:1205-1218): + # output, target, index_q, weights, index_k_comp, topk_indices, + # topk_probs, loss_coeff. It matters even for a forward-only check, + # because ``num_rows`` is derived from ``target.shape``. attached = TileLangCSAIndexerLossAutoScaler.apply( output, + target, index_q, weights, index_k, topk_indices, topk_probs, - target, 0.01, ) self.assertIsNot(attached, output) diff --git a/tests/single_card_tests/transformer/test_vha_dsv4.py b/tests/single_card_tests/transformer/test_vha_dsv4.py index 6f412f1919..733db19c17 100644 --- a/tests/single_card_tests/transformer/test_vha_dsv4.py +++ b/tests/single_card_tests/transformer/test_vha_dsv4.py @@ -120,6 +120,21 @@ def _make_config( def _build(config, layer_number=0): + """Build the CSA-family attention for one layer. + + Always ``"dsv4_hybrid_attention"``, i.e. the ``DSv4HybridSelfAttention`` / + ``CompressedSparseAttention`` branch of ``get_attention_spec``. The + hybrid-MLA (``csa_compress_ratios == -2``) branch, where + ``dsa_indexer_use_sparse_loss`` selects dense ``MHADSAWarmupAttention`` + (phase 2) or ``MQALatentAttention`` (phase 3), is a *different* logical + attention type and is never reached from here -- no config in this file puts + a ``-2`` in ``csa_compress_ratios`` or sets ``hybrid_mla_attention``, so the + ``dsa_indexer_use_sparse_loss=False`` above only reaches the CSA layers' + own ``_resolve_topk_effective``. In production VHA does co-exist with the + ``-2`` layers, but ``MultiLatentAttention`` applies the postmix to + ``core_attn_out`` *after* the core attention returns, so it is agnostic to + which of the two the phase selected. + """ model_parallel_cuda_manual_seed(_SEED) spec = get_attention_spec( config=config, From a9088fee4aab1ca6c23174859a3b7a36e62d1ad4 Mon Sep 17 00:00:00 2001 From: shenliang03 Date: Wed, 12 Aug 2026 16:58:12 +0800 Subject: [PATCH 2/3] [release/0.4][Improvements] Fall back to pad_token_id=0 in the indexer loss mask Review of #1721 flagged the `assert pad_token_id is not None` copied into `HybridMLAIndexerMixin._indexer_loss_mask` from `csa_attention.py:2387-2390` (upstream #1291): `TransformerConfig.from_config` can copy a `None` straight out of an external/HF config, so a run the embedding and the MoE router both accept would abort at its first indexer loss -- and under `python -O`, where asserts are stripped, the ids would have been compared against `None` and no row masked at all. Fold `None` to `0`, the convention every other consumer of the field already follows (`gpt_embedding.py:214-216`, `mtp_embedding_layer.py:105-107`, `moe_router.py:605-607` and four more sites), and pin it with a regression test that asserts the `None` config produces the same mask, the same denominator and the same logged loss as `pad_token_id=0`. The line references in the two warmup tests move with the edit. --- .../transformer/hybrid_mla_indexer.py | 13 ++++-- .../transformer/test_mqa_dsa_warmup_cp.py | 6 +-- .../test_hybrid_mla_warmup_doc_mask_loss.py | 43 ++++++++++++++++++- 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/src/paddlefleet/transformer/hybrid_mla_indexer.py b/src/paddlefleet/transformer/hybrid_mla_indexer.py index d9257b6ce4..c51867a3a0 100644 --- a/src/paddlefleet/transformer/hybrid_mla_indexer.py +++ b/src/paddlefleet/transformer/hybrid_mla_indexer.py @@ -200,13 +200,20 @@ def _indexer_loss_mask(self, input_ids: Tensor | None, b: int, s: int): **global** valid-row count, so summing the per-rank losses reproduces the single-rank reduction. ``input_ids`` arrives sharded unless ``experimental_dataflow``, exactly as at ``csa_attention.py:2419-2428``. + + An unset ``pad_token_id`` falls back to ``0``, the convention every other + consumer of it in this repository already follows + (``gpt_embedding.py:214-216``, ``mtp_embedding_layer.py:105-107``, + ``moe_router.py:605-607`` and four more sites): ``from_config`` can copy a + ``None`` straight out of an external/HF config, and treating that as a + fatal configuration error here would abort the first indexer loss of a + run that the embedding and the router accepted. """ if input_ids is None: return None, None pad_token_id = getattr(self.config, "pad_token_id", 0) - assert pad_token_id is not None, ( - "pad_token_id must be set in config when input_ids is provided" - ) + if pad_token_id is None: + pad_token_id = 0 if self.cp_enabled: if not getattr(self.config, "experimental_dataflow", False): input_ids = ContextParallelGatherOp.apply( diff --git a/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py b/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py index 752bf3c29a..3ec9487577 100644 --- a/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py +++ b/tests/multi_card_tests/transformer/test_mqa_dsa_warmup_cp.py @@ -42,7 +42,7 @@ sequence and row-sliced (``hybrid_mla_indexer._indexer_valid_range``). 3. The full-candidate KL still has to normalise across the CP group: the masked branch divides by the **global** valid-row count - (``hybrid_mla_indexer.py:215-220``) and the unmasked one folds ``/cp_size`` + (``hybrid_mla_indexer.py:222-227``) and the unmasked one folds ``/cp_size`` into the coefficient handed to the *backward* (``mha_dsa_warmup_attention.py:322-326``). 4. A layout with genuine row-validity pad rows. ``_STRADDLE`` sums to exactly @@ -222,7 +222,7 @@ def _warmup_cfg(cp_size, loss_coeff=0.0): cfg.cp_balance_mode = "contiguous_allgather" # Production EB dataflow hands every rank the *global* ``input_ids``, which # is the branch ``_indexer_loss_mask`` takes when this flag is set - # (``hybrid_mla_indexer.py:210-214``). + # (``hybrid_mla_indexer.py:217-221``). cfg.experimental_dataflow = True cfg.pad_token_id = 0 cfg.context_parallel_size = cp_size @@ -730,7 +730,7 @@ def test_4_indexer_loss_cp_normalisation(self): Read straight out of ``DSAIndexerLossLoggingHelper``, so it observes the denominator rather than its shadow in the gradients. Masked divides by - the **global** valid-row count (``hybrid_mla_indexer.py:215-220``); + the **global** valid-row count (``hybrid_mla_indexer.py:222-227``); unmasked takes the plain local mean with ``/cp_size`` folded into the coefficient, which ``_assert_loss_coeff`` checks reaches the backward. """ diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py index f3d1a981d0..8cd6382ce6 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_warmup_doc_mask_loss.py @@ -1164,11 +1164,52 @@ def test_pad_tail_excluded_from_indexer_loss_warmup(self): # unscaled. The ``/cp_size`` branch is a multi-card concern. self.assertEqual(cap["coeff"], module.indexer_loss_coeff) + def test_unset_pad_token_id_falls_back_to_zero(self): + """``pad_token_id=None`` masks the same rows as ``0``; it is not fatal. + + ``TransformerConfig.from_config`` can copy a ``None`` straight out of an + external/HF config, and every other consumer of the field in this + repository folds it to ``0`` (``gpt_embedding.py:214-216``, + ``mtp_embedding_layer.py:105-107``, ``moe_router.py:605-607`` and four + more sites). The earlier revision of ``_indexer_loss_mask`` asserted + ``is not None`` instead, which aborted such a run at its *first* indexer + loss -- and under ``python -O``, where asserts are stripped, would have + compared the ids against ``None`` and masked nothing. Flagged in the + upstream review of PR #1721; this pins the fallback. + """ + seqlen, real_tokens = 256, 200 + ids = np.zeros([1, seqlen], dtype="int64") + ids[0, :real_tokens] = np.arange(1, real_tokens + 1) + ids_t = paddle.to_tensor(ids) + + module = _warmup_module() + module.config.pad_token_id = None + mask_none, rows_none = module._indexer_loss_mask(ids_t, 1, seqlen) + + reference = _warmup_module() + self.assertEqual(reference.config.pad_token_id, 0) + mask_zero, rows_zero = reference._indexer_loss_mask(ids_t, 1, seqlen) + + self.assertEqual(rows_none, float(real_tokens)) + self.assertEqual(rows_none, rows_zero) + np.testing.assert_array_equal(mask_none.numpy(), mask_zero.numpy()) + + # And the whole loss path, not just the helper: the fallback has to reach + # both the KL sum and the denominator the backward divides by. + logged, cap = self._step(module, seqlen, [seqlen], input_ids=ids_t) + self.assertEqual(float(cap["mask"].sum()), float(real_tokens)) + self.assertEqual(cap["num_rows"], float(real_tokens)) + kl_per_row = self._kl_per_row(cap) + ref = float( + (kl_per_row * cap["mask"]).sum() / real_tokens * cap["coeff"] + ) + self.assertLess(abs(logged - ref) / abs(ref), 1e-5) + def test_no_input_ids_uses_the_plain_row_mean(self): """Without ``input_ids`` the reduction is ``kl.mean() * coeff``. ``_indexer_loss_mask`` returns ``(None, None)`` - (``hybrid_mla_indexer.py:203-205``) and ``_attach_indexer_loss`` passes + (``hybrid_mla_indexer.py:212-213``) and ``_attach_indexer_loss`` passes that straight down, which is the same unmasked branch ``csa_attention._compute_fused_indexer_target`` takes: the backward then falls back to the kernel's own ``1/(B*Sq)``, and only the *logged* scalar From ea48c10abeb5985fbdd8e2ed7cd2b3ddb8dc4d2b Mon Sep 17 00:00:00 2001 From: shenliang03 Date: Wed, 12 Aug 2026 17:19:31 +0800 Subject: [PATCH 3/3] [release/0.4][Improvements] Drop the resynced recompute/GA keys from the phase-2 drift sentinel `TestConfigDeltas` pins, key by key, how each experiment YAML may differ from the phase-1 baseline, and it fails both ways: an unlisted difference fails, and an allowlisted key that stops differing fails too (so the allowlist cannot rot into a blanket exemption). Now that the warmup phase runs the baseline's dense attention -- its `-2` layers build `MHADSAWarmupAttention`, a `DotProductAttention` subclass -- its memory/compute picture is phase 1's, and the phase-2 YAML was resynced to the baseline's `recompute_granularity: selective` + module list and `gradient_accumulation_steps: 2`. Those five keys therefore no longer differ for that one config and the sentinel's second assertion fires. Remove them from its allowlist; the three latent-MQA variants keep them, because their picture really does differ. Only reachable where the parent config repository is checked out above PaddleFleet (the YAMLs live there), so CI skips it. --- .../test_hybrid_mla_config_pipeline.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py b/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py index 3441a67a4e..5ffd041f05 100644 --- a/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py +++ b/tests/single_card_tests/transformer/test_hybrid_mla_config_pipeline.py @@ -999,6 +999,22 @@ def _yaml_allowlist(self, name): # ``dsa_indexer_rope_fusion`` (``dsa_attention.py:458-461``). if name != _DSA: allowed["apply_rope_fusion"] = (True, False) + else: + # Phase 2's memory/compute picture *is* phase 1's now that its ``-2`` + # layers run the same dense attention (``MHADSAWarmupAttention`` + # subclasses ``DotProductAttention``), so the more conservative + # full/uniform recompute and GA=1 this yaml forked with have no + # reason left to exist and it was resynced to the baseline. The keys + # stay in ``_YAML_COMMON_DELTA`` for the three latent-MQA variants, + # whose picture really is different. + for key in ( + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "gradient_accumulation_steps", + ): + allowed.pop(key) if name in _MQA_DSA_CFGS: # ``indexer_init_from_scratch`` is mandatory once an indexer exists # (``modeling.py`` hard-errors on ``None``), so both mqa_dsa phases