diff --git a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py index 8162c9f3d5..8a972bf77d 100644 --- a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py +++ b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py @@ -22,7 +22,7 @@ import paddle import paddle.nn.functional as F -from paddlefleet.fusions.csa_sparse_attn_utils import _local_to_global_flat +from paddlefleet.fusions.csa_sparse_attn_utils import local_to_global_flat try: from paddlefleet_ops.flash_mla import ( @@ -55,6 +55,7 @@ def flash_mla_sparse_attn( indexer_topk: int = 0, d_v=None, topk_length=None, + global_kv_idx_remap_fusion: bool = False, ): if _flash_mla_sparse_fwd is None: raise RuntimeError("flash_mla is not available") @@ -69,7 +70,9 @@ def flash_mla_sparse_attn( q_flat = q.reshape([b * sq, h, d]) kv_flat = kv.reshape([b * skv, d]) - global_idxs = _local_to_global_flat(topk_idxs, skv) + global_idxs = local_to_global_flat( + topk_idxs, skv, fused=global_kv_idx_remap_fusion + ) # [b, sq] -> [b * sq]: one valid-prefix length per flattened query row. topk_length_flat = ( None diff --git a/src/paddlefleet/fusions/csa_sparse_attn.py b/src/paddlefleet/fusions/csa_sparse_attn.py index 752a3e88b0..1ffaabb3a5 100644 --- a/src/paddlefleet/fusions/csa_sparse_attn.py +++ b/src/paddlefleet/fusions/csa_sparse_attn.py @@ -160,6 +160,7 @@ def forward( backend, topk_length=None, indexer_topk=0, + global_kv_idx_remap_fusion=False, ): from paddlefleet.fusions.csa_sparse_attn_utils import prepare_inputs @@ -168,6 +169,7 @@ def forward( ctx.softmax_scale = float(softmax_scale) ctx.attn_sink_dtype = attn_sink.dtype ctx.backend = backend + ctx.global_kv_idx_remap_fusion = global_kv_idx_remap_fusion # ``topk_length`` is a forward-only early-stop hint: correctness comes # from the ``-1`` padding in ``topk_idxs``, which backward already turns # into its own bound via ``_csa_compute_topk_length``. Only remember @@ -200,6 +202,7 @@ def forward( sm_scale=ctx.softmax_scale, topk_length=topk_length, indexer_topk=indexer_topk, + global_kv_idx_remap_fusion=global_kv_idx_remap_fusion, ) CSASparseAttention._lse_indexer = lse_indexer else: @@ -228,7 +231,7 @@ def backward(ctx, grad_output): if ctx.backend == "cudnn": from paddlefleet.cudnn_ops import csa_sparse_attn_bwd_cudnn from paddlefleet.fusions.csa_sparse_attn_utils import ( - _local_to_global_flat, + local_to_global_flat, ) _, s_kv, dkv_dim = kv_full.shape @@ -238,7 +241,9 @@ def backward(ctx, grad_output): do_flat = grad_output.reshape([b * sq, np_heads, hn]) kv_flat = kv_full.reshape([b * s_kv, dkv_dim]) lse_flat = lse.reshape([b * sq, np_heads]) - topk_idxs_flat = _local_to_global_flat(topk_idxs, s_kv) + topk_idxs_flat = local_to_global_flat( + topk_idxs, s_kv, fused=ctx.global_kv_idx_remap_fusion + ) # SM90's backward kernel is unsafe with a trailing topk_length bound # (see _csa_bwd_honours_topk_length_holes), so fall back to None there. @@ -306,6 +311,7 @@ def csa_sparse_attn( backend="tilelang", topk_length=None, indexer_topk=0, + global_kv_idx_remap_fusion=False, ): """Unified CSA sparse attention entry point. @@ -315,6 +321,11 @@ def csa_sparse_attn( row. Only the "cudnn" and "unfused" backends support it; it lets the kernel stop early instead of walking all ``topk`` slots, which is what makes the full-causal MQA layers affordable. + global_kv_idx_remap_fusion: use the fused Triton local->global KV + column index remap instead of the eager elementwise chain. + Bit-identical either way; wired from the + ``sparse_attn_global_kv_idx_remap_fusion`` config field. Only the + "cudnn" backend builds that table, so it is ignored otherwise. """ if backend == "unfused": return unfused_compressed_sparse_attn( @@ -339,6 +350,7 @@ def csa_sparse_attn( backend, topk_length, indexer_topk, + global_kv_idx_remap_fusion, ) if CSASparseAttention._lse_indexer is not None: lse_indexer = CSASparseAttention._lse_indexer diff --git a/src/paddlefleet/fusions/csa_sparse_attn_utils.py b/src/paddlefleet/fusions/csa_sparse_attn_utils.py index d18c6b9499..0d3c729fea 100644 --- a/src/paddlefleet/fusions/csa_sparse_attn_utils.py +++ b/src/paddlefleet/fusions/csa_sparse_attn_utils.py @@ -68,3 +68,18 @@ def _local_to_global_flat(local_idxs, seqlen_kv: int): return paddle.where(valid, idxs_flat + batch_offsets, idxs_flat).cast( "int32" ) + + +def local_to_global_flat(local_idxs, seqlen_kv: int, *, fused: bool = False): + """Dispatch ``_local_to_global_flat`` between eager and fused Triton. + + ``fused`` comes from the ``sparse_attn_global_kv_idx_remap_fusion`` config + field, threaded down through the sparse-attn PyLayers. It is bit-identical + to the eager reference (one kernel instead of seven), so this switch only + trades kernel count for a Triton dependency -- it never changes values. + """ + if fused: + from paddlefleet.triton_ops import local_to_global_flat_triton + + return local_to_global_flat_triton(local_idxs, seqlen_kv) + return _local_to_global_flat(local_idxs, seqlen_kv) diff --git a/src/paddlefleet/fusions/mqa_sparse_attn.py b/src/paddlefleet/fusions/mqa_sparse_attn.py index b864068d0f..efdd27fc99 100644 --- a/src/paddlefleet/fusions/mqa_sparse_attn.py +++ b/src/paddlefleet/fusions/mqa_sparse_attn.py @@ -92,6 +92,7 @@ def forward( attn_sink=None, indexer_topk=0, sink_grad_fusion=False, + global_kv_idx_remap_fusion=False, backward_backend="cudnn", ): from paddlefleet.cudnn_ops.attn.csa_sparse_attn_fwd_cudnn import ( @@ -138,6 +139,7 @@ def forward( # output / gradient). The sink gradient is routed back to the parameter. ctx.learnable_sink = attn_sink is not None ctx.sink_grad_fusion = sink_grad_fusion + ctx.global_kv_idx_remap_fusion = global_kv_idx_remap_fusion if attn_sink is None: sink = paddle.full([_DSA_HEADS], _NEG_SINK, dtype="float32") else: @@ -179,6 +181,7 @@ def forward( d_v=d_v, topk_length=topk_len_flat.reshape([b, s]), indexer_topk=int(indexer_topk), + global_kv_idx_remap_fusion=global_kv_idx_remap_fusion, ) # out [b, s, 64, d_v], lse [b, s, 64] _MQASparseAttention._lse_indexer = lse_indexer @@ -261,14 +264,19 @@ def backward(ctx, grad_output): else: from paddlefleet.cudnn_ops import csa_sparse_attn_bwd_cudnn from paddlefleet.fusions.csa_sparse_attn_utils import ( - _local_to_global_flat, + local_to_global_flat, ) q_flat = q_pad.reshape([b * s, hpad, dk]) o_flat = out.reshape([b * s, hpad, d_v]) do_flat = do.reshape([b * s, hpad, d_v]) kv_flat = kv.reshape([b * skv, dk]) - gidx_flat = _local_to_global_flat(token_indices, skv) + # Only the cuDNN backward needs the flat-global column ids; the + # tilelang kernel above indexes ``token_indices`` per batch itself, + # so the remap -- fused or eager -- has no place on that branch. + gidx_flat = local_to_global_flat( + token_indices, skv, fused=ctx.global_kv_idx_remap_fusion + ) # dq/dkv softmax normalization for the finite-sink absorbed-MQA path. # @@ -400,9 +408,10 @@ def backward(ctx, grad_output): d_attn_sink = d_attn_sink.cast(ctx.attn_sink_dtype) # One returned grad per **tensor** input, in order. Non-tensor inputs - # (sm_scale, d_v, backward_backend) occupy no slot. ``attn_sink`` - # occupies a slot only when it was passed as a tensor (sinkless -> None - # -> no slot), so the returned count is 3 (sinkless) or 4 (learnable sink). + # (sm_scale, d_v, global_kv_idx_remap_fusion, backward_backend) occupy + # no slot. ``attn_sink`` occupies a slot only when it was passed as a + # tensor (sinkless -> None -> no slot), so the returned count is 3 + # (sinkless) or 4 (learnable sink). grads = [dq, dkv, None] # query, kv, token_indices if ctx.learnable_sink: grads.append(d_attn_sink) @@ -418,6 +427,7 @@ def mqa_sparse_attn( attn_sink=None, indexer_topk=0, sink_grad_fusion=False, + global_kv_idx_remap_fusion=False, backward_backend="cudnn", ): """Absorbed-MQA sparse attention (FlashMLA sparse fwd + selectable bwd). @@ -446,6 +456,11 @@ def mqa_sparse_attn( from the ``dsa_sink_grad_fusion`` config field; HySparse's ``block_sparse_mqa_attention_dsa`` leaves it at the default and keeps the eager epilogue by design. + global_kv_idx_remap_fusion: use the fused Triton local->global KV + column index remap in the forward and in the ``"cudnn"`` + backward instead of the eager elementwise chain. + Bit-identical either way; wired from the + ``sparse_attn_global_kv_idx_remap_fusion`` config field. backward_backend: ``"cudnn"`` (default, fast, non-deterministic dkv) or ``"tilelang"`` (deterministic, ~14x slower on SM100). @@ -462,6 +477,7 @@ def mqa_sparse_attn( attn_sink, int(indexer_topk), sink_grad_fusion, + global_kv_idx_remap_fusion, str(backward_backend), ) lse_indexer = _MQASparseAttention._lse_indexer diff --git a/src/paddlefleet/transformer/csa_attention.py b/src/paddlefleet/transformer/csa_attention.py index a705183f9e..19e156037a 100644 --- a/src/paddlefleet/transformer/csa_attention.py +++ b/src/paddlefleet/transformer/csa_attention.py @@ -127,18 +127,32 @@ def _derive_csa_doc_boundaries( mask = startend_row_indices.flatten().cast("int64") positions = paddle.arange(seqlen, dtype="int64") - is_boundary = paddle.zeros([seqlen], dtype="bool") - is_boundary[0] = True - is_boundary[1:] = (positions[1:] == mask[:-1]) & (mask[1:] != mask[:-1]) + # Concat rather than ``is_boundary[0] = True``: assigning a Python bool into + # a device tensor issues a 1-byte pageable ``cudaMemcpy``, which blocks the + # host until the device queue drains. On the layer43 config that lands behind + # a DeepEP combine and costs ~2.9 ms per ``-2`` layer. + is_boundary = paddle.concat( + [ + paddle.ones([1], dtype="bool"), + (positions[1:] == mask[:-1]) & (mask[1:] != mask[:-1]), + ] + ) + + # ``doc_start_per_pos`` is the most recent boundary at or before t. Boundary + # positions increase, so the running max of ``is_boundary * positions`` is a + # forward fill; once the boundaries are materialised -- which this function + # needs anyway for ``doc_lens`` / ``doc_starts`` -- a forward fill is cumsum + # + gather. ``paddle.cummax`` over a single long row falls into a one-block + # scan (``KernelScanInnerWithIndices``): 2.3 ms at seqlen 65536 and 5.3 ms at + # 131072, versus ~0.04 ms for cumsum + gather at either length. + doc_starts_i64 = paddle.nonzero(is_boundary).flatten() + doc_id = paddle.cumsum(is_boundary.cast("int32"), axis=0) - 1 + doc_start_per_pos = paddle.gather(doc_starts_i64, doc_id, axis=0) - doc_start_per_pos = paddle.cummax( - is_boundary.cast("int64") * positions, axis=0 - ).values pos_in_doc = positions - doc_start_per_pos doc_len_per_pos = mask - doc_start_per_pos is_valid = pos_in_doc < doc_len_per_pos - doc_starts_i64 = paddle.nonzero(is_boundary).flatten() doc_lens = (mask[doc_starts_i64] - doc_starts_i64).cast("int32") doc_starts = doc_starts_i64 @@ -2204,6 +2218,9 @@ def __init__( config, "csa_indexer_backend", "tilelang" ) self.indexer_loss_coeff = getattr(config, "dsa_indexer_loss_coeff", 0.0) + self.global_kv_idx_remap_fusion = getattr( + config, "sparse_attn_global_kv_idx_remap_fusion", False + ) def _resolve_topk_effective(self, n_compressed: int): """Return the CSA indexer top-k width for current phase. @@ -3385,4 +3402,5 @@ def compressed_sparse_attn( backend=sparse_attn_backend, topk_length=topk_length, indexer_topk=indexer_topk, + global_kv_idx_remap_fusion=self.global_kv_idx_remap_fusion, ) diff --git a/src/paddlefleet/transformer/dsv4_hybrid_attention.py b/src/paddlefleet/transformer/dsv4_hybrid_attention.py index f888270e4c..d50fdc7465 100644 --- a/src/paddlefleet/transformer/dsv4_hybrid_attention.py +++ b/src/paddlefleet/transformer/dsv4_hybrid_attention.py @@ -1077,6 +1077,7 @@ def _full_attn_forward( b, sq, _ = core_attn_out.shape pos_dim = self.qk_pos_emb_head_dim nope_dim = self.v_head_dim - pos_dim + postmix_done = False if pos_dim > 0: core_attn_out = core_attn_out.reshape( @@ -1091,7 +1092,16 @@ def _full_attn_forward( # DSv4 reference uses pure norm-preserving RoPE; YaRN's mscale is not applied. mscale = 1.0 - if ( + if self._can_fuse_inv_rope_postmix(_in_full_recompute): + # Fused inverse RoPE + ungrouped VHA postmix. Bitwise identical + # to running the two separately, but never materialises the + # full-width rotated output. It consumes the postmix, so the + # standalone postmix block below must be skipped. + core_attn_out = self._apply_inv_rope_vha_postmix( + core_attn_out, freqs, nope_dim, pos_dim, mscale + ) + postmix_done = True + elif ( self.config.apply_rope_fusion and not self.config.high_precision_rope ): @@ -1127,7 +1137,7 @@ def _full_attn_forward( # output projection. When the whole block is already wrapped in a # full_attn RecomputeWithoutOutput, skip the nested selective recompute # (the full block recompute already frees these activations). - if self.use_vha_postmix: + if self.use_vha_postmix and not postmix_done: if ( self.recompute_vha_postmix and self.training @@ -1197,6 +1207,59 @@ def _gate(self, gate_source: Tensor, core_attn_out: Tensor) -> Tensor: core_attn_out = core_attn_out * paddle.nn.functional.sigmoid(gate) return core_attn_out + def _can_fuse_inv_rope_postmix(self, in_full_recompute: bool) -> bool: + """Whether the inverse RoPE can be folded into the postmix GEMM. + + Every rejected case falls back to the unfused pair, so this is a pure + performance switch: the eager RoPE path, high_precision_rope, the + grouped postmix topology (einsum, no [nh,nh] GEMM to split) and the + postmix's own selective recompute wrapper all keep working unchanged. + """ + if not getattr(self.config, "fuse_inv_rope_into_vha_postmix", False): + return False + if not self.use_vha_postmix or self.vha_postmix_grouped: + return False + if not self.config.apply_rope_fusion: + return False + if self.config.high_precision_rope: + return False + # Re-entering the fused PyLayer from a nested recompute wrapper buys + # nothing (the fusion already avoids the intermediate it would free). + if ( + self.recompute_vha_postmix + and self.training + and not in_full_recompute + ): + return False + return True + + def _apply_inv_rope_vha_postmix( + self, + attn_out: Tensor, + freqs: Tensor, + nope_dim: int, + pe_dim: int, + mscale: float, + ) -> Tensor: + """Inverse RoPE + ungrouped VHA postmix in one pass. + + attn_out: [b, sq, nh, v_head_dim]. Returns [b, sq, nh * v_head_dim], + matching what the unfused RoPE followed by ``_apply_vha_postmix`` would + return, bit for bit. The postmix matrix is rebuilt inside the fused op + exactly as ``_apply_vha_postmix``'s ungrouped branch builds it. + """ + from paddlefleet.triton_ops import fused_inv_rope_vha_postmix + + return fused_inv_rope_vha_postmix( + attn_out, + freqs, + self.vha_postmix_U, + self.vha_postmix_V, + nope_dim, + pe_dim, + mscale, + ) + def _apply_vha_postmix(self, attn_out: Tensor) -> Tensor: """Low-rank cross-head mixing of the attention output. diff --git a/src/paddlefleet/transformer/mqa_latent_attention.py b/src/paddlefleet/transformer/mqa_latent_attention.py index 40642cf0e0..1ef8b9d332 100644 --- a/src/paddlefleet/transformer/mqa_latent_attention.py +++ b/src/paddlefleet/transformer/mqa_latent_attention.py @@ -805,6 +805,10 @@ def __init__( # Fused Triton epilogue for the analytic sink gradient. Only reachable # when ``softmax_offset`` exists; a sinkless layer has no sink gradient. self.sink_grad_fusion = getattr(config, "dsa_sink_grad_fusion", False) + # Fused Triton local->global KV column index remap (bit-identical). + self.global_kv_idx_remap_fusion = getattr( + config, "sparse_attn_global_kv_idx_remap_fusion", False + ) # Row layout the indexer forward runs on; see ``mqa_indexer_cp_mode`` in # ``transformer_config`` and ``_indexer_topk_dualchunk``. Gated on a real # CP group: with ``cp_size == 1`` there is nothing to rebalance. @@ -1807,6 +1811,7 @@ def _sparse_attn( attn_sink=self.softmax_offset, indexer_topk=indexer_topk, sink_grad_fusion=self.sink_grad_fusion, + global_kv_idx_remap_fusion=self.global_kv_idx_remap_fusion, backward_backend=self.sparse_attn_backward_backend, ) diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index 8875e6464b..2ad0783a70 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -235,6 +235,24 @@ class TransformerConfig(ModelParallelConfig): True: within-group block-diagonal mixing that only recombines heads inside each o_group (mixing stays within a group).""" + fuse_inv_rope_into_vha_postmix: bool = False + """Fuse the HCA inverse RoPE into the ungrouped VHA postmix GEMM (DSv4 hybrid). + + The unfused path materialises ``inv_rope(O)`` as a full-width tensor and feeds + it to the postmix ``[nh,nh]`` GEMM, which costs one extra read+write of the + whole attention output plus a second live copy of it. Because RoPE only + touches the trailing ``qk_pos_emb_head_dim`` channels while the GEMM + contracts the head axis, the same result can be assembled from a full-width + GEMM on the *unrotated* output plus a narrow GEMM on the rotated pe channels, + which never needs the wide intermediate. + + Bitwise identical to the unfused path -- forward, activation gradient and the + postmix U/V gradients -- and asserted as such in + ``tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py``. Requires + ``use_vha_attention`` and ``apply_rope_fusion``, and is skipped for + ``vha_postmix_grouped``, ``high_precision_rope`` and when the postmix has its + own selective recompute wrapper.""" + use_vha_premix: bool = False """If True (and use_vha_attention is also True), replaces the DSv4 hybrid Q up-projection (linear_q_up_proj) with a structured VHA premix: the compressed Q is reshaped into @@ -1415,6 +1433,31 @@ class TransformerConfig(ModelParallelConfig): the switch otherwise rather than let it be a silent no-op. """ + sparse_attn_global_kv_idx_remap_fusion: bool = False + """Whether to fuse the per-batch-local -> flat-global KV column index remap + (``idx + b * seqlen_kv``) consumed by the cuDNN / FlashMLA sparse-attention + kernels (``csa_sparse_attn_utils._local_to_global_flat``). + + Not about MoE routing: these are KV *column* indices of the sparse-attention + support (window + compressed slots), not expert top-k ids. + + The eager version spends seven elementwise kernels on the full + ``[b * sq, topk]`` table (``full`` + ``greater_equal`` + ``arange`` + + ``expand`` + ``scale`` + ``add`` + ``where``) to express a single pass; the + Triton kernel does it in one. The result is bit-identical, so this only + trades kernel count for a Triton dependency and can be flipped freely. + + Scope: every ``_local_to_global_flat`` call site -- the ``"cudnn"`` + sparse-attention forward and backward of both + ``CompressedSparseAttention`` (HCA ``ratio=128`` and CSA/DSA + ``1 < ratio < 128`` layers) and ``MQALatentAttention``. No effect on the + ``"tilelang"`` / ``"unfused"`` backends, which never build the flat global + index table, nor on ``block_sparse_mqa_attention_dsa``, which leaves it at + the default. ``MQALatentAttention``'s forward is always FlashMLA, so it + remaps regardless of ``mqa_sparse_attn_backward_backend``; only its + backward follows that switch. + """ + stage1_overlap: bool = False """ overlap backward with sharding gradient reduce for non-pipeline parallelism diff --git a/src/paddlefleet/triton_ops/__init__.py b/src/paddlefleet/triton_ops/__init__.py index 64ce900d0b..59df0636f0 100644 --- a/src/paddlefleet/triton_ops/__init__.py +++ b/src/paddlefleet/triton_ops/__init__.py @@ -19,6 +19,11 @@ from .fused_sink_grad import fused_sink_grad from .fused_yarn_rope_freqs import fused_yarn_rope_freqs from .grouped_matmul_fusion import GroupedMatmulTriton, fused_grouped_matmul +from .inv_rope_vha_postmix_fusion import ( + InvRopeVhaPostmixFusion, + fused_inv_rope_vha_postmix, +) +from .local_to_global_idxs_fusion import local_to_global_flat_triton from .mla_rope_inplace_fusion import ( fused_apply_mla_rope_inplace, fused_apply_rope_half, @@ -54,6 +59,9 @@ "fused_apply_mla_rope_inplace", "fused_apply_rope_half", "fused_rope_cat_key", + "fused_inv_rope_vha_postmix", + "InvRopeVhaPostmixFusion", + "local_to_global_flat_triton", "ulysses_alltoall_fused_supported", "ulysses_single_all_to_all_fused", ] diff --git a/src/paddlefleet/triton_ops/inv_rope_vha_postmix_fusion.py b/src/paddlefleet/triton_ops/inv_rope_vha_postmix_fusion.py new file mode 100644 index 0000000000..cd58078c77 --- /dev/null +++ b/src/paddlefleet/triton_ops/inv_rope_vha_postmix_fusion.py @@ -0,0 +1,455 @@ +# 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. + +"""Fused HCA inverse RoPE + ungrouped VHA postmix. + +Computes ``out = M @ inv_rope(O)`` -- the ``pos_dim > 0`` inverse-RoPE block of +``DSv4HybridAttention`` followed by the ``grouped=False`` branch of +``_apply_vha_postmix`` -- without ever materialising ``inv_rope(O)``. + +Why a channel split is exact. RoPE only mixes the trailing ``pe_dim`` channels of +each head while the postmix GEMM contracts the *head* axis, so the channel axis +is a pure N dimension of the GEMM: + + out[t,h,c] = sum_h' M[h,h'] * roped(O)[t,h',c] + +Splitting that N axis leaves the accumulation order untouched, so +``matmul(M, X[..., a:b]) == matmul(M, X)[..., a:b]`` bit for bit. That is a +cuBLAS property rather than a documented guarantee, so it is asserted in +``tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py`` alongside every +other equality this module leans on. + +All arithmetic is shared with ``mla_rope_inplace_fusion`` rather than duplicated: +the rotation reuses that module's ``_mul_round_bf16`` (whose inline-PTX +``cvt.rn.bf16.f32`` is what blocks FFMA folding and keeps the fused result equal +to eager Paddle), its ``_fused_cos_sin`` prelude, and its in-place forward / +backward kernels. Duplicating any of it would let the two drift apart silently. +""" + +import paddle +import triton +import triton.language as tl + +from .mla_rope_inplace_fusion import ( + _fused_cos_sin, + _get_block_h, + _mul_round_bf16, + _rope_mla_inplace_bwd_kernel, + _rope_mla_inplace_fwd_kernel, +) + + +def _check_shape(t, nope_dim, pe_dim): + if t.dim() != 3: + raise ValueError(f"t must be [B*S, H, D]; got {t.shape}") + if t.stride(-1) != 1: + raise ValueError("t must have a contiguous last dim") + if t.shape[-1] != nope_dim + pe_dim: + raise ValueError( + f"t last dim {t.shape[-1]} != nope_dim + pe_dim " + f"({nope_dim} + {pe_dim})" + ) + if t.shape[1] % _get_block_h(t.shape[1]) != 0: + raise ValueError(f"head_num {t.shape[1]} not divisible by its BLOCK_H") + + +def _check_rope_args(t, cos, sin, nope_dim, pe_dim): + _check_shape(t, nope_dim, pe_dim) + if pe_dim % 4 != 0: + raise ValueError(f"pe_dim must be a multiple of 4; got {pe_dim}") + if not cos.is_contiguous() or not sin.is_contiguous(): + raise ValueError("cos/sin must be contiguous") + if cos.shape[-1] != pe_dim or sin.shape[-1] != pe_dim: + raise ValueError( + f"cos/sin last dim must be pe_dim={pe_dim}; got " + f"{cos.shape[-1]}/{sin.shape[-1]}" + ) + + +def build_mla_rope_cos_sin( + freqs: paddle.Tensor, + b: int, + s: int, + pe_dim: int, + mscale: float, + inverse: bool, + dtype: paddle.dtype, +) -> tuple[paddle.Tensor, paddle.Tensor]: + """cos/sin prelude of ``fused_apply_mla_rope_inplace``, exposed separately. + + Driving the raw kernels only stays bitwise equal to the standalone op if the + cos/sin tensors are built identically, so keep the checks here in sync with + that wrapper's. + """ + if freqs.dim() != 4 or freqs.shape[2] != 1: + raise ValueError(f"freqs must be [B,S,1,D]; got {freqs.shape}") + b_f, s_f, _, d_f = freqs.shape + if s_f != s or d_f != pe_dim: + raise ValueError( + f"freqs {freqs.shape} mismatches [B,S]=[{b},{s}], pe_dim={pe_dim}" + ) + if not (b_f == 1 or b_f == b): + raise ValueError(f"freqs B {b_f} must be 1 or {b}") + if b_f < b: + freqs = freqs.broadcast_to([b, s, 1, pe_dim]) + return _fused_cos_sin(freqs, mscale, inverse, dtype) + + +@triton.jit +def _rope_pe_gather_kernel( + T, + T_OUT, + COS, + SIN, + nope_dim, + pe_dim: tl.constexpr, + head_num: tl.constexpr, + stride_in_seq, + stride_in_nheads, + stride_out_seq, + stride_out_nheads, + BLOCK_H: tl.constexpr, +): + """Rotate ``T[..., nope_dim:]`` into a *compact* ``[B*S, H, pe_dim]`` buffer. + + Same loads and same ``_mul_round_bf16`` sequence as + ``_rope_mla_inplace_fwd_kernel``; only the output addressing differs (the + input keeps the wide row stride and the ``nope_dim`` channel offset, the + output is densely packed). The forward needs the rotated pe channels as a + standalone GEMM operand and never wants the nope channels copied. + """ + pid_m = tl.program_id(axis=0).to(tl.int64) + pid_head = tl.program_id(axis=1).to(tl.int64) + + half: tl.constexpr = pe_dim // 2 + + cos_left = tl.load(COS + pid_m * pe_dim + tl.arange(0, half)) + sin_left = tl.load(SIN + pid_m * pe_dim + tl.arange(0, half)) + cos_right = tl.load(COS + pid_m * pe_dim + half + tl.arange(0, half)) + sin_right = tl.load(SIN + pid_m * pe_dim + half + tl.arange(0, half)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, half) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, half) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, half) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, half) + + T = T + pid_m * stride_in_seq + pid_head * BLOCK_H * stride_in_nheads + T_OUT = ( + T_OUT + pid_m * stride_out_seq + pid_head * BLOCK_H * stride_out_nheads + ) + head_mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + in_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_in_nheads.to(tl.int64) + + nope_dim + + tl.arange(0, pe_dim)[None, :] + ) + out_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_out_nheads.to(tl.int64) + + tl.arange(0, pe_dim)[None, :] + ) + + x = tl.load(T + in_off, mask=head_mask) + x = tl.reshape(x, (BLOCK_H, half, 2)) + x_1, x_2 = tl.split(x) + + y_left = _mul_round_bf16(x_1, cos_left).to(tl.float32) - _mul_round_bf16( + x_2, sin_left + ).to(tl.float32) + y_right = _mul_round_bf16(x_2, cos_right).to(tl.float32) + _mul_round_bf16( + x_1, sin_right + ).to(tl.float32) + + y = tl.join(y_left, y_right) + y = tl.reshape(y, (BLOCK_H, pe_dim)) + tl.store(T_OUT + out_off, y, mask=head_mask) + + +@triton.jit +def _pe_scatter_kernel( + SRC, + DST, + pe_dim: tl.constexpr, + head_num: tl.constexpr, + src_seq_stride, + src_head_stride, + dst_seq_stride, + dst_head_stride, + dst_chan_off, + BLOCK_H: tl.constexpr, +): + """``DST[..., dst_chan_off:] = SRC`` for a compact ``[M, H, pe_dim]`` source. + + Paddle can express this as ``t[..., nope:] = compact``, but its strided + elementwise copy does not vectorise the pattern: at nope=448/pe=64 it spends + 56 us on a 128 MiB slice against 49 us here, and the matching gather is 7x + off (332 us). One contiguous ``pe_dim``-wide load/store per head fixes it. + """ + pid_m = tl.program_id(axis=0).to(tl.int64) + pid_head = tl.program_id(axis=1).to(tl.int64) + + head_mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + rows = tl.arange(0, BLOCK_H)[:, None] + chan = tl.arange(0, pe_dim)[None, :] + + src = SRC + pid_m * src_seq_stride + pid_head * BLOCK_H * src_head_stride + dst = DST + pid_m * dst_seq_stride + pid_head * BLOCK_H * dst_head_stride + src_off = rows * src_head_stride.to(tl.int64) + chan + dst_off = rows * dst_head_stride.to(tl.int64) + dst_chan_off + chan + + tl.store( + dst + dst_off, tl.load(src + src_off, mask=head_mask), mask=head_mask + ) + + +def rope_pe_to_compact( + t: paddle.Tensor, + cos: paddle.Tensor, + sin: paddle.Tensor, + nope_dim: int, + pe_dim: int, +) -> paddle.Tensor: + """Rotate the pe channels of ``t`` [B*S, H, D] into a fresh compact buffer. + + Returns [B*S, H, pe_dim], bitwise equal to + ``fused_apply_mla_rope_inplace(t, ...)[..., nope_dim:]``. + """ + _check_rope_args(t, cos, sin, nope_dim, pe_dim) + m, h, _ = t.shape + out = paddle.empty([m, h, pe_dim], dtype=t.dtype) + block_h = _get_block_h(h) + _rope_pe_gather_kernel[(m, triton.cdiv(h, block_h))]( + t, + out, + cos, + sin, + nope_dim, + pe_dim, + h, + t.stride(0), + t.stride(1), + out.stride(0), + out.stride(1), + block_h, + ) + return out + + +def rope_full_out_of_place( + t: paddle.Tensor, + cos: paddle.Tensor, + sin: paddle.Tensor, + nope_dim: int, + pe_dim: int, +) -> paddle.Tensor: + """RoPE the pe channels of ``t`` [B*S, H, D] into a fresh [B*S, H, D] buffer. + + Same kernel and arguments as ``RoPEMLAInplaceFusion.forward`` with + ``clone_input=True``, hence bitwise equal to it. + """ + _check_rope_args(t, cos, sin, nope_dim, pe_dim) + m, h, d = t.shape + out = paddle.empty([m, h, d], dtype=t.dtype) + block_h = _get_block_h(h) + _rope_mla_inplace_fwd_kernel[(m, triton.cdiv(h, block_h))]( + t, + out, + cos, + sin, + nope_dim, + pe_dim, + h, + t.stride(0), + t.stride(1), + block_h, + triton.next_power_of_2(max(nope_dim, 1)), + True, + ) + return out + + +def rope_pe_transpose_inplace_( + t: paddle.Tensor, + cos: paddle.Tensor, + sin: paddle.Tensor, + nope_dim: int, + pe_dim: int, +) -> paddle.Tensor: + """Apply the transpose rotation to ``t[..., nope_dim:]`` in place. + + Same kernel and arguments as ``RoPEMLAInplaceFusion.backward``, so a gradient + pushed through here matches the standalone op bit for bit. + """ + _check_rope_args(t, cos, sin, nope_dim, pe_dim) + m, h, _ = t.shape + block_h = _get_block_h(h) + _rope_mla_inplace_bwd_kernel[(m, triton.cdiv(h, block_h))]( + t, + cos, + sin, + nope_dim, + pe_dim, + h, + t.stride(0), + t.stride(1), + block_h, + ) + return t + + +def scatter_pe_slice_( + t: paddle.Tensor, compact: paddle.Tensor, nope_dim: int, pe_dim: int +) -> paddle.Tensor: + """``t[..., nope_dim:] = compact`` for a [B*S, H, D] tensor, vectorised.""" + _check_shape(t, nope_dim, pe_dim) + if list(compact.shape) != [t.shape[0], t.shape[1], pe_dim]: + raise ValueError( + f"compact must be {[t.shape[0], t.shape[1], pe_dim]}; " + f"got {compact.shape}" + ) + m, h, _ = t.shape + block_h = _get_block_h(h) + _pe_scatter_kernel[(m, triton.cdiv(h, block_h))]( + compact, + t, + pe_dim, + h, + compact.stride(0), + compact.stride(1), + t.stride(0), + t.stride(1), + nope_dim, + block_h, + ) + return t + + +class InvRopeVhaPostmixFusion(paddle.autograd.PyLayer): + """``out = M @ inv_rope(O)`` without ever materialising ``inv_rope(O)``. + + The unfused pair costs 4N of traffic and keeps two full-width copies of the + attention output alive; this costs 2.75N and keeps one. + + Forward: + pe_roped = rope(O[..., nope:]) -> compact [B*S, nh, pe] + out = matmul(M, O) -> full width; the pe block it computes + from unrotated data is discarded + out[..., nope:] = matmul(M, pe_roped) + + Backward. The composite really is ``out = M @ O_roped``, so: + + - The activation gradient needs no rotated operand at all: ``dO_roped = + M^T @ dOut`` full width, then the transpose rotation in place on the pe + block. That is exactly the sequence the unfused path runs, hence bitwise + identical. + - The weight gradient ``dM[h,h'] = sum_{t,c} dOut[t,h,c] * O_roped[t,h',c]`` + contracts over the *full* channel axis, so it cannot be assembled from two + partial sums without changing the reduction tree. The rotated tensor is + therefore rebuilt here and both gradients are taken from + ``paddle._C_ops.matmul_grad`` -- the same op the unfused postmix GEMM's + backward calls, so they match by construction on every architecture. + + Hand-rolling that GEMM over head-major ``[H, B*S, D]`` operands (which is + what ``matmul_grad`` internally builds with two ``TilingSwapDim1And2`` + passes at 2.5 TB/s) is ~870us/layer faster at the production shape, and + was bitwise equal across every shape tested on sm10.3 -- but not on sm90, + where CI caught ``head-major wgrad (128,4,64)`` differing in 8/16 elements. + For a small ``[nh,nh] x K`` GEMM cuBLAS selects its algorithm per + architecture, so no shape-based gate can make that route safe and it was + dropped. Only the *forward* split still rests on a cuBLAS property + (``matmul(M, X[..., a:b]) == matmul(M, X)[..., a:b]``), which is asserted + over a wide shape sweep in the unit test. + + Nothing is ever mutated in place on the saved attention output, which the CSA + backward also reads for its ``delta = rowsum(dO * O)``. + """ + + @staticmethod + def forward(ctx, o_flat, m, cos, sin, nope_dim, pe_dim): + pe_roped = rope_pe_to_compact(o_flat, cos, sin, nope_dim, pe_dim) + out = paddle.matmul(m, o_flat) + # The pe block `out` just computed came from unrotated channels; replace + # it with the narrow GEMM on the rotated operand. + scatter_pe_slice_(out, paddle.matmul(m, pe_roped), nope_dim, pe_dim) + ctx.save_for_backward(o_flat, m, cos, sin) + ctx.nope_dim = nope_dim + ctx.pe_dim = pe_dim + ctx.m_needs_grad = not m.stop_gradient + return out + + @staticmethod + def backward(ctx, d_out): + o_flat, m, cos, sin = ctx.saved_tensors + nope_dim, pe_dim = ctx.nope_dim, ctx.pe_dim + if not d_out.is_contiguous(): + d_out = d_out.contiguous() + + if not ctx.m_needs_grad: + # Frozen postmix (an indexer-only warmup stage, say): only the + # activation gradient is wanted, and that never needs the rotated + # operand, so skip rebuilding it. + d_o = paddle.matmul(m, d_out, transpose_x=True) + rope_pe_transpose_inplace_(d_o, cos, sin, nope_dim, pe_dim) + return d_o, None, None, None + + # Rebuild the rotated tensor and hand both gradients to the very op the + # unfused path used, so they come out of the same kernels by + # construction rather than by a cuBLAS coincidence. Hand-rolling the + # weight-gradient GEMM over head-major operands is ~870us/layer faster + # at the production shape and was bitwise equal on sm10.3, but *not* on + # sm90 (CI caught `head-major wgrad (128,4,64)` differing by 8/16 + # elements): for a small [nh,nh] x K GEMM cuBLAS picks its algorithm per + # architecture, so no shape-based gate can make that route safe. + o_roped = rope_full_out_of_place(o_flat, cos, sin, nope_dim, pe_dim) + d_m, d_o = paddle._C_ops.matmul_grad(m, o_roped, d_out, False, False) + del o_roped + + # d_o is the gradient wrt the rotated tensor; push it back through the + # rotation exactly as the standalone RoPE op's backward does. + rope_pe_transpose_inplace_(d_o, cos, sin, nope_dim, pe_dim) + return d_o, d_m, None, None + + +def fused_inv_rope_vha_postmix( + attn_out: paddle.Tensor, + freqs: paddle.Tensor, + postmix_u: paddle.Tensor, + postmix_v: paddle.Tensor, + nope_dim: int, + pe_dim: int, + mscale: float = 1.0, +) -> paddle.Tensor: + """Inverse RoPE + ungrouped VHA postmix in one pass. + + Args: + attn_out: [b, sq, nh, v_head_dim], contiguous bf16 attention output. + freqs: [b_or_1, sq, 1, pe_dim] fp32 angle tensor. + postmix_u, postmix_v: the [nh, rank] postmix factors. + + Returns: + [b, sq, nh * v_head_dim], bitwise equal to what the unfused RoPE followed + by ``_apply_vha_postmix``'s ungrouped branch returns. + """ + b, sq, nh, d = attn_out.shape + if not attn_out.is_contiguous(): + attn_out = attn_out.contiguous() + cos, sin = build_mla_rope_cos_sin( + freqs, b, sq, pe_dim, mscale, True, attn_out.dtype + ) + # Same construction as _apply_vha_postmix's ungrouped branch; keep the two + # in sync or the fused path stops being bitwise equal. + m = paddle.matmul(postmix_v, postmix_u, transpose_y=True) + m = m + paddle.eye(nh, dtype=m.dtype) + out = InvRopeVhaPostmixFusion.apply( + attn_out.reshape([b * sq, nh, d]), m, cos, sin, nope_dim, pe_dim + ) + return out.reshape([b, sq, nh * d]) diff --git a/src/paddlefleet/triton_ops/local_to_global_idxs_fusion.py b/src/paddlefleet/triton_ops/local_to_global_idxs_fusion.py new file mode 100644 index 0000000000..13d85790a0 --- /dev/null +++ b/src/paddlefleet/triton_ops/local_to_global_idxs_fusion.py @@ -0,0 +1,137 @@ +# 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. + +""" +Fused Triton replacement for ``csa_sparse_attn_utils._local_to_global_flat``. + +The eager reference spends five elementwise kernels on the full +``[b * sq, topk]`` index table -- ``full`` (the scalar in ``idxs >= 0`` is +materialised at full size), ``greater_equal``, ``add``, ``where``, ``cast`` +-- to express a single pass. At sq=16384 / topk=640 (ernielite HCA layer, +cp=4) that is ~380 MB of traffic per call for 40 MB of useful output. + +This module does the same thing in one kernel: one read, one write. + +Not differentiable by design -- inputs and outputs are integer index tables. +All three call sites are inside ``paddle.autograd.PyLayer`` bodies +(``fusions/csa_sparse_attn.py`` forward+backward, ``fusions/mqa_sparse_attn.py`` +backward, ``cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py``), which run with +grad tracking disabled and return ``None`` for the ``topk_idxs`` gradient. +""" + +import paddle +from paddle import Tensor + +from .utils import is_torch_compat_available + +if is_torch_compat_available(): + paddle.enable_compat(scope={"triton"}) + +import triton +import triton.language as tl + + +@triton.jit +def local_to_global_flat_kernel( + Idxs_ptr, # [n_rows, topk] local idxs, int32/int64, negative == invalid + Out_ptr, # [n_rows, topk] int32 global indices + topk, + sq, # rows per batch; flat row == b * sq + s + seqlen_kv, # KV length of one batch entry + BLOCK_K: tl.constexpr, # power-of-2 block along the topk axis +): + """One program per ``(row, topk-block)``. + + Computes ``out = idx + (row // sq) * seqlen_kv`` where ``idx >= 0`` and + passes ``idx`` through untouched otherwise (the reference keeps the + original negative value, it does not normalise it to -1). + + Arithmetic runs in int64 and is truncated on store. That matches the eager + reference bit-for-bit for both input dtypes: the low 32 bits of a + two's-complement sum are identical whether the reference accumulated in + int32 (wrapping ``add``, then a no-op ``cast``) or in int64 (exact ``add``, + then a truncating ``cast``). + + ``base`` is int64 so addressing stays correct past 2**31 elements. That is + defensive only: the widest real table is ``65536 * 2176`` (~142M elements), + so the unit tests cannot reach the 32-bit overflow point (it would need an + ~8.6 GB index tensor) and a 32-bit ``base`` passes them unchanged. + """ + row = tl.program_id(0) + kblk = tl.program_id(1) + + batch_offset = (row // sq).to(tl.int64) * seqlen_kv + + offs = kblk * BLOCK_K + tl.arange(0, BLOCK_K) + mask = offs < topk + base = row.to(tl.int64) * topk + + idx = tl.load(Idxs_ptr + base + offs, mask=mask, other=0).to(tl.int64) + result = tl.where(idx >= 0, idx + batch_offset, idx) + + tl.store( + Out_ptr + base + offs, + result.to(Out_ptr.dtype.element_ty), + mask=mask, + ) + + +def local_to_global_flat_triton( + local_idxs: Tensor, + seqlen_kv: int, + *, + allow_alias: bool = False, +) -> Tensor: + """Drop-in fused replacement for ``_local_to_global_flat``. + + Args: + local_idxs: ``[b, sq, topk]`` int32/int64 indices into one batch + entry's KV, negative == invalid slot. + seqlen_kv: KV sequence length per batch entry. + allow_alias: when ``b == 1`` the batch offset is 0 and the reference + reduces to the identity, so the result can be returned as a + reshaped *view* of ``local_idxs`` with no kernel at all. Off by + default because the reference always returns fresh storage; only + enable it where the caller treats the result as read-only. + + Returns: + ``[b * sq, topk]`` int32, bit-identical to the eager reference. + """ + assert local_idxs.ndim == 3, ( + f"local_idxs must be [b, sq, topk], got {local_idxs.shape}" + ) + b, sq, topk = local_idxs.shape + n_rows = b * sq + + if allow_alias and b == 1: + flat = local_idxs.reshape([n_rows, topk]) + return flat if flat.dtype == paddle.int32 else flat.cast("int32") + + out = paddle.empty([n_rows, topk], dtype="int32") + if n_rows == 0 or topk == 0: + return out + + idxs = local_idxs.contiguous() + BLOCK_K = min(triton.next_power_of_2(topk), 1024) + grid = (n_rows, triton.cdiv(topk, BLOCK_K)) + local_to_global_flat_kernel[grid]( + idxs, + out, + topk, + sq, + int(seqlen_kv), + BLOCK_K=BLOCK_K, + num_warps=4, + ) + return out diff --git a/src/paddlefleet/triton_ops/mla_rope_inplace_fusion.py b/src/paddlefleet/triton_ops/mla_rope_inplace_fusion.py index 0c885cf247..a4f6e7c85e 100644 --- a/src/paddlefleet/triton_ops/mla_rope_inplace_fusion.py +++ b/src/paddlefleet/triton_ops/mla_rope_inplace_fusion.py @@ -101,6 +101,7 @@ def _cos_sin_kernel( @triton.jit def _rope_mla_inplace_fwd_kernel( T, + T_OUT, COS, SIN, nope_dim, @@ -109,8 +110,19 @@ def _rope_mla_inplace_fwd_kernel( stride_x_seq, stride_x_nheads, BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + OUT_OF_PLACE: tl.constexpr, ): - """Forward: rotate t[..., nope_dim:] in place (interleaved in/out).""" + """Forward: rotate t[..., nope_dim:] (interleaved in/out). + + Reads from ``T`` and writes to ``T_OUT``. With ``OUT_OF_PLACE=False`` the + caller passes the same pointer for both and the kernel behaves exactly as + the original in-place version: the nope channels are never touched, so + there is no extra traffic and no extra allocation. With + ``OUT_OF_PLACE=True`` the nope channels are additionally copied across, + which lets the caller keep the input buffer intact without paying for a + separate ``clone()`` pass over the whole tensor. + """ pid_m = tl.program_id(axis=0).to(tl.int64) pid_head = tl.program_id(axis=1).to(tl.int64) @@ -130,19 +142,32 @@ def _rope_mla_inplace_fwd_kernel( # Pointer to the start of this token's (head_block_first) row, then advance # past the nope channels to land on the rope slice. - T = T + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads + row_off = pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads + T = T + row_off + T_OUT = T_OUT + row_off + head_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads.to(tl.int64) + head_mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + + # Out-of-place only: carry the untouched nope channels over to the output + # buffer. BLOCK_D is next_power_of_2(nope_dim) on the host, so this is a + # single masked pass; the in-place path compiles this branch away entirely. + if OUT_OF_PLACE: + offs_d = tl.arange(0, BLOCK_D) + nope_off = head_off + offs_d[None, :] + nope_mask = head_mask & (offs_d[None, :] < nope_dim) + tl.store( + T_OUT + nope_off, + tl.load(T + nope_off, mask=nope_mask), + mask=nope_mask, + ) + # Offsets into the rope slice: [BLOCK_H, pe_dim] with last dim STRIDE=1. # We deliberately load the whole pe_dim contiguously instead of poking # at 2k / 2k+1 with stride-2 offsets — Triton's lowering for stride-2 # int64 offsets has historically been flaky (extra sector requests, no # vectorization), and explicit contiguous loads compile down to # `ld.global.v4.b32` which is the theoretical optimum for bf16. - flat_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads.to(tl.int64) - + nope_dim - + tl.arange(0, pe_dim)[None, :] - ) - head_mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + flat_off = head_off + nope_dim + tl.arange(0, pe_dim)[None, :] # One contiguous load per program, then de-interleave in registers. x = tl.load(T + flat_off, mask=head_mask) # [BLOCK_H, pe_dim] bf16 @@ -166,7 +191,7 @@ def _rope_mla_inplace_fwd_kernel( # same 2k / 2k+1 positions) and store in one contiguous write. y = tl.join(y_left, y_right) # [BLOCK_H, half, 2] y = tl.reshape(y, (BLOCK_H, pe_dim)) - tl.store(T + flat_off, y, mask=head_mask) + tl.store(T_OUT + flat_off, y, mask=head_mask) @triton.jit @@ -228,8 +253,6 @@ class RoPEMLAInplaceFusion(paddle.autograd.PyLayer): @staticmethod def forward(ctx, t, cos, sin, nope_dim, pe_dim, clone_input): - # Clone input if the upstream depends on it. - t = t.clone() if clone_input else t assert t.stride(-1) == 1 assert cos.is_contiguous() assert sin.is_contiguous() @@ -246,9 +269,21 @@ def forward(ctx, t, cos, sin, nope_dim, pe_dim, clone_input): f"head_num must be divisible by BLOCK_H ({BLOCK_H}), got {H}" ) + # When the upstream still needs `t`, write to a fresh buffer instead of + # cloning first: `clone()` would read+write the whole tensor and the + # kernel would then read+write the rope slice again. Letting the kernel + # read `t` and write `out` (carrying the nope channels across on the + # way) is a single pass, and leaves `t` untouched just the same. + # clone_input=False keeps the true in-place behaviour: same pointer in + # and out, nope branch compiled away, no allocation. + out = paddle.empty(t.shape, dtype=t.dtype) if clone_input else t + out_flat = out.reshape([B * S, H, D]) if clone_input else t_flat + BLOCK_D = triton.next_power_of_2(max(nope_dim, 1)) + grid = (B * S, triton.cdiv(H, BLOCK_H)) _rope_mla_inplace_fwd_kernel[grid]( t_flat, + out_flat, cos, sin, nope_dim, @@ -257,6 +292,8 @@ def forward(ctx, t, cos, sin, nope_dim, pe_dim, clone_input): t_flat.stride(0), t_flat.stride(1), BLOCK_H, + BLOCK_D, + clone_input, ) ctx.save_for_backward(cos, sin) @@ -264,8 +301,9 @@ def forward(ctx, t, cos, sin, nope_dim, pe_dim, clone_input): ctx.pe_dim = pe_dim ctx.block_h = BLOCK_H ctx.shape = (B, S, H, D) - # Return the reshape-back view; storage is identical to input t. - return t + # clone_input=False returns the reshape-back view of the input, whose + # storage is identical to `t`; clone_input=True returns the new buffer. + return out @staticmethod def backward(ctx, grad): @@ -343,17 +381,21 @@ def fused_apply_mla_rope_inplace( Args: t: [B, S, H, nope_dim + pe_dim], contiguous, bf16 (or fp16/fp32). - Mutated in place. + Mutated in place unless clone_input=True. freqs: [B, S, 1, pe_dim], fp32 angle tensor. May be non-contiguous. nope_dim: number of leading nope channels left untouched. mscale: scaling factor for rotary embedding. inverse: if True, apply the inverse rotation (used by the inv_rope post-attention canonicalisation step). - clone_input: if True, clone the input t before applying rope. + clone_input: if True, leave `t` untouched and return a new tensor + instead (needed when the upstream still reads `t`, e.g. an + attention output that its own backward has saved). Returns: - t (same storage as the input). Channels [..., :nope_dim] are - unchanged; channels [..., nope_dim:] are rotated. + With clone_input=False, `t` itself (same storage as the input). + With clone_input=True, a freshly allocated tensor; `t` is not + modified. Either way channels [..., :nope_dim] carry the input's + nope values unchanged and [..., nope_dim:] are rotated. """ # Check t assert t.is_contiguous(), ( diff --git a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py index 98d01fa6d7..f89150acd9 100644 --- a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py +++ b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py @@ -577,6 +577,7 @@ def _fake_ctx(self): softmax_scale=0.125, attn_sink_dtype=attn_sink.dtype, backend="cudnn", + global_kv_idx_remap_fusion=False, has_topk_length=False, query_needs_grad=True, kv_full_needs_grad=True, diff --git a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_utils.py b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_utils.py index fb00b6b9eb..d8fe50f6e8 100644 --- a/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_utils.py +++ b/tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_utils.py @@ -312,6 +312,7 @@ def fake_fwd( sm_scale=None, indexer_topk=0, topk_length=None, + **kwargs, ): bb, ss, hh, dd = q.shape out = paddle.ones([bb, ss, hh, dd], dtype=q.dtype) diff --git a/tests/single_card_tests/ai_edited_test/transformer/test_ai_dsv4_hybrid_attention_recompute.py b/tests/single_card_tests/ai_edited_test/transformer/test_ai_dsv4_hybrid_attention_recompute.py index 9b1a8bad7c..f60ddeb8d0 100644 --- a/tests/single_card_tests/ai_edited_test/transformer/test_ai_dsv4_hybrid_attention_recompute.py +++ b/tests/single_card_tests/ai_edited_test/transformer/test_ai_dsv4_hybrid_attention_recompute.py @@ -181,6 +181,14 @@ def _fake_get_qkv( DSv4HybridAttention._full_attn_forward, inst ) inst._gate = types.MethodType(DSv4HybridAttention._gate, inst) + # _full_attn_forward consults this before the inverse-RoPE block. Bind the + # real gate rather than a stub so a change to its conditions surfaces here; + # with use_vha_postmix False it returns False, keeping these tests on the + # unfused path they were written for. + inst._can_fuse_inv_rope_postmix = types.MethodType( + DSv4HybridAttention._can_fuse_inv_rope_postmix, inst + ) + inst.vha_postmix_grouped = False return inst diff --git a/tests/single_card_tests/ai_edited_test/triton_ops/test_fused_sink_grad.py b/tests/single_card_tests/ai_edited_test/triton_ops/test_fused_sink_grad.py index 2e0f056019..f408eebbd8 100644 --- a/tests/single_card_tests/ai_edited_test/triton_ops/test_fused_sink_grad.py +++ b/tests/single_card_tests/ai_edited_test/triton_ops/test_fused_sink_grad.py @@ -361,6 +361,7 @@ def _make_bwd_case( attn_sink_dtype=paddle.empty([0], dtype=attn_sink_dtype).dtype, learnable_sink=learnable_sink, sink_grad_fusion=fusion, + global_kv_idx_remap_fusion=False, # The sink-gradient epilogue under test lives on the cuDNN branch only: # the tilelang backward takes ``d_sink`` from its own kernel and never # reaches ``fused_sink_grad`` (see ``mqa_sparse_attn.backward``). diff --git a/tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py b/tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py new file mode 100644 index 0000000000..c3d7418c49 --- /dev/null +++ b/tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py @@ -0,0 +1,743 @@ +# 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. + +"""Tests for the fused HCA inverse RoPE + ungrouped VHA postmix. + +`InvRopeVhaPostmixFusion` computes `out = M @ inv_rope(O)` without ever +materialising `inv_rope(O)`. It is only worth having if it is *bitwise* +identical to the unfused pair it replaces, so every check here is exact +equality with no tolerance: + + - forward output + - activation gradient (grad wrt the attention output) + - postmix weight gradients (grad wrt vha_postmix_U / vha_postmix_V) + - `O` byte-identical after backward: the wgrad rebuilds the rotated tensor + in place on the CSA-saved attention output and must restore it, and Paddle + has no version counter that would catch a failure to do so + - no full-width rotated intermediate is allocated + +Two cuBLAS/Paddle properties the fusion leans on are asserted directly, so a +library upgrade fails the test instead of silently drifting the loss curve: + + - splitting a GEMM along its N (channel) axis is bitwise equal to slicing the + full-width result -- the whole reason a channel split can be exact + - the explicit dgrad/wgrad formulas used in the fused backward reproduce what + Paddle's own matmul backward emits +""" + +import unittest + +import paddle + +from paddlefleet.transformer.dsv4_hybrid_attention import ( + DSv4HybridSelfAttention, +) +from paddlefleet.triton_ops import fused_apply_mla_rope_inplace +from paddlefleet.triton_ops.inv_rope_vha_postmix_fusion import ( + InvRopeVhaPostmixFusion, + build_mla_rope_cos_sin, + fused_inv_rope_vha_postmix, + rope_full_out_of_place, + rope_pe_to_compact, + scatter_pe_slice_, +) + +# ernielite_layer43_pretrain_non_absorbed_mqa_hca_dsa_sparse_loss.yaml at 64K +# context with context_parallel_size=4: b=1, sq=65536/4=16384, and +# num_attention_heads=64, v_head_dim=512, qk_pos_emb_head_dim=64 from +# model_config.json. vha_postmix_rank defaults to num_attention_heads // 4. +NH, VD, PE = 64, 512, 64 +NOPE = VD - PE +RANK = NH // 4 +PROD_BS = 16384 +# Smaller default so the whole suite does not need ~10 GiB; the production +# shape gets its own test. +BS = 2048 + + +def _check_equal(a: paddle.Tensor, b: paddle.Tensor, what: str = "") -> None: + """Binary-exact equality check (no tolerance allowed).""" + a32, b32 = a.astype("float32"), b.astype("float32") + if bool(paddle.all(a32 == b32)): + return + diff = (a32 - b32).abs() + raise AssertionError( + f"{what or 'tensor'} not bitwise equal: " + f"{int((a32 != b32).sum())}/{a.numel().item()} elements differ, " + f"max|diff|={float(diff.max()):.6e}" + ) + + +def _make_inputs(bs: int, seed: int = 0): + paddle.seed(seed) + o = paddle.randn([bs, NH, VD], "bfloat16") + u = (paddle.randn([NH, RANK], "float32") * 0.01).astype("bfloat16") + v = (paddle.randn([NH, RANK], "float32") * 0.01).astype("bfloat16") + freqs = paddle.randn([1, bs, 1, PE], "float32") + freqs.stop_gradient = True + d_out = paddle.randn([bs, NH, VD], "bfloat16") + return o, u, v, freqs, d_out + + +def _make_m(u: paddle.Tensor, v: paddle.Tensor, nh: int = NH) -> paddle.Tensor: + """Same construction as _apply_vha_postmix's ungrouped branch.""" + m = paddle.matmul(v, u, transpose_y=True) + return m + paddle.eye(nh, dtype=m.dtype) + + +def _leaves(o, u, v): + out = [] + for t in (o, u, v): + leaf = t.detach() + leaf.stop_gradient = False + out.append(leaf) + return out + + +def _run_reference(o, u, v, freqs, d_out): + """Unfused path: full-width inverse RoPE, then the postmix GEMM.""" + bs = o.shape[0] + o_l, u_l, v_l = _leaves(o, u, v) + # non-leaf, mirroring core_attn_out being an op output + x = o_l * 1.0 + roped = fused_apply_mla_rope_inplace( + x.reshape([1, bs, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + out = paddle.matmul(_make_m(u_l, v_l), roped.reshape([bs, NH, VD])) + out.backward(d_out.clone()) + return out.detach(), o_l.grad, u_l.grad, v_l.grad + + +def _run_fused(o, u, v, freqs, d_out): + bs = o.shape[0] + o_l, u_l, v_l = _leaves(o, u, v) + x = o_l * 1.0 + cos, sin = build_mla_rope_cos_sin( + freqs, 1, bs, PE, 1.0, True, paddle.bfloat16 + ) + out = InvRopeVhaPostmixFusion.apply( + x.reshape([bs, NH, VD]), _make_m(u_l, v_l), cos, sin, NOPE, PE + ) + out.backward(d_out.clone()) + return out.detach(), o_l.grad, u_l.grad, v_l.grad + + +class TestInvRopeVhaPostmixFusion(unittest.TestCase): + def _compare(self, bs: int, seed: int = 0) -> None: + args = _make_inputs(bs, seed) + ref = _run_reference(*args) + fused = _run_fused(*args) + for name, a, b in zip( + ("forward out", "grad_O", "grad_U", "grad_V"), fused, ref + ): + _check_equal(a, b, f"{name} (B*S={bs}, seed={seed})") + + def test_bitwise_default_shape(self) -> None: + self._compare(BS) + + def test_bitwise_multiple_seeds(self) -> None: + for seed in (1, 2, 3): + self._compare(BS, seed=seed) + + def test_bitwise_production_shape(self) -> None: + """64K context / CP=4 -> sq=16384, a 1 GiB attention output.""" + self._compare(PROD_BS) + + def test_bitwise_odd_token_count(self) -> None: + """Token count that is not a multiple of anything in particular.""" + self._compare(1023) + + def test_o_restored_after_backward(self) -> None: + """The saved attention output must come back untouched. + + Nothing downstream would complain if it did not: the CSA backward reads + the same buffer to build `delta = rowsum(dO * O)` and Paddle silently + accepts a mutated saved tensor. The fused backward builds its rotated + operand into a fresh transposed buffer precisely so that this holds. + """ + o, u, v, freqs, d_out = _make_inputs(BS) + snapshot = o.clone() + o_l, u_l, v_l = _leaves(o, u, v) + cos, sin = build_mla_rope_cos_sin( + freqs, 1, BS, PE, 1.0, True, paddle.bfloat16 + ) + # Pass the leaf's own storage through so the PyLayer saves a view of it. + x = o_l.reshape([BS, NH, VD]) + out = InvRopeVhaPostmixFusion.apply( + x, _make_m(u_l, v_l), cos, sin, NOPE, PE + ) + out.backward(d_out.clone()) + _check_equal(o_l, snapshot, "O after backward") + + def test_frozen_postmix_skips_wgrad(self) -> None: + """With U/V frozen only the activation gradient is produced. + + That path skips both transposed operands, so it must not allocate them + either -- checked indirectly by the gradient still matching exactly. + """ + o, u, v, freqs, d_out = _make_inputs(BS) + cos, sin = build_mla_rope_cos_sin( + freqs, 1, BS, PE, 1.0, True, paddle.bfloat16 + ) + + o_r = o.detach() + o_r.stop_gradient = False + m_frozen = _make_m(u, v) # built from non-leaf constants -> no grad + m_frozen.stop_gradient = True + roped = fused_apply_mla_rope_inplace( + (o_r * 1.0).reshape([1, BS, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + paddle.matmul(m_frozen, roped.reshape([BS, NH, VD])).backward( + d_out.clone() + ) + + o_f = o.detach() + o_f.stop_gradient = False + snapshot = o_f.clone() + InvRopeVhaPostmixFusion.apply( + (o_f * 1.0).reshape([BS, NH, VD]), m_frozen, cos, sin, NOPE, PE + ).backward(d_out.clone()) + + _check_equal(o_f.grad, o_r.grad, "grad_O with frozen postmix") + _check_equal(o_f, snapshot, "O with frozen postmix") + + def test_no_wide_rotated_intermediate(self) -> None: + """Forward must leave one full-width tensor live, not two. + + The unfused pair keeps the rotated attention output alive alongside its + own result, which is exactly the copy this fusion exists to remove. + """ + o, u, v, freqs, d_out = _make_inputs(BS) + cos, sin = build_mla_rope_cos_sin( + freqs, 1, BS, PE, 1.0, True, paddle.bfloat16 + ) + m = _make_m(u, v) + wide = BS * NH * VD * 2 + with paddle.no_grad(): + flat = o.reshape([BS, NH, VD]) + # warm up the JIT and any lazy allocator growth + InvRopeVhaPostmixFusion.apply(flat, m, cos, sin, NOPE, PE) + fused_apply_mla_rope_inplace( + o.reshape([1, BS, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + paddle.device.synchronize() + + before = paddle.device.cuda.memory_allocated() + out = InvRopeVhaPostmixFusion.apply(flat, m, cos, sin, NOPE, PE) + paddle.device.synchronize() + delta = paddle.device.cuda.memory_allocated() - before + del out + + before = paddle.device.cuda.memory_allocated() + roped = fused_apply_mla_rope_inplace( + o.reshape([1, BS, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + out_ref = paddle.matmul(m, roped.reshape([BS, NH, VD])) + paddle.device.synchronize() + delta_ref = paddle.device.cuda.memory_allocated() - before + del roped, out_ref + + # Fused: only the output survives; the compact pe buffers (N/8 each) + # are freed inside the call. + self.assertEqual( + delta, + wide, + f"fused forward left {delta} B live, expected exactly one " + f"full-width output ({wide} B)", + ) + # Unfused: rotated copy + output. + self.assertEqual(delta_ref, 2 * wide) + + +class TestFusionAssumptions(unittest.TestCase): + """Pin the library behaviour the fusion's exactness rests on.""" + + def test_split_n_gemm_is_bitwise(self) -> None: + """matmul(M, X[..., a:b]) == matmul(M, X)[..., a:b]. + + The postmix GEMM contracts the head axis, so the channel axis is a pure + N dimension and splitting it cannot reorder the accumulation. This is the + *only* cuBLAS property the fusion still relies on (the weight gradient + goes through `matmul_grad` precisely because a hand-rolled GEMM there + turned out to be architecture-dependent), so sweep it widely: small and + large head counts, powers of two and not, and both split points. + """ + for bs, nh, vd, pe in ( + (BS, NH, VD, PE), + (PROD_BS, NH, VD, PE), + (1023, NH, VD, PE), + (1, NH, VD, PE), + (128, 4, 64, 32), + (128, 4, 64, 16), + (333, 8, 128, 32), + (97, 32, 96, 32), + (7, 3, 40, 8), + (64, 16, 256, 64), + ): + nope = vd - pe + rank = max(1, nh // 4) + paddle.seed(bs + nh) + m = _make_m( + paddle.randn([nh, rank], "bfloat16"), + paddle.randn([nh, rank], "bfloat16"), + nh, + ) + x = paddle.randn([bs, nh, vd], "bfloat16") + tag = f"({bs},{nh},{vd},pe={pe})" + with paddle.no_grad(): + full = paddle.matmul(m, x) + _check_equal( + paddle.matmul(m, x[..., :nope]), + full[..., :nope], + f"nope split {tag}", + ) + _check_equal( + paddle.matmul(m, x[..., nope:]), + full[..., nope:], + f"pe split {tag}", + ) + # the pe operand is a compact buffer in the fused path, not a view + _check_equal( + paddle.matmul(m, x[..., nope:].contiguous()), + full[..., nope:], + f"pe split from a compact operand {tag}", + ) + + def test_backward_formulas_match_matmul_grad(self) -> None: + """`matmul_grad` and the explicit dgrad == autograd's own results. + + The fused backward takes both gradients from `matmul_grad` and computes + the frozen-postmix dgrad explicitly, so pin both against autograd. + """ + paddle.seed(0) + u = paddle.randn([NH, RANK], "bfloat16") + v = paddle.randn([NH, RANK], "bfloat16") + x0 = paddle.randn([BS, NH, VD], "bfloat16") + d_out = paddle.randn([BS, NH, VD], "bfloat16") + + m_l = _make_m(u, v).detach() + m_l.stop_gradient = False + x_l = x0.detach() + x_l.stop_gradient = False + paddle.matmul(m_l, x_l).backward(d_out.clone()) + + m_d, x_d = m_l.detach(), x0.detach() + with paddle.no_grad(): + _check_equal( + paddle.matmul(m_d, d_out, transpose_x=True), + x_l.grad, + "explicit dgrad", + ) + d_m, d_x = paddle._C_ops.matmul_grad(m_d, x_d, d_out, False, False) + _check_equal(d_m, m_l.grad, "matmul_grad wgrad") + _check_equal(d_x, x_l.grad, "matmul_grad dgrad") + _check_equal( + paddle.einsum("bhd,bkd->hk", d_out, x_d), + m_l.grad, + "einsum wgrad", + ) + + def test_pe_scatter_matches_paddle(self) -> None: + """The vectorised pe scatter == Paddle's strided slice assignment.""" + paddle.seed(0) + t = paddle.randn([BS, NH, VD], "bfloat16") + compact = paddle.randn([BS, NH, PE], "bfloat16") + with paddle.no_grad(): + a, b = t.clone(), t.clone() + scatter_pe_slice_(a, compact, NOPE, PE) + b[..., NOPE:] = compact + _check_equal(a, b, "pe scatter") + + def test_compact_rope_matches_wide_rope(self) -> None: + """rope_pe_to_compact == the standalone op's pe channels.""" + paddle.seed(0) + t = paddle.randn([BS, NH, VD], "bfloat16") + freqs = paddle.randn([1, BS, 1, PE], "float32") + freqs.stop_gradient = True + with paddle.no_grad(): + for inverse in (True, False): + cos, sin = build_mla_rope_cos_sin( + freqs, 1, BS, PE, 1.0, inverse, paddle.bfloat16 + ) + ref = fused_apply_mla_rope_inplace( + t.reshape([1, BS, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=inverse, + clone_input=True, + ) + _check_equal( + rope_pe_to_compact(t, cos, sin, NOPE, PE), + ref.reshape([BS, NH, VD])[..., NOPE:].contiguous(), + f"compact rope (inverse={inverse})", + ) + + +class _GateStub: + """Minimal stand-in for the attention layer's gating attributes.""" + + def __init__(self, **kw): + self.config = type("C", (), {})() + self.config.fuse_inv_rope_into_vha_postmix = True + self.config.apply_rope_fusion = True + self.config.high_precision_rope = False + self.use_vha_postmix = True + self.vha_postmix_grouped = False + self.recompute_vha_postmix = False + self.training = True + for k, val in kw.items(): + if hasattr(self.config, k): + setattr(self.config, k, val) + else: + setattr(self, k, val) + + +class TestFusionGating(unittest.TestCase): + """Every rejected combination must fall back, not raise.""" + + def _gate(self, in_full_recompute=False, **kw) -> bool: + return DSv4HybridSelfAttention._can_fuse_inv_rope_postmix( + _GateStub(**kw), in_full_recompute + ) + + def test_enabled_by_default_config(self) -> None: + self.assertTrue(self._gate()) + + def test_disabled_cases(self) -> None: + for kw in ( + {"fuse_inv_rope_into_vha_postmix": False}, + {"use_vha_postmix": False}, + {"vha_postmix_grouped": True}, + {"apply_rope_fusion": False}, + {"high_precision_rope": True}, + {"recompute_vha_postmix": True}, + ): + self.assertFalse(self._gate(**kw), f"should be disabled for {kw}") + + def test_nested_recompute_allowed_inside_full_recompute(self) -> None: + self.assertTrue( + self._gate(in_full_recompute=True, recompute_vha_postmix=True) + ) + + def test_eval_mode_ignores_postmix_recompute(self) -> None: + self.assertTrue(self._gate(recompute_vha_postmix=True, training=False)) + + +class _LayerStub: + """Just enough of the attention layer to drive the two postmix methods.""" + + def __init__(self, u, v, nh=NH, vd=VD): + self.vha_postmix_U = u + self.vha_postmix_V = v + self.num_attention_heads = nh + self.v_head_dim = vd + self.vha_postmix_grouped = False + self.o_local_groups = 8 + + +class TestWiring(unittest.TestCase): + """`_apply_inv_rope_vha_postmix` must match RoPE + `_apply_vha_postmix`.""" + + def test_method_matches_unfused_pair(self) -> None: + b, sq = 1, BS + o, u, v, freqs, d_out = _make_inputs(b * sq) + + o_r, u_r, v_r = _leaves(o, u, v) + roped = fused_apply_mla_rope_inplace( + (o_r * 1.0).reshape([b, sq, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + ref = DSv4HybridSelfAttention._apply_vha_postmix( + _LayerStub(u_r, v_r), roped + ) + ref.backward(d_out.reshape([b, sq, NH * VD]).clone()) + + o_f, u_f, v_f = _leaves(o, u, v) + got = DSv4HybridSelfAttention._apply_inv_rope_vha_postmix( + _LayerStub(u_f, v_f), + (o_f * 1.0).reshape([b, sq, NH, VD]), + freqs, + NOPE, + PE, + 1.0, + ) + self.assertEqual(list(got.shape), [b, sq, NH * VD]) + self.assertEqual(list(got.shape), list(ref.shape)) + got.backward(d_out.reshape([b, sq, NH * VD]).clone()) + + _check_equal(got.detach(), ref.detach(), "wired forward") + _check_equal(o_f.grad, o_r.grad, "wired grad_O") + _check_equal(u_f.grad, u_r.grad, "wired grad_U") + _check_equal(v_f.grad, v_r.grad, "wired grad_V") + + +class TestSmallHeadCounts(unittest.TestCase): + """End-to-end bitwise equality at small head counts. + + The production config has nh=64, but small head counts are where cuBLAS + algorithm selection gets unstable: an earlier revision hand-rolled the + weight-gradient GEMM over head-major operands and CI caught it diverging at + ``(bs=128, nh=4, d=64)`` on sm90 while it was bitwise equal on sm10.3. The + weight gradient now goes through ``matmul_grad`` on every path, so this must + hold on every architecture -- powers of two and not. + """ + + def _run(self, fused, bs, nh, vd, pe, seed): + nope = vd - pe + rank = max(1, nh // 4) + paddle.seed(seed) + o = paddle.randn([bs, nh, vd], "bfloat16") + u = (paddle.randn([nh, rank], "float32") * 0.05).astype("bfloat16") + v = (paddle.randn([nh, rank], "float32") * 0.05).astype("bfloat16") + freqs = paddle.randn([1, bs, 1, pe], "float32") + freqs.stop_gradient = True + d_out = paddle.randn([bs, nh, vd], "bfloat16") + + o_l, u_l, v_l = _leaves(o, u, v) + m = paddle.matmul(v_l, u_l, transpose_y=True) + paddle.eye( + nh, dtype="bfloat16" + ) + x = o_l * 1.0 + if fused: + cos, sin = build_mla_rope_cos_sin( + freqs, 1, bs, pe, 1.0, True, paddle.bfloat16 + ) + out = InvRopeVhaPostmixFusion.apply( + x.reshape([bs, nh, vd]), m, cos, sin, nope, pe + ) + else: + roped = fused_apply_mla_rope_inplace( + x.reshape([1, bs, nh, vd]), + freqs, + nope, + 1.0, + inverse=True, + clone_input=True, + ) + out = paddle.matmul(m, roped.reshape([bs, nh, vd])) + snapshot = o_l.clone() + out.backward(d_out.clone()) + return out.detach(), o_l.grad, u_l.grad, v_l.grad, o_l, snapshot + + def test_bitwise(self) -> None: + for bs, nh, vd, pe, seed in ( + (128, 4, 64, 32, 0), # the shape CI flagged on sm90 + (128, 4, 64, 16, 1), + (7, 3, 40, 8, 0), # not a power of two + (64, 3, 40, 8, 1), + (1023, 3, 40, 8, 2), + (333, 8, 128, 32, 0), + (97, 32, 96, 32, 0), + (64, 16, 256, 64, 0), + ): + tag = f"(bs={bs}, nh={nh}, d={vd}, pe={pe})" + ref = self._run(False, bs, nh, vd, pe, seed) + got = self._run(True, bs, nh, vd, pe, seed) + for name, a, b in zip( + ("out", "grad_O", "grad_U", "grad_V"), got, ref + ): + _check_equal(a, b, f"{name} {tag}") + _check_equal(got[4], got[5], f"O untouched {tag}") + + +class TestArgumentValidation(unittest.TestCase): + """Every helper rejects malformed input instead of reading out of bounds.""" + + def setUp(self) -> None: + paddle.seed(0) + self.t = paddle.randn([32, NH, VD], "bfloat16") + self.freqs = paddle.randn([1, 32, 1, PE], "float32") + self.freqs.stop_gradient = True + self.cos, self.sin = build_mla_rope_cos_sin( + self.freqs, 1, 32, PE, 1.0, True, paddle.bfloat16 + ) + + def test_rank_must_be_three(self) -> None: + with self.assertRaisesRegex(ValueError, r"\[B\*S, H, D\]"): + rope_pe_to_compact( + self.t.reshape([1, 32, NH, VD]), self.cos, self.sin, NOPE, PE + ) + + def test_channel_split_must_add_up(self) -> None: + with self.assertRaisesRegex(ValueError, r"nope_dim \+ pe_dim"): + rope_pe_to_compact(self.t, self.cos, self.sin, NOPE + 1, PE) + + def test_cos_sin_must_be_contiguous(self) -> None: + with self.assertRaisesRegex(ValueError, "cos/sin must be contiguous"): + rope_pe_to_compact(self.t, self.cos[..., :-4], self.sin, NOPE, PE) + + def test_cos_sin_width_checked(self) -> None: + narrow = self.cos[..., :-4].contiguous() + with self.assertRaisesRegex(ValueError, "cos/sin last dim"): + rope_pe_to_compact(self.t, narrow, self.sin, NOPE, PE) + + def test_freqs_rank_checked(self) -> None: + with self.assertRaisesRegex(ValueError, r"freqs must be \[B,S,1,D\]"): + build_mla_rope_cos_sin( + self.freqs[0], 1, 32, PE, 1.0, True, paddle.bfloat16 + ) + + def test_freqs_seqlen_checked(self) -> None: + with self.assertRaisesRegex(ValueError, "mismatches"): + build_mla_rope_cos_sin( + self.freqs, 1, 33, PE, 1.0, True, paddle.bfloat16 + ) + + def test_freqs_batch_checked(self) -> None: + with self.assertRaisesRegex(ValueError, "must be 1 or"): + build_mla_rope_cos_sin( + paddle.randn([3, 32, 1, PE], "float32"), + 2, + 32, + PE, + 1.0, + True, + paddle.bfloat16, + ) + + def test_freqs_broadcast_over_batch(self) -> None: + """B=1 freqs are broadcast, matching a manually tiled version.""" + cos_b, sin_b = build_mla_rope_cos_sin( + self.freqs, 2, 32, PE, 1.0, True, paddle.bfloat16 + ) + tiled = self.freqs.broadcast_to([2, 32, 1, PE]) + cos_t, sin_t = build_mla_rope_cos_sin( + tiled, 2, 32, PE, 1.0, True, paddle.bfloat16 + ) + _check_equal(cos_b, cos_t, "broadcast cos") + _check_equal(sin_b, sin_t, "broadcast sin") + + def test_scatter_shape_checked(self) -> None: + with self.assertRaisesRegex(ValueError, "compact must be"): + scatter_pe_slice_( + self.t.clone(), + paddle.randn([32, NH, PE + 4], "bfloat16"), + NOPE, + PE, + ) + + +class TestOutOfPlaceRopeHelper(unittest.TestCase): + """``rope_full_out_of_place`` rebuilds the operand the wgrad needs.""" + + def test_matches_standalone_op(self) -> None: + for bs, nh, vd, pe in ( + (BS, NH, VD, PE), + (97, 8, 128, 32), + (7, 3, 40, 8), + ): + nope = vd - pe + paddle.seed(bs) + t = paddle.randn([bs, nh, vd], "bfloat16") + before = t.clone() + freqs = paddle.randn([1, bs, 1, pe], "float32") + freqs.stop_gradient = True + with paddle.no_grad(): + for inverse in (True, False): + cos, sin = build_mla_rope_cos_sin( + freqs, 1, bs, pe, 1.0, inverse, paddle.bfloat16 + ) + ref = fused_apply_mla_rope_inplace( + t.reshape([1, bs, nh, vd]), + freqs, + nope, + 1.0, + inverse=inverse, + clone_input=True, + ).reshape([bs, nh, vd]) + _check_equal( + rope_full_out_of_place(t, cos, sin, nope, pe), + ref, + f"rope_full_out_of_place ({bs},{nh},{vd},{pe}," + f"inverse={inverse})", + ) + _check_equal(t, before, "input unchanged by rope_full_out_of_place") + + +class TestEntryPoint(unittest.TestCase): + """``fused_inv_rope_vha_postmix`` is what the attention layer calls.""" + + def test_matches_unfused_pair(self) -> None: + sq = 64 + o, u, v, freqs, d_out = _make_inputs(sq) + g = d_out.reshape([1, sq, NH * VD]) + + o_r, u_r, v_r = _leaves(o, u, v) + roped = fused_apply_mla_rope_inplace( + (o_r * 1.0).reshape([1, sq, NH, VD]), + freqs, + NOPE, + 1.0, + inverse=True, + clone_input=True, + ) + ref = paddle.matmul( + _make_m(u_r, v_r), roped.reshape([sq, NH, VD]) + ).reshape([1, sq, NH * VD]) + ref.backward(g.clone()) + + o_f, u_f, v_f = _leaves(o, u, v) + got = fused_inv_rope_vha_postmix( + (o_f * 1.0).reshape([1, sq, NH, VD]), freqs, u_f, v_f, NOPE, PE + ) + self.assertEqual(list(got.shape), [1, sq, NH * VD]) + got.backward(g.clone()) + + _check_equal(got.detach(), ref.detach(), "entry point forward") + _check_equal(o_f.grad, o_r.grad, "entry point grad_O") + _check_equal(u_f.grad, u_r.grad, "entry point grad_U") + _check_equal(v_f.grad, v_r.grad, "entry point grad_V") + + def test_accepts_non_contiguous_input(self) -> None: + """A sliced attention output must be handled, not silently mis-read.""" + sq = 64 + o, u, v, freqs, _ = _make_inputs(sq) + wide = paddle.concat([o, o], axis=-1) # [sq, NH, 2 * VD] + view = wide[..., :VD].reshape([1, sq, NH, VD]) + self.assertFalse(view.is_contiguous()) + with paddle.no_grad(): + got = fused_inv_rope_vha_postmix(view, freqs, u, v, NOPE, PE) + want = fused_inv_rope_vha_postmix( + o.reshape([1, sq, NH, VD]), freqs, u, v, NOPE, PE + ) + _check_equal(got, want, "non-contiguous input") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/test_local_to_global_idxs_fusion.py b/tests/single_card_tests/test_local_to_global_idxs_fusion.py new file mode 100644 index 0000000000..1bb3c1554d --- /dev/null +++ b/tests/single_card_tests/test_local_to_global_idxs_fusion.py @@ -0,0 +1,597 @@ +# 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. + +""" +Unit tests for ``paddlefleet.triton_ops.local_to_global_idxs_fusion``. + +The eager ``_local_to_global_flat`` is the ground truth; the fused kernel must +match it *bit-for-bit* (these are index tables -- there is no tolerance). +""" + +import unittest + +import numpy as np +import paddle + +from paddlefleet.fusions.csa_sparse_attn_utils import ( + _local_to_global_flat, + local_to_global_flat, +) +from paddlefleet.triton_ops.local_to_global_idxs_fusion import ( + local_to_global_flat_triton, +) + + +def _require_sm100_sparse_kernels(testcase): + """Skip unless the FlashMLA / DSA sparse-attention kernels can actually run. + + Importing ``paddlefleet_ops.flash_mla`` is not a capability check: the + library loads on Hopper (the CI runner is SM 9.0 and the log shows + "Successfully loaded ecosystem library: flash_mla") while the kernels these + tests reach are SM100-only -- the backward the docstrings quote lives in + ``sparse_attention_backward/dsa_bwd_sm100.py``. Calling them on SM 9.x is + not a meaningful test, so gate on the capability the way + ``test_hysparse_online_tilelang_train_step`` does for the TileLang kernels. + """ + if not paddle.device.is_compiled_with_cuda(): + testcase.skipTest("CUDA build of Paddle is required") + if paddle.device.cuda.device_count() == 0: + testcase.skipTest("No CUDA device available") + major = paddle.device.cuda.get_device_capability()[0] + if major < 10: + testcase.skipTest( + f"FlashMLA / DSA sparse kernels require SM 10.x; got SM {major}.x" + ) + try: + from paddlefleet_ops.flash_mla import ( # noqa: F401 + flash_mla_sparse_fwd, + ) + except (ImportError, RuntimeError): + testcase.skipTest("flash_mla is not available") + + +def _assert_bitwise(case, ref, got): + assert ref.dtype == got.dtype, f"{case}: dtype {ref.dtype} vs {got.dtype}" + assert ref.shape == got.shape, f"{case}: shape {ref.shape} vs {got.shape}" + r, g = ref.numpy(), got.numpy() + if not np.array_equal(r, g): + bad = np.argwhere(r != g) + p = tuple(bad[0]) + raise AssertionError( + f"{case}: {len(bad)}/{r.size} mismatches, " + f"first at {p}: ref={r[p]} got={g[p]}" + ) + + +def _make_idxs(b, sq, topk, seqlen_kv, dtype, rng, pad_ratio=0.3): + """Random indices in [0, seqlen_kv), ``pad_ratio`` of them set to -1. + + The pad mask is drawn as int16 rather than float64: these tables reach tens + of millions of entries and the mask would otherwise be the single largest + host allocation in the file. + """ + shape = (b, sq, topk) + idxs = rng.integers(0, max(seqlen_kv, 1), size=shape).astype(np.int64) + pad = rng.integers(0, 1000, size=shape, dtype=np.int16) + idxs[pad < int(pad_ratio * 1000)] = -1 + return paddle.to_tensor(idxs, dtype=dtype) + + +def _low32(x): + """Two's-complement truncation to int32, not relying on numpy cast rules.""" + y = np.asarray(x, dtype=np.int64) + return ((y + 2**31) % 2**32 - 2**31).astype(np.int32) + + +def _numpy_oracle(idxs, seqlen_kv): + """Independent reference derived from the spec, not from the Paddle code. + + ``out[b * sq + s, k] = idx + b * seqlen_kv`` when ``idx >= 0``, else ``idx`` + unchanged, truncated to int32. + + One oracle covers both input dtypes: the eager int32 path (wrapping ``add`` + then a no-op ``cast``) and the eager int64 path (exact ``add`` then a + truncating ``cast``) both reduce to the low 32 bits of the exact sum, + because low32(a + low32(p)) == low32(a + p). + """ + np_idxs = idxs.numpy().astype(np.int64) + b, sq, topk = np_idxs.shape + flat = np_idxs.reshape(b * sq, topk) + off = np.repeat(np.arange(b, dtype=np.int64) * int(seqlen_kv), sq)[:, None] + return _low32(np.where(flat >= 0, flat + off, flat)) + + +class TestLocalToGlobalFlatFusion(unittest.TestCase): + def setUp(self): + self.rng = np.random.default_rng(0) + + def _check(self, case, idxs, seqlen_kv): + ref = _local_to_global_flat(idxs, seqlen_kv) + got = local_to_global_flat_triton(idxs, seqlen_kv) + _assert_bitwise(case, ref, got) + return ref + + def test_smallest_possible_launch(self): + """One row, one column: the cheapest possible kernel launch. + + Deliberately first in the file so the CI log distinguishes "the very + first Triton compile stalls" (no progress at all) from "the volume of a + later test is the problem" (this one's dot appears, then the stall). + """ + idxs = paddle.to_tensor([[[3]]], dtype="int32") + got = local_to_global_flat_triton(idxs, 8) + self.assertEqual(list(got.shape), [1, 1]) + self.assertEqual(int(got[0, 0]), 3) + + def test_ernielite_hca_shape(self): + """Real ernielite HCA widths: window+compress=640, seqlen_kv=65536+512. + + ``sq`` is only a grid multiplier -- the kernel is one program per row and + ``cdiv(topk, BLOCK_K)`` along topk -- so it is kept small here. The real + sq_local=16384 costs 10.5M entries per dtype and, with the eager + reference plus the host round-trip in the comparison, dominated this + file's runtime. + """ + for dtype in ("int32", "int64"): + idxs = _make_idxs(1, 1024, 640, 65536 + 512, dtype, self.rng) + self._check(f"hca[{dtype}]", idxs, 65536 + 512) + + def test_shapes_and_dtypes(self): + cases = [ + (1, 1, 1), + (1, 7, 5), + (1, 16, 640), + (2, 4, 6), + (3, 5, 1), + (4, 33, 64), + (2, 8, 1025), # topk > BLOCK_K -> multi-block along topk + (5, 3, 2047), # non power-of-2, multi-block + ] + for dtype in ("int32", "int64"): + for b, sq, topk in cases: + seqlen_kv = max(topk, 8) * 2 + 3 + idxs = _make_idxs(b, sq, topk, seqlen_kv, dtype, self.rng) + case = f"[{dtype}] b={b},sq={sq},topk={topk}" + self._check(case, idxs, seqlen_kv) + + def test_edge_values(self): + """All-padding rows, all-valid rows, index 0, and non -1 negatives.""" + for dtype in ("int32", "int64"): + all_pad = paddle.full([3, 4, 5], -1, dtype=dtype) + self._check(f"all_pad[{dtype}]", all_pad, 32) + + all_zero = paddle.zeros([3, 4, 5], dtype=dtype) + self._check(f"all_zero[{dtype}]", all_zero, 32) + + # the reference passes negatives through unchanged, it does not + # normalise them to -1; make sure the fusion does the same. + other_neg = paddle.to_tensor( + [[[-5, 0, -1, 7], [-2, 3, -9, 0]]], dtype=dtype + ) + ref = self._check(f"other_neg[{dtype}]", other_neg, 16) + np.testing.assert_array_equal( + ref.numpy(), np.array([[-5, 0, -1, 7], [-2, 3, -9, 0]]) + ) + + def test_non_contiguous_input(self): + """CSA slices the index table before this call; keep that path exact.""" + for dtype in ("int32", "int64"): + full = _make_idxs(2, 16, 128, 256, dtype, self.rng) + sliced = full[:, 4:12, 8:72] + self.assertEqual(sliced.shape, [2, 8, 64]) + self._check(f"sliced[{dtype}]", sliced, 256) + + def test_large_batch_offset_int32_wraparound(self): + """b * seqlen_kv beyond int32 must wrap the same way in both paths.""" + seqlen_kv = 2**30 + rows = [[[0, 5, -1]], [[1, -1, 9]], [[3, 4, 0]]] + idxs = paddle.to_tensor(rows, dtype="int32") + self._check("int32_wrap", idxs, seqlen_kv) + + idxs64 = paddle.to_tensor(rows, dtype="int64") + self._check("int64_no_wrap", idxs64, seqlen_kv) + + def test_allow_alias_matches_when_batch_is_one(self): + for dtype in ("int32", "int64"): + idxs = _make_idxs(1, 64, 96, 512, dtype, self.rng) + ref = _local_to_global_flat(idxs, 512) + got = local_to_global_flat_triton(idxs, 512, allow_alias=True) + _assert_bitwise(f"alias[{dtype}]", ref, got) + + def test_not_differentiable(self): + """Integer index tables carry no gradient; fusion must not add one.""" + idxs = _make_idxs(2, 4, 8, 32, "int32", self.rng) + self.assertTrue(idxs.stop_gradient) + out = local_to_global_flat_triton(idxs, 32) + self.assertTrue(out.stop_gradient) + + def test_dispatcher_matches_eager(self): + """``local_to_global_flat(fused=...)`` must agree with the reference.""" + for dtype in ("int32", "int64"): + idxs = _make_idxs(3, 32, 48, 256, dtype, self.rng) + ref = _local_to_global_flat(idxs, 256) + for fused in (False, True): + got = local_to_global_flat(idxs, 256, fused=fused) + _assert_bitwise(f"dispatch[{dtype},fused={fused}]", ref, got) + + +class TestConfigSwitch(unittest.TestCase): + """``sparse_attn_global_kv_idx_remap_fusion`` must be inert numerically.""" + + def test_config_field_defaults_off(self): + from paddlefleet.transformer.transformer_config import TransformerConfig + + field = TransformerConfig.__dataclass_fields__[ + "sparse_attn_global_kv_idx_remap_fusion" + ] + self.assertIs(field.default, False) + + def test_cudnn_sparse_attn_end_to_end(self): + """Flipping the switch must not move out / dq / d_sink at all. + + ``dkv`` is deliberately excluded: ``reduce_dKV_from_reg`` in + ``paddlefleet_ops/.../sparse_attention_backward/dsa_bwd_sm100.py`` + accumulates it with ``cute.arch.atomic_add`` scattered by the top-k + index, so it is not bitwise reproducible even eager-vs-eager whenever a + KV column has more than one writer (measured: flakes on 19/19 repeats, + up to ~100 bf16 elements, 1 ULP). The switch adds no deviation beyond + that pre-existing nondeterminism. + + ``d_sink`` IS checked, but only because ``sq <= dSink_block_q``: the + separate ``sum_dSink`` kernel launches ``ceil_div(seqlen_q, 256)`` + q-blocks and each one does its own ``atomic_add``, so a single block + has no contention and is exactly reproducible. Raising ``sq`` above 256 + makes ``d_sink`` flake too (measured: 19/19 repeats at sq=2048, + max|diff| ~1e-6) -- keep the assert honest by keeping sq <= 256. + + ``dq`` needs no such caveat: ``store_dQ`` writes it with a TMA store + from the single CTA that owns each query row. + """ + _require_sm100_sparse_kernels(self) + + from paddlefleet.fusions.csa_sparse_attn import csa_sparse_attn + + b, sq, h, d = 1, 256, 64, 512 # FlashMLA sparse fixes h_q=64, d_v=512 + assert sq <= 256, "sq must stay within one sum_dSink block (see above)" + s_kv, topk = 384, 128 + rng = np.random.default_rng(0) + idx = rng.integers(0, s_kv, (b, sq, topk)).astype(np.int32) + idx[rng.random((b, sq, topk)) < 0.3] = -1 + q_np = rng.standard_normal((b, sq, h, d)).astype(np.float32) + kv_np = rng.standard_normal((b, s_kv, d)).astype(np.float32) + sink_np = rng.standard_normal((h,)).astype(np.float32) + go_np = rng.standard_normal((b, sq, h * d)).astype(np.float32) + + def run(fused): + q = paddle.to_tensor(q_np, dtype="bfloat16") + kv = paddle.to_tensor(kv_np, dtype="bfloat16") + sink = paddle.to_tensor(sink_np, dtype="float32") + for t in (q, kv, sink): + t.stop_gradient = False + out = csa_sparse_attn( + q, + kv, + sink, + paddle.to_tensor(idx, dtype="int32"), + d**-0.5, + backend="cudnn", + global_kv_idx_remap_fusion=fused, + ) + out.backward(paddle.to_tensor(go_np, dtype=out.dtype)) + return out, q.grad, sink.grad + + ref, got = run(False), run(True) + for name, a, c in zip(("out", "dq", "d_sink"), ref, got): + _assert_bitwise(f"switch[{name}]", a, c) + + def test_cudnn_sparse_attn_unaligned_topk(self): + """topk not a multiple of the arch alignment exercises the F.pad path. + + ``flash_mla_sparse_attn`` pads ``global_idxs`` up to a multiple of 64 + (SM100) with ``-1`` right after the remap, so the fused output has to + survive that padding untouched. + """ + _require_sm100_sparse_kernels(self) + + from paddlefleet.fusions.csa_sparse_attn import csa_sparse_attn + + b, sq, h, d = 1, 256, 64, 512 + s_kv, topk = 320, 100 # 100 -> padded to 128 + rng = np.random.default_rng(3) + idx = rng.integers(0, s_kv, (b, sq, topk)).astype(np.int32) + idx[rng.random((b, sq, topk)) < 0.3] = -1 + q_np = rng.standard_normal((b, sq, h, d)).astype(np.float32) + kv_np = rng.standard_normal((b, s_kv, d)).astype(np.float32) + sink_np = rng.standard_normal((h,)).astype(np.float32) + go_np = rng.standard_normal((b, sq, h * d)).astype(np.float32) + + def run(fused): + q = paddle.to_tensor(q_np, dtype="bfloat16") + kv = paddle.to_tensor(kv_np, dtype="bfloat16") + sink = paddle.to_tensor(sink_np, dtype="float32") + for t in (q, kv, sink): + t.stop_gradient = False + out = csa_sparse_attn( + q, + kv, + sink, + paddle.to_tensor(idx, dtype="int32"), + d**-0.5, + backend="cudnn", + global_kv_idx_remap_fusion=fused, + ) + out.backward(paddle.to_tensor(go_np, dtype=out.dtype)) + return out, q.grad, sink.grad + + ref, got = run(False), run(True) + for name, a, c in zip(("out", "dq", "d_sink"), ref, got): + _assert_bitwise(f"unaligned[{name}]", a, c) + + def test_mqa_sparse_attn_end_to_end(self): + """The switch is also wired into the absorbed-MQA path; pin that too. + + Swept over ``mqa_sparse_attn_backward_backend``, because the remap has a + different reach on each branch: the FlashMLA forward always builds the + flat-global table, while only the ``"cudnn"`` backward consumes one -- + the ``"tilelang"`` backward indexes ``token_indices`` per batch itself. + + ``dkv`` is only checked on ``"tilelang"``. On ``"cudnn"`` it comes from + the atomic epilogue and is not reproducible even against itself, which + is the same exclusion the CSA case makes; the deterministic kernel is + what lets this test pin the one output that branch has to leave out. + ``s`` stays within one ``sum_dSink`` block so ``d_sink`` is exact. + """ + _require_sm100_sparse_kernels(self) + + from paddlefleet.fusions.mqa_sparse_attn import mqa_sparse_attn + + b, s, h, d_qk, d_v = 1, 128, 64, 576, 512 # absorbed MQA layout + s_kv, width = 256, 64 + assert s <= 256, "keep s within one sum_dSink block" + rng = np.random.default_rng(11) + tok = rng.integers(0, s_kv, (b, s, width)).astype(np.int32) + tok[rng.random((b, s, width)) < 0.3] = -1 + q_np = rng.standard_normal((b, s, h, d_qk)).astype(np.float32) + kv_np = rng.standard_normal((b, s_kv, d_qk)).astype(np.float32) + sink_np = rng.standard_normal((h,)).astype(np.float32) + go_np = rng.standard_normal((b, s, h * d_v)).astype(np.float32) + + def run(fused, backend): + q = paddle.to_tensor(q_np, dtype="bfloat16") + kv = paddle.to_tensor(kv_np, dtype="bfloat16") + sink = paddle.to_tensor(sink_np, dtype="float32") + for t in (q, kv, sink): + t.stop_gradient = False + out = mqa_sparse_attn( + q, + kv, + paddle.to_tensor(tok, dtype="int32"), + d_qk**-0.5, + d_v, + attn_sink=sink, + global_kv_idx_remap_fusion=fused, + backward_backend=backend, + ) + out.backward(paddle.to_tensor(go_np, dtype=out.dtype)) + return { + "out": out, + "dq": q.grad, + "dkv": kv.grad, + "d_sink": sink.grad, + } + + for backend in ("cudnn", "tilelang"): + names = ("out", "dq", "d_sink") + if backend == "tilelang": + names += ("dkv",) + with self.subTest(backward_backend=backend): + ref, got = run(False, backend), run(True, backend) + for name in names: + _assert_bitwise( + f"mqa[{backend}][{name}]", ref[name], got[name] + ) + + +class TestBitwiseAgainstOracle(unittest.TestCase): + """eager == independent numpy oracle == fused, in every scenario. + + Comparing the fusion only against the eager code would pass even if both + were wrong, so every case is pinned to ``_numpy_oracle`` as well. + """ + + def setUp(self): + self.rng = np.random.default_rng(1234) + + def _check3(self, case, idxs, seqlen_kv): + oracle = _numpy_oracle(idxs, seqlen_kv) + ref = _local_to_global_flat(idxs, seqlen_kv) + got = local_to_global_flat_triton(idxs, seqlen_kv) + self.assertEqual(ref.dtype, paddle.int32, case) + self.assertEqual(got.dtype, paddle.int32, case) + np.testing.assert_array_equal( + ref.numpy(), oracle, err_msg=f"{case}: eager vs oracle" + ) + np.testing.assert_array_equal( + got.numpy(), oracle, err_msg=f"{case}: fused vs oracle" + ) + + def test_topk_block_boundaries(self): + """BLOCK_K = min(next_pow2(topk), 1024); probe both sides of edges.""" + for topk in ( + 1, + 2, + 3, + 7, + 31, + 32, + 33, + 63, + 64, + 65, + 127, + 128, + 129, + 255, + 256, + 511, + 512, + 513, + 640, + 1023, + 1024, + 1025, + 1536, + 2047, + 2048, + 3072, + 4097, + ): + for dtype in ("int32", "int64"): + skv = topk * 2 + 5 + idxs = _make_idxs(2, 3, topk, skv, dtype, self.rng) + self._check3(f"topk={topk} [{dtype}]", idxs, skv) + + def test_batch_and_seq_shapes(self): + for b in (1, 2, 3, 7, 8, 16, 33): + for sq in (1, 2, 7, 64): + for dtype in ("int32", "int64"): + idxs = _make_idxs(b, sq, 96, 512, dtype, self.rng) + self._check3(f"b={b} sq={sq} [{dtype}]", idxs, 512) + + def test_pad_ratios(self): + """No padding, partial padding, all padding.""" + for ratio in (0.0, 0.25, 0.5, 0.75, 1.0): + for dtype in ("int32", "int64"): + idxs = _make_idxs( + 4, 16, 130, 1024, dtype, self.rng, pad_ratio=ratio + ) + self._check3(f"pad={ratio} [{dtype}]", idxs, 1024) + + def test_value_boundaries(self): + """0, seqlen_kv-1, -1 and other negatives, hand-written per row.""" + skv = 97 + rows = [ + [0, skv - 1, -1, 1, skv - 2], + [-1, -1, -1, -1, -1], + [0, 0, 0, 0, 0], + [-5, -2, -9, -128, -1], + [skv - 1, skv - 1, 0, -1, 3], + ] + for dtype in ("int32", "int64"): + idxs = paddle.to_tensor([rows, rows, rows], dtype=dtype) + self._check3(f"boundaries [{dtype}]", idxs, skv) + + def test_int32_overflow_matrix(self): + """b * seqlen_kv near and beyond 2**31 must wrap identically.""" + rows = [[[0, 5, -1, 17]], [[1, -1, 9, 0]], [[3, 4, 0, -7]]] + for skv in (2**20, 2**28, 2**30, 2**31 - 1, 2**31, 3 * 10**8): + for dtype in ("int32", "int64"): + idxs = paddle.to_tensor(rows, dtype=dtype) + self._check3(f"skv={skv} [{dtype}]", idxs, skv) + + def test_ernielite_real_shapes(self): + """The two index-table widths this actually runs on in ernielite.""" + # sq trimmed for the same reason as test_ernielite_hca_shape; the two + # topk widths and the seqlen_kv that sets the int32 offset range -- the + # things this test is about -- are the real ones. + cases = [ + (1, 1024, 128 + 512, 65536 + 512), # HCA: window + compressed + (1, 1024, 128 + 2048, 65536 + 512), # DSA: window + index_topk + ] + for b, sq, topk, skv in cases: + for dtype in ("int32", "int64"): + idxs = _make_idxs(b, sq, topk, skv, dtype, self.rng) + self._check3(f"real b={b} topk={topk} [{dtype}]", idxs, skv) + + def test_non_contiguous_and_strided(self): + """Sliced / stepped views must give the same answer as a dense copy.""" + for dtype in ("int32", "int64"): + full = _make_idxs(4, 32, 256, 1024, dtype, self.rng) + for name, view in ( + ("slice", full[:, 4:20, 8:200]), + ("step2", full[:, ::2, ::2]), + ("tail", full[1:, -8:, -65:]), + ): + self._check3(f"{name} [{dtype}]", view, 1024) + dense = paddle.to_tensor(view.numpy(), dtype=dtype) + _assert_bitwise( + f"{name} view vs dense [{dtype}]", + local_to_global_flat_triton(dense, 1024), + local_to_global_flat_triton(view, 1024), + ) + + +class TestFusedKernelProperties(unittest.TestCase): + """Properties of the fused kernel itself, independent of the reference.""" + + def setUp(self): + self.rng = np.random.default_rng(7) + + def test_repeatable_across_runs(self): + """No atomics in the kernel, so repeated launches must be identical.""" + idxs = _make_idxs(3, 512, 640, 4096, "int32", self.rng) + base = local_to_global_flat_triton(idxs, 4096).numpy() + for i in range(8): + again = local_to_global_flat_triton(idxs, 4096).numpy() + np.testing.assert_array_equal(base, again, err_msg=f"run {i}") + + def test_output_is_fresh_storage_by_default(self): + """The reference always allocates; the default path must match that.""" + idxs = _make_idxs(1, 32, 64, 256, "int32", self.rng) + out = local_to_global_flat_triton(idxs, 256) + self.assertNotEqual(out.data_ptr(), idxs.data_ptr()) + out[0, 0] = 12345 # writing the result must not touch the input + self.assertNotEqual(int(idxs[0, 0, 0]), 12345) + + def test_alias_shares_storage_only_when_safe(self): + """``allow_alias`` aliases exactly for b == 1 and int32, never else.""" + i32 = _make_idxs(1, 32, 64, 256, "int32", self.rng) + self.assertEqual( + local_to_global_flat_triton(i32, 256, allow_alias=True).data_ptr(), + i32.data_ptr(), + ) + i64 = _make_idxs(1, 32, 64, 256, "int64", self.rng) + self.assertNotEqual( + local_to_global_flat_triton(i64, 256, allow_alias=True).data_ptr(), + i64.data_ptr(), + ) + many = _make_idxs(3, 32, 64, 256, "int32", self.rng) + self.assertNotEqual( + local_to_global_flat_triton(many, 256, allow_alias=True).data_ptr(), + many.data_ptr(), + ) + + def test_rejects_wrong_rank(self): + for shape in ([8, 16], [2, 4, 8, 16]): + with self.assertRaises(AssertionError): + local_to_global_flat_triton( + paddle.zeros(shape, dtype="int32"), 64 + ) + + def test_degenerate_shapes_skip_the_kernel(self): + """An empty row or topk axis must return early, not launch a 0-size grid. + + ``triton.cdiv(0, BLOCK_K)`` is 0, so without the guard the launch is a + no-op on some Triton versions and an error on others. + """ + for b, sq, topk in ((0, 8, 16), (2, 0, 16), (2, 8, 0)): + out = local_to_global_flat_triton( + paddle.zeros([b, sq, topk], dtype="int32"), 64 + ) + self.assertEqual(list(out.shape), [b * sq, topk]) + self.assertEqual(out.dtype, paddle.int32) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/test_mla_rope_inplace_fusion.py b/tests/single_card_tests/test_mla_rope_inplace_fusion.py index 886d8cb5b4..5a5a1d7585 100644 --- a/tests/single_card_tests/test_mla_rope_inplace_fusion.py +++ b/tests/single_card_tests/test_mla_rope_inplace_fusion.py @@ -19,7 +19,10 @@ - Forward output equivalent (bit-exact match) to the slice + rope + concat baseline used in dsv4_hybrid_attention.py. - - Truly in-place: q.data_ptr() preserved. + - Truly in-place when clone_input=False: q.data_ptr() preserved and not a + single byte of extra device memory allocated. + - With clone_input=True: the input is left completely untouched and the + result is bit-identical to the clone_input=False result. - Backward grads match the autograd reference. """ @@ -31,7 +34,10 @@ _apply_rotary_pos_emb_bshd, ) from paddlefleet.triton_ops import fused_apply_mla_rope_inplace -from paddlefleet.triton_ops.mla_rope_inplace_fusion import _fused_cos_sin +from paddlefleet.triton_ops.mla_rope_inplace_fusion import ( + RoPEMLAInplaceFusion, + _fused_cos_sin, +) # Shapes from DeepSeek-V4-Flash B, S, H, D = 1, 4096, 64, 512 @@ -74,6 +80,7 @@ def _run_case( s: int, freqs: paddle.Tensor, inverse: bool = False, + clone_input: bool = False, ) -> None: """Shared driver: q is contiguous bf16; freqs supplied by caller.""" x = paddle.randn([b, s, H, D], "bfloat16") @@ -93,16 +100,24 @@ def _run_case( x_fused.stop_gradient = False q_fused = x_fused.clone() # non-leaf — safe target for in-place kernel self.assertTrue(q_fused.is_contiguous()) - nope_before = q_fused[..., :NOPE_DIM].clone() + input_before = q_fused.clone() ptr_before = q_fused.data_ptr() out_fused = fused_apply_mla_rope_inplace( - q_fused, freqs, NOPE_DIM, inverse=inverse + q_fused, freqs, NOPE_DIM, inverse=inverse, clone_input=clone_input ) - # ---- in-place storage invariants ---- - self.assertIs(out_fused, q_fused) - self.assertEqual(out_fused.data_ptr(), ptr_before) - _check_equal(q_fused[..., :NOPE_DIM], nope_before) + # ---- storage invariants ---- + if clone_input: + # A fresh buffer, and the input must survive completely intact — + # this is what the attention backward relies on. + self.assertIsNot(out_fused, q_fused) + self.assertNotEqual(out_fused.data_ptr(), ptr_before) + _check_equal(q_fused, input_before) + else: + self.assertIs(out_fused, q_fused) + self.assertEqual(out_fused.data_ptr(), ptr_before) + # The nope channels of the result always carry the input's values. + _check_equal(out_fused[..., :NOPE_DIM], input_before[..., :NOPE_DIM]) # ---- forward parity ---- _check_equal(out_fused, out_ref) @@ -120,6 +135,12 @@ def test_forward_backward(self) -> None: freqs.stop_gradient = True self._run_case(B, S, freqs) + def test_forward_backward_clone_input(self) -> None: + """Same as above but with clone_input=True (o inv-rope call site).""" + freqs = paddle.randn([B, S, 1, ROPE_DIM]) + freqs.stop_gradient = True + self._run_case(B, S, freqs, clone_input=True) + def test_freqs_noncontiguous_b_gt_1(self) -> None: """Test multi-batch and non-contiguous freqs.""" b = 2 @@ -140,6 +161,89 @@ def test_inverse(self) -> None: freqs.stop_gradient = True self._run_case(B, S, freqs, inverse=True) + def test_inverse_clone_input(self) -> None: + """Inverse rope with clone_input=True: the production o path.""" + freqs = paddle.randn([B, S, 1, ROPE_DIM]) + freqs.stop_gradient = True + self._run_case(B, S, freqs, inverse=True, clone_input=True) + + def test_clone_input_matches_inplace_bitwise(self) -> None: + """clone_input must not change a single bit of the result. + + Runs both modes on identical inputs and compares outputs and grads + bit-for-bit, so the out-of-place kernel path cannot silently diverge + from the in-place one (e.g. by reordering the bf16 rounding). + """ + s = 256 + freqs = paddle.randn([B, s, 1, ROPE_DIM]) + freqs.stop_gradient = True + x = paddle.randn([B, s, H, D], "bfloat16") + out_grad = paddle.randn([B, s, H, D], "bfloat16") + + results = {} + for clone_input in (False, True): + leaf = x.detach() + leaf.stop_gradient = False + q = leaf.clone() + out = fused_apply_mla_rope_inplace( + q, freqs, NOPE_DIM, inverse=True, clone_input=clone_input + ) + out.backward(out_grad.clone()) + results[clone_input] = (out.clone(), leaf.grad.clone()) + + _check_equal(results[False][0], results[True][0]) + _check_equal(results[False][1], results[True][1]) + + def test_inplace_allocates_nothing(self) -> None: + """clone_input=False must cost zero extra device memory. + + Measured around `RoPEMLAInplaceFusion.apply` rather than the public + wrapper, so the cos/sin buffers `_fused_cos_sin` allocates do not + pollute the reading. clone_input=True is checked in the same way to + confirm it allocates exactly one output tensor and nothing more. + """ + s = 256 + freqs = paddle.randn([B, s, 1, ROPE_DIM]) + freqs.stop_gradient = True + cos, sin = _fused_cos_sin(freqs, 1.0, False, paddle.bfloat16) + nbytes = B * s * H * D * 2 # bf16 + + with paddle.no_grad(): + t = paddle.randn([B, s, H, D], "bfloat16") + # Warm up the JIT / any lazy allocator growth first. + RoPEMLAInplaceFusion.apply(t, cos, sin, NOPE_DIM, ROPE_DIM, False) + RoPEMLAInplaceFusion.apply(t, cos, sin, NOPE_DIM, ROPE_DIM, True) + paddle.device.synchronize() + + before = paddle.device.cuda.memory_allocated() + out_ip = RoPEMLAInplaceFusion.apply( + t, cos, sin, NOPE_DIM, ROPE_DIM, False + ) + paddle.device.synchronize() + delta_inplace = paddle.device.cuda.memory_allocated() - before + self.assertIs(out_ip, t) + + before = paddle.device.cuda.memory_allocated() + out_oop = RoPEMLAInplaceFusion.apply( + t, cos, sin, NOPE_DIM, ROPE_DIM, True + ) + paddle.device.synchronize() + delta_clone = paddle.device.cuda.memory_allocated() - before + + self.assertEqual( + delta_inplace, + 0, + f"clone_input=False allocated {delta_inplace} bytes; the in-place " + "path must not allocate", + ) + self.assertEqual( + delta_clone, + nbytes, + f"clone_input=True allocated {delta_clone} bytes, expected exactly " + f"one output tensor ({nbytes})", + ) + del out_oop + def test_fused_cos_sin(self) -> None: freqs = paddle.randn([B, S, 1, ROPE_DIM]) dtype = paddle.bfloat16 diff --git a/tests/single_card_tests/transformer/test_inv_rope_vha_postmix_layer.py b/tests/single_card_tests/transformer/test_inv_rope_vha_postmix_layer.py new file mode 100644 index 0000000000..1d44b8df03 --- /dev/null +++ b/tests/single_card_tests/transformer/test_inv_rope_vha_postmix_layer.py @@ -0,0 +1,255 @@ +# 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. + +"""Layer-level bitwise check for ``fuse_inv_rope_into_vha_postmix``. + +``test_inv_rope_vha_postmix_fusion`` pins the fused op in isolation. This runs a +real ``DSv4HybridSelfAttention`` end to end with the flag off and on and requires +the layer output, the input gradient and *every* parameter gradient to be +bitwise identical, so nothing about the surrounding layer (grouped o-proj, gate, +CSA backward, RoPE freq construction) can quietly change when the flag flips. + +The postmix is deliberately moved off its identity initialisation (``V`` is +zero-initialised, which would make every postmix gradient exactly zero and the +weight-gradient comparison vacuous). +""" + +import unittest + +import paddle +from paddle.distributed.fleet.meta_parallel import build_spec_layer + +from paddlefleet.models.gpt.gpt_layer_specs import get_attention_spec +from paddlefleet.tensor_parallel.random import model_parallel_cuda_manual_seed +from paddlefleet.training.initialize import initialize_fleet +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.transformer_config import TransformerConfig +from paddlefleet.utils import init_method_normal, scaled_init_method_normal + +initialize_fleet(strategy=paddle.distributed.fleet.DistributedStrategy()) + +_SEED = 42 +HIDDEN = 256 +NH, VD, PE = 8, 32, 16 +NOPE = VD - PE +RANK = 2 + + +def _make_config(**overrides): + kwargs = { + "num_hidden_layers": 4, + "hidden_size": HIDDEN, + "num_attention_heads": NH, + "params_dtype": paddle.bfloat16, + "bf16": True, + "use_bias": False, + "multi_latent_attention": True, + "experimental_attention_variant": "dsv4_hybrid", + "q_lora_rank": 64, + "kv_lora_rank": NOPE, + "qk_nope_head_dim": NOPE, + "qk_rope_head_dim": PE, + "qk_pos_emb_head_dim": PE, + "v_head_dim": VD, + "hybrid_mla_q_lora_rank": 1536, + "hybrid_mla_kv_lora_rank": 512, + "hybrid_mla_qk_nope_head_dim": 192, + "hybrid_mla_qk_rope_head_dim": 64, + "hybrid_mla_v_head_dim": 256, + "hybrid_mla_num_attention_heads": 64, + "hybrid_mla_num_key_value_heads": 64, + "o_groups": 4, + "o_lora_rank": 32, + "rope_type": "rope", + "rotary_base": 10000.0, + "rotary_percent": 1.0, + "normalization": "RMSNorm", + "use_qk_norm": True, + # all-128 => every layer is an HCA layer, which is where the inverse + # RoPE + postmix pair lives + "csa_compress_ratios": [128, 128, 128, 128], + "csa_window_size": 16, + "dsa_index_n_heads": 4, + "dsa_index_head_dim": 32, + "dsa_index_topk": 8, + "dsa_indexer_loss_coeff": 1.0, + "dsa_indexer_use_sparse_loss": False, + "dsa_indexer_rotary_interleaved": False, + "apply_rope_fusion": True, + "attention_dropout": 0.0, + "attention_softmax_in_fp32": True, + "masked_softmax_fusion": False, + "softmax_type": "vanilla", + "csa_indexer_backend": "unfused", + "csa_sparse_attn_backend": "unfused", + "tensor_model_parallel_size": 1, + "context_parallel_size": 1, + "csa_dense_mode": False, + "init_method": init_method_normal(0.02), + "output_layer_init_method": scaled_init_method_normal(0.02, 1, 2.0), + "rms_norm_eps": 1e-5, + "use_vha_attention": True, + "vha_postmix_rank": RANK, + } + kwargs.update(overrides) + config = TransformerConfig(**kwargs) + config.dtype = "bfloat16" + return config + + +def _build(config, layer_number=0): + """Build the layer with bf16 parameters (the fused RoPE path is bf16 only).""" + model_parallel_cuda_manual_seed(_SEED) + prev = paddle.get_default_dtype() + paddle.set_default_dtype("bfloat16") + try: + spec = get_attention_spec( + config=config, + attention_layer_type="dsv4_hybrid_attention", + attn_mask_type=AttnMaskType.causal, + ) + return build_spec_layer(spec, config=config, layer_number=layer_number) + finally: + paddle.set_default_dtype(prev) + + +def _check_equal(a, b, what): + a32, b32 = a.astype("float32"), b.astype("float32") + if bool(paddle.all(a32 == b32)): + return + diff = (a32 - b32).abs() + raise AssertionError( + f"{what} not bitwise equal: " + f"{int((a32 != b32).sum())}/{a.numel().item()} elements differ, " + f"max|diff|={float(diff.max()):.6e}" + ) + + +def _postmix_init(seed): + """Non-identity U/V so the postmix parameter gradients are non-trivial.""" + paddle.seed(seed) + u = (paddle.randn([NH, RANK], "float32") * 0.05).astype("bfloat16") + v = (paddle.randn([NH, RANK], "float32") * 0.05).astype("bfloat16") + return u, v + + +def _run_layer(fuse, sq, layer_number=0, seed=0, expect_gate=None): + config = _make_config(fuse_inv_rope_into_vha_postmix=fuse) + attn = _build(config, layer_number=layer_number) + gate = attn._can_fuse_inv_rope_postmix(False) + if expect_gate is not None and gate is not expect_gate: + raise AssertionError( + f"fuse={fuse} should give gate={expect_gate}, got {gate}" + ) + u, v = _postmix_init(_SEED + 1) + attn.vha_postmix_U.set_value(u) + attn.vha_postmix_V.set_value(v) + attn.train() + + paddle.seed(seed) + x = paddle.randn([1, sq, HIDDEN], "bfloat16") + x.stop_gradient = False + g_out = paddle.randn([1, sq, HIDDEN], "bfloat16") + + out, _bias = attn(x, attention_mask=None) + out.backward(g_out) + grads = { + name: (None if p.grad is None else p.grad.clone()) + for name, p in attn.named_parameters() + } + return out.detach(), x.grad.clone(), grads + + +class TestInvRopePostmixLayerBitwise(unittest.TestCase): + def _compare(self, sq, layer_number=0, seed=0): + ref = _run_layer(False, sq, layer_number, seed, expect_gate=False) + got = _run_layer(True, sq, layer_number, seed, expect_gate=True) + tag = f"sq={sq} layer={layer_number} seed={seed}" + + _check_equal(got[0], ref[0], f"layer output ({tag})") + _check_equal(got[1], ref[1], f"grad hidden_states ({tag})") + + self.assertEqual(sorted(got[2]), sorted(ref[2])) + checked = 0 + for name in sorted(ref[2]): + a, b = got[2][name], ref[2][name] + self.assertEqual( + a is None, b is None, f"grad presence differs for {name}" + ) + if a is None: + continue + _check_equal(a, b, f"grad {name} ({tag})") + checked += 1 + # The postmix parameters must be among the compared gradients and must + # actually carry signal, otherwise the weight-gradient half of this test + # proves nothing. + for name in ("vha_postmix_U", "vha_postmix_V"): + self.assertIsNotNone(ref[2][name], f"{name} got no gradient") + self.assertGreater( + float(ref[2][name].astype("float32").abs().sum()), + 0.0, + f"{name} gradient is all zero; test is vacuous", + ) + self.assertGreaterEqual(checked, 8, "suspiciously few parameter grads") + + def test_determinism_precondition(self) -> None: + """Two unfused runs must agree, or nothing else here is meaningful.""" + a = _run_layer(False, 32, expect_gate=False) + b = _run_layer(False, 32, expect_gate=False) + _check_equal(a[0], b[0], "unfused output, run twice") + _check_equal(a[1], b[1], "unfused grad hidden_states, run twice") + for name in sorted(a[2]): + if a[2][name] is not None: + _check_equal(a[2][name], b[2][name], f"unfused grad {name}") + + def test_bitwise_sq32(self) -> None: + self._compare(32) + + def test_bitwise_sq128(self) -> None: + self._compare(128) + + def test_bitwise_odd_seq(self) -> None: + self._compare(97) + + def test_bitwise_other_seed(self) -> None: + self._compare(64, seed=7) + + def test_bitwise_other_layer(self) -> None: + self._compare(64, layer_number=2) + + def test_grouped_postmix_falls_back(self) -> None: + """grouped=True has no [nh,nh] GEMM to split, so it must not fuse.""" + config = _make_config( + fuse_inv_rope_into_vha_postmix=True, + vha_postmix_grouped=True, + vha_postmix_rank=1, + ) + attn = _build(config) + self.assertFalse(attn._can_fuse_inv_rope_postmix(False)) + attn.train() + paddle.seed(0) + x = paddle.randn([1, 32, HIDDEN], "bfloat16") + out, _ = attn(x, attention_mask=None) + self.assertEqual(out.shape, [1, 32, HIDDEN]) + + def test_high_precision_rope_falls_back(self) -> None: + config = _make_config( + fuse_inv_rope_into_vha_postmix=True, high_precision_rope=True + ) + attn = _build(config) + self.assertFalse(attn._can_fuse_inv_rope_postmix(False)) + + +if __name__ == "__main__": + unittest.main()