[release/0.4] Fuse HCA inverse RoPE into the VHA postmix, and the sparse-attn KV index remap

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

src/paddlefleet/fusions/csa_sparse_attn.py

@@ -434,6 +434,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
 
@@ -442,6 +443,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
@@ -521,6 +523,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,
             )
             if head_tile != np_heads:
                 lse_real = lse[:, :, :np_heads].contiguous()
@@ -586,7 +589,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
@@ -610,7 +613,9 @@ def backward(ctx, grad_output):
             do_flat = grad_output.reshape([b * sq, kh, hn])
             kv_flat = kv_full.reshape([b * s_kv, dkv_dim])
             lse_flat = lse.reshape([b * sq, kh])
-            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
+            )
 
             if ctx.kernel_hn != hn:
                 # Same exact zero-padding the forward used (see
@@ -705,6 +710,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.
 
@@ -714,6 +720,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.
 
     ``query`` may carry any head count up to 128 on the "cudnn" backend; counts
     that are not one of the kernel's head tiles (e.g. 32 or 24) are handled
@@ -745,6 +756,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

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)

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:
@@ -175,6 +177,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
 
@@ -252,14 +255,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.
             #
@@ -390,9 +398,10 @@ def backward(ctx, grad_output):
                     d_attn_sink = d_attn_sink.cast("float32")
 
         # 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)
@@ -408,6 +417,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).
@@ -436,6 +446,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).
 
@@ -452,6 +467,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

src/paddlefleet/transformer/csa_attention.py

@@ -128,18 +128,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
 
@@ -2130,6 +2144,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.
@@ -3121,4 +3138,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,
         )

src/paddlefleet/transformer/dsv4_hybrid_attention.py

@@ -796,6 +796,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(
@@ -810,7 +811,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
             ):
@@ -846,7 +856,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
@@ -903,6 +913,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.
 

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,
         )