chore: import upstream snapshot with attribution

python/paddle/incubate/nn/__init__.py

@@ -0,0 +1,34 @@
+#   Copyright (c) 2021 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.
+
+from .layer.fused_dropout_add import FusedDropoutAdd
+from .layer.fused_dropout_nd import FusedDropout  # noqa: F401
+from .layer.fused_linear import FusedLinear
+from .layer.fused_transformer import (
+    FusedBiasDropoutResidualLayerNorm,
+    FusedFeedForward,
+    FusedMultiHeadAttention,
+    FusedMultiTransformer,
+    FusedTransformerEncoderLayer,
+)
+
+__all__ = [
+    'FusedMultiHeadAttention',
+    'FusedFeedForward',
+    'FusedTransformerEncoderLayer',
+    'FusedMultiTransformer',
+    'FusedLinear',
+    'FusedBiasDropoutResidualLayerNorm',
+    'FusedDropoutAdd',
+]

python/paddle/incubate/nn/attn_bias.py

@@ -0,0 +1,275 @@
+# Copyright (c) 2023 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.
+
+
+# The following codes are from https://github.com/facebookresearch/xformers
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import paddle
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class AttentionBias(ABC):
+    @abstractmethod
+    def materialize(self, shape, dtype=paddle.float32):
+        raise NotImplementedError
+
+
+class LowerTriangularMask(AttentionBias):
+    def materialize(self, shape, dtype=paddle.float32):
+        create_as = dtype if dtype is not paddle.bfloat16 else paddle.float32
+        tensor = paddle.full(
+            shape=shape, fill_value=float("-inf"), dtype=create_as
+        )
+        return paddle.triu(tensor, diagonal=1).astype(dtype)
+
+    def add_bias(self, bias):
+        return LowerTriangularMaskWithTensorBias(bias)
+
+
+class LowerTriangularMaskWithTensorBias(LowerTriangularMask):
+    def __init__(self, bias):
+        self._bias = bias
+
+    def materialize(self, shape, dtype=paddle.float32):
+        return super().materialize(shape, dtype) + self._bias
+
+
+@dataclass
+class SeqLenInfo:
+    seqstart: paddle.Tensor
+    max_seqlen: int
+    seqstart_py: list[int]
+
+    def intervals(self):
+        yield from zip(self.seqstart_py, self.seqstart_py[1:])
+
+    @classmethod
+    def from_seqlens(cls, seqlens):
+        seqstart_py = [0]
+        max_seqlen = -1
+        for seqlen in seqlens:
+            max_seqlen = max(max_seqlen, seqlen)
+            seqstart_py.append(seqstart_py[-1] + seqlen)
+        seqstart = paddle.to_tensor(seqstart_py, dtype=paddle.int32)
+        return cls(
+            max_seqlen=max_seqlen, seqstart=seqstart, seqstart_py=seqstart_py
+        )
+
+    def split(self, x, batch_sizes=None):
+        assert self.seqstart_py[-1] == x.shape[1] and x.shape[0] == 1
+        if batch_sizes is None:
+            batch_sizes = [1] * (len(self.seqstart_py) - 1)
+        split_chunks = []
+        it = 0
+        for batch_size in batch_sizes:
+            split_chunks.append(
+                self.seqstart_py[it + batch_size] - self.seqstart_py[it]
+            )
+            it += batch_size
+        return [
+            tensor.reshape([bs, -1, *tensor.shape[2:]])
+            for bs, tensor in zip(batch_sizes, x.split(split_chunks, axis=1))
+        ]
+
+
+@dataclass
+class PaddedSeqLenInfo(SeqLenInfo):
+    seqlen: paddle.Tensor
+    seqlen_py: Sequence[int]
+
+    def intervals(self):
+        for (start, _), length in zip(super().intervals(), self.seqlen_py):
+            yield start, start + length
+
+    @classmethod
+    def from_seqlens(cls, seqlens):
+        raise NotImplementedError(
+            "Please use SeqLenInfo.from_seq_lens() or PaddedSeqLenInfo.from_seq_lens_padded()."
+        )
+
+    @classmethod
+    def from_seqlens_padded(cls, seqlens, padding):
+        assert all(seqlen <= padding for seqlen in seqlens)
+        seqstart_py = list(range(0, len(seqlens) * padding + 1, padding))
+        return cls(
+            seqlen=paddle.to_tensor(seqlens, dtype=paddle.int32),
+            seqlen_py=seqlens,
+            max_seqlen=max(seqlens),
+            seqstart=paddle.to_tensor(seqstart_py, dtype=paddle.int32),
+            seqstart_py=seqstart_py,
+        )
+
+    def split(self, x, batch_sizes=None):
+        raise NotImplementedError
+
+
+@dataclass
+class BlockDiagonalMask(AttentionBias):
+    q_seqinfo: SeqLenInfo
+    k_seqinfo: SeqLenInfo
+    _batch_sizes: Sequence[int] | None = None
+
+    def _create_block_mask(self, shape, dtype=paddle.float32):
+        return paddle.zeros(shape=shape, dtype=dtype)
+
+    def materialize(self, shape, dtype=paddle.float32):
+        assert shape[-1] == self.k_seqinfo.seqstart_py[-1]
+        assert shape[-2] == self.q_seqinfo.seqstart_py[-1]
+        mask = paddle.full(shape[-2:], fill_value=float('-inf'), dtype=dtype)
+        for (q_start, q_end), (k_start, k_end) in zip(
+            self.q_seqinfo.intervals(), self.k_seqinfo.intervals()
+        ):
+            sub_shape = [q_end - q_start, k_end - k_start]
+            mask[q_start:q_end, k_start:k_end] = self._create_block_mask(
+                sub_shape, dtype
+            )
+        for _ in range(len(shape) - 2):
+            mask = mask.unsqueeze(0)
+        return mask.expand(shape)
+
+    @classmethod
+    def from_seqlens(cls, q_seqlen, kv_seqlen=None):
+        assert kv_seqlen is None or len(q_seqlen) == len(kv_seqlen)
+        q_seqinfo = SeqLenInfo.from_seqlens(q_seqlen)
+        if kv_seqlen is None or q_seqlen == kv_seqlen:
+            k_seqinfo = q_seqinfo
+        else:
+            k_seqinfo = SeqLenInfo.from_seqlens(kv_seqlen)
+        return cls(q_seqinfo=q_seqinfo, k_seqinfo=k_seqinfo)
+
+    @classmethod
+    def from_tensor_list(cls, tensors):
+        batch_sizes = [tensor.shape[0] for tensor in tensors]
+        seqlens = []
+        for x in tensors:
+            for _ in range(x.shape[0]):
+                seqlens.append(x.shape[1])
+        block_diag = cls.from_seqlens(seqlens)
+        block_diag._batch_sizes = batch_sizes
+        concated_tensor = paddle.concat(
+            [x.reshape([1, -1, *x.shape[2:]]) for x in tensors], axis=1
+        )
+        return block_diag, concated_tensor
+
+    @classmethod
+    def from_tensor_lists_qkv(cls, tensors_q, tensors_k, tensors_v=None):
+        assert len(tensors_q) == len(tensors_k)
+        assert tensors_v is None or len(tensors_v) == len(tensors_q)
+        batch_sizes = [tensor.shape[0] for tensor in tensors_q]
+        q_seqlens, kv_seqlens = [], []
+        for i, (q, k) in enumerate(zip(tensors_q, tensors_k)):
+            assert q.shape[0] == k.shape[0]
+            q_seqlens.extend([q.shape[1]] * q.shape[0])
+            kv_seqlens.extend([k.shape[1]] * k.shape[0])
+            assert tensors_v is None or tensors_v[i].shape[:2] == k.shape[:2]
+        block_diag = cls.from_seqlens(q_seqlens, kv_seqlens)
+        block_diag._batch_sizes = [x.shape[0] for x in tensors_q]
+        return (
+            block_diag,
+            paddle.concat(
+                [x.reshape([1, -1, *x.shape[2:]]) for x in tensors_q], axis=1
+            ),
+            paddle.concat(
+                [x.reshape([1, -1, *x.shape[2:]]) for x in tensors_k], axis=1
+            ),
+            (
+                paddle.concat(
+                    [x.reshape([1, -1, *x.shape[2:]]) for x in tensors_v],
+                    axis=1,
+                )
+                if tensors_v is not None
+                else None
+            ),
+        )
+
+    def split_queries(self, tensor):
+        return self.q_seqinfo.split(tensor, self._batch_sizes)
+
+    def split_kv(self, tensor):
+        return self.k_seqinfo.split(tensor, self._batch_sizes)
+
+    def split(self, tensor):
+        assert self.q_seqinfo is self.k_seqinfo
+        return self.q_seqinfo.split(tensor, self._batch_sizes)
+
+    def make_causal(self):
+        return BlockDiagonalCausalMask(
+            q_seqinfo=self.q_seqinfo,
+            k_seqinfo=self.k_seqinfo,
+            _batch_sizes=self._batch_sizes,
+        )
+
+
+@dataclass
+class BlockDiagonalCausalMask(BlockDiagonalMask):
+    def _create_block_mask(self, shape, dtype=paddle.float32):
+        return LowerTriangularMask().materialize(shape=shape, dtype=dtype)
+
+
+@dataclass
+class BlockDiagonalCausalWithOffsetPaddedKeysMask(AttentionBias):
+    q_seqinfo: SeqLenInfo
+    k_seqinfo: PaddedSeqLenInfo
+    causal_diagonal: paddle.Tensor | None = None
+
+    def _create_block_mask(self, shape, offset=0, dtype=paddle.float32):
+        create_as = dtype if dtype is not paddle.bfloat16 else paddle.float32
+        tensor = paddle.full(shape, dtype=create_as, fill_value=float('-inf'))
+        return paddle.triu(tensor, diagonal=1 + offset).astype(dtype)
+
+    def materialize(self, shape, dtype=paddle.float32):
+        assert shape[-1] == self.k_seqinfo.seqstart_py[-1]
+        assert shape[-2] == self.q_seqinfo.seqstart_py[-1]
+        mask = paddle.full(shape[-2:], dtype=dtype, fill_value=float('-inf'))
+        for i, ((q_start, q_end), (k_start, k_end)) in enumerate(
+            zip(self.q_seqinfo.intervals(), self.k_seqinfo.intervals())
+        ):
+            mask[q_start:q_end, k_start:k_end] = self._create_block_mask(
+                (q_end - q_start, k_end - k_start),
+                offset=(
+                    0
+                    if self.causal_diagonal is None
+                    else int(self.causal_diagonal[i].item())
+                ),
+                dtype=dtype,
+            )
+        for _ in range(len(shape) - 2):
+            mask = mask.unsqueeze(0)
+        return mask.expand(shape)
+
+    @classmethod
+    def from_seqlens(
+        cls, q_seqlen, kv_padding, kv_seqlen, causal_diagonal=None
+    ):
+        assert kv_seqlen is None or len(q_seqlen) == len(kv_seqlen)
+        q_seqinfo = SeqLenInfo.from_seqlens(q_seqlen)
+        k_seqinfo = PaddedSeqLenInfo.from_seqlens_padded(kv_seqlen, kv_padding)
+        return cls(
+            q_seqinfo=q_seqinfo,
+            k_seqinfo=k_seqinfo,
+            causal_diagonal=causal_diagonal,
+        )

python/paddle/incubate/nn/functional/__init__.py

@@ -0,0 +1,123 @@
+# ruff: noqa: F401
+# Copyright (c) 2021 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.
+
+
+from .batched_gemm import batched_gemm
+from .blha_get_max_len import blha_get_max_len
+from .block_multihead_attention import (
+    block_multihead_attention,
+    block_multihead_attention_xpu,
+)
+
+# from .moe_gate_dispatch_permute import moe_gate_dispatch_permute
+from .build_src_rank_and_local_expert_id import (
+    build_src_rank_and_local_expert_id,
+)
+from .cal_aux_loss import cal_aux_loss
+from .cross_entropy_with_softmax_bwd_w_downcast import (
+    cross_entropy_with_softmax_bwd_w_downcast,
+)
+from .embedding_grad_add_to import embedding_grad_add_to_
+from .expand_modality_expert_id import expand_modality_expert_id
+from .fast_ln import fast_ln
+from .fast_rms_norm import fast_rms_norm
+from .fp8 import (
+    fp8_gemm_blockwise,
+    fp8_quant_blockwise,
+    fused_act_dequant,
+    fused_stack_transpose_quant,
+    fused_swiglu_weighted_bwd,
+    fused_transpose_split_quant,
+    fused_transpose_wlch_split_quant,
+    fused_weighted_swiglu_act_quant,
+)
+from .fused_bias_act import fused_bias_act
+from .fused_dot_product_attention import (
+    cudnn_flash_attention,
+    fused_dot_product_attention,
+)
+from .fused_dropout_add import fused_dropout_add
+from .fused_gate_attention import fused_gate_attention
+from .fused_layer_norm import fused_layer_norm
+from .fused_matmul_bias import (
+    fused_linear,
+    fused_linear_activation,
+    fused_matmul_bias,
+)
+from .fused_partial_rope import fused_partial_rope
+from .fused_rms_norm import fused_rms_norm
+from .fused_rms_norm_ext import fused_rms_norm_ext
+from .fused_rotary_position_embedding import fused_rotary_position_embedding
+from .fused_transformer import (
+    fused_bias_dropout_residual_layer_norm,
+    fused_feedforward,
+    fused_multi_head_attention,
+    fused_multi_transformer,
+)
+from .int_bincount import int_bincount
+from .masked_multihead_attention import masked_multihead_attention
+from .moe_combine import moe_combine
+from .moe_combine_no_weight import moe_combine_no_weight
+from .moe_gate_dispatch import moe_gate_dispatch
+from .moe_gate_dispatch_and_quant import moe_gate_dispatch_and_quant
+from .moe_gate_dispatch_partial_nosoftmaxtopk import (
+    moe_gate_dispatch_partial_nosoftmaxtopk,
+)
+from .moe_gate_dispatch_permute import moe_gate_dispatch_permute
+from .swiglu import swiglu
+from .variable_length_memory_efficient_attention import (
+    variable_length_memory_efficient_attention,
+)
+
+__all__ = [
+    'embedding_grad_add_to_',
+    'fp8_gemm_blockwise',
+    'cross_entropy_with_softmax_bwd_w_downcast',
+    'fp8_quant_blockwise',
+    'fast_ln',
+    'fast_rms_norm',
+    'fused_act_dequant',
+    'fused_multi_head_attention',
+    'fused_feedforward',
+    'fused_multi_transformer',
+    'fused_matmul_bias',
+    'fused_linear',
+    'fused_linear_activation',
+    'fused_bias_dropout_residual_layer_norm',
+    'fused_dropout_add',
+    'fused_rotary_position_embedding',
+    'fused_stack_transpose_quant',
+    'fused_transpose_split_quant',
+    'fused_transpose_wlch_split_quant',
+    'variable_length_memory_efficient_attention',
+    "fused_rms_norm",
+    "fused_layer_norm",
+    "fused_bias_act",
+    'fused_swiglu_weighted_bwd',
+    'fused_weighted_swiglu_act_quant',
+    "masked_multihead_attention",
+    "blha_get_max_len",
+    "block_multihead_attention",
+    "moe_combine",
+    "expand_modality_expert_id",
+    "cal_aux_loss",
+    "build_src_rank_and_local_expert_id",
+    "int_bincount",
+    "fused_rms_norm_ext",
+    "moe_gate_dispatch",
+    "moe_gate_dispatch_permute",
+    "moe_gate_dispatch_partial_nosoftmaxtopk",
+    "moe_gate_dispatch_and_quant",
+]

python/paddle/incubate/nn/functional/batched_gemm.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from paddle import Tensor, _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+
+
+def batched_gemm(
+    lhs: Tensor,
+    rhs: Tensor,
+    batch_sizes: list,
+    trans_lhs: bool = False,
+    trans_rhs: bool = False,
+) -> tuple[Tensor]:
+    """
+    Cluster launched gemm into one op, which can be further fused and optimized.
+
+    Args:
+        lhs (Tensor): A tensor shaped in (total_seq_len, input_hidden_size), meant to be
+        perform gemm operation according to batch range.
+        rhs (Tensor): A tensor shaped in (num_batches, input_hidden_size, output_hidden_size).
+        batch_sizes(list): A list of integers representing the number of rows in each batch.
+        trans_lhs (bool): Whether view lhs matrix as last 2D-transposed. Default: False.
+        trans_rhs (bool): Whether view rhs matrix as last 2D-transposed. Default: False.
+
+    Returns:
+        tuple:
+            - out (Tensor): The result of batched gemm operation.
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.batched_gemm(lhs, rhs, batch_sizes, trans_lhs, trans_rhs)

python/paddle/incubate/nn/functional/blha_get_max_len.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2024 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def blha_get_max_len(
+    seq_lens_encoder: Tensor, seq_lens_decoder: Tensor, batch_size: Tensor
+) -> tuple[Tensor, Tensor]:
+    """
+    Apply Fused BlhaGetMaxLen kernel. Typically used before the block_multihead_attention operator.
+
+    Args:
+        seq_lens_encoder (Tensor): Sentence length of the encoder.
+        seq_lens_decoder (Tensor): Sentence length of the decoder.
+        batch_size (Tensor): the batch size.
+
+    Returns:
+        Tensor|(max_enc_len_this_time, max_dec_len_this_time)
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+
+            >>> seq_lens_encoder = paddle.cast(paddle.randn(shape=[10]), dtype=paddle.int32)
+            >>> seq_lens_decoder = paddle.cast(paddle.randn(shape=[10]), dtype=paddle.int32)
+            >>> bsz = 10
+            >>> batch_size = paddle.ones(shape=[bsz])
+            >>> max_enc_len_this_time, max_dec_len_this_time = paddle.incubate.nn.functional.blha_get_max_len(
+            ...     seq_lens_encoder, seq_lens_decoder, batch_size
+            ... )
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.blha_get_max_len(
+            seq_lens_encoder, seq_lens_decoder, batch_size
+        )
+
+    helper = LayerHelper('blha_get_max_len', **locals())
+    max_enc_len_this_time = helper.create_variable_for_type_inference(
+        dtype="int32"
+    )
+    max_dec_len_this_time = helper.create_variable_for_type_inference(
+        dtype="int32"
+    )
+
+    inputs = {}
+    inputs['seq_lens_encoder'] = seq_lens_encoder
+    inputs['seq_lens_decoder'] = seq_lens_decoder
+    inputs['batch_size'] = batch_size
+
+    outputs = {
+        'max_enc_len_this_time': max_enc_len_this_time,
+        'max_dec_len_this_time': max_dec_len_this_time,
+    }
+    helper.append_op(
+        type='blha_get_max_len',
+        inputs=inputs,
+        outputs=outputs,
+    )
+    return max_enc_len_this_time, max_dec_len_this_time

python/paddle/incubate/nn/functional/block_multihead_attention.py

@@ -0,0 +1,583 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+from paddle import _C_ops
+from paddle.framework import (
+    LayerHelper,
+    in_dynamic_or_pir_mode,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+    _QuantRoundType: TypeAlias = Literal[0, 1]
+
+
+def block_multihead_attention(
+    qkv: Tensor,
+    key_cache: Tensor,
+    value_cache: Tensor,
+    seq_lens_encoder: Tensor,
+    seq_lens_decoder: Tensor,
+    seq_lens_this_time: Tensor,
+    padding_offsets: Tensor,
+    cum_offsets: Tensor,
+    cu_seqlens_q: Tensor,
+    cu_seqlens_k: Tensor,
+    block_tables: Tensor,
+    pre_key_cache: Tensor | None = None,
+    pre_value_cache: Tensor | None = None,
+    cache_k_quant_scales: Tensor | None = None,
+    cache_v_quant_scales: Tensor | None = None,
+    cache_k_dequant_scales: Tensor | None = None,
+    cache_v_dequant_scales: Tensor | None = None,
+    qkv_out_scale: Tensor | None = None,
+    qkv_bias: Tensor | None = None,
+    out_shift: Tensor | None = None,
+    out_smooth: Tensor | None = None,
+    max_enc_len_this_time: Tensor | None = None,
+    max_dec_len_this_time: Tensor | None = None,
+    rope_emb: Tensor | None = None,
+    mask: Tensor | None = None,
+    tgt_mask: Tensor | None = None,
+    max_seq_len: int = -1,
+    block_size: int = 64,
+    use_neox_style: bool = False,
+    use_dynamic_cachekv_quant: bool = False,
+    quant_round_type: _QuantRoundType = 1,
+    quant_max_bound: float = 127.0,
+    quant_min_bound: float = -127.0,
+    out_scale: float = -1,
+    compute_dtype: str = "default",
+    rope_theta: float = 10000.0,
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+    """
+    Block Multi-head attention for text summarization.
+
+    Args:
+        qkv (Tensor): The qkv Tensor. Its shape is [token_num, 3 * num_head * head_size].
+        key_cache (Tensor): The key_cache Tensor. Its shape is [max_block_num, num_head, block_size, head_size].
+        value_cache (Tensor): The value_cache Tensor. Its shape is [max_block_num, num_head, block_size, head_size].
+        seq_lens_encoder (Tensor): The encoder sequence lengths of the sequences in the batch. Its shape is [batchsize, 1].
+        seq_lens_decoder (Tensor): The decoder sequence lengths of the sequences in the batch. Its shape is [batchsize, 1].
+        seq_lens_this_time (Tensor): The real sequence lengths of the sequences in the batch. Its shape is [batchsize, 1].
+        padding_offsets (Tensor): The offsets from unpadding to padding. Its shape is [token_num].
+        cum_offsets (Tensor): The offsets from padding to unpadding. Its shape is [batchsize].
+        cu_seqlens_q (Tensor): The cum sequence lengths of query. Its shape is [batchsize + 1, 1].
+        cu_seqlens_k (Tensor): The cum sequence lengths of key. Its shape is [batchsize + 1, 1].
+        block_tables (Tensor): The block tables, used to index the cache. Its shape is [batchsize, block_num_per_seq].
+        pre_key_cache (Tensor): The pre caches of key. Its shape is [batchsize, num_head, pre_cache_length, head_size].
+        pre_value_cache (Tensor): The pre caches of value. Its shape is [batchsize, num_head, pre_cache_length, head_size].
+        cache_k_quant_scales (Tensor): The quant scales of cache key. Its shape depends on quant mode (dynamic or static). If dynamic quantization is enabled, its shape is [batchsize, num_head], otherwise its shape is [num_head].
+        cache_v_quant_scales (Tensor): The quant scales of cache value. Its shape depends on quant mode (dynamic or static). If dynamic quantization is enabled, its shape is [batchsize, num_head], otherwise its shape is [num_head].
+        cache_k_dequant_scales (Tensor): The dequant scales of cache key. Its shape depends on quant mode (dynamic or static). If dynamic quantization is enabled, its shape is [batchsize, num_head], otherwise its shape is [num_head].
+        cache_v_dequant_scales (Tensor): The dequant scales of cache value. Its shape depends on quant mode (dynamic or static). If dynamic quantization is enabled, its shape is [batchsize, num_head], otherwise its shape is [num_head].
+        qkv_out_scale (Tensor): The dequant scale of qkv, which is the input of BLHA. If the dtype of qkv is `int32`, this input will be applied. Its shape is [3 * num_head * head_size], and its dtype should be `float32`.
+        qkv_bias (Tensor): The bias of qkv. Its shape is [3 * num_head * head_size].
+        out_shift (Tensor): Shift bias of fmha_out, which is the 1st return value. Its shape is [num_head * head_size].
+        out_smooth (Tensor): Smooth weight of fmha_out. Its shape is [num_head * head_size].
+        max_enc_len_this_time (Tensor): Sentence length of the encoder this time. Its shape is [1].
+        max_dec_len_this_time (Tensor): Sentence length of the decoder this time. Its shape is [1].
+        rope_emb (Tensor): The RoPE embedding. Its shape is [2, batchsize, max_seq_len, 1, head_size // 2].
+        mask (Tensor): The mask of qk_matmul in encoder. Its shape is [batchsize, 1, max_seq_len, max_seq_len].
+        tgt_mask (Tensor): The mask of qk_matmul in decoder. Its shape is [batchsize, 1, 1, max_seq_len].
+        max_seq_len (Int): The max length of the input. Default is -1.
+        block_size (Int): The block_size of cache. Default is 64.
+        use_neox_style (Bool): Whether neox_style RoPE is used or not. Default is False.
+        use_dynamic_cachekv_quant (Bool): Whether dynamic cache kv quantization is applied or not. Default is False.
+        quant_round_type (Int): The quant round type in cache kv quantization and fmha_out quantization. If 0 is set, value will be rounding to nearest ties to even. If 1 is set, value will be rounding to nearest ties away from zero.
+        quant_max_bound (Float32): The max bound of float type to int type.
+        quant_min_bound (Float32): The min bound of float type to int type.
+        out_scale (Float32): The quant scale of fmha_out. Default is -1, which means do not apply quantization for fmha_out.
+        compute_dtype (Str): A compute dtype, is used to represent the input data type. Default is "default", which means compute dtype is determined by input dtype. However, if the dtype of input is Int32, this value should be set to actual dtype of the model.
+        rope_theta (Float32): The theta of RoPE. Default is 10000.0.
+    Returns:
+        Tensor|(output, qkv_out, cache_k_out, cache_v_out), which output is the output of
+        block_multihead_attention layers, qkv_out is inplace with input `qkv`, cache_k_out and cache_v_out are inplace with input `cache_k` and `cache_v`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('Need compile flash attention')
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import numpy as np
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import (
+            ...     block_multihead_attention,
+            ... )
+            >>> paddle.device.set_device('gpu')
+
+            >>> def get_padding_offset(bsz, max_seq_len, seq_lens_this_time):
+            ...     cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time)
+            ...     cum_offsets = paddle.zeros(shape=(bsz + 1), dtype="int32")
+            ...     cum_offsets[1:] = cum_offsets_now
+            ...     token_num = paddle.sum(seq_lens_this_time)
+            ...     padding_offsets = paddle.zeros(shape=(token_num), dtype="int32")
+            ...     cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32")
+            ...     cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32")
+            ...     for i in range(bsz):
+            ...         seq_len_now = seq_lens_this_time[i]
+            ...         cum_offset = cum_offsets[i]
+            ...         for j in range(seq_len_now):
+            ...             padding_offsets[i * max_seq_len - cum_offset + j] = cum_offset
+            ...         cum_seq_len = (i + 1) * max_seq_len - cum_offsets[i + 1]
+            ...         cu_seqlens_q[i + 1] = cum_seq_len
+            ...         cu_seqlens_k[i + 1] = cum_seq_len
+            ...     return (
+            ...         padding_offsets,
+            ...         cum_offsets[:-1],
+            ...         cu_seqlens_q,
+            ...         cu_seqlens_k,
+            ...     )
+
+            >>> def remove_padding(seq_lens, cu_seq_lens, inputs, token_num):
+            ...     bsz, num_head, seq_len, head_size = inputs.shape
+            ...     output = paddle.zeros(
+            ...         shape=[token_num, num_head * head_size],
+            ...         dtype=inputs.dtype,
+            ...     )
+            ...     inputs = inputs.transpose([0, 2, 1, 3]).reshape([bsz, seq_len, -1])
+            ...     for i in range(bsz):
+            ...         seq_len_now = seq_lens[i]
+            ...         start_idx = cu_seq_lens[i]
+            ...         end_idx = cu_seq_lens[i + 1]
+            ...         output[start_idx:end_idx, :] = inputs[i, :seq_len_now, :]
+            ...     return output
+
+            >>> def create_attn_mask(
+            ...     mask_type,
+            ...     batch_size,
+            ...     seq_lens,
+            ...     pre_cache_length=0,
+            ... ):
+            ...     max_seq_len = max(seq_lens)
+            ...     mask = paddle.zeros(
+            ...         [batch_size, 1, max_seq_len, max_seq_len + pre_cache_length],
+            ...         dtype=mask_type,
+            ...     )
+            ...     mask[:, :, :, :pre_cache_length] = 1
+            ...     for i in range(batch_size):
+            ...         seq_len = seq_lens[i]
+            ...         mask[i, 0, :seq_len, :seq_len] = (paddle.tril(paddle.ones(shape=(seq_len, seq_len), dtype=mask_type)) - 1) * 1e4
+            ...     return mask
+
+            >>> def naive_attention_impl(
+            ...     query,
+            ...     key,
+            ...     value,
+            ...     cache_k,
+            ...     cache_v,
+            ...     pre_cache_k,
+            ...     pre_cache_v,
+            ...     mask,
+            ...     scale=1.0,
+            ... ):
+            ...     batch = query.shape[0]
+            ...     heads = query.shape[1]
+            ...     seq_len = query.shape[2]
+            ...     head_dim = query.shape[3]
+            ...     kv_head = key.shape[1]
+            ...     key = key.reshape([batch, kv_head, 1, seq_len, head_dim])
+            ...     key = paddle.tile(key, [1, 1, heads // kv_head, 1, 1])
+            ...     key = key.reshape([batch, heads, seq_len, head_dim])
+            ...     if pre_cache_k is not None:
+            ...         key = paddle.concat([pre_cache_k, key], axis=2)
+            ...     if cache_k is not None:
+            ...         key = paddle.concat([cache_k, key], axis=2)
+            ...     value = value.reshape([batch, kv_head, 1, seq_len, head_dim])
+            ...     value = paddle.tile(value, [1, 1, heads // kv_head, 1, 1])
+            ...     value = value.reshape([batch, heads, seq_len, head_dim])
+            ...     if pre_cache_v is not None:
+            ...         value = paddle.concat([pre_cache_v, value], axis=2)
+            ...     if cache_v is not None:
+            ...         value = paddle.concat([cache_v, value], axis=2)
+            ...     qk_res = paddle.matmul(query, key, transpose_y=True)
+            ...     attention = qk_res * scale
+            ...     if mask is not None:
+            ...         attention = attention + mask
+            ...     softmax_result = paddle.nn.functional.softmax(attention, -1)
+            ...     result = paddle.matmul(softmax_result, value)
+            ...     return result
+
+            >>> batch_size = 2
+            >>> num_head = 8
+            >>> seq_len = 64
+            >>> max_dec_len = 64
+            >>> head_size = 64
+            >>> hid_dim = num_head * head_size
+            >>> block_size = 64
+            >>> block_num_per_seq = (seq_len + max_dec_len + block_size - 1) // block_size
+            >>> max_block_num = block_num_per_seq * batch_size
+            >>> free_list = list(range(max_block_num - 1, -1, -1))
+            >>> token_num = seq_len * batch_size
+
+            >>> dtype = paddle.float16
+
+            >>> seq_lens_encoder = paddle.to_tensor(
+            ...     [
+            ...         seq_len,
+            ...     ]
+            ...     * batch_size,
+            ...     "int32",
+            ... )
+            >>> seq_lens_decoder = paddle.to_tensor(
+            ...     [0] * batch_size,
+            ...     "int32",
+            ... )
+            >>> seq_lens_this_time = seq_lens_encoder
+            >>> qkv_shape = (
+            ...     batch_size,
+            ...     num_head,
+            ...     seq_len,
+            ...     head_size,
+            ... )
+            >>> q = paddle.randn(shape=qkv_shape, dtype=dtype)
+            >>> k = paddle.randn(shape=qkv_shape, dtype=dtype)
+            >>> v = paddle.randn(shape=qkv_shape, dtype=dtype)
+            >>> qkv = paddle.stack(
+            ...     [
+            ...         q.transpose([0, 2, 1, 3]).reshape([token_num, hid_dim]),
+            ...         k.transpose([0, 2, 1, 3]).reshape([token_num, hid_dim]),
+            ...         v.transpose([0, 2, 1, 3]).reshape([token_num, hid_dim]),
+            ...     ],
+            ...     axis=1,
+            ... ).reshape([token_num, -1])
+            >>> scale = 1.0 / np.sqrt(head_size)
+            >>> cache_shape = (max_block_num, num_head, block_size, head_size)
+            >>> cache_k = paddle.zeros(cache_shape, dtype=dtype)
+            >>> cache_v = paddle.zeros(cache_shape, dtype=dtype)
+            >>> block_tables = paddle.zeros(shape=(batch_size, block_num_per_seq), dtype="int32")
+            >>> for i in range(batch_size):
+            ...     need_block_num = (seq_len + max_dec_len + block_size - 1) // block_size
+            ...     for j in range(need_block_num):
+            ...         block_tables[i, j] = free_list.pop()
+            >>> padding_offset, cum_offset, cu_seqlens_q, cu_seqlens_k = get_padding_offset(batch_size, seq_len, seq_lens_this_time)
+            >>> out = block_multihead_attention(
+            ...     qkv,
+            ...     cache_k,
+            ...     cache_v,
+            ...     seq_lens_encoder,
+            ...     seq_lens_decoder,
+            ...     seq_lens_this_time,
+            ...     padding_offset,
+            ...     cum_offset,
+            ...     cu_seqlens_q,
+            ...     cu_seqlens_k,
+            ...     block_tables,
+            ...     None,  # pre_key_cache
+            ...     None,  # pre_value_cache
+            ...     None,  # cache_k_quant_scales
+            ...     None,  # cache_v_quant_scales
+            ...     None,  # cache_k_dequant_scales
+            ...     None,  # cache_v_dequant_scales
+            ...     None,  # qkv_out_scale
+            ...     None,  # qkv_bias
+            ...     None,  # out_shift
+            ...     None,  # out_smooth
+            ...     None,  # max_enc_len_this_time
+            ...     None,  # max_dec_len_this_time
+            ...     None,  # rotary_embs
+            ...     None,  # attn_mask
+            ...     None,  # tgt_mask
+            ...     seq_len,
+            ...     block_size,
+            ... )[0]
+
+            >>> attention_mask = create_attn_mask(
+            ...     dtype,
+            ...     batch_size,
+            ...     [
+            ...         seq_len,
+            ...     ]
+            ...     * batch_size,
+            ... )
+
+            >>> out_ref = naive_attention_impl(q, k, v, None, None, None, None, attention_mask, scale)
+            >>> out_ref = remove_padding(seq_lens_this_time, cu_seqlens_q, out_ref, token_num)
+            >>> # equals to: out_ref = out
+
+            >>> print(out.shape)  # [token_num, hid_dim]
+            paddle.Size([128, 512])
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.block_multihead_attention_(
+            qkv,
+            key_cache,
+            value_cache,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            seq_lens_this_time,
+            padding_offsets,
+            cum_offsets,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            block_tables,
+            pre_key_cache,
+            pre_value_cache,
+            rope_emb,
+            mask,
+            tgt_mask,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            qkv_out_scale,
+            qkv_bias,
+            out_shift,
+            out_smooth,
+            max_enc_len_this_time,
+            max_dec_len_this_time,
+            max_seq_len,
+            block_size,
+            use_neox_style,
+            use_dynamic_cachekv_quant,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+            out_scale,
+            compute_dtype,
+            rope_theta,
+        )
+
+    helper = LayerHelper('block_multihead_attention', **locals())
+    out = helper.create_variable_for_type_inference(dtype=qkv.dtype)
+
+    inputs = {}
+    inputs['qkv'] = qkv
+    inputs['key_cache'] = key_cache
+    inputs['value_cache'] = value_cache
+    inputs['seq_lens_encoder'] = seq_lens_encoder
+    inputs['seq_lens_decoder'] = seq_lens_decoder
+    inputs['seq_lens_this_time'] = seq_lens_this_time
+    inputs['padding_offsets'] = padding_offsets
+    inputs['cum_offsets'] = cum_offsets
+    inputs['cu_seqlens_q'] = cu_seqlens_q
+    inputs['cu_seqlens_k'] = cu_seqlens_k
+    inputs['block_tables'] = block_tables
+    if pre_key_cache is not None:
+        inputs['pre_key_cache'] = pre_key_cache
+    if pre_value_cache is not None:
+        inputs['pre_value_cache'] = pre_value_cache
+    if rope_emb is not None:
+        inputs['rope_emb'] = rope_emb
+    if mask is not None:
+        inputs['mask'] = mask
+    if tgt_mask is not None:
+        inputs['tgt_mask'] = tgt_mask
+    if cache_k_quant_scales is not None:
+        inputs["cache_k_quant_scales"] = cache_k_quant_scales
+    if cache_v_quant_scales is not None:
+        inputs["cache_v_quant_scales"] = cache_v_quant_scales
+    if cache_k_dequant_scales is not None:
+        inputs["cache_k_dequant_scales"] = cache_k_dequant_scales
+    if cache_v_dequant_scales is not None:
+        inputs["cache_v_dequant_scales"] = cache_v_dequant_scales
+    if qkv_out_scale is not None:
+        inputs["qkv_out_scale"] = qkv_out_scale
+    if qkv_bias is not None:
+        inputs["qkv_bias"] = qkv_bias
+    if out_shift is not None:
+        inputs["out_shift"] = out_shift
+    if out_smooth is not None:
+        inputs["out_smooth"] = out_smooth
+    if max_enc_len_this_time is not None:
+        inputs["max_enc_len_this_time"] = max_enc_len_this_time
+    if max_dec_len_this_time is not None:
+        inputs["max_dec_len_this_time"] = max_dec_len_this_time
+
+    outputs = {
+        'fmha_out': out,
+        'qkv_out': qkv,
+        'key_cache_out': key_cache,
+        'value_cache_out': value_cache,
+    }
+    helper.append_op(
+        type='block_multihead_attention',
+        inputs=inputs,
+        outputs=outputs,
+        attrs={
+            'max_seq_len': max_seq_len,
+            'block_size': block_size,
+            'use_neox_style': use_neox_style,
+            'dynamic_cachekv_quant': use_dynamic_cachekv_quant,
+            'quant_round_type': quant_round_type,
+            'quant_max_bound': quant_max_bound,
+            'quant_min_bound': quant_min_bound,
+            'out_scale': out_scale,
+            'compute_dtype': compute_dtype,
+            'rope_theta': rope_theta,
+        },
+    )
+    return out, qkv, key_cache, value_cache
+
+
+def block_multihead_attention_xpu(
+    qkv: Tensor,
+    key_cache: Tensor,
+    value_cache: Tensor,
+    seq_lens_encoder: Tensor,
+    seq_lens_decoder: Tensor,
+    seq_lens_this_time: Tensor,
+    padding_offsets: Tensor,
+    cum_offsets: Tensor,
+    cu_seqlens_q: Tensor,
+    cu_seqlens_k: Tensor,
+    block_tables: Tensor,
+    cache_k_per_batch_maxs: Tensor,
+    cache_v_per_batch_maxs: Tensor,
+    pre_key_cache: Tensor | None = None,
+    pre_value_cache: Tensor | None = None,
+    cache_k_quant_scales: Tensor | None = None,
+    cache_v_quant_scales: Tensor | None = None,
+    cache_k_dequant_scales: Tensor | None = None,
+    cache_v_dequant_scales: Tensor | None = None,
+    qkv_out_scale: Tensor | None = None,
+    qkv_bias: Tensor | None = None,
+    out_shift: Tensor | None = None,
+    out_smooth: Tensor | None = None,
+    max_enc_len_this_time: Tensor | None = None,
+    max_dec_len_this_time: Tensor | None = None,
+    rope_emb: Tensor | None = None,
+    mask: Tensor | None = None,
+    tgt_mask: Tensor | None = None,
+    max_seq_len: int = -1,
+    block_size: int = 64,
+    use_neox_style: bool = False,
+    use_dynamic_cachekv_quant: bool = False,
+    quant_round_type: _QuantRoundType = 1,
+    quant_max_bound: float = 127.0,
+    quant_min_bound: float = -127.0,
+    out_scale: float = -1,
+    compute_dtype: str = "default",
+    rope_theta: float = 10000.0,
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+    if in_dynamic_or_pir_mode():
+        return _C_ops.block_multihead_attention_xpu(
+            qkv,
+            key_cache,
+            value_cache,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            seq_lens_this_time,
+            padding_offsets,
+            cum_offsets,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            block_tables,
+            cache_k_per_batch_maxs,
+            cache_v_per_batch_maxs,
+            pre_key_cache,
+            pre_value_cache,
+            rope_emb,
+            mask,
+            tgt_mask,
+            cache_k_quant_scales,
+            cache_v_quant_scales,
+            cache_k_dequant_scales,
+            cache_v_dequant_scales,
+            qkv_out_scale,
+            qkv_bias,
+            out_shift,
+            out_smooth,
+            max_enc_len_this_time,
+            max_dec_len_this_time,
+            max_seq_len,
+            block_size,
+            use_neox_style,
+            use_dynamic_cachekv_quant,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+            out_scale,
+            compute_dtype,
+            rope_theta,
+        )
+
+    helper = LayerHelper('block_multihead_attention_xpu', **locals())
+    out = helper.create_variable_for_type_inference(dtype=qkv.dtype)
+
+    inputs = {}
+    inputs['qkv'] = qkv
+    inputs['key_cache'] = key_cache
+    inputs['value_cache'] = value_cache
+    inputs['seq_lens_encoder'] = seq_lens_encoder
+    inputs['seq_lens_decoder'] = seq_lens_decoder
+    inputs['seq_lens_this_time'] = seq_lens_this_time
+    inputs['padding_offsets'] = padding_offsets
+    inputs['cum_offsets'] = cum_offsets
+    inputs['cu_seqlens_q'] = cu_seqlens_q
+    inputs['cu_seqlens_k'] = cu_seqlens_k
+    inputs['block_tables'] = block_tables
+    inputs['cache_k_per_batch_maxs'] = cache_k_per_batch_maxs
+    inputs['cache_v_per_batch_maxs'] = cache_v_per_batch_maxs
+    if pre_key_cache is not None:
+        inputs['pre_key_cache'] = pre_key_cache
+    if pre_value_cache is not None:
+        inputs['pre_value_cache'] = pre_value_cache
+    if rope_emb is not None:
+        inputs['rope_emb'] = rope_emb
+    if mask is not None:
+        inputs['mask'] = mask
+    if tgt_mask is not None:
+        inputs['tgt_mask'] = tgt_mask
+    if cache_k_quant_scales is not None:
+        inputs["cache_k_quant_scales"] = cache_k_quant_scales
+    if cache_v_quant_scales is not None:
+        inputs["cache_v_quant_scales"] = cache_v_quant_scales
+    if cache_k_dequant_scales is not None:
+        inputs["cache_k_dequant_scales"] = cache_k_dequant_scales
+    if cache_v_dequant_scales is not None:
+        inputs["cache_v_dequant_scales"] = cache_v_dequant_scales
+    if qkv_out_scale is not None:
+        inputs["qkv_out_scale"] = qkv_out_scale
+    if qkv_bias is not None:
+        inputs["qkv_bias"] = qkv_bias
+    if out_shift is not None:
+        inputs["out_shift"] = out_shift
+    if out_smooth is not None:
+        inputs["out_smooth"] = out_smooth
+    if max_enc_len_this_time is not None:
+        inputs["max_enc_len_this_time"] = max_enc_len_this_time
+    if max_dec_len_this_time is not None:
+        inputs["max_dec_len_this_time"] = max_dec_len_this_time
+
+    outputs = {
+        'fmha_out': out,
+        'qkv_out': qkv,
+        'key_cache_out': key_cache,
+        'value_cache_out': value_cache,
+    }
+    helper.append_op(
+        type='block_multihead_attention_xpu',
+        inputs=inputs,
+        outputs=outputs,
+        attrs={
+            'max_seq_len': max_seq_len,
+            'block_size': block_size,
+            'use_neox_style': use_neox_style,
+            'dynamic_cachekv_quant': use_dynamic_cachekv_quant,
+            'quant_round_type': quant_round_type,
+            'quant_max_bound': quant_max_bound,
+            'quant_min_bound': quant_min_bound,
+            'out_scale': out_scale,
+            'compute_dtype': compute_dtype,
+            'rope_theta': rope_theta,
+        },
+    )
+    return out, qkv, key_cache, value_cache

python/paddle/incubate/nn/functional/cal_aux_loss.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def math_cal_aux_loss(
+    gate_prob: Tensor,
+    dispatch_mask: Tensor,
+    tokens_mask: Tensor,
+    dispatch_tokens_mask: Tensor,
+    num_experts: int,
+    use_group: bool,
+    moe_k: int,
+    clip_min: float,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        gate_prob
+        dispatch_mask
+        tokens_mask
+        dispatch_tokens_mask
+        num_experts
+        use_group
+        moe_k
+        clip_min
+
+    Returns:
+        l_aux
+        seqlen_float
+        ce
+    """
+    if tokens_mask is not None and tokens_mask.dtype != gate_prob.dtype:
+        tokens_mask = tokens_mask.astype(gate_prob.dtype)
+
+    scale = None
+    if dispatch_tokens_mask is not None:
+        seqlen_float = dispatch_tokens_mask.astype(gate_prob.dtype).sum()
+        if (
+            tokens_mask is not None
+            and gate_prob.shape[0] != dispatch_tokens_mask.shape[0]
+        ):
+            scale = seqlen_float / paddle.clip(tokens_mask.sum(), min=clip_min)
+    elif tokens_mask is not None:
+        seqlen_float = tokens_mask.sum()
+    else:
+        seqlen_float = gate_prob.numel().astype(gate_prob.dtype) / num_experts
+    seqlen_float = paddle.clip(seqlen_float, min=clip_min)
+
+    if len(dispatch_mask.shape) == 2:
+        dispatch_mask = dispatch_mask.sum(0)
+    ce = dispatch_mask.astype(gate_prob.dtype).detach() / seqlen_float
+    me = paddle.sum(gate_prob, axis=0) / seqlen_float
+
+    l_aux = paddle.sum(me * ce) * num_experts
+    if use_group:
+        l_aux = l_aux / moe_k
+
+    if scale is not None:
+        # forward local me, backward global me
+        one = paddle.ones([], dtype="float32")
+        l_aux = l_aux + (scale - one) * l_aux.detach()
+
+    return l_aux, seqlen_float, ce
+
+
+def cal_aux_loss(
+    gate_prob: Tensor,
+    dispatch_mask: Tensor,
+    tokens_mask: Tensor,
+    dispatch_tokens_mask: Tensor,
+    num_experts: int,
+    use_group: bool,
+    moe_k: int,
+    clip_min: float,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        gate_prob:
+        dispatch_mask:
+        tokens_mask:
+        dispatch_tokens_mask:
+        num_experts:
+        use_group:
+        moe_k:
+        clip_min:
+
+    Returns:
+    """
+    if in_dynamic_or_pir_mode():
+        if paddle.is_compiled_with_xpu():
+            return math_cal_aux_loss(
+                gate_prob,
+                dispatch_mask,
+                tokens_mask,
+                dispatch_tokens_mask,
+                num_experts,
+                use_group,
+                moe_k,
+                clip_min,
+            )
+        else:
+            return _C_ops.cal_aux_loss(
+                gate_prob,
+                dispatch_mask,
+                tokens_mask,
+                dispatch_tokens_mask,
+                num_experts,
+                use_group,
+                moe_k,
+                clip_min,
+            )
+
+    helper = LayerHelper('cal_aux_loss', **locals())
+    l_aux_loss = helper.create_variable_for_type_inference(
+        dtype=gate_prob.dtype
+    )
+    seqlen_float = helper.create_variable_for_type_inference(
+        dtype=gate_prob.dtype
+    )
+    ce = helper.create_variable_for_type_inference(dtype=gate_prob.dtype)
+
+    inputs = {
+        'gate_prob': gate_prob,
+        'dispatch_mask': dispatch_mask,
+        'tokens_mask': tokens_mask,
+        'dispatch_tokens_mask': dispatch_tokens_mask,
+    }
+    attrs = {
+        'num_experts': num_experts,
+        'use_group': use_group,
+        'moe_k': moe_k,
+        'clip_min': clip_min,
+    }
+    outputs = {'l_aux_loss': l_aux_loss, 'seqlen_float': seqlen_float, 'ce': ce}
+    helper.append_op(
+        type='cal_aux_loss', inputs=inputs, attrs=attrs, outputs=outputs
+    )
+    return l_aux_loss, seqlen_float, ce

python/paddle/incubate/nn/functional/cross_entropy_with_softmax_bwd_w_downcast.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def cross_entropy_with_softmax_bwd_w_downcast(
+    label: Tensor,
+    softmax: Tensor,
+    loss_grad: Tensor,
+    name: str | None = None,
+) -> Tensor:
+    if in_dynamic_or_pir_mode():
+        return _C_ops.cross_entropy_with_softmax_bwd_w_downcast(
+            label,
+            softmax,
+            loss_grad,
+        )

python/paddle/incubate/nn/functional/embedding_grad_add_to.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def embedding_grad_add_to_(
+    token_indices: Tensor,
+    main_grad_: Tensor,
+    out_grad: Tensor,
+    name: str | None = None,
+) -> Tensor:
+    if in_dynamic_or_pir_mode():
+        return _C_ops.embedding_grad_add_to_(
+            token_indices,
+            main_grad_,
+            out_grad,
+        )

python/paddle/incubate/nn/functional/expand_modality_expert_id.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def expand_modality_expert_id(
+    expert_id: Tensor,
+    num_expert_per_modality: int,
+    group_size: int,
+    modality_offset: int,
+    is_group_expert: bool,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        expert_id:
+        num_expert_per_modality:
+        group_size:
+        modality_offset:
+        is_group_expert:
+
+    Returns:
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.expand_modality_expert_id(
+            expert_id,
+            num_expert_per_modality,
+            group_size,
+            modality_offset,
+            is_group_expert,
+        )
+    helper = LayerHelper('expand_modality_expert_id', **locals())
+    expert_id_out = helper.create_variable_for_type_inference(
+        dtype=expert_id.dtype
+    )
+    inputs = {'expert_id': expert_id}
+    attrs = {
+        'num_expert_per_modality': num_expert_per_modality,
+        'group_size': group_size,
+        'modality_offset': modality_offset,
+        'is_group_expert': is_group_expert,
+    }
+    helper.append_op(
+        type='expand_modality_expert_id',
+        inputs=inputs,
+        attrs=attrs,
+        outputs={'expert_id_out': expert_id_out},
+    )
+    return expert_id_out

python/paddle/incubate/nn/functional/fast_ln.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fast_ln(
+    x: Tensor,
+    scale: Tensor,
+    bias: Tensor,
+    epsilon: float = 1e-5,
+) -> tuple[Tensor, Tensor, Tensor]:
+    r"""
+    Apply Fast LayerNorm kernel.
+
+    Args:
+        x (Tensor): the input Tensor..
+        scale (Tensor): the weight Tensor to affine output.
+        bias (Tensor): the bias Tensor to affine output.
+        epsilon (float): a small float number to avoid divide 0.
+
+    Returns:
+        y: the Tensor after performing layernorm.
+        mean: the mean of input tensor
+        invvar: the invert variance(scaling factor) of y
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fast_ln(
+            x,
+            scale,
+            bias,
+            epsilon,
+        )

python/paddle/incubate/nn/functional/fast_rms_norm.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fast_rms_norm(
+    x: Tensor,
+    scale: Tensor,
+    epsilon: float = 1e-5,
+) -> tuple[Tensor, Tensor]:
+    r"""
+    Apply Fast LayerNorm kernel.
+
+    Args:
+        x (Tensor): the input Tensor..
+        scale (Tensor): the weight Tensor to affine output.
+        epsilon (float): a small float number to avoid divide 0.
+
+    Returns:
+        y: the Tensor after performing layernorm.
+        invvar: the invert variance(scaling factor) of y
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fast_rms_norm(
+            x,
+            scale,
+            epsilon,
+        )

python/paddle/incubate/nn/functional/fp8.py

@@ -0,0 +1,486 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import Tensor, _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+# special re-use of empty to reduce launch cost.
+@functools.cache
+def _empty_tensor() -> Tensor:
+    """Get tensor with no entries and no data"""
+    return Tensor()
+
+
+def fused_stack_transpose_quant(
+    x: Sequence[Tensor], transpose: bool = True
+) -> tuple[Tensor, Tensor]:
+    """
+    Fused operation that performs stacking, optional transposition, and quantization
+    on a list of bfloat16 tensors.
+
+    Args:
+        x (list[Tensor] or tuple[Tensor]): A list or tuple of bfloat16 tensors, where each tensor
+            has shape `[M, K]`. All tensors should have the same shape and dtype.
+        transpose (bool, optional): If True, applies a transpose before quantization.
+            Default is True.
+
+    Returns:
+        tuple:
+            - out (Tensor): The quantized output tensor with dtype `float8_e4m3fn`.
+            - scale (Tensor): A float32 tensor representing the quantization scale.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('BF16 requires SM80 or higher env')
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.set_device('gpu')
+
+            >>> x_vec = []
+            >>> num_experts = 1
+            >>> seq_len = 2048
+            >>> hidden_size = 128
+            >>> for _ in range(num_experts):
+            ...     x = paddle.randn([seq_len, hidden_size], dtype='bfloat16')
+            ...     x = paddle.clip(x, min=-50, max=50)
+            ...     x_vec.append(x)
+
+            >>> out, scale = F.fused_stack_transpose_quant(x_vec, transpose=True)
+            >>> print(out.shape)
+            paddle.Size([128, 2048])
+            >>> print(scale.shape)
+            paddle.Size([1, 16])
+    """
+    if in_dynamic_or_pir_mode():
+        if transpose:
+            return _C_ops.fused_stack_transpose_quant(x)
+        else:
+            return _C_ops.fused_stack_quant(x)
+
+
+def fused_act_dequant(
+    x: Tensor,
+    x_scale: Tensor,
+) -> Tensor:
+    """
+    Applies fused activation and dequantization operation to convert float8 quantized data back to bfloat16.
+
+    Args:
+        x (Tensor): Input quantized tensor with dtype float8_e4m3fn and shape [M, N]. This tensor contains the quantized
+            activations from previous layers.
+        x_scale (Tensor): Dequantization scale tensor with dtype float32 and shape [M, (N + 127) // 128] or int32 and shape [M, (N + 511) // 512].
+            Each scale value corresponds to a 128-column in the input tensor.
+
+    Returns:
+        Tensor. Dequantized output tensor with dtype bfloat16 and shape [M, N]. The values are
+            computed as input * scale for each corresponding 128-column block.
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_act_dequant(x, x_scale)
+
+
+def fused_swiglu_weighted_bwd(
+    o1: Tensor,
+    do2_s: Tensor,
+    unzipped_probs: Tensor,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """
+    Computes gradients for fused weighted SwiGLU activation function in backward pass.
+
+    Note:
+        This function performs the backward propagation for the SwiGLU (Swish-Gated Linear Unit)
+        activation with probability weighting. It computes gradients with respect to both the
+        input activations and the probability weights, while also recomputing forward pass values
+        for memory efficiency. The kernel automatically selects between vectorized and standard
+        implementations based on input dimensions.
+
+    Args:
+        o1 (Tensor): Forward pass input tensor with dtype bfloat16 and shape
+            [..., intermediate_size * 2]. The tensor is split into two halves:
+            - Left half [0:intermediate_size]: x1 values (gate inputs)
+            - Right half [intermediate_size:]: x2 values (activation inputs)
+            This is the same input used in the forward SwiGLU computation.
+        do2_s (Tensor): Upstream gradient tensor with dtype bfloat16 and shape
+            [..., intermediate_size]. Contains gradients flowing back from
+            the next layer, representing ∂L/∂output before probability weighting.
+            Each element corresponds to the gradient of one output element.
+        unzipped_probs (Tensor): Probability weighting tensor with dtype float32 and
+            shape matching the batch dimensions of o1 and do2_s
+            [...]. Each probability value was used to weight the
+            corresponding row's output in the forward pass.
+
+    Returns:
+        tuple:
+            - do1 (Tensor). Input gradients with dtype bfloat16 and shape
+              [..., intermediate_size * 2]. Layout matches o1:
+              - [0:intermediate_size]: ∂L/∂x1 (gradients w.r.t. gate inputs)
+              - [intermediate_size:]: ∂L/∂x2 (gradients w.r.t. activation inputs)
+            - probs_grad (Tensor). Probability gradients with dtype float32 and
+              shape [...]. Each element is ∂L/∂prob for the corresponding batch item,
+              computed as the sum of (∂L/∂output_i * SwiGLU_output_i) across the
+              intermediate dimension.
+            - o2_s (Tensor). Recomputed forward output with dtype bfloat16 and
+              shape [..., intermediate_size]. Contains SwiGLU(x1, x2) * prob values.
+              This avoids storing forward activations, trading computation for memory.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('BF16 requires SM80 or higher env')
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.set_device('gpu')
+
+            >>> batch_size, seq_len = 32, 128
+            >>> intermediate_size = 2048
+
+            >>> o1 = paddle.randn(
+            ...     [batch_size, seq_len, intermediate_size * 2],
+            ...     dtype='bfloat16',
+            ... )
+            >>> do2_s = paddle.randn([batch_size, seq_len, intermediate_size], dtype='bfloat16')
+            >>> expert_probs = paddle.rand([batch_size, seq_len, 1], dtype='float32')
+
+            >>> do1, probs_grad, o2_s = F.fused_swiglu_weighted_bwd(o1, do2_s, expert_probs)
+            >>> print(do1.shape)
+            paddle.Size([32, 128, 4096])
+            >>> print(probs_grad.shape)
+            paddle.Size([32, 128, 1])
+            >>> print(o2_s.shape)
+            paddle.Size([32, 128, 2048])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_swiglu_weighted_bwd(o1, do2_s, unzipped_probs)
+
+
+def fused_transpose_split_quant(
+    x, input_scales, tokens_per_expert, pow_2_scales=False
+):
+    """
+    Applies fused transpose, split, and quantization operation for Mixture of Experts (MoE) models.
+
+    Note:
+        This function performs three operations in a single optimized CUDA kernel:
+        1. Quantizes input from bfloat16 to float8_e4m3fn format using column-wise scaling
+        2. Transposes the matrix from [M, K] to [K, M] layout
+        3. Splits the transposed data across multiple experts based on token distribution
+
+    Args:
+        x (Tensor): Input tensor of shape [M, K] with dtype bfloat16, where M is the total
+            number of tokens and K is the feature dimension. M must be divisible by 128
+            for optimal performance.
+        tokens_per_expert (List[int]): List containing the number of tokens assigned to each expert.
+            Each value should be a multiple of 128 for optimal performance.
+            The sum should equal M (total tokens). Values can be 0 for
+            unused experts.
+        pow_2_scales (bool, optional): Whether to constrain quantization scales to powers of 2
+            for better hardware efficiency. If True, scales will be
+            rounded to the nearest power of 2. Default: False.
+
+    Returns:
+        tuple:
+            - outs (List[Tensor]). List of quantized and transposed output tensors, one per expert.
+              Each tensor has shape [K, tokens_per_expert[i]] and dtype float8_e4m3fn.
+              Empty tensors are included for experts with 0 tokens.
+            - scales (List[Tensor]). List of dequantization scale tensors, one per expert.
+              Each tensor has shape [K // 128, tokens_per_expert[i] // 128]
+              and dtype float32. These are the reciprocal of quantization scales.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('BF16 requires SM80 or higher env')
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.set_device('gpu')
+
+            >>> x = paddle.randn([384, 512], dtype='bfloat16')
+            >>> x = paddle.clip(x, min=-50, max=50)
+            >>> tokens_per_expert = [128, 128, 128]
+            >>> outs, scales = F.fused_transpose_split_quant(x, None, tokens_per_expert, pow_2_scales=True)
+            >>> print(outs[0].shape)
+            paddle.Size([512, 128])
+            >>> print(scales[0].shape)
+            paddle.Size([1, 512])
+    """
+    tokens_per_expert = [int(t) for t in tokens_per_expert]
+
+    if x.shape[0] == 0 or x.shape[1] == 0:
+        return [], []
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_transpose_split_quant(
+            x, input_scales, tokens_per_expert, pow_2_scales
+        )
+
+
+def fused_transpose_wlch_split_quant(
+    x: Tensor, tokens_per_expert: Sequence[int], pow_2_scales: bool = False
+) -> tuple[list[Tensor], list[Tensor]]:
+    tokens_per_expert = [int(t) for t in tokens_per_expert]
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_transpose_wlch_split_quant(
+            x, tokens_per_expert, pow_2_scales
+        )
+
+
+def fused_weighted_swiglu_act_quant(
+    x: Tensor,
+    prob: Tensor | None = None,
+    using_pow2_scaling: bool = False,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor]:
+    """
+    Applies fused weighted SwiGLU activation followed by quantization to float8_e4m3fn format.
+
+    Note:
+        This function combines four operations into a single optimized CUDA kernel:
+        1. SwiGLU activation: SwiGLU(x1, x2) = SiLU(x1) * x2 = (x1 * sigmoid(x1)) * x2
+        2. Probability weighting: multiply by optional probability factors
+        3. Activation computation: compute final activation values in float32 precision
+        4. Quantization: convert results to float8_e4m3fn with computed scaling factors
+
+        The input tensor is split into two halves along the last dimension:
+        - Left half [0, cols/2): first input to SwiGLU (gate values)
+        - Right half [cols/2, cols): second input to SwiGLU (activation values)
+
+    Args:
+        x (Tensor): Input tensor with dtype bfloat16 and shape [..., cols], where cols
+            must be even. The tensor is interpreted as two concatenated matrices:
+            gate values [0:cols/2] and activation values [cols/2:cols].
+            Typical shapes: [batch_size, sequence_length, hidden_dim] or
+            [tokens, expert_dim] in MoE scenarios.
+        prob (Tensor, optional): Probability weighting tensor with dtype float32 and
+            shape matching x's batch dimensions [...]. Each value
+            multiplies the corresponding row's activation output.
+        using_pow2_scaling (bool, optional): Whether to use power-of-2 quantization
+            scaling for hardware efficiency.
+
+    Returns:
+        tuple:
+            - out (Tensor). Quantized activation output with dtype float8_e4m3fn
+              and shape [..., cols/2]. Contains the quantized SwiGLU results.
+            - scale (Tensor). Dequantization scales with dtype float32 and shape
+              [..., (cols/2 + 127) // 128]. Each scale corresponds to a 128-element
+              block in the output tensor. To dequantize: original_value = quantized_value / scale.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('BF16 requires SM80 or higher env')
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.set_device('gpu')
+
+            >>> batch_size, seq_len, expert_dim = 32, 128, 2048
+            >>> x = paddle.randn([batch_size, seq_len, expert_dim], dtype='bfloat16')
+            >>> quantized_out, scales = F.fused_weighted_swiglu_act_quant(x)
+            >>> print(x.shape)
+            paddle.Size([32, 128, 2048])
+            >>> print(quantized_out.shape)
+            paddle.Size([4096, 1024])
+            >>> print(scales.shape)
+            paddle.Size([4096, 8])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_weighted_swiglu_act_quant(
+            x, prob, using_pow2_scaling
+        )
+
+
+def fp8_gemm_blockwise(
+    a,
+    a_decode_scale,
+    b,
+    b_decode_scale,
+    out_dtype,
+    out: Tensor | None = None,
+    bias: Tensor | None = None,
+    accumulate: bool = False,
+    use_split_accumulator: bool = True,
+    is_a_1d_scaled: bool = True,
+    is_b_1d_scaled: bool = True,
+):
+    assert bias is None, "Bias is not supported"
+
+    if bias is None:
+        bias = _empty_tensor()
+    else:
+        assert bias.dtype in (
+            paddle.float16,
+            paddle.bfloat16,
+        ), "Only fp16 and bfloat16 bias are supported."
+
+    M, K = a.shape
+    N, K_b = b.shape
+
+    if out is None:
+        out = paddle.empty((M, N), dtype=out_dtype)
+    else:
+        assert out.shape == [
+            M,
+            N,
+        ], f"Expected shape {(M, N)}, got {out.shape}"
+        assert out.is_contiguous(), "Output tensor is not contiguous."
+
+    if in_dynamic_or_pir_mode():
+        # Create workspace tensor for cuBLAS
+        workspace_size = (
+            33_554_432
+            if paddle.device.cuda.get_device_properties().major >= 9
+            else 4_194_304
+        )
+        workspace = paddle.empty([workspace_size], dtype=paddle.uint8)
+        transa, transb = True, False
+        grad = False
+        math_sm_count = 112
+
+        # Call the C++ operator - it returns (output, pre_gelu_out, workspace_out)
+        output, _, _ = _C_ops.fp8_gemm_blockwise_(
+            b,
+            b_decode_scale,
+            a,
+            a_decode_scale,
+            out,
+            bias,
+            _empty_tensor(),
+            workspace,
+            transa,
+            transb,
+            grad,
+            accumulate,
+            use_split_accumulator,
+            math_sm_count,
+            is_b_1d_scaled,
+            is_a_1d_scaled,
+        )
+        return output
+
+
+def fp8_quant_blockwise(
+    x: Tensor,
+    epsilon: float = 0.0,
+    input_transpose: bool = False,
+    output_scale_transpose: bool = True,
+    return_transpose_only: bool = False,
+    using_pow2_scale: bool = True,
+    using_ue8m0_scale: bool = False,
+    quant_method: str = "1x128",
+    output_type: str = "e4m3",
+    name: str | None = None,
+):
+    """
+    Applies blockwise FP8 quantization to input tensor with flexible configuration options.
+    Note:
+        This function performs blockwise quantization from higher precision formats (typically bfloat16)
+        to FP8 format (float8_e4m3fn by default). The quantization is performed in blocks (128x128 or 1x128)
+        for better numerical stability and hardware efficiency.
+    Args:
+        x (Tensor): Input tensor to be quantized. Typically has dtype bfloat16 and shape [M, N].
+        epsilon (float, optional): Small constant added to avoid division by zero when computing scales.
+            Default: 0.0.
+        input_transpose (bool, optional): Whether to transpose the input before quantization.
+            If True, input shape [M, N] becomes [N, M]. Default: False.
+        output_scale_transpose (bool, optional): Whether to transpose the output scale tensor.
+            Default: True.
+        return_transpose_only (bool, optional): If True and input_transpose is True, returns only
+            the transposed quantized output and scale. Default: False.
+        using_pow2_scale (bool, optional): Whether to use power-of-2 quantization scaling for
+            better hardware efficiency. Default: True.
+        using_ue8m0_scale (bool, optional): Whether to use unsigned 8-bit with mantissa 0 scaling format.
+            If True, the output scale tensor has dtype int32, where each element contains 4 packed ue8m0 scales.
+            If False, the output scale tensor has dtype float32.
+            Default: False.
+        quant_method (str, optional): Quantization block size method. Options: "1x128" or "128x128".
+            "1x128" uses 1x128 blocks, "128x128" uses 128x128 blocks. Default: "1x128".
+        output_type (str, optional): Output FP8 format. Currently only "e4m3" (float8_e4m3fn) is supported.
+            Default: "e4m3".
+        name (str, optional): Name for the operation. Default: None.
+    Returns:
+        tuple: The return value depends on the configuration:
+            - If not input_transpose: returns (quantized_output, scale)
+            - If return_transpose_only and input_transpose: returns (transposed_quantized_output, transposed_scale)
+            - Otherwise: returns (quantized_output, scale, transposed_quantized_output, transposed_scale)
+        Where:
+            - quantized_output (Tensor): Quantized output tensor with dtype float8_e4m3fn.
+            - scale (Tensor): Dequantization scale tensor with dtype float32.
+            - transposed_quantized_output (Tensor): Transposed quantized output (if input_transpose).
+            - transposed_scale (Tensor): Transposed scale tensor (if input_transpose).
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('BF16 requires SM80 or higher env')
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.set_device('gpu')
+            >>> x = paddle.randn([1024, 512], dtype='bfloat16')
+            >>> x = paddle.clip(x, min=-50, max=50)
+            # Basic quantization
+            >>> quantized, scale = F.fp8_quant_blockwise(x)
+            >>> print(quantized.shape)
+            paddle.Size([1024, 512])
+            >>> print(scale.shape)
+            paddle.Size([1024, 4])
+            # With transpose
+            >>> quantized, scale, transposed_quantized, transposed_scale = F.fp8_quant_blockwise(
+            ...     x, input_transpose=True, return_transpose_only=False
+            ... )
+            >>> print(transposed_quantized.shape)
+            paddle.Size([512, 1024])
+    """
+
+    if quant_method == "1x128":
+        using_1x128 = True
+    elif quant_method == "128x128":
+        using_1x128 = False
+    else:
+        raise ValueError("Unsupported quantization method")
+
+    if output_type == "e4m3":
+        using_e5m2 = False
+    else:
+        raise ValueError("Unsupported output type")
+
+    if in_dynamic_or_pir_mode():
+        x_fp8, scale, x_fp8_t, scale_t = _C_ops.fp8_quant_blockwise(
+            x,
+            epsilon,
+            using_1x128,
+            input_transpose,
+            output_scale_transpose,
+            return_transpose_only,
+            using_e5m2,
+            using_pow2_scale,
+            using_ue8m0_scale,
+        )
+        # Aligned with kitchen's logic
+        if not input_transpose:
+            return x_fp8, scale
+        elif return_transpose_only:
+            return x_fp8_t, scale_t
+        else:
+            return x_fp8, scale, x_fp8_t, scale_t

python/paddle/incubate/nn/functional/fused_bias_act.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2024 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+
+def fused_bias_act(
+    x: Tensor,
+    bias: Tensor | None = None,
+    dequant_scales: Tensor | None = None,
+    shift: Tensor | None = None,
+    smooth: Tensor | None = None,
+    act_method: str = "gelu",
+    compute_dtype: str = "default",
+    quant_scale: float = -1,
+    quant_round_type: int = 0,
+    quant_max_bound: float = 0,
+    quant_min_bound: float = 0,
+) -> Tensor:
+    """
+    Applies fused_bias_act kernel
+
+    Args:
+        x (Tensor): the input Tensor.
+        bias (Tensor, optional): the input bias Tensor. If it is None, no bias addition would be performed. Otherwise, the bias will be added before activation function. Default: None.
+        dequant_scales (Tensor, optional): the dequantization scale tensor, If it is None, no dequantization will be performed. Default: None.
+        shift (Tensor, optional): the shift tensor, used to shift the input tensor before activation function. If None, no translation will be performed. Default: None.
+        smooth (Tensor, optional): the smooth tensor, used to smooth the input tensor before activation function. If None, no smoothing processing will be performed. Default: None.
+        act_method (Str, optional): the activation method, specify the activation function to be used. Default: gelu.
+        compute_dtype (Str, optional): a compute dtype, is used to represent the input data type. Default is "default", which means compute dtype is determined by input dtype.
+        quant_scale (Float, optional): the quant scale. Default: -1.
+        quant_round_type (Int, optional): the quant round type, if 0 is set, value will be rounding to nearest ties to even. If 1 is set, value will be rounding to nearest ties away from zero. Default: 0.
+        quant_max_bound (Float, optional): the max bound of float type to int type. Default: 0.
+        quant_min_bound (Float, optional): the min bound of float type to int type. Default: 0.
+
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_bias_act
+
+            >>> paddle.set_device('gpu')
+            >>> x = paddle.randn([3, 5])
+            >>> bias = paddle.randn([5])
+            >>> out = fused_bias_act(x, bias)
+            >>> print(out.shape)
+            paddle.Size([3, 5])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_bias_act(
+            x,
+            bias,
+            dequant_scales,
+            shift,
+            smooth,
+            act_method,
+            compute_dtype,
+            quant_scale,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+        )
+    helper = LayerHelper("fused_bias_act")
+    if x.dtype == "int32":
+        if compute_dtype == "bf16":
+            dtype = "uint16"
+        elif compute_dtype == "fp16":
+            dtype = "float16"
+        elif compute_dtype == "fp32":
+            dtype = "float32"
+        out = helper.create_variable_for_type_inference(dtype=dtype)
+    else:
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+    inputs = {}
+    inputs["x"] = x
+    if bias is not None:
+        inputs["bias"] = bias
+    if dequant_scales is not None:
+        inputs["dequant_scales"] = dequant_scales
+
+    if shift is not None:
+        inputs["shift"] = shift
+
+    if smooth is not None:
+        inputs["smooth"] = smooth
+
+    attrs = {
+        "act_method": act_method,
+        "compute_dtype": compute_dtype,
+        "quant_scale": quant_scale,
+        "quant_round_type": quant_round_type,
+        "quant_max_bound": quant_max_bound,
+        "quant_min_bound": quant_min_bound,
+    }
+
+    helper.append_op(
+        type="fused_bias_act",
+        inputs=inputs,
+        outputs={"out": out},
+        attrs=attrs,
+    )
+    return out

python/paddle/incubate/nn/functional/fused_dot_product_attention.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from paddle import Tensor, _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+
+def cudnn_flash_attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    bias: Tensor | None = None,
+    cu_seqlen_q: Tensor | None = None,
+    cu_seqlen_k: Tensor | None = None,
+    scaling_factor: float = 1.0,
+    dropout_prob: float = 0.0,
+    training: bool = True,
+    mask_type: str | None = None,
+    bias_type: str | None = None,
+    name: str | None = None,
+) -> Tensor:
+    r"""
+    Fused Dot Product Attention. This is a fusion operator to compute scaled dot product attention in transformer
+    model architecture. This operator only supports running on Ampere and Hopper GPU and need cudnn version >= 8906.
+
+    Args:
+        q (Tensor): The query tensor. The data type is bfloat16, float16.
+        k (Tensor): The key tensor. The data type is bfloat16, float16.
+        v (Tensor): The value tensor. The data type is bfloat16, float16.
+        bias (Tensor, optional): The bias tensor. The data type needs to be the same as `q`, `k`, `v`.
+            The shape of the bias tensor should be [b,1,s,s] or [b,h,s,s] or [1,1,s,s] or [1, h, s, s].
+        cu_seqlen_q (Tensor, optional): The cu_seqlen_q tensor. The data type is int32.
+        cu_seqlen_k (Tensor, optional): The cu_seqlen_k tensor. The data type is int32.
+        scaling_factor (float): The scaling factor for the attention scores.
+        dropout_prob (float): The dropout probability.
+        training (bool): A flag indicating whether it is in train phrase or not.
+        mask_type (str, optional): The mask type. It can be 'none', 'padding', 'causal', 'paddle_causal'. Default is None.
+        bias_type (str, optional): The bias type. It can be 'none', 'pre_scale_bias', 'post_scale_bias'. Default is None.
+
+
+    Returns:
+        A Tensor representing the fused dot product attention, has same shape and data type as `q` .
+
+    Note:
+        This API is provides the full functionality of the cuDNN flash attention. Such as dbias, padding_causal mask, etc.
+    """
+
+    if mask_type is None:
+        mask_type = "none"
+    if bias_type is None:
+        bias_type = "none"
+    assert mask_type in [
+        'none',
+        'padding',
+        'causal',
+        'paddle_causal',
+    ], "mask_type should be 'none', 'padding', 'causal', 'paddle_causal'"
+    assert bias_type in [
+        'none',
+        'pre_scale_bias',
+        'post_scale_bias',
+    ], "bias_type should be 'none', 'pre_scale_bias', 'post_scale_bias'"
+
+    if in_dynamic_or_pir_mode():
+        out, _, _ = _C_ops.fused_dot_product_attention(
+            q,
+            k,
+            v,
+            bias,
+            cu_seqlen_q,
+            cu_seqlen_k,
+            scaling_factor,
+            dropout_prob,
+            training,
+            mask_type,
+            bias_type,
+        )
+        return out
+    else:
+        helper = LayerHelper('fused_dot_product_attention', **locals())
+        out = helper.create_variable_for_type_inference(dtype=q.dtype)
+        softmax_out = helper.create_variable_for_type_inference(
+            dtype="float", stop_gradient=True
+        )
+        rng_state = helper.create_variable_for_type_inference(
+            dtype='int64', stop_gradient=True
+        )
+
+        attrs = {
+            "scaling_factor": scaling_factor,
+            "dropout_probability": dropout_prob,
+            "is_training": training,
+            "mask_type_str": mask_type,
+            "bias_type_str": bias_type,
+        }
+        helper.append_op(
+            type='fused_dot_product_attention',
+            inputs={
+                'q': q,
+                'k': k,
+                'v': v,
+                'bias': bias,
+                'cu_seqlen_q': cu_seqlen_q,
+                'cu_seqlen_k': cu_seqlen_k,
+            },
+            outputs={
+                'out': [out],
+                'softmax_out': [softmax_out],
+                'rng_state': [rng_state],
+            },
+            attrs=attrs,
+        )
+        return out
+
+
+def fused_dot_product_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Tensor | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scaling_factor: float | None = None,
+    training: bool = True,
+    name: str | None = None,
+) -> Tensor:
+    r"""
+    Fused Dot Product Attention. This is a fusion operator to compute scaled dot product attention in transformer
+    model architecture. This operator only supports running on Ampere and Hopper GPU and need cudnn version >= 8906.
+
+    Args:
+        query(Tensor): The query tensor in the Attention module.
+                        4-D tensor with shape:
+                        [batch_size, seq_len, num_heads, head_dim].
+                        The dtype can be float16 or bfloat16.
+        key(Tensor): The key tensor in the Attention module.
+                        4-D tensor with shape:
+                        [batch_size, seq_len, num_heads, head_dim].
+                        The dtype can be float16 or bfloat16.
+        value(Tensor): The value tensor in the Attention module.
+                        4-D tensor with shape:
+                        [batch_size, seq_len, num_heads, head_dim].
+                        The dtype can be float16 or bfloat16.
+        attn_mask(Tensor,optional): A float mask of the same type as query,
+                key, value that is added to the attention score.
+        dropout_p (float): The dropout probability.
+        is_causal (bool): A flag indicating whether it is causal masking or not. If True, the mask will be ignored.
+        scaling_factor (float): The scaling factor for the attention scores.
+        training (bool): A flag indicating whether it is in train phrase or not.
+        name(str|None, optional): The default value is None. Normally there is no need for user
+                        to set this property. For more information, please refer to
+                        :ref:`api_guide_Name`.
+
+    Returns:
+        A Tensor representing the fused dot product attention, has same shape and data type as `q` .
+
+    Note:
+        This API is designed to be aligned with `nn.functional.scaled_dot_product_attention` API. So the
+            arguments are almost the same as `nn.functional.scaled_dot_product_attention`. This difference
+            is that `nn.functional.scaled_dot_product_attention` calls the open source flash attention
+            (https://github.com/Dao-AILab/flash-attention). While this API calls the kernel implemented
+            by cuDNN, which achieves better performance on the latest GPU architectures (Hopper and After).
+
+        The mask is passed as a post scale bias in this API. So this mask can be a arbitrary mask, not just padding mask.
+    """
+
+    if is_causal:
+        assert attn_mask is None, "mask must be None when is_causal is True"
+        mask_type = "causal"
+        bias_type = "none"
+    elif attn_mask is not None:
+        mask_type = "none"
+        bias_type = "post_scale_bias"  # pass mask as a post scale bias
+    else:
+        mask_type = "none"
+        bias_type = "none"
+
+    if attn_mask is not None:
+        assert attn_mask.dtype == query.dtype, (
+            "attn_mask dtype should be the same as qkv dtype"
+        )
+
+    cu_seqlen_q = None
+    cu_seqlen_k = None
+    if scaling_factor is None:
+        head_dim = query.shape[3]
+        scaling_factor = head_dim**-0.5
+    if in_dynamic_or_pir_mode():
+        out, _, _ = _C_ops.fused_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask,
+            cu_seqlen_q,
+            cu_seqlen_k,
+            scaling_factor,
+            dropout_p,
+            training,
+            mask_type,
+            bias_type,
+        )
+        return out
+    else:
+        helper = LayerHelper('fused_dot_product_attention', **locals())
+        out = helper.create_variable_for_type_inference(dtype=query.dtype)
+        softmax_out = helper.create_variable_for_type_inference(dtype="float")
+        rng_state = helper.create_variable_for_type_inference(dtype='int64')
+
+        attrs = {
+            "scaling_factor": scaling_factor,
+            "dropout_probability": dropout_p,
+            "is_training": training,
+            "mask_type_str": mask_type,
+            "bias_type_str": bias_type,
+        }
+        helper.append_op(
+            type='fused_dot_product_attention',
+            inputs={
+                'q': query,
+                'k': key,
+                'v': value,
+                'bias': attn_mask,
+                'cu_seqlen_q': cu_seqlen_q,
+                'cu_seqlen_k': cu_seqlen_k,
+            },
+            outputs={
+                'out': [out],
+                'softmax_out': [softmax_out],
+                'rng_state': [rng_state],
+            },
+            attrs=attrs,
+        )
+        return out

python/paddle/incubate/nn/functional/fused_dropout_add.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+import warnings
+
+import paddle
+from paddle import _C_ops
+from paddle.base import core
+from paddle.common_ops_import import default_main_program
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+flag = [False]
+
+
+def paddle_dropout_add(x, y, p=0.5, training=True, mode="upscale_in_train"):
+    tmp = paddle.nn.functional.dropout(x, p, training=training, mode=mode)
+    return tmp + y
+
+
+def fused_dropout_add(
+    x: Tensor,
+    y: Tensor,
+    p: float = 0.5,
+    training: bool = True,
+    mode: Literal[
+        'upscale_in_train', 'downscale_in_infer'
+    ] = 'upscale_in_train',
+    name: str | None = None,
+) -> Tensor:
+    r"""
+    Fused Dropout and Add.
+
+    Args:
+        x (Tensor): The input tensor. The data type is bfloat16, float16, float32 or float64.
+        y (Tensor): The input tensor. The data type is bfloat16, float16, float32 or float64.
+
+        p (float|int, optional): Probability of setting units to zero. Default: 0.5.
+        training (bool, optional): A flag indicating whether it is in train phrase or not. Default: True.
+        mode(str, optional): ['upscale_in_train'(default) | 'downscale_in_infer'].
+
+            1. upscale_in_train (default), upscale the output at training time
+
+                - train: :math:`out = x \times \frac{mask}{(1.0 - dropout\_prob)} + y`
+                - inference: :math:`out = x + y`
+
+            2. downscale_in_infer, downscale the output at inference
+
+                - train: :math:`out = input \times mask + y`
+                - inference: :math:`out = input \times (1.0 - dropout\_prob) + y`
+
+        name (str, optional): Name for the operation, Default: None. For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Tensor representing the fused dropout and add, has same shape and data type as `x` .
+
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_dropout_add
+
+            >>> paddle.set_device('gpu')
+            >>> paddle.seed(2023)
+            >>> x = paddle.randn([4, 10], dtype="float32")
+            >>> y = paddle.randn([4, 10], dtype="float32")
+            >>> out = fused_dropout_add(x, y, p=0.5)
+            >>> print(out)
+            Tensor(shape=[4, 10], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+            [[-0.49133155,  0.53819323, -2.58393312,  0.06336236, -1.09908366,
+               0.22085167,  2.19751787,  0.05034769,  0.53417486,  0.84864247],
+             [ 0.78248203, -1.59652555, -0.14399840, -0.77985179, -0.17006736,
+              -0.30991879, -0.36593807, -0.51025450,  1.46401680,  0.61627960],
+             [ 4.50472546, -0.48472026,  0.60729283,  0.33509624, -0.25593102,
+              -1.45173049,  1.06727099,  0.00440830, -0.77340341,  0.67393088],
+             [ 1.29453969,  0.07568165,  0.71947742, -0.71768606, -2.57172823,
+               1.89179027,  3.26482797,  1.10493207, -1.04569530, -1.04862499]])
+    """
+    if not flag[0]:
+        flag[0] = True
+        warnings.warn(
+            "Currently, fused_dropout_add maybe has precision problem, so it falls back to dropout + add. "
+        )
+    return paddle_dropout_add(x, y, p, training=training, mode=mode)
+
+    if isinstance(p, (int, float)):
+        # fast return for p == 0
+        if p == 0:
+            return x + y
+        elif p < 0 or p > 1:
+            raise ValueError("p argument should between 0 and 1")
+    if mode not in ('downscale_in_infer', 'upscale_in_train'):
+        raise ValueError(
+            "mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
+        )
+    seed = None
+    if in_dynamic_or_pir_mode():
+        if default_main_program().random_seed != 0:
+            seed = default_main_program().random_seed
+        out, seed_offset = _C_ops.fused_dropout_add(
+            x,
+            y,
+            None,
+            p,
+            not training,
+            mode,
+            seed if seed is not None else 0,
+            seed is not None,
+        )
+        return out
+    else:
+        helper = LayerHelper('fused_dropout_add', **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+        seed_offset = helper.create_variable_for_type_inference(
+            dtype=core.VarDesc.VarType.INT64, stop_gradient=True
+        )
+
+        def get_attrs(prog, dropout_prob, is_test, seed):
+            if (seed is None or seed == 0) and prog.random_seed != 0:
+                seed = prog.random_seed
+            attrs = {
+                'p': dropout_prob,
+                'is_test': is_test,
+                'mode': mode,
+                'seed': seed if seed is not None else 0,
+                'fix_seed': seed is not None,
+            }
+            return attrs
+
+        attrs = get_attrs(helper.main_program, p, not training, seed)
+
+        helper.append_op(
+            type='fused_dropout_add',
+            inputs={'x': x, 'y': y, 'seed_tensor': None},
+            outputs={'out': [out], 'seed_offset': [seed_offset]},
+            attrs=attrs,
+        )
+        return out

python/paddle/incubate/nn/functional/fused_gate_attention.py

@@ -0,0 +1,174 @@
+#   Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _legacy_C_ops
+from paddle.framework import in_dynamic_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fused_gate_attention(
+    query: Tensor,
+    key: Tensor | None = None,
+    query_weight: Tensor | None = None,
+    key_weight: Tensor | None = None,
+    value_weight: Tensor | None = None,
+    qkv_weight: Tensor | None = None,
+    gate_linear_weight: Tensor | None = None,
+    gate_linear_bias: Tensor | None = None,
+    out_linear_weight: Tensor | None = None,
+    out_linear_bias: Tensor | None = None,
+    nonbatched_bias: Tensor | None = None,
+    attn_mask: Tensor | None = None,
+    has_gating: bool = True,
+    merge_qkv: bool = True,
+    use_flash_attn: bool = False,
+) -> Tensor:
+    r"""
+    Attention maps queries and a set of key-value pairs to outputs, and
+    Gate Attention performs multiple parallel attention to jointly attending
+    to information from different representation subspaces. This API only
+    support self_attention. The pseudo code is as follows:
+
+    .. code-block:: text
+
+        c = c ** (-0.5)
+        q = paddle.einsum('nbqa,ahc->nbqhc', q_data, query_w) * c
+        k = paddle.einsum('nbka,ahc->nbkhc', m_data, key_w)
+        v = paddle.einsum('nbka,ahc->nbkhc', m_data, value_w)
+        logits = paddle.einsum('nbqhc,nbkhc->nbhqk', q, k) + bias
+
+        if nonbatched_bias is not None:
+            logits += paddle.unsqueeze(nonbatched_bias, axis=1)
+
+        weights = paddle.nn.functional.softmax(logits)
+        weighted_avg = paddle.einsum('nbhqk,nbkhc->nbqhc', weights, v)
+
+        if has_gating:
+            gate_values = paddle.einsum('nbqc,chv->nbqhv', q_data, gating_w) + gating_b
+            gate_values = paddle.nn.functional.sigmoid(gate_values)
+            weighted_avg *= gate_values
+
+        output = paddle.einsum('nbqhc,hco->nbqo', weighted_avg, output_w) + output_b
+
+
+    Args:
+        query (Tensor): The input query tensor. The shape is [batch_size, msa_len, res_len, q_dim].
+        key (Tensor, optional): The input key tensor, which can be set when
+            merge_qkv is False. The shape is [batch_size, msa_len, m_size, kv_dim]. Default None.
+        query_weight (Tensor, optional): The weight of query linear, which should be set when input
+            key is not None. The shape is [q_dim, num_heads, head_dim]. Default None.
+        key_weight (Tensor, optional): The weight of key linear, which should be set when input key
+            is not None. The shape is [kv_dim, num_heads, head_dim]. Default None.
+        value_weight (Tensor, optional): The weight of value linear, which should be set when input
+            key is not None. The shape is [kv_dim, num_heads, head_dim]. Default None.
+        qkv_weight (Tensor, optional): The weight of qkv linear, which should be set when merge_qkv
+            is True. The shape is [3, num_heads, head_dim, q_dim]. Default None.
+        gate_linear_weight (Tensor, optional): The weight of gating linear, which should be set when
+            has_gating is True. The shape is [q_dim, num_heads, head_dim]. Default None.
+        gate_linear_bias (Tensor, optional): The bias of gating linear, which should be set when
+            has_gating is True. The shape is [num_heads, head_dim]. Default None.
+        out_linear_weight (Tensor, optional): The weight of output linear. The shape is [num_heads, head_dim, q_dim]. Default None.
+        out_linear_bias (Tensor): The bias of output linear, the shape is [q_dim]. Default None.
+        nonbatched_bias (Tensor, optional): The extra bias. The shape is [batch_size, 1, num_heads, res_len, m_size]. Default None.
+        attn_mask (Tensor, optional):  The attention mask. The shape is [batch_size, msa_len, 1, 1, res_len]. Default None.
+        has_gating (bool, optional): Whether has the gating linear. Default True.
+        merge_qkv (bool, optional): Whether has the gating linear. Default True.
+        use_flash_attn (bool, optional): Whether use flash-attention to speedup. Default False.
+
+    Returns:
+        Tensor: The output Tensor, the data type and shape is same as `query`.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+
+            >>> # batch_size = 2
+            >>> # msa_len = 4
+            >>> # res_len = 2
+            >>> # q_dim = 4
+            >>> # num_heads = 8
+            >>> # head_dim = 4
+            >>> # m_size = res_len (when merge_qkv is True)
+
+            >>> # query: [batch_size, msa_len, res_len, q_dim]
+            >>> query = paddle.rand(shape=[2, 4, 2, 4], dtype="float32")
+
+            >>> # qkv_weight:  [3, n_heads, head_dim, q_dim]
+            >>> qkv_weight = paddle.rand(shape=[3, 8, 4, 4], dtype="float32")
+
+            >>> # nonbatched_bias: [batch_size, 1, num_heads, res_len, m_size]
+            >>> nonbatched_bias = paddle.rand(shape=[2, 1, 8, 2, 2], dtype="float32")
+
+            >>> # attn_mask: [batch_size, msa_len, 1, 1, m_size]
+            >>> attn_mask = paddle.rand(shape=[2, 4, 1, 1, 2], dtype="float32")
+
+            >>> # gate_linear_weight: [q_dim, num_heads, head_dim]
+            >>> gate_linear_weight = paddle.rand(shape=[4, 8, 4], dtype="float32")
+            >>> # gate_bias: [num_heads, head_dim]
+            >>> gate_linear_bias = paddle.rand(shape=[8, 4], dtype="float32")
+
+            >>> # out_linear_weight: [num_heads, head_dim, q_dim]
+            >>> out_linear_weight = paddle.rand(shape=[8, 4, 4], dtype="float32")
+            >>> # out_linear_bias: [q_dim]
+            >>> out_linear_bias = paddle.rand(shape=[4], dtype="float32")
+
+            >>> # output: [batch_size, msa_len, res_len, q_dim]
+            >>> output = F.fused_gate_attention(
+            ...     query=query,
+            ...     qkv_weight=qkv_weight,
+            ...     gate_linear_weight=gate_linear_weight,
+            ...     gate_linear_bias=gate_linear_bias,
+            ...     out_linear_weight=out_linear_weight,
+            ...     out_linear_bias=out_linear_bias,
+            ...     nonbatched_bias=nonbatched_bias,
+            ...     attn_mask=attn_mask,
+            ...     has_gating=True,
+            ...     merge_qkv=True,
+            ... )
+            >>> print(output.shape)
+            paddle.Size([2, 4, 2, 4])
+
+    """
+    if in_dynamic_mode():
+        _, _, _, _, _, _, _, _, out = _legacy_C_ops.fused_gate_attention(
+            query,
+            key,
+            query_weight,
+            key_weight,
+            value_weight,
+            qkv_weight,
+            nonbatched_bias,
+            attn_mask,
+            gate_linear_weight,
+            gate_linear_bias,
+            out_linear_weight,
+            out_linear_bias,
+            'has_gating',
+            has_gating,
+            'merge_qkv',
+            merge_qkv,
+            "use_flash_attn",
+            use_flash_attn,
+        )
+        return out

python/paddle/incubate/nn/functional/fused_layer_norm.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import paddle
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@overload
+def fused_layer_norm(
+    x: Tensor,
+    norm_weight: Tensor,
+    norm_bias: Tensor,
+    epsilon: float,
+    residual_alpha: float = ...,
+    begin_norm_axis: int = ...,
+    bias: Tensor | None = ...,
+    residual: None = ...,
+    quant_scale: float = ...,
+    quant_round_type: float = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> Tensor: ...
+
+
+@overload
+def fused_layer_norm(
+    x: Tensor,
+    norm_weight: Tensor,
+    norm_bias: Tensor,
+    epsilon: float,
+    residual_alpha: float = ...,
+    begin_norm_axis: int = ...,
+    bias: Tensor | None = ...,
+    residual: Tensor = ...,
+    quant_scale: float = ...,
+    quant_round_type: float = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+def fused_layer_norm(
+    x,
+    norm_weight,
+    norm_bias,
+    epsilon,
+    residual_alpha=1.0,
+    begin_norm_axis=1,
+    bias=None,
+    residual=None,
+    quant_scale=-1,
+    quant_round_type=0,
+    quant_max_bound=0,
+    quant_min_bound=0,
+):
+    r"""
+    Apply Fused LayerNorm kernel. Also support LayerNorm(bias + residual_alpha * residual + x) fused pattern.
+
+    when norm_weight and norm_bias is None, it return fused (bias + residual_alpha * residual + x)
+
+    Args:
+        x (Tensor): the input Tensor..
+        norm_weight (Tensor): the weight Tensor to affine output.
+        norm_bias (Tensor): the bias Tensor to affine output.
+        epsilon (float): a small float number to avoid divide 0.
+        residual_alpha (float): a scale factor for residual. default is 1.
+        begin_norm_axis (int): the begin axis to normalize. default is 1.
+        bias (optional|Tensor): the previous layers's bias to fused.
+        residual (optional|Tensor): the residual input to fused.
+        quant_scale (float): the quant scale.
+        quant_round_type (float): the quant round type.
+        quant_max_bound (float): the quant max bound to clip.
+        quant_min_bound (float): the quant min bound to clip.
+
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+
+            >>> paddle_x = paddle.cast(paddle.randn(shape=[32, 256]), dtype=paddle.float16)
+            >>> paddle_weight = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float32)
+            >>> paddle_bias = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float32)
+            >>> epsilon = 1e-6
+            >>> paddle_layernorm = paddle.incubate.nn.functional.fused_layer_norm(paddle_x, paddle_weight, paddle_bias, epsilon, 1)
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_bias_residual_layernorm(
+            x,
+            bias,
+            residual,
+            norm_weight,
+            norm_bias,
+            epsilon,
+            residual_alpha,
+            begin_norm_axis,
+            quant_scale,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+        )
+    # static mode
+    helper = LayerHelper('fused_layernorm', **locals())
+    out = None
+    if quant_scale <= 0:
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    else:
+        out = helper.create_variable_for_type_inference(dtype=paddle.int8)
+    outputs_dict = {}
+    outputs_dict['out'] = out
+    outputs_dict['mean'] = helper.create_variable_for_type_inference(
+        dtype=paddle.float32
+    )
+    outputs_dict['variance'] = helper.create_variable_for_type_inference(
+        dtype=paddle.float32
+    )
+
+    residual_out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    outputs_dict['residual_out'] = residual_out
+
+    inputs = {'x': x}
+    if norm_weight is not None:
+        inputs['norm_weight'] = norm_weight
+    if norm_bias is not None:
+        inputs['norm_bias'] = norm_bias
+    if residual is not None:
+        inputs['residual'] = residual
+    if bias is not None:
+        inputs['bias'] = bias
+
+    helper.append_op(
+        type='fused_bias_residual_layernorm',
+        inputs=inputs,
+        attrs={
+            "epsilon": epsilon,
+            "residual_alpha": residual_alpha,
+            "begin_norm_axis": begin_norm_axis,
+            "quant_scale": quant_scale,
+            "quant_round_type": quant_round_type,
+            "quant_max_bound": quant_max_bound,
+            "quant_min_bound": quant_min_bound,
+        },
+        outputs=outputs_dict,
+    )
+    return (out, residual_out, outputs_dict['mean'], outputs_dict['variance'])

python/paddle/incubate/nn/functional/fused_matmul_bias.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2022 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from paddle import _C_ops, _legacy_C_ops
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import (
+    in_dynamic_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
+from paddle.tensor.linalg import matmul
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fused_matmul_bias(
+    x: Tensor,
+    y: Tensor,
+    bias: Tensor | None = None,
+    transpose_x: bool = False,
+    transpose_y: bool = False,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Applies matrix multiplication of two tensors and then bias addition if provided.
+    This method requires CUDA version >= 11.6.
+
+    Args:
+        x (Tensor): the first input Tensor to be multiplied.
+        y (Tensor): the second input Tensor to be multiplied. Its rank must be 2.
+        bias (Tensor, optional): the input bias Tensor. If it is None, no bias addition would
+            be performed. Otherwise, the bias is added to the matrix multiplication result. Default: None.
+        transpose_x (bool, optional): Whether to transpose :math:`x` before multiplication. Default: False.
+        transpose_y (bool, optional): Whether to transpose :math:`y` before multiplication. Default: False.
+        name (str, optional): For detailed information, please refer to
+            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('fused_gemm_epilogue is only supported when CUDA version >= 11.6')
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_matmul_bias
+
+            >>> paddle.set_device('gpu')
+            >>> x = paddle.randn([3, 5])
+            >>> y = paddle.randn([4, 5])
+            >>> bias = paddle.randn([5])
+            >>> out = fused_matmul_bias(x, y, bias)
+            >>> print(out.shape)
+            paddle.Size([3, 5])
+    """
+    if bias is None:
+        return matmul(x, y, transpose_x, transpose_y, name)
+    if in_dynamic_or_pir_mode():
+        out, _ = _C_ops.fused_gemm_epilogue(
+            x, y, bias, transpose_x, transpose_y, "none"
+        )
+        return out
+
+    helper = LayerHelper('fused_matmul_bias', **locals())
+    out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    helper.append_op(
+        type='fused_gemm_epilogue',
+        inputs={'X': x, 'Y': y, 'Bias': bias},
+        outputs={'Out': out},
+        attrs={'trans_x': transpose_x, 'trans_y': transpose_y},
+    )
+    return out
+
+
+def fused_linear(
+    x: Tensor,
+    weight: Tensor,
+    bias: Tensor | None = None,
+    transpose_weight: bool = False,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Fully-connected linear transformation operator. This method requires CUDA version >= 11.6.
+
+    Args:
+        x (Tensor): the input Tensor to be multiplied.
+        weight (Tensor): the weight Tensor to be multiplied. Its rank must be 2.
+        bias (Tensor, optional): the input bias Tensor. If it is None, no bias addition would
+            be performed. Otherwise, the bias is added to the matrix multiplication result. Default: None.
+        transpose_weight (bool, optional): Whether to transpose :math:`weight` before multiplication. Default: False.
+        name (str, optional): For detailed information, please refer to
+            :ref:`api_guide_Name` . Usually name is no need to set and None by default.
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('fused_gemm_epilogue is only supported when CUDA version >= 11.6')
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_linear
+
+            >>> paddle.set_device('gpu')
+            >>> x = paddle.randn([3, 4])
+            >>> weight = paddle.randn([4, 5])
+            >>> bias = paddle.randn([5])
+            >>> out = fused_linear(x, weight, bias)
+            >>> print(out.shape)
+            paddle.Size([3, 5])
+    """
+    return fused_matmul_bias(x, weight, bias, False, transpose_weight, name)
+
+
+def fused_linear_activation(
+    x: Tensor,
+    y: Tensor,
+    bias: Tensor,
+    trans_x: bool = False,
+    trans_y: bool = False,
+    activation: Literal['gelu', 'relu'] | None = None,
+) -> Tensor:
+    """
+    Fully-connected linear and activation transformation operator. This method requires CUDA version >= 11.6.
+
+    Args:
+        x (Tensor): the input Tensor to be multiplied.
+        y (Tensor): the weight Tensor to be multiplied. Its rank must be 2.
+        bias (Tensor): the input bias Tensor, the bias is added to the matrix multiplication result.
+        trans_x (bool, optional): Whether to transpose :math:`x` before multiplication.
+        trans_y (bool, optional): Whether to transpose :math:`y` before multiplication.
+        activation (str, optional): Activation function, Currently, the available activation functions are
+            limited to "gelu" (Gaussian Error Linear Unit) and "relu" (Rectified Linear Unit).
+            These activation functions are applied to the output of the bias add. Default: None.
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('fused_gemm_epilogue is only supported when CUDA version >= 11.6')
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import (
+            ...     fused_linear_activation,
+            ... )
+
+            >>> paddle.set_device('gpu')
+            >>> x = paddle.randn([3, 4])
+            >>> weight = paddle.randn([4, 5])
+            >>> bias = paddle.randn([5])
+            >>> out = fused_linear_activation(x, weight, bias)
+            >>> print(out.shape)
+            paddle.Size([3, 5])
+    """
+    if activation is None:
+        activation = "none"
+
+    if in_dynamic_mode():
+        return _legacy_C_ops.fused_gemm_epilogue(
+            x,
+            y,
+            bias,
+            'trans_x',
+            trans_x,
+            'trans_y',
+            trans_y,
+            'activation',
+            activation,
+        )
+
+    if in_pir_mode():
+        out, _ = _C_ops.fused_gemm_epilogue(
+            x,
+            y,
+            bias,
+            trans_x,
+            trans_y,
+            activation,
+        )
+        return out
+
+    helper = LayerHelper('fused_matmul_bias', **locals())
+    out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    helper.append_op(
+        type='fused_gemm_epilogue',
+        inputs={'X': x, 'Y': y, 'Bias': bias},
+        outputs={'Out': out},
+        attrs={
+            'trans_x': trans_x,
+            'trans_y': trans_y,
+            'activation': activation,
+        },
+    )
+
+    return out

python/paddle/incubate/nn/functional/fused_partial_rope.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fused_partial_rope(
+    x: Tensor,
+    cos: Tensor,
+    sin: Tensor,
+) -> Tensor:
+    r"""
+    Applies partial rotary position embedding on the pe_head_dim portion of input.
+
+    Args:
+        x (Tensor): The input tensor. The data type is bfloat16. The shape of x must be [batch_size, seq_len, num_heads, head_dim].
+        cos (Tensor): The input tensor. The data type is bfloat16. The shape of cos must be [1, seq_len, 1, pe_head_dim] and pe_head_dim must be a multiple of 2 and mustn't exceed head_dim.
+        sin (Tensor): The input tensor. The data type is bfloat16. The shape of sin must be [1, seq_len, 1, pe_head_dim] and pe_head_dim must be a multiple of 2 and mustn't exceed head_dim.
+
+    Returns:
+        out: Tensor representing the fused rotary position embedding, has same shape and data type as `x` .
+
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_partial_rope
+
+            >>> paddle.set_device('gpu')
+            >>> paddle.seed(2025)
+
+            >>> # x: [batch_size, seq_len, num_heads, head_dim]
+            >>> x = paddle.randn([2, 2, 2, 4], dtype='bfloat16')
+
+            >>> # sin, cos: [1, seq_len, 1, pe_head_dim]
+            >>> cos = paddle.randn([1, 2, 1, 2], dtype='bfloat16')
+            >>> sin = paddle.randn([1, 2, 1, 2], dtype='bfloat16')
+
+            >>> # out: [batch_size, seq_len, num_heads, head_dim]
+            >>> out = fused_partial_rope(x, cos, sin)
+            >>> print(out)
+            Tensor(shape=[2, 2, 2, 4], dtype=bfloat16, place=Place(gpu:0), stop_gradient=True,
+                   [[[[-0.17968750,  0.28125000, -0.34765625, -0.92187500],
+                      [-0.83593750,  2.        , -0.13476562, -0.67187500]],
+                     [[ 0.38281250, -0.63281250,  0.25000000, -1.03125000],
+                      [-1.92187500,  2.12500000,  1.92968750, -4.21875000]]],
+                    [[[-0.90625000, -1.62500000, -0.22167969, -0.68359375],
+                      [-0.76562500,  0.23828125,  0.36523438,  0.53515625]],
+                     [[ 0.92578125, -0.85156250, -0.75000000,  1.50000000],
+                      [ 0.41992188, -1.13281250,  0.73437500, -2.18750000]]]])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_partial_rope(x, cos, sin)

python/paddle/incubate/nn/functional/fused_rms_norm.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import paddle
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@overload
+def fused_layer_norm(
+    x: Tensor,
+    norm_weight: Tensor,
+    norm_bias: Tensor,
+    epsilon: float,
+    begin_norm_axis: int,
+    bias: Tensor | None = ...,
+    residual: None = ...,
+    quant_scale: float = ...,
+    quant_round_type: float = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> Tensor: ...
+
+
+@overload
+def fused_layer_norm(
+    x: Tensor,
+    norm_weight: Tensor,
+    norm_bias: Tensor,
+    epsilon: float,
+    begin_norm_axis: int,
+    bias: Tensor | None = ...,
+    residual: Tensor = ...,
+    quant_scale: float = ...,
+    quant_round_type: float = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+def fused_rms_norm(
+    x,
+    norm_weight,
+    norm_bias,
+    epsilon,
+    begin_norm_axis=1,
+    bias=None,
+    residual=None,
+    quant_scale=-1,
+    quant_round_type=0,
+    quant_max_bound=0,
+    quant_min_bound=0,
+):
+    r"""
+    Apply Fused RMSNorm kernel. Also support RMSNorm(bias + residual + x) fused pattern.
+
+    Args:
+        x (Tensor): the input Tensor..
+        norm_weight (Tensor): the weight Tensor to affine output.
+        norm_bias (Tensor): the bias Tensor to affine output.
+        epsilon (float): a small float number to avoid divide 0.
+        begin_norm_axis (int): the begin axis to normalize.
+        bias (optional|Tensor): the previous layers's bias to fused.
+        residual (optional|Tensor): the residual input to fused.
+        quant_scale (float): the quant scale.
+        quant_round_type (float): the quant round type.
+        quant_max_bound (float): the quant max bound to clip.
+        quant_min_bound (float): the quant min bound to clip.
+
+
+    Returns:
+        Tensor: the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+
+            >>> paddle_x = paddle.cast(paddle.randn(shape=[32, 256]), dtype=paddle.float16)
+            >>> paddle_weight = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float16)
+            >>> paddle_bias = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float16)
+            >>> epsilon = 1e-6
+            >>> paddle_rmsnorm = paddle.incubate.nn.functional.fused_rms_norm(paddle_x, paddle_weight, paddle_bias, epsilon, 1)
+    """
+    input_rank = len(x.shape)
+    if begin_norm_axis < 0:
+        begin_norm_axis += input_rank
+
+    if begin_norm_axis < 0 or begin_norm_axis >= input_rank:
+        raise ValueError(
+            f"begin_norm_axis must be in range [0, {input_rank}), "
+            f"but got {begin_norm_axis}"
+            + (
+                f" (originally {begin_norm_axis - input_rank})"
+                if begin_norm_axis < 0
+                else ""
+            )
+        )
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_rms_norm_quant(
+            x,
+            bias,
+            residual,
+            norm_weight,
+            norm_bias,
+            epsilon,
+            begin_norm_axis,
+            quant_scale,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+        )
+    # static mode
+    helper = LayerHelper('fused_rms_norm_quant', **locals())
+    out = None
+    if quant_scale <= 0:
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    else:
+        out = helper.create_variable_for_type_inference(dtype=paddle.int8)
+    outputs_dict = {}
+    outputs_dict['out'] = out
+
+    residual_out = helper.create_variable_for_type_inference(dtype=x.dtype)
+    outputs_dict['residual_out'] = residual_out
+
+    inv_var = helper.create_variable_for_type_inference(dtype=paddle.float32)
+    outputs_dict['inv_var'] = inv_var
+
+    inputs = {'x': x, 'norm_weight': norm_weight}
+    if norm_bias is not None:
+        inputs['norm_bias'] = norm_bias
+    if residual is not None:
+        inputs['residual'] = residual
+    if bias is not None:
+        inputs['bias'] = bias
+
+    helper.append_op(
+        type='fused_rms_norm_quant',
+        inputs=inputs,
+        attrs={
+            "epsilon": epsilon,
+            "begin_norm_axis": begin_norm_axis,
+            "quant_scale": quant_scale,
+            "quant_round_type": quant_round_type,
+            "quant_max_bound": quant_max_bound,
+            "quant_min_bound": quant_min_bound,
+        },
+        outputs=outputs_dict,
+    )
+    return (out, residual_out, outputs_dict['inv_var'])

python/paddle/incubate/nn/functional/fused_rms_norm_ext.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2025 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.
+
+# File: python/paddle/incubate/nn/functional/layer_norm_cuda.py
+from paddle import _C_ops
+from paddle.base.data_feeder import convert_dtype
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+
+def fused_rms_norm_ext(x, scale, epsilon=1e-5, name=None):
+    """
+    Applies Layer Normalization over the last dimension of the input tensor using CUDA implementation.
+    Args:
+        x (Tensor): Input tensor of shape [rows, cols] or higher dimensions (flattened to 2D).
+        scale (Tensor): Scale tensor of shape [cols].
+        bias (Tensor, optional): Bias tensor of shape [cols]. If None, no bias is added.
+        epsilon (float): Small constant to avoid division by zero.
+        name (str, optional): Name of the operator.
+    Returns:
+        y (Tensor): Normalized tensor of same shape as x.
+        mean (Tensor): Tensor of shape [rows], the mean of each row.
+        invvar (Tensor): Tensor of shape [rows], the inverse standard deviation of each row.
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_rms_norm_ext(x, scale, epsilon)
+    helper = LayerHelper('fused_rms_norm_ext', **locals())
+    dtype = convert_dtype(x.dtype)
+    y = helper.create_variable_for_type_inference(dtype)
+    invvar = helper.create_variable_for_type_inference('float32')
+
+    inputs = {'x': x, 'scale': scale}
+
+    helper.append_op(
+        type='fused_rms_norm_ext',
+        inputs=inputs,
+        outputs={'y': y, 'invvar': invvar},
+        attrs={'epsilon': epsilon},
+    )
+    return y, invvar

python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def fused_rotary_position_embedding(
+    q: Tensor,
+    k: Tensor | None = None,
+    v: Tensor | None = None,
+    sin: Tensor | None = None,
+    cos: Tensor | None = None,
+    position_ids: Tensor | None = None,
+    use_neox_rotary_style: bool = True,
+    time_major: bool = False,
+    rotary_emb_base: float = 10000.0,
+) -> tuple[Tensor, Tensor, Tensor]:
+    r"""
+    Fused rotary position embedding.
+
+    Args:
+        q (Tensor): The input tensor. The data type is bfloat16, float16, float32 or float64. The shape of q must be [batch_size, seq_len, num_heads, head_dim] or [seq_len, batch_size, num_heads, head_dim] and head_dim must be a multiple of 2.
+        k (Tensor, optional): The input tensor. The data type is bfloat16, float16, float32 or float64. The shape of k must be [batch_size, seq_len, num_heads, head_dim] or [seq_len, batch_size, num_heads, head_dim] and head_dim must be a multiple of 2.
+        v (Tensor, optional): The input tensor. The data type is bfloat16, float16, float32 or float64. The shape of v must be [batch_size, seq_len, num_heads, head_dim] or [seq_len, batch_size, num_heads, head_dim] and head_dim must be a multiple of 2.
+        sin (Tensor, optional): The input tensor. The data type is bfloat16, float16, float32 or float64. The shape of sin must be [seq_len, head_dim] or [1, seq_len, 1, head_dim] and head_dim must be a multiple of 2.
+        cos (Tensor, optional): The input tensor. The data type is bfloat16, float16, float32 or float64. The shape of cos must be [seq_len, head_dim] or [1, seq_len, 1, head_dim] and head_dim must be a multiple of 2.
+        position_ids (Tensor, optional): The input tensor. The data type is int64. The shape of position_ids must be [batch_size, seq_len].
+        use_neox_rotary_style(optional|bool): When the use_neox_rotary_style is True, every two adjacent numbers are calculated. When the use_neox_rotary_style is False, the numbers corresponding to the positions of the front half and back half segments are calculated. Default True.
+        time_major(optional|bool): Whether the first dimension of the q, k, v input means the time steps. If time_major is True, the shape of Tensor is [seq_len, batch_size, num_heads, head_dim], otherwise [batch_size, seq_len, num_heads, head_dime]. Defaults to False. `time_steps` means the length of input sequence.
+        rotary_emb_base(optional|float): the base of the rotary embedding. Default 10000.
+
+    Returns:
+        out_q/out_k/out_v Tensor representing the fused rotary position embedding, has same shape and data type as `q` .
+
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import fused_rotary_position_embedding
+
+            >>> paddle.set_device('gpu')
+
+            >>> # batch_size = 2
+            >>> # seq_len = 2
+            >>> # num_heads = 2
+            >>> # head_dim = 2
+
+            >>> paddle.seed(1204)
+
+            >>> # q, k, v: [batch_size, seq_len, num_heads, head_dim]
+            >>> q = paddle.randn([2, 2, 2, 2], dtype='float16')
+            >>> k = paddle.randn([2, 2, 2, 2], dtype='float16')
+            >>> v = paddle.randn([2, 2, 2, 2], dtype='float16')
+
+            >>> # sin, cos: [1, seq_len, 1, head_dim]
+            >>> x = paddle.randn([1, 2, 1, 2], dtype='float16')
+            >>> y = paddle.randn([1, 2, 1, 2], dtype='float16')
+            >>> sin = paddle.sin(x)
+            >>> cos = paddle.cos(y)
+
+            >>> # position_ids: [batch_size, seq_len]
+            >>> position_ids = paddle.randint(high=2, size=[2, 2], dtype='int64')
+
+            >>> # out_q, out_k, out_v: [batch_size, seq_len, num_heads, head_dim]
+            >>> out_q, out_k, out_v = fused_rotary_position_embedding(
+            ...     q, k, v, sin=sin, cos=cos, position_ids=position_ids, use_neox_rotary_style=False
+            ... )
+            >>> print(out_q)
+            >>> # doctest: +SKIP("Random output")
+            Tensor(shape=[2, 2, 2, 2], dtype=float16, place=Place(gpu:0), stop_gradient=True,
+            [[[[-0.54931641,  0.64990234],
+               [-1.08691406,  1.18261719]],
+              [[ 0.57812500,  0.11749268],
+               [-0.63281250,  0.15551758]]],
+             [[[-0.77050781,  0.07733154],
+               [-0.73730469, -0.16735840]],
+              [[ 0.07116699, -0.90966797],
+               [-0.03628540, -0.20202637]]]])
+            >>> # doctest: -SKIP
+    """
+    if (sin is None) or (cos is None):
+        assert position_ids is None, (
+            "position_ids without sin/cos is not correctly supported now."
+        )
+        assert use_neox_rotary_style, (
+            "rotate_half without sin/cos is not correctly supported now."
+        )
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.fused_rotary_position_embedding(
+            q,
+            k,
+            v,
+            sin,
+            cos,
+            position_ids,
+            use_neox_rotary_style,
+            time_major,
+            rotary_emb_base,
+        )
+
+    helper = LayerHelper('fused_rotary_position_embedding', **locals())
+    out_q = helper.create_variable_for_type_inference(dtype=q.dtype)
+    out_k = (
+        helper.create_variable_for_type_inference(dtype=k.dtype) if k else None
+    )
+    out_v = (
+        helper.create_variable_for_type_inference(dtype=v.dtype) if v else None
+    )
+
+    outputs = {'out_q': out_q}
+    if out_k:
+        outputs.update({'out_k': out_k})
+    if out_v:
+        outputs.update({'out_v': out_v})
+
+    helper.append_op(
+        type='fused_rotary_position_embedding',
+        inputs={
+            'q': q,
+            'k': k,
+            'v': v,
+            'sin': sin,
+            'cos': cos,
+            'position_ids': position_ids,
+        },
+        outputs=outputs,
+        attrs={
+            'use_neox_rotary_style': use_neox_rotary_style,
+            'time_major': time_major,
+            'rotary_emb_base': rotary_emb_base,
+        },
+    )
+
+    return out_q, out_k, out_v

python/paddle/incubate/nn/functional/int_bincount.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2025 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.
+
+import paddle
+from paddle import _C_ops
+from paddle.base.data_feeder import convert_dtype
+from paddle.base.framework import (
+    convert_nptype_to_datatype_or_vartype,
+    core,
+    in_dynamic_or_pir_mode,
+)
+from paddle.base.layer_helper import LayerHelper
+
+
+def math_int_bincount(x, low, high, dtype):
+    """
+    A mathematically equivalent implementation of int_bincount using scatter and sum
+
+    Args:
+        x (Tensor): A 1D or 2D int64 tensor containing category indices.
+        low (int): The minimum possible category index (usually 0).
+        high (int): One past the maximum category index (i.e., number of categories).
+        dtype (paddle.dtype): Data type of the output tensor (e.g., paddle.int64).
+
+    Returns:
+        Tensor: A 1D tensor of shape [high - low], where each element is
+                the count of occurrences of that category in `x`.
+    """
+    if x.ndim not in [0, 1, 2]:
+        raise ValueError(
+            f"x must be a 0D, 1D or 2D tensor, but got ndim={x.ndim}"
+        )
+
+    if x.dtype not in [paddle.int32, paddle.int64]:
+        raise ValueError(f"x.dtype must be int32 or int64, but got {x.dtype}")
+
+    if dtype not in ['int32', 'int64', paddle.int32, paddle.int64]:
+        raise ValueError(f"dtype must be 'int32' or 'int64', but got '{dtype}'")
+
+    if high < low:
+        raise ValueError(
+            f"'high' ({high}) must be greater than or equal to 'low' ({low})"
+        )
+
+    if x.numel().item() == 0:
+        return paddle.zeros([high - low], dtype=dtype)
+
+    if x.ndim == 0:
+        x = x.reshape([-1]).unsqueeze(0)  # Shape: [1, N]
+    elif x.ndim == 1:
+        x = x.unsqueeze(0)  # Shape: [1, N]
+
+    x_min = x.min().item()
+    x_max = x.max().item()
+
+    if x_min < 0:
+        raise ValueError(
+            f"Elements of x must be non-negative, but got min={x_min}"
+        )
+
+    max_val = max(x_max + 1, high)
+    mask = paddle.zeros([x.shape[0], max_val], dtype=x.dtype)
+    mask = mask.put_along_axis(
+        x, paddle.to_tensor(1.0, dtype=x.dtype), axis=1, reduce='add'
+    )
+
+    count = paddle.sum(mask, axis=0).cast(dtype)
+    return count[low:high]
+
+
+def int_bincount(x, low, high, dtype=None, name=None):
+    if in_dynamic_or_pir_mode():
+        if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
+            dtype = convert_nptype_to_datatype_or_vartype(dtype)
+
+        if paddle.is_compiled_with_xpu():
+            return math_int_bincount(x, low, high, dtype)
+        else:
+            return _C_ops.int_bincount(x, low, high, dtype)
+
+    helper = LayerHelper("int_bincount", **locals())
+    out_dtype = dtype if dtype is not None else x.dtype
+    y = helper.create_variable_for_type_inference(dtype=out_dtype)
+    dtype_attr = convert_dtype(out_dtype)
+
+    helper.append_op(
+        type="int_bincount",
+        inputs={"x": x},
+        outputs={"y": y},
+        attrs={
+            "low": low,
+            "high": high,
+            "dtype": dtype_attr,
+        },
+    )
+    return y

python/paddle/incubate/nn/functional/masked_multihead_attention.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.utils.deprecated import deprecated
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@overload
+def masked_multihead_attention(
+    x: Tensor,
+    cache_kv: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    src_mask: Tensor | None = ...,
+    cum_offsets: Tensor | None = ...,
+    sequence_lengths: Tensor | None = ...,
+    rotary_tensor: Tensor | None = ...,
+    beam_cache_offset: None = ...,
+    qkv_out_scale: Tensor | None = ...,
+    out_shift: Tensor | None = ...,
+    out_smooth: Tensor | None = ...,
+    seq_len: int = ...,
+    rotary_emb_dims: int = ...,
+    use_neox_rotary_style: bool = ...,
+    compute_dtype: str = ...,
+    out_scale: float = ...,
+    quant_round_type: int = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def masked_multihead_attention(
+    x: Tensor,
+    cache_kv: Tensor | None = ...,
+    bias: Tensor | None = ...,
+    src_mask: Tensor | None = ...,
+    cum_offsets: Tensor | None = ...,
+    sequence_lengths: Tensor | None = ...,
+    rotary_tensor: Tensor | None = ...,
+    beam_cache_offset: Tensor = ...,
+    qkv_out_scale: Tensor | None = ...,
+    out_shift: Tensor | None = ...,
+    out_smooth: Tensor | None = ...,
+    seq_len: int = ...,
+    rotary_emb_dims: int = ...,
+    use_neox_rotary_style: bool = ...,
+    compute_dtype: str = ...,
+    out_scale: float = ...,
+    quant_round_type: int = ...,
+    quant_max_bound: float = ...,
+    quant_min_bound: float = ...,
+) -> tuple[Tensor, Tensor, Tensor]: ...
+
+
+@deprecated(
+    since="3.4.0",
+    level=1,
+    update_to="paddle.nn.functional.scaled_dot_product_attention",
+)
+def masked_multihead_attention(
+    x,
+    cache_kv=None,
+    bias=None,
+    src_mask=None,
+    cum_offsets=None,
+    sequence_lengths=None,
+    rotary_tensor=None,
+    beam_cache_offset=None,
+    qkv_out_scale=None,
+    out_shift=None,
+    out_smooth=None,
+    seq_len=1,
+    rotary_emb_dims=0,
+    use_neox_rotary_style=False,
+    compute_dtype='default',
+    out_scale=-1,
+    quant_round_type=1,
+    quant_max_bound=127.0,
+    quant_min_bound=-127.0,
+):
+    r"""
+    Masked Multi-head attention for text summarization.
+    This is a fusion operator to compute masked multi-head attention in transformer model architecture.
+    This operator only supports running on GPU.
+
+    Args:
+        x (Tensor): The input tensor could be 2-D tensor. Its shape is [batch_size, 3 * num_head * head_dim].
+        cache_kv (Tensor): The cache structure tensors for the generation model. Its shape is [2, batch_size, num_head, max_seq_len, head_dim].
+        bias (Tensor, optional): The bias tensor. Its shape is [3, num_head, head_dim].
+        src_mask (Tensor, optional): The src_mask tensor. Its shape is [batch_size, 1, 1, sequence_length].
+        sequence_lengths (Tensor, optional): The sequence_lengths tensor, used to index input. Its shape is [batch_size, 1].
+        rotary_tensor (Tensor, optional): The rotary_tensor tensor. The dtype must be float. Its shape is [batch_size, 1, 1, sequence_length, head_dim].
+        beam_cache_offset (Tensor, optional): The beam_cache_offset tensor. Its shape is [batch_size, beam_size, max_seq_len + max_dec_len].
+        qkv_out_scale (Tensor, optional): The qkv_out_scale tensor, used in quant. Its shape is [3, num_head, head_dim].
+        out_shift (Tensor, optional): The out_shift tensor, used in quant.
+        out_smooth (Tensor, optional): The out_smooth tensor, used in quant.
+        seq_len (int, optional): The seq_len, used to get input length. Default 1.
+        rotary_emb_dims (int, optional): The rotary_emb_dims. Default 1.
+        use_neox_rotary_style (bool, optional): A flag indicating whether neox_rotary_style is needed or not. Default False.
+        compute_dtype (string): A compute dtype, used to represent the input data type.
+        out_scale (float, optional): The out_scale, used in quant.
+        quant_round_type (int, optional): The quant_round_type, used in quant. Default 1.
+        quant_max_bound (float, optional): The quant_max_bound, used in quant. Default 127.0.
+        quant_min_bound (float, optional): The quant_min_bound, used in quant. Default -127.0.
+
+    Returns:
+        Tensor|tuple: If "beam_cache_offset_out" is not none, return the
+        tuple (output, cache_kvs_out, beam_cache_offset_out), which output is the output of
+        masked_multihead_attention layers, cache_kvs_out is inplace with input `cache_kvs`.
+        If "beam_cache_offset_out" is none, return the tuple (output, cache_kvs_out).
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> paddle.device.set_device('gpu')
+
+            >>> # input: [batch_size, 3 * num_head * dim_head]
+            >>> x = paddle.rand(shape=(2, 3 * 32 * 128), dtype="float32")
+
+            >>> # src_mask: [batch_size, 1, 1, sequence_length]
+            >>> src_mask = paddle.rand(shape=(2, 1, 1, 10), dtype="float32")
+
+            >>> # cache_kv: [2, batch_size, num_head, max_seq_len, dim_head]
+            >>> cache_kv = paddle.rand(shape=(2, 2, 32, 64, 128), dtype="float32")
+
+            >>> output = F.masked_multihead_attention(
+            ...     x,
+            ...     src_mask=src_mask,
+            ...     cache_kv=cache_kv,
+            ... )
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.masked_multihead_attention_(
+            x,
+            cache_kv,
+            bias,
+            src_mask,
+            cum_offsets,
+            sequence_lengths,
+            rotary_tensor,
+            beam_cache_offset,
+            qkv_out_scale,
+            out_shift,
+            out_smooth,
+            seq_len,
+            rotary_emb_dims,
+            use_neox_rotary_style,
+            compute_dtype,
+            out_scale,
+            quant_round_type,
+            quant_max_bound,
+            quant_min_bound,
+        )
+
+    helper = LayerHelper('masked_multihead_attention', **locals())
+    if x.dtype == "int32":
+        if compute_dtype == "bf16":
+            dtype = "uint16"
+        elif compute_dtype == "fp16":
+            dtype = "float16"
+        elif compute_dtype == "fp32":
+            dtype = "float32"
+        out = helper.create_variable_for_type_inference(dtype=dtype)
+    else:
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+    inputs = {}
+    inputs['x'] = x
+    inputs['cache_kv'] = cache_kv
+    if bias is not None:
+        inputs['bias'] = bias
+    if src_mask is not None:
+        inputs['src_mask'] = src_mask
+    if cum_offsets is not None:
+        inputs['cum_offsets'] = cum_offsets
+    if sequence_lengths is not None:
+        inputs['sequence_lengths'] = sequence_lengths
+    if rotary_tensor is not None:
+        inputs['rotary_tensor'] = rotary_tensor
+    beam_cache_offset_flag = False
+    if beam_cache_offset is not None:
+        inputs['beam_cache_offset'] = beam_cache_offset
+        beam_cache_offset_flag = True
+    else:
+        beam_cache_offset = helper.create_variable_for_type_inference(
+            dtype="int"
+        )
+    if qkv_out_scale is not None:
+        inputs['qkv_out_scale'] = qkv_out_scale
+    if out_shift is not None:
+        inputs['out_shift'] = out_shift
+    if out_smooth is not None:
+        inputs['out_smooth'] = out_smooth
+
+    outputs = {
+        'out': out,
+        'cache_kv_out': cache_kv,
+        'beam_cache_offset_out': beam_cache_offset,
+    }
+    helper.append_op(
+        type='masked_multihead_attention',
+        inputs=inputs,
+        outputs=outputs,
+        attrs={
+            'seq_len': seq_len,
+            'rotary_emb_dims': rotary_emb_dims,
+            'use_neox_rotary_style': use_neox_rotary_style,
+            'compute_dtype': compute_dtype,
+            'out_scale': out_scale,
+            'quant_round_type': quant_round_type,
+            'quant_max_bound': quant_max_bound,
+            'quant_min_bound': quant_min_bound,
+        },
+    )
+    return (
+        (out, cache_kv, beam_cache_offset)
+        if beam_cache_offset_flag is not None
+        else (out, cache_kv)
+    )

python/paddle/incubate/nn/functional/moe_combine.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_combine(
+    x: Tensor,
+    combine_weights: Tensor,
+    scatter_index: Tensor,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        x: Input tensor [seq, dim]
+        combine_weights: Combination weights [s, k]
+        scatter_index: Scatter indices [k, s] dtype=int32
+
+    Returns:
+        Output Combined output [s, dim]
+    """
+    if in_dynamic_or_pir_mode():
+        if not (
+            x.process_mesh is None
+            and combine_weights.process_mesh is None
+            and scatter_index.process_mesh is None
+        ):
+            # auto parallel mode
+            return _C_ops.moe_combine_auto(x, combine_weights, scatter_index)
+        return _C_ops.moe_combine(x, combine_weights, scatter_index)
+    helper = LayerHelper('moe_combine', **locals())
+    y = helper.create_variable_for_type_inference(dtype=x.dtype)
+    inputs = {
+        'x': x,
+        'combine_weights': combine_weights,
+        'scatter_index': scatter_index,
+    }
+    helper.append_op(type='moe_combine', inputs=inputs, outputs={'y': y})
+    return y

python/paddle/incubate/nn/functional/moe_combine_no_weight.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_combine_no_weight(
+    x: Tensor,
+    combine_weight: Tensor,
+    scatter_index: Tensor,
+    epsilon: float = 1e-15,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        x: Input tensor [num_tokens, hidden_size]
+        scatter_index: Scatter indices [seq_len, k] dtype=int32
+
+    Returns:
+        Output Combined output [seq_len, hidden_size]
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.moe_combine_no_weight(
+            x, combine_weight, scatter_index, epsilon
+        )
+    helper = LayerHelper('moe_combine_no_weight', **locals())
+    y = helper.create_variable_for_type_inference(dtype=x.dtype)
+    inputs = {
+        'x': x,
+        'combine_weight': combine_weight,
+        'scatter_index': scatter_index,
+        'epsilon': epsilon,
+    }
+    helper.append_op(
+        type='moe_combine_no_weight', inputs=inputs, outputs={'y': y}
+    )
+    return y

python/paddle/incubate/nn/functional/moe_gate_dispatch.py

@@ -0,0 +1,280 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import _C_ops
+
+# from ....framework import LayerHelper, in_dynamic_or_pir_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_gate_dispatch(
+    x: Tensor,
+    gate_logits: Tensor,
+    corr_bias: Tensor,
+    k: int,
+    capacity: int,
+    use_pad: bool,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Args:
+        x:
+        gate_logits:
+        corr_bias:
+        k:
+        capacity:
+        use_pad:
+
+    Returns:
+        y:
+        combine_weights:
+        scatter_index:
+        expert_offset:
+        expert_id:
+    """
+    if in_dynamic_or_pir_mode():
+        if paddle.device.is_compiled_with_custom_device('npu'):
+            return math_moe_gate_dispatch(
+                x, gate_logits, corr_bias, k, capacity, use_pad
+            )
+        else:
+            if not (
+                x.process_mesh is None and gate_logits.process_mesh is None
+            ):
+                return _C_ops.moe_gate_dispatch_auto(
+                    x, gate_logits, corr_bias, k, capacity, use_pad
+                )
+            return _C_ops.moe_gate_dispatch(
+                x, gate_logits, corr_bias, k, capacity, use_pad
+            )
+
+    helper = LayerHelper('moe_gate_dispatch', **locals())
+    y = helper.create_variable_for_type_inference(dtype=x.dtype)
+    combine_weights = helper.create_variable_for_type_inference(
+        dtype=paddle.float32
+    )
+    scatter_index = helper.create_variable_for_type_inference(
+        dtype=paddle.int32
+    )
+    expert_offset = helper.create_variable_for_type_inference(
+        dtype=paddle.int64
+    )
+    expert_id = helper.create_variable_for_type_inference(dtype=paddle.int32)
+
+    inputs = {
+        'x': x,
+        'gate_logits': gate_logits,
+        'corr_bias': corr_bias,
+    }
+    attrs = {
+        'k': k,
+        'capacity': capacity,
+        'use_pad': use_pad,
+    }
+    outputs = {
+        'y': y,
+        'combine_weights': combine_weights,
+        'scatter_index': scatter_index,
+        'expert_offset': expert_offset,
+        'expert_id': expert_id,
+    }
+    helper.append_op(
+        type='moe_gate_dispatch',
+        inputs=inputs,
+        attrs=attrs,
+        outputs=outputs,
+    )
+    return y, combine_weights, scatter_index, expert_offset, expert_id
+
+
+def topk_gating_softmax(
+    gate_logits,
+    corr_bias,
+    topk,
+):
+    # Calculate scores with bias added (used for Top-K selection)
+    scores_for_selection = (
+        gate_logits + corr_bias if corr_bias is not None else gate_logits
+    )
+    # Get Top-K indices
+    combine_weights, expert_id = paddle.topk(
+        scores_for_selection, k=topk, axis=1
+    )
+
+    # Initialize source_rows: for each column, increment by 1 across rows,
+    # then move to the next column after finishing one full column
+    source_rows = paddle.to_tensor(
+        [
+            k_idx * gate_logits.shape[0] + row_idx
+            for row_idx in range(gate_logits.shape[0])
+            for k_idx in range(topk)
+        ]
+    )
+    return combine_weights, expert_id, source_rows
+
+
+def sorter_kernel(expert_id, source_rows):
+    # Flatten all data
+    flat_expert = expert_id.flatten()
+
+    # Global sorting index (sorted by expert_id in ascending order)
+    sort_idx = paddle.argsort(flat_expert)
+
+    # Apply the sorting
+    sorted_expert = paddle.gather(flat_expert, sort_idx)
+    sorted_source = paddle.gather(source_rows, sort_idx)
+
+    # Reshape back to [num_rows, k]
+    return (sorted_expert.reshape(expert_id.shape), sorted_source)
+
+
+def compute_total_rows_before_expert(permuted_experts, num_experts):
+    expert_offset = paddle.searchsorted(
+        permuted_experts.flatten(), paddle.arange(num_experts), right=True
+    )
+    return expert_offset
+
+
+def initialize_moe_routing_matrix(
+    unpermuted_input,
+    gate_logits,
+    expanded_dest_row_to_expanded_source_row,
+    permuted_experts,
+    expert_offset,
+    combine_weights,
+    capacity,
+    use_pad=False,
+):
+    splits = paddle.concat(
+        [
+            paddle.to_tensor([0]),
+            expert_offset,
+            paddle.to_tensor([len(expanded_dest_row_to_expanded_source_row)]),
+        ]
+    )
+    expanded_dest_row_to_expanded_source_row = paddle.concat(
+        [
+            paddle.sort(
+                expanded_dest_row_to_expanded_source_row[
+                    splits[i] : splits[i + 1]
+                ]
+            )
+            for i in range(len(splits) - 1)
+        ]
+    )
+
+    expanded_source_row_to_expanded_dest_row = paddle.scatter_nd(
+        index=expanded_dest_row_to_expanded_source_row.unsqueeze(1),
+        updates=paddle.arange(
+            expanded_dest_row_to_expanded_source_row.shape[0]
+        ),
+        shape=[expanded_dest_row_to_expanded_source_row.shape[0]],
+    )
+
+    y = paddle.zeros(
+        [gate_logits.shape[1] * capacity, unpermuted_input.shape[1]],
+        dtype=unpermuted_input.dtype,
+    )
+
+    if use_pad:
+        iexpert = paddle.gather(
+            permuted_experts.flatten(),
+            expanded_source_row_to_expanded_dest_row.flatten(),
+        )
+
+        extended_offset = paddle.concat(
+            [paddle.zeros([1], dtype='int64'), expert_offset]
+        )
+        offset = paddle.gather(extended_offset, iexpert)
+
+        iexpert_cap = iexpert * capacity
+        row_in_expert = (
+            expanded_source_row_to_expanded_dest_row.flatten() - offset
+        )
+
+        input_indices = (
+            paddle.arange(row_in_expert.shape[0]) % unpermuted_input.shape[0]
+        )
+
+        y = paddle.scatter(
+            x=y,
+            index=row_in_expert + iexpert_cap,
+            updates=unpermuted_input[input_indices],
+            overwrite=True,
+        )
+
+        expanded_source_row_to_expanded_dest_row = (
+            expanded_source_row_to_expanded_dest_row
+            + iexpert_cap.reshape(
+                expanded_source_row_to_expanded_dest_row.shape
+            )
+            - offset.reshape(expanded_source_row_to_expanded_dest_row.shape)
+        )
+
+        expanded_source_row_to_expanded_dest_row = (
+            expanded_source_row_to_expanded_dest_row.reshape(
+                [combine_weights.shape[1], combine_weights.shape[0]]
+            )
+        )
+        mask = (
+            row_in_expert.reshape(
+                [combine_weights.shape[1], combine_weights.shape[0]]
+            )
+            < capacity
+        )
+
+        expanded_source_row_to_expanded_dest_row = paddle.where(
+            mask,
+            expanded_source_row_to_expanded_dest_row,
+            paddle.zeros_like(expanded_source_row_to_expanded_dest_row),
+        )
+        combine_weights = paddle.where(
+            mask.T, combine_weights, paddle.zeros_like(combine_weights)
+        )
+
+    return y, expanded_source_row_to_expanded_dest_row, combine_weights
+
+
+def math_moe_gate_dispatch(x, gate_logits, corr_bias, k, capacity, use_pad):
+    combine_weights, expert_id, source_rows = topk_gating_softmax(
+        gate_logits, corr_bias, k
+    )
+
+    permuted_experts, permuted_rows = sorter_kernel(expert_id, source_rows)
+
+    expert_offset = compute_total_rows_before_expert(
+        permuted_experts, gate_logits.shape[1]
+    )
+
+    y, scatter_index, combine_weights = initialize_moe_routing_matrix(
+        x,
+        gate_logits,
+        permuted_rows,
+        permuted_experts,
+        expert_offset,
+        combine_weights,
+        capacity,
+        use_pad,
+    )
+
+    return y, combine_weights, scatter_index, expert_offset, expert_id

python/paddle/incubate/nn/functional/moe_gate_dispatch_and_quant.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.incubate.nn.functional import (
+    fp8,
+    moe_gate_dispatch,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_gate_dispatch_and_quant(
+    x: Tensor,
+    gate_logits: Tensor,
+    corr_bias: Tensor,
+    k: int,
+    capacity: int,
+    use_pad: bool,
+    use_pow2_scale: bool,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+    """
+    Args:
+        x: Input tensor [batch_size, seq_len, hidden_dim].
+        gate_logits: Gate logits for choosing experts [batch_size, seq_len, num_experts].
+        corr_bias: Optional correction bias to adjust gate logits.
+        k: Top-k experts to be selected.
+        capacity: The maximum number of tokens an expert can handle.
+        use_pad: Boolean indicating if padding is used for uniform input length.
+        use_pow2_scale: Boolean indicating if power-of-two scaling is applied for quantization.
+
+    Returns:
+        Tuple of Tensors containing:
+        - fp8_out: Processed output tensor in FP8 format.
+        - scale: Scaling factors used during processing.
+        - combine_weights: Weights for combining experts' outputs.
+        - scatter_index: Indices for scattering inputs to experts.
+        - expert_offset: Offset indices for each expert.
+        - expert_id: IDs of selected experts for each position.
+    """
+    if not in_dynamic_or_pir_mode():
+        raise NotImplementedError('Static graph mode not implemented')
+    else:
+        if paddle.device.is_compiled_with_custom_device('npu'):
+            return math_moe_gate_dispatch_and_quant(
+                x, gate_logits, corr_bias, k, capacity, use_pad
+            )
+        else:
+            return _C_ops.moe_gate_dispatch_and_quant(
+                x, gate_logits, corr_bias, k, capacity, use_pad, use_pow2_scale
+            )
+
+
+def math_moe_gate_dispatch_and_quant(
+    x, gate_logits, corr_bias, k, capacity, use_pad, use_pow2_scale
+):
+    y, combine_weights, scatter_index, expert_offset, expert_id = (
+        moe_gate_dispatch(x, gate_logits, corr_bias, k, capacity, use_pad)
+    )
+    y_fp8, scale = fp8.fp8_quant_blockwise(
+        y,
+        quant_method="1x128",
+        output_scale_transpose=False,
+        using_pow2_scale=use_pow2_scale,
+    )
+    return (
+        y_fp8,
+        scale,
+        combine_weights,
+        scatter_index,
+        expert_offset,
+        expert_id,
+    )

python/paddle/incubate/nn/functional/moe_gate_dispatch_partial_nosoftmaxtopk.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_gate_dispatch_partial_nosoftmaxtopk(
+    x: Tensor,
+    combine_weights: Tensor,
+    expert_id: Tensor,
+    k: int,
+    capacity: int,
+    num_experts: int,
+    use_pad: bool,
+    expert_start_index: int,
+    expert_end_index: int,
+    reverse_token_drop: bool,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+    if in_dynamic_or_pir_mode():
+        return _C_ops.moe_gate_dispatch_partial_nosoftmaxtopk(
+            x,
+            combine_weights,
+            expert_id,
+            k,
+            capacity,
+            num_experts,
+            use_pad,
+            expert_start_index,
+            expert_end_index,
+            reverse_token_drop,
+        )
+    helper = LayerHelper("moe_gate_dispatch_partial_nosoftmaxtopk", **locals())
+    y = helper.create_variable_for_type_inference(dtype=x.dtype)
+    combine_weights_out = helper.create_variable_for_type_inference(
+        dtype=combine_weights.dtype
+    )
+    scatter_index = helper.create_variable_for_type_inference(dtype='int32')
+    scatter_index_rev = helper.create_variable_for_type_inference(dtype='int32')
+    expert_offset = helper.create_variable_for_type_inference(dtype='int64')
+    expert_nums_local = helper.create_variable_for_type_inference(dtype='int64')
+    inputs = {
+        "x": x,
+        "combine_weights": combine_weights,
+        "expert_id": expert_id,
+    }
+    outputs = {
+        "y": y,
+        "combine_weights_out": combine_weights_out,
+        "scatter_index": scatter_index,
+        "scatter_index_rev": scatter_index_rev,
+        "expert_offset": expert_offset,
+        "expert_nums_local": expert_nums_local,
+    }
+    attrs = {
+        "k": k,
+        "capacity": capacity,
+        "num_experts": num_experts,
+        "use_pad": use_pad,
+        "expert_start_index": expert_start_index,
+        "expert_end_index": expert_end_index,
+        "reverse_token_drop": reverse_token_drop,
+    }
+    helper.append_op(
+        type="moe_gate_dispatch_partial_nosoftmaxtopk",
+        inputs=inputs,
+        outputs=outputs,
+        attrs=attrs,
+    )
+    return (
+        y,
+        combine_weights_out,
+        scatter_index,
+        scatter_index_rev,
+        expert_offset,
+        expert_nums_local,
+    )

python/paddle/incubate/nn/functional/moe_gate_dispatch_permute.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.base.framework import in_dynamic_or_pir_mode
+from paddle.base.layer_helper import LayerHelper
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def moe_gate_dispatch_permute(
+    x: Tensor,
+    gate_logits: Tensor,
+    corr_bias: Tensor,
+    k: int,
+    capacity: int,
+    world_size: int,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    """
+    Dispatch and permute for Mixture of Experts (MoE).
+
+    Args:
+        x: Input tensor [batch_size, seq_len, hidden_dim].
+        gate_logits: Gate logits for choosing experts [batch_size, seq_len, num_experts].
+        corr_bias: Optional correction bias to adjust gate logits.
+        k: Top-k experts to be selected.
+        capacity: The maximum number of tokens an expert can handle.
+        world_size: Number of distributed processes.
+        name: Optional name for the operation.
+
+    Returns:
+        Tuple of Tensors containing:
+        - y: Output tensor after dispatch and permute.
+        - combine_weights: Weights for combining experts' outputs.
+        - scatter_index: Indices for scattering inputs to experts.
+        - expert_offset: Offset indices for each expert.
+        - expert_id: IDs of selected experts for each position.
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.moe_gate_dispatch_permute(
+            x, gate_logits, corr_bias, k, capacity, world_size
+        )
+
+    helper = LayerHelper('moe_gate_dispatch_permute', **locals())
+    y = helper.create_variable_for_type_inference(dtype=x.dtype)
+    combine_weights = helper.create_variable_for_type_inference(dtype='float')
+    scatter_index = helper.create_variable_for_type_inference(dtype='int32')
+    expert_offset = helper.create_variable_for_type_inference(dtype='int32')
+    expert_id = helper.create_variable_for_type_inference(dtype='int32')
+
+    inputs = {
+        'x': x,
+        'gate_logits': gate_logits,
+        'corr_bias': corr_bias if corr_bias is not None else None,
+    }
+    attrs = {'k': k, 'capacity': capacity, 'world_size': world_size}
+    outputs = {
+        'y': y,
+        'combine_weights': combine_weights,
+        'scatter_index': scatter_index,
+        'expert_offset': expert_offset,
+        'expert_id': expert_id,
+    }
+
+    helper.append_op(
+        type='moe_gate_dispatch_permute',
+        inputs=inputs,
+        outputs=outputs,
+        attrs=attrs,
+    )
+    return y, combine_weights, scatter_index, expert_offset, expert_id

python/paddle/incubate/nn/functional/swiglu.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.utils import deprecated
+
+from ....framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@deprecated(
+    since="3.3.0",
+    update_to="paddle.nn.functional.swiglu",
+    level=1,
+    reason="paddle.incubate.nn.functional.swiglu will be removed in future. Please use paddle.nn.functional.swiglu instead.",
+)
+def swiglu(
+    x: Tensor, y: Tensor | None = None, name: str | None = None
+) -> Tensor:
+    """
+    This function performs SwiGLU activation to the input Tensor.
+
+    .. math::
+
+        out = silu(x) * y when y is not None
+        out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+
+    Args:
+        x (Tensor): The first input Tensor of SwiGLU.
+        y (Tensor, optional): The second input Tensor of SwiGLU. Default: None.
+        name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None.
+
+    Returns:
+        A Tensor with the same data type with x and y.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.incubate.nn.functional as F
+            >>> x = paddle.to_tensor([1, 2], dtype='float32')
+            >>> out1, out2 = F.swiglu(x), F.swiglu(x, x)
+            >>> print(out1, out2)
+            Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
+                   [1.46211720]) Tensor(shape=[2], dtype=float32, place=Place(cpu), stop_gradient=True,
+                   [0.73105860, 3.52318811])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.swiglu(x, y)

python/paddle/incubate/nn/functional/variable_length_memory_efficient_attention.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2023 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.
+
+# The following codes are from https://github.com/facebookresearch/xformers
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def variable_length_memory_efficient_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    seq_lens: Tensor,
+    kv_seq_lens: Tensor,
+    mask: Tensor | None = None,
+    scale: float | None = None,
+    causal: bool = False,
+    pre_cache_length: int = 0,
+) -> Tensor:
+    """
+    Cutlass Memory Efficient Variable Attention.
+    This method requires SM_ARCH in sm70, sm75, sm80.
+
+    Args:
+        query (Tensor): The Query Tensor. Its shape is [batchsize, num_head, seq_len, head_size].
+        key (Tensor): The Key Tensor. Its shape is [batchsize, num_head, seq_len, head_size].
+        value (Tensor): The Value Tensor. Its shape is [batchsize, num_head, seq_len, head_size].
+        seq_lens (Tensor): The sequence lengths of the sequences in the batch, used to index query. Its shape is [batchsize, 1].
+        kv_seq_lens (Tensor): The sequence lengths of the sequences in the batch, used to index key and value. Its shape is [batchsize, 1].
+        mask (Tensor|None, optional): The Mask Tensor. Its shape is [batchsize, 1, query_seq_len, key_seq_len].
+        scale (Float|None, optional): The attention matrix's scale. Default is sqrt(1.0 / head_size).
+        causal (Bool): Whether causal masking is used or not. Default is False.
+        pre_cache_length (Int): The length of the pre-cache. Default is 0.
+    Returns:
+        Tensor, the output Tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import math
+            >>> import paddle
+            >>> from paddle.incubate.nn.functional import (
+            ...     variable_length_memory_efficient_attention,
+            ... )
+            >>> paddle.device.set_device('gpu')
+
+            >>> batch = 1
+            >>> num_head = 8
+            >>> seq_len = 256
+            >>> head_size = 32
+
+            >>> dtype = paddle.float16
+
+            >>> query = paddle.randn([batch, num_head, seq_len, head_size], dtype=dtype)
+            >>> key = paddle.randn([batch, num_head, seq_len, head_size], dtype=dtype)
+            >>> value = paddle.randn([batch, num_head, seq_len, head_size], dtype=dtype)
+            >>> seq_lens = paddle.to_tensor(
+            ...     [seq_len] * batch,
+            ...     dtype='int32',
+            ... )
+            >>> mask = paddle.randn([batch, 1, seq_len, seq_len], dtype=dtype)
+
+            >>> scale = float(1.0 / math.sqrt(head_size))
+            >>> pre_cache_length = 0
+
+            >>> def naive_attention_impl(query, key, value, mask, scale):
+            ...     qk_res = paddle.matmul(query, key, transpose_y=True)
+            ...     attention = qk_res * scale
+            ...     attention = attention + mask
+            ...     softmax_result = paddle.nn.functional.softmax(attention, -1)
+            ...     result = paddle.matmul(softmax_result, value)
+            ...     return result
+
+            >>> out = naive_attention_impl(query, key, value, mask, scale)
+            >>> # equals to: out = variable_length_memory_efficient_attention(query, key, value, seq_lens, seq_lens, mask, scale, pre_cache_length)
+
+            >>> out.shape  # [batch, num_head, seq_len, head_size]
+            paddle.Size([1, 8, 256, 32])
+    """
+    if scale is None:
+        head_size = query.shape[3]
+        scale = float(1.0 / math.sqrt(head_size))
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.variable_length_memory_efficient_attention(
+            query,
+            key,
+            value,
+            seq_lens,
+            kv_seq_lens,
+            mask,
+            scale,
+            causal,
+            pre_cache_length,
+        )
+
+    helper = LayerHelper(
+        'variable_length_memory_efficient_attention', **locals()
+    )
+    out = helper.create_variable_for_type_inference(dtype=query.dtype)
+    helper.append_op(
+        type='variable_length_memory_efficient_attention',
+        inputs={
+            'query': query,
+            'key': key,
+            'value': value,
+            'seq_lens': seq_lens,
+            'kv_seq_lens': kv_seq_lens,
+            "mask": mask,
+        },
+        attrs={
+            "scale": scale,
+            "causal": causal,
+            "pre_cache_length": pre_cache_length,
+        },
+        outputs={'out': out},
+    )
+    return out

python/paddle/incubate/nn/layer/fused_dropout_add.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2023 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from paddle.incubate.nn import functional as F
+from paddle.nn import Layer
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+class FusedDropoutAdd(Layer):
+    r"""
+    Fused Dropout and Add.
+
+    Parameters:
+        p (float|int, optional): Probability of setting units to zero. Default: 0.5
+        mode(str, optional): ['upscale_in_train'(default) | 'downscale_in_infer']
+
+                               1. upscale_in_train (default), upscale the output at training time
+
+                                  - train: :math:`out = x \times \frac{mask}{(1.0 - p)} + y`
+                                  - inference: :math:`out = x + y`
+
+                               2. downscale_in_infer, downscale the output at inference
+
+                                  - train: :math:`out = x \times mask + y`
+                                  - inference: :math:`out = x \times (1.0 - p) + y`
+        name (str, optional): Name for the operation, Default: None. For more information, please refer to :ref:`api_guide_Name`.
+
+    Shape:
+        - x: N-D tensor.
+        - y: N-D tensor.
+        - output: N-D tensor, the same shape as x.
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+            >>> from paddle.incubate.nn.layer.fused_dropout_add import FusedDropoutAdd
+
+            >>> x = paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype="float32")
+            >>> y = paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype="float32")
+
+            >>> m = FusedDropoutAdd(p=0.5)
+
+            >>> out = m(x, y)
+    """
+
+    def __init__(
+        self,
+        p: float = 0.5,
+        mode: Literal[
+            'upscale_in_train', 'downscale_in_infer'
+        ] = "upscale_in_train",
+        name: str | None = None,
+    ) -> None:
+        super().__init__()
+        self.p = p
+        self.mode = mode
+        self.name = name
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        out = F.fused_dropout_add(
+            x,
+            y,
+            p=self.p,
+            training=self.training,
+            mode=self.mode,
+            name=self.name,
+        )
+        return out
+
+    def extra_repr(self) -> str:
+        name_str = f', name={self.name}' if self.name else ''
+        return f'p={self.p}, mode={self.mode}{name_str}'

python/paddle/incubate/nn/layer/fused_dropout_nd.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2023 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.
+
+import paddle
+from paddle import _legacy_C_ops
+from paddle.framework import in_dynamic_mode
+
+
+class FusedDropout(paddle.nn.Layer):
+    r"""
+    Dropout is a regularization technique for reducing overfitting by preventing
+    neuron co-adaption during training as described in the paper:
+    `Improving neural networks by preventing co-adaptation of feature detectors <https://arxiv.org/abs/1207.0580>`_
+    The dropout operator randomly sets the outputs of some units to zero, while upscale others
+    according to the given dropout probability.
+
+    It is an optimized implementation for ``paddle.nn.Dropout``.
+
+    In dygraph mode, please use ``eval()`` to switch to evaluation mode, where dropout is disabled.
+
+    Parameters:
+        p (float|int, optional): Probability of setting units to zero. Default: 0.5
+        axis (int|list|tuple, optional): The axis along which the dropout is performed. Default: None.
+        mode(str, optional): ['upscale_in_train'(default) | 'downscale_in_infer']
+
+                               1. upscale_in_train (default), upscale the output at training time
+
+                                  - train: :math:`out = input \times \frac{mask}{(1.0 - p)}`
+                                  - inference: :math:`out = input`
+
+                               2. downscale_in_infer, downscale the output at inference
+
+                                  - train: :math:`out = input \times mask`
+                                  - inference: :math:`out = input \times (1.0 - p)`
+        name (str, optional): Name for the operation, Default: None. For more information, please refer to :ref:`api_guide_Name`.
+
+    Shape:
+        - input: N-D tensor.
+        - output: N-D tensor, the same shape as input.
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.seed(2023)
+
+            >>> x = paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype="float32")
+            >>> m = paddle.incubate.nn.FusedDropout(p=0.5)
+
+            >>> y_train = m(x)
+            >>> print(y_train)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[0., 0., 6.],
+             [0., 0., 0.]])
+
+            >>> m.eval()  # switch the model to test phase
+            >>> y_test = m(x)
+            >>> print(y_test)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[1., 2., 3.],
+             [4., 5., 6.]])
+    """
+
+    def __init__(self, p=0.5, axis=None, mode="upscale_in_train", name=None):
+        super().__init__()
+
+        if not isinstance(p, (float, int)):
+            raise TypeError("p argument should be a number")
+        if p < 0 or p > 1:
+            raise ValueError("p argument should between 0 and 1")
+
+        mode = (
+            'downgrade_in_infer' if mode == 'downscale_in_infer' else mode
+        )  # semantic transfer
+        if mode not in ('downscale_in_infer', 'upscale_in_train'):
+            raise ValueError(
+                "mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
+            )
+
+        if axis and not isinstance(axis, (int, list, tuple)):
+            raise TypeError("datatype of axis argument should be int or list")
+
+        self.p = p
+        self.mode = mode
+        self.name = name
+
+        self.axis = None
+        if axis is not None:
+            self.axis = [axis] if isinstance(axis, int) else list(axis)
+
+    def forward(self, input):
+        # fast return for p == 0
+        if self.p == 0:
+            return input
+
+        if self.axis is not None and in_dynamic_mode():
+            seed = None
+            if paddle.static.default_main_program().random_seed != 0:
+                seed = paddle.static.default_main_program().random_seed
+
+            out, mask = _legacy_C_ops.dropout_nd(
+                input,
+                'dropout_prob',
+                self.p,
+                'is_test',
+                not self.training,
+                'fix_seed',
+                seed is not None,
+                'seed',
+                seed if seed is not None else 0,
+                'dropout_implementation',
+                self.mode,
+                'axis',
+                self.axis,
+            )
+        else:
+            out = paddle.nn.functional.dropout(
+                input,
+                p=self.p,
+                axis=self.axis,
+                training=self.training,
+                mode=self.mode,
+                name=self.name,
+            )
+        return out
+
+    def extra_repr(self):
+        name_str = f', name={self.name}' if self.name else ''
+        return f'p={self.p}, axis={self.axis}, mode={self.mode}{name_str}'

python/paddle/incubate/nn/layer/fused_linear.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2022 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle.incubate.nn import functional as F
+from paddle.nn import Layer
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import ParamAttrLike
+
+
+class FusedLinear(Layer):
+    r"""
+    Linear layer takes only one multi-dimensional tensor as input with the
+    shape :math:`[batch\_size, *, in\_features]` , where :math:`*` means any
+    number of additional dimensions. It multiplies input tensor with the weight
+    (a 2-D tensor of shape :math:`[in\_features, out\_features]` ) and produces
+    an output tensor of shape :math:`[batch\_size, *, out\_features]` .
+    If :math:`bias\_attr` is not False, the bias (a 1-D tensor of
+    shape :math:`[out\_features]` ) will be created and added to the output.
+
+    Parameters:
+        in_features (int): The number of input units.
+        out_features (int): The number of output units.
+        weight_attr (ParamAttr|None, optional): The attribute for the learnable
+            weight of this layer. The default value is None and the weight will be
+            initialized to zero. For detailed information, please refer to
+            paddle.ParamAttr.
+        transpose_weight (bool): Whether to transpose the `weight` Tensor before
+            multiplication.
+        bias_attr (ParamAttr|bool|None, optional): The attribute for the learnable bias
+            of this layer. If it is set to False, no bias will be added to the output.
+            If it is set to None or one kind of ParamAttr, a bias parameter will
+            be created according to ParamAttr. For detailed information, please refer
+            to paddle.ParamAttr. The default value is None and the bias will be
+            initialized to zero.
+        name (str|None, optional): Normally there is no need for user to set this parameter.
+            For detailed information, please refer to :ref:`api_guide_Name` .
+
+    Attribute:
+        **weight** (Parameter): the learnable weight of this layer.
+
+        **bias** (Parameter): the learnable bias of this layer.
+
+    Shape:
+        - input: Multi-dimensional tensor with shape :math:`[batch\_size, *, in\_features]` .
+        - output: Multi-dimensional tensor with shape :math:`[batch\_size, *, out\_features]` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+            >>> from paddle.incubate.nn import FusedLinear
+
+            >>> x = paddle.randn([3, 4])
+            >>> linear = FusedLinear(4, 5)
+            >>> y = linear(x)
+            >>> print(y.shape)
+            paddle.Size([3, 5])
+    """
+
+    weight: Tensor
+    bias: Tensor
+    transpose_weight: bool
+    name: str | None
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        transpose_weight: bool = False,
+        name: str | None = None,
+    ) -> None:
+        super().__init__()
+        if transpose_weight:
+            weight_shape = [out_features, in_features]
+        else:
+            weight_shape = [in_features, out_features]
+        dtype = self._helper.get_default_dtype()
+        self.weight = self.create_parameter(
+            shape=weight_shape, attr=weight_attr, dtype=dtype, is_bias=False
+        )
+        self.bias = self.create_parameter(
+            shape=[out_features], attr=bias_attr, dtype=dtype, is_bias=True
+        )
+        self.transpose_weight = transpose_weight
+        self.name = name
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.fused_linear(
+            input, self.weight, self.bias, self.transpose_weight, self.name
+        )

python/paddle/incubate/nn/layer/io.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2023 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.
+import paddle
+
+from ....base.framework import Variable
+from ....framework import LayerHelper, core
+
+
+class BlockGuardServ(paddle.static.nn.control_flow.BlockGuard):
+    """
+    BlockGuardServ class.
+
+    BlockGuardServ class is used to create an op with a block in a program.
+    """
+
+    def __init__(self, server):
+        if not (isinstance(server, ListenAndServ)):
+            raise TypeError("BlockGuardServ takes a ListenAndServ")
+        super().__init__(server.helper.main_program)
+        self.server = server
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            return False
+
+        self.server.complete_op()
+        return super().__exit__(exc_type, exc_val, exc_tb)
+
+
+class ListenAndServ:
+    """
+    **ListenAndServ Layer**
+
+    ListenAndServ is used to create a rpc server bind and listen
+    on specific TCP port, this server will run the sub-block when
+    received variables from clients.
+
+    Args:
+        endpoint(string): IP:port string which the server will listen on.
+        inputs(list): a list of variables that the server will get from clients.
+        fan_in(int): how many client are expected to report to this server, default: 1.
+        optimizer_mode(bool): whether to run the server as a parameter server, default: True.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> from paddle.incubate.nn.layer.io import ListenAndServ
+            >>> import paddle
+            >>> paddle.enable_static()
+            >>> place = paddle.CPUPlace()
+            >>> main = paddle.static.Program()
+            >>> with paddle.static.program_guard(main):
+            ...     serv = ListenAndServ("127.0.0.1:6170", ["X"], optimizer_mode=False)
+            ...     with serv.do():
+            ...         x = paddle.static.data(shape=[32, 32], dtype='float32', name="X")
+            ...         paddle.nn.initializer.Constant(value=1.0)(x, main.global_block())
+            ...         paddle.scale(x=x, scale=10.0)
+
+            >>> exe = paddle.static.Executor(place)
+            >>> exe.run(main)
+    """
+
+    def __init__(self, endpoint, inputs, fan_in=1, optimizer_mode=True):
+        self.helper = LayerHelper("listen_and_serv")
+        self.inputs = inputs
+        self.outputs = []
+        self.endpoint = endpoint
+        self.fan_in = fan_in
+        # FIXME(typhoonzero): add optimizer_mode is stupid, should make it more
+        # general.
+        self.optimizer_mode = optimizer_mode
+
+    def do(self):
+        return BlockGuardServ(self)
+
+    def get_params_and_grads(self):
+        main_program = self.helper.main_program
+        current_block = main_program.current_block()
+        parent_block = self.parent_block()
+        # params and grads in the same order.
+        params = []
+        grads = []
+        for op in current_block.ops:
+            # FIXME(typhoonzero): op.inputs is None if it's cloned.
+            if self.optimizer_mode:
+                if "Grad" in op.inputs and "Param" in op.inputs:
+                    params.append(op.inputs["Param"].name)
+                    grads.append(op.inputs["Grad"].name)
+            else:
+                # simple recv mode, recv operators inputs.
+                for iname in op.input_names:
+                    for in_var_name in op.input(iname):
+                        params.append(parent_block.var(in_var_name))
+                        grads.append(parent_block.var(in_var_name))
+
+        return params, grads
+
+    def parent_block(self):
+        prog = self.helper.main_program
+        parent_idx = prog.current_block().parent_idx
+        assert parent_idx >= 0
+        parent_block = prog.block(parent_idx)
+        return parent_block
+
+    def complete_op(self):
+        from paddle.incubate.distributed.fleet.parameter_server.mode import (
+            DistributedMode,
+        )
+
+        main_program = self.helper.main_program
+        current_block = main_program.current_block()
+        parent_block = self.parent_block()
+
+        parent_block.append_op(
+            type='listen_and_serv',
+            inputs={"X": self.inputs},
+            outputs={},
+            attrs={
+                'endpoint': self.endpoint,
+                'Fanin': self.fan_in,
+                'optimize_blocks': [
+                    current_block
+                ],  # did not support multiple optimize blocks in layers
+                'distributed_mode': DistributedMode.SYNC,  # did not support async now in layers
+                'grad_to_block_id': [""],
+            },
+        )
+
+
+def Send(endpoints, send_vars, dummy_output=None, sync=True):
+    """
+    Send variables to the server side, and get vars from server
+    side when server have finished running server side program.
+
+    Args:
+        endpoints (str): comma separated IP:PORT pairs in the order
+                   of send_vars to send
+        send_vars (list): variables to send to server
+        sync (bool): whether to wait the request finish
+
+    """
+    assert type(send_vars) == list
+
+    if dummy_output is None:
+        dummy_output = []
+    elif isinstance(dummy_output, Variable):
+        dummy_output = [dummy_output]
+
+    assert type(dummy_output) == list
+
+    epmap = endpoints.split(",")
+    endpoints = list(set(epmap))
+
+    helper = LayerHelper("Send", **locals())
+    rpc_op_role_name = core.op_proto_and_checker_maker.kOpRoleAttrName()
+
+    helper.append_op(
+        type="send",
+        inputs={"X": send_vars},
+        outputs={"Out": dummy_output},
+        attrs={
+            "endpoints": endpoints,
+            "epmap": epmap,
+            rpc_op_role_name: core.op_proto_and_checker_maker.OpRole.RPC,
+        },
+    )
+    if sync:
+        helper.append_op(
+            type="send_barrier",
+            inputs={"X": dummy_output},
+            outputs={"Out": []},
+            attrs={"endpoints": endpoints},
+        )
+
+
+def Recv(endpoints, get_vars, dummy_input=None, sync=True):
+    """
+    Receive variables from server side
+
+    Args:
+        endpoints (str): comma separated IP:PORT pairs in the order
+                   of send_vars to send
+        get_vars (list): vars to get from server after send completes.
+        sync (bool): whether to wait the request finish
+
+    Returns:
+        list: list of received variables
+    """
+    assert type(get_vars) == list
+
+    if dummy_input is None:
+        dummy_input = []
+    elif isinstance(dummy_input, Variable):
+        dummy_input = [dummy_input]
+
+    assert type(dummy_input) == list
+
+    epmap = endpoints.split(",")
+    endpoints = list(set(epmap))
+
+    helper = LayerHelper("Recv", **locals())
+    helper.append_op(
+        type="recv",
+        inputs={"X": dummy_input},
+        outputs={"Out": get_vars},
+        attrs={"endpoints": endpoints, "epmap": epmap},
+    )
+    if sync:
+        helper.append_op(
+            type="fetch_barrier",
+            outputs={"Out": get_vars},
+            attrs={"endpoints": endpoints},
+        )
+    return get_vars

python/paddle/incubate/nn/loss.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2019 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from paddle import _C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from typing import Literal, TypeAlias
+
+    from paddle import Tensor
+
+    _ReduceModeStringLiteral: TypeAlias = Literal['mean', 'sum', 'none']
+    _ReduceModeNumberLiteral: TypeAlias = Literal[0, 1, 2]
+    _ReduceMode: TypeAlias = _ReduceModeStringLiteral | _ReduceModeNumberLiteral
+
+
+def identity_loss(x: Tensor, reduction: _ReduceMode = "none") -> Tensor:
+    r"""Marks a tensor as being part of the loss calculation for IPU.
+
+    This operator is used to handle on the (final) loss of a model so that
+    it is used as the start of backpropagation.
+
+    When `reduction` is `none`, return raw `Out`.
+
+    When `reduction` is `mean`, return
+
+    .. math::
+        Out = MEAN(Out)
+
+    When `reduction` is `sum`, return
+
+    .. math::
+        Out = SUM(Out)
+
+    Parameters:
+        x (Variable): The input tensor. The shapes is [N, *], where N is batch size and `*` means any number of
+             additional dimensions. It's data type should be float32, float64 on CPU and float16, float32 on IPU.
+        reduction(str|int, optional): Reduce the loss output. Supported string values are: 'sum', 'mean', 'none'
+                            the corresponding int values are 0, 1, 2 respectively. The default value is "none".
+
+    Returns:
+        Variable: The loss ``Tensor`` with the specified reduction applied.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.enable_static()
+            >>> loss = paddle.static.data(name="loss", shape=[-1, 1], dtype="float32")
+            >>> out = paddle.incubate.identity_loss(loss, reduction=1)
+    """
+    if isinstance(reduction, str):
+        reduction = {"sum": 0, "mean": 1, "none": 2}.get(reduction.lower())
+        if reduction is None:
+            raise TypeError("Unsupported reduction type.")
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.identity_loss(x, reduction)
+
+    check_variable_and_dtype(x, 'x', ['float32', 'float64'], "identity_loss")
+    attrs = {'reduction': reduction}
+    helper = LayerHelper('identity_loss', **locals())
+    dtype = helper.input_dtype(input_param_name='x')
+    out = helper.create_variable_for_type_inference(dtype)
+    helper.append_op(
+        type="identity_loss", inputs={"X": x}, outputs={"Out": out}, attrs=attrs
+    )
+    return out

python/paddle/incubate/nn/memory_efficient_attention.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2023 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.
+
+# The following codes are from https://github.com/facebookresearch/xformers
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import paddle
+from paddle import _C_ops
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+from .attn_bias import (
+    BlockDiagonalCausalMask,
+    BlockDiagonalCausalWithOffsetPaddedKeysMask,
+    BlockDiagonalMask,
+    LowerTriangularMask,
+    LowerTriangularMaskWithTensorBias,
+)
+
+SUPPORTED_ATTN_BIAS_TYPES = {
+    type(None),
+    paddle.Tensor,
+    paddle.pir.Value,
+    LowerTriangularMask,
+    LowerTriangularMaskWithTensorBias,
+    BlockDiagonalMask,
+    BlockDiagonalCausalMask,
+    BlockDiagonalCausalWithOffsetPaddedKeysMask,
+}
+
+
+def _get_seqlen_info(attn_bias):
+    if isinstance(
+        attn_bias,
+        (BlockDiagonalMask, BlockDiagonalCausalWithOffsetPaddedKeysMask),
+    ):
+        return (
+            attn_bias.k_seqinfo.seqstart,
+            attn_bias.q_seqinfo.seqstart,
+            attn_bias.q_seqinfo.max_seqlen,
+            attn_bias.k_seqinfo.max_seqlen,
+        )
+    else:
+        return None, None, -1, -1
+
+
+def _get_tensor_bias(attn_bias):
+    if isinstance(attn_bias, (paddle.Tensor, paddle.pir.Value)):
+        return attn_bias
+    elif isinstance(attn_bias, LowerTriangularMaskWithTensorBias):
+        return attn_bias._bias
+    else:
+        return None
+
+
+def memory_efficient_attention(
+    query, key, value, attn_bias=None, p=0.0, scale=None, training=True
+):
+    assert type(attn_bias) in SUPPORTED_ATTN_BIAS_TYPES
+    causal = isinstance(
+        attn_bias,
+        (
+            LowerTriangularMask,
+            BlockDiagonalCausalMask,
+            BlockDiagonalCausalWithOffsetPaddedKeysMask,
+        ),
+    )
+    seqstart_k, seqstart_q, max_seqlen_q, max_seqlen_k = _get_seqlen_info(
+        attn_bias
+    )
+    # NOTE: compute_logsumexp = training
+    causal_diagonal = (
+        attn_bias.causal_diagonal
+        if isinstance(attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask)
+        else None
+    )
+    seqlen_k = (
+        attn_bias.k_seqinfo.seqlen
+        if isinstance(attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask)
+        else None
+    )
+    if scale is None:
+        scale = -1.0
+
+    bias = _get_tensor_bias(attn_bias)
+    is_test = not training
+
+    if in_dynamic_or_pir_mode():
+        output, logsumexp, seed_and_offset = _C_ops.memory_efficient_attention(
+            query,
+            key,
+            value,
+            bias,
+            seqstart_q,
+            seqstart_k,
+            causal_diagonal,
+            seqlen_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            causal,
+            p,
+            scale,
+            is_test,
+        )
+        return output
+
+    helper = LayerHelper('memory_efficient_attention', **locals())
+    output = helper.create_variable_for_type_inference(dtype=query.dtype)
+    logsumexp = helper.create_variable_for_type_inference(dtype='float')
+    seed_and_offset = helper.create_variable_for_type_inference(dtype='int32')
+    helper.append_op(
+        type='memory_efficient_attention',
+        inputs={
+            'query': query,
+            'key': key,
+            'value': value,
+            'bias': bias,
+            "cu_seqlens_q": seqstart_q,
+            "cu_seqlens_k": seqstart_k,
+            "causal_diagonal": causal_diagonal,
+            "seqlen_k": seqlen_k,
+        },
+        attrs={
+            "max_seqlen_q": max_seqlen_q,
+            "max_seqlen_k": max_seqlen_k,
+            "causal": causal,
+            "dropout_p": p,
+            "scale": scale,
+            "is_test": is_test,
+        },
+        outputs={
+            'output': output,
+            'logsumexp': logsumexp,
+            "seed_and_offset": seed_and_offset,
+        },
+    )
+
+    return output