Cherry-pick PRs from incubate/paddlenlp_fleety (#11049)

* add offload opt util (#10607)

* fix offload optimizer (#10608)

* fix sharding reshard bug (#10613)

* reset ema (#10790)

* Move fused_quanted_ops and token_dispatcher_utils to FleetY (#10803)

* move setup_fp8.py (#10820)

* Add tokens_zip_unique_add_subbatch and merge_subbatch_cast ops (#10822)

* add tokens_zip_unique_add_subbatch and merge_subbatch_cast ops

* enhance ut

* enhance ut again

* remove duplicate codes

* enhance ut again

* Sharding reshard supports mismatch parameter name (#10479)

* Sharding Reshard Supports N -> 1 (#10532)

* cherry-pick https://github.com/PaddlePaddle/PaddleNLP/pull/10531

* support tpdp-ep sharding reshard (#10568)

* fix typo

* revert sharding_first

.pre-commit-config.yaml

@@ -62,3 +62,12 @@ repos:
         language: python
         files: \.(md|markdown|rst)$
         pass_filenames: true
+
+-   repo: local
+    hooks:
+    -   id: clang-format
+        name: clang-format
+        description: Format files with ClangFormat.
+        entry: bash ./tools/codestyle/clang_format.sh -i
+        language: system
+        files: \.(c|cc|cxx|cpp|cu|h|hpp|hxx|xpu|kps)$

paddlenlp/trainer/trainer.py

@@ -372,7 +372,13 @@ class Trainer:
         self.optimizer_grouped_parameters = None
         self.sharding_io = None
         if self.args.should_save_sharding_stage1_model or self.args.should_load_sharding_stage1_model:
-            self.sharding_io = ShardingIO(self.args, self.model, self.optimizer)
+            self.sharding_io = ShardingIO(
+                self.args,
+                self.model,
+                self.optimizer,
+                remap_parameter_name=self.args.load_sharded_model_remap_parameter_name,
+            )
+
         if self.args.unified_checkpoint:
             self.unified_checkpoint_handler = UnifiedCheckpointHandler(self.args)
 
@@ -805,9 +811,16 @@ class Trainer:
         if resume_from_checkpoint is not None:
             path = _add_variant(PADDLE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
             path = os.path.join(resume_from_checkpoint, path).replace("optimizer", "ema")
+            if self.args.zcc_save_ema_coef is not None and self.sharding_io is not None:
+                success, err_msg = self.sharding_io.check_same_strategy(resume_from_checkpoint)
+            else:
+                success, err_msg = True, None
             if os.path.exists(path):
-                logger.info(f"ZCC EMA load from {path}")
-                self.zcc_manager.set_ema_state_dict(path)
+                if success:
+                    logger.info(f"ZCC EMA load from {path}")
+                    self.zcc_manager.set_ema_state_dict(path)
+                else:
+                    logger.info(f"ZCC EMA does not load {path} because {err_msg}")
             else:
                 logger.info(f"ZCC EMA state dict not found, in: {path}")
 

paddlenlp/trainer/training_args.py

@@ -633,6 +633,11 @@ class TrainingArguments:
         },
     )
 
+    load_sharded_model_remap_parameter_name: bool = field(
+        default=False,
+        metadata={"help": "Whether to remap parameter name when load_sharded_model = true."},
+    )
+
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={
@@ -2039,6 +2044,11 @@ class TrainingArguments:
                 sharding_parallel_degree * tensor_parallel_degree * sep_parallel_degree * pipeline_parallel_degree
             )
 
+            if expert_parallel_degree > 1:
+                assert (
+                    self.expert_tensor_parallel_degree <= 1
+                ), "expert_tensor_parallel_degree > 1 is not supported when expert_parallel_degree > 1"
+
             assert not (
                 self.data_parallel_degree > 1 and expert_parallel_degree > 1
             ), f"Currently only support use expert_data_parallel strategy together with sharding_parallel strategy, but not with data_parallel strategy. Currently data_parallel_degree is {self.data_parallel_degree}."
@@ -2227,6 +2237,17 @@ class TrainingArguments:
         else:
             return 0
 
+    @property
+    def expert_parallel_rank(self):
+        if self.use_hybrid_parallel:
+            hcg = fleet.get_hybrid_communicate_group()
+            if hasattr(hcg, "get_expert_parallel_rank"):
+                return max(hcg.get_expert_parallel_rank(), 0)
+            else:
+                return 0
+        else:
+            return 0
+
     @property
     def context_parallel_rank(self):
         if self.use_hybrid_parallel:
@@ -2252,7 +2273,7 @@ class TrainingArguments:
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
             if self.sharding_parallel_degree > 1:
                 name.append(self._format_name("shard", self.sharding_parallel_rank, self.sharding_parallel_degree))
-            if self.use_expert_parallel:
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
             return "_".join(name)
         else:
@@ -2268,7 +2289,7 @@ class TrainingArguments:
                 name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
-            if self.use_expert_parallel:
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
             return "_".join(name)
 
@@ -2277,7 +2298,9 @@ class TrainingArguments:
                 return self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree)
             return None
 
-    def sharded_name_suffix(self, shard_id=None, pp_id=None, moe_id=None):
+    def sharded_name_suffix(self, shard_id=None, pp_id=None, moe_id=None, sharding_parallel_degree=None):
+        if sharding_parallel_degree is None:
+            sharding_parallel_degree = self.sharding_parallel_degree
         if self.use_hybrid_parallel:
             name = []
             if self.tensor_parallel_degree > 1:
@@ -2287,12 +2310,12 @@ class TrainingArguments:
                     pp_id = self.pipeline_parallel_rank
                 assert isinstance(pp_id, int)
                 name.append(self._format_name("pp", pp_id, self.pipeline_parallel_degree))
-            if self.sharding_parallel_degree > 1:
+            if sharding_parallel_degree > 1:
                 if shard_id is None:
                     shard_id = self.sharding_parallel_rank
                 assert isinstance(shard_id, int)
-                name.append(self._format_name("shard", shard_id, self.sharding_parallel_degree))
-            if self.use_expert_parallel:
+                name.append(self._format_name("shard", shard_id, sharding_parallel_degree))
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 if moe_id is None:
                     moe_id = self.data_parallel_rank
                 assert isinstance(moe_id, int)
@@ -2418,9 +2441,7 @@ class TrainingArguments:
     def should_load_sharding_stage1_model(self):
         if self.enable_auto_parallel:
             return False
-        return (
-            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.load_sharded_model
-        )
+        return self.load_sharded_model
 
     @property
     def should_load_dataset(self):

paddlenlp/trainer/utils/offload_optimizer.py

@@ -0,0 +1,81 @@
+# 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.distributed.fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer import (
+    HybridParallelOptimizer,
+)
+from paddle.optimizer import Optimizer
+
+from .sharding_io import to_device
+
+
+def offload(tensor):
+    if paddle.is_compiled_with_cuda():
+        place = paddle.CUDAPinnedPlace()
+    else:
+        place = paddle.CPUPlace()
+
+    new_tensor = to_device(tensor, place)
+    assert new_tensor is tensor, "to_device must be inplace operation"
+
+
+def reload(tensor):
+    new_tensor = to_device(tensor)
+    assert new_tensor is tensor, "to_device must be inplace operation"
+
+
+def hack_offload_optimizer():
+    # Step 1: mock _add_accumulator
+    origin_add_accumulator = getattr(Optimizer, "_add_accumulator")
+
+    def new_add_accumulator(self, *args, **kwargs):
+        x = origin_add_accumulator(self, *args, **kwargs)
+        offload(x)
+        return x
+
+    setattr(Optimizer, "_add_accumulator", new_add_accumulator)
+
+    # Step 2: mock _C_ops.adamw_ and _C_ops.adamw
+    for name in ["adam_", "adamw_"]:
+        origin_op = getattr(_C_ops, name)
+
+        def new_opt_op(*args):
+            for arg in args:
+                if isinstance(arg, paddle.Tensor):
+                    reload(arg)
+
+            ret = origin_op(*args)
+
+            for i, arg in enumerate(args):
+                if i >= 2 and isinstance(arg, paddle.Tensor):  # do not offload parameter and gradient
+                    offload(arg)
+            return ret
+
+        setattr(_C_ops, name, new_opt_op)
+
+    # Step 3: mock _insert_sync
+    opt_type = HybridParallelOptimizer
+    origin_insert_sync = getattr(opt_type, "_insert_sync")
+
+    def new_insert_sync(self, sync_var, *args, **kwargs):
+        origin_place = sync_var.place
+        reload(sync_var)
+        ret = origin_insert_sync(self, sync_var, *args, **kwargs)
+        new_sync_var = to_device(sync_var, origin_place)
+        assert new_sync_var is sync_var, "to_device must be inplace operation"
+        return ret
+
+    setattr(opt_type, "_insert_sync", new_insert_sync)

paddlenlp/trainer/utils/reshard/__init__.py

@@ -18,6 +18,14 @@ from .common import (
     SHARDING_STRATEGY_V2,
     NodeModelState,
     all_gather_state_dict,
+    convert_opt_name_to_tname,
+    get_moe_sharding_group,
+    get_param_sharding_group,
     get_sharding_strategy,
     is_sharding_opt,
+    merge_model_state,
+    merge_opt_state,
+    split_model_state,
+    split_opt_state,
+    split_structure_name_mapping,
 )

paddlenlp/trainer/utils/reshard/common.py

@@ -16,11 +16,14 @@ from collections import OrderedDict
 
 import numpy as np
 import paddle
+import paddle.distributed.fleet as fleet
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
     DygraphShardingOptimizer,
 )
 from paddle.distributed.fleet.utils.log_util import logger
 
+from paddlenlp.utils.tools import get_env_device
+
 try:
     from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
         DygraphShardingOptimizerV2,
@@ -60,18 +63,60 @@ def get_sharding_strategy(optimizer):
     return SHARDING_STRATEGY_V1
 
 
+def convert_opt_name_to_tname(tensor_names, opt_names):
+    tensor_names = set(tensor_names)
+    all_names = []
+    all_names.extend(list(tensor_names))
+    all_names.extend(opt_names)
+    all_names.sort()
+    pre_t_name = ""
+    suffix = [
+        "_fp32_master_0_beta1_pow_acc_0",
+        "_fp32_master_0_beta2_pow_acc_0",
+        "_fp32_master_0_moment1_0",
+        "_fp32_master_0_moment2_0",
+        "_beta1_pow_acc_0",
+        "_beta2_pow_acc_0",
+        "_moment1_0",
+        "_moment2_0",
+    ]
+    opt_to_t = {}
+    for n in all_names:
+        if n in tensor_names:
+            # we get a param
+            pre_t_name = n
+        else:
+            assert pre_t_name
+            opt_to_t[n] = pre_t_name
+
+    for t in opt_names:
+        _find = False
+        for s in suffix:
+            if get_env_device() == "xpu" and t.endswith(s + ".SCALE_VALUE"):
+                # NOTE: for xpu adamw, all optimizer state will have an extra attribute end with SCALE_VALUE.
+                # This extra attribute won't be used, just skip it.
+                _find = True
+                break
+            if t.endswith(s):
+                logger.info(f"{t}-{t[:-len(s)]}--{t[:-len(s)] in tensor_names}")
+                opt_to_t[t] = t[: -len(s)]
+                _find = True
+                break
+        assert _find
+    return opt_to_t
+
+
 class NodeModelState:
-    def __init__(self, mp_rank=None, sharding_rank=None, pp_rank=None):
+    def __init__(self, group):
         self._model_weights = OrderedDict()
         self._opt_state = OrderedDict()
         self._master_weights = OrderedDict()
         self._lr_scheduler = None
-        self.set_node_rank(mp_rank, sharding_rank, pp_rank)
+        self._group = group
 
-    def set_node_rank(self, mp_rank, sharding_rank, pp_rank):
-        self._mp_rank = mp_rank
-        self._sharding_rank = sharding_rank
-        self._pp_rank = pp_rank
+    @property
+    def group(self):
+        return self._group
 
     def _add_kv(self, d, k, v):
         assert k not in d
@@ -259,43 +304,6 @@ class NodeModelState:
         change the key of master weights dict from param_name to (structure_name, param_name)
         """
         # pack key for pp convert
-        def _opt_name_to_tname(tensor_names, opt_names):
-            tensor_names = set(tensor_names)
-            all_names = []
-            all_names.extend(list(tensor_names))
-            all_names.extend(opt_names)
-            all_names.sort()
-            pre_t_name = ""
-            suffix = [
-                "_fp32_master_0_beta1_pow_acc_0",
-                "_fp32_master_0_beta2_pow_acc_0",
-                "_fp32_master_0_moment1_0",
-                "_fp32_master_0_moment2_0",
-                "_beta1_pow_acc_0",
-                "_beta2_pow_acc_0",
-                "_moment1_0",
-                "_moment2_0",
-            ]
-            opt_to_t = {}
-            for n in all_names:
-                if n in tensor_names:
-                    # we get a param
-                    pre_t_name = n
-                else:
-                    assert pre_t_name
-                    opt_to_t[n] = pre_t_name
-
-            for t in opt_names:
-                _find = False
-                for s in suffix:
-                    if t.endswith(s):
-                        logger.info(f"{t}-{t[:-len(s)]}--{t[:-len(s)] in tensor_names}")
-                        opt_to_t[t] = t[: -len(s)]
-                        _find = True
-                        break
-                assert _find
-            return opt_to_t
-
         if structure_name_mapping is not None:
             tname_to_structure_name = {v: k for (k, v) in structure_name_mapping.items()}
         else:
@@ -304,7 +312,7 @@ class NodeModelState:
 
         tensor_names = list(tname_to_structure_name.keys())
         opt_names = list(self._opt_state.keys())
-        opt_name_to_tname = _opt_name_to_tname(tensor_names, opt_names)
+        opt_name_to_tname = convert_opt_name_to_tname(tensor_names, opt_names)
 
         # model state
         model_weights_tmp = OrderedDict()
@@ -399,12 +407,13 @@ class NodeModelState:
 
         return node_model_states
 
-    def even_distribute(self, group):
+    def even_distribute(self):
         """
         distribute the node state evenly among all workers in group， and make sure
         in the dicts of (key, rank)=>tensor, items keys of the same key but different rank are distributed to the
         same worker
         """
+        group = self.group
         # sharding degree == 1
         if group is None or group.nranks < 2:
             return self
@@ -438,7 +447,7 @@ class NodeModelState:
             def filter_func(key):
                 assert key[0] in key_to_rank, key
                 dst_rank = key_to_rank[key[0]]
-                return dst_rank == group.rank
+                return dst_rank == max(group.rank, 0)
 
             return _all_gather_state_dict(state_dict, filter_func, group)
 
@@ -447,10 +456,11 @@ class NodeModelState:
         self._master_weights = distribute(self._master_weights)
         return self
 
-    def reshard(self, group, filter_func):
+    def reshard(self, filter_func):
         """
         reshard according to the passed in filter_func
         """
+        group = self.group
         self._model_weights = _all_gather_state_dict(self._model_weights, filter_func, group)
         self._opt_state = _all_gather_state_dict(self._opt_state, filter_func, group)
         self._master_weights = _all_gather_state_dict(self._master_weights, filter_func, group)
@@ -503,6 +513,7 @@ class NodeModelState:
         return self
 
     def merge_from(self, other, rank=None):
+        assert other.group is self.group
         self.add_weights(other.model_weights, rank)
         self.add_opts(other.opt_state, rank)
         self.add_master_weights(other.master_weights, rank)
@@ -520,6 +531,68 @@ class NodeModelState:
         return opt_state_dict
 
 
+def split_model_state(model_state, group_getter):
+    res = OrderedDict()
+    for k, v in model_state.items():
+        group = group_getter.get_group(k)
+        if group.id not in res:
+            res[group.id] = OrderedDict()
+        res[group.id][k] = v
+    return res
+
+
+def merge_model_state(model_state_map):
+    res = OrderedDict()
+    for gid, model_state in model_state_map.items():
+        res.update(model_state)
+    return res
+
+
+def split_opt_state(opt_state, group_getter):
+    res = OrderedDict()
+    lr_scheduler = opt_state.get("LR_Scheduler", None)
+    for k, v in opt_state.items():
+        if k == "LR_Scheduler":
+            continue
+        elif k == "master_weights":
+            for kk, vv in v.items():
+                group = group_getter.get_group(kk)
+                if group.id not in res:
+                    res[group.id] = {"master_weights": OrderedDict(), "LR_Scheduler": lr_scheduler}
+                res[group.id]["master_weights"][kk] = vv
+        else:
+            assert isinstance(v, paddle.Tensor), type(v)
+            group = group_getter.get_group(k)
+            if group.id not in res:
+                res[group.id] = {"master_weights": OrderedDict(), "LR_Scheduler": lr_scheduler}
+            res[group.id][k] = v
+    return res
+
+
+def merge_opt_state(opt_state_map):
+    res = {"LR_Scheduler": None, "master_weights": OrderedDict()}
+    for gid, opt_state in opt_state_map.items():
+        for k, v in opt_state.items():
+            if k == "LR_Scheduler":
+                if v is not None:
+                    res["LR_Scheduler"] = v
+            elif k == "master_weights":
+                res["master_weights"].update(v)
+            else:
+                res[k] = v
+    return res
+
+
+def split_structure_name_mapping(structure_name_mapping, group_getter):
+    res = OrderedDict()
+    for k, v in structure_name_mapping.items():
+        group = group_getter.get_group(k)
+        if group.id not in res:
+            res[group.id] = OrderedDict()
+        res[group.id][k] = v
+    return res
+
+
 def all_gather_simple_object(obj, group):
     res = []
     if group.nranks < 2:
@@ -536,12 +609,13 @@ def all_gather_state_dict(state_dict, filter_func, group):
             weight = weight.numpy()
         return weight
 
+    group_rank = max(group.rank, 0)
     state_dict = {k: map_func(v) for (k, v) in state_dict.items()}
 
     meta_dict = {}
     for (k, v) in state_dict.items():
         # src rank
-        meta_dict[k] = (v.dtype, v.shape, group.rank)
+        meta_dict[k] = (v.dtype, v.shape, group_rank)
 
     meta_dict_list = all_gather_simple_object(meta_dict, group)
 
@@ -555,20 +629,21 @@ def all_gather_state_dict(state_dict, filter_func, group):
     meta_list = sorted(meta_list, key=lambda x: x[0])
     for (k, meta) in meta_list:
         dtype, shape, rank = meta
-        if rank == group.rank:
+        if rank == group_rank:
             assert k in state_dict
             tensor = paddle.to_tensor(state_dict[k])
             del state_dict[k]
         else:
             tensor = paddle.to_tensor(np.empty(shape, dtype))
-        logger.info(f"broadcast {k} from {rank}")
+        logger.info(f"broadcast {k} from {rank}, group {group}")
         # broadcast the tensor
-        paddle.distributed.broadcast(
-            tensor,
-            src=group.ranks[rank],
-            group=group,
-            sync_op=True,
-        )
+        if group.nranks > 1:
+            paddle.distributed.broadcast(
+                tensor,
+                src=group.ranks[rank],
+                group=group,
+                sync_op=True,
+            )
         if filter_func(k):
             res[k] = tensor.cpu()
         del tensor
@@ -585,3 +660,29 @@ def _all_gather_state_dict(state_dict, filter_func, group):
     for (k, v) in tmp_state_dict.items():
         state_dict[k] = v
     return state_dict
+
+
+def get_moe_sharding_group(hcg=None):
+    if hcg is None:
+        hcg = fleet.get_hybrid_communicate_group()
+    if hasattr(hcg, "get_moe_sharding_parallel_group"):
+        return hcg.get_moe_sharding_parallel_group()
+    else:
+        return None
+
+
+def get_param_sharding_group(param, hcg=None):
+    if hcg is None:
+        hcg = fleet.get_hybrid_communicate_group()
+    default_group = hcg.get_sharding_parallel_group()
+    ep_sharding_group = get_moe_sharding_group(hcg)
+
+    if not hasattr(param, "color"):
+        return default_group
+    color = getattr(param, "color")
+    if isinstance(color, dict):
+        group = color.get("group", default_group)
+        assert group is default_group or group is ep_sharding_group, f"unsupported group: {group}"
+        return group
+    else:
+        return default_group

paddlenlp/trainer/utils/reshard/sharding_v1.py

@@ -17,26 +17,32 @@ from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding
 )
 
 from ....transformers.model_utils import unwrap_optimizer
+from .common import is_sharding_opt
 
 
-def shard(node_model_state, model, optimizer, hcg):
-    group = hcg.get_sharding_parallel_group()
-    cur_rank = group.rank
-    optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
-    assert optimizer is not None
-    param2rank = optimizer._param2rank
+def shard(node_model_state, model, optimizer):
+    cur_rank = max(node_model_state.group.rank, 0)
+    unwrapped_optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
+    if unwrapped_optimizer is not None:
+        optimizer = unwrapped_optimizer
+        assert not is_sharding_opt(optimizer)
+        param2rank = optimizer._param2rank
 
-    def filter_func(key):
-        names = key
-        param_name = names[1]
-        assert param_name in param2rank
-        dst_rank = param2rank[param_name]
-        return dst_rank == cur_rank
+        def filter_func(key):
+            names = key
+            param_name = names[1]
+            assert param_name in param2rank
+            dst_rank = param2rank[param_name]
+            return dst_rank == cur_rank
 
-    node_model_state.reshard(group, filter_func)
+    else:
+        assert not is_sharding_opt(optimizer)
+        filter_func = lambda key: True
+
+    node_model_state.reshard(filter_func)
     return node_model_state
 
 
-def restore(node_model_state, model, optimizer, hcg):
+def restore(node_model_state, model, optimizer):
     node_model_state.drop_rank()
     return node_model_state

paddlenlp/trainer/utils/reshard/sharding_v2.py

@@ -34,13 +34,16 @@ except:
 
 from paddle.distributed.communication.reduce import ReduceOp
 
+from .common import get_moe_sharding_group
 
-def shard(node_model_state, model, optimizer, hcg):
+
+def shard(node_model_state, model, optimizer):
     assert DygraphShardingOptimizerV2 is not None
-    group = hcg.get_sharding_parallel_group()
-    cur_rank = group.rank
     split_infos = collect_split_info(optimizer, model)
 
+    group = node_model_state.group
+    cur_rank = max(group.rank, 0)
+
     def split_func(k, v):
         param_name = k[1]
         opt_name = k[-1]
@@ -87,15 +90,14 @@ def shard(node_model_state, model, optimizer, hcg):
         return rank == cur_rank
 
     # reshard
-    node_model_state.reshard(group, filter_func)
+    node_model_state.reshard(filter_func)
     node_model_state.drop_rank()
     return node_model_state
 
 
-def restore(node_model_state, model, optimizer, hcg):
-    group = hcg.get_sharding_parallel_group()
+def restore(node_model_state, model, optimizer):
     # evenly distribute param
-    node_model_state.even_distribute(group)
+    node_model_state.even_distribute()
     param_shapes = {k: v.shape for (k, v) in model.state_dict().items()}
 
     def merge_func(k, v):
@@ -175,7 +177,7 @@ def collect_split_info(optimizer, model, only_return_lengths=False):
         for comm_buffer in optimizer._comm_buffer_list:
             gather_infos(comm_buffer)
 
-    assert len(split_infos)
+    assert len(split_infos) > 0
     return split_infos
 
 
@@ -211,11 +213,16 @@ def is_matched_optimizer_state_dict(opt_state_dict, optimizer, model, hcg=None,
     if need_allgather:
         if hcg is None:
             hcg = fleet.get_hybrid_communicate_group()
-        group = hcg.get_sharding_parallel_group()
-        if group is not None and group.nranks > 1:
-            x = paddle.to_tensor([is_matched], dtype=paddle.int32)
-            paddle.distributed.stream.all_reduce(x, op=ReduceOp.MIN, group=group, sync_op=True, use_calc_stream=True)
-            global_is_matched = int(x.numpy()[0])
+        sharding_group = hcg.get_sharding_parallel_group()
+        moe_sharding_group = get_moe_sharding_group(hcg)
+        for group in [sharding_group, moe_sharding_group]:
+            if group is not None and group.nranks > 1:
+                x = paddle.to_tensor([is_matched], dtype=paddle.int32)
+                paddle.distributed.stream.all_reduce(
+                    x, op=ReduceOp.MIN, group=group, sync_op=True, use_calc_stream=True
+                )
+                is_matched = int(x.numpy()[0])
+        global_is_matched = is_matched
     else:
         global_is_matched = is_matched
 

paddlenlp/trainer/utils/sharding_io.py

@@ -17,6 +17,7 @@ import json
 import os
 from collections import OrderedDict
 
+import numpy
 import paddle
 import paddle.distributed as dist
 from paddle.distributed import fleet
@@ -42,7 +43,17 @@ from paddlenlp.utils.log import logger
 from paddlenlp.utils.tools import get_env_device
 
 from . import reshard as reshard_util
-from .reshard import SHARDING_STRATEGY_V1, SHARDING_STRATEGY_V2, pp_reshard
+from .reshard import (
+    SHARDING_STRATEGY_V1,
+    SHARDING_STRATEGY_V2,
+    get_param_sharding_group,
+    merge_model_state,
+    merge_opt_state,
+    pp_reshard,
+    split_model_state,
+    split_opt_state,
+    split_structure_name_mapping,
+)
 
 
 def to_device(tensor, place=None):
@@ -63,7 +74,7 @@ def to_device(tensor, place=None):
 
 def filter_sharded_params(state_dict, optimizer, sharding_group, include_freeze_params=False):
 
-    sharding_rank = sharding_group.rank
+    sharding_rank = max(sharding_group.rank, 0)
     sharding_world_size = sharding_group.nranks
     from paddlenlp.trainer.utils import reshard as reshard_util
 
@@ -112,7 +123,10 @@ def exclude_parameters_in_state_dict(
     )
     # allgather parameter names in sharding group
     tmp = []
-    paddle.distributed.all_gather_object(tmp, param_names_in_master_weights, group=sharding_group)
+    if sharding_group.nranks > 1:
+        paddle.distributed.all_gather_object(tmp, param_names_in_master_weights, group=sharding_group)
+    else:
+        tmp = [param_names_in_master_weights]
     param_names_in_master_weights = set([v for item in tmp for v in item])
     logger.info("sharding_group_param_names:{}".format(param_names_in_master_weights))
     non_parameters_state_dict = copy.copy(model_state_dict)
@@ -123,8 +137,140 @@ def exclude_parameters_in_state_dict(
     return non_parameters_state_dict
 
 
+class ParameterNameRemapper:
+    def __init__(self, old_mapping, new_mapping, checkpoint):
+        self.checkpoint = checkpoint
+        self.p_name_map = {}
+        for k, v in old_mapping.items():
+            assert k in new_mapping, f"structure name not found: {k} {new_mapping.keys()}"
+            new_v = new_mapping[k]
+            if v not in self.p_name_map:
+                self.p_name_map[v] = new_v
+            else:
+                old_v = self.p_name_map[v]
+                assert old_v == new_v, f"structure name {k} has different parameter name {new_v} {self.p_name_map[v]}"
+        self.old_p_names = list(self.p_name_map.keys())
+        self.new_mapping = dict([(k, new_mapping[k]) for k in old_mapping.keys()])
+
+    def _map_tensor(self, tensor, old_p_name=None):
+        if old_p_name is not None:
+            new_p_name = self.p_name_map.get(old_p_name)
+            assert new_p_name is not None, f"parameter name {old_p_name} not found"
+        else:
+            new_p_name = None
+
+        def _map_name(old_name):
+            if new_p_name is not None:
+                assert old_name.startswith(old_p_name)
+                return new_p_name + old_name[len(old_p_name) :]
+            else:
+                new_name = self.p_name_map.get(old_name)
+                assert new_name is not None, f"parameter name {old_name} not found"
+                return new_name
+
+        if isinstance(tensor, paddle.Tensor):
+            new_name = _map_name(tensor.name)
+            tensor.name = new_name
+            return new_name, tensor
+        else:
+            assert isinstance(tensor, (list, tuple)), type(tensor)
+            old_name, value = tensor
+            new_name = self._map_name(old_name)
+            return new_name, (new_name, value)
+
+    def remap_model_state(self, model_state):
+        for k, v in model_state.items():
+            if not isinstance(v, numpy.ndarray):
+                model_state[k] = self._map_tensor(v)[1]
+        return model_state
+
+    def remap_optimizer_state(self, opt_state):
+        lr_scheduler_key = "LR_Scheduler"
+        master_weight_key = "master_weights"
+
+        new_opt_state = {}
+        new_master_weights = None
+        opt_names = []
+        for k, v in opt_state.items():
+            if k == lr_scheduler_key:
+                new_opt_state[k] = v
+            elif k == master_weight_key:
+                for kk, vv in v.items():
+                    new_kk = self.p_name_map[kk]
+                    if new_master_weights is None:
+                        new_opt_state[master_weight_key] = {}
+                        new_master_weights = new_opt_state[master_weight_key]
+                    new_master_weights[new_kk] = self._map_tensor(vv, kk)[1]
+            else:
+                assert isinstance(v, paddle.Tensor), type(v)
+                opt_names.append(v.name)
+
+        opt_to_pname = reshard_util.convert_opt_name_to_tname(self.old_p_names, opt_names)
+        for opt_name in opt_names:
+            v = opt_state[opt_name]
+            new_opt_name, new_v = self._map_tensor(v, opt_to_pname[opt_name])
+            new_opt_state[new_opt_name] = new_v
+
+        opt_state.clear()
+        opt_state.update(new_opt_state)
+        return opt_state
+
+
+class GroupGetter:
+    def __init__(self, model, hcg=None):
+        self.structure_name_mapping = {}
+        self.structure_name_to_group = {}
+        self.tensor_name_to_group = {}
+        self.parameter_names = []
+        self.group_map = OrderedDict()
+        self.hcg = hcg or fleet.get_hybrid_communicate_group()
+        for k, v in model.state_dict().items():
+            self.structure_name_mapping[k] = v.name
+            group = get_param_sharding_group(v, self.hcg)
+            self.structure_name_to_group[k] = group
+            self.tensor_name_to_group[v.name] = group
+            self.group_map[group.id] = group
+            self.parameter_names.append(v.name)
+
+    def _get_parameter_name(self, name):
+        if name in self.tensor_name_to_group:
+            return name
+
+        suffix = [
+            "_fp32_master_0_beta1_pow_acc_0",
+            "_fp32_master_0_beta2_pow_acc_0",
+            "_fp32_master_0_moment1_0",
+            "_fp32_master_0_moment2_0",
+            "_beta1_pow_acc_0",
+            "_beta2_pow_acc_0",
+            "_moment1_0",
+            "_moment2_0",
+        ]
+
+        for s in suffix:
+            if name.endswith(s):
+                tmp = name[: -len(s)]
+                assert tmp in self.tensor_name_to_group, f"cannot find {name}"
+                return tmp
+
+        raise ValueError(f"cannot find {name}")
+
+    def get_group(self, name):
+        if name in self.structure_name_to_group:
+            assert name not in self.tensor_name_to_group, name
+            return self.structure_name_to_group[name]
+        else:
+            return self.tensor_name_to_group[self._get_parameter_name(name)]
+
+    def get_group_by_id(self, gid):
+        return self.group_map[gid]
+
+    def get_group_ids(self):
+        return list(self.group_map.keys())
+
+
 class ShardingIO:
-    def __init__(self, args, model, optimizer=None, hcg=None):
+    def __init__(self, args, model, optimizer=None, hcg=None, remap_parameter_name=False):
         self.args = args
         self.model = model
         self.optimizer = optimizer
@@ -134,6 +280,33 @@ class ShardingIO:
             self.hcg = fleet.get_hybrid_communicate_group()
             self.sharding_group = self.hcg.get_sharding_parallel_group()
 
+        self.remap_parameter_name = remap_parameter_name
+        self.remapper = None
+
+    def _get_remapper(self, checkpoint):
+        if not self.remap_parameter_name:
+            return None
+
+        if self.remapper is None or self.remapper.checkpoint != checkpoint:
+            new_mapping = {}
+            for k, v in self.model.state_dict().items():
+                new_mapping[k] = v.name
+
+            suffix = self._sharding_meta_suffix()
+            model_meta = self._load_model_meta_impl(checkpoint)
+            old_mapping = model_meta["sharding_metas"][suffix]["structure_name_mapping"]
+            self.remapper = ParameterNameRemapper(old_mapping, new_mapping, checkpoint)
+        return self.remapper
+
+    def _remap_parameter_name(self, checkpoint, state_dict, is_opt):
+        remapper = self._get_remapper(checkpoint)
+        if remapper is None:
+            return state_dict
+        if is_opt:
+            return remapper.remap_optimizer_state(state_dict)
+        else:
+            return remapper.remap_model_state(state_dict)
+
     def set_optimizer(self, optimizer):
         self.optimizer = optimizer
 
@@ -146,6 +319,9 @@ class ShardingIO:
             base_weight_name (str): The name of the checkpoint file.
             model_wrapped (nn.Layer): The wrapped model.
         """
+        group_getter = GroupGetter(self.model)
+        gids = group_getter.get_group_ids()
+
         parallel_config = self._load_distributed_strategy(checkpoint)
         pp_degree = parallel_config["pp_degree"]
         mp_degree = parallel_config["mp_degree"]
@@ -161,19 +337,25 @@ class ShardingIO:
             assert cur_pp_degree <= 1, "can not reshard from non pp to pp"
 
         def load_model_slices():
-            model_state = reshard_util.NodeModelState()
+            model_state = {gid: reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid)) for gid in gids}
             for j in range(self.args.pipeline_parallel_rank, pp_degree, cur_pp_degree):
                 cur_sharding_meta = self._load_sharding_meta(checkpoint, j)
                 assert "structure_name_mapping" in cur_sharding_meta
                 structure_name_map = cur_sharding_meta["structure_name_mapping"]
+                structure_name_map = split_structure_name_mapping(structure_name_map, group_getter)
                 for i in range(self.args.sharding_parallel_rank, sharding_degree, cur_sharding_degree):
                     tmp = self._load_one_state_dict_from_checkpoint(
-                        checkpoint, base_weight_name, self.args.sharded_name_suffix(i, j)
+                        checkpoint,
+                        base_weight_name,
+                        self.args.sharded_name_suffix(i, j, sharding_parallel_degree=sharding_degree),
                     )
-                    node_model_state_tmp = reshard_util.NodeModelState()
-                    node_model_state_tmp.add_weights(tmp)
-                    node_model_state_tmp.pack_keys(structure_name_map)
-                    model_state.merge_from(node_model_state_tmp, i)
+                    tmp = split_model_state(tmp, group_getter)
+                    for gid in gids:
+                        sub_tmp = tmp.get(gid, {})
+                        node_model_state_tmp = reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid))
+                        node_model_state_tmp.add_weights(sub_tmp)
+                        node_model_state_tmp.pack_keys(structure_name_map.get(gid, {}))
+                        model_state[gid].merge_from(node_model_state_tmp, i)
             return model_state
 
         node_model_state = load_model_slices()
@@ -183,19 +365,31 @@ class ShardingIO:
             reshard_context = pp_reshard.build_pipeline_context(meta, model_wrapped)
             node_model_state = pp_reshard.reshard(node_model_state, reshard_context, self.hcg)
 
-        node_model_state.drop_rank()
-        node_model_state.unpack_keys()
-        state_dict = node_model_state.model_weights
+        if opt_state_dict is None:
+            opt_state_dict = self.optimizer.state_dict()
+        opt_state_dict = split_opt_state(opt_state_dict, group_getter)
 
-        def filter_func(name):
-            return True
+        res_state_dict = OrderedDict()
+        for gid, nms in node_model_state.items():
+            nms.drop_rank()
+            nms.unpack_keys()
+            state_dict = nms.model_weights
 
-        state_dict = reshard_util.all_gather_state_dict(state_dict, filter_func, self.sharding_group)
+            def filter_func(name):
+                return True
 
-        if self.args.bf16:
-            state_dict = self._recover_params_from_master_weights(state_dict, opt_state_dict=opt_state_dict)
+            state_dict = reshard_util.all_gather_state_dict(state_dict, filter_func, nms.group)
 
-        return state_dict
+            if self.args.bf16:
+                state_dict = self._recover_params_from_master_weights(
+                    state_dict,
+                    opt_state_dict=opt_state_dict.get(gid, {}),
+                    group=group_getter.get_group_by_id(gid),
+                )
+
+            res_state_dict.update(state_dict)
+
+        return res_state_dict
 
     def _load_one_state_dict_from_checkpoint(self, resume_from_checkpoint, base_weight_name, weight_name_suffix):
         """
@@ -208,14 +402,19 @@ class ShardingIO:
         logger.info(f"Loading model from {resume_from_checkpoint} .")
         # We load the model state dict on the CPU to avoid an OOM error.
         state_dict = paddle.load(file_path, return_numpy=True)
+        state_dict = self._remap_parameter_name(resume_from_checkpoint, state_dict, is_opt=False)
         return state_dict
 
-    def _load_optimizer_state_of_one_shard(self, checkpoint, base_opt_name, optimizer_name_suffix):
+    def _load_optimizer_state_of_one_shard(self, checkpoint, base_opt_name, optimizer_name_suffix, group_getter=None):
         optimizer_name = _add_variant(base_opt_name, optimizer_name_suffix)
         path = os.path.join(checkpoint, optimizer_name)
         logger.info(f"load optimizer state from {path}")
         if os.path.isfile(path):
-            return self._modify_ckpt_for_compatibility(paddlenlp_load(path, map_location="cpu"))
+            return self._remap_parameter_name(
+                checkpoint,
+                self._modify_ckpt_for_compatibility(paddlenlp_load(path, map_location="cpu")),
+                is_opt=True,
+            )
         logger.info(f"{path} not exists")
         return None
 
@@ -246,9 +445,12 @@ class ShardingIO:
         if sharding_strategy == SHARDING_STRATEGY_V1:
             param2rank = sharding_meta["param2rank"]
             optimizer = unwrap_optimizer(self.optimizer, DygraphShardingOptimizer)
-            assert optimizer
-            if len(param2rank) == 0:
-                logger.warning("The param2rank is empty. Force reshard would be performed.")
+            if self.args.sharding_parallel_degree > 1:
+                assert optimizer is not None
+            else:
+                assert optimizer is None
+            if len(param2rank) == 0 or optimizer is None:
+                logger.warning("The param2rank is empty or sharding degree is 1. Force reshard would be performed.")
                 return True
             assert len(param2rank) == len(optimizer._param2rank)
             for (k, v) in param2rank.items():
@@ -279,10 +481,11 @@ class ShardingIO:
         """load state_dict of multiple shard from_checkpoint, Only load model state dict."""
 
         parallel_config = self._load_distributed_strategy(checkpoint)
-        sharding_meta = self._load_sharding_meta(checkpoint, 0)
+        sharding_meta = self._load_sharding_meta(checkpoint)
         pp_degree = parallel_config["pp_degree"]
         mp_degree = parallel_config["mp_degree"]
         sharding_degree = parallel_config["sharding_degree"]
+        assert sharding_degree > 1, "sharding degree of the checkpoint should be larger than 1"
         sharding_strategy = SHARDING_STRATEGY_V1
         if "sharding_strategy" in sharding_meta:
             sharding_strategy = sharding_meta["sharding_strategy"]
@@ -297,9 +500,14 @@ class ShardingIO:
         cur_sharding_degree = self.args.sharding_parallel_degree
         cur_sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
 
+        group_getter = GroupGetter(self.model)
+
         if not self._need_reshard(checkpoint):
             one_shard_opt_state_dict = self._load_optimizer_state_of_one_shard(
-                checkpoint, base_opt_name, self.args.optimizer_name_suffix
+                checkpoint,
+                base_opt_name,
+                self.args.sharded_name_suffix(sharding_parallel_degree=sharding_degree),
+                group_getter=group_getter,
             )
 
             if sharding_strategy == SHARDING_STRATEGY_V2 and cur_sharding_strategy == SHARDING_STRATEGY_V2:
@@ -316,15 +524,17 @@ class ShardingIO:
             one_shard_opt_state_dict = None
 
         logger.info("reshard optimizer state")
+        gids = group_getter.get_group_ids()
 
         def load_model_slices():
-            model_state = reshard_util.NodeModelState()
+            model_state = {gid: reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid)) for gid in gids}
             for j in range(self.args.pipeline_parallel_rank, pp_degree, cur_pp_degree):
                 cur_sharding_meta = self._load_sharding_meta(checkpoint, j)
                 assert "structure_name_mapping" in cur_sharding_meta
                 structure_name_map = cur_sharding_meta["structure_name_mapping"]
+                structure_name_map = split_structure_name_mapping(structure_name_map, group_getter)
                 for i in range(self.args.sharding_parallel_rank, sharding_degree, cur_sharding_degree):
-                    sharded_name_suffix = self.args.sharded_name_suffix(i, j)
+                    sharded_name_suffix = self.args.sharded_name_suffix(i, j, sharding_parallel_degree=sharding_degree)
                     if one_shard_opt_state_dict is None:
                         tmp = self._load_optimizer_state_of_one_shard(checkpoint, base_opt_name, sharded_name_suffix)
                     else:
@@ -332,18 +542,25 @@ class ShardingIO:
                             self.args.optimizer_name_suffix == sharded_name_suffix
                         ), f"{self.args.optimizer_name_suffix} vs {sharded_name_suffix}"
                         tmp = one_shard_opt_state_dict
-                    node_model_state_tmp = reshard_util.NodeModelState()
-                    node_model_state_tmp.add_opts(tmp)
-                    node_model_state_tmp.pack_keys(structure_name_map)
-                    model_state.merge_from(node_model_state_tmp, i)
+
+                    tmp = split_opt_state(tmp, group_getter)
+                    for gid in gids:
+                        sub_tmp = tmp.get(gid, {})
+                        node_model_state_tmp = reshard_util.NodeModelState(group=group_getter.get_group_by_id(gid))
+                        node_model_state_tmp.add_opts(sub_tmp)
+                        node_model_state_tmp.pack_keys(structure_name_map.get(gid, {}))
+                        model_state[gid].merge_from(node_model_state_tmp, i)
             return model_state
 
         def reshard_pp(model_state):
             # pp reshard
             if self._need_reshard_pp(checkpoint):
+                assert len(model_state) == 1, "only support one group reshard"
+                key = list(model_state.keys())[0]
+                tmp = model_state[key]
                 meta = self._load_model_meta(checkpoint)
                 reshard_context = pp_reshard.build_pipeline_context(meta, model_wrapped)
-                model_state = pp_reshard.reshard(model_state, reshard_context, self.hcg)
+                model_state = {key: pp_reshard.reshard(tmp, reshard_context, self.hcg)}
             return model_state
 
         def reshard_sharding(node_model_state):
@@ -353,17 +570,23 @@ class ShardingIO:
                 if sharding_strategy == SHARDING_STRATEGY_V1
                 else reshard_util.sharding_v2.restore
             )
-            node_model_state = restore_func(node_model_state, self.model, self.optimizer, self.hcg)
+
+            for gid in gids:
+                node_model_state[gid] = restore_func(node_model_state[gid], self.model, self.optimizer)
 
             shard_func = (
                 reshard_util.sharding_v1.shard
                 if cur_sharding_strategy == SHARDING_STRATEGY_V1
                 else reshard_util.sharding_v2.shard
             )
-            node_model_state = shard_func(node_model_state, model_wrapped, self.optimizer, self.hcg)
-            # drop structural name in the key
-            node_model_state.unpack_keys()
-            return node_model_state.get_opt_state_dict()
+
+            ret_opt_state_dict = OrderedDict()
+            for gid in gids:
+                node_model_state[gid] = shard_func(node_model_state[gid], model_wrapped, self.optimizer)
+                # drop structural name in the key
+                node_model_state[gid].unpack_keys()
+                ret_opt_state_dict[gid] = node_model_state[gid].get_opt_state_dict()
+            return merge_opt_state(ret_opt_state_dict)
 
         node_model_state = load_model_slices()
         node_model_state = reshard_pp(node_model_state)
@@ -371,16 +594,21 @@ class ShardingIO:
 
     def manipulate_state_dict_and_config(self, model_to_save, merge_tensor_parallel=False, state_dict=None):
         weight_name_suffix = self.args.sharded_name_suffix()
+        group_getter = GroupGetter(model_to_save)
+        gids = group_getter.get_group_ids()
 
         if state_dict is None:
             state_dict = model_to_save.state_dict()
             if self.args.should_save_sharding_stage1_model:
-                state_dict = filter_sharded_params(
-                    state_dict,
-                    self.optimizer,
-                    self.sharding_group,
-                    self.args.save_sharding_stage1_model_include_freeze_params,
-                )
+                state_dict = split_model_state(state_dict, group_getter)
+                for gid in gids:
+                    state_dict[gid] = filter_sharded_params(
+                        state_dict.get(gid, {}),
+                        self.optimizer,
+                        self.sharding_group,
+                        self.args.save_sharding_stage1_model_include_freeze_params,
+                    )
+                state_dict = merge_model_state(state_dict)
 
         config_to_save = None
         merge_tensor_parallel = merge_tensor_parallel and self.args.use_hybrid_parallel
@@ -402,11 +630,17 @@ class ShardingIO:
         if self.args.bf16 and self.args.should_save_sharding_stage1_model:
             param_names_in_master_weights = []
             optimzier_state_dict = self.optimizer.state_dict()
-            assert "master_weights" in optimzier_state_dict
-            param_names_in_master_weights = list(optimzier_state_dict["master_weights"].keys())
-            state_dict = exclude_parameters_in_state_dict(
-                state_dict, param_names_in_master_weights, self.sharding_group
-            )
+            optimzier_state_dict = split_opt_state(optimzier_state_dict, group_getter)
+            state_dict = split_model_state(state_dict, group_getter)
+            for gid in gids:
+                sub_opt_state = optimzier_state_dict.get(gid, {})
+                param_names_in_master_weights = list(sub_opt_state.get("master_weights", {}).keys())
+                state_dict[gid] = exclude_parameters_in_state_dict(
+                    state_dict.get(gid, {}),
+                    param_names_in_master_weights,
+                    group_getter.get_group_by_id(gid),
+                )
+            state_dict = merge_model_state(state_dict)
             logger.info(
                 "param_names_in_master_weights len:{}, bf16 state_dict len:{}, :{}".format(
                     len(param_names_in_master_weights), len(state_dict), state_dict.keys()
@@ -431,31 +665,59 @@ class ShardingIO:
 
         return model_meta
 
+    def _check_distributed_strategy(self, parallel_config):
+        ep_degree = parallel_config.get("ep_degree", 1)
+        if ep_degree > 1:
+            tp_degree = parallel_config["mp_degree"]
+            sharding_degree = parallel_config["sharding_degree"]
+            moe_sharding_degree = parallel_config.get("moe_sharding_degree", 1)
+            assert tp_degree * sharding_degree == ep_degree * moe_sharding_degree, "mismatch parallel degree settings"
+
+    def check_same_strategy(self, resume_from_checkpoint=None):
+        if resume_from_checkpoint:
+            cur_config = self._get_distributed_strategy()
+            old_config = self._load_model_meta_impl(resume_from_checkpoint)["parallel_config"]
+            keys = list(old_config.keys())
+            for key in keys:
+                if key not in cur_config:
+                    return False, f"missing {key}"
+                else:
+                    old_value = old_config[key]
+                    cur_value = cur_config[key]
+                    if old_value != cur_value:
+                        return False, f"{key} not match: {old_value} vs {cur_value}"
+        return True, None
+
     def _get_distributed_strategy(self):
         pp_degree = 1
         mp_degree = 1
         sharding_degree = 1
-        vpp_degree = 1
+        ep_degree = 1
+        moe_sharding_degree = 1
         nranks = dist.get_world_size()
         if self.args.use_hybrid_parallel and nranks > 1:
             hcg = fleet.get_hybrid_communicate_group()
             mp_degree = hcg.get_model_parallel_world_size()
             pp_degree = hcg.get_pipe_parallel_world_size()
             sharding_degree = hcg.get_sharding_parallel_world_size()
-            """
-            if pp_degree > 1:
-                assert isinstance(model, fleet.meta_parallel.PipelineParallel), "must be pipeline model"
-                vpp_degree = model._layers.get_num_virtual_stages()
-            """
+            if hasattr(hcg, "get_expert_parallel_world_size"):
+                ep_degree = hcg.get_expert_parallel_world_size()
+            if hasattr(hcg, "get_moe_sharding_parallel_world_size"):
+                moe_sharding_degree = hcg.get_moe_sharding_parallel_world_size()
+
         parallel_config = {
             "pp_degree": pp_degree,
             "mp_degree": mp_degree,
             "sharding_degree": sharding_degree,
-            "vpp_degree": vpp_degree,
+            "ep_degree": ep_degree,
+            "moe_sharding_degree": moe_sharding_degree,
         }
+        self._check_distributed_strategy(parallel_config)
         return parallel_config
 
-    def _recover_params_from_master_weights(self, state_dict, opt_state_dict=None):
+    def _recover_params_from_master_weights(self, state_dict, opt_state_dict=None, group=None):
+        if group is None:
+            group = self.sharding_group
         if opt_state_dict is None:
             opt_state_dict = self.optimizer.state_dict()
         assert "master_weights" in opt_state_dict, opt_state_dict.keys()
@@ -469,27 +731,26 @@ class ShardingIO:
             master_weights[k].name = name
 
         structure_name_map = {k: v.name for (k, v) in self.model.state_dict().items()}
-        node_model_state = reshard_util.NodeModelState()
-        node_model_state_tmp = reshard_util.NodeModelState()
+        node_model_state = reshard_util.NodeModelState(group=group)
+        node_model_state_tmp = reshard_util.NodeModelState(group=group)
         node_model_state_tmp.add_master_weights(master_weights)
         node_model_state_tmp.pack_keys(structure_name_map)
-        node_model_state.merge_from(node_model_state_tmp, self.sharding_group.rank)
+        node_model_state.merge_from(node_model_state_tmp, max(group.rank, 0))
         del node_model_state_tmp
-        assert reshard_util.is_sharding_opt(self.optimizer)
         sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
         restore_func = (
             reshard_util.sharding_v1.restore
             if sharding_strategy == SHARDING_STRATEGY_V1
             else reshard_util.sharding_v2.restore
         )
-        node_model_state = restore_func(node_model_state, self.model, self.optimizer, self.hcg)
+        node_model_state = restore_func(node_model_state, self.model, self.optimizer)
         node_model_state.unpack_keys()
         master_weights = node_model_state.master_weights
 
         def filter_func(name):
             return True
 
-        master_weights = reshard_util.all_gather_state_dict(master_weights, filter_func, self.sharding_group)
+        master_weights = reshard_util.all_gather_state_dict(master_weights, filter_func, group)
         model_state_dict = self.model.state_dict()
         logger.info(f"state-dict-keys: {state_dict.keys()}, nums: {len(state_dict.keys())}")
         logger.info("before recover, model_state_dict number: {}".format(len(model_state_dict)))
@@ -517,14 +778,46 @@ class ShardingIO:
         paddle.distributed.all_gather_object(res, obj, group)
         return res
 
-    def _load_model_meta(self, dir):
+    def _load_model_meta_impl(self, dir):
         meta_path = os.path.join(dir, MODEL_META_NAME)
         assert os.path.exists(meta_path), f"{meta_path} not exist"
         with open(meta_path, "r") as handle:
             model_dist_meta = json.load(handle)
+
         assert "parallel_config" in model_dist_meta
+        self._check_distributed_strategy(model_dist_meta["parallel_config"])
         return model_dist_meta
 
+    def _load_model_meta(self, dir):
+        model_meta = self._load_model_meta_impl(dir)
+        remapper = self._get_remapper(dir)
+        if remapper is not None:
+            suffix = self._sharding_meta_suffix()
+            sharding_metas = model_meta["sharding_metas"]
+            cur_sharding_metas = sharding_metas.pop(suffix)
+            sharding_metas.clear()
+            sharding_metas[suffix] = cur_sharding_metas
+            cur_sharding_metas["structure_name_mapping"] = remapper.new_mapping
+            if "param2rank" in cur_sharding_metas:
+                new_param2rank = {}
+                for k, rank in cur_sharding_metas["param2rank"].items():
+                    new_k = remapper.p_name_map[k]
+                    new_param2rank[new_k] = rank
+                cur_sharding_metas["param2rank"] = new_param2rank
+        return model_meta
+
+    def _sharding_meta_suffix(self, tp_rank=None, pp_rank=None):
+        if tp_rank is None:
+            tp_rank = self.args.tensor_parallel_rank
+        if pp_rank is None:
+            pp_rank = self.args.pipeline_parallel_rank
+        suffix = f"tp{tp_rank:0>2d}_pp{pp_rank:0>2d}"
+        if self.args.expert_parallel_degree > 1:
+            ep_rank = self.args.expert_parallel_rank
+            return f"{suffix}_ep{ep_rank:0>2d}"
+        else:
+            return suffix
+
     def _load_distributed_strategy(self, dir):
         model_dist_meta = self._load_model_meta(dir)
         parallel_config = model_dist_meta["parallel_config"]
@@ -534,9 +827,7 @@ class ShardingIO:
         return parallel_config
 
     def _load_sharding_meta(self, dir, pp_rank=None):
-        if pp_rank is None:
-            pp_rank = self.args.pipeline_parallel_rank
-        suffix = f"tp{self.args.tensor_parallel_rank:0>2d}_pp{pp_rank:0>2d}"
+        suffix = self._sharding_meta_suffix(pp_rank=pp_rank)
         distributed_model_meta = self._load_model_meta(dir)
         if "sharding_metas" in distributed_model_meta:
             sharding_metas = distributed_model_meta["sharding_metas"]
@@ -600,12 +891,16 @@ class ShardingIO:
         sharding_meta["param2rank"] = param2rank
         sharding_meta["structure_name_mapping"] = structure_name_mapping
         sharding_meta["param_meta"] = param_meta
+        sharding_meta["param_meta_keys"] = ["shape", "dtype", "is_distributed", "no_sync"]
         sharding_meta["sharding_strategy"] = sharding_strategy
         sharding_meta["enable_overlap"] = pp_overlap
-        suffix = f"tp{self.args.tensor_parallel_rank:0>2d}_pp{self.args.pipeline_parallel_rank:0>2d}"
+        suffix = self._sharding_meta_suffix()
         sharding_metas[suffix] = sharding_meta
         sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_model_parallel_group())
         sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
         sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_pipe_parallel_group())
         sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
+        if self.args.expert_parallel_degree > 1:
+            sharding_metas_list = self._all_gather_simple_object(sharding_metas, self.hcg.get_expert_parallel_group())
+            sharding_metas = {k: v for e in sharding_metas_list for (k, v) in e.items()}
         return sharding_metas

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_act_dequant.cu

@@ -0,0 +1,123 @@
+#include "quant_utils.h"
+
+
+// cuda结构体拷贝
+template <typename T, int N>
+struct alignas(16) VectorType {
+  T data[N];
+};
+
+__global__ void FusedActDequant(
+    const phi::float8_e4m3fn*__restrict__ Xin,
+    const float *__restrict__ Xscale,
+    phi::bfloat16*__restrict__ out,
+    const int rows,
+    const int cols
+ ) {
+  const int this_row_idx = blockIdx.x;
+  if (this_row_idx >= rows) return;
+
+  const int Xscale_stride = (cols + 127) / 128;  // 计算缩放因子的步长
+
+  const int vector_size = 16;  // 向量的元素数量，处理16个元素
+
+  // 每行的向量数量
+  const int num_vectors = cols / vector_size;
+  const int remaining_elements = cols % vector_size;
+
+  const int tid = threadIdx.x;
+
+  for (int vec_idx = tid; vec_idx < num_vectors; vec_idx += blockDim.x) {
+    int x_offset = vec_idx * vector_size;
+    int64_t X_idx = (int64_t)this_row_idx * (int64_t)cols + (int64_t)x_offset;
+
+    // 加载16个 __nv_fp8_e4m3 元素到向量中
+    const VectorType<__nv_fp8_e4m3, vector_size>* X_vec_ptr =
+      reinterpret_cast<const VectorType<__nv_fp8_e4m3, vector_size>*>(Xin + X_idx);
+    VectorType<__nv_fp8_e4m3, vector_size> X_vec = X_vec_ptr[0];
+
+    // 获取对应的缩放因子
+    int64_t scale_idx = (int64_t)this_row_idx * (int64_t)Xscale_stride + (x_offset / 128);
+    float this_scale = Xscale[scale_idx];
+
+    // 初始化输出向量
+    VectorType<__nv_bfloat16, vector_size> out_vec;
+
+    // 逐元素处理向量中的数据
+    #pragma unroll
+    for (int i = 0; i < vector_size; ++i) {
+      // 将fp8转换为float
+      float X_value = static_cast<float>(X_vec.data[i]);
+      // 乘以缩放因子
+      X_value *= this_scale;
+      // 转换为bfloat16并存储
+      out_vec.data[i] = __float2bfloat16(X_value);
+    }
+
+    // 将输出向量存储到全局内存
+    VectorType<__nv_bfloat16, vector_size>* out_vec_ptr =
+        reinterpret_cast<VectorType<__nv_bfloat16, vector_size>*>(out + X_idx);
+    out_vec_ptr[0] = out_vec;
+  }
+
+  // 处理剩余不能被向量化的元素
+  if (remaining_elements > 0) {
+    int x_offset = num_vectors * vector_size;
+    int64_t X_idx = (int64_t)this_row_idx * (int64_t)cols + (int64_t)x_offset;
+    int64_t idx = X_idx + tid;
+    if (tid < remaining_elements) {
+      float X_value = static_cast<float>(Xin[idx]);
+      X_value *= Xscale[(int64_t)this_row_idx * (int64_t)Xscale_stride + (x_offset / 128)];
+      out[idx] = __float2bfloat16(X_value);
+    }
+  }
+}
+
+void dispatch_fused_act_dequant(
+    const paddle::Tensor &X,
+    const paddle::Tensor &Xscale,
+    paddle::Tensor &out,
+    const int rows,
+    const int cols) {
+  dim3 grid;
+  dim3 block;
+  grid.x = X.shape()[0];
+  block.x=256;
+  FusedActDequant<<<grid, block, 0, X.stream()>>>(
+    X.data<phi::float8_e4m3fn>(),
+    Xscale.data<float>(),
+    out.data<phi::bfloat16>(),
+    rows,
+    cols
+  );
+}
+std::vector<paddle::Tensor> fused_act_dequant(
+    const paddle::Tensor &X,
+    const paddle::Tensor &Xscale) {
+  PD_CHECK(X.dtype() == paddle::DataType::FLOAT8_E4M3FN);
+  PD_CHECK(Xscale.dtype() == paddle::DataType::FLOAT32);
+  int rows, cols;
+  rows = X.shape()[0];
+  cols = X.shape()[1];
+  paddle::Tensor out;
+  
+  out = paddle::empty({rows, cols}, paddle::DataType::BFLOAT16, X.place());
+  auto out_ptr =
+    reinterpret_cast<void *>(out.data<phi::bfloat16>());
+  cudaMemsetAsync(out_ptr,
+                  0,
+                  sizeof(phi::bfloat16) * rows * cols,
+                  out.stream());
+  dispatch_fused_act_dequant(
+      X,
+      Xscale,
+      out,
+      rows,
+      cols);
+  return {out};
+}
+
+PD_BUILD_OP(fused_act_dequant)
+    .Inputs({"X", "Xscale"})
+    .Outputs({"output"})
+    .SetKernelFn(PD_KERNEL(fused_act_dequant));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_act_dequant_transpose_act_quant.cu

@@ -0,0 +1,202 @@
+#include "quant_utils.h"
+
+constexpr int64_t TILE_SIZE = 128;  // 每个 block 处理 128x128 的元素块
+#define LAUNCH_FUSED_ACT_DEQUANT_TRANSPOSE_ACT_QUANT(__using_pow2_scaling,     \
+                                                     __padding_last_dim_to_8x) \
+  do {                                                                         \
+    auto kernel = FusedActDequantTransposeActQuant<outT,                       \
+                                                   __using_pow2_scaling,       \
+                                                   __padding_last_dim_to_8x>;  \
+    int smem_size = 128 * 129 * sizeof(float);                                 \
+    PD_CHECK(cudaFuncSetAttribute(kernel,                                      \
+                                  cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                                  smem_size) == cudaSuccess);                  \
+    kernel<<<grid, block, smem_size, X.stream()>>>(                            \
+        X.data<phi::float8_e4m3fn>(),                                          \
+        Xscale.data<float>(),                                                  \
+        out.data<outT>(),                                                      \
+        scale.data<float>(),                                                   \
+        rows,                                                                  \
+        cols,                                                                  \
+        TYPE_MAX);                                                             \
+  } while (0)
+
+
+#define BLOCK_SIZE 128
+template <typename OutT, bool using_pow2_scaling, bool padding_last_dim_to_8x>
+__global__ void FusedActDequantTransposeActQuant(
+    const phi::float8_e4m3fn *__restrict__ Xin,
+    const float *__restrict__ Xscale,
+    OutT *__restrict__ out,
+    float *__restrict__ scales,
+    const int rows,  // 此处的row未带padding
+    const int cols,  // 同上
+    const int TYPE_MAX) {
+  // 共享内存用于:1.直接读取+存储转置后的数据  2.用于1x128规约+量化
+  // 共享内存布局,前128x128为计算结果，后128为scale(复用)
+  extern __shared__ __align__(1024) uint8_t smem_buffer[];
+  float *smem_tile = reinterpret_cast<float *>(smem_buffer);
+  float *smem_max =
+      reinterpret_cast<float *>(smem_buffer) + BLOCK_SIZE * BLOCK_SIZE;
+
+  const int g_block_y_offset = blockIdx.y * BLOCK_SIZE;  // 块内行坐标偏移
+  const int g_block_x_offset = blockIdx.x * BLOCK_SIZE;  // 块内列坐标偏移
+  const __nv_fp8_e4m3 *X = reinterpret_cast<const __nv_fp8_e4m3 *>(Xin);
+  // 阶段0:
+  // 原始fp8 scale读入smem_max，用于后续dequant
+  // ------------------------------
+  if (threadIdx.y == 0) {
+    for (int i = threadIdx.x; i < BLOCK_SIZE; i += blockDim.x) {
+      smem_max[i] = Xscale[i];
+    }
+  }
+  __syncthreads();  // smem_tile中的Xscale数据已ready
+
+  // 阶段1:
+  // Elementwise加载数据、计算反量化、并将结果直接以特定布局存入共享内存‌(32x32)
+  // ------------------------------
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    float row_scale = smem_max[y_offset];
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const int in_y_idx = g_block_y_offset + y_offset;
+      const int in_x_idx = g_block_x_offset + x_offset;
+      const int src_idx = in_y_idx * cols + in_x_idx;
+      // 类型提升、dequant、共享内存swizzle写数
+      if (in_y_idx < rows && in_x_idx < cols) {
+        const int src_idx = in_y_idx * cols + in_x_idx;
+        smem_tile[swizzled_2d_idx(x_offset, BLOCK_SIZE, y_offset)] =
+            static_cast<float>(X[src_idx]) * row_scale;
+      }
+    }
+  }
+  __syncthreads();  // smem_tile中的swiglu数据(按需transpose)已ready
+
+  // 阶段2: ‌
+  // Shared 两级reduce, 给出每行的absmax_f
+  // ------------------------------------------------------
+  float local_max = 0.0f;
+  // 每个warp计算一行最大值，每个warp处理4行
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    // 行quantize max计算, 每行128个元素，每个线程处理4个元素
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      bool is_output_inner_OOB;
+      is_output_inner_OOB = (g_block_y_offset + x_offset) >= rows;
+      if (is_output_inner_OOB) break;  // 列越界则不取值、不影响max;
+      local_max = fmaxf(
+          local_max,
+          fabsf(smem_tile[swizzled_2d_idx(
+              y_offset,
+              BLOCK_SIZE,
+              x_offset)]));  // 正常情况下提供绝对值给其他线程,每个线程的local_max最多为4数最大值
+    }
+    bool is_output_outer_OOB;
+    is_output_outer_OOB = (g_block_x_offset + y_offset) >= cols;
+    if (is_output_outer_OOB) break;  // 行越界，不对该行做量化
+    local_max =
+        warpReduceMax(local_max);  // 无论该block有多少个有效线程，max均合法
+    if (threadIdx.x == 0)
+      smem_max[y_offset] = local_max;  // x0 顺序写，复用，无conflict
+  }
+
+  // 阶段3:
+  // Output放缩强转 + Scale写回‌
+  // ------------------------------------------------------------------
+  // 无论是否padding_last_dim_to_8x，输出连续维的stride均为合法值
+  const int padded_rows = (padding_last_dim_to_8x) ? ((rows + 7) & -8) : rows;
+  const int padded_cols = (padding_last_dim_to_8x) ? ((cols + 7) & -8) : cols;
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const float scale_on_fp32_to_outputT =
+          ComputeScale<__nv_bfloat16, OutT, using_pow2_scaling>(
+              smem_max[y_offset], 0.0f);
+      const float scale_on_fp8_to_inputT = __frcp_rn(scale_on_fp32_to_outputT);
+      float output_scaled_fp32 =
+          smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] *
+          scale_on_fp32_to_outputT;
+      const OutT output_scaled_fp8 = static_cast<OutT>(output_scaled_fp32);
+      const int g_output_y_offset = g_block_x_offset + y_offset;
+      const int g_output_x_offset = g_block_y_offset + x_offset;
+      // 如果padding连续维为8的倍数，则将连续维的stride使用位运算向上取整为最近的8倍数
+      const int g_output_inner_stride = padded_rows;
+      const int g_scale_inner_stride = (rows + 127) / 128;
+      if (g_output_y_offset < cols &&
+          g_output_x_offset < g_output_inner_stride) {
+        out[g_output_y_offset * g_output_inner_stride + g_output_x_offset] =
+            (g_output_x_offset < rows) ? output_scaled_fp8 : (OutT)0;
+        scales[g_output_y_offset * g_scale_inner_stride +
+               g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+      }
+    }
+  }
+}
+
+template <typename outT>
+void dispatch_fused_act_dequant_transpose_act_quant(
+    const paddle::Tensor &X,
+    const paddle::Tensor &Xscale,
+    paddle::Tensor &out,
+    paddle::Tensor &scale,
+    const int rows,
+    const int cols,
+    const bool &using_pow2_scaling,
+    const bool &padding_last_dim_to_8x,
+    const float TYPE_MAX) {
+  dim3 grid;
+  dim3 block(32, 32);
+  DISPATCH_BOOL(
+      using_pow2_scaling,
+      k_using_pow2_scaling,
+      DISPATCH_BOOL(padding_last_dim_to_8x, k_padding_last_dim_to_8x, {
+        grid.y = (rows + 127) / 128;
+        grid.x = (cols + 127) / 128;
+        LAUNCH_FUSED_ACT_DEQUANT_TRANSPOSE_ACT_QUANT(k_using_pow2_scaling,
+                                                     k_padding_last_dim_to_8x);
+      }))
+}
+std::vector<paddle::Tensor> fused_act_dequant_transpose_act_quant(
+    const paddle::Tensor &X,
+    const paddle::Tensor &Xscale,
+    const bool &padding_last_dim_to_8x,
+    const bool &using_pow2_scaling) {
+  PD_CHECK(X.dtype() == paddle::DataType::FLOAT8_E4M3FN);
+  PD_CHECK(Xscale.dtype() == paddle::DataType::FLOAT32);
+
+  int64_t data_rows = size_to_dim(X.shape().size() - 1, X.shape());
+  int64_t data_cols = X.shape().back();
+  int64_t rows = data_rows, cols = data_cols;
+
+  paddle::Tensor out;
+  paddle::Tensor scale;
+
+  if (padding_last_dim_to_8x) {
+    rows = ((data_rows + 7) / 8) *
+           8;  // 向上padding到8的倍数, 因为128为8的倍数，不影响scale shape
+  }
+  out = paddle::empty({cols, rows}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+  scale = paddle::empty(
+      {cols, (rows + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+
+  dispatch_fused_act_dequant_transpose_act_quant<phi::float8_e4m3fn>(
+      X,
+      Xscale,
+      out,
+      scale,
+      data_rows,
+      data_cols,
+      using_pow2_scaling,
+      padding_last_dim_to_8x,
+      F8LimitsTrait<__nv_fp8_e4m3>::max);
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_act_dequant_transpose_act_quant)
+    .Inputs({"X", "Xscale"})
+    .Outputs({"output", "scale"})
+    .Attrs({"padding_last_dim_to_8x: bool", "using_pow2_scaling: bool"})
+    .SetKernelFn(PD_KERNEL(fused_act_dequant_transpose_act_quant));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_act_quant.cu

@@ -0,0 +1,235 @@
+#include "quant_utils.h"
+
+constexpr int64_t TILE_SIZE = 128;  // 每个 block 处理 128x128 的元素块
+#define LAUNCH_FUSED_ACT_QUANT(                                                \
+    __transpose, __using_pow2_scaling, __padding_last_dim_to_8x)               \
+  do {                                                                         \
+    auto kernel = FusedActQuant<outT,                                          \
+                                __transpose,                                   \
+                                __using_pow2_scaling,                          \
+                                __padding_last_dim_to_8x>;                     \
+    int smem_size = 128 * 129 * sizeof(float);                                 \
+    PD_CHECK(cudaFuncSetAttribute(kernel,                                      \
+                                  cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                                  smem_size) == cudaSuccess);                  \
+    kernel<<<grid, block, smem_size, X.stream()>>>(X.data<phi::bfloat16>(),    \
+                                                   out.data<outT>(),           \
+                                                   scale.data<float>(),        \
+                                                   rows,                       \
+                                                   cols,                       \
+                                                   TYPE_MAX);                  \
+  } while (0)
+
+
+#define BLOCK_SIZE 128
+template <typename OutT,
+          bool transpose_output,
+          bool using_pow2_scaling,
+          bool padding_last_dim_to_8x>
+__global__ void FusedActQuant(const phi::bfloat16 *__restrict__ Xin,
+                              OutT *__restrict__ out,
+                              float *__restrict__ scales,
+                              const int rows,  // 此处的row未带padding
+                              const int cols,  // 同上
+                              const int TYPE_MAX) {
+  // 共享内存用于:1.直接读取+存储转置后的数据  2.用于1x128规约+量化
+  // 共享内存布局,前128x128为计算结果，后128为scale
+  extern __shared__ __align__(1024) uint8_t smem_buffer[];
+  float *smem_tile = reinterpret_cast<float *>(smem_buffer);
+  float *smem_max =
+      reinterpret_cast<float *>(smem_buffer) + BLOCK_SIZE * BLOCK_SIZE;
+
+  const int g_block_y_offset = blockIdx.y * BLOCK_SIZE;  // 块内行坐标偏移
+  const int g_block_x_offset = blockIdx.x * BLOCK_SIZE;  // 块内列坐标偏移
+  const __nv_bfloat16 *X = reinterpret_cast<const __nv_bfloat16 *>(Xin);
+
+  // 阶段1:
+  // Elementwise加载数据、计算、并将结果直接以特定布局存入共享内存‌(32x32)
+  // ------------------------------
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const int in_y_idx = g_block_y_offset + y_offset;
+      const int in_x_idx = g_block_x_offset + x_offset;
+      const int src_idx = in_y_idx * cols + in_x_idx;
+      if constexpr (transpose_output) {  // shared-mem conflict free写数
+        if (in_y_idx < rows && in_x_idx < cols) {
+          const int src_idx = in_y_idx * cols + in_x_idx;
+          smem_tile[swizzled_2d_idx(x_offset, BLOCK_SIZE, y_offset)] =
+              X[src_idx];
+        }
+      } else {
+        if (in_y_idx < rows && in_x_idx < cols) {
+          smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] =
+              X[src_idx];
+        }
+      }
+    }
+  }
+  __syncthreads();  // smem_tile中的swiglu数据(按需transpose)已ready
+
+  // 阶段2: ‌
+  // Shared 两级reduce, 给出每行的absmax_f
+  // ------------------------------------------------------
+  float local_max = 0.0f;
+  // 每个warp计算一行最大值，每个warp处理4行
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    // 行quantize max计算, 每行128个元素，每个线程处理4个元素
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      bool is_output_inner_OOB;
+      if constexpr (transpose_output) {
+        is_output_inner_OOB = (g_block_y_offset + x_offset) >= rows;
+      } else {
+        is_output_inner_OOB = (g_block_x_offset + x_offset) >= cols;
+      }
+      if (is_output_inner_OOB) break;  // 列越界则不取值、不影响max;
+      local_max = fmaxf(
+          local_max,
+          fabsf(smem_tile[swizzled_2d_idx(
+              y_offset,
+              BLOCK_SIZE,
+              x_offset)]));  // 正常情况下提供绝对值给其他线程,每个线程的local_max最多为4数最大值
+    }
+    bool is_output_outer_OOB;
+    if constexpr (transpose_output) {
+      is_output_outer_OOB = (g_block_x_offset + y_offset) >= cols;
+    } else {
+      is_output_outer_OOB = (g_block_y_offset + y_offset) >= rows;
+    }
+    if (is_output_outer_OOB) break;  // 行越界，不对该行做量化
+    local_max =
+        warpReduceMax(local_max);  // 无论该block有多少个有效线程，max均合法
+    if (threadIdx.x == 0)
+      smem_max[y_offset] = local_max;  // x0 顺序写，无conflict
+  }
+  __syncthreads();  // smem_max中的scale数据ready，128个元素对应128行的scale
+
+  // 阶段3:
+  // Output放缩强转 + Scale写回‌
+  // ------------------------------------------------------------------
+  // 无论是否padding_last_dim_to_8x，输出连续维的stride均为合法值
+  const int padded_rows = (padding_last_dim_to_8x) ? ((rows + 7) & -8) : rows;
+  const int padded_cols = (padding_last_dim_to_8x) ? ((cols + 7) & -8) : cols;
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const float scale_on_fp32_to_outputT =
+          ComputeScale<__nv_bfloat16, OutT, using_pow2_scaling>(
+              smem_max[y_offset], 0.0f);
+      const float scale_on_fp8_to_inputT = __frcp_rn(scale_on_fp32_to_outputT);
+      float output_scaled_fp32 =
+          smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] *
+          scale_on_fp32_to_outputT;
+      const OutT output_scaled_fp8 = static_cast<OutT>(output_scaled_fp32);
+      if constexpr (transpose_output) {
+        const int g_output_y_offset = g_block_x_offset + y_offset;
+        const int g_output_x_offset = g_block_y_offset + x_offset;
+        // 如果padding连续维为8的倍数，则将连续维的stride使用位运算向上取整为最近的8倍数
+        const int g_output_inner_stride = padded_rows;
+        const int g_scale_inner_stride = (rows + 127) / 128;
+        if (g_output_y_offset < cols &&
+            g_output_x_offset < g_output_inner_stride) {
+          out[g_output_y_offset * g_output_inner_stride + g_output_x_offset] =
+              (g_output_x_offset < rows) ? output_scaled_fp8 : (OutT)0;
+          scales[g_output_y_offset * g_scale_inner_stride +
+                 g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+        }
+      } else {
+        const int g_output_y_offset = g_block_y_offset + y_offset;
+        const int g_output_x_offset = g_block_x_offset + x_offset;
+        const int g_scale_inner_stride = (cols + 127) / 128;
+        const int g_output_inner_stride = padded_cols;
+        if (g_output_y_offset < rows &&
+            g_output_x_offset < g_output_inner_stride) {
+          out[g_output_y_offset * g_output_inner_stride + g_output_x_offset] =
+              (g_output_x_offset < cols) ? output_scaled_fp8 : (OutT)0;
+          scales[g_output_y_offset * g_scale_inner_stride +
+                 g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+        }
+      }
+    }
+  }
+}
+
+template <typename outT>
+void dispatch_fused_act_quant(const paddle::Tensor &X,
+                              paddle::Tensor &out,
+                              paddle::Tensor &scale,
+                              const int rows,
+                              const int cols,
+                              const bool &transpose_output,
+                              const bool &using_pow2_scaling,
+                              const bool &padding_last_dim_to_8x,
+                              const float TYPE_MAX) {
+  dim3 grid;
+  dim3 block(32, 32);
+  DISPATCH_BOOL(
+      transpose_output,
+      k_transpose_output,
+      DISPATCH_BOOL(
+          using_pow2_scaling,
+          k_using_pow2_scaling,
+          DISPATCH_BOOL(padding_last_dim_to_8x, k_padding_last_dim_to_8x, {
+            grid.y = (rows + 127) / 128;
+            grid.x = (cols + 127) / 128;
+            LAUNCH_FUSED_ACT_QUANT(k_transpose_output,
+                                   k_using_pow2_scaling,
+                                   k_padding_last_dim_to_8x);
+          })))
+}
+std::vector<paddle::Tensor> fused_act_quant(const paddle::Tensor &X,
+                                            const bool &transpose_output,
+                                            const bool &padding_last_dim_to_8x,
+                                            const bool &using_pow2_scaling) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);
+  int64_t data_rows = size_to_dim(X.shape().size() - 1, X.shape());
+  int64_t data_cols = X.shape().back();
+  int64_t rows = data_rows, cols = data_cols;
+
+  paddle::Tensor out;
+  paddle::Tensor scale;
+
+  if (transpose_output) {
+    if (padding_last_dim_to_8x) {
+      rows = ((data_rows + 7) / 8) *
+             8;  // 向上padding到8的倍数, 因为128为8的倍数，不影响scale shape
+    }
+    out =
+        paddle::empty({cols, rows}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+    scale = paddle::empty(
+        {cols, (rows + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+  } else {
+    if (padding_last_dim_to_8x) {
+      cols = ((data_cols + 7) / 8) *
+             8;  // 向上padding到8的倍数, 因为128为8的倍数，不影响scale shape
+    }
+    out =
+        paddle::empty({rows, cols}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+    scale = paddle::empty(
+        {rows, (cols + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+  }
+
+  dispatch_fused_act_quant<phi::float8_e4m3fn>(
+      X,
+      out,
+      scale,
+      data_rows,
+      data_cols,
+      transpose_output,
+      using_pow2_scaling,
+      padding_last_dim_to_8x,
+      F8LimitsTrait<__nv_fp8_e4m3>::max);
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_act_quant)
+    .Inputs({"X"})
+    .Outputs({"output", "scale"})
+    .Attrs({"transpose_output: bool",
+            "padding_last_dim_to_8x: bool",
+            "using_pow2_scaling: bool"})
+    .SetKernelFn(PD_KERNEL(fused_act_quant));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_spaq.cu

@@ -0,0 +1,369 @@
+#include "quant_utils.h"
+
+#define LAUNCH_FUSED_SPAQ(__using_pow2_scaling, __with_prob)          \
+  do {                                                                \
+    auto kernel = FusedSPAQKernel<__using_pow2_scaling, __with_prob>; \
+    kernel<<<grid, block, 0, X.stream()>>>(                           \
+        X.data<phi::bfloat16>(),                                      \
+        prob ? prob->data<float>() : nullptr,                         \
+        out.data<phi::float8_e4m3fn>(),                               \
+        scale.data<float>(),                                          \
+        rows,                                                         \
+        cols);                                                        \
+  } while (0)
+
+#define LAUNCH_FUSED_SPAQ_VEC4(__using_pow2_scaling, __with_prob) \
+  do {                                                            \
+    auto kernel = FusedSPAQKernelVec4<__using_pow2_scaling,       \
+                                      __with_prob,                \
+                                      thread_per_block>;          \
+    kernel<<<grid, block, 0, X.stream()>>>(                       \
+        X.data<phi::bfloat16>(),                                  \
+        prob ? prob->data<float>() : nullptr,                     \
+        out.data<phi::float8_e4m3fn>(),                           \
+        scale.data<float>(),                                      \
+        rows,                                                     \
+        cols,                                                     \
+        scale_cols);                                              \
+  } while (0)
+
+
+typedef struct __align__(8) {
+  __nv_bfloat16 x;
+  __nv_bfloat16 y;
+  __nv_bfloat16 z;
+  __nv_bfloat16 w;
+}
+bfloat16x4_t;
+
+typedef struct __align__(4) {
+  __nv_fp8_e4m3 x;
+  __nv_fp8_e4m3 y;
+  __nv_fp8_e4m3 z;
+  __nv_fp8_e4m3 w;
+}
+fp8_e4m3x4_t;
+
+__device__ __forceinline__ float fast_swiglu(const __nv_bfloat16 x,
+                                             const __nv_bfloat16 y) {
+  const float x_f = __bfloat162float(x);
+  const float y_f = __bfloat162float(y);
+  const float silu = x_f * __frcp_rn(1.0f + __expf(-x_f));
+  const float result = silu * y_f;
+  return result;
+}
+__device__ __forceinline__ float4 fast_swiglu_vec4(const bfloat16x4_t &lhs,
+                                                   const bfloat16x4_t &rhs) {
+  const float x_f_x = __bfloat162float(lhs.x);
+  const float x_f_y = __bfloat162float(lhs.y);
+  const float x_f_z = __bfloat162float(lhs.z);
+  const float x_f_w = __bfloat162float(lhs.w);
+
+  const float y_f_x = __bfloat162float(rhs.x);
+  const float y_f_y = __bfloat162float(rhs.y);
+  const float y_f_z = __bfloat162float(rhs.z);
+  const float y_f_w = __bfloat162float(rhs.w);
+
+  const float silu_x = x_f_x * __frcp_rn(1.0f + __expf(-x_f_x));
+  const float silu_y = x_f_y * __frcp_rn(1.0f + __expf(-x_f_y));
+  const float silu_z = x_f_z * __frcp_rn(1.0f + __expf(-x_f_z));
+  const float silu_w = x_f_w * __frcp_rn(1.0f + __expf(-x_f_w));
+
+  return {silu_x * y_f_x, silu_y * y_f_y, silu_z * y_f_z, silu_w * y_f_w};
+}
+__device__ __forceinline__ float amax_float4(const float4 &vec) {
+  return fmaxf(fmaxf(fabsf(vec.x), fabsf(vec.y)),
+               fmaxf(fabsf(vec.z), fabsf(vec.w)));
+}
+
+__device__ __forceinline__ fp8_e4m3x4_t
+scale_fp32x4_to_fp8x4(const float4 &vec, const float scale) {
+  return {static_cast<__nv_fp8_e4m3>(vec.x * scale),
+          static_cast<__nv_fp8_e4m3>(vec.y * scale),
+          static_cast<__nv_fp8_e4m3>(vec.z * scale),
+          static_cast<__nv_fp8_e4m3>(vec.w * scale)};
+}
+
+
+template <bool using_pow2_scaling, bool with_prob, int thread_per_block>
+__global__ void FusedSPAQKernelVec4(const phi::bfloat16 *__restrict__ Xin,
+                                    const float *__restrict__ prob,
+                                    phi::float8_e4m3fn *__restrict__ out,
+                                    float *__restrict__ scales,
+                                    const int64_t rows,
+                                    const int64_t cols,
+                                    const int64_t scale_cols) {
+  constexpr int elements_per_thread = 4;
+  constexpr int warp_size = 32;
+  constexpr int warp_num = thread_per_block / warp_size;
+  const int64_t scale_stride = scale_cols;
+  const int lane = threadIdx.x % warp_size;
+  const int64_t x_offset =
+      static_cast<int64_t>(threadIdx.x) * elements_per_thread;
+  const unsigned int mask = 0xffffffff;  // whole warp mask
+
+  // 使用grid stride循环处理所有行
+  for (int64_t base_y = blockIdx.y; base_y < rows; base_y += gridDim.y) {
+    const int64_t in_y_idx = base_y;
+    const int64_t in_x_idx = static_cast<int64_t>(blockIdx.x) *
+                                 static_cast<int64_t>(blockDim.x) *
+                                 elements_per_thread +
+                             x_offset;
+    const int64_t src_idx = in_y_idx * cols + in_x_idx;
+
+    float p_t0;
+
+    // 边界检查
+    if (in_x_idx >= cols / 2) [[unlikely]]
+      continue;
+
+    if constexpr (with_prob) {
+      // Prefetch prob
+      if (lane == 0) p_t0 = prob[in_y_idx];
+    }
+
+    const __nv_bfloat16 *X = reinterpret_cast<const __nv_bfloat16 *>(Xin);
+
+    // Initialize activation storage
+    float4 act_f32x4;
+    bfloat16x4_t lhs_bf16x4, rhs_bf16x4;
+
+    // Reinterpret input pointer as bfloat16x4_t* for vectorized loading
+    const bfloat16x4_t *X_lhs_vec =
+        reinterpret_cast<const bfloat16x4_t *>(X + src_idx);
+    const bfloat16x4_t *X_rhs_vec =
+        reinterpret_cast<const bfloat16x4_t *>(X + src_idx + cols / 2);
+
+    lhs_bf16x4 = *X_lhs_vec;
+    rhs_bf16x4 = *X_rhs_vec;
+
+    act_f32x4 = fast_swiglu_vec4(lhs_bf16x4, rhs_bf16x4);
+
+    if constexpr (with_prob) {
+      // Warp level sync to avoid syncthreads
+      const float p = __shfl_sync(mask, p_t0, 0);
+      act_f32x4.x *= p;
+      act_f32x4.y *= p;
+      act_f32x4.z *= p;
+      act_f32x4.w *= p;
+    }
+
+    // Phase 2: Block Reduction to find per-quant block absolute maxima
+    // Compute absolute values
+    float thread_amax = amax_float4(act_f32x4);
+
+// All-Reduce within the warp
+#pragma unroll
+    for (int offset = 16; offset > 0; offset /= 2) {
+      const float val = __shfl_down_sync(mask, thread_amax, offset);
+      thread_amax = fmaxf(thread_amax, val);
+    }
+    const float final_amax = __shfl_sync(mask, thread_amax, 0);
+
+    // Phase 3: Compute scales and quantize the outputs
+    const float scale = ComputeScale<float, __nv_fp8_e4m3, using_pow2_scaling>(
+        final_amax, 0.0f);
+    const float inv_scale = __frcp_rn(scale);
+
+    const fp8_e4m3x4_t act_fp8x4 = scale_fp32x4_to_fp8x4(act_f32x4, scale);
+    fp8_e4m3x4_t *const out_vec_addr =
+        reinterpret_cast<fp8_e4m3x4_t *>(out + in_y_idx * cols / 2 + in_x_idx);
+    *out_vec_addr = act_fp8x4;
+
+    if (lane == 0) {
+      const int64_t scale_idx = in_y_idx * scale_stride + in_x_idx / 128;
+      scales[scale_idx] = inv_scale;
+    }
+  }
+}
+
+template <bool using_pow2_scaling, bool with_prob>
+__global__ void FusedSPAQKernel(const phi::bfloat16 *__restrict__ Xin,
+                                const float *__restrict__ prob,
+                                phi::float8_e4m3fn *__restrict__ out,
+                                float *__restrict__ scales,
+                                const int rows,
+                                const int cols) {
+  // Configure shared memory
+  __shared__ float smem_tile[256];  // Shared memory for activation values
+  __shared__ float warp_max[2][4];  // Shared memory for warp maxima (2 quant
+                                    // blocks x 4 warps)
+  __shared__ __nv_bfloat16
+      quant_block_amax[2];  // Shared memory for quant block maxima
+
+  const __nv_bfloat16 *X = reinterpret_cast<const __nv_bfloat16 *>(Xin);
+  const int x_offset = threadIdx.x;
+  const int quant_block_idx =
+      threadIdx.x / 128;  // 0 or 1, two quant blocks per block
+  const int in_y_idx = blockIdx.y;
+  const int in_x_idx = blockIdx.x * blockDim.x + x_offset;
+  const int src_idx = in_y_idx * cols + in_x_idx;
+
+  // Load data and compute swiGLU activation
+  if (in_x_idx < cols / 2) [[likely]] {
+    __nv_bfloat16 x1 = X[src_idx];             // First half of the input
+    __nv_bfloat16 x2 = X[src_idx + cols / 2];  // Second half of the input
+
+    if constexpr (with_prob) {
+      float row_prob = prob[in_y_idx];
+      smem_tile[x_offset] = fast_swiglu(x1, x2) * row_prob;
+    } else {
+      smem_tile[x_offset] = fast_swiglu(x1, x2);
+    }
+  }
+
+  __syncthreads();  // Ensure all threads have loaded their data
+
+  // Phase 2: Block Reduction to find per-quant block absolute maximums
+  float local_max = (in_x_idx < (cols / 2)) ? fabsf(smem_tile[x_offset]) : 0.0f;
+
+
+  // Warp-level reduction
+  unsigned int mask = 0xffffffff;
+  int lane = threadIdx.x % 32;
+  int warp_id =
+      (threadIdx.x % 128) / 32;  // Warp ID within the quant block (0-3)
+
+  // Reduce within the warp
+  for (int offset = 16; offset > 0; offset /= 2) {
+    float val = __shfl_down_sync(mask, local_max, offset);
+    local_max = fmaxf(local_max, val);
+  }
+
+  // Store warp maxima
+  if (lane == 0) {
+    warp_max[quant_block_idx][warp_id] = local_max;
+  }
+
+  __syncthreads();
+
+  // Reduce warp maxima to get quant block maxima
+  if (warp_id == 0 && lane < 4) {
+    if (threadIdx.x < 256) {  // Ensure only valid threads participate
+      float block_max = warp_max[quant_block_idx][lane];
+      // Reduce over the 4 warp maxima
+      if (lane == 0) {
+        block_max = fmaxf(block_max, warp_max[quant_block_idx][1]);
+        block_max = fmaxf(block_max, warp_max[quant_block_idx][2]);
+        block_max = fmaxf(block_max, warp_max[quant_block_idx][3]);
+        quant_block_amax[quant_block_idx] = __float2bfloat16(block_max);
+      }
+    }
+  }
+
+  __syncthreads();
+
+  // Phase 3: Compute scales and quantize the outputs
+  const float block_max_float = (float)quant_block_amax[quant_block_idx];
+  const int scale_stride = (cols / 2 + 127) / 128;
+
+  float scale = ComputeScale<float, __nv_fp8_e4m3, using_pow2_scaling>(
+      block_max_float, 0.0f);
+  float inv_scale = __frcp_rn(scale);
+
+  // Quantize
+  float output_scaled_fp32 = smem_tile[x_offset] * scale;
+
+
+  const int g_output_y_offset = in_y_idx;
+  const int g_output_x_offset = in_x_idx;
+
+  // Write output and scales
+  if (g_output_y_offset < rows && g_output_x_offset < cols / 2) {
+    out[g_output_y_offset * (cols / 2) + g_output_x_offset] =
+        static_cast<phi::float8_e4m3fn>(output_scaled_fp32);
+    if (x_offset % 128 == 0) {
+      // Only one thread per quant block writes the scale
+      scales[g_output_y_offset * scale_stride + in_x_idx / 128] = inv_scale;
+    }
+  }
+}
+
+
+void dispatch_fused_spaq(const paddle::Tensor &X,
+                         const paddle::optional<paddle::Tensor> &prob,
+                         paddle::Tensor &out,
+                         paddle::Tensor &scale,
+                         const int rows,
+                         const int cols,
+                         const bool &using_pow2_scaling,
+                         const bool &with_prob) {
+  constexpr int thread_per_block = 256;
+  dim3 grid;
+  dim3 block;
+  if (cols % 8 == 0) {
+    // Use mixed vectorizing strategy, while cols size be 8x (4x2)
+    // Each thread process 4 bfloat16 element in same row, each warp handles
+    // 1x128 vector Each block handles several sub-row (numel = 4 x blockDim.x)
+    // of input vector
+    block.x = thread_per_block;
+    constexpr int vec_numel = 4;
+    const int scale_cols = scale.shape().back();
+    DISPATCH_BOOL(
+        using_pow2_scaling,
+        k_using_pow2_scaling,
+        DISPATCH_BOOL(
+            with_prob, k_with_prob, grid.y = rows > 65535 ? 65535 : rows;
+            grid.x =
+                ((cols / 2) + block.x * vec_numel - 1) / (block.x * vec_numel);
+            LAUNCH_FUSED_SPAQ_VEC4(k_using_pow2_scaling, k_with_prob);))
+
+  } else {
+    // Plain elementwise strategy:
+    // Each block processing a sub-row (numel = blockDim.x) of the input tensor.
+    block.x = thread_per_block;
+    DISPATCH_BOOL(
+        using_pow2_scaling,
+        k_using_pow2_scaling,
+        DISPATCH_BOOL(
+            with_prob, k_with_prob, grid.y = rows > 65535 ? 65535 : rows;
+            grid.x = ((cols / 2) + block.x - 1) / block.x;
+            LAUNCH_FUSED_SPAQ(k_using_pow2_scaling, k_with_prob);))
+  }
+}
+
+
+std::vector<paddle::Tensor> fused_spaq(
+    const paddle::Tensor &X,
+    const paddle::optional<paddle::Tensor> &prob,
+    const bool &using_pow2_scaling) {
+  // ---------------- Arguments check --------------------
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);
+  if (prob) PD_CHECK(prob.get().dtype() == paddle::DataType::FLOAT32);
+  int64_t rows = size_to_dim(X.shape().size() - 1, X.shape());
+  int64_t cols = X.shape().back();
+  PADDLE_ENFORCE_EQ(cols % 2,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The last dim of Input(X) should be exactly divided "
+                        "by 2 , but got %d",
+                        cols));
+  if (prob) {
+    PADDLE_ENFORCE_EQ(prob.get().shape()[0],
+                      rows,
+                      common::errors::InvalidArgument(
+                          "The first dim of Input(X) should be equal to the "
+                          "first dim of Input(prob) but got X.shape[0]: %d, "
+                          "prob.shape[0]: %d",
+                          rows,
+                          prob.get().shape()[0]));
+  }
+
+  paddle::Tensor out;
+  paddle::Tensor scale;
+
+  out = paddle::empty(
+      {rows, cols / 2}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+  scale = paddle::empty(
+      {rows, ((cols / 2) + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+
+  dispatch_fused_spaq(
+      X, prob, out, scale, rows, cols, using_pow2_scaling, !!prob);
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_spaq)
+    .Inputs({"X", paddle::Optional("prob")})
+    .Outputs({"output", "scale"})
+    .Attrs({"using_pow2_scaling: bool"})
+    .SetKernelFn(PD_KERNEL(fused_spaq));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_stack_transpose_quant.cu

@@ -0,0 +1,320 @@
+#include "quant_utils.h"
+
+template <typename T, int VecSize>
+struct __align__(sizeof(T) * VecSize) VecType {
+  T val[VecSize];
+  __host__ __device__ inline T& operator[](size_t i) { return val[i]; }
+  __host__ __device__ inline const T& operator[](size_t i) const {
+    return val[i];
+  }
+};
+
+struct FastDiv {
+  FastDiv() {}
+  FastDiv(uint64_t d) {
+    for (shift_val = 0; shift_val < 64; ++shift_val) {
+      uint64_t shift_limit = uint64_t(1) << shift_val;
+      if (shift_limit >= d) break;
+    }
+
+    // quotient = ((uint128_t)n_hi << 64) / d
+    uint64_t quotient = 0;
+    uint64_t n_hi = (uint64_t(1) << shift_val) - d, n_lo = 0;
+    for (int i = 63; i >= 0; --i) {
+      uint64_t d_hi = i == 0 ? 0 : d >> (64 - i);
+      uint64_t d_lo = d << i;
+      if (n_hi == 0 && n_lo == 0) break;
+      if ((d_hi < n_hi) || (d_hi <= n_hi && d_lo <= n_lo)) {
+        quotient |= uint64_t(1) << i;
+        n_hi -= d_hi + (d_lo > n_lo);
+        n_lo -= d_lo;
+      }
+    }
+    multiplier = quotient + 1;
+  }
+
+  __device__ uint64_t Div(uint64_t n) const {
+    uint64_t t = __umul64hi(n, multiplier);
+    return (t + n) >> shift_val;
+  }
+
+  int shift_val;
+  uint64_t multiplier;
+};
+
+__device__ void BlockLoad(const int64_t* __restrict__ X_ptrs,
+                          __nv_bfloat16 input[4][4],
+                          size_t K,
+                          size_t block_y,
+                          size_t block_x) {
+  const __nv_bfloat16* X =
+      reinterpret_cast<const __nv_bfloat16*>(X_ptrs[blockIdx.z]);
+
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_m = block_y * 128 + threadIdx.y + i * 32;
+    size_t idx_k = block_x * 128 + threadIdx.x * 4;
+    size_t idx = idx_m * K + idx_k;
+
+    using LoadT = VecType<__nv_bfloat16, 4>;
+    LoadT data = *reinterpret_cast<const LoadT*>(X + idx);
+    for (int j = 0; j < 4; j++) {
+      input[i][j] = data[j];
+    }
+  }
+}
+
+__device__ __nv_bfloat16 WarpReduceMax(__nv_bfloat16 x) {
+  for (int offset = 16; offset > 0; offset /= 2) {
+    __nv_bfloat16 t = __shfl_down_sync(0xffffffff, x, offset);
+    x = __hmax(x, t);
+  }
+  return x;
+}
+
+__device__ __nv_bfloat16 BlockReduceMax(__nv_bfloat16 input[4][4]) {
+  // [(4), 32, 32, (4)] => [32, 32]
+  __nv_bfloat16 local_max;
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      __nv_bfloat16 t = __habs(input[i][j]);
+      local_max = (i == 0 && j == 0) ? t : __hmax(local_max, t);
+    }
+  }
+
+  // [32, (32)] => [32]
+  __nv_bfloat16 warp_max = WarpReduceMax(local_max);
+
+  // [(32)] => [1]
+  __shared__ __nv_bfloat16 block_max[32];
+  if (threadIdx.x == 0) {
+    block_max[threadIdx.y] = warp_max;
+  }
+  __syncthreads();
+  if (threadIdx.y == 0) {
+    warp_max = WarpReduceMax(block_max[threadIdx.x]);
+    if (threadIdx.x == 0) {
+      block_max[0] = warp_max;
+    }
+  }
+  __syncthreads();
+
+  return block_max[0];
+}
+
+template <typename OutT>
+__global__ void __launch_bounds__(1024)
+    FusedStackQuantKernel(const int64_t* __restrict__ X_ptrs,
+                          OutT* __restrict__ out,
+                          float* __restrict__ scale,
+                          size_t M,
+                          size_t K,
+                          FastDiv K_div_128) {
+  size_t block_y = K_div_128.Div(blockIdx.x);
+  size_t block_x = blockIdx.x - block_y * (K / 128);
+
+  // Load 128x128 elements from X
+  __nv_bfloat16 input[4][4];
+  BlockLoad(X_ptrs, input, K, block_y, block_x);
+
+  // Find the maximum in all elements
+  __nv_bfloat16 amax = BlockReduceMax(input);
+
+  // Compute scale and store back
+  float scale_inv = ComputeScale<__nv_bfloat16, OutT>(amax, 0.0f);
+  float scale_out = __frcp_rn(scale_inv);
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_m = block_y;
+    size_t idx_k = block_x;
+    size_t idx = (idx_n * (M / 128) + idx_m) * (K / 128) + idx_k;
+    scale[idx] = scale_out;
+  }
+
+  // Scale X and store to out
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_m = block_y * 128 + threadIdx.y + i * 32;
+    size_t idx_k = block_x * 128 + threadIdx.x * 4;
+    size_t idx = (idx_n * M + idx_m) * K + idx_k;
+
+    using StoreT = VecType<OutT, 4>;
+    StoreT data;
+    for (int j = 0; j < 4; j++) {
+      float input_fp32 = static_cast<float>(input[i][j]);
+      float output_scaled = input_fp32 * scale_inv;
+      data[j] = static_cast<OutT>(output_scaled);
+    }
+    *reinterpret_cast<StoreT*>(out + idx) = data;
+  }
+}
+
+template <typename OutT>
+__global__ void __launch_bounds__(1024)
+    FusedStackTransposeQuantKernel(const int64_t* __restrict__ X_ptrs,
+                                   OutT* __restrict__ out,
+                                   float* __restrict__ scale,
+                                   size_t M,
+                                   size_t K,
+                                   FastDiv K_div_128) {
+  size_t block_y = K_div_128.Div(blockIdx.x);
+  size_t block_x = blockIdx.x - block_y * (K / 128);
+
+  // Load 128x128 elements from X
+  __nv_bfloat16 input[4][4];
+  BlockLoad(X_ptrs, input, K, block_y, block_x);
+
+  // Find the maximum in all elements
+  __nv_bfloat16 amax = BlockReduceMax(input);
+
+  // Compute scale and store back
+  float scale_inv = ComputeScale<__nv_bfloat16, OutT>(amax, 0.0f);
+  float scale_out = __frcp_rn(scale_inv);
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_k = block_x;
+    size_t idx_m = block_y;
+    size_t idx = (idx_n * (K / 128) + idx_k) * (M / 128) + idx_m;
+    scale[idx] = scale_out;
+  }
+
+  // Scale X and transpose in shared memory
+  __shared__ OutT shm[128][129];
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      float input_fp32 = static_cast<float>(input[i][j]);
+      float output_scaled = input_fp32 * scale_inv;
+      shm[threadIdx.x * 4 + j][i * 32 + threadIdx.y] =
+          static_cast<OutT>(output_scaled);
+    }
+  }
+  __syncthreads();
+
+  // Store X back to out
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_k = block_x * 128 + threadIdx.y + i * 32;
+    size_t idx_m = block_y * 128 + threadIdx.x * 4;
+    size_t idx = (idx_n * K + idx_k) * M + idx_m;
+
+    using StoreT = VecType<OutT, 4>;
+    StoreT data;
+    for (int j = 0; j < 4; j++) {
+      data[j] = shm[i * 32 + threadIdx.y][threadIdx.x * 4 + j];
+    }
+    *reinterpret_cast<StoreT*>(out + idx) = data;
+  }
+}
+
+/**
+ * Stack tensors in X, optionally transpose dim[-1] and dim[-2], and do
+ * quantization on both dim[-1] and dim[-2].
+ *
+ * Inputs:
+ *   X    : N tensors of [M, K], bfloat16
+ *
+ * Outputs:
+ *   if Transpose:
+ *     out  : [N * K, M], float8_e4m3fn
+ *     scale: [N * K / 128, M / 128], float
+ *   else:
+ *     out  : [N * M, K], float8_e4m3fn
+ *     scale: [N * M / 128, K / 128], float
+ *
+ * Requirements:
+ *   1) N <= 65535
+ *   2) M % 128 == 0
+ *   3) K % 128 == 0
+ */
+template <bool Transpose>
+std::vector<paddle::Tensor> fused_stack_transpose_quant(
+    const std::vector<paddle::Tensor>& X) {
+  int64_t N = X.size();
+  PD_CHECK(N > 0);
+  for (int64_t i = 0; i < N; i++) {
+    PD_CHECK(X[i].dtype() == paddle::DataType::BFLOAT16);
+  }
+
+  std::vector<int64_t> shape = X[0].shape();
+  PD_CHECK(shape.size() == 2);
+  int64_t M = shape[0];
+  int64_t K = shape[1];
+
+  for (int64_t i = 1; i < N; i++) {
+    std::vector<int64_t> shape = X[i].shape();
+    PD_CHECK(shape.size() == 2);
+    PD_CHECK(shape[0] == M);
+    PD_CHECK(shape[1] == K);
+  }
+
+  PADDLE_ENFORCE_LE(N,
+                    65535,
+                    common::errors::InvalidArgument(
+                        "The batch size (N) must be no larger than 65535."));
+  PADDLE_ENFORCE_EQ(M % 128,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The upper dim (M) must be multiple of 128."));
+  PADDLE_ENFORCE_EQ(K % 128,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The lower dim (K) must be multiple of 128."));
+
+  // Allocate for out and scale
+  std::vector<int64_t> out_shape, scale_shape;
+  if (Transpose) {
+    out_shape = {N * K, M};
+    scale_shape = {N * K / 128, M / 128};
+  } else {
+    out_shape = {N * M, K};
+    scale_shape = {N * M / 128, K / 128};
+  }
+
+  const auto& place = X[0].place();
+  paddle::Tensor out =
+      paddle::empty(out_shape, paddle::DataType::FLOAT8_E4M3FN, place);
+  paddle::Tensor scale =
+      paddle::empty(scale_shape, paddle::DataType::FLOAT32, place);
+
+  // Skip 0-size
+  if (M == 0 || K == 0) {
+    return {out, scale};
+  }
+
+  // Copy the pointers in X to device
+  paddle::Tensor X_ptrs_cpu = paddle::empty({N}, paddle::DataType::INT64);
+  int64_t* X_ptrs_data = X_ptrs_cpu.data<int64_t>();
+  for (int64_t i = 0; i < N; i++) {
+    X_ptrs_data[i] = reinterpret_cast<int64_t>(X[i].data());
+  }
+  paddle::Tensor X_ptrs_gpu = X_ptrs_cpu.copy_to(place, /* blocking= */ false);
+
+  // Launch kernel
+  dim3 grid((M / 128) * (K / 128), 1, N);
+  dim3 block(32, 32);
+
+#define LAUNCH_KERNEL(KERNEL)                                               \
+  KERNEL<<<grid, block, 0, X[0].stream()>>>(X_ptrs_gpu.data<int64_t>(),     \
+                                            out.data<phi::float8_e4m3fn>(), \
+                                            scale.data<float>(),            \
+                                            M,                              \
+                                            K,                              \
+                                            FastDiv(K / 128))
+  if (Transpose) {
+    LAUNCH_KERNEL(FusedStackTransposeQuantKernel);
+  } else {
+    LAUNCH_KERNEL(FusedStackQuantKernel);
+  }
+#undef LAUNCH_KERNEL
+
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_stack_quant)
+    .Inputs({paddle::Vec("X")})
+    .Outputs({"output", "scale"})
+    .SetKernelFn(PD_KERNEL(fused_stack_transpose_quant<false>));
+
+PD_BUILD_OP(fused_stack_transpose_quant)
+    .Inputs({paddle::Vec("X")})
+    .Outputs({"output", "scale"})
+    .SetKernelFn(PD_KERNEL(fused_stack_transpose_quant<true>));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_swiglu_act_quant.cu

@@ -0,0 +1,400 @@
+#include "quant_utils.h"
+
+#define LAUNCH_FUSED_SWIGLU_ACT_QUANT(                                         \
+    __is_combine, __transpose, __using_pow2_scaling, __padding_last_dim_to_8x) \
+  do {                                                                         \
+    auto kernel = FusedSwigluActQuantKernel<outT,                              \
+                                            __is_combine,                      \
+                                            __transpose,                       \
+                                            __using_pow2_scaling,              \
+                                            __padding_last_dim_to_8x>;         \
+    int smem_size = 128 * 129 * sizeof(float);                                 \
+    PD_CHECK(cudaFuncSetAttribute(kernel,                                      \
+                                  cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                                  smem_size) == cudaSuccess);                  \
+    kernel<<<grid, block, smem_size, X.stream()>>>(                            \
+        X.data<phi::bfloat16>(),                                               \
+        Y ? Y->data<phi::bfloat16>() : nullptr,                                \
+        out.data<outT>(),                                                      \
+        scale.data<float>(),                                                   \
+        rows,                                                                  \
+        cols,                                                                  \
+        TYPE_MAX);                                                             \
+  } while (0)
+
+
+__device__ __forceinline__ float fast_swiglu(const __nv_bfloat16 x,
+                                             const __nv_bfloat16 y) {
+  const float x_f = __bfloat162float(x);
+  const float y_f = __bfloat162float(y);
+  const float silu = x_f * __frcp_rn(1.0f + __expf(-x_f));
+  const float result = silu * y_f;
+  return result;
+}
+
+constexpr int BLOCK_SIZE = 128;
+template <typename OutT,
+          bool is_combine,
+          bool transpose_output,
+          bool using_pow2_scaling,
+          bool padding_last_dim_to_8x>
+__global__ void FusedSwigluActQuantKernel(const phi::bfloat16 *__restrict__ Xin,
+                                          const phi::bfloat16 *__restrict__ Yin,
+                                          OutT *__restrict__ out,
+                                          float *__restrict__ scales,
+                                          const int rows,
+                                          const int cols,
+                                          const int TYPE_MAX) {
+  // 如果is_combine,
+  // 则Xin中存储了X和Y的数据，此时gridDim.x应为cols/BLOCK_SIZE的一半
+  // 共享内存用于:1.计算后直接存储转置后的结果数据  2.用于1x128规约+量化
+  // 共享内存布局,前128x128为计算结果，后128为scale
+  extern __shared__ __align__(1024) uint8_t smem_buffer[];
+  // extern __shared__  uint8_t smem_buffer[];
+  float *smem_tile = reinterpret_cast<float *>(smem_buffer);
+  float *smem_max =
+      reinterpret_cast<float *>(smem_buffer) + BLOCK_SIZE * BLOCK_SIZE;
+
+  const int g_block_y_offset = blockIdx.y * BLOCK_SIZE;  // 块内行坐标偏移
+  const int g_block_x_offset = blockIdx.x * BLOCK_SIZE;  // 块内列坐标偏移
+  const __nv_bfloat16 *X = reinterpret_cast<const __nv_bfloat16 *>(Xin);
+  const __nv_bfloat16 *Y = reinterpret_cast<const __nv_bfloat16 *>(Yin);
+
+  // 阶段1:
+  // Elementwise加载数据、计算、并将结果直接以特定布局存入共享内存‌(32x32)
+  // ------------------------------
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const int in_y_idx = g_block_y_offset + y_offset;
+      const int in_x_idx = g_block_x_offset + x_offset;
+      const int src_idx = in_y_idx * cols + in_x_idx;
+      if constexpr (transpose_output) {  // shared-mem conflict free写数
+        if constexpr (is_combine) {
+          if (in_y_idx < rows && in_x_idx < cols / 2) {
+            smem_tile[swizzled_2d_idx(x_offset, BLOCK_SIZE, y_offset)] =
+                fast_swiglu(X[src_idx], X[src_idx + cols / 2]);
+          }
+        } else {
+          if (in_y_idx < rows && in_x_idx < cols) {
+            const int src_idx = in_y_idx * cols + in_x_idx;
+            smem_tile[swizzled_2d_idx(x_offset, BLOCK_SIZE, y_offset)] =
+                fast_swiglu(X[src_idx], Y[src_idx]);
+          }
+        }
+      } else {
+        if constexpr (is_combine) {
+          if (in_y_idx < rows && in_x_idx < cols / 2) {
+            smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] =
+                fast_swiglu(X[src_idx], X[src_idx + cols / 2]);
+          }
+        } else {
+          if (in_y_idx < rows && in_x_idx < cols) {
+            smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] =
+                fast_swiglu(X[src_idx], Y[src_idx]);
+          }
+        }
+      }
+    }
+  }
+  __syncthreads();  // smem_tile中的swiglu数据(按需transpose)已ready
+
+  // 阶段2: ‌
+  // Shared 两级reduce, 给出每行的absmax_f
+  // ------------------------------------------------------
+  float local_max = 0.0f;
+  // 每个warp计算一行最大值，每个warp处理4行
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    // 行quantize max计算, 每行128个元素，每个线程处理4个元素
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      bool is_output_inner_OOB;
+      if constexpr (transpose_output) {
+        is_output_inner_OOB = (g_block_y_offset + x_offset) >= rows;
+      } else {
+        is_output_inner_OOB = (g_block_x_offset + x_offset) >= cols;
+      }
+      if (is_output_inner_OOB) break;  // 列越界则不取值、不影响max;
+      local_max = fmaxf(
+          local_max,
+          fabsf(smem_tile[swizzled_2d_idx(
+              y_offset,
+              BLOCK_SIZE,
+              x_offset)]));  // 正常情况下提供绝对值给其他线程,每个线程的local_max最多为4数最大值
+    }
+    bool is_output_outer_OOB;
+    if constexpr (transpose_output) {
+      is_output_outer_OOB = (g_block_x_offset + y_offset) >= cols;
+    } else {
+      is_output_outer_OOB = (g_block_y_offset + y_offset) >= rows;
+    }
+    if (is_output_outer_OOB) break;  // 行越界，不对该行做量化
+    local_max =
+        warpReduceMax(local_max);  // 无论该block有多少个有效线程，max均合法
+    if (threadIdx.x == 0)
+      smem_max[y_offset] = local_max;  // x0 顺序写，无conflict
+  }
+  __syncthreads();  // smem_max中的scale数据ready，128个元素对应128行的scale
+
+  // 阶段3:
+  // Output放缩强转 + Scale写回‌
+  // ------------------------------------------------------------------
+  /* 输出缓冲区伪代码：
+  if(transpose_output){
+      if(is_combine){
+          out = {cols / 2, rows}, scale = {cols / 2, (rows + 127) / 128}
+      }else{
+          out = {cols, rows}, scale = {cols, (rows + 127) / 128}
+      }
+  }else{
+      if(is_combine){
+          out {rows, cols / 2}, scale = {rows, ((cols / 2) + 127) / 128}
+      }else{
+          out = {rows, cols}, scale = {rows, (cols + 127) / 128}
+      }
+  }
+  */
+  for (int y_offset = threadIdx.y; y_offset < BLOCK_SIZE;
+       y_offset += blockDim.y) {
+    for (int x_offset = threadIdx.x; x_offset < BLOCK_SIZE;
+         x_offset += blockDim.x) {
+      const float scale_on_fp32_to_outputT =
+          ComputeScale<__nv_bfloat16, OutT, using_pow2_scaling>(
+              smem_max[y_offset], 0.0f);
+      const float scale_on_fp8_to_inputT = __frcp_rn(scale_on_fp32_to_outputT);
+      float output_scaled_fp32 =
+          smem_tile[swizzled_2d_idx(y_offset, BLOCK_SIZE, x_offset)] *
+          scale_on_fp32_to_outputT;
+      const OutT output_scaled_fp8 = static_cast<OutT>(output_scaled_fp32);
+      if constexpr (transpose_output) {
+        const int g_output_y_offset = g_block_x_offset + y_offset;
+        const int g_output_x_offset = g_block_y_offset + x_offset;
+        // 如果padding连续维为8的倍数，则将连续维的rank使用位运算向上取整为最近的8倍数
+        const int g_output_inner_rank =
+            (padding_last_dim_to_8x) ? (rows + 7) & -8 : rows;
+        if constexpr (is_combine) {
+          const int g_scale_inner_rank = (rows + 127) / 128;
+          if (g_output_y_offset < cols / 2 &&
+              g_output_x_offset < g_output_inner_rank) {
+            // 如果超过了原有的rank，则必然是padding情况，顺手使用0进行padding
+            out[g_output_y_offset * g_output_inner_rank + g_output_x_offset] =
+                (g_output_x_offset < rows) ? output_scaled_fp8 : (OutT)0;
+            scales[g_output_y_offset * g_scale_inner_rank +
+                   g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+          }
+        } else {
+          const int g_scale_inner_rank = (rows + 127) / 128;
+          if (g_output_y_offset < cols &&
+              g_output_x_offset < g_output_inner_rank) {
+            out[g_output_y_offset * g_output_inner_rank + g_output_x_offset] =
+                (g_output_x_offset < rows) ? output_scaled_fp8 : (OutT)0;
+            scales[g_output_y_offset * g_scale_inner_rank +
+                   g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+          }
+        }
+      } else {
+        const int g_output_y_offset = g_block_y_offset + y_offset;
+        const int g_output_x_offset = g_block_x_offset + x_offset;
+        if constexpr (is_combine) {
+          const int g_scale_inner_rank = (cols / 2 + 127) / 128;
+          const int g_output_inner_rank =
+              (padding_last_dim_to_8x) ? ((cols / 2 + 7) & -8) : (cols / 2);
+          if (g_output_y_offset < rows &&
+              g_output_x_offset < g_output_inner_rank) {
+            out[g_output_y_offset * g_output_inner_rank + g_output_x_offset] =
+                (g_output_x_offset < cols / 2) ? output_scaled_fp8 : (OutT)0;
+            scales[g_output_y_offset * g_scale_inner_rank +
+                   g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+          }
+        } else {
+          const int g_scale_inner_rank = (cols + 127) / 128;
+          const int g_output_inner_rank =
+              (padding_last_dim_to_8x) ? ((cols + 7) & -8) : (cols);
+          if (g_output_y_offset < rows &&
+              g_output_x_offset < g_output_inner_rank) {
+            out[g_output_y_offset * g_output_inner_rank + g_output_x_offset] =
+                (g_output_x_offset < cols) ? output_scaled_fp8 : (OutT)0;
+            scales[g_output_y_offset * g_scale_inner_rank +
+                   g_output_x_offset / 128] = scale_on_fp8_to_inputT;
+          }
+        }
+      }
+    }
+  }
+}
+
+template <typename outT>
+void dispatch_fused_swiglu_act_quant(const paddle::Tensor &X,
+                                     const paddle::optional<paddle::Tensor> &Y,
+                                     paddle::Tensor &out,
+                                     paddle::Tensor &scale,
+                                     const int rows,
+                                     const int cols,
+                                     const bool &transpose_output,
+                                     const bool &using_pow2_scaling,
+                                     const bool &padding_last_dim_to_8x,
+                                     const float TYPE_MAX) {
+  const bool is_combine = Y ? false : true;
+  dim3 grid;
+  dim3 block(32, 32);
+  DISPATCH_BOOL(
+      is_combine,
+      k_is_combine,
+      DISPATCH_BOOL(
+          transpose_output,
+          k_transpose_output,
+          DISPATCH_BOOL(
+              using_pow2_scaling,
+              k_using_pow2_scaling,
+              DISPATCH_BOOL(
+                  padding_last_dim_to_8x,
+                  k_padding_last_dim_to_8x,
+                  if constexpr (k_is_combine) {
+                    grid.y = (rows + 127) / 128;
+                    grid.x = ((cols / 2) + 127) / 128;
+                  } else {
+                    grid.y = (rows + 127) / 128;
+                    grid.x = (cols + 127) / 128;
+                  } LAUNCH_FUSED_SWIGLU_ACT_QUANT(k_is_combine,
+                                                  k_transpose_output,
+                                                  k_using_pow2_scaling,
+                                                  k_padding_last_dim_to_8x);))))
+}
+std::vector<paddle::Tensor> fused_swiglu_act_quant(
+    const paddle::Tensor &X,
+    const paddle::optional<paddle::Tensor> &Y,
+    const bool &transpose_output,
+    const bool &to_e4m3,
+    const bool &using_pow2_scaling,
+    const bool &padding_last_dim_to_8x) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);
+  int64_t data_rows = size_to_dim(X.shape().size() - 1, X.shape());
+  int64_t data_cols = X.shape().back();
+  int64_t rows = data_rows, cols = data_cols;
+
+  paddle::Tensor out;
+  paddle::Tensor scale;
+
+  if (Y) {
+    // Y存在
+    PD_CHECK(Y.get().dtype() == paddle::DataType::BFLOAT16);
+    auto Xdims = X.dims();
+    const auto &y_tensor = Y.get();
+    const auto &y_dims = y_tensor.dims();
+    PADDLE_ENFORCE_EQ(y_dims,
+                      Xdims,
+                      common::errors::InvalidArgument(
+                          "The shape of Input(Y):[%s] must be equal "
+                          "to the shape of Input(X):[%s].",
+                          y_dims,
+                          Xdims));
+    if (transpose_output) {
+      if (padding_last_dim_to_8x) {
+        rows = ((data_rows + 7) / 8) *
+               8;  // 向上padding到8的倍数, 因为128为8的倍数，不影响scale shape
+      }
+      if (to_e4m3) {
+        out = paddle::empty(
+            {cols, rows}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+      } else {
+        out = paddle::empty(
+            {cols, rows}, paddle::DataType::FLOAT8_E5M2, X.place());
+      }
+      scale = paddle::empty(
+          {cols, (rows + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+    } else {
+      if (padding_last_dim_to_8x) {
+        cols = ((data_cols + 7) / 8) *
+               8;  // 向上padding到8的倍数, 因为128为8的倍数，不影响scale shape
+      }
+      if (to_e4m3) {
+        out = paddle::empty(
+            {rows, cols}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+      } else {
+        out = paddle::empty(
+            {rows, cols}, paddle::DataType::FLOAT8_E5M2, X.place());
+      }
+      scale = paddle::empty(
+          {rows, (cols + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+    }
+  } else {
+    // Y不存在时，X的column是输出column的两倍
+    auto Xdims = X.dims();
+    auto n = Xdims.at(Xdims.size() - 1);
+    PADDLE_ENFORCE_EQ(n % 2,
+                      0,
+                      common::errors::InvalidArgument(
+                          "The last dim of Input(X) should be exactly divided "
+                          "by 2 when Input(Y) is None, but got %d",
+                          n));
+    if (transpose_output) {
+      if (padding_last_dim_to_8x) {
+        rows =
+            ((data_rows + 7) / 8) * 8;  // rows 向上padding到8的倍数,
+                                        // 因为128为8的倍数，不影响scale shape
+      }
+      if (to_e4m3) {
+        out = paddle::empty(
+            {cols / 2, rows}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+      } else {
+        out = paddle::empty(
+            {cols / 2, rows}, paddle::DataType::FLOAT8_E5M2, X.place());
+      }
+      scale = paddle::empty(
+          {cols / 2, (rows + 127) / 128}, paddle::DataType::FLOAT32, X.place());
+    } else {
+      if (padding_last_dim_to_8x) {
+        cols = ((data_cols / 2 + 7) / 8) *
+               16;  // col/2 向上padding到8的倍数,
+                    // 因为128为8的倍数，不影响scale shape
+      }
+      if (to_e4m3) {
+        out = paddle::empty(
+            {rows, cols / 2}, paddle::DataType::FLOAT8_E4M3FN, X.place());
+      } else {
+        out = paddle::empty(
+            {rows, cols / 2}, paddle::DataType::FLOAT8_E5M2, X.place());
+      }
+      scale = paddle::empty({rows, ((cols / 2) + 127) / 128},
+                            paddle::DataType::FLOAT32,
+                            X.place());
+    }
+  }
+  if (to_e4m3) {
+    dispatch_fused_swiglu_act_quant<phi::float8_e4m3fn>(
+        X,
+        Y,
+        out,
+        scale,
+        data_rows,
+        data_cols,
+        transpose_output,
+        using_pow2_scaling,
+        padding_last_dim_to_8x,
+        F8LimitsTrait<__nv_fp8_e4m3>::max);
+  } else {
+    dispatch_fused_swiglu_act_quant<phi::float8_e5m2>(
+        X,
+        Y,
+        out,
+        scale,
+        data_rows,
+        data_cols,
+        transpose_output,
+        using_pow2_scaling,
+        padding_last_dim_to_8x,
+        F8LimitsTrait<__nv_fp8_e5m2>::max);
+  }
+  return {out, scale};
+}
+
+PD_BUILD_OP(fused_swiglu_act_quant)
+    .Inputs({"X", paddle::Optional("Y")})
+    .Outputs({"output", "scale"})
+    .Attrs({"transpose_output: bool",
+            "to_e4m3: bool",
+            "using_pow2_scaling: bool",
+            "padding_last_dim_to_8x: bool"})
+    .SetKernelFn(PD_KERNEL(fused_swiglu_act_quant));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_swiglu_probs_bwd.cu

@@ -0,0 +1,316 @@
+// 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.
+
+// swiglu_probs_grad_op.cu
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+#include "paddle/extension.h"
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+#include <cuda_bf16.h>
+using BFloat16 = __nv_bfloat16;
+#else
+struct BFloat16 {
+  uint16_t x;
+
+  __host__ __device__ BFloat16() : x(0) {}
+
+  __host__ __device__ BFloat16(float val) {
+    uint32_t* val_bits = reinterpret_cast<uint32_t*>(&val);
+    x = static_cast<uint16_t>(*val_bits >> 16);
+  }
+
+  __host__ __device__ operator float() const {
+    uint32_t bits = static_cast<uint32_t>(x) << 16;
+    return *reinterpret_cast<float*>(&bits);
+  }
+};
+#endif
+
+template <int thread_per_block>
+__global__ void SwigluProbsGradKernel(
+    const BFloat16* o1,           // [seq_len*topk, moe_intermediate_size*2]
+    const BFloat16* do2_s,        // [seq_len*topk, moe_intermediate_size]
+    const float* unzipped_probs,  // [seq_len*topk, 1]
+    BFloat16* do1,                // [seq_len*topk, moe_intermediate_size*2]
+    float* probs_grad,            // [seq_len*topk, 1]
+    BFloat16* o2_s,               // [seq_len*topk, moe_intermediate_size]
+    int moe_intermediate_size) {
+  const int row_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+
+  const BFloat16* o1_row = o1 + row_idx * moe_intermediate_size * 2;
+  const BFloat16* do2_s_row = do2_s + row_idx * moe_intermediate_size;
+  BFloat16* do1_row = do1 + row_idx * moe_intermediate_size * 2;
+  BFloat16* o2s_row = o2_s + row_idx * moe_intermediate_size;
+
+  float prob = unzipped_probs[row_idx];
+
+  __shared__ float sum_buffer[thread_per_block];
+
+  float local_probs_grad = 0.0f;
+
+  for (int i = tid; i < moe_intermediate_size; i += blockDim.x) {
+    float lhs = static_cast<float>(o1_row[i]);
+    float rhs = static_cast<float>(o1_row[i + moe_intermediate_size]);
+
+    float sig = 1.0f / (1.0f + expf(-lhs));
+    float tmp = sig * lhs;
+    float o2_val = tmp * rhs;
+
+    float do2_s_val = static_cast<float>(do2_s_row[i]);
+    float do2_val = do2_s_val * prob;
+
+    float x0_grad = do2_val * rhs * sig * (1.0f + lhs - tmp);
+    float x1_grad = do2_val * tmp;
+
+    do1_row[i] = BFloat16(x0_grad);
+    do1_row[i + moe_intermediate_size] = BFloat16(x1_grad);
+    o2s_row[i] = BFloat16(o2_val * prob);
+
+    local_probs_grad += do2_s_val * o2_val;
+  }
+
+  sum_buffer[tid] = local_probs_grad;
+  __syncthreads();
+
+  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+    if (tid < stride) {
+      sum_buffer[tid] += sum_buffer[tid + stride];
+    }
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    probs_grad[row_idx] = sum_buffer[0];
+  }
+}
+
+typedef struct __align__(8) {
+  __nv_bfloat16 x;
+  __nv_bfloat16 y;
+  __nv_bfloat16 z;
+  __nv_bfloat16 w;
+}
+bfloat16x4_t;
+
+__device__ __forceinline__ float4 fast_swiglu_vec4(const bfloat16x4_t& lhs,
+                                                   const bfloat16x4_t& rhs) {
+  const float x_f_x = __bfloat162float(lhs.x);
+  const float x_f_y = __bfloat162float(lhs.y);
+  const float x_f_z = __bfloat162float(lhs.z);
+  const float x_f_w = __bfloat162float(lhs.w);
+
+  const float y_f_x = __bfloat162float(rhs.x);
+  const float y_f_y = __bfloat162float(rhs.y);
+  const float y_f_z = __bfloat162float(rhs.z);
+  const float y_f_w = __bfloat162float(rhs.w);
+
+  const float silu_x = x_f_x * __frcp_rn(1.0f + __expf(-x_f_x));
+  const float silu_y = x_f_y * __frcp_rn(1.0f + __expf(-x_f_y));
+  const float silu_z = x_f_z * __frcp_rn(1.0f + __expf(-x_f_z));
+  const float silu_w = x_f_w * __frcp_rn(1.0f + __expf(-x_f_w));
+
+  return {silu_x * y_f_x, silu_y * y_f_y, silu_z * y_f_z, silu_w * y_f_w};
+}
+
+__device__ __forceinline__ float4 f4_prod(const float4& x_f,
+                                          const float4& y_f) {
+  return {x_f.x * y_f.x, x_f.y * y_f.y, x_f.z * y_f.z, x_f.w * y_f.w};
+}
+__device__ __forceinline__ float4 f4_prod(const float4& x_f, const float& y_f) {
+  return {x_f.x * y_f, x_f.y * y_f, x_f.z * y_f, x_f.w * y_f};
+}
+__device__ __forceinline__ float4 f4_add(const float4& x_f, const float& y_f) {
+  return {x_f.x + y_f, x_f.y + y_f, x_f.z + y_f, x_f.w + y_f};
+}
+__device__ __forceinline__ float4 f4_add(const float4& x_f, const float4& y_f) {
+  return {x_f.x + y_f.x, x_f.y + y_f.y, x_f.z + y_f.z, x_f.w + y_f.w};
+}
+__device__ __forceinline__ float4 f4_sub(const float4& x_f, const float4& y_f) {
+  return {x_f.x - y_f.x, x_f.y - y_f.y, x_f.z - y_f.z, x_f.w - y_f.w};
+}
+__device__ __forceinline__ float4 fast_sig_vec4(const float4& x_vec4) {
+  const float sig_x = __frcp_rn(1.0f + __expf(-x_vec4.x));
+  const float sig_y = __frcp_rn(1.0f + __expf(-x_vec4.y));
+  const float sig_z = __frcp_rn(1.0f + __expf(-x_vec4.z));
+  const float sig_w = __frcp_rn(1.0f + __expf(-x_vec4.w));
+  return {sig_x, sig_y, sig_z, sig_w};
+}
+__device__ __forceinline__ float4
+load_and_cast_float4(const bfloat16x4_t* x_vec4_ptr) {
+  bfloat16x4_t x_vec4 = *x_vec4_ptr;
+  return {
+      static_cast<float>(x_vec4.x),
+      static_cast<float>(x_vec4.y),
+      static_cast<float>(x_vec4.z),
+      static_cast<float>(x_vec4.w),
+  };
+}
+__device__ __forceinline__ void cast_and_store_bf16x4(bfloat16x4_t* dst_ptr,
+                                                      const float4& x_vec4) {
+  *dst_ptr = {static_cast<__nv_bfloat16>(x_vec4.x),
+              static_cast<__nv_bfloat16>(x_vec4.y),
+              static_cast<__nv_bfloat16>(x_vec4.z),
+              static_cast<__nv_bfloat16>(x_vec4.w)};
+}
+__device__ __forceinline__ float mreduce_f4(const float4& x_f4,
+                                            const float4& y_f4) {
+  float x_m = x_f4.x * y_f4.x;
+  float y_m = x_f4.y * y_f4.y;
+  float z_m = x_f4.z * y_f4.z;
+  float w_m = x_f4.w * y_f4.w;
+  return {x_m + y_m + z_m + w_m};
+}
+
+template <int thread_per_block>
+__global__ void SwigluProbsGradKernelVec4(
+    const BFloat16* o1,           // [seq_len*topk, moe_intermediate_size*2]
+    const BFloat16* do2_s,        // [seq_len*topk, moe_intermediate_size]
+    const float* unzipped_probs,  // [seq_len*topk, 1]
+    BFloat16* do1,                // [seq_len*topk, moe_intermediate_size*2]
+    float* probs_grad,            // [seq_len*topk, 1]
+    BFloat16* o2_s,               // [seq_len*topk, moe_intermediate_size]
+    int moe_intermediate_size) {
+  constexpr int numel_per_thread = 4;
+  constexpr int k_warp_size = 32;
+  const int64_t row_idx = blockIdx.x;
+  const int64_t tid = threadIdx.x;
+
+  const BFloat16* o1_row = o1 + row_idx * (int64_t)moe_intermediate_size * 2;
+  const BFloat16* do2_s_row = do2_s + row_idx * (int64_t)moe_intermediate_size;
+  const bfloat16x4_t* o1_row_left_half_vec4 =
+      reinterpret_cast<const bfloat16x4_t*>(o1_row);
+  const bfloat16x4_t* do2_s_row_vec4 =
+      reinterpret_cast<const bfloat16x4_t*>(do2_s_row);
+  const bfloat16x4_t* o1_row_right_half_vec4 =
+      reinterpret_cast<const bfloat16x4_t*>(o1_row +
+                                            (int64_t)moe_intermediate_size);
+  BFloat16* do1_row = do1 + row_idx * (int64_t)moe_intermediate_size * 2;
+  BFloat16* o2s_row = o2_s + row_idx * (int64_t)moe_intermediate_size;
+  bfloat16x4_t* do1_row_vec4 = reinterpret_cast<bfloat16x4_t*>(do1_row);
+  bfloat16x4_t* o2s_row_vec4 = reinterpret_cast<bfloat16x4_t*>(o2s_row);
+
+  float prob = unzipped_probs[row_idx];
+  __shared__ float sum_buffer[thread_per_block];
+
+  float local_probs_grad = 0.0f;
+
+  const int vec_numel = (int64_t)moe_intermediate_size / numel_per_thread;
+  for (int64_t i = tid; i < vec_numel; i += blockDim.x) {
+    float4 lhs_vec4 = load_and_cast_float4(o1_row_left_half_vec4 + i);
+    float4 rhs_vec4 = load_and_cast_float4(o1_row_right_half_vec4 + i);
+    float4 do2_s_val_vec4 = load_and_cast_float4(do2_s_row_vec4 + i);
+    float4 sig_vec4 = fast_sig_vec4(lhs_vec4);
+    float4 tmp_vec4 = f4_prod(sig_vec4, lhs_vec4);
+    float4 o2_val_vec4 = f4_prod(tmp_vec4, rhs_vec4);
+    float4 o2s_val_vec4 = f4_prod(o2_val_vec4, prob);
+    float4 do2_val_vec4 = f4_prod(do2_s_val_vec4, prob);
+    float4 x0_grad_vec4 = f4_prod(
+        do2_val_vec4,
+        f4_prod(rhs_vec4,
+                f4_prod(sig_vec4, (f4_sub(f4_add(lhs_vec4, 1.0f), tmp_vec4)))));
+    float4 x1_grad_vec4 = f4_prod(do2_val_vec4, tmp_vec4);
+    cast_and_store_bf16x4(do1_row_vec4 + i, x0_grad_vec4);
+    cast_and_store_bf16x4(do1_row_vec4 + i + vec_numel, x1_grad_vec4);
+    cast_and_store_bf16x4(o2s_row_vec4 + i, o2s_val_vec4);
+    local_probs_grad += mreduce_f4(do2_s_val_vec4, o2_val_vec4);
+  }
+
+  sum_buffer[tid] = local_probs_grad;
+  __syncthreads();
+
+#pragma unroll
+  for (int stride = blockDim.x / 2; stride >= k_warp_size; stride >>= 1) {
+    if (tid < stride) {
+      sum_buffer[tid] += sum_buffer[tid + stride];
+    }
+    __syncthreads();
+  }
+
+  if (tid < k_warp_size) {
+    local_probs_grad = sum_buffer[tid];
+#pragma unroll
+    for (int offset = k_warp_size / 2; offset > 0; offset >>= 1) {
+      local_probs_grad +=
+          __shfl_down_sync(0xFFFFFFFF, local_probs_grad, offset);
+    }
+  }
+
+  if (tid == 0) {
+    probs_grad[row_idx] = local_probs_grad;
+  }
+}
+
+std::vector<paddle::Tensor> SwigluProbsGradCUDABackward(
+    const paddle::Tensor& o1,
+    const paddle::Tensor& do2_s,
+    const paddle::Tensor& unzipped_probs,
+    bool inplace) {
+  auto o1_dims = o1.dims();
+  int o1_outer_dim = 1;
+  for (int i = 0; i < o1_dims.size() - 1; i++) {
+    o1_outer_dim *= o1_dims[i];
+  }
+
+  const int moe_intermediate_size_2 = o1_dims[o1_dims.size() - 1];
+  const int moe_intermediate_size = moe_intermediate_size_2 / 2;
+
+  auto do1 = inplace ? o1 : paddle::empty_like(o1);
+  auto probs_grad =
+      paddle::empty({o1_outer_dim}, paddle::DataType::FLOAT32, o1.place());
+  auto o2_s = inplace ? do2_s : paddle::empty_like(do2_s);
+
+  const BFloat16* o1_ptr =
+      reinterpret_cast<const BFloat16*>(o1.data<phi::bfloat16>());
+  const BFloat16* do2_s_ptr =
+      reinterpret_cast<const BFloat16*>(do2_s.data<phi::bfloat16>());
+  const float* unzipped_probs_ptr = unzipped_probs.data<float>();
+  BFloat16* do1_ptr = reinterpret_cast<BFloat16*>(do1.data<phi::bfloat16>());
+  float* probs_grad_ptr = probs_grad.data<float>();
+  BFloat16* o2_s_ptr = reinterpret_cast<BFloat16*>(o2_s.data<phi::bfloat16>());
+
+  constexpr int block_size = 256;
+  if (moe_intermediate_size % 4 != 0) {
+    SwigluProbsGradKernel<block_size>
+        <<<o1_outer_dim, block_size, 0, o1.stream()>>>(o1_ptr,
+                                                       do2_s_ptr,
+                                                       unzipped_probs_ptr,
+                                                       do1_ptr,
+                                                       probs_grad_ptr,
+                                                       o2_s_ptr,
+                                                       moe_intermediate_size);
+  } else {
+    SwigluProbsGradKernelVec4<block_size>
+        <<<o1_outer_dim, block_size, 0, o1.stream()>>>(o1_ptr,
+                                                       do2_s_ptr,
+                                                       unzipped_probs_ptr,
+                                                       do1_ptr,
+                                                       probs_grad_ptr,
+                                                       o2_s_ptr,
+                                                       moe_intermediate_size);
+  }
+
+
+  return {do1, probs_grad, o2_s};
+}
+
+PD_BUILD_OP(fused_swiglu_probs_bwd)
+    .Inputs({"o1", "do2_s", "unzipped_probs"})
+    .Outputs({"do1", "probs_grad", "o2_s"})
+    .Attrs({"inplace : bool"})
+    .SetKernelFn(PD_KERNEL(SwigluProbsGradCUDABackward));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/fused_transpose_split_quant.cu

@@ -0,0 +1,301 @@
+#include "quant_utils.h"
+
+template <typename T, int VecSize>
+struct __align__(sizeof(T) * VecSize) VecType {
+  T val[VecSize];
+  __host__ __device__ inline T& operator[](size_t i) { return val[i]; }
+  __host__ __device__ inline const T& operator[](size_t i) const {
+    return val[i];
+  }
+};
+
+template <int VecSize>
+__device__ void BlockLoad(const phi::bfloat16* X,
+                          __nv_bfloat16 input[4][4],
+                          size_t K) {
+  for (size_t i = 0; i < 4; i++) {
+    size_t off_m = blockIdx.x * 128 + threadIdx.y + i * 32;
+    size_t off_k = blockIdx.y * 128 + threadIdx.x * VecSize;
+    size_t offset = off_m * K + off_k;
+
+    for (size_t j = 0; j < 4; j += VecSize) {
+      if (off_k + j * 32 < K) {
+        size_t idx = offset + j * 32;
+        using LoadT = VecType<__nv_bfloat16, VecSize>;
+        LoadT data = *reinterpret_cast<const LoadT*>(X + idx);
+        for (int k = 0; k < VecSize; k++) {
+          input[i][j + k] = data[k];
+        }
+      }
+    }
+  }
+}
+
+__device__ void BlockColumnMax(const __nv_bfloat16 input[4][4],
+                               __nv_bfloat16 amax[4],
+                               __nv_bfloat16* shm) {
+  // Reduce [(4), 32, 32, 4] => [32, 32, 4]
+  __nv_bfloat16 warp_max[4];
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      __nv_bfloat16 t = __habs(input[i][j]);
+      warp_max[j] = i == 0 ? t : __hmax(warp_max[j], t);
+    }
+  }
+
+  // Reduce [(32), 32, 4] => [32, 4]
+  for (int i = 0; i < 4; i++) {
+    shm[threadIdx.y * 128 + i * 32 + threadIdx.x] = warp_max[i];
+  }
+  __syncthreads();
+  for (int offset = 16; offset > 0; offset /= 2) {
+    if (threadIdx.y < offset) {
+      for (int i = 0; i < 4; i++) {
+        shm[threadIdx.y * 128 + i * 32 + threadIdx.x] =
+            __hmax(shm[threadIdx.y * 128 + i * 32 + threadIdx.x],
+                   shm[(threadIdx.y + offset) * 128 + i * 32 + threadIdx.x]);
+      }
+    }
+    __syncthreads();
+  }
+
+  for (int i = 0; i < 4; i++) {
+    amax[i] = shm[i * 32 + threadIdx.x];
+  }
+  __syncthreads();
+}
+
+template <typename OutT, bool Pow2Scales, int VecSize>
+__device__ void BlockStoreScale(float* scale,
+                                size_t off_m,
+                                __nv_bfloat16 amax[4],
+                                float scale_inv[4],
+                                size_t K) {
+  float scale_out[4];
+  for (int i = 0; i < 4; i++) {
+    scale_inv[i] = ComputeScale<__nv_bfloat16, OutT, Pow2Scales>(
+        static_cast<float>(amax[i]), 0.0f);
+    scale_out[i] = __frcp_rn(scale_inv[i]);
+  }
+  if (threadIdx.y == 0) {
+    size_t idx_m = blockIdx.x - off_m / 128;
+    size_t off_k = blockIdx.y * 128 + threadIdx.x * VecSize;
+    size_t offset = idx_m * K + off_k;
+
+    for (size_t j = 0; j < 4; j += VecSize) {
+      if (off_k + j * 32 < K) {
+        size_t idx = offset + j * 32;
+        using StoreT = VecType<float, VecSize>;
+        StoreT data;
+        for (int k = 0; k < VecSize; k++) {
+          data[k] = scale_out[j + k];
+        }
+        *reinterpret_cast<StoreT*>(scale + idx) = data;
+      }
+    }
+  }
+}
+
+template <typename OutT, int VecSize>
+__device__ void BlockStoreOut(OutT* out,
+                              size_t off_m,
+                              size_t cur_tokens,
+                              const OutT shm[128][129],
+                              size_t K) {
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_m = blockIdx.x * 128 + threadIdx.x * 4;
+    size_t idx_k = blockIdx.y * 128 + threadIdx.y + i * 32;
+    size_t idx = idx_k * cur_tokens + (idx_m - off_m);
+
+    if (idx_k < K) {
+      using StoreT = VecType<OutT, VecSize>;
+      StoreT data;
+      for (int j = 0; j < VecSize; j++) {
+        data[j] = shm[i * 32 + threadIdx.y][threadIdx.x * 4 + j];
+      }
+      *reinterpret_cast<StoreT*>(out + idx) = data;
+    }
+  }
+}
+
+__device__ std::pair<size_t, size_t> GetExpertIdx(int64_t* tokens_per_expert,
+                                                  size_t num_experts) {
+  __shared__ size_t expert_idx_, off_m_;
+
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    size_t idx_m = blockIdx.x * 128;
+    size_t off_m = 0, next_off_m = 0;
+    size_t expert_idx;
+    for (expert_idx = 0; expert_idx < num_experts; expert_idx++) {
+      next_off_m += tokens_per_expert[expert_idx];
+      if (idx_m >= off_m && idx_m < next_off_m) {
+        break;
+      }
+      off_m = next_off_m;
+    }
+    expert_idx_ = expert_idx;
+    off_m_ = off_m;
+  }
+  __syncthreads();
+
+  return {expert_idx_, off_m_};
+}
+
+template <typename OutT, bool Pow2Scales, int VecSize>
+__global__ void __launch_bounds__(1024)
+    FusedTransposeSplitQuantKernel(const phi::bfloat16* __restrict__ X,
+                                   int64_t* __restrict__ meta,
+                                   size_t num_experts,
+                                   size_t K) {
+  __shared__ OutT shm[128][129];
+  int64_t* tokens_per_expert = meta;
+  OutT** out_ptrs = reinterpret_cast<OutT**>(meta + num_experts);
+  float** scale_ptrs = reinterpret_cast<float**>(meta + num_experts * 2);
+
+  // Get expert_idx and offset at the M dim of the current block
+  auto expert_info = GetExpertIdx(tokens_per_expert, num_experts);
+  size_t expert_idx = expert_info.first;
+  size_t off_m = expert_info.second;
+
+  // Load 128x128 elements from X
+  __nv_bfloat16 input[4][4];
+  BlockLoad<VecSize>(X, input, K);
+
+  // Find the maximum of each 128 elements on the M axis
+  __nv_bfloat16 amax[4];
+  BlockColumnMax(input, amax, reinterpret_cast<__nv_bfloat16*>(shm));
+
+  // Compute scale and scale_inv, then store scale back
+  float scale_inv[4];
+  BlockStoreScale<OutT, Pow2Scales, VecSize>(
+      scale_ptrs[expert_idx], off_m, amax, scale_inv, K);
+
+  // Scale X and save into shared memory with transposed layout
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j += VecSize) {
+      for (int k = 0; k < VecSize; k++) {
+        float input_fp32 = static_cast<float>(input[i][j + k]);
+        float output_scaled = input_fp32 * scale_inv[j + k];
+        shm[threadIdx.x * VecSize + j * 32 + k][i * 32 + threadIdx.y] =
+            static_cast<OutT>(output_scaled);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Store 128x128 elements back
+  // Note: out is always 4x vectorizable.
+  BlockStoreOut<OutT, 4>(
+      out_ptrs[expert_idx], off_m, tokens_per_expert[expert_idx], shm, K);
+}
+
+/**
+ * Quantize on dim[0] of X, transpose dim[0] and dim[1] of X, then
+ * split the result into out and scale.
+ *
+ * Inputs:
+ *   X     : [SUM(M_1...M_N), K], bfloat16
+ *
+ * Outputs:
+ *   out   : {[K, M_1], [K, M_2], ..., [K, M_N]}, float8_e4m3fn
+ *   scale : {[M_1/128, K], [M_2/128, K], ..., [M_N/128, K]}, float
+ *
+ * Attrs:
+ *   pow_2_scales
+ *         : bool that indicates whether to use power-of-2 scaling
+ *
+ * Requirements:
+ *   1) M_i % 128 == 0 for M_i in [M_1, M_2, ..., M_N]
+ *   2) K <= 65535 * 128
+ */
+void fused_transpose_split_quant(const paddle::Tensor& X,
+                                 std::vector<paddle::Tensor>& outs,
+                                 std::vector<paddle::Tensor>& scales,
+                                 bool pow_2_scales) {
+  // Check X
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);
+
+  std::vector<int64_t> shape = X.shape();
+  PD_CHECK(shape.size() == 2);
+  const int64_t M = shape[0];
+  const int64_t K = shape[1];
+
+  // Check outs and scales
+  const size_t num_experts = outs.size();
+  PD_CHECK(scales.size() == num_experts);
+
+  std::vector<int64_t> tokens_per_expert;
+  int64_t sum_tokens = 0;
+  for (size_t i = 0; i < num_experts; i++) {
+    PD_CHECK(outs[i].dtype() == paddle::DataType::FLOAT8_E4M3FN);
+    PD_CHECK(scales[i].dtype() == paddle::DataType::FLOAT32);
+
+    std::vector<int64_t> out_shape = outs[i].shape();
+    PD_CHECK(out_shape.size() == 2);
+    PD_CHECK(out_shape[0] == K);
+    PD_CHECK(out_shape[1] % 128 == 0);
+    tokens_per_expert.push_back(out_shape[1]);
+    sum_tokens += out_shape[1];
+
+    std::vector<int64_t> scale_shape = scales[i].shape();
+    PD_CHECK(scale_shape.size() == 2);
+    PD_CHECK(scale_shape[0] == out_shape[1] / 128);
+    PD_CHECK(scale_shape[1] == K);
+  }
+
+  PD_CHECK(sum_tokens == M,
+           "sum of out[i].shape[1] must be equal to X.shape[0]");
+  PD_CHECK(K <= 65535 * 128, "only supports K <= 65535 * 128");
+
+  // Skip 0-size
+  if (M == 0 || K == 0) {
+    return;
+  }
+
+  // Copy meta (tokens_per_expert, out_ptrs, scale_ptrs) to device
+  paddle::Tensor meta_cpu = paddle::empty(
+      {static_cast<int64_t>(num_experts * 3)}, paddle::DataType::INT64);
+  int64_t* meta_ptr = meta_cpu.data<int64_t>();
+  for (size_t i = 0; i < num_experts; i++) {
+    meta_ptr[i] = static_cast<int64_t>(tokens_per_expert[i]);
+  }
+  for (size_t i = 0; i < num_experts; i++) {
+    meta_ptr[num_experts + i] =
+        reinterpret_cast<int64_t>(outs[i].data<phi::float8_e4m3fn>());
+  }
+  for (size_t i = 0; i < num_experts; i++) {
+    meta_ptr[num_experts * 2 + i] =
+        reinterpret_cast<int64_t>(scales[i].data<float>());
+  }
+  paddle::Tensor meta_gpu = meta_cpu.copy_to(X.place(), /*blocking=*/false);
+
+  // Launch kernel
+  dim3 grid(M / 128, (K + 127) / 128);
+  dim3 block(32, 32);
+
+#define LAUNCH_KERNEL(POW_2_SCALES, VEC_SIZE)                                \
+  FusedTransposeSplitQuantKernel<phi::float8_e4m3fn, POW_2_SCALES, VEC_SIZE> \
+      <<<grid, block, 0, X.stream()>>>(                                      \
+          X.data<phi::bfloat16>(), meta_gpu.data<int64_t>(), num_experts, K);
+#define LAUNCH_KERNEL_PARTIAL(VEC_SIZE) \
+  if (pow_2_scales) {                   \
+    LAUNCH_KERNEL(true, VEC_SIZE);      \
+  } else {                              \
+    LAUNCH_KERNEL(false, VEC_SIZE);     \
+  }
+
+  if (K % 4 == 0) {
+    LAUNCH_KERNEL_PARTIAL(4);
+  } else if (K % 2 == 0) {
+    LAUNCH_KERNEL_PARTIAL(2);
+  } else {
+    LAUNCH_KERNEL_PARTIAL(1);
+  }
+#undef LAUNCH_KERNEL_PARTIAL
+#undef LAUNCH_KERNEL
+}
+
+PD_BUILD_OP(fused_transpose_split_quant)
+    .Inputs({"X", paddle::Vec("outs"), paddle::Vec("scales")})
+    .Attrs({"pow_2_scales: bool"})
+    .SetKernelFn(PD_KERNEL(fused_transpose_split_quant));

slm/model_zoo/gpt-3/external_ops/fused_quanted_ops/quant_utils.h

@@ -0,0 +1,169 @@
+#pragma once
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <limits>
+
+#include "paddle/extension.h"
+#include "paddle/phi/api/all.h"
+#include "paddle/phi/common/float8_e4m3fn.h"
+#include "paddle/phi/common/float8_e5m2.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+#define DISPATCH_BOOL(condition, ConstName, ...) \
+  {                                              \
+    if (condition) {                             \
+      constexpr bool ConstName = true;           \
+      { __VA_ARGS__ }                            \
+    } else {                                     \
+      constexpr bool ConstName = false;          \
+      { __VA_ARGS__ }                            \
+    }                                            \
+  }
+
+// 对二维坐标进行swizzle变换，提供相对offset,避免bank conflict
+__device__ __forceinline__ int swizzled_2d_idx(const int outer_dim,
+                                               const int inner_rank,
+                                               const int inner_dim) {
+  return outer_dim * inner_rank + outer_dim ^ inner_dim;
+}
+// ------------------------------ Numerical Part (from
+// kitchen)--------------------------- Type trait for extreme values of fp8
+// types. Used in the calculation of scale factors as a constexpr lookup from
+// e4m3 or e5m2 to the max finite value.
+template <typename T>
+struct F8LimitsTrait;
+
+template <>
+struct F8LimitsTrait<__nv_fp8_e4m3> {
+  static constexpr float max = 448.0f;
+};
+template <>
+struct F8LimitsTrait<phi::float8_e4m3fn> {
+  static constexpr float max = 448.0f;
+};
+
+template <>
+struct F8LimitsTrait<__nv_fp8_e5m2> {
+  static constexpr float max = 57344.0f;
+};
+template <>
+struct F8LimitsTrait<phi::float8_e5m2> {
+  static constexpr float max = 57344.0f;
+};
+
+// Type trait to resolve the max finite value
+// represented by a input type to quantization.
+// Or to represent max representable power of 2
+// finite value.
+template <typename T, bool ForcePow2>
+struct HighPrecisionFloatScaleLimitsTrait;
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<float, false> {
+  static constexpr float max = std::numeric_limits<float>::max();
+};
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<float, true> {
+  // Hex float format of 1.0 * 2 ^ 127
+  static constexpr float max = 0x1.0p127;
+};
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<nv_bfloat16, false> {
+  // Hex float format of 1.(7 bits of 1) * 2 ^ 127
+  static constexpr float max = 0x1.FEp127;
+};
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<nv_bfloat16, true> {
+  // Hex float format of 1.0 * 2 ^ 127
+  static constexpr float max = 0x1.0p127;
+};
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<half, false> {
+  // Hex float format of 1.(10 bits of 1) * 2 ^ 15
+  static constexpr float max = 0x1.FFCp15;
+};
+
+template <>
+struct HighPrecisionFloatScaleLimitsTrait<half, true> {
+  // Hex float format of 1.0 * 2 ^ 15
+  static constexpr float max = 0x1.0p15;
+};
+// ----------------------------- Scale Part ---------------------------
+// Calculate the quantization scale for an individual data element
+// given the amax(abs(tile)) value for a given quantization tile.
+//
+//
+// Arguments:
+// IType: data type of the tensor being quantized (float or bf16)
+// OType: quantized data type (e4m3 or e5m2)
+// pow_2_scaling: Whether to force the scale to be a power of 2.
+// amax: The evaluation of amax(abs(tile)) for the quantization tile.
+// eps: An epsilon used as a floor for amax.
+template <typename IType, typename OType, bool Power2Scaling = false>
+__device__ __forceinline__ float ComputeScale(const float amax,
+                                              const float eps) {
+  constexpr float fp8_max = F8LimitsTrait<OType>::max;
+
+  // Clamping amax to avoid division by small numbers
+  float amax_mod = fmaxf(amax, eps);
+
+  // Handle overflow cases for non-clamped amax (eps is 0 or very small)
+  if (amax_mod == 0.f) {
+    // If amax is 0, return 1
+    return 1.f;
+  }
+  // Compute scale factor
+  float scale = fp8_max / amax_mod;
+
+  if (isinf(scale)) {
+    // If scale is infinity, return max value of IType
+    return HighPrecisionFloatScaleLimitsTrait<IType, Power2Scaling>::max;
+  }
+  if (scale == 0.0) {
+    return scale;
+  }
+  if constexpr (Power2Scaling) {
+    uint32_t scale_bits = *reinterpret_cast<uint32_t *>(&scale);
+    // Scale must be positive, shift it
+    uint8_t exp = scale_bits >> 23;
+
+    // inf scales already early returned, as did nan scales.
+    // The cases to consider here are normals, zero, and subnormals.
+    // zero is not possible with current math as
+    // 448.0 / float_max == 1.31655e-36, which is the smallest
+    // possible scale given current dtypes. It is still in the normal
+    // fp32 range with an exponent of -120, so subnormals are also
+    // not possible.
+    int32_t normal_biased_exp = static_cast<int32_t>(exp) - 127;
+    __builtin_assume(exp != 0);
+    // Normal numbers case.
+
+    scale = ldexpf(1.0f, normal_biased_exp);
+  }
+  return scale;
+}
+// -------------------------------------- From Kitchen
+// ----------------------------------
+
+inline int64_t size_to_dim(size_t k, std::vector<int64_t> dims) {
+  PD_CHECK(k >= 0 && k <= dims.size());
+  int64_t r = 1;
+  for (size_t i = 0; i < k; ++i) {
+    r *= dims[i];
+  }
+  return r;
+}
+
+__device__ __forceinline__ float warpReduceMax(float val) {
+  for (int offset = 16; offset > 0; offset /= 2)
+    val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
+  return val;
+}

slm/model_zoo/gpt-3/external_ops/setup_fp8.py

@@ -0,0 +1,112 @@
+# 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 multiprocessing
+import os
+
+def run(func):
+    p = multiprocessing.Process(target=func)
+    p.start()
+    p.join()
+
+
+def change_pwd():
+    path = os.path.dirname(__file__)
+    if path:
+        os.chdir(path)
+
+
+def setup_fused_quant_ops():
+    """setup_fused_fp8_ops"""
+    from paddle.utils.cpp_extension import CUDAExtension, setup
+
+    change_pwd()
+    setup(
+        name="FusedQuantOps",
+        ext_modules=CUDAExtension(
+            sources=[
+                "fused_quanted_ops/fused_swiglu_act_quant.cu",
+                "fused_quanted_ops/fused_act_quant.cu",
+                "fused_quanted_ops/fused_act_dequant.cu",
+                "fused_quanted_ops/fused_act_dequant_transpose_act_quant.cu",
+                "fused_quanted_ops/fused_swiglu_probs_bwd.cu",
+                "fused_quanted_ops/fused_spaq.cu",
+                "fused_quanted_ops/fused_stack_transpose_quant.cu",
+                "fused_quanted_ops/fused_transpose_split_quant.cu",
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-w", "-Wno-abi", "-fPIC", "-std=c++17"],
+                "nvcc": [
+                    "-O3",
+                    "-U__CUDA_NO_HALF_OPERATORS__",
+                    "-U__CUDA_NO_HALF_CONVERSIONS__",
+                    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+                    "-DCUTE_ARCH_MMA_SM90A_ENABLE",
+                    "--expt-relaxed-constexpr",
+                    "--expt-extended-lambda",
+                    "-lineinfo",
+                    "-DCUTLASS_DEBUG_TRACE_LEVEL=0",
+                    "-maxrregcount=50",
+                    "-gencode=arch=compute_90a,code=sm_90a",
+                    "-DNDEBUG",
+                ]
+            },
+        ),
+    )
+
+
+def setup_token_dispatcher_utils():
+    from paddle.utils.cpp_extension import CUDAExtension, setup
+
+    change_pwd()
+    setup(
+        name="TokenDispatcherUtils",
+        ext_modules=CUDAExtension(
+            sources=[
+                "token_dispatcher_utils/topk_to_multihot.cu",
+                "token_dispatcher_utils/topk_to_multihot_grad.cu",
+                "token_dispatcher_utils/tokens_unzip_and_zip.cu",
+                "token_dispatcher_utils/tokens_stable_unzip.cu",
+                "token_dispatcher_utils/tokens_guided_unzip.cu",
+                "token_dispatcher_utils/regroup_tokens.cu",
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-w", "-Wno-abi", "-fPIC", "-std=c++17"],
+                "nvcc": [
+                    "-O3",
+                    "-U__CUDA_NO_HALF_OPERATORS__",
+                    "-U__CUDA_NO_HALF_CONVERSIONS__",
+                    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+                    "-DCUTE_ARCH_MMA_SM90A_ENABLE",
+                    "--expt-relaxed-constexpr",
+                    "--expt-extended-lambda",
+                    "-maxrregcount=32",
+                    "-lineinfo",
+                    "-DCUTLASS_DEBUG_TRACE_LEVEL=0",
+                    "-gencode=arch=compute_90a,code=sm_90a",
+                    "-DNDEBUG",
+                ],
+            },
+        ),
+    )
+
+
+run(setup_token_dispatcher_utils)
+run(setup_fused_quant_ops)

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/regroup_tokens.cu

@@ -0,0 +1,114 @@
+#include "utils.h"
+
+__global__ void regroup_tokens_kernel(const phi::bfloat16* __restrict__ X_in,
+                                      const phi::bfloat16* __restrict__ dout_in,
+                                      const int* __restrict__ expert_idx,
+                                      const int original_token_num,
+                                      const int h1,
+                                      const int h2,
+                                      const int expert_token_num,
+                                      int* __restrict__ atomic_offset_counters,
+                                      phi::bfloat16* grouped_X_out,
+                                      phi::bfloat16* grouped_dout_out) {
+  int row_idx = blockIdx.x;
+  if (row_idx >= original_token_num) return;
+
+  const __nv_bfloat16* X = reinterpret_cast<const __nv_bfloat16*>(X_in);
+  const __nv_bfloat16* dout = reinterpret_cast<const __nv_bfloat16*>(dout_in);
+  __nv_bfloat16* grouped_X = reinterpret_cast<__nv_bfloat16*>(grouped_X_out);
+  __nv_bfloat16* grouped_dout =
+      reinterpret_cast<__nv_bfloat16*>(grouped_dout_out);
+
+  extern __shared__ __nv_bfloat16* target_rowbase_X;
+  extern __shared__ __nv_bfloat16* target_rowbase_dout;
+  if (threadIdx.x == 0) {
+    int target_group = expert_idx[row_idx];
+    int offset_X = target_group * expert_token_num * h1;  // stride = token_len
+    int offset_dout =
+        target_group * expert_token_num * h2;  // stride = feature_len
+    __nv_bfloat16* group_base_X = grouped_X + offset_X;
+    __nv_bfloat16* group_base_dout = grouped_dout + offset_dout;
+    int target_row_idx = atomicAdd(&(atomic_offset_counters[target_group]), 1);
+    target_rowbase_X = group_base_X + target_row_idx * h1;
+    target_rowbase_dout = group_base_dout + target_row_idx * h2;
+  }
+  __syncthreads();
+
+  // X相关的数据搬移
+  /*
+  for (int col_offset = threadIdx.x; col_offset < h1;
+       col_offset += blockDim.x) {
+    target_rowbase_X[col_offset] = X[row_idx * h1 + col_offset];
+  }
+  */
+  vectorized_memcpy(&X[row_idx * h1], target_rowbase_X, h1);
+  // dout相关的数据搬移
+  /*
+  for (int col_offset = threadIdx.x; col_offset < h2;
+       col_offset += blockDim.x) {
+    target_rowbase_dout[col_offset] = dout[row_idx * h2 + col_offset];
+  }
+  */
+  vectorized_memcpy(&dout[row_idx * h2], target_rowbase_dout, h2);
+}
+
+void dispatch_regroup_tokens_kernel(const paddle::Tensor& X,
+                                    const paddle::Tensor& dout,
+                                    const paddle::Tensor& expert_idx,
+                                    const int expert_token_num,
+                                    paddle::Tensor& atomic_offset_counters,
+                                    paddle::Tensor grouped_X,
+                                    paddle::Tensor grouped_dout) {
+  dim3 grid;
+  dim3 block;
+  int original_token_num = X.shape()[0];
+  int h1 = X.shape()[1];
+  int h2 = dout.shape()[1];
+  grid.x = original_token_num;
+  block.x = 256;  // 单block处理单token+dout
+  regroup_tokens_kernel<<<grid, block, 0, X.stream()>>>(
+      X.data<phi::bfloat16>(),
+      dout.data<phi::bfloat16>(),
+      expert_idx.data<int>(),
+      original_token_num,
+      h1,
+      h2,
+      expert_token_num,
+      atomic_offset_counters.data<int>(),
+      grouped_X.data<phi::bfloat16>(),
+      grouped_dout.data<phi::bfloat16>());
+}
+std::vector<paddle::Tensor> regroup_tokens(const paddle::Tensor& X,
+                                           const paddle::Tensor& dout,
+                                           const paddle::Tensor& expert_idx,
+                                           const int& expert_num,
+                                           const int& token_max_per_expert) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16);  // 当前只支持BFLOAT16
+  int seqlen = X.shape()[0];
+  int h1 = X.shape()[1];
+  int h2 = dout.shape()[1];
+  paddle::Tensor grouped_X, grouped_dout;
+
+  // 待优化，padding时使用full算子默认将写不到的元素归为0
+  grouped_X = paddle::zeros(
+      {expert_num * token_max_per_expert, h1}, X.dtype(), X.place());
+  grouped_dout = paddle::zeros(
+      {expert_num * token_max_per_expert, h2}, X.dtype(), X.place());
+  // 将原子计数数组初始化为0
+  auto atomic_offset_counters =
+      paddle::zeros({expert_num}, paddle::DataType::INT32, X.place());
+  dispatch_regroup_tokens_kernel(X,
+                                 dout,
+                                 expert_idx,
+                                 token_max_per_expert,
+                                 atomic_offset_counters,
+                                 grouped_X,
+                                 grouped_dout);
+  return {grouped_X, grouped_dout};
+}
+
+PD_BUILD_OP(regroup_tokens)
+    .Inputs({"X", "dout", "expert_idx"})
+    .Outputs({"grouped_X", "grouped_dout"})
+    .Attrs({"expert_num: int", "token_max_per_expert: int"})
+    .SetKernelFn(PD_KERNEL(regroup_tokens));

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/tokens_guided_unzip.cu

@@ -0,0 +1,195 @@
+#include "utils.h"
+
+// --------------------------- kernels ------------------------
+template <int topk, int num_experts>
+__global__ void probs_topk_guided_unzip_kernel(
+    const phi::bfloat16 *__restrict__ probs_topk_in,
+    const int *__restrict__ expert_routemap_topk,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    phi::bfloat16 *__restrict__ guided_unzipped_probs_1d_out,
+    const int total_zipped_tokens_num) {
+  const int this_row =
+      blockIdx.x * blockDim.x + threadIdx.x;  // 一个线程处理一行topk
+  if (this_row >= total_zipped_tokens_num) return;
+  const __nv_bfloat16 *probs_topk =
+      reinterpret_cast<const __nv_bfloat16 *>(probs_topk_in);
+
+  __nv_bfloat16 token_prob_topk[topk];
+  int token_expert_rowmap[num_experts];
+  int token_expert_route[topk];
+  __nv_bfloat16 *guided_unzipped_probs_1d =
+      reinterpret_cast<__nv_bfloat16 *>(guided_unzipped_probs_1d_out);
+// 使用该行的num_expert规模的行映射信息填充寄存器组，非法值为-1
+#pragma unroll
+  for (int i = 0; i < num_experts; i++) {
+    token_expert_rowmap[i] =
+        zipped_expertwise_rowmap[this_row * num_experts + i];
+  }
+
+// 使用该行的topk规模的prob和route信息填充寄存器组
+#pragma unroll
+  for (int i = 0; i < topk; ++i) {
+    token_prob_topk[i] = probs_topk[this_row * topk + i];
+    token_expert_route[i] = expert_routemap_topk[this_row * topk + i];
+  }
+
+#pragma unroll
+  for (int i = 0; i < topk; ++i) {
+    // 如果route不为-1，则probs亦合法，将该token的prob放入对应的unzipped_probs_1d中
+    if (token_expert_route[i] != -1) {
+      const int routed_expert = token_expert_route[i];
+      const int mapped_row = token_expert_rowmap[routed_expert];
+      guided_unzipped_probs_1d[mapped_row] = token_prob_topk[i];
+    }
+  }
+}
+
+
+template <typename T, int num_experts>
+__global__ void tokens_guided_unzip_kernel(
+    const T *__restrict__ X_in,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    T *__restrict__ guided_unzipped_X_out,
+    const int total_zipped_tokens_num,
+    const int token_length) {
+  const int this_row = blockIdx.x;
+  if (this_row >= total_zipped_tokens_num) return;
+  /*
+  const __nv_bfloat16 *X = reinterpret_cast<const __nv_bfloat16 *>(X_in);
+  __nv_bfloat16 *guided_unzipped_X =
+      reinterpret_cast<__nv_bfloat16 *>(guided_unzipped_X_out);
+  */
+
+  int local_row_pushlist[num_experts];
+// 填充该行token被广播到的rows和对应的概率
+#pragma unroll
+  for (int expert = 0; expert < num_experts; ++expert) {
+    local_row_pushlist[expert] =
+        zipped_expertwise_rowmap[this_row * num_experts + expert];
+  }
+  for (int expert = 0; expert < num_experts; ++expert) {
+    int push_row = local_row_pushlist[expert];
+    // 该专家没被发到，跳过
+    if (push_row == -1) continue;
+    /*
+    for (int i = threadIdx.x; i < token_length; i += blockDim.x) {
+      guided_unzipped_X_out[push_row * token_length + i] =
+          X_in[this_row * token_length + i];
+    }
+    */
+    vectorized_memcpy(&X_in[this_row * token_length],
+                      &guided_unzipped_X_out[push_row * token_length],
+                      token_length);
+  }
+}
+// ------------------------------ Dispatch -----------------------------
+void dispatch_probs_topk_guided_unzip(
+    const paddle::Tensor &probs_topk,
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    paddle::Tensor &guided_unzipped_probs_1d,
+    const int num_experts,
+    const int total_zipped_tokens_num,
+    const int topk) {
+  dim3 grid, block;
+  block.x = 256;
+  grid.x = (total_zipped_tokens_num + block.x - 1) /
+           block.x;  // 每一个thread处理一个probs_topk
+  if (num_experts == 4 && topk == 8) {
+    probs_topk_guided_unzip_kernel<8, 4>
+        <<<grid, block, 0, probs_topk.stream()>>>(
+            probs_topk.data<phi::bfloat16>(),
+            expert_routemap_topk.data<int>(),
+            zipped_expertwise_rowmap.data<int>(),
+            guided_unzipped_probs_1d.data<phi::bfloat16>(),
+            total_zipped_tokens_num);
+  }
+}
+
+void dispatch_tokens_guided_unzip(
+    const paddle::Tensor &X,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    paddle::Tensor &guided_unzipped_X,
+    const int num_experts,
+    const int total_zipped_tokens_num,
+    const int token_length) {
+  dim3 grid, block;
+  grid.x = total_zipped_tokens_num;
+  block.x = 256;
+  if (num_experts == 4) {
+    if (X.dtype() == paddle::DataType::BFLOAT16) {
+      tokens_guided_unzip_kernel<phi::bfloat16, 4>
+          <<<grid, block, 0, X.stream()>>>(
+              X.data<phi::bfloat16>(),
+              zipped_expertwise_rowmap.data<int>(),
+              guided_unzipped_X.data<phi::bfloat16>(),
+              total_zipped_tokens_num,
+              token_length);
+    } else if (X.dtype() == paddle::DataType::FLOAT32) {
+      tokens_guided_unzip_kernel<float, 4>
+          <<<grid, block, 0, X.stream()>>>(X.data<float>(),
+                                           zipped_expertwise_rowmap.data<int>(),
+                                           guided_unzipped_X.data<float>(),
+                                           total_zipped_tokens_num,
+                                           token_length);
+    }
+  }
+}
+
+// -----------------------------------  API ----------------------------------
+std::vector<paddle::Tensor> probs_topk_guided_unzip(
+    const paddle::Tensor &probs_topk,
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    const int &total_unzipped_tokens_num,
+    const int &num_experts,
+    const int &topk) {
+  PD_CHECK(probs_topk.dtype() == paddle::DataType::BFLOAT16);
+  int rows = probs_topk.shape()[0];  // seqlen
+  int cols = probs_topk.shape()[1];  // 一般为8
+  PD_CHECK(topk == cols);
+
+  //------------------------ 输出1张量 ------------------------
+  auto guided_unzipped_probs_1d = paddle::empty(
+      {total_unzipped_tokens_num}, probs_topk.dtype(), probs_topk.place());
+
+  dispatch_probs_topk_guided_unzip(probs_topk,
+                                   expert_routemap_topk,
+                                   zipped_expertwise_rowmap,
+                                   guided_unzipped_probs_1d,
+                                   num_experts,
+                                   rows,
+                                   cols);
+  return {guided_unzipped_probs_1d};
+}
+std::vector<paddle::Tensor> tokens_guided_unzip(
+    const paddle::Tensor &X,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    const int &total_unzipped_tokens_num,
+    const int &num_experts) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16 ||
+           X.dtype() == paddle::DataType::FLOAT32);
+  int rows = X.shape()[0];  // seqlen
+  int cols = X.shape()[1];  // 一般为7168
+
+  //------------------------ 输出1张量 ------------------------
+  auto guided_unzipped_X =
+      paddle::empty({total_unzipped_tokens_num, cols}, X.dtype(), X.place());
+
+  dispatch_tokens_guided_unzip(
+      X, zipped_expertwise_rowmap, guided_unzipped_X, num_experts, rows, cols);
+  return {guided_unzipped_X};
+}
+
+// -----------------------------------  注册 ----------------------------------
+PD_BUILD_OP(tokens_guided_unzip)
+    .Inputs({"X", "zipped_expertwise_rowmap"})
+    .Outputs({"guided_unzipped_X"})
+    .Attrs({"total_unzipped_token_num: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(tokens_guided_unzip));
+
+PD_BUILD_OP(probs_topk_guided_unzip)
+    .Inputs({"probs_topk", "expert_routemap_topk", "zipped_expertwise_rowmap"})
+    .Outputs({"guided_unzipped_probs_1d"})
+    .Attrs({"total_unzipped_token_num: int", "num_experts: int", "topk: int"})
+    .SetKernelFn(PD_KERNEL(probs_topk_guided_unzip));

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/tokens_unzip_and_zip.cu

@@ -0,0 +1,901 @@
+// 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.
+
+#include "utils.h"
+
+template <typename X_T,
+          typename routemap_T,
+          typename probs_T,
+          int topk,
+          int num_experts>
+__global__ void token_unzip_kernel(
+    const X_T *__restrict__ X,
+    const routemap_T *__restrict__ routemap_topk,
+    const probs_T *__restrict__ probs_topk,
+    X_T *__restrict__ X_unzipped,
+    int *__restrict__ zipped_expertwise_rowmap,
+    probs_T *__restrict__ probs_unzipped,
+    int *__restrict__ expert_idx_unzipped,
+    int *__restrict__ atomic_extended_offset_counter,
+    int *__restrict__ row_valid,
+    const int total_zipped_tokens_num,
+    const int total_unzipped_tokens_num,
+    const int token_length) {
+  // 每个线程处理一行数据
+  const int row_idx = blockIdx.x;
+  // 仅在线程组2中被更新，不初始化
+  extern __shared__ int shared_original_row;
+
+  if (row_idx < total_unzipped_tokens_num) [[likely]] {
+    // 线程组0，
+    // 主要处理topk和增广部分的行索引、处理专家广播后的行表、一对一搬移
+    if (row_idx < total_zipped_tokens_num) [[likely]] {
+      // ----------------- 增广行的任务派发逻辑，交给thread0 --------------
+      if (threadIdx.x == 0) [[unlikely]] {
+        // 寄存器加载、存储，消耗2xtopk 个reg
+        // 每行只有一次非广播的机会
+        bool isFirst = true;
+        int local_expert_rowmap[num_experts];
+// 寄存器填入非法值，避免误用（0为合法rowidx）
+#pragma unroll
+        for (int i = 0; i < num_experts; i++) {
+          local_expert_rowmap[i] = -1;
+        }
+        for (int i = 0; i < topk; i++) {
+          routemap_T this_expert_idx = routemap_topk[row_idx * topk + i];
+          probs_T this_expert_prob = probs_topk[row_idx * topk + i];
+          if (this_expert_idx < 0) [[likely]]
+            continue;
+          // 第一次出现，直接搬入
+          if (isFirst) [[likely]] {
+            isFirst = false;
+            probs_unzipped[row_idx] = this_expert_prob;
+            expert_idx_unzipped[row_idx] = this_expert_idx;
+            local_expert_rowmap[this_expert_idx] = row_idx;
+          } else {  // 增广部分, 原子更新行偏置,并计算扩展行索引
+            int extended_row_offset;
+            extended_row_offset =
+                atomicAdd(&atomic_extended_offset_counter[0], 1);
+            int extended_row_idx =
+                total_zipped_tokens_num + extended_row_offset;
+            probs_unzipped[extended_row_idx] = this_expert_prob;
+            expert_idx_unzipped[extended_row_idx] = this_expert_idx;
+            // 处理专家广播后的行表，用于zip进行收集
+            local_expert_rowmap[this_expert_idx] = extended_row_idx;
+          }
+        }
+// ------------------ 更新专家广播后的行表，用于zip进行收集 -----------
+// 将合法值和未被触碰的非法值返回给zipped_expertwise_rowmap
+#pragma unroll
+        for (int i = 0; i < num_experts; i++) {
+          zipped_expertwise_rowmap[row_idx * num_experts + i] =
+              local_expert_rowmap[i];
+          int valid_offset = local_expert_rowmap[i] - total_zipped_tokens_num;
+          // 只给增广行传递信号量，非法值保持为0
+          if (valid_offset >= 0) {
+            atomicExch(&row_valid[valid_offset], row_idx);  // 发送任务信号量
+          }
+        }
+      }
+      // 这个syncthread可能并不必要，但尽可能为了不让线程间差太多，还是这样吧。
+      __syncthreads();
+      // 处理完增广事务，对位搬搬移第一次出现的数据,可用inplace优化
+      vectorized_memcpy(&X[(int64_t)row_idx * (int64_t)token_length],
+                        &X_unzipped[(int64_t)row_idx * (int64_t)token_length],
+                        token_length);
+    } else {  // 线程组1， 忙等、并发处理数据搬移
+      if (threadIdx.x == 0) {
+        int extended_row_offset = row_idx - total_zipped_tokens_num;
+        int local_original_row = -1;
+        // 忙等该行的 row_valid变为非-1的合法值
+        while (local_original_row == -1) {
+          local_original_row = atomicExch(&row_valid[extended_row_offset], -1);
+        }
+        // 传递给同组线程共享
+        shared_original_row = local_original_row;
+      }
+      __syncthreads();  // 所有该组线程都等0号取任务，再搬移数据
+      int original_row = shared_original_row;
+      // 搬
+      vectorized_memcpy(&X[(int64_t)original_row * (int64_t)token_length],
+                        &X_unzipped[(int64_t)row_idx * (int64_t)token_length],
+                        token_length);
+    }
+  }
+}
+
+
+template <int num_experts, bool MP = true>
+__global__ void tokens_weighted_zip_kernel(
+    const phi::bfloat16 *__restrict__ unzipped_tokens_in,
+    const phi::bfloat16 *__restrict__ unzipped_token_probs_in,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    phi::bfloat16 *__restrict__ weighted_zipped_tokens_out,
+    const int total_zipped_tokens_num,
+    const int token_length) {
+  const int this_row = blockIdx.x;
+  if (this_row >= total_zipped_tokens_num) return;
+
+  const __nv_bfloat16 *unzipped_tokens =
+      reinterpret_cast<const __nv_bfloat16 *>(unzipped_tokens_in);
+  const __nv_bfloat16 *probs_unzipped =
+      reinterpret_cast<const __nv_bfloat16 *>(unzipped_token_probs_in);
+  __nv_bfloat16 *weighted_zipped_tokens =
+      reinterpret_cast<__nv_bfloat16 *>(weighted_zipped_tokens_out);
+
+  int local_row_fetchlist[num_experts];
+  __nv_bfloat16 local_expert_problist[num_experts];
+
+// -------------------------初始化任务表 ------------------------
+#pragma unroll
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int fetch_row =
+        zipped_expertwise_rowmap[this_row * num_experts + expert];
+    local_row_fetchlist[expert] = fetch_row;
+    if (fetch_row >= 0) {
+      local_expert_problist[expert] = probs_unzipped[fetch_row];
+    }
+  }
+
+  constexpr int vecSize = 2;  // __nv_bfloat162 = 2 x bfloat16
+  const int num_full_vec = token_length / vecSize;
+  const int remaining_elems = token_length % vecSize;
+  const int thread_stride = blockDim.x * vecSize;
+
+  if constexpr (MP) {
+    // ------------------------ 手动混合精度 ---------------------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      float2 sum = {0.0f, 0.0f};
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &weighted_zipped_tokens[(int64_t)this_row * (int64_t)token_length +
+                                  x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        const int fetch_row_index = fetch_row >= 0 ? fetch_row : 0;
+        // 手动类型提升
+        float2 token_vec =
+            __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162 *>(
+                &unzipped_tokens[(int64_t)fetch_row_index *
+                                     (int64_t)token_length +
+                                 x_offset]));
+        float prob = fetch_row >= 0
+                         ? __bfloat162float(local_expert_problist[expert])
+                         : 0.0f;
+        float2 prob_vec = {prob, prob};
+        sum.x = __fmaf_rn(token_vec.x, prob_vec.x, sum.x);
+        sum.y = __fmaf_rn(token_vec.y, prob_vec.y, sum.y);
+      }
+      // 类型下降为原有精度
+      *out_ptr = __float22bfloat162_rn(sum);
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      float sum = 0.0f;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        int fetch_row_index = fetch_row >= 0 ? fetch_row : 0;
+        float token_val = __bfloat162float(
+            unzipped_tokens[(int64_t)fetch_row_index * (int64_t)token_length +
+                            i]);
+        float prob = fetch_row >= 0
+                         ? __bfloat162float(local_expert_problist[expert])
+                         : 0.0f;
+        sum += prob * token_val;
+      }
+      weighted_zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] =
+          __float2bfloat16_rn(sum);
+    }
+  } else {
+    // ------------------------ BF16 intrinsics 加权累加 -----------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      __nv_bfloat162 sum = {0, 0};
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &weighted_zipped_tokens[(int64_t)this_row * (int64_t)token_length +
+                                  x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        const int fetch_row_index = fetch_row >= 0 ? fetch_row : 0;
+        __nv_bfloat162 token_vec = *reinterpret_cast<const __nv_bfloat162 *>(
+            &unzipped_tokens[(int64_t)fetch_row_index * (int64_t)token_length +
+                             x_offset]);
+        __nv_bfloat16 prob =
+            fetch_row >= 0 ? local_expert_problist[expert] : (__nv_bfloat16)0;
+        __nv_bfloat162 prob_vec = {prob, prob};
+        sum = __hfma2(token_vec, prob_vec, sum);
+      }
+      *out_ptr = sum;
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      __nv_bfloat16 sum = (__nv_bfloat16)0;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        int fetch_row_index = fetch_row >= 0 ? fetch_row : 0;
+        __nv_bfloat16 token_val =
+            unzipped_tokens[(int64_t)fetch_row_index * (int64_t)token_length +
+                            i];
+        __nv_bfloat16 prob =
+            fetch_row >= 0 ? local_expert_problist[expert] : (__nv_bfloat16)0;
+        sum += prob * token_val;
+      }
+      weighted_zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] =
+          sum;
+    }
+  }
+}
+
+template <int MAX_NUM_EXPERTS_C, bool MP = true>
+__global__ void tokens_zip_kernel(
+    const phi::bfloat16 *__restrict__ unzipped_tokens_in,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    const int *__restrict__ expert_routemap_topk,
+    const phi::bfloat16 *__restrict__ unzipped_token_probs,
+    phi::bfloat16 *__restrict__ zipped_tokens_out,
+    phi::bfloat16 *__restrict__ zipped_probs_topk,
+    const int total_zipped_tokens_num,
+    const int token_length,
+    const int num_experts,
+    const int topk) {
+  const int this_row = blockIdx.x;
+  if (this_row >= total_zipped_tokens_num) return;
+
+  const __nv_bfloat16 *unzipped_tokens =
+      reinterpret_cast<const __nv_bfloat16 *>(unzipped_tokens_in);
+  __nv_bfloat16 *zipped_tokens =
+      reinterpret_cast<__nv_bfloat16 *>(zipped_tokens_out);
+
+  int local_row_fetchlist[MAX_NUM_EXPERTS_C];
+
+// -------------------------初始化任务表 ------------------------
+#pragma unroll
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int fetch_row =
+        zipped_expertwise_rowmap[this_row * num_experts + expert];
+    local_row_fetchlist[expert] = fetch_row;
+  }
+
+#pragma unroll
+  for (int k = 0; k < topk; ++k) {
+    const int expert_idx = expert_routemap_topk[this_row * topk + k];
+    if (expert_idx < 0) [[likely]]
+      continue;
+    const int expert_fetch_row = local_row_fetchlist[expert_idx];
+    zipped_probs_topk[this_row * topk + k] =
+        unzipped_token_probs[expert_fetch_row];
+  }
+
+  constexpr int vecSize = 2;  // __nv_bfloat162 = 2 x bfloat16
+  const int num_full_vec = token_length / vecSize;
+  const int remaining_elems = token_length % vecSize;
+  const int thread_stride = blockDim.x * vecSize;
+
+  if constexpr (MP) {
+    // ------------------------ 手动混合精度 ---------------------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      float2 sum = {0.0f, 0.0f};
+      __nv_bfloat162 raw = {0, 0};
+      int aggreg_cnt = 0;
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &zipped_tokens[(int64_t)this_row * (int64_t)token_length + x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        aggreg_cnt++;
+        // 手动类型提升
+        raw = *reinterpret_cast<const __nv_bfloat162 *>(
+            &unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length +
+                             x_offset]);
+        float2 token_vec = __bfloat1622float2(raw);
+        sum.x = __fadd_rn(token_vec.x, sum.x);
+        sum.y = __fadd_rn(token_vec.y, sum.y);
+      }
+      // 选择性类型下降为原有精度
+      *out_ptr = (aggreg_cnt > 1) ? __float22bfloat162_rn(sum) : raw;
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      float sum = 0.0f;
+      __nv_bfloat16 raw = 0;
+      int aggreg_cnt = 0;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        aggreg_cnt++;
+        raw = unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length + i];
+        float token_val = __bfloat162float(raw);
+        sum = __fadd_rn(token_val, sum);
+      }
+      zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] =
+          (aggreg_cnt > 1) ? __float2bfloat16_rn(sum) : raw;
+    }
+  } else {
+    // ------------------------ BF16 intrinsics 加权累加 -----------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      __nv_bfloat162 sum = {0, 0};
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &zipped_tokens[(int64_t)this_row * (int64_t)token_length + x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        __nv_bfloat162 token_vec = *reinterpret_cast<const __nv_bfloat162 *>(
+            &unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length +
+                             x_offset]);
+        sum = __hadd2(sum, token_vec);
+      }
+      *out_ptr = sum;
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      __nv_bfloat16 sum = (__nv_bfloat16)0;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        __nv_bfloat16 token_val =
+            unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length + i];
+        sum = __hadd(sum, token_val);
+      }
+      zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] = sum;
+    }
+  }
+}
+
+template <int MAX_NUM_EXPERTS_C, bool MP = true>
+__global__ void tokens_zip_kernel(
+    const phi::bfloat16 *__restrict__ unzipped_tokens_in,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    const int *__restrict__ expert_routemap_topk,
+    const float *__restrict__ unzipped_token_probs,
+    phi::bfloat16 *__restrict__ zipped_tokens_out,
+    float *__restrict__ zipped_probs_topk,
+    const int total_zipped_tokens_num,
+    const int token_length,
+    const int num_experts,
+    const int topk) {
+  const int this_row = blockIdx.x;
+  if (this_row >= total_zipped_tokens_num) return;
+
+  const __nv_bfloat16 *unzipped_tokens =
+      reinterpret_cast<const __nv_bfloat16 *>(unzipped_tokens_in);
+  __nv_bfloat16 *zipped_tokens =
+      reinterpret_cast<__nv_bfloat16 *>(zipped_tokens_out);
+
+  int local_row_fetchlist[MAX_NUM_EXPERTS_C];
+
+// -------------------------初始化任务表 ------------------------
+#pragma unroll
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int fetch_row =
+        zipped_expertwise_rowmap[this_row * num_experts + expert];
+    local_row_fetchlist[expert] = fetch_row;
+  }
+
+#pragma unroll
+  for (int k = 0; k < topk; ++k) {
+    const int expert_idx = expert_routemap_topk[this_row * topk + k];
+    if (expert_idx < 0) [[likely]]
+      continue;
+    const int expert_fetch_row = local_row_fetchlist[expert_idx];
+    zipped_probs_topk[this_row * topk + k] =
+        unzipped_token_probs[expert_fetch_row];
+  }
+
+  constexpr int vecSize = 2;  // __nv_bfloat162 = 2 x bfloat16
+  const int num_full_vec = token_length / vecSize;
+  const int remaining_elems = token_length % vecSize;
+  const int thread_stride = blockDim.x * vecSize;
+
+  if constexpr (MP) {
+    // ------------------------ 手动混合精度 ---------------------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      float2 sum = {0.0f, 0.0f};
+      __nv_bfloat162 raw = {0, 0};
+      int aggreg_cnt = 0;
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &zipped_tokens[(int64_t)this_row * (int64_t)token_length + x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        aggreg_cnt++;
+        // 手动类型提升
+        raw = *reinterpret_cast<const __nv_bfloat162 *>(
+            &unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length +
+                             x_offset]);
+        float2 token_vec = __bfloat1622float2(raw);
+        sum.x = __fadd_rn(token_vec.x, sum.x);
+        sum.y = __fadd_rn(token_vec.y, sum.y);
+      }
+      // 选择性类型下降为原有精度
+      *out_ptr = (aggreg_cnt > 1) ? __float22bfloat162_rn(sum) : raw;
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      float sum = 0.0f;
+      __nv_bfloat16 raw = 0;
+      int aggreg_cnt = 0;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        aggreg_cnt++;
+        raw = unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length + i];
+        float token_val = __bfloat162float(raw);
+        sum = __fadd_rn(token_val, sum);
+      }
+      zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] =
+          (aggreg_cnt > 1) ? __float2bfloat16_rn(sum) : raw;
+    }
+  } else {
+    // ------------------------ BF16 intrinsics 加权累加 -----------------------
+    // 齐整区域向量化搬移
+    for (int x_offset = threadIdx.x * vecSize;
+         x_offset < num_full_vec * vecSize;
+         x_offset += thread_stride) {
+      __nv_bfloat162 sum = {0, 0};
+      __nv_bfloat162 *out_ptr = reinterpret_cast<__nv_bfloat162 *>(
+          &zipped_tokens[(int64_t)this_row * (int64_t)token_length + x_offset]);
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        const int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        __nv_bfloat162 token_vec = *reinterpret_cast<const __nv_bfloat162 *>(
+            &unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length +
+                             x_offset]);
+        sum = __hadd2(sum, token_vec);
+      }
+      *out_ptr = sum;
+    }
+
+    // 剩余元素处理
+    for (int i = num_full_vec * vecSize + threadIdx.x; i < token_length;
+         i += blockDim.x) {
+      __nv_bfloat16 sum = (__nv_bfloat16)0;
+#pragma unroll
+      for (int expert = 0; expert < num_experts; ++expert) {
+        int fetch_row = local_row_fetchlist[expert];
+        if (fetch_row < 0) continue;
+        __nv_bfloat16 token_val =
+            unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length + i];
+        sum = __hadd(sum, token_val);
+      }
+      zipped_tokens[(int64_t)this_row * (int64_t)token_length + i] = sum;
+    }
+  }
+}
+
+template <int MAX_NUM_EXPERTS_C>
+__global__ void tokens_zip_kernel(
+    const float *__restrict__ unzipped_tokens,
+    const int *__restrict__ zipped_expertwise_rowmap,
+    const int *__restrict__ expert_routemap_topk,
+    const float *__restrict__ unzipped_token_probs,
+    float *__restrict__ zipped_tokens,
+    float *__restrict__ zipped_probs_topk,
+    const int total_zipped_tokens_num,
+    const int token_length,
+    const int num_experts,
+    const int topk) {
+  const int this_row = blockIdx.x;
+  if (this_row >= total_zipped_tokens_num) return;
+  int local_row_fetchlist[MAX_NUM_EXPERTS_C];
+
+// -------------------------初始化任务表 ------------------------
+#pragma unroll
+  for (int expert = 0; expert < num_experts; ++expert) {
+    const int fetch_row =
+        zipped_expertwise_rowmap[this_row * num_experts + expert];
+    local_row_fetchlist[expert] = fetch_row;
+  }
+
+#pragma unroll
+  for (int k = 0; k < topk; ++k) {
+    const int expert_idx = expert_routemap_topk[this_row * topk + k];
+    if (expert_idx < 0) [[likely]]
+      continue;
+    const int expert_fetch_row = local_row_fetchlist[expert_idx];
+    zipped_probs_topk[this_row * topk + k] =
+        unzipped_token_probs[expert_fetch_row];
+  }
+
+  const int thread_stride = blockDim.x;
+
+  // ------------------------ 手动混合精度 ---------------------------------
+  // 齐整区域向量化搬移
+  for (int x_offset = threadIdx.x; x_offset < token_length;
+       x_offset += thread_stride) {
+    float sum = 0.0f;
+#pragma unroll
+    for (int expert = 0; expert < num_experts; ++expert) {
+      const int fetch_row = local_row_fetchlist[expert];
+      if (fetch_row < 0) continue;
+      // 手动类型提升
+      sum += unzipped_tokens[(int64_t)fetch_row * (int64_t)token_length +
+                             x_offset];
+    }
+    zipped_tokens[(int64_t)this_row * (int64_t)token_length + x_offset] = sum;
+  }
+}
+
+// ---------------------------- Dispatch ---------------------------------
+void dispatch_tokens_unzip(const paddle::Tensor &X,
+                           const paddle::Tensor &expert_routemap_topk,
+                           const paddle::Tensor &expert_prob_topk,
+                           paddle::Tensor &X_unzipped,
+                           paddle::Tensor &zipped_expertwise_rowmap,
+                           paddle::Tensor &token_prob_unzipped,
+                           paddle::Tensor &expert_idx_unzipped,
+                           paddle::Tensor &atomic_extended_offset_counter,
+                           paddle::Tensor &row_valid,
+                           const int total_zipped_tokens_num,
+                           const int total_unzipped_tokens_num,
+                           const int token_length,
+                           const int topk,
+                           const int num_experts) {
+  dim3 grid, block;
+  grid.x = total_unzipped_tokens_num;
+  block.x = 256;
+
+// 定义类型获取宏
+#define DTYPE_CASE(dtype, type) dtype == paddle::DataType::type
+#define GET_DATA(tensor, type) tensor.data<type>()
+
+// 分发处理不同的类型组合
+#define DISPATCH_CASE(TOKEN_T, PROB_T, INT_T, TOPK, NUM_EXPERTS)               \
+  auto kernel = token_unzip_kernel<TOKEN_T, INT_T, PROB_T, TOPK, NUM_EXPERTS>; \
+  kernel<<<grid, block, 0, X.stream()>>>(                                      \
+      GET_DATA(X, TOKEN_T),                                                    \
+      GET_DATA(expert_routemap_topk, INT_T),                                   \
+      GET_DATA(expert_prob_topk, PROB_T),                                      \
+      GET_DATA(X_unzipped, TOKEN_T),                                           \
+      GET_DATA(zipped_expertwise_rowmap, INT_T),                               \
+      GET_DATA(token_prob_unzipped, PROB_T),                                   \
+      expert_idx_unzipped.data<int>(),                                         \
+      atomic_extended_offset_counter.data<int>(),                              \
+      row_valid.data<int>(),                                                   \
+      total_zipped_tokens_num,                                                 \
+      total_unzipped_tokens_num,                                               \
+      token_length);
+
+// 可扩展：处理特定的topk和num_experts组合,可根据之后需求进行扩展
+#define HANDLE_EXPERT_CASE(TOKEN_T, PROB_T, INT_T) \
+  if (topk == 8 && num_experts == 4) {             \
+    DISPATCH_CASE(TOKEN_T, PROB_T, INT_T, 8, 4)    \
+  } else {                                         \
+    /* 超出当前任务范围，*/               \
+    std::__throw_invalid_argument;                 \
+  }
+
+#define HANDLE_TOKEN_TYPE(PROB_T, INT_T)                  \
+  if (DTYPE_CASE(X.dtype(), BFLOAT16)) {                  \
+    HANDLE_EXPERT_CASE(phi::bfloat16, PROB_T, INT_T)      \
+  } else if (DTYPE_CASE(X.dtype(), FLOAT8_E4M3FN)) {      \
+    HANDLE_EXPERT_CASE(phi::float8_e4m3fn, PROB_T, INT_T) \
+  }
+
+#define HANDLE_PROB_TYPE(INT_T)                               \
+  if (DTYPE_CASE(expert_prob_topk.dtype(), BFLOAT16)) {       \
+    HANDLE_TOKEN_TYPE(phi::bfloat16, INT_T)                   \
+  } else if (DTYPE_CASE(expert_prob_topk.dtype(), FLOAT32)) { \
+    HANDLE_TOKEN_TYPE(float, INT_T)                           \
+  }
+
+  // 可扩展：根据整型类型控制派发，未来可支持int8，但int64不行，因为下标开销太重了，建议直接cast到int32
+  if (DTYPE_CASE(zipped_expertwise_rowmap.dtype(), INT32)) {
+    HANDLE_PROB_TYPE(int)
+  }
+
+#undef DTYPE_CASE
+#undef GET_DATA
+#undef DISPATCH_CASE
+#undef HANDLE_EXPERT_CASE
+#undef HANDLE_TOKEN_TYPE
+#undef HANDLE_PROB_TYPE
+}
+
+template <int MAX_NUM_EXPERTS_C>
+void dispatch_tokens_zip(const paddle::Tensor &unzipped_tokens,
+                         const paddle::Tensor &zipped_expertwise_rowmap,
+                         const paddle::Tensor &expert_routemap_topk,
+                         const paddle::Tensor &unzipped_token_probs,
+                         paddle::Tensor &zipped_tokens,
+                         paddle::Tensor &zipped_probs_topk,
+                         const int total_zipped_tokens_num,
+                         const int num_experts,
+                         const int token_length,
+                         const int topk) {
+  dim3 grid, block;
+  grid.x = total_zipped_tokens_num;
+  block.x = 256;
+
+  // Map data types to C++ types
+
+  if (unzipped_tokens.dtype() == paddle::DataType::BFLOAT16) {
+    if (zipped_probs_topk.dtype() == paddle::DataType::FLOAT32) {
+      tokens_zip_kernel<MAX_NUM_EXPERTS_C>
+          <<<grid, block, 0, unzipped_tokens.stream()>>>(
+              unzipped_tokens.data<phi::bfloat16>(),
+              zipped_expertwise_rowmap.data<int>(),
+              expert_routemap_topk.data<int>(),
+              unzipped_token_probs.data<float>(),
+              zipped_tokens.data<phi::bfloat16>(),
+              zipped_probs_topk.data<float>(),
+              total_zipped_tokens_num,
+              token_length,
+              num_experts,
+              topk);
+    } else if (zipped_probs_topk.dtype() == paddle::DataType::BFLOAT16) {
+      tokens_zip_kernel<MAX_NUM_EXPERTS_C>
+          <<<grid, block, 0, unzipped_tokens.stream()>>>(
+              unzipped_tokens.data<phi::bfloat16>(),
+              zipped_expertwise_rowmap.data<int>(),
+              expert_routemap_topk.data<int>(),
+              unzipped_token_probs.data<phi::bfloat16>(),
+              zipped_tokens.data<phi::bfloat16>(),
+              zipped_probs_topk.data<phi::bfloat16>(),
+              total_zipped_tokens_num,
+              token_length,
+              num_experts,
+              topk);
+    }
+  } else if (unzipped_tokens.dtype() == paddle::DataType::FLOAT32) {
+    tokens_zip_kernel<MAX_NUM_EXPERTS_C>
+        <<<grid, block, 0, unzipped_tokens.stream()>>>(
+            unzipped_tokens.data<float>(),
+            zipped_expertwise_rowmap.data<int>(),
+            expert_routemap_topk.data<int>(),
+            unzipped_token_probs.data<float>(),
+            zipped_tokens.data<float>(),
+            zipped_probs_topk.data<float>(),
+            total_zipped_tokens_num,
+            token_length,
+            num_experts,
+            topk);
+  }
+}
+
+void dispatch_tokens_weighted_zip(
+    const paddle::Tensor &unzipped_tokens,
+    const paddle::Tensor &unzipped_token_probs,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    paddle::Tensor &weighted_zipped_tokens,
+    const int total_zipped_tokens_num,
+    const int num_experts,
+    const int token_length) {
+  dim3 grid, block;
+  grid.x = total_zipped_tokens_num;
+  block.x = 256;
+
+  // Map data types to C++ types
+  if (num_experts == 4) {
+    tokens_weighted_zip_kernel<4><<<grid, block, 0, unzipped_tokens.stream()>>>(
+        unzipped_tokens.data<phi::bfloat16>(),
+        unzipped_token_probs.data<phi::bfloat16>(),
+        zipped_expertwise_rowmap.data<int>(),
+        weighted_zipped_tokens.data<phi::bfloat16>(),
+        total_zipped_tokens_num,
+        token_length);
+  }
+}
+
+// -------------------------------- API -----------------------------------
+std::vector<paddle::Tensor> tokens_weighted_zip(
+    const paddle::Tensor &unzipped_tokens,
+    const paddle::Tensor &unzipped_token_probs,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    const int &total_zipped_tokens_num,
+    const int &num_experts) {
+  PD_CHECK(unzipped_tokens.dtype() == paddle::DataType::BFLOAT16);
+  int rows = unzipped_tokens.shape()[0];  // seqlen
+  int cols = unzipped_tokens.shape()[1];  // 一般为7168
+
+  //------------------------ 输出1张量 ------------------------
+  auto weighted_zipped_tokens = paddle::empty({total_zipped_tokens_num, cols},
+                                              unzipped_tokens.dtype(),
+                                              unzipped_tokens.place());
+
+  dispatch_tokens_weighted_zip(unzipped_tokens,
+                               unzipped_token_probs,
+                               zipped_expertwise_rowmap,
+                               weighted_zipped_tokens,
+                               total_zipped_tokens_num,
+                               num_experts,
+                               cols);
+  return {weighted_zipped_tokens};
+}
+
+/*
+PD_BUILD_OP(tokens_zip)
+    .Inputs({"unzipped_tokens", "zipped_expertwise_rowmap",
+"expert_routemap_topk","unzipped_token_probs"}) .Outputs({"zipped_tokens",
+"zipped_probs_topk"}) .Attrs({"total_zipped_tokens: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(tokens_zip));
+*/
+std::vector<paddle::Tensor> tokens_zip(
+    const paddle::Tensor &unzipped_tokens,
+    const paddle::Tensor &zipped_expertwise_rowmap,
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::Tensor &unzipped_token_probs,
+    const int &total_zipped_tokens_num,
+    const int &num_experts) {
+  PD_CHECK(unzipped_tokens.dtype() == paddle::DataType::BFLOAT16 ||
+           unzipped_tokens.dtype() == paddle::DataType::FLOAT32);
+  const int rows = unzipped_tokens.shape()[0];       // seqlen
+  const int cols = unzipped_tokens.shape()[1];       // 一般为7168
+  const int topk = expert_routemap_topk.shape()[1];  // 一般为8
+
+
+  //------------------------ 输出1张量 ------------------------
+  auto zipped_tokens = paddle::empty({total_zipped_tokens_num, cols},
+                                     unzipped_tokens.dtype(),
+                                     unzipped_tokens.place());
+  auto zipped_probs_topk = paddle::empty({total_zipped_tokens_num, topk},
+                                         unzipped_token_probs.dtype(),
+                                         unzipped_token_probs.place());
+  // ----------------------- 0初始化 zipped_probs_topk ------------------
+  if (unzipped_token_probs.dtype() == paddle::DataType::FLOAT32) {
+    void *zipped_probs_topk_ptr =
+        reinterpret_cast<void *>(zipped_probs_topk.data<float>());
+    cudaMemsetAsync(zipped_probs_topk_ptr,
+                    0,
+                    sizeof(float) * total_zipped_tokens_num * topk,
+                    unzipped_token_probs.stream());
+  } else if (unzipped_token_probs.dtype() == paddle::DataType::BFLOAT16) {
+    void *zipped_probs_topk_ptr =
+        reinterpret_cast<void *>(zipped_probs_topk.data<phi::bfloat16>());
+    cudaMemsetAsync(zipped_probs_topk_ptr,
+                    0,
+                    sizeof(phi::bfloat16) * total_zipped_tokens_num * topk,
+                    unzipped_token_probs.stream());
+  }
+  if (rows != 0) {
+    PD_SWITCH_NUM_EXPERTS(num_experts, ([&] {
+                            dispatch_tokens_zip<MAX_NUM_EXPERTS_C>(
+                                unzipped_tokens,
+                                zipped_expertwise_rowmap,
+                                expert_routemap_topk,
+                                unzipped_token_probs,
+                                zipped_tokens,
+                                zipped_probs_topk,
+                                total_zipped_tokens_num,
+                                num_experts,
+                                cols,
+                                topk);
+                          }));
+  }
+  return {zipped_tokens, zipped_probs_topk};
+}
+
+std::vector<paddle::Tensor> tokens_unzip(
+    const paddle::Tensor &X,
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::Tensor &expert_prob_topk,
+    const int &total_unzipped_tokens_num,
+    const int &topk,
+    const int &num_experts) {
+  PD_CHECK(X.dtype() == paddle::DataType::BFLOAT16 ||
+           X.dtype() == paddle::DataType::FLOAT8_E4M3FN);
+  PD_CHECK(expert_prob_topk.dtype() == paddle::DataType::BFLOAT16 ||
+           expert_prob_topk.dtype() == paddle::DataType::FLOAT32);
+  PD_CHECK(expert_routemap_topk.dtype() == paddle::DataType::INT32);
+  int rows = X.shape()[0];  // seqlen
+  int cols = X.shape()[1];  // 一般为7168
+  int original_token_num = rows;
+
+  //------------------------ 输出四张量 ------------------------
+  auto X_unzipped =
+      paddle::empty({total_unzipped_tokens_num, cols}, X.dtype(), X.place());
+  // seqlen x num_experts, 每个token的每个专家(如果被发到)对应的行索引, 未初始化
+  auto zipped_expertwise_rowmap = paddle::empty(
+      {original_token_num, num_experts}, paddle::DataType::INT32, X.place());
+  auto token_prob_unzipped = paddle::empty({total_unzipped_tokens_num},
+                                           expert_prob_topk.dtype(),
+                                           expert_prob_topk.place());
+  auto expert_idx_unzipped = paddle::empty(
+      {total_unzipped_tokens_num}, paddle::DataType::INT32, X.place());
+
+  //------------------------ 辅助二张量 ------------------------
+  // 用于原子记录当前以增广的行数，其上限应为 total_unzipped_tokens_num - rows
+  auto atomic_extended_offset_counter =
+      paddle::zeros({1}, paddle::DataType::INT32, X.place());
+  // 增广行数的合法性向量，用于线程组1唤起
+  int extended_row_num = total_unzipped_tokens_num - rows;
+  auto row_valid =
+      paddle::empty({extended_row_num}, paddle::DataType::INT32, X.place());
+  void *row_valid_gpu = reinterpret_cast<void *>(row_valid.data<int>());
+  cudaMemsetAsync(
+      row_valid_gpu, -1, sizeof(int) * extended_row_num, X.stream());
+
+
+  dispatch_tokens_unzip(X,
+                        expert_routemap_topk,
+                        expert_prob_topk,
+                        X_unzipped,
+                        zipped_expertwise_rowmap,
+                        token_prob_unzipped,
+                        expert_idx_unzipped,
+                        atomic_extended_offset_counter,
+                        row_valid,
+                        rows,
+                        total_unzipped_tokens_num,
+                        cols,
+                        topk,
+                        num_experts);
+  return {X_unzipped,
+          zipped_expertwise_rowmap,
+          token_prob_unzipped,
+          expert_idx_unzipped};
+}
+
+PD_BUILD_OP(tokens_unzip)
+    .Inputs({"X", "expert_routemap_topk", "expert_prob_topk"})
+    .Outputs({"X_unzipped",
+              "zipped_expertwise_rowmap",
+              "token_prob_unzipped",
+              "expert_idx_unzipped"})
+    .Attrs({"total_unzipped_tokens_num: int", "topk: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(tokens_unzip));
+
+PD_BUILD_OP(tokens_zip)
+    .Inputs({"unzipped_tokens",
+             "zipped_expertwise_rowmap",
+             "expert_routemap_topk",
+             "unzipped_token_probs"})
+    .Outputs({"zipped_tokens", "zipped_probs_topk"})
+    .Attrs({"total_zipped_tokens: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(tokens_zip));
+
+PD_BUILD_OP(tokens_weighted_zip)
+    .Inputs({"unzipped_tokens",
+             "unzipped_token_probs",
+             "zipped_expertwise_rowmap"})
+    .Outputs({"weighted_zipped_tokens"})
+    .Attrs({"total_zipped_tokens: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(tokens_weighted_zip));
+
+#undef DISPATCH_CASE
+#undef DISPATCH_TOKEN_TYPE
+#undef DISPATCH_PROB_TYPE

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/topk_to_multihot.cu

@@ -0,0 +1,160 @@
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <limits>
+
+#include "paddle/extension.h"
+#include "paddle/phi/api/all.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+template <bool containsProb, int topk, int expert>
+__global__ void convert_to_multihot_specialized_kernel(
+    const int *__restrict__ routemap_topk,  // 输入矩阵 [seqlen, topk]
+    const float *__restrict__ probs_topk,   // 输入矩阵 [seqlen, topk]
+    int *routemap_multihot,                 // 输出矩阵 [seqlen, expert]
+    float *probs_multihot,                  // 输出矩阵 [seqlen, expert]
+    int seqlen                              // 序列长度
+) {
+  // 每个线程处理一行数据
+  const int row_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (row_idx < seqlen) [[likely]] {
+    // 寄存器加载、存储，消耗2x(topk + expert)个reg
+    int local_routemap_topk[topk];
+    float local_probs_topk[topk];
+    int local_routemap_multihot[expert];
+    float local_probs_multihot[expert];
+    // 合并初始化 thread local
+#pragma unroll
+    for (int i = 0; i < expert; i++) {
+      local_routemap_multihot[i] = 0;
+      if constexpr (containsProb) {
+        local_probs_multihot[i] = 0.0f;
+      }
+    }
+    // local_topk 数据加载
+#pragma unroll
+    for (int i = 0; i < topk; i++) {
+      local_routemap_topk[i] = routemap_topk[row_idx * topk + i];
+      if constexpr (containsProb) {
+        local_probs_topk[i] = probs_topk[row_idx * topk + i];
+      }
+    }
+    // local_multihot 数据填充
+#pragma unroll
+    for (int i = 0; i < topk; i++) {
+      int expert_idx = local_routemap_topk[i];
+      if (expert_idx >= 0 && expert_idx < expert) {
+        local_routemap_multihot[expert_idx] = 1;
+        if constexpr (containsProb) {
+          local_probs_multihot[expert_idx] = local_probs_topk[i];
+        }
+      }
+    }
+
+    // 连续访存结果写入，拆两个for以获取更好的写入带宽
+#pragma unroll
+    for (int i = 0; i < expert; i++) {
+      routemap_multihot[row_idx * expert + i] = local_routemap_multihot[i];
+    }
+
+#pragma unroll
+    for (int i = 0; i < expert; i++) {
+      probs_multihot[row_idx * expert + i] = local_probs_multihot[i];
+    }
+  }
+}
+
+// 通用版本的 kernel dispatch逻辑， 覆盖两种case，可通过牺牲编译时间来扩展
+template <bool containsProb>
+void dispatch_fused_topk_to_multihot(
+    const paddle::Tensor &routemap_topk,
+    const paddle::optional<paddle::Tensor> &probs_topk,
+    paddle::Tensor &routemap_multihot,
+    paddle::Tensor &probs_multihot,
+    int seqlen,
+    int topk,
+    int num_experts) {
+  dim3 block, grid;
+  block.x = 256;
+  grid.x = (seqlen + block.x - 1) / block.x;
+
+#define HANDLE_CASE(tk, exp)                                 \
+  if (topk == tk && num_experts == exp) {                    \
+    if constexpr (containsProb) {                            \
+      convert_to_multihot_specialized_kernel<true, tk, exp>  \
+          <<<grid, block, 0, routemap_topk.stream()>>>(      \
+              routemap_topk.data<int>(),                     \
+              probs_topk.get().data<float>(),                \
+              routemap_multihot.data<int>(),                 \
+              probs_multihot.data<float>(),                  \
+              seqlen);                                       \
+    } else {                                                 \
+      convert_to_multihot_specialized_kernel<false, tk, exp> \
+          <<<grid, block, 0, routemap_topk.stream()>>>(      \
+              routemap_topk.data<int>(),                     \
+              nullptr,                                       \
+              routemap_multihot.data<int>(),                 \
+              nullptr,                                       \
+              seqlen);                                       \
+    }                                                        \
+    return;                                                  \
+  }
+
+  // 处理常见的topk, expert组合
+  HANDLE_CASE(8, 8);
+  HANDLE_CASE(8, 4);
+#undef HANDLE_CASE
+}
+
+std::vector<paddle::Tensor> fused_topk_to_multihot(
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::optional<paddle::Tensor> &expert_probability_topk,
+    const int &seqlen,
+    const int &topk,
+    const int &num_experts) {
+  PD_CHECK(expert_routemap_topk.dtype() == paddle::DataType::INT32);
+  if (expert_probability_topk)
+    PD_CHECK(expert_probability_topk.get().dtype() ==
+             paddle::DataType::FLOAT32);
+
+  paddle::Tensor expert_routemap_multihot =
+      paddle::empty({seqlen, num_experts},
+                    expert_routemap_topk.dtype(),
+                    expert_routemap_topk.place());
+  paddle::Tensor expert_probability_multihot;
+
+  if (expert_probability_topk) {  // 如果包含prob_topk，则为prob_multihot留出空间
+    expert_probability_multihot =
+        paddle::empty({seqlen, num_experts},
+                      expert_probability_topk->dtype(),
+                      expert_probability_topk->place());
+    dispatch_fused_topk_to_multihot<true>(expert_routemap_topk,
+                                          expert_probability_topk,
+                                          expert_routemap_multihot,
+                                          expert_probability_multihot,
+                                          seqlen,
+                                          topk,
+                                          num_experts);
+  } else {
+    dispatch_fused_topk_to_multihot<false>(expert_routemap_topk,
+                                           expert_probability_topk,
+                                           expert_routemap_multihot,
+                                           expert_probability_multihot,
+                                           seqlen,
+                                           topk,
+                                           num_experts);
+  }
+  return {expert_routemap_multihot, expert_probability_multihot};
+}
+
+PD_BUILD_OP(fused_topk_to_multihot)
+    .Inputs({"expert_routemap_topk",
+             paddle::Optional("expert_probability_topk")})
+    .Outputs({"expert_routemap_multihot",
+              paddle::Optional("expert_probability_multihot")})
+    .Attrs({"seqlen: int", "topk: int", "num_experts: int"})
+    .SetKernelFn(PD_KERNEL(fused_topk_to_multihot));

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/topk_to_multihot_grad.cu

@@ -0,0 +1,107 @@
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <limits>
+
+#include "paddle/extension.h"
+#include "paddle/phi/api/all.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+template <int topk, int expert>
+__global__ void convert_to_topk_specialized_kernel(
+    const int *__restrict__ routemap_topk,     // 输入矩阵 [seqlen, topk]
+    const float *__restrict__ probs_multihot,  // 输入矩阵 [seqlen, expert]
+    float *probs_topk,                         // 输出矩阵 [seqlen, topk]
+    int seqlen                                 // 序列长度
+) {
+  // 每个线程处理一行数据
+  const int row_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (row_idx < seqlen) [[likely]] {
+    // 寄存器加载、存储，消耗2xtopk + expert 个reg
+    int local_routemap_topk[topk];
+    float local_probs_multihot[expert];
+    float local_probs_topk[topk];
+    // 合并初始化 thread local
+#pragma unroll
+    for (int i = 0; i < topk; i++) {
+      local_probs_topk[i] = 0.0f;
+      local_routemap_topk[i] = routemap_topk[row_idx * topk + i];
+    }
+    // local_multihot 数据填充
+#pragma unroll
+    for (int i = 0; i < topk; i++) {
+      const int expert_idx = local_routemap_topk[i];
+      if (expert_idx > -1 && expert_idx < expert) {
+        local_probs_topk[i] = local_probs_multihot[expert_idx];
+      }
+    }
+
+    // 写回
+#pragma unroll
+    for (int i = 0; i < topk; i++) {
+      probs_topk[row_idx * topk + i] = local_probs_topk[i];
+    }
+  }
+}
+
+// 通用版本的 kernel dispatch逻辑， 覆盖两种case，可通过牺牲编译时间来扩展
+void dispatch_multihot_prob_backto_topk(const paddle::Tensor &routemap_topk,
+                                        const paddle::Tensor &probs_multihot,
+                                        paddle::Tensor &probs_topk,
+                                        int seqlen,
+                                        int topk,
+                                        int num_experts) {
+  dim3 block, grid;
+  block.x = 256;
+  grid.x = (seqlen + block.x - 1) / block.x;
+
+#define HANDLE_CASE(tk, exp)                          \
+  if (topk == tk && num_experts == exp) {             \
+    convert_to_topk_specialized_kernel<tk, exp>       \
+        <<<grid, block, 0, routemap_topk.stream()>>>( \
+            routemap_topk.data<int>(),                \
+            probs_multihot.data<float>(),             \
+            probs_topk.data<float>(),                 \
+            seqlen);                                  \
+    return;                                           \
+  }
+
+  // 处理常见的topk, expert组合
+  HANDLE_CASE(8, 8);
+  HANDLE_CASE(8, 4);
+#undef HANDLE_CASE
+}
+
+std::vector<paddle::Tensor> multihot_prob_backto_topk(
+    const paddle::Tensor &expert_routemap_topk,
+    const paddle::Tensor &expert_probability_multihot,
+    const int &seqlen,
+    const int &topk,
+    const int &num_experts) {
+  PD_CHECK(expert_routemap_topk.dtype() == paddle::DataType::INT32);
+  PD_CHECK(expert_probability_multihot.dtype() == paddle::DataType::FLOAT32);
+
+  paddle::Tensor expert_probability_topk =
+      paddle::empty({seqlen, topk},
+                    expert_probability_multihot.dtype(),
+                    expert_probability_multihot.place());
+
+  dispatch_multihot_prob_backto_topk(expert_routemap_topk,
+                                     expert_probability_multihot,
+                                     expert_probability_topk,
+                                     seqlen,
+                                     topk,
+                                     num_experts);
+
+  return {expert_probability_topk};
+}
+
+PD_BUILD_OP(fused_multihot_prob_backto_topk)
+    .Inputs({"expert_routemap_topk", "expert_probability_multihot"})
+    .Attrs({"seqlen: int", "topk: int", "num_experts: int"})
+    .Outputs({"expert_probability_topk"})
+    .SetKernelFn(PD_KERNEL(multihot_prob_backto_topk));

slm/model_zoo/gpt-3/external_ops/token_dispatcher_utils/utils.h

@@ -0,0 +1,128 @@
+// 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.
+
+#pragma once
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include <iostream>
+#include <limits>
+
+#include "paddle/extension.h"
+#include "paddle/phi/api/all.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+// 简单的类型提取
+template <paddle::DataType DType>
+struct TypeMap;
+template <>
+struct TypeMap<paddle::DataType::BFLOAT16> {
+  using type = phi::bfloat16;
+};
+template <>
+struct TypeMap<paddle::DataType::FLOAT16> {
+  using type = phi::float16;
+};
+template <>
+struct TypeMap<paddle::DataType::FLOAT32> {
+  using type = float;
+};
+template <>
+struct TypeMap<paddle::DataType::INT32> {
+  using type = int;
+};
+template <>
+struct TypeMap<paddle::DataType::INT64> {
+  using type = int64_t;
+};
+
+
+// cuda结构体拷贝
+template <typename T, int N>
+struct alignas(16) VectorType {
+  T data[N];
+};
+
+// 128Byte对齐的结构体
+template <>
+struct alignas(16) VectorType<float, 4> {
+  float4 data;  // Built-in CUDA vector type
+};
+
+template <>
+struct alignas(16) VectorType<__nv_bfloat16, 8> {
+  __nv_bfloat16 data[8];
+};
+
+template <>
+struct alignas(16) VectorType<__nv_fp8_e4m3, 16> {
+  __nv_fp8_e4m3 data[16];
+};
+
+template <>
+struct alignas(16) VectorType<uint8_t, 16> {
+  uint8_t data[16];
+};
+
+// Helper function to perform vectorized memory copy
+template <typename T>
+__device__ __forceinline__ void vectorized_memcpy(const T* src,
+                                                  T* dst,
+                                                  int num_elements) {
+  constexpr int vector_size_in_bytes = 16;
+  const int elements_per_vector = vector_size_in_bytes / sizeof(T);
+
+  // 已知单行token向量化不会超过4G大小，用int节省整数开销
+  int num_vectors = num_elements / elements_per_vector;
+  int remaining_elements = num_elements % elements_per_vector;
+
+  using VecType = VectorType<T, elements_per_vector>;
+  const VecType* src_vec = reinterpret_cast<const VecType*>(src);
+  VecType* dst_vec = reinterpret_cast<VecType*>(dst);
+
+// 已知paddle框架中的显存分配均为256Bytes对齐，所以默认align
+#pragma unroll
+  for (int idx = threadIdx.x; idx < num_vectors; idx += blockDim.x) {
+    dst_vec[idx] = src_vec[idx];
+  }
+
+  // 剩余无法向量化处理的元素
+  if (remaining_elements > 0) {
+    int offset = num_vectors * elements_per_vector;
+    for (int i = threadIdx.x; i < remaining_elements; i += blockDim.x) {
+      dst[offset + i] = src[offset + i];
+    }
+  }
+}
+
+#define PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, __max_num_experts, ...) \
+  if (__num_expert <= __max_num_experts) {                               \
+    constexpr auto MAX_NUM_EXPERTS_C = __max_num_experts;                \
+    do {                                                                 \
+      __VA_ARGS__();                                                     \
+    } while (0);                                                         \
+    break;                                                               \
+  }
+
+
+#define PD_SWITCH_NUM_EXPERTS(__num_experts_expr, ...)           \
+  do {                                                           \
+    auto __num_expert = (__num_experts_expr);                    \
+    PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 8, __VA_ARGS__);    \
+    PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 16, __VA_ARGS__);   \
+    PD_SWITCH_NUM_EXPERTS_IMPL(__num_expert, 32, __VA_ARGS__);   \
+    PD_THROW("Unsupported expert number %d", int(__num_expert)); \
+  } while (0)

tests/ops/test_merge_subbatch_cast.py

@@ -0,0 +1,36 @@
+# 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 numpy as np
+import paddle
+import TokenDispatcherUtils as TDU
+
+subbatch_rows = 30
+remainder_row = 0
+left_shape = [100, 200]
+
+for num in range(33):
+    for dtype in [paddle.bfloat16, paddle.float32]:
+        x = []
+        for i in range(num - 1):
+            x.append(paddle.randn([subbatch_rows] + left_shape, dtype=paddle.float32))
+        x.append(
+            paddle.randn([subbatch_rows if remainder_row == 0 else remainder_row] + left_shape, dtype=paddle.float32)
+        )
+
+        y1 = paddle.concat(x, axis=0).astype(dtype)
+        y2 = TDU.merge_subbatch_cast(x, dtype)
+
+        diff = np.abs(y1.numpy() - y2.numpy()).max()
+        assert diff == 0, diff

tests/ops/test_tokens_zip_unique_add_subbatch.py

@@ -0,0 +1,89 @@
+# 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 random
+
+import numpy as np
+import paddle
+import TokenDispatcherUtils as TDU
+
+
+def tokens_zip_unique_add_with_subbatch(zipped, unzipped, index_unzipped, zipped_rows, subbatch_rows=None):
+    if subbatch_rows is None or subbatch_rows <= 0 or zipped_rows <= 0:
+        return TDU.tokens_zip_unique_add(zipped, unzipped, index_unzipped, zipped_rows)
+    else:
+        if isinstance(zipped, paddle.Tensor):
+            num_split = (zipped_rows + subbatch_rows - 1) // subbatch_rows
+            remainder = zipped_rows % subbatch_rows
+            if remainder == 0:
+                rows = [subbatch_rows] * num_split
+            else:
+                rows = [subbatch_rows] * (num_split - 1) + [remainder]
+
+            if zipped.shape[0] == 0:
+                dtype = zipped.dtype
+                hidden_size = zipped.shape[1]
+                zipped = [paddle.zeros([r, hidden_size], dtype=dtype) for r in rows]
+            else:
+                zipped = paddle.split(zipped, rows, axis=0)
+        return TDU.tokens_zip_unique_add_subbatch(zipped, unzipped, index_unzipped, zipped_rows, subbatch_rows)
+
+
+def generate_index_unzipped(zipped_rows, unzipped_rows):
+    index = random.sample(range(zipped_rows), unzipped_rows)
+    assert len(index) == len(set(index))
+    return paddle.to_tensor(index, dtype=paddle.int64)
+
+
+def main():
+    seed = 2048
+    hidden_size = 7168
+    zipped_rows = 5800
+    unzipped_rows = 4788
+    subbatch_rows = 380
+    dtype = paddle.bfloat16
+
+    paddle.seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+    zipped_origin = paddle.randn([zipped_rows, hidden_size], dtype=paddle.float32)
+    unzipped = paddle.randn([unzipped_rows, hidden_size], dtype=dtype)
+    index_unzipped = generate_index_unzipped(zipped_rows, unzipped_rows)
+
+    md5sum = None
+    for use_subbatch in [False, True]:
+        random.seed(seed + 100)
+        args = [zipped_origin.clone(), unzipped, index_unzipped, zipped_rows]
+        if use_subbatch and hasattr(TDU, "tokens_zip_unique_add_subbatch"):
+            args.append(subbatch_rows)
+
+        for _ in range(4):
+            output = tokens_zip_unique_add_with_subbatch(*args)
+            args[0] = output
+            args[2] = generate_index_unzipped(zipped_rows, unzipped_rows)
+
+        if isinstance(output, (list, tuple)):
+            output = paddle.concat(output, axis=0)
+
+        cur_md5sum = output._md5sum()
+        if md5sum is None:
+            md5sum = output._md5sum()
+            print(f"MD5SUM: {md5sum}")
+        else:
+            assert md5sum == cur_md5sum, f"{md5sum} vs {cur_md5sum}"
+
+
+if __name__ == "__main__":
+    main()

tools/codestyle/clang_format.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+
+# 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.
+
+set -e
+
+readonly VERSION="13.0.0"
+
+version=$(clang-format -version)
+
+if ! [[ $(python -V 2>&1 | awk '{print $2}' | awk -F '.' '{print $1$2}') -ge 36 ]]; then
+    echo "clang-format installation by pip need python version great equal 3.6,
+          please change the default python to higher version."
+    exit 1
+fi
+
+if ! [[ $version == *"$VERSION"* ]]; then
+    # low version of pip may not have the source of clang-format whl
+    pip install --upgrade pip
+    pip install clang-format==13.0.0
+fi
+
+clang-format $@
+