update using_post_norm_recompute

* add_stepped_rc

* polish code

* remove stepevent

* fix bug

.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)$

llm/config/deepseek-v3/pretrain_argument.json

@@ -1,25 +1,29 @@
 {
-  "model_name_or_path": "deepseek-ai/DeepSeek-V3",
+  "model_name_or_path": "./model_config/DeepSeek-V3",
   "tokenizer_name_or_path": "deepseek-ai/DeepSeek-V3",
   "input_dir": "./data",
   "output_dir": "./checkpoints/pretrain_ckpts",
+  "resume_from_huggingface_ckpt": "./huggingface_ckpt/",
   "per_device_train_batch_size": 1,
-  "gradient_accumulation_steps": 4,
+  "gradient_accumulation_steps": 60,
   "per_device_eval_batch_size": 1,
   "tensor_parallel_degree": 1,
-  "pipeline_parallel_degree": 1,
-  "sharding_parallel_degree": 8,
-  "expert_parallel_degree": 4,
+  "pipeline_parallel_degree": 8,
+  "pipeline_parallel_config": "use_dualpipev enable_overlap_p2p_comm",
+  "sharding_parallel_degree": 32,
+  "sharding_parallel_config": "split_param",
+  "sharding_comm_buffer_size_MB": 2048,
+  "expert_parallel_degree": 32,
   "sharding": "stage1",
   "virtual_pp_degree": 1,
   "sequence_parallel": 0,
   "use_flash_attention": true,
-  "max_seq_length": 4096,
+  "max_seq_length": 4097,
   "learning_rate": 3e-05,
   "min_learning_rate": 3e-06,
   "warmup_steps": 30,
   "logging_steps": 1,
-  "max_steps": 10000,
+  "max_steps": 200,
   "save_steps": 5000,
   "eval_steps": 1000,
   "weight_decay": 0.01,
@@ -27,15 +31,24 @@
   "fp16_opt_level": "O2",
   "warmup_ratio": 0.01,
   "max_grad_norm": 1.0,
-  "dataloader_num_workers": 1,
+  "amp_master_grad": 1,
+  "dataloader_num_workers": 8,
   "continue_training": 0,
   "do_train": true,
   "do_eval": true,
   "do_predict": false,
   "disable_tqdm": true,
-  "recompute": true,
+  "recompute": false,
   "distributed_dataloader": 1,
-  "recompute_granularity": "full",
   "unified_checkpoint": true,
-  "save_total_limit": 2
+  "save_total_limit": 2,
+  "skip_profile_timer": false,
+  "use_fused_rms_norm": true,
+  "fuse_attention_ffn": true,
+  "use_fused_rope": true,
+  "save_sharded_model": false,
+  "load_sharded_model": false,
+  "use_expert_parallel": true,
+  "unified_checkpoint_config": "skip_save_model_weight",
+  "offload_optim": true
 }

llm/model_config/DeepSeek-V3-test/config.json

@@ -0,0 +1,67 @@
+{
+    "architectures": [
+      "DeepseekV3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+      "AutoConfig": "configuration_deepseek.DeepseekV3Config",
+      "AutoModel": "modeling_deepseek.DeepseekV3Model",
+      "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM"
+    },
+    "aux_loss_alpha": 0.001,
+    "bos_token_id": 0,
+    "eos_token_id": 1,
+    "ep_size": 1,
+    "first_k_dense_replace": 0,
+    "hidden_act": "silu",
+    "hidden_size": 7168,
+    "initializer_range": 0.02,
+    "intermediate_size": 18432,
+    "kv_lora_rank": 512,
+    "max_position_embeddings": 163840,
+    "model_type": "deepseek_v3",
+    "moe_intermediate_size": 2048,
+    "moe_layer_freq": 1,
+    "n_group": 8,
+    "n_routed_experts": 256,
+    "n_shared_experts": 1,
+    "norm_topk_prob": true,
+    "num_attention_heads": 128,
+    "num_experts_per_tok": 8,
+    "num_hidden_layers": 2,
+    "num_key_value_heads": 128,
+    "num_nextn_predict_layers": 1,
+    "pretraining_tp": 1,
+    "q_lora_rank": 1536,
+    "qk_nope_head_dim": 128,
+    "qk_rope_head_dim": 64,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "beta_fast": 32,
+      "beta_slow": 1,
+      "factor": 40,
+      "mscale": 1.0,
+      "mscale_all_dim": 1.0,
+      "original_max_position_embeddings": 4096,
+      "type": "yarn"
+    },
+    "rope_theta": 10000,
+    "routed_scaling_factor": 2.5,
+    "scoring_func": "sigmoid",
+    "seq_aux": true,
+    "tie_word_embeddings": false,
+    "topk_group": 4,
+    "topk_method": "noaux_tc",
+    "dtype": "bfloat16",
+    "transformers_version": "4.33.1",
+    "use_cache": true,
+    "v_head_dim": 128,
+    "vocab_size": 129280,
+    "using_flex_token": true,
+    "using_fake_gate": true,
+    "use_fused_rms_norm": true,
+    "fuse_attention_ffn": true,
+    "use_fused_rope": true,
+    "token_drop_steps": true
+  }

llm/model_config/DeepSeek-V3-test/tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 131072,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "LlamaTokenizerFast",
+  "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}{%- for message in messages %}{%- if message['role'] == 'system' %}{%- if ns.is_first_sp %}{% set ns.system_prompt = ns.system_prompt + message['content'] %}{% set ns.is_first_sp = false %}{%- else %}{% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}{%- endif %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{{'<｜Assistant｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}"
+}

llm/model_config/DeepSeek-V3/config.json

@@ -0,0 +1,78 @@
+{
+    "architectures": [
+      "DeepseekV3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+      "AutoConfig": "configuration_deepseek.DeepseekV3Config",
+      "AutoModel": "modeling_deepseek.DeepseekV3Model",
+      "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM"
+    },
+    "aux_loss_alpha": 0.001,
+    "bos_token_id": 0,
+    "eos_token_id": 1,
+    "ep_size": 1,
+    "first_k_dense_replace": 3,
+    "hidden_act": "silu",
+    "hidden_size": 7168,
+    "initializer_range": 0.02,
+    "intermediate_size": 18432,
+    "kv_lora_rank": 512,
+    "max_position_embeddings": 163840,
+    "model_type": "deepseek_v3",
+    "moe_intermediate_size": 2048,
+    "moe_layer_freq": 1,
+    "n_group": 8,
+    "n_routed_experts": 256,
+    "n_shared_experts": 1,
+    "norm_topk_prob": true,
+    "num_attention_heads": 128,
+    "num_experts_per_tok": 8,
+    "num_hidden_layers": 61,
+    "num_key_value_heads": 128,
+    "num_nextn_predict_layers": 1,
+    "pretraining_tp": 1,
+    "q_lora_rank": 1536,
+    "qk_nope_head_dim": 128,
+    "qk_rope_head_dim": 64,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "beta_fast": 32,
+      "beta_slow": 1,
+      "factor": 40,
+      "mscale": 1.0,
+      "mscale_all_dim": 1.0,
+      "original_max_position_embeddings": 4096,
+      "type": "yarn"
+    },
+    "rope_theta": 10000,
+    "routed_scaling_factor": 2.5,
+    "scoring_func": "sigmoid",
+    "seq_aux": true,
+    "tie_word_embeddings": false,
+    "topk_group": 4,
+    "topk_method": "noaux_tc",
+    "dtype": "bfloat16",
+    "transformers_version": "4.33.1",
+    "use_cache": true,
+    "v_head_dim": 128,
+    "vocab_size": 129280,
+    "using_flex_token": true,
+    "using_fake_gate": false,
+    "use_fused_rms_norm": true,
+    "fuse_attention_ffn": true,
+    "use_fused_rope": true,
+    "token_drop_steps": 0,
+    "recompute_fwd_gate_up": true,
+    "adaptive_remained_O1_recompute_ratio": 0,
+    "using_post_norm_recompute": true,
+    "is_split_group_gemm": false,
+    "use_dualpipev": true,
+    "send_mtp_embed": false,
+    "mlp_fwd_subbatch_rows": 0,
+    "mlp_bwd_subbatch_rows": 65536,    
+    "output_subbatch_rows": 2048,
+    "recompute_fa3": true,
+    "stepped_recompute_fwd_gate_up": true
+  }

llm/model_config/DeepSeek-V3/tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 131072,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "LlamaTokenizerFast",
+  "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}{%- for message in messages %}{%- if message['role'] == 'system' %}{%- if ns.is_first_sp %}{% set ns.system_prompt = ns.system_prompt + message['content'] %}{% set ns.is_first_sp = false %}{%- else %}{% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}{%- endif %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{{'<｜Assistant｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}"
+}

llm/run.sh

@@ -0,0 +1,23 @@
+# 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.
+
+# llama 模型数据下载
+# mkdir -p data
+# wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.bin
+# wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k.idx
+
+# mpirun sh script/kill_process.sh 
+# mpirun rm -rf output
+nohup bash script/train_gpu.sh config/deepseek-v3/pretrain_argument.json &
+

llm/run_pretrain.py

@@ -27,7 +27,9 @@ from paddlenlp.data.causal_dataset import (
     print_rank_0,
 )
 from paddlenlp.trainer import (
+    FP8QuantWeightCallback,
     PdArgumentParser,
+    StepFlexToken,
     Trainer,
     TrainingArguments,
     get_last_checkpoint,
@@ -402,6 +404,7 @@ def main():
 
     tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
     config = AutoConfig.from_pretrained(model_args.model_name_or_path)
+
     # set all llm config
     LlmMetaConfig.set_llm_config(config, training_args)
     config.use_fast_layer_norm = model_args.use_fast_layer_norm
@@ -502,6 +505,21 @@ def main():
                 dtype=dtype,
             )
     else:
+        # 修改这里降低模型层数，deepseek前3层为dense层，之后才有稀疏层
+        # config.num_hidden_layers = 4  # v3是61
+        # config.first_k_dense_replace = 0  # v3是3
+        # # 修改这里降低模型专家数量，如果希望进行EP并行，专家数量要能够被并行度整除
+        # config.n_routed_experts = 64  # v3是256
+        # config.num_experts_per_tok = 8  # v3是8
+        # config.topk_group = 4  # v3是4
+
+        # config.using_flex_token = True
+        # config.num_nextn_predict_layers = 1
+        # config.using_fake_gate = True
+        # config.use_fused_rms_norm = True
+        # config.fuse_attention_ffn = True
+        # config.use_fused_rope = True
+        # config.token_drop_steps = 0
         model = model_class.from_config(config, dtype=dtype)
 
     if training_args.recompute:
@@ -551,6 +569,8 @@ def main():
         * data_args.max_seq_length
     )
 
+    callbacks = [StepFlexToken(), FP8QuantWeightCallback()]
+
     trainer = PretrainingTrainer(
         model=model,
         args=training_args,
@@ -559,6 +579,7 @@ def main():
         eval_dataset=eval_dataset if training_args.do_eval else None,
         optimizers=(None, lr_scheduler),
         tokenizer=tokenizer,
+        callbacks=callbacks,
     )
 
     checkpoint = None

llm/script/kill_process.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+# 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.
+
+set -x
+skip_kill_time=${1:-"False"}
+function kill_impl() {
+    skip_kill_time=$1
+    # kill aadiff test finally.
+    pids=`ps -ef | grep pretrain.py | grep -v grep | awk '{print $2}'`
+    if [[ "$pids" != "" ]] ; then
+        echo $pids
+        echo $pids | xargs kill -9
+    fi
+
+    echo "Killing processes on gpu"
+    lsof /dev/nvidia* | awk '{print $2}' | xargs -I {} kill -9 {}
+}
+
+kill_impl $skip_kill_time || true

llm/script/selective_launch.py

@@ -0,0 +1,73 @@
+# 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.
+
+"""
+Selective launch script.
+
+Usage: python script/selective_launch.py <port> <ranks> <ranks> <ranks> ...
+"""
+import os
+import sys
+
+
+def parse_ranks(ranks_strs):
+    """
+    parse_ranks
+    """
+    # NOTE: You can return ranks directly here to change script/train_gpu.sh
+    # and script/kill_process.sh together
+
+    # Example 1: Use contiguous nodes [8, 16)
+    # return range(8, 16)
+
+    # Example 2: Use non-contiguous nodes [4, 8) + {10} + [30, 32), i.e., [4, 5, 6, 7, 10, 30, 31]
+    # return list(range(0, 16)) + list(range(24, 40))
+
+    # Example 3:
+    # Just Python code, return any nodes you want!
+
+    if not ranks_strs:
+        return None
+
+    ranks = []
+    for r in ranks_strs:
+        r = eval(r)
+        if isinstance(r, int):
+            ranks.append(r)
+        else:
+            ranks.extend(r)
+    return ranks
+
+
+def main(port, ranks):
+    """
+    main
+    """
+    ips = [ip.strip() for ip in os.getenv("TRAINER_INSTANCES").split(",") if ip.strip()]
+    if ranks is None:
+        ranks = list(range(len(ips)))
+    ranks = sorted(list(set(ranks)))
+    my_rank = int(os.getenv("POD_INDEX", "0"))
+    if my_rank not in ranks:
+        return
+
+    rank = ranks.index(my_rank)
+    nranks = len(ranks)
+
+    master = ips[ranks[0]]
+    print(f"--master {master}:{port} --rank {rank} --nnodes {nranks}")
+
+
+if __name__ == "__main__":
+    main(int(sys.argv[1]), parse_ranks(sys.argv[2:]))

llm/script/sync.py

@@ -0,0 +1,132 @@
+# -*- coding: utf-8 -*-
+# !/usr/bin/env python3
+
+# 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.
+
+"""
+@version: 1.0
+@file: sync.py
+@time: 2023/07/03 17:22:19
+@Copyright (c) 2023 Baidu.com, Inc. All Rights Reserved
+
+这一行开始写关于本文件的说明与解释
+
+
+"""
+import distutils.util as util
+import multiprocessing
+import os
+import sys
+from shlex import quote
+
+
+def get_self_ip():
+    """
+    get_self_ip
+    """
+    # ip = subprocess.check_output(['hostname', '-i']).decode().strip()
+    ip = os.getenv("POD_IP")
+    return ip
+
+
+def get_ips():
+    """
+    get_ips
+    """
+    path = os.path.join(os.path.dirname(__file__), "iplist.txt")
+    if not os.path.exists(path):
+        path = "/root/paddlejob/workspace/hostfile"
+        assert os.path.exists(path), "iplist file not exists!"
+
+    with open(path, "r") as f:
+        return [ip.strip().split(" ")[0] for ip in f.readlines() if ip.strip() and not ip.strip().startswith("#")]
+
+
+def concat_cmd(cmd):
+    """
+    concat_cmd
+    """
+    return " ".join([quote(c) for c in cmd])
+
+
+def sync_file(path):
+    """
+    sync_file
+    """
+    if isinstance(path, (list, tuple)):
+        for p in path:
+            sync_file(p)
+        return
+
+    path = os.path.abspath(path)
+    # dirname = os.path.dirname(path)
+
+    self_ip = get_self_ip()
+    ips = get_ips()
+    all_cmds = []
+    for ip in ips:
+        if ip != self_ip:
+            # cmd = "mkdir -p {0}; scp -r {2} {1}:{2}".format(dirname, ip, path)
+            # ssh_cmd = concat_cmd(['ssh', ip, cmd])
+            ssh_cmd = "scp -r {1} {0}:{1}".format(ip, path)
+            all_cmds.append(ssh_cmd)
+    if len(all_cmds) > 0:
+        with multiprocessing.Pool(min(len(all_cmds), 20)) as pool:
+            pool.map(execute_cmd, all_cmds)
+
+
+def execute_cmd(cmd):
+    """
+    execute_cmd
+    """
+    print(cmd)
+    # assert os.system(cmd) == 0
+    return os.system(cmd)
+
+
+def sync_cmd(cmd):
+    """
+    sync_cmd
+    """
+    self_ip = get_self_ip()
+    ips = get_ips()
+    joined_cmd = "cd {0}; {1}".format(os.path.abspath("."), concat_cmd(cmd))
+
+    all_cmds = []
+    for ip in ips:
+        is_self = ip == self_ip
+        if not is_self:
+            exe_cmd = concat_cmd(["ssh", ip, joined_cmd])
+        else:
+            exe_cmd = joined_cmd
+
+        if is_self and not util.strtobool(os.getenv("HAS_SELF", "1")):
+            continue
+
+        all_cmds.append(exe_cmd)
+
+    with multiprocessing.Pool(len(all_cmds)) as pool:
+        pool.map(execute_cmd, all_cmds)
+
+
+if __name__ == "__main__":
+    mode = sys.argv[1]
+    assert mode in ["file", "cmd"]
+    if mode == "file":
+        sync_file(sys.argv[2:])
+    elif mode == "cmd":
+        sync_cmd(sys.argv[2:])
+    else:
+        raise ValueError("Invalid mode: {}".format(mode))

llm/script/train_gpu.sh

@@ -0,0 +1,75 @@
+#!/bin/bash
+
+# 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.
+
+unset PADDLE_ELASTIC_JOB_ID
+unset PADDLE_TRAINER_ENDPOINTS
+unset DISTRIBUTED_TRAINER_ENDPOINTS
+unset FLAGS_START_PORT
+unset PADDLE_ELASTIC_TIMEOUT
+
+nnodes=$PADDLE_TRAINERS_NUM
+rank=$PADDLE_TRAINER_ID
+
+for name in `env | grep -E 'PADDLE|ENDPOINT' | awk -F'=' '{print $1}'`; do
+  unset ${name}
+done
+
+#export FLAGS_shard_bypass_dygraph_optimizer=1
+export NCCL_IB_GID_INDEX=3
+export NVSHMEM_IB_GID_INDEX=3
+export NVSHMEM_IB_TRAFFIC_CLASS=162
+
+#export NVSHMEM_IB_ENABLE_IBGDA=true
+##export NVSHMEM_DISABLE_P2P=1
+export NVSHMEM_BOOTSTRAP=UID
+
+unset NVSHMEM_HCA_LIST 
+unset NVSHMEM_ENABLE_NIC_PE_MAPPING
+
+LAUNCH_CMD=`python script/selective_launch.py 36677`
+if [[ -z "$LAUNCH_CMD" ]]; then
+    exit 0
+fi
+
+export PYTHONPATH=../:$PYTHONPATH
+export CUDA_PATH=/usr/local/cuda-12.9
+
+# Flags for best performance
+export DSV3_USE_FP8_GEMM=true
+export DSV3_USE_ATTEN_RECOMPUTE=true
+export FA_VERSION=3
+export FLAGS_share_tensor_for_grad_tensor_holder=1
+export FLAGS_use_default_stream=false
+export DSV3_USE_FP8_DISPATCH=true
+export USE_DS_GEMM=true
+
+# Flags for allocator
+export FLAGS_large_pool_auto_growth_chunk_size_in_mb=500
+export FLAGS_small_pool_auto_growth_chunk_size_in_mb=20
+export FLAGS_small_pool_size_in_mb=10
+export FLAGS_samll_pool_pre_alloc_in_mb=500
+export FLAGS_large_pool_pre_alloc_in_mb=61440
+export FLAGS_deep_ep_comm_prealloc_in_mb=1000
+
+
+bash script/kill_process.sh 
+
+python3.10 -m paddle.distributed.launch \
+    --log_dir output/paddle_distributed_logs \
+    $LAUNCH_CMD \
+    --run_mode=collective \
+    ${script:-run_pretrain.py}  \
+    $@

ops/csrc/fp8/deep_gemm/__init__.py

@@ -23,8 +23,10 @@ from .jit_kernels import (
     get_col_major_tma_aligned_tensor,
     get_m_alignment_for_contiguous_layout,
     get_num_sms,
+    k_grouped_wgrad_gemm_fp8_fp8_fp32_nt,
     m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
     m_grouped_gemm_fp8_fp8_bf16_nt_masked,
     set_num_sms,
+    wgrad_gemm_fp8_fp8_fp32_nt,
 )
-from .utils import bench, calc_diff, get_cuda_home
+from .utils import calc_diff

ops/csrc/fp8/deep_gemm/include/deep_gemm/fp8_gemm.cuh

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// The file has been adapted from DeepSeek DeepEP project
+// The file has been adapted from DeepSeek DeepGEMM project
 // Copyright (c) 2025 DeepSeek
-// Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wunknown-attributes"
@@ -34,25 +34,27 @@
 
 namespace deep_gemm {
 
-enum class Layout {
-    RowMajor,
-    ColMajor
-};
+template <uint32_t kNumFormerIters, uint32_t kGap, uint32_t kEnd>
+__device__ __host__ void outer_launch_k_iterations(const auto& inner_launch_k_iterations, const auto& func, uint32_t num_former_iters) {
+    if (num_former_iters == kNumFormerIters) {
+        inner_launch_k_iterations(func, cute::Int<kNumFormerIters>{});
+        return;
+    }
 
-template <uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup>
-__device__ __host__ constexpr int get_num_threads_per_sm(int block_m) {
-    DG_STATIC_ASSERT(kNumMathThreadsPerGroup == 128, "Only support 128 threads per math group");
-    return (block_m == 64 ? 1 : 2) * kNumMathThreadsPerGroup + kNumTMAThreads;
+    if constexpr (kNumFormerIters + kGap <= kEnd)
+        outer_launch_k_iterations<kNumFormerIters + kGap, kGap, kEnd>(inner_launch_k_iterations, func, num_former_iters);
 }
 
 template <uint32_t SHAPE_N, uint32_t SHAPE_K,
           uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
+          uint32_t BLOCK_N_PADDING,
+          uint32_t kSwizzleDMode,
           uint32_t kNumGroups, uint32_t kNumStages,
           uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup,
-          uint32_t kNumTMAMulticast,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           GemmType kGemmType>
 __global__ void __launch_bounds__(get_num_threads_per_sm<kNumTMAThreads, kNumMathThreadsPerGroup>(BLOCK_M), 1)
-fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
+fp8_gemm_kernel(float* scales_b, int* grouped_layout,
                 uint32_t shape_m,
                 const __grid_constant__ CUtensorMap tensor_map_a,
                 const __grid_constant__ CUtensorMap tensor_map_b,
@@ -61,15 +63,16 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
     // Scaling checks
     DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
-    DG_STATIC_ASSERT(ceil_div(BLOCK_N, BLOCK_K) == 1, "Too much B scales in a single block");
+    DG_STATIC_ASSERT(ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
 
     // Types
     using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
     using Barrier = cutlass::arch::ClusterTransactionBarrier;
+    DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
 
     // Shared memory
-    static constexpr int kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
-    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
+    static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * (BLOCK_N + BLOCK_N_PADDING) * sizeof(__nv_bfloat16);
     static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
     static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
     static constexpr uint32_t SMEM_SCALES_A_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
@@ -84,12 +87,17 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
     const uint32_t lane_idx = get_lane_id();
 
-    // Prefetch TMA descriptors at very beginning
+    // Prefetch TMA descriptors at the very beginning
     if (threadIdx.x == kNumMathThreads) {
-        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_a));
-        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_b));
-        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_scales_a));
-        cute::prefetch_tma_descriptor(reinterpret_cast<cute::TmaDescriptor const*>(&tensor_map_d));
+        // NOTES: `reinterpret_cast` must be here, or NVRTC will fail
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_a));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_b));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_scales_a));
+
+        // `tensor_map_d` is only used in swizzling mode
+        // For the `kSwizzleDMode == 0 and BLOCK_N_PADDING == 0` case, it will be treated as padding mode
+        if constexpr (kSwizzleDMode > 0)
+            cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_d));
     }
     __syncwarp();
 
@@ -110,7 +118,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 
     // Fill shared memory pointers
     #pragma unroll
-    for (int i = 0; i < kNumStages; ++ i) {
+    for (uint32_t i = 0; i < kNumStages; ++ i) {
         smem_a[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
         smem_b[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
         smem_scales_a[i] = reinterpret_cast<float*>(smem_buffer + SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) + i * SMEM_SCALES_A_SIZE_PER_STAGE);
@@ -120,7 +128,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     // Fill barriers
     auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_scales_b) + SMEM_SCALES_B_SIZE);
     #pragma unroll
-    for (int i = 0; i < kNumStages; ++ i) {
+    for (uint32_t i = 0; i < kNumStages; ++ i) {
         full_barriers[i] = barrier_start_ptr + i;
         empty_barriers[i] = barrier_start_ptr + kNumStages + i;
     }
@@ -128,8 +136,10 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     // Initialize barriers
     DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
     if (threadIdx.x == kNumMathThreads) {
+        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
+        // even with TMA multicast disabled, we want to make the behavior aligned
         #pragma unroll
-        for (int i = 0; i < kNumStages; ++ i) {
+        for (uint32_t i = 0; i < kNumStages; ++ i) {
             full_barriers[i]->init(1);
             empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
         }
@@ -145,24 +155,37 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
     // For pipeline unrolling
     struct DivisibleK {};
     struct NotDivisibleK {};
-    auto launch_k_iterations = [](const auto& func) {
-        if constexpr (SHAPE_K % kFullKOfAllStages == 0) {
-            for (int k_iter = 0; k_iter < kNumIterations; ++ k_iter)
-                func(k_iter, DivisibleK{});
-        } else {
-            for (int k_iter = 0; k_iter < kNumIterations - 1; ++ k_iter)
-                func(k_iter, DivisibleK{});
-            func(kNumIterations - 1, NotDivisibleK{});
-        }
+    struct SkipComputation {};
+    struct NotSkipComputation {};
+    auto launch_k_iterations = [](const auto& func, bool skip_computation, uint32_t num_former_iters) {
+        constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
+        constexpr uint32_t kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
+        constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
+
+        // NOTES: for too-many branches (> 5), we disable this optimization
+        // Otherwise, the compiler must know the dynamic variable `num_former_iters`'s real value
+        outer_launch_k_iterations<0, kGap, kEnd>([=](const auto& func, auto num_former_iters_type) {
+            if (skip_computation) {
+                for (uint32_t k_iter = 0; k_iter < kNumIterations; ++ k_iter)
+                    func(k_iter, DivisibleK{}, SkipComputation{}, num_former_iters_type);
+            } else if (SHAPE_K % kFullKOfAllStages == 0) {
+                for (uint32_t k_iter = 0; k_iter < kNumIterations; ++ k_iter)
+                    func(k_iter, DivisibleK{}, NotSkipComputation{}, num_former_iters_type);
+            } else {
+                for (uint32_t k_iter = 0; k_iter < kNumIterations - 1; ++ k_iter)
+                    func(k_iter, DivisibleK{}, NotSkipComputation{}, num_former_iters_type);
+                func(kNumIterations - 1, NotDivisibleK{}, NotSkipComputation{}, num_former_iters_type);
+            }
+        }, func, kShouldOptimize ? num_former_iters : 0);
     };
 
     // Register reconfigurations
-    constexpr int kNumTMARegisters = 40;
-    constexpr int kNumMathRegisters = 232;
+    constexpr uint32_t kNumTMARegisters = 40;
+    constexpr uint32_t kNumMathRegisters = 232;
 
     // Block scheduler
     uint32_t m_block_idx, n_block_idx;
-    auto scheduler = Scheduler<kGemmType, SHAPE_N, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast>(shape_m, grouped_layout);
+    auto scheduler = Scheduler<kGemmType, SHAPE_N, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA>(shape_m, grouped_layout);
 
     if (threadIdx.x >= kNumMathThreads) {
         // TMA warp-group for loading data
@@ -172,28 +195,39 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
         if (threadIdx.x == kNumMathThreads) {
             // Persistently schedule over blocks
             while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
-                launch_k_iterations([&](int k_iter, auto type) {
-                    constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
-                    constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
-                    DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
+                launch_k_iterations([&](uint32_t k_iter, auto divisible_type, auto _, auto __) {
+                    constexpr bool kHasDivisibleStages = std::is_same_v<decltype(divisible_type), DivisibleK>;
+                    constexpr uint32_t kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
 
+                    // Assign TMA multicast number into A and B
+                    // NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
+                    const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
+                    const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                    const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                    DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
+
+                    // NOTES: unrolling and `kNumInnerStages` are vital for performance, NVCC will try to eliminate all
+                    // shared memory pointers, e.g. `full_barriers` registers, if all the access indices are constant
                     #pragma unroll
                     for (uint32_t s = 0; s < kNumInnerStages; ++ s) {
                         // Wait consumer release
                         empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
 
-                        // Issue TMA A with broadcasting
+                        // Issue TMA A
                         auto& full_barrier = *full_barriers[s];
-                        int k_idx = k_iter * kFullKOfAllStages + s * BLOCK_K;
-                        tma_copy<kNumTMAMulticast>(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
-                                                   smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
-                        tma_copy<kNumTMAMulticast>(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
-                                                   smem_scales_a[s], m_block_idx * BLOCK_M,
-                                                   scheduler.get_global_idx(SHAPE_K_SCALES, 1, k_idx / BLOCK_K));
+                        uint32_t k_idx = k_iter * kFullKOfAllStages + s * BLOCK_K;
+                        tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx),
+                                 num_tma_multicast_a);
+                        tma_copy(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_scales_a[s], m_block_idx * BLOCK_M,
+                                 scheduler.get_global_idx(SHAPE_K_SCALES, 1, k_idx / BLOCK_K),
+                                 num_tma_multicast_a);
 
-                        // Issue TMA B without broadcasting
+                        // Issue TMA B
                         tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
-                                 smem_b[s], k_idx, scheduler.get_global_idx<false>(SHAPE_N, BLOCK_N, n_block_idx, m_block_idx));
+                                 smem_b[s], k_idx, scheduler.get_global_idx<false>(SHAPE_N, BLOCK_N, n_block_idx, m_block_idx),
+                                 num_tma_multicast_b);
                         full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE);
                     }
 
@@ -203,7 +237,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                         empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);
                         full_barriers[s]->arrive();
                     }
-                });
+                }, false, 0);
             }
 
             // To safely deconstruct distributed shared barriers, we need another round of empty waits
@@ -244,25 +278,29 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Accumulation for WGMMA or CUDA promotion
-            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
+            constexpr uint32_t WAVE_BLOCK_M = WGMMA::M * get_num_math_warpgroups(BLOCK_M);
+            DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
+            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
 
             // Empty barrier arrival
-            auto empty_barrier_arrive = [&](int s) {
+            auto empty_barrier_arrive = [&](uint32_t s) {
                 if constexpr (kNumTMAMulticast == 1) {
                     lane_idx == 0 ? empty_barriers[s]->arrive() : void();
                 } else {
-                    lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(lane_idx) : void();
+                    auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
+                    lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta) : void();
                 }
             };
 
             // Launch MMAs
-            launch_k_iterations([&](int k_iter, auto type) {
-                constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
-                constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K;
-                DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
+            launch_k_iterations([&](uint32_t k_iter, auto divisible_type, auto skip_type, auto _) {
+                constexpr bool kSkipComputation = std::is_same_v<decltype(skip_type), SkipComputation>;
+                constexpr bool kHasDivisibleStages = std::is_same_v<decltype(divisible_type), DivisibleK>;
+                constexpr uint32_t kNumInnerStages = kSkipComputation ? 0 :
+                    (kHasDivisibleStages ? kNumStages : (SHAPE_K % kFullKOfAllStages) / BLOCK_K);
 
                 #pragma unroll
-                for (int s = 0; s < kNumInnerStages; ++ s) {
+                for (uint32_t s = 0; s < kNumInnerStages; ++ s) {
                     // Read B scales
                     float scale_b_0 = ld_shared(smem_scales_b + k_iter * kNumStages + s), scale_b_1;
                     // NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks
@@ -272,42 +310,54 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                     // Wait TMA arrivals
                     full_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter) & 1);
 
-                    // Read A scales
-                    // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
-                    auto scale_a_0 = ld_shared(smem_scales_a[s] + r_0), scale_a_1 = ld_shared(smem_scales_a[s] + r_1);
+                    // TODO: remove some useless computation for unaligned Ms
+                    #pragma unroll
+                    for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                      	auto m_offset = local_idx * WAVE_BLOCK_M;
 
-                    // Commit WGMMA instructions
-                    #pragma unroll
-                    for (int i = 0; i < WGMMA::kNumAccum; ++ i)
-                        warpgroup_fence_operand(accum[i]);
-                    warpgroup_arrive();
-                    #pragma unroll
-                    for (int k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
-                        auto desc_a = make_smem_desc(smem_a[s] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
-                        auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
-                        WGMMA::wgmma(desc_a, desc_b, accum, k);
-                    }
-                    warpgroup_commit_batch();
-                    #pragma unroll
-                    for (int i = 0; i < WGMMA::kNumAccum; ++ i)
-                        warpgroup_fence_operand(accum[i]);
-                    warpgroup_wait<0>();
+                    	// Read A scales
+                    	// NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
+                    	auto scale_a_0 = ld_shared(smem_scales_a[s] + r_0 + m_offset);
+                        auto scale_a_1 = ld_shared(smem_scales_a[s] + r_1 + m_offset);
 
-                    // Notify barrier arrival
-                    empty_barrier_arrive(s);
+                    	// Commit WGMMA instructions
+                    	#pragma unroll
+                    	for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                            warpgroup_fence_operand(accum[i]);
+                    	warpgroup_arrive();
+                    	#pragma unroll
+                    	for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
+                            auto desc_a = make_smem_desc(smem_a[s] + (math_wg_idx * WGMMA::M + m_offset) * BLOCK_K + k * WGMMA::K, 1);
+                            auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
+                            WGMMA::wgmma(desc_a, desc_b, accum, k);
+                    	}
+                    	warpgroup_commit_batch();
+                    	#pragma unroll
+                    	for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
+                            warpgroup_fence_operand(accum[i]);
+                    	warpgroup_wait<0>();
 
-                    // Promote with scales
-                    float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
-                    float scale_0_1, scale_1_1;
-                    if constexpr (not kMustUseUniformedScaleB)
-                        scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
-                    #pragma unroll
-                    for (int i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
-                        bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
-                        final_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
-                        final_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
-                        final_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
-                        final_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
+                    	// Notify barrier arrival at the last warpgroup wave
+                        if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
+                    	    empty_barrier_arrive(s);
+
+                    	// Promote with scales
+                    	// NOTES: making it as predicates is very important for performance, comparing to two loops
+                    	float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
+                    	float scale_0_1, scale_1_1;
+                    	if constexpr (not kMustUseUniformedScaleB)
+                            scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
+
+                        auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                    	#pragma unroll
+                    	for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                            // NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
+                            bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
+                            shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
+                            shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
+                            shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
+                            shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
+                    	}
                     }
                 }
 
@@ -317,36 +367,85 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
                     full_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter) & 1);
                     empty_barrier_arrive(s);
                 }
-            });
+            }, not scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M), num_former_iters);
 
-            // Write back to shared memory using STSM
+            // TMA checks
+            constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
+            constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
+            constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
+            DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
+            DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
+                            "Unaligned TMA store or too many TMA store instructions");
+            DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
+            DG_STATIC_ASSERT(static_cast<uint32_t>(kSwizzleDMode > 0) + static_cast<uint32_t>(BLOCK_N_PADDING > 0) <= 1,
+                            "Swizzling and padding are not compatible");
+
+            // Wait last TMA store to be finished
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N)
+                cute::tma_store_wait<0>();
+            cutlass::arch::NamedBarrier(kNumMathThreads).sync();
+
+            // Write back to shared memory using STSM and issue TMA stores
             DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
             #pragma unroll
-            for (auto i = 0; i < WGMMA::kNumAccum / 8; ++ i) {
-                SM90_U32x4_STSM_N<nv_bfloat162>::copy(
-                    __float22bfloat162_rn({final_accum[i * 8 + 0], final_accum[i * 8 + 1]}),
-                    __float22bfloat162_rn({final_accum[i * 8 + 2], final_accum[i * 8 + 3]}),
-                    __float22bfloat162_rn({final_accum[i * 8 + 4], final_accum[i * 8 + 5]}),
-                    __float22bfloat162_rn({final_accum[i * 8 + 6], final_accum[i * 8 + 7]}),
-                    smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + i * 16 + 8 * (lane_idx / 16)
-                );
-            }
-            if constexpr (WGMMA::kNumAccum % 8 != 0) {
-                SM90_U32x2_STSM_N<nv_bfloat162>::copy(
-                    __float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 0], final_accum[WGMMA::kNumAccum / 8 * 8 + 1]}),
-                    __float22bfloat162_rn({final_accum[WGMMA::kNumAccum / 8 * 8 + 2], final_accum[WGMMA::kNumAccum / 8 * 8 + 3]}),
-                    smem_d + (warp_idx * 16 + lane_idx % 16) * BLOCK_N + WGMMA::kNumAccum / 8 * 16
-                );
+            for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                auto m_offset = local_idx * WAVE_BLOCK_M;
+                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                #pragma unroll
+                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                    // Swizzle or padding into the correct address
+                    uint8_t* smem_ptr = nullptr;
+                    if constexpr (kSwizzleDMode > 0) {
+                        // Calculate the swizzling atom offset and in-atom offset
+                        constexpr uint32_t kNumBankGroupBytes = 16;
+                        auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
+
+                        // Calculate the index of the bank group to be written in the atom
+                        auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
+
+                        // Reshape the atom in another view and swizzle
+                        //  - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
+                        //  - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
+                        constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
+                        auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
+                        auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
+                        col ^= row % (kSwizzleDMode / 16);
+
+                        // Add back into the base pointer
+                        // NOTES: think twice before modifying this, as changes may affect the number of instructions
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d) +                // Base pointer
+                            warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) +            // Warp offset
+                            m_offset * kSwizzleDMode +                                 // Wave offset
+                            atom_offset * BLOCK_M * kSwizzleDMode +                    // Swizzle atom offset (constants)
+                            row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
+                    } else {
+                        // No swizzling, just padding
+                        // NOTES: padding must be zero for BF16 output
+                        DG_STATIC_ASSERT(BLOCK_N_PADDING == 0, "Padding must be zero for BF16 output");
+                        smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * (BLOCK_N + BLOCK_N_PADDING) + i * 8);
+                    }
+
+                    // NOTES: only 16 lanes' addresses are used
+                    SM90_U32x2_STSM_N<nv_bfloat162>::copy(
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
+                        __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
+                        smem_ptr
+                    );
+                }
             }
             cute::tma_store_fence();
             cutlass::arch::NamedBarrier(kNumMathThreads).sync();
 
             // Use TMA store to write back to global memory
-            if (threadIdx.x == 0) {
-                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_d, n_block_idx * BLOCK_N,
+            // TODO: compatible with FP32 output
+            DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
+            if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
+                auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
+                auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
+                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
+                                              n_block_idx * BLOCK_N + in_block_n_offset,
                                               scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx));
                 cute::tma_store_arrive();
-                cute::tma_store_wait<0>();
             }
             __syncwarp();
         }
@@ -357,106 +456,6 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
 #endif
 }
 
-template <uint32_t SHAPE_N, uint32_t SHAPE_K,
-          uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
-          uint32_t kNumGroups, uint32_t kNumStages,
-          uint32_t kNumTMAMulticast,
-          GemmType kGemmType>
-class Gemm {
-private:
-    using Barrier = cuda::barrier<cuda::thread_scope_block>;
-
-public:
-    Gemm() = default;
-
-    static void run(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
-                    uint32_t shape_m,
-                    const CUtensorMap& tma_a_desc,
-                    const CUtensorMap& tma_b_desc,
-                    const CUtensorMap& tma_scales_a_desc,
-                    const CUtensorMap& tma_d_desc,
-                    cudaStream_t stream,
-                    int num_sms, uint32_t smem_size) {
-        // NOTES: we must use 4 warps to do TMA, because `setmaxnreg.aligned` requires 4 warps
-        constexpr uint32_t kNumTMAThreads = 128;
-        constexpr uint32_t kNumMathThreadsPerGroup = 128;
-        auto kernel = fp8_gemm_kernel<SHAPE_N, SHAPE_K, BLOCK_M, BLOCK_N, BLOCK_K,
-                                      kNumGroups, kNumStages, kNumTMAThreads, kNumMathThreadsPerGroup,
-                                      kNumTMAMulticast, kGemmType>;
-        DG_HOST_ASSERT(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) == cudaSuccess);
-
-        // Cluster launch
-        cudaLaunchConfig_t config;
-        config.gridDim = num_sms;
-        config.blockDim = get_num_threads_per_sm<kNumTMAThreads, kNumMathThreadsPerGroup>(BLOCK_M);
-        config.dynamicSmemBytes = smem_size;
-        config.stream = stream;
-
-        // Clusters for TMA multicast
-        // NOTES: `>= 4` cluster size will cause performance degradation
-        cudaLaunchAttribute attr;
-        attr.id = cudaLaunchAttributeClusterDimension;
-        attr.val.clusterDim = {kNumTMAMulticast, 1, 1};
-        config.attrs = &attr;
-        config.numAttrs = 1;
-
-        // Launch
-        auto status = cudaLaunchKernelEx(&config, kernel,
-                                         gmem_d, scales_b, grouped_layout,
-                                         shape_m,
-                                         tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc);
-        DG_HOST_ASSERT(status == cudaSuccess);
-    }
-
-    template <typename T>
-    static CUtensorMap make_2d_tma_a_desc(T* global_address, uint32_t shape_m) {
-        return make_2d_tma_desc(global_address, Layout::RowMajor,
-                                shape_m * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), SHAPE_K, BLOCK_M, BLOCK_K);
-    }
-
-    template <typename T>
-    static CUtensorMap make_2d_tma_b_desc(T* global_address) {
-        return make_2d_tma_desc(global_address, Layout::ColMajor,
-                                SHAPE_K, SHAPE_N * (kGemmType != GemmType::Normal ? kNumGroups : 1), BLOCK_K, BLOCK_N);
-    }
-
-    template <typename T>
-    static CUtensorMap make_2d_tma_d_desc(T* global_address, uint32_t shape_m) {
-        return make_2d_tma_desc(global_address, Layout::RowMajor,
-                                shape_m * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), SHAPE_N,
-                                min(BLOCK_M, shape_m), BLOCK_N,
-                                CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
-    }
-
-    template <typename T>
-    static CUtensorMap make_2d_tma_scales_a_desc(T* global_address, uint32_t shape_m) {
-        // Make TMA aligned to 16 bytes
-        constexpr uint32_t kAlignment = 16 / sizeof(T);
-        shape_m = ceil_div(shape_m, kAlignment) * kAlignment;
-
-        return make_2d_tma_desc(global_address, Layout::ColMajor,
-                                shape_m, ceil_div(SHAPE_K, BLOCK_K) * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), BLOCK_M, 1,
-                                CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
-    }
-
-    template <typename T>
-    static CUtensorMap make_2d_tma_desc(
-            T* global_address, Layout layout,
-            uint32_t gmem_rows, uint32_t gmem_cols,
-            uint32_t smem_rows, uint32_t smem_cols,
-            CUtensorMapSwizzle swizzle_type = CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B) {
-        if (layout == Layout::RowMajor) {
-            uint64_t gmem_dim[2] = {gmem_cols, gmem_rows};
-            uint32_t smem_dim[2] = {smem_cols, smem_rows};
-            return make_2d_tma_copy_desc(global_address, gmem_dim, gmem_cols * sizeof(T), smem_dim, swizzle_type);
-        } else {
-            uint64_t gmem_dim[2] = {gmem_rows, gmem_cols};
-            uint32_t smem_dim[2] = {smem_rows, smem_cols};
-            return make_2d_tma_copy_desc(global_address, gmem_dim, gmem_rows * sizeof(T), smem_dim, swizzle_type);
-        }
-    }
-};
-
 };  // namespace deep_gemm
 
 #pragma clang diagnostic pop

ops/csrc/fp8/deep_gemm/include/deep_gemm/fp8_wgrad_gemm.cuh

@@ -0,0 +1,381 @@
+// 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.
+
+// The file has been adapted from DeepSeek DeepGEMM project
+// Copyright (c) 2025 DeepSeek
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+#pragma once
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wunknown-attributes"
+
+#include <cutlass/arch/barrier.h>
+#include <cutlass/arch/reg_reconfig.h>
+
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/arch/copy_sm90_desc.hpp>
+#include <cute/arch/copy_sm90_tma.hpp>
+
+#include "mma_utils.cuh"
+#include "scheduler.cuh"
+#include "tma_utils.cuh"
+#include "utils.cuh"
+
+namespace deep_gemm {
+
+template <uint32_t SHAPE_M, uint32_t SHAPE_N,
+          uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
+          uint32_t kNumStages, uint32_t kNumLastStages,
+          uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA>
+__global__ void __launch_bounds__(get_num_threads_per_sm<kNumTMAThreads, kNumMathThreadsPerGroup>(BLOCK_M), 1)
+fp8_wgrad_gemm_kernel(uint32_t shape_k, 
+                      const __grid_constant__ CUtensorMap tensor_map_a,
+                      const __grid_constant__ CUtensorMap tensor_map_b,
+                      const __grid_constant__ CUtensorMap tensor_map_scales_a,
+                      const __grid_constant__ CUtensorMap tensor_map_scales_b,
+                      const __grid_constant__ CUtensorMap tensor_map_d) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) || defined(__CLION_IDE__)
+    // Scaling checks
+    DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
+
+    // Types
+    using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
+    using Barrier = cutlass::arch::ClusterTransactionBarrier;
+    DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
+
+    // Shared memory
+    static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(float);
+    static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
+    static constexpr uint32_t SMEM_SCALES_A_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
+    static constexpr uint32_t SMEM_SCALES_B_SIZE_PER_STAGE = BLOCK_N * sizeof(float);
+    static constexpr uint32_t ALIGNED_SMEM_SCALES_B_SIZE_PER_STAGE = ceil_div(SMEM_SCALES_B_SIZE_PER_STAGE, 128U) * 128U;
+
+    // Configs
+    constexpr uint32_t kFullKOfAllStages = kNumStages * BLOCK_K;
+    constexpr uint32_t kNumThreads = get_num_threads_per_sm<kNumTMAThreads, kNumMathThreadsPerGroup>(BLOCK_M);
+    constexpr uint32_t kNumMathThreads = kNumThreads - kNumTMAThreads;
+
+    const uint32_t shape_k_scales = ceil_div(shape_k, BLOCK_K);
+    const uint32_t num_iterations = ceil_div(shape_k, kFullKOfAllStages);
+    const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+    const uint32_t lane_idx = get_lane_id();
+
+    // Prefetch TMA descriptors at the very beginning
+    if (threadIdx.x == kNumMathThreads) {
+        // NOTES: `reinterpret_cast` must be here, or NVRTC will fail
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_a));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_b));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_scales_a));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_scales_b));
+        cute::prefetch_tma_descriptor(reinterpret_cast<const cute::TmaDescriptor*>(&tensor_map_d));
+    }
+    __syncwarp();
+
+    // Align to 1024 bytes for swizzle-128B
+    extern __shared__ __align__(1024) uint8_t smem_buffer[];
+    DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes");
+
+    // Data on shared memory
+    auto smem_d = reinterpret_cast<float*>(smem_buffer);
+    __nv_fp8_e4m3* smem_a[kNumStages];
+    __nv_fp8_e4m3* smem_b[kNumStages];
+    float* smem_scales_a[kNumStages];
+    float* smem_scales_b[kNumStages];
+
+    // TMA Barrier for both divisible and non-divisible cases
+    Barrier* full_barriers[kNumStages + 1];
+    Barrier* empty_barriers[kNumStages + 1];
+
+    // Fill shared memory pointers
+    #pragma unroll
+    for (int i = 0; i < kNumStages; ++ i) {
+        smem_a[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
+        smem_b[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
+        smem_scales_a[i] = reinterpret_cast<float*>(smem_buffer + SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) 
+                                                    + i * SMEM_SCALES_A_SIZE_PER_STAGE);
+        smem_scales_b[i] = reinterpret_cast<float*>(smem_buffer + SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE) 
+                                                    + i * ALIGNED_SMEM_SCALES_B_SIZE_PER_STAGE);
+    }
+
+    // Fill barriers
+    DG_STATIC_ASSERT(sizeof(Barrier) % sizeof(float) == 0, "Misaligned barriers");
+    auto barrier_start_ptr = reinterpret_cast<Barrier*>(smem_buffer + SMEM_D_SIZE + kNumStages 
+                                * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE + ALIGNED_SMEM_SCALES_B_SIZE_PER_STAGE));
+    #pragma unroll
+    for (int i = 0; i < kNumStages + 1; ++ i) {
+        full_barriers[i] = barrier_start_ptr + i;
+        empty_barriers[i] = barrier_start_ptr + kNumStages + 1 + i;
+    }
+
+    // Initialize barriers
+    DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "To many TMA multicast");
+    if (threadIdx.x == kNumMathThreads) {
+        // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
+        // even with TMA multicast disabled, we want to make the behavior aligned
+        #pragma unroll
+        for (int i = 0; i < kNumStages; ++ i) {
+            full_barriers[i]->init(1);
+            empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
+        }
+        full_barriers[kNumStages]->init(1);
+        empty_barriers[kNumStages]->init(1);
+
+        // Make initialized barrier visible in async proxy
+        cutlass::arch::fence_view_async_shared();
+        (kNumTMAMulticast > 1) ? cutlass::arch::fence_barrier_init() : void();
+    }
+
+    // Synchronize all threads to make barrier visible in normal memory model
+    (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads();
+
+    // For pipeline unrolling
+    struct DivisibleK {};
+    struct NotDivisibleK {};
+    auto launch_k_iterations = [&](const auto& func) {
+        if constexpr (kNumLastStages == 0) {
+            for (int k_iter = 0; k_iter < num_iterations; ++ k_iter)
+                func(k_iter, DivisibleK{});
+        } else {
+            for (int k_iter = 0; k_iter < num_iterations - 1; ++ k_iter)
+                func(k_iter, DivisibleK{});
+            func(num_iterations - 1, NotDivisibleK{});
+        }
+    };
+
+    // Register reconfigurations
+    constexpr int kNumTMARegisters = 40;
+    constexpr int kNumMathRegisters = 232;
+
+    // Block scheduler
+    uint32_t m_block_idx, n_block_idx;
+    auto scheduler = Scheduler<GemmType::Normal, SHAPE_N, BLOCK_M, BLOCK_N, 1, kNumTMAMulticast, kIsTMAMulticastOnA>(SHAPE_M);
+
+    if (threadIdx.x >= kNumMathThreads) {
+        // TMA warp-group for loading data
+        cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
+
+        // NOTES: only one thread (or warp) will be used
+        if (threadIdx.x == kNumMathThreads) {
+            // Persistently schedule over blocks
+            while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+                launch_k_iterations([&](int k_iter, auto type) {
+                    constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
+                    constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : kNumLastStages;
+                    DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
+
+                    // Assign TMA multicast number into A and B
+                    // NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
+                    const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
+                    const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                    const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
+                    DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
+
+                    #pragma unroll
+                    for (uint32_t s = 0; s < kNumInnerStages; ++ s) {
+                        // Wait consumer release
+                        empty_barriers[s]->wait((scheduler.current_iter * num_iterations + k_iter + 1) & 1);
+
+                        // Issue TMA A
+                        auto& full_barrier = *full_barriers[s];
+                        int k_idx = k_iter * kFullKOfAllStages + s * BLOCK_K;
+                        tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_a[s], k_idx, m_block_idx * BLOCK_M, num_tma_multicast_a);
+                        tma_copy(&tensor_map_scales_a, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_scales_a[s], m_block_idx * BLOCK_M,
+                                 k_idx / BLOCK_K, num_tma_multicast_a);
+
+                        // Issue TMA B
+                        tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_b[s], k_idx, n_block_idx * BLOCK_N, num_tma_multicast_b);
+                        tma_copy(&tensor_map_scales_b, reinterpret_cast<uint64_t*>(&full_barrier),
+                                 smem_scales_b[s], n_block_idx * BLOCK_N, k_idx / BLOCK_K, num_tma_multicast_b);
+
+                        full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SCALES_A_SIZE_PER_STAGE + SMEM_SCALES_B_SIZE_PER_STAGE);
+                    }
+
+                    // Wait unaligned cases
+                    #pragma unroll
+                    for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) {
+                        empty_barriers[s]->wait((scheduler.current_iter * num_iterations + k_iter + 1) & 1);
+                        full_barriers[s]->arrive();
+                    }
+                });
+
+                // Issue TMA D
+                empty_barriers[kNumStages]->wait((scheduler.current_iter + 1) & 1);
+                auto& full_barrier = *full_barriers[kNumStages];
+                tma_copy(&tensor_map_d, reinterpret_cast<uint64_t*>(&full_barrier),
+                         smem_d, n_block_idx * BLOCK_N, m_block_idx * BLOCK_M, 1);
+                full_barrier.arrive_and_expect_tx(SMEM_D_SIZE);
+            }
+
+            // To safely deconstruct distributed shared barriers, we need another round of empty waits
+            if constexpr (kNumTMAMulticast > 1) {
+                #pragma unroll
+                for (uint32_t s = 0; s < kNumStages; ++ s)
+                    empty_barriers[s]->wait((scheduler.current_iter * num_iterations + 1) & 1);
+            }
+        }
+    } else {
+        // Math warp-groups for WGMMA
+        cutlass::arch::warpgroup_reg_alloc<kNumMathRegisters>();
+
+        // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers
+        const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / kNumMathThreadsPerGroup, 0);
+        const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4;
+        const auto r_0 = warp_idx * 16 + row_idx, r_1 = r_0 + 8;
+
+        // Empty barrier arrival
+        auto empty_barrier_arrive = [&](int s) {
+            if constexpr (kNumTMAMulticast == 1) {
+                lane_idx == 0 ? empty_barriers[s]->arrive() : void();
+            } else {
+                auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
+                lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta) : void();
+            }
+        };
+
+        // Persistently schedule over blocks
+        while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
+            // Decide the number of scales B to load
+            DG_STATIC_ASSERT(SHAPE_N % 8 == 0, "Invalid shape N");
+            cutlass::arch::NamedBarrier(kNumMathThreads).sync();
+
+            // Accumulation for WGMMA or CUDA promotion
+            constexpr int WAVE_BLOCK_M = WGMMA::M * get_num_math_warpgroups(BLOCK_M);
+            float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
+            float2 scales_b[WGMMA::kNumAccum / 4];
+
+            // Launch MMAs
+            launch_k_iterations([&](int k_iter, auto type) {
+                constexpr bool kHasDivisibleStages = std::is_same_v<decltype(type), DivisibleK>;
+                constexpr int kNumInnerStages = kHasDivisibleStages ? kNumStages : kNumLastStages;
+                DG_STATIC_ASSERT(kNumInnerStages != 0, "Invalid number of inner stages");
+
+                #pragma unroll
+                for (int s = 0; s < kNumInnerStages; ++ s) {
+                    // Wait TMA arrivals
+                    full_barriers[s]->wait((scheduler.current_iter * num_iterations + k_iter) & 1);
+
+                    #pragma unroll
+                    for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+                      	auto m_offset = local_idx * WAVE_BLOCK_M;
+
+                        // Read A scales
+                        auto scale_a_0 = ld_shared(smem_scales_a[s] + r_0 + m_offset);
+                        auto scale_a_1 = ld_shared(smem_scales_a[s] + r_1 + m_offset);
+
+                        // Commit WGMMA instructions
+                        #pragma unroll
+                        for (int i = 0; i < WGMMA::kNumAccum; ++ i)
+                            warpgroup_fence_operand(accum[i]);
+                        warpgroup_arrive();
+                        #pragma unroll
+                        for (int k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
+                            auto desc_a = make_smem_desc(smem_a[s] + (math_wg_idx * WGMMA::M + m_offset) * BLOCK_K + k * WGMMA::K, 1);
+                            auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);
+                            WGMMA::wgmma(desc_a, desc_b, accum, k);
+                        }
+                        warpgroup_commit_batch();
+
+                        // Read B scales at the first warpgroup wave
+                        if (local_idx == 0) {
+                            #pragma unroll
+                            for (int i = 0; i < WGMMA::kNumAccum / 4; ++i)
+                                scales_b[i] = ld_shared(reinterpret_cast<float2*>(smem_scales_b[s] + i * 8 + col_idx * 2));
+                            __syncwarp();
+                        }
+
+                        #pragma unroll
+                        for (int i = 0; i < WGMMA::kNumAccum; ++ i)
+                            warpgroup_fence_operand(accum[i]);
+                        warpgroup_wait<0>();
+
+                        // Notify barrier arrival at the last warpgroup wave
+                        if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
+                            empty_barrier_arrive(s);
+
+                        // Promote with scales
+                        auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                        #pragma unroll
+                        for (int i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                            const float &scale_b_0 = scales_b[i].x;
+                            const float &scale_b_1 = scales_b[i].y;
+                            shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
+                            shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
+                            shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
+                            shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
+                        }
+                    }
+                }
+
+                // Wait last TMA store to be finished
+                if (k_iter == 0 and scheduler.current_iter > 0) {
+                    if (threadIdx.x == 0) {
+                        cute::tma_store_wait<0>();
+                        empty_barriers[kNumStages]->arrive();
+                    }
+                    __syncwarp();
+                }
+
+                // Wait unaligned cases
+                #pragma unroll
+                for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) {
+                    full_barriers[s]->wait((scheduler.current_iter * num_iterations + k_iter) & 1);
+                    empty_barrier_arrive(s);
+                }
+            });
+
+            // Wait TMA D arrivals
+            full_barriers[kNumStages]->wait(scheduler.current_iter & 1);
+
+            // Accumulate to D shared memory
+            #pragma unroll
+            for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
+            	auto m_offset = local_idx * WAVE_BLOCK_M;
+                auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
+                auto smem_d_0 = reinterpret_cast<float2*>(smem_d + (m_offset + r_0) * BLOCK_N + col_idx * 2);
+                auto smem_d_1 = reinterpret_cast<float2*>(smem_d + (m_offset + r_1) * BLOCK_N + col_idx * 2);
+                #pragma unroll
+                for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
+                    float2 d_0 = ld_shared(smem_d_0 + i * 4);
+                    st_shared(smem_d_0 + i * 4, {d_0.x + shifted_accum[i * 4 + 0], d_0.y + shifted_accum[i * 4 + 1]});
+                    float2 d_1 = ld_shared(smem_d_1 + i * 4);
+                    st_shared(smem_d_1 + i * 4, {d_1.x + shifted_accum[i * 4 + 2], d_1.y + shifted_accum[i * 4 + 3]});
+                }
+            }
+
+            cute::tma_store_fence();
+            cutlass::arch::NamedBarrier(kNumMathThreads).sync();
+
+            // Use TMA store to write back to global memory
+            if (threadIdx.x == 0) {
+                cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_d, n_block_idx * BLOCK_N, m_block_idx * BLOCK_M);
+                cute::tma_store_arrive();
+            }
+            __syncwarp();
+        }
+    }
+#else
+    if (blockIdx.x == 0 and threadIdx.x == 0)
+        DG_DEVICE_ASSERT(false && "This kernel only support sm_90a");
+#endif
+}
+
+}; // namespace deep_gemm
+
+#pragma clang diagnostic pop

ops/csrc/fp8/deep_gemm/include/deep_gemm/mma_utils.cuh

@@ -12,750 +12,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// The file has been adapted from DeepSeek DeepEP project
+// The file has been adapted from DeepSeek DeepGEMM project
 // Copyright (c) 2025 DeepSeek
-// Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 #pragma once
 
+#ifndef __CUDACC_RTC__
 #include <cuda.h>
+#endif
+
+#include <cute/arch/mma_sm90_gmma.hpp>
+#include <cute/arch/mma_sm90_gmma_ext.hpp>
 
 #include "utils.cuh"
 
 namespace deep_gemm {
 
-struct SM90_64x16x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %10, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n16k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7},"
-                     " %8,"
-                     " %9,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07)
-                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 16;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x24x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %14, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n24k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
-                     " %8,   %9,   %10,  %11},"
-                     " %12,"
-                     " %13,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11)
-                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 24;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x32x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %18, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15},"
-                     " %16,"
-                     " %17,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15)
-                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 32;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x40x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %22, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n40k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15, "
-                     " %16,  %17,  %18,  %19},"
-                     " %20,"
-                     " %21,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 40;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x48x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %26, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n48k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7, "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15, "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23},"
-                     " %24,"
-                     " %25,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 48;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x56x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %30, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n56k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27}, "
-                     " %28,"
-                     " %29,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                     : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                       "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                       "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                       "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27)
-                     : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 56;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x64x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %34, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31}, "
-                     " %32,"
-                     " %33,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 64;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x72x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %38, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n72k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35}, "
-                     " %36,"
-                     " %37,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 72;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x80x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %42, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n80k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39}, "
-                     " %40,"
-                     " %41,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 80;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x88x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %46, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n88k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43}, "
-                     " %44,"
-                     " %45,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 88;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x96x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %50, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n96k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47}, "
-                     " %48,"
-                     " %49,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 96;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x104x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 float& d48, float& d49, float& d50, float& d51,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %54, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n104k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
-                     " %48,  %49,  %50,  %51}, "
-                     " %52,"
-                     " %53,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
-                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              d[48], d[49], d[50], d[51],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 104;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x112x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %58, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n112k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
-                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55}, "
-                     " %56,"
-                     " %57,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
-                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 112;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x120x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
-                                 float& d56, float& d57, float& d58, float& d59,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %62, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n120k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
-                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
-                     " %56,  %57,  %58,  %59}, "
-                     " %60,"
-                     " %61,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
-                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
-                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
-              d[56], d[57], d[58], d[59],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 120;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x128x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
-                                 float& d56, float& d57, float& d58, float& d59, float& d60, float& d61, float& d62, float& d63,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %66, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
-                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
-                     " %56,  %57,  %58,  %59,  %60,  %61,  %62,  %63}, "
-                     " %64,"
-                     " %65,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
-                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
-                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
-              d[56], d[57], d[58], d[59], d[60], d[61], d[62], d[63],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 128;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
-struct SM90_64x192x32_F32E4M3E4M3_SS {
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b,
-                                 float& d00, float& d01, float& d02, float& d03, float& d04, float& d05, float& d06, float& d07,
-                                 float& d08, float& d09, float& d10, float& d11, float& d12, float& d13, float& d14, float& d15,
-                                 float& d16, float& d17, float& d18, float& d19, float& d20, float& d21, float& d22, float& d23,
-                                 float& d24, float& d25, float& d26, float& d27, float& d28, float& d29, float& d30, float& d31,
-                                 float& d32, float& d33, float& d34, float& d35, float& d36, float& d37, float& d38, float& d39,
-                                 float& d40, float& d41, float& d42, float& d43, float& d44, float& d45, float& d46, float& d47,
-                                 float& d48, float& d49, float& d50, float& d51, float& d52, float& d53, float& d54, float& d55,
-                                 float& d56, float& d57, float& d58, float& d59, float& d60, float& d61, float& d62, float& d63,
-                                 float& d64, float& d65, float& d66, float& d67, float& d68, float& d69, float& d70, float& d71,
-                                 float& d72, float& d73, float& d74, float& d75, float& d76, float& d77, float& d78, float& d79,
-                                 float& d80, float& d81, float& d82, float& d83, float& d84, float& d85, float& d86, float& d87,
-                                 float& d88, float& d89, float& d90, float& d91, float& d92, float& d93, float& d94, float& d95,
-                                 bool scale_d) {
-        asm volatile("{\n"
-                     ".reg .pred p;\n"
-                     "setp.ne.b32 p, %98, 0;\n"
-                     "wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3"
-                     "{%0,   %1,   %2,   %3,   %4,   %5,   %6,   %7,   "
-                     " %8,   %9,   %10,  %11,  %12,  %13,  %14,  %15,  "
-                     " %16,  %17,  %18,  %19,  %20,  %21,  %22,  %23,  "
-                     " %24,  %25,  %26,  %27,  %28,  %29,  %30,  %31,  "
-                     " %32,  %33,  %34,  %35,  %36,  %37,  %38,  %39,  "
-                     " %40,  %41,  %42,  %43,  %44,  %45,  %46,  %47,  "
-                     " %48,  %49,  %50,  %51,  %52,  %53,  %54,  %55,  "
-                     " %56,  %57,  %58,  %59,  %60,  %61,  %62,  %63,  "
-                     " %64,  %65,  %66,  %67,  %68,  %69,  %70,  %71,  "
-                     " %72,  %73,  %74,  %75,  %76,  %77,  %78,  %79,  "
-                     " %80,  %81,  %82,  %83,  %84,  %85,  %86,  %87,  "
-                     " %88,  %89,  %90,  %91,  %92,  %93,  %94,  %95},  "
-                     " %96,"
-                     " %97,"
-                     " p   , 1,    1;\n"
-                     "}\n"
-                    : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07),
-                      "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15),
-                      "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23),
-                      "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31),
-                      "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39),
-                      "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47),
-                      "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55),
-                      "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63),
-                      "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71),
-                      "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79),
-                      "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87),
-                      "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95)
-                    : "l"(desc_a), "l"(desc_b), "r"(int32_t(scale_d)));
-    }
-
-    __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
-        wgmma(desc_a, desc_b,
-              d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7],
-              d[8], d[9], d[10], d[11], d[12], d[13], d[14], d[15],
-              d[16], d[17], d[18], d[19], d[20], d[21], d[22], d[23],
-              d[24], d[25], d[26], d[27], d[28], d[29], d[30], d[31],
-              d[32], d[33], d[34], d[35], d[36], d[37], d[38], d[39],
-              d[40], d[41], d[42], d[43], d[44], d[45], d[46], d[47],
-              d[48], d[49], d[50], d[51], d[52], d[53], d[54], d[55],
-              d[56], d[57], d[58], d[59], d[60], d[61], d[62], d[63],
-              d[64], d[65], d[66], d[67], d[68], d[69], d[70], d[71],
-              d[72], d[73], d[74], d[75], d[76], d[77], d[78], d[79],
-              d[80], d[81], d[82], d[83], d[84], d[85], d[86], d[87],
-              d[88], d[89], d[90], d[91], d[92], d[93], d[94], d[95],
-              scale_d);
-    }
-
-    static constexpr int M = 64;
-    static constexpr int N = 192;
-    static constexpr int K = 32;
-    static constexpr int kNumAccum = M * N / 128;
-};
-
 template <typename dtype_t>
 struct SM90_U32x2_STSM_N {
     __device__ __forceinline__ static void
@@ -777,15 +50,15 @@ struct SM90_U32x4_STSM_N {
     }
 };
 
-__device__ void warpgroup_arrive() {
+__forceinline__ __device__ void warpgroup_arrive() {
     asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
 }
 
-__device__ void warpgroup_commit_batch() {
+__forceinline__ __device__ void warpgroup_commit_batch() {
     asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
 }
 
-__device__ void warpgroup_fence_operand(float& reg) {
+__forceinline__ __device__ void warpgroup_fence_operand(float& reg) {
     asm volatile("" : "+f"(reg) :: "memory");
 }
 
@@ -813,6 +86,12 @@ __device__  __forceinline__ float ld_shared(const float* __restrict__ ptr) {
     return ret;
 }
 
+__device__  __forceinline__ float2 ld_shared(const float2* __restrict__ ptr) {
+    float2 ret;
+    asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];" : "=f"(ret.x), "=f"(ret.y) : "l"(ptr));
+    return ret;
+}
+
 __device__ __forceinline__ void st_shared(const float* ptr, float val) {
     asm volatile("st.shared.f32 [%0], %1;" :: "l"(ptr), "f"(val));
 }
@@ -821,6 +100,10 @@ __device__ __forceinline__ void st_shared(const uint32_t* ptr, uint32_t val) {
     asm volatile("st.shared.u32 [%0], %1;" :: "l"(ptr), "r"(val));
 }
 
+__device__ __forceinline__ void st_shared(const float2* ptr, float2 val) {
+    asm volatile("st.shared.v2.f32 [%0], {%1, %2};" :: "l"(ptr), "f"(val.x), "f"(val.y));
+}
+
 template <int N>
 __device__ void warpgroup_wait() {
     DG_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]");
@@ -876,28 +159,72 @@ __device__ GmmaDescriptor make_smem_desc(PointerType smem_ptr, int layout_type,
     return desc;
 }
 
+template <int N_, typename MMA>
+struct FP8MMA {
+    
+    template <size_t ...Idx>
+    __forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, std::index_sequence<Idx...>) {
+        using namespace cute::SM90::GMMA;
+        MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
+    }
+
+    __forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
+        call_fma_impl(desc_a, desc_b, d, scale_d, std::make_index_sequence<N_/2>{});
+    }
+
+    static constexpr int M = 64;
+    static constexpr int N = N_;
+    static constexpr int K = 32;
+    static constexpr int kNumAccum = M * N / 128;
+};
+
 template <int N>
 struct FP8MMASelector {
+
+    static constexpr auto select_mma() {
+        using namespace cute::SM90::GMMA;
+        if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_SS_TN();
+        if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_SS_TN();
+    }
+
     static constexpr auto select_type() {
-        if constexpr (N == 16) return SM90_64x16x32_F32E4M3E4M3_SS();
-        if constexpr (N == 24) return SM90_64x24x32_F32E4M3E4M3_SS();
-        if constexpr (N == 32) return SM90_64x32x32_F32E4M3E4M3_SS();
-        if constexpr (N == 40) return SM90_64x40x32_F32E4M3E4M3_SS();
-        if constexpr (N == 48) return SM90_64x48x32_F32E4M3E4M3_SS();
-        if constexpr (N == 56) return SM90_64x56x32_F32E4M3E4M3_SS();
-        if constexpr (N == 64) return SM90_64x64x32_F32E4M3E4M3_SS();
-        if constexpr (N == 72) return SM90_64x72x32_F32E4M3E4M3_SS();
-        if constexpr (N == 80) return SM90_64x80x32_F32E4M3E4M3_SS();
-        if constexpr (N == 88) return SM90_64x88x32_F32E4M3E4M3_SS();
-        if constexpr (N == 96) return SM90_64x96x32_F32E4M3E4M3_SS();
-        if constexpr (N == 104) return SM90_64x104x32_F32E4M3E4M3_SS();
-        if constexpr (N == 112) return SM90_64x112x32_F32E4M3E4M3_SS();
-        if constexpr (N == 120) return SM90_64x120x32_F32E4M3E4M3_SS();
-        if constexpr (N == 128) return SM90_64x128x32_F32E4M3E4M3_SS();
-        if constexpr (N == 192) return SM90_64x192x32_F32E4M3E4M3_SS();
+        return FP8MMA<N, decltype(select_mma())>();
     }
 
     using type = decltype(select_type());
 };
 
+enum class Layout {
+    RowMajor,
+    ColMajor
+};
+
+__device__ __host__ constexpr int get_num_math_warpgroups(int block_m) {
+    return block_m == 64 ? 1 : 2;
+}
+
+template <uint32_t kNumTMAThreads, uint32_t kNumMathThreadsPerGroup>
+__device__ __host__ constexpr int get_num_threads_per_sm(int block_m) {
+    DG_STATIC_ASSERT(kNumMathThreadsPerGroup == 128, "Only support 128 threads per math group");
+    return get_num_math_warpgroups(block_m) * kNumMathThreadsPerGroup + kNumTMAThreads;
+}
+
 } // namespace deep_gemm

ops/csrc/fp8/deep_gemm/include/deep_gemm/nvrtc_std.cuh

@@ -0,0 +1,104 @@
+// 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.
+
+// The file has been adapted from DeepSeek DeepGEMM project
+// Copyright (c) 2025 DeepSeek
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+#pragma once
+
+#ifdef __CUDACC_RTC__
+
+using int8_t = signed char;
+using uint8_t = unsigned char;
+using int16_t = signed short;
+using uint16_t = unsigned short;
+using int32_t = signed int;
+using uint32_t = unsigned int;
+using int64_t = signed long long;
+using uint64_t = unsigned long long;
+using cuuint64_t = unsigned long long;
+
+#ifndef CU_TENSOR_MAP_NUM_QWORDS
+#define CU_TENSOR_MAP_NUM_QWORDS 16
+
+struct CUtensorMap_st {
+#if defined(__cplusplus) && (__cplusplus >= 201103L)
+    alignas(64)
+#elif __STDC_VERSION__ >= 201112L
+    _Alignas(64)
+#endif
+        cuuint64_t opaque[CU_TENSOR_MAP_NUM_QWORDS];
+};
+
+using CUtensorMap = CUtensorMap_st;
+#endif
+
+namespace std {
+
+template <class T, T v> struct integral_constant {
+  static constexpr T value = v;
+
+  using value_type = T;
+  using type = integral_constant;
+
+  __device__ constexpr operator value_type() const noexcept { return value; }
+
+  __device__ constexpr value_type operator()() const noexcept { return value; }
+};
+
+using false_type = integral_constant<bool, false>;
+using true_type = integral_constant<bool, true>;
+
+template <class T, class U> struct is_same : false_type {};
+
+template <class T> struct is_same<T, T> : true_type {};
+
+template <class T, class U>
+inline constexpr bool is_same_v = is_same<T, U>::value;
+
+namespace index_sequence_impl {
+
+// Based on https://stackoverflow.com/a/32223343/11717224
+template <size_t... Ints> struct index_sequence {
+  using type = index_sequence;
+  using value_type = size_t;
+  static constexpr size_t size() noexcept { return sizeof...(Ints); }
+};
+
+template <class Sequence1, class Sequence2> struct _merge_and_renumber;
+
+template <size_t... I1, size_t... I2>
+struct _merge_and_renumber<index_sequence<I1...>, index_sequence<I2...>>
+    : index_sequence<I1..., (sizeof...(I1) + I2)...> {};
+
+template <size_t N>
+struct make_index_sequence
+    : _merge_and_renumber<typename make_index_sequence<N / 2>::type,
+                          typename make_index_sequence<N - N / 2>::type> {};
+
+template <> struct make_index_sequence<0> : index_sequence<> {};
+template <> struct make_index_sequence<1> : index_sequence<0> {};
+
+} // namespace index_sequence_impl
+
+template <size_t... Ns>
+using index_sequence = index_sequence_impl::index_sequence<Ns...>;
+
+template <size_t N>
+using make_index_sequence = index_sequence_impl::make_index_sequence<N>;
+
+} // namespace std
+
+#endif

ops/csrc/fp8/deep_gemm/include/deep_gemm/scheduler.cuh

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// The file has been adapted from DeepSeek DeepEP project
+// The file has been adapted from DeepSeek DeepGEMM project
 // Copyright (c) 2025 DeepSeek
-// Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
+#pragma once
 #include "utils.cuh"
 
 namespace deep_gemm {
@@ -30,9 +31,10 @@ enum class GemmType {
 #pragma ide diagnostic ignored "cppcoreguidelines-pro-type-member-init"
 template <GemmType kGemmType,
           uint32_t SHAPE_N, uint32_t BLOCK_M, uint32_t BLOCK_N,
-          uint32_t kNumGroups, uint32_t kNumTMAMulticast,
+          uint32_t kNumGroups,
+          uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
           uint32_t kNumNBlocks = ceil_div(SHAPE_N, BLOCK_N),
-          uint32_t kNumNBlocksPerGroup = 16>
+          uint32_t kNum1DBlocksPerGroup = 16>
 struct Scheduler {
     int current_iter = -1;
     uint32_t num_aligned_m_blocks;
@@ -40,13 +42,16 @@ struct Scheduler {
     // For normal GEMM
     // Maybe not used in the masked grouped GEMM
     uint32_t num_blocks;
+    uint32_t num_blocks_in_group;
+    bool is_peer_cta_alive = true;
 
     // For grouped GEMM
     int* grouped_layout;
+
     // Only used for masked layout
     uint32_t curr_group_idx, curr_cumsum;
 
-    __device__ __forceinline__ explicit Scheduler(const uint32_t shape_m,
+    __device__ __forceinline__ explicit Scheduler(const uint32_t& shape_m,
                                                   int* grouped_layout = nullptr) {
         num_aligned_m_blocks = ceil_div(shape_m, BLOCK_M);
         if constexpr (kGemmType == GemmType::Normal) {
@@ -60,28 +65,77 @@ struct Scheduler {
         }
     }
 
-    __device__ __forceinline__ void get_swizzled_block_idx(const uint32_t num_m_blocks, int block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
-        DG_STATIC_ASSERT(kNumNBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
+    // ReSharper disable once CppNotAllPathsReturnValue
+    __device__ __forceinline__ bool is_computation_valid(const uint32_t& m_block_idx, const uint32_t& m_offset) const {
+        if constexpr (kGemmType == GemmType::Normal) {
+            return true;
+        } else if constexpr (kGemmType == GemmType::GroupedContiguous) {
+            return __ldg(grouped_layout + m_offset + m_block_idx * BLOCK_M) >= 0;
+        } else if constexpr (kGemmType == GemmType::GroupedMasked) {
+            return m_offset + m_block_idx * BLOCK_M < __ldg(grouped_layout + curr_group_idx);
+        }
+    }
+
+    __device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) const {
+        if (num_blocks_in_group == 1)
+            return false;
+        if constexpr (kGemmType == GemmType::Normal or kGemmType == GemmType::GroupedMasked) {
+            return true;
+        } else {
+            DG_STATIC_ASSERT(kGemmType == GemmType::GroupedContiguous, "Invalid Gemm type");
+            if constexpr (kIsTMAMulticastOnA) {
+                return true;
+            } else {
+                auto group_idx = __ldg(grouped_layout + m_block_idx * BLOCK_M);
+                auto peer_group_idx = __ldg(grouped_layout + (m_block_idx ^ 1) * BLOCK_M);
+                return group_idx == peer_group_idx;
+            }
+        }
+    }
+
+    __device__ __forceinline__ void get_swizzled_block_idx(const uint32_t& num_m_blocks, const uint32_t& block_idx,
+                                                           uint32_t& m_block_idx, uint32_t& n_block_idx) {
+        DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");
 
         // Swizzle for better L2 usages
-        auto num_blocks_per_group = num_m_blocks * kNumNBlocksPerGroup;
+        auto primary_num_blocks = kIsTMAMulticastOnA ? kNumNBlocks : num_m_blocks;
+        auto secondary_num_blocks = kIsTMAMulticastOnA ? num_m_blocks : kNumNBlocks;
+        auto num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup;
         auto group_idx = block_idx / num_blocks_per_group;
-        auto first_n_block_idx = group_idx * kNumNBlocksPerGroup;
-        auto num_n_blocks_in_group = min(kNumNBlocksPerGroup, kNumNBlocks - first_n_block_idx);
+        auto first_block_idx = group_idx * kNum1DBlocksPerGroup;
         auto in_group_idx = block_idx % num_blocks_per_group;
-        m_block_idx = in_group_idx / num_n_blocks_in_group;
-        n_block_idx = first_n_block_idx + in_group_idx % num_n_blocks_in_group;
+        num_blocks_in_group = min(kNum1DBlocksPerGroup, primary_num_blocks - first_block_idx);
+
+        // Fix unaligned TMA multicast
+        if (kNumTMAMulticast > 1 and num_blocks_in_group % 2 != 0) {
+            if (in_group_idx < (num_blocks_in_group ^ 1) * secondary_num_blocks) {
+                num_blocks_in_group = num_blocks_in_group ^ 1;
+            } else {
+                in_group_idx = in_group_idx - (num_blocks_in_group ^ 1) * secondary_num_blocks;
+                first_block_idx += num_blocks_in_group ^ 1;
+                num_blocks_in_group = 1;
+            }
+        }
+
+        // Convert to final M/N block indices
+        if constexpr (kIsTMAMulticastOnA) {
+            m_block_idx = in_group_idx / num_blocks_in_group;
+            n_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
+        } else {
+            m_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
+            n_block_idx = in_group_idx / num_blocks_in_group;
+        }
     }
 
     template <bool kIgnoreGroupedForGroupedContiguous=true>
-    __device__ __forceinline__ uint32_t get_global_idx(const uint32_t shape_dim, const uint32_t block_size,
+    __device__ __forceinline__ uint32_t get_global_idx(const uint32_t& shape_dim, const uint32_t& block_size,
                                                        const uint32_t& block_idx, const uint32_t& m_block_idx=0) {
         if constexpr (kGemmType == GemmType::Normal) {
             return block_idx * block_size;
-        } else if (kGemmType == GemmType::GroupedContiguous) {
+        } else if constexpr (kGemmType == GemmType::GroupedContiguous) {
             auto offset = kIgnoreGroupedForGroupedContiguous ? 0 : __ldg(grouped_layout + m_block_idx * BLOCK_M);
             return offset * shape_dim + block_idx * block_size;
-        } else if (kGemmType == GemmType::GroupedMasked) {
+        } else if constexpr (kGemmType == GemmType::GroupedMasked) {
             return curr_group_idx * shape_dim + block_idx * block_size;
         }
     }
@@ -96,7 +150,7 @@ struct Scheduler {
                 if (curr_group_idx == kNumGroups)
                     return false;
 
-                // Within current group
+                // Within the current group
                 num_m_blocks = ceil_div(static_cast<uint32_t>(__ldg(grouped_layout + curr_group_idx)), BLOCK_M);
                 auto current_m_block_cumsum = curr_cumsum + num_m_blocks;
                 if (next_block_idx < current_m_block_cumsum * kNumNBlocks)
@@ -111,11 +165,16 @@ struct Scheduler {
             if (next_block_idx >= num_blocks)
                 return false;
 
+            // NOTES: we don't have to set `is_peer_cta_alive` for masked grouped GEMM, as it must be aligned
+            is_peer_cta_alive = kNumNBlocks % kNumTMAMulticast == 0 or          // Always aligned on N (constant bypass)
+                                num_aligned_m_blocks % kNumTMAMulticast == 0 or // Always aligned on M (constant bypass)
+                                (next_block_idx ^ 1) < num_blocks;              // Peer CTA in bound
             get_swizzled_block_idx(num_aligned_m_blocks, next_block_idx, m_block_idx, n_block_idx);
         }
         return true;
     }
 };
+
 #pragma clang diagnostic pop
 
 } // namespace deep_gemm

ops/csrc/fp8/deep_gemm/include/deep_gemm/tma_utils.cuh

@@ -12,104 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// The file has been adapted from DeepSeek DeepEP project
+// The file has been adapted from DeepSeek DeepGEMM project
 // Copyright (c) 2025 DeepSeek
-// Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 #pragma once
 
-#include <cassert>
-#include <cuda.h>
-#include <cudaTypedefs.h>
-#include <cuda_fp8.h>
-#include <cuda_runtime.h>
-#include <cuda/barrier>
-
 #include "utils.cuh"
 
 namespace deep_gemm {
 
-template <class T>
-constexpr CUtensorMapDataType get_CUtensorMapDataType() {
-    if constexpr (std::is_same<T, uint8_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT8;
-    } else if constexpr (std::is_same<T, __nv_fp8_e4m3>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT8;
-    } else if constexpr (std::is_same<T, __nv_fp8_e5m2>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT8;
-    } else if constexpr (std::is_same<T, uint16_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT16;
-    } else if constexpr (std::is_same<T, uint32_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT32;
-    } else if constexpr (std::is_same<T, uint64_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_UINT64;
-    } else if constexpr (std::is_same<T, int32_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_INT32;
-    } else if constexpr (std::is_same<T, int64_t>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_INT64;
-    } else if constexpr (std::is_same<T, __half>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_FLOAT16;
-    } else if constexpr (std::is_same<T, float>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
-    } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
-    }  else if constexpr (std::is_same<T, double>::value) {
-        return CU_TENSOR_MAP_DATA_TYPE_FLOAT64;
-    }
-}
-
-PFN_cuTensorMapEncodeTiled get_cuTensorMapEncodeTiled() {
-    // Get pointer to `cuTensorMapEncodeTiled`
-    cudaDriverEntryPointQueryResult driver_status;
-    void* cuTensorMapEncodeTiled_ptr = nullptr;
-
-/*
-#if CUDA_VERSION >= 12050
-    cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr, 12000,
-                                     cudaEnableDefault, &driver_status);
-#else
-*/
-    cudaGetDriverEntryPoint("cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr,
-                            cudaEnableDefault, &driver_status);
-//#endif
-
-    if (driver_status != cudaDriverEntryPointSuccess)
-        throw std::runtime_error("driver_status != cudaDriverEntryPointSuccess");
-    return reinterpret_cast<PFN_cuTensorMapEncodeTiled>(cuTensorMapEncodeTiled_ptr);
-}
-
-template <typename T>
-CUtensorMap make_2d_tma_copy_desc(T* global_address, uint64_t gmem_dim[2],
-                                  uint64_t stride_in_bytes, uint32_t smem_dim[2],
-                                  CUtensorMapSwizzle swizzle_type,
-                                  PFN_cuTensorMapEncodeTiled encode_func = nullptr) {
-    CUtensorMap tensor_map{};
-    constexpr uint32_t rank = 2;
-    uint64_t global_stride[rank - 1] = {stride_in_bytes};
-    uint32_t elem_strides[rank] = {1, 1};
-
-    if (encode_func == nullptr)
-        encode_func = get_cuTensorMapEncodeTiled();
-
-    auto result = encode_func(
-            &tensor_map, get_CUtensorMapDataType<typename std::remove_cv<T>::type>(), rank,
-            global_address, gmem_dim, global_stride, smem_dim, elem_strides,
-            CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle_type,
-            CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
-            CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-    DG_HOST_ASSERT(result == CUDA_SUCCESS);
-    return tensor_map;
-}
-
-template <uint32_t kNumTMAMulticast = 1>
+// TODO: move this function to other files
 __device__ __forceinline__ void
 tma_copy(void const* desc_ptr, uint64_t* barrier_ptr, void* smem_ptr,
-         int32_t const& crd_0, int32_t const& crd_1) {
+         int32_t const& crd_0, int32_t const& crd_1, uint32_t num_tma_multicast) {
     constexpr auto cache_hint = static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
-    if constexpr (kNumTMAMulticast == 1) {
+    if (num_tma_multicast == 1) {
         cute::SM90_TMA_LOAD_2D::copy(desc_ptr, barrier_ptr, cache_hint, smem_ptr, crd_0, crd_1);
     } else if (cute::block_rank_in_cluster() == 0) {
-        cute::SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, barrier_ptr, (1 << kNumTMAMulticast) - 1, cache_hint, smem_ptr, crd_0, crd_1);
+        cute::SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, barrier_ptr, (1 << num_tma_multicast) - 1, cache_hint, smem_ptr, crd_0, crd_1);
     }
 }
 

ops/csrc/fp8/deep_gemm/include/deep_gemm/utils.cuh

@@ -12,40 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// The file has been adapted from DeepSeek DeepEP project
+// The file has been adapted from DeepSeek DeepGEMM project
 // Copyright (c) 2025 DeepSeek
-// Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+// Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 #pragma once
 
-#include <exception>
-
 #ifdef __CLION_IDE__
-__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) { asm volatile("trap;"); }
+
+__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) {
+    asm volatile("trap;");
+}
+
 #define printf host_device_printf
 #endif
 
-class AssertionException : public std::exception {
-private:
-    std::string message{};
-
-public:
-    explicit AssertionException(const std::string& message) : message(message) {}
-
-    const char *what() const noexcept override { return message.c_str(); }
-};
-
-#ifndef DG_HOST_ASSERT
-#define DG_HOST_ASSERT(cond)                                        \
-do {                                                                \
-    if (not (cond)) {                                               \
-        printf("Assertion failed: %s:%d, condition: %s\n",          \
-               __FILE__, __LINE__, #cond);                          \
-        throw AssertionException("Assertion failed: " #cond);       \
-    }                                                               \
-} while (0)
-#endif
-
 #ifndef DG_DEVICE_ASSERT
 #define DG_DEVICE_ASSERT(cond)                                                          \
 do {                                                                                    \
@@ -63,4 +44,9 @@ do {
 template <typename T>
 __device__ __host__ constexpr T ceil_div(T a, T b) {
     return (a + b - 1) / b;
+}
+
+template <typename T>
+__device__ __host__ constexpr T constexpr_gcd(T a, T b) {
+    return b == 0 ? a : constexpr_gcd(b, a % b);
 }

ops/csrc/fp8/deep_gemm/jit/__init__.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
-from .compiler import build, get_nvcc_compiler
-from .runtime import Runtime
-from .template import cpp_format, generate
+from .compiler import NVCCCompiler, NVRTCCompiler, build, get_nvcc_compiler
+from .runtime import FP8GemmRuntime, FP8WGradGemmRuntime, Runtime

ops/csrc/fp8/deep_gemm/jit/compiler.py

@@ -12,49 +12,59 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 import functools
 import hashlib
 import os
 import re
 import subprocess
+import time
 import uuid
-from typing import Tuple
+from typing import Any, Dict, List, Tuple, Type
+
+import cuda.bindings
+import cuda.bindings.nvrtc as nvrtc
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
 
-from ..utils import get_cuda_home
 from . import interleave_ffma
 from .runtime import Runtime, RuntimeCache
-from .template import typename_map
 
 runtime_cache = RuntimeCache()
 
 
 def hash_to_hex(s: str) -> str:
-    md5 = hashlib.md5()
-    md5.update(s.encode("utf-8"))
-    return md5.hexdigest()[0:12]
+    """Transforms input into a hexadecimal hash string"""
+    if isinstance(s, int):
+        return f"{s & 0xFFFFFFFFFFFFFFFF:016x}"
+    elif isinstance(s, str):
+        md5 = hashlib.md5()
+        md5.update(s.encode("utf-8"))
+        return md5.hexdigest()[0:12]
+    else:
+        raise TypeError(f"Unsupported type for hashing: {type(s)}")
 
 
 @functools.lru_cache(maxsize=None)
 def get_jit_include_dir() -> str:
-    return f"{os.path.dirname(os.path.abspath(__file__))}/../include"
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "include")
 
 
 @functools.lru_cache(maxsize=None)
 def get_deep_gemm_version() -> str:
-    # Update include directories
-    include_dir = f"{get_jit_include_dir()}/deep_gemm"
-    assert os.path.exists(include_dir), f"Cannot find GEMM include directory {include_dir}"
     md5 = hashlib.md5()
+
+    # Update include directories
+    include_dir = os.path.join(get_jit_include_dir(), "deep_gemm")
+    assert os.path.exists(include_dir), f"Cannot find GEMM include directory {include_dir}"
     for filename in filter(lambda x: x.endswith(".cuh"), sorted(os.listdir(include_dir))):
-        with open(f"{include_dir}/{filename}", "rb") as f:
+        with open(os.path.join(include_dir, filename), "rb") as f:
             md5.update(f.read())
 
     # Update `interleave_ffma.py`
-    with open(f"{os.path.dirname(os.path.realpath(__file__))}/interleave_ffma.py", "rb") as f:
+    with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "interleave_ffma.py"), "rb") as f:
         md5.update(f.read())
     return md5.hexdigest()[0:12]
 
@@ -62,17 +72,18 @@ def get_deep_gemm_version() -> str:
 @functools.lru_cache(maxsize=None)
 def get_nvcc_compiler() -> Tuple[str, str]:
     paths = []
-    if os.getenv("DG_NVCC_COMPILER"):
-        paths.append(os.getenv("DG_NVCC_COMPILER"))
-    CUDA_HOME = get_cuda_home()
-    paths.append(f"{CUDA_HOME}/bin/nvcc")
+    if os.getenv("DG_JIT_NVCC_COMPILER"):
+        paths.append(os.getenv("DG_JIT_NVCC_COMPILER"))
+    paths.append(os.path.join(CUDA_HOME, "bin", "nvcc"))
 
     # Try to find the first available NVCC compiler
     least_version_required = "12.3"
     version_pattern = re.compile(r"release (\d+\.\d+)")
     for path in paths:
         if os.path.exists(path):
-            match = version_pattern.search(os.popen(f"{path} --version").read())
+            command = [path, "--version"]
+            result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+            match = version_pattern.search(result.stdout)
             version = match.group(1)
             assert match, f"Cannot get the version of NVCC compiler {path}"
             assert (
@@ -84,21 +95,21 @@ def get_nvcc_compiler() -> Tuple[str, str]:
 
 @functools.lru_cache(maxsize=None)
 def get_default_user_dir():
-    if "DG_CACHE_DIR" in os.environ:
-        path = os.getenv("DG_CACHE_DIR")
+    if "DG_JIT_CACHE_DIR" in os.environ:
+        path = os.getenv("DG_JIT_CACHE_DIR")
         os.makedirs(path, exist_ok=True)
         return path
-    return os.path.expanduser("~") + "/.deep_gemm"
+    return os.path.join(os.path.expanduser("~"), ".deep_gemm")
 
 
 @functools.lru_cache(maxsize=None)
 def get_tmp_dir():
-    return f"{get_default_user_dir()}/tmp"
+    return os.path.join(get_default_user_dir(), "tmp")
 
 
 @functools.lru_cache(maxsize=None)
 def get_cache_dir():
-    return f"{get_default_user_dir()}/cache"
+    return os.path.join(get_default_user_dir(), "cache")
 
 
 def make_tmp_dir():
@@ -107,68 +118,204 @@ def make_tmp_dir():
     return tmp_dir
 
 
-def put(path, data, is_binary=False):
+def put(path, data):
     # Write and do POSIX atomic replace
-    tmp_file_path = f"{make_tmp_dir()}/file.tmp.{str(uuid.uuid4())}.{hash_to_hex(path)}"
-    with open(tmp_file_path, "wb" if is_binary else "w") as f:
+    tmp_file_path = os.path.join(make_tmp_dir(), f"file.tmp.{str(uuid.uuid4())}.{hash_to_hex(path)}")
+    with open(tmp_file_path, "wb" if isinstance(data, bytes) else "w") as f:
         f.write(data)
     os.replace(tmp_file_path, path)
 
 
-def build(name: str, arg_defs: tuple, code: str) -> Runtime:
-    # Compiler flags
-    nvcc_flags = [
-        "-std=c++17",
-        "-shared",
-        "-O3",
-        "--expt-relaxed-constexpr",
-        "--expt-extended-lambda",
-        "-gencode=arch=compute_90a,code=sm_90a",
-        "--ptxas-options=--register-usage-level=10" + (",--verbose" if "DG_PTXAS_VERBOSE" in os.environ else ""),
-        # Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases
-        "--diag-suppress=177,174,940",
-    ]
-    cxx_flags = ["-fPIC", "-O3", "-Wno-deprecated-declarations", "-Wno-abi"]
-    flags = [*nvcc_flags, f'--compiler-options={",".join(cxx_flags)}']
-    include_dirs = [get_jit_include_dir()]
+class Compiler:
+    @classmethod
+    def signature(cls) -> str:
+        pass
 
-    # Build signature
-    enable_sass_opt = get_nvcc_compiler()[1] <= "12.8" and int(os.getenv("DG_DISABLE_FFMA_INTERLEAVE", 0)) == 0
-    signature = f"{name}$${get_deep_gemm_version()}$${code}$${get_nvcc_compiler()}$${flags}$${enable_sass_opt}"
-    name = f"kernel.{name}.{hash_to_hex(signature)}"
-    path = f"{get_cache_dir()}/{name}"
+    @staticmethod
+    def __version__() -> Tuple[int, int]:
+        pass
 
-    # Check runtime cache or file system hit
-    global runtime_cache
-    if runtime_cache[path] is not None:
-        if os.getenv("DG_JIT_DEBUG", None):
-            print(f"Using cached JIT runtime {name} during build")
-        return runtime_cache[path]
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        pass
 
-    # Write the code
-    os.makedirs(path, exist_ok=True)
-    args_path = f"{path}/kernel.args"
-    src_path = f"{path}/kernel.cu"
-    put(args_path, ", ".join([f"('{arg_def[0]}', {typename_map[arg_def[1]]})" for arg_def in arg_defs]))
-    put(src_path, code)
+    @staticmethod
+    def flags() -> List[str]:
+        cpp_standard = int(os.getenv("DG_JIT_OVERRIDE_CPP_STANDARD", 20))
+        return [
+            f"-std=c++{cpp_standard}",
+            "--ptxas-options=--register-usage-level=10"
+            + (",--verbose" if "DG_JIT_PTXAS_VERBOSE" in os.environ else ""),
+            # Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases
+            "--diag-suppress=39,161,174,177,186,940",
+        ]
 
-    # Compile into a temporary SO file
-    so_path = f"{path}/kernel.so"
-    tmp_so_path = f"{make_tmp_dir()}/nvcc.tmp.{str(uuid.uuid4())}.{hash_to_hex(so_path)}.so"
-    # Compile
-    command = [get_nvcc_compiler()[0], src_path, "-o", tmp_so_path, *flags, *[f"-I{d}" for d in include_dirs]]
-    if os.getenv("DG_JIT_DEBUG", None) or os.getenv("DG_JIT_PRINT_NVCC_COMMAND", False):
-        print(f"Compiling JIT runtime {name} with command {command}")
-    return_code = subprocess.check_call(command)
-    assert return_code == 0, f"Failed to compile {src_path}"
+    @staticmethod
+    def include_dirs() -> List[str]:
+        return [get_jit_include_dir()]
 
-    # Interleave FFMA reuse
-    if enable_sass_opt:
-        interleave_ffma.process(tmp_so_path)
+    @classmethod
+    def build(cls, name: str, runtime_cls: Type[Runtime], kwargs: Dict[str, Any] = None) -> Runtime:
+        code = runtime_cls.generate(kwargs)
+        # Compiler flags
+        flags = cls.flags()
 
-    # Atomic replace SO file
-    os.replace(tmp_so_path, so_path)
+        # Build signature
+        enable_sass_opt = cls.__version__() <= (12, 8) and not int(os.getenv('DG_JIT_DISABLE_FFMA_INTERLEAVE', 0))
+        signature = f'{name}$${get_deep_gemm_version()}$${cls.signature()}$${flags}$${enable_sass_opt}$${code}'
+        name = f"kernel.{name}.{hash_to_hex(signature)}"
+        path = os.path.join(get_cache_dir(), name)
 
-    # Put cache and return
-    runtime_cache[path] = Runtime(path)
-    return runtime_cache[path]
+        # Check runtime cache or file system hit
+        global runtime_cache
+        cached_runtime = runtime_cache.get(path, runtime_cls, name, kwargs)
+        if cached_runtime is not None:
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                print(f"Using cached JIT runtime {name} during build")
+            return cached_runtime
+
+        # Compile into a temporary CU file
+        os.makedirs(path, exist_ok=True)
+        cubin_path = os.path.join(path, "kernel.cubin")
+        tmp_cubin_path = os.path.join(make_tmp_dir(), f"nvcc.tmp.{str(uuid.uuid4())}.{hash_to_hex(cubin_path)}.cubin")
+
+        start_time = time.time()
+        cls.compile(name, code, tmp_cubin_path)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(f"Compilation of JIT runtime {name} took {elapsed_time:.2f} seconds.")
+
+        # Interleave FFMA reuse
+        if enable_sass_opt:
+            interleave_ffma.process(tmp_cubin_path)
+
+        # Atomic replace files
+        os.replace(tmp_cubin_path, cubin_path)
+
+        # Put cache and return
+        runtime = runtime_cache.get(path, runtime_cls, name, kwargs, force_enable_cache=True)
+        assert runtime is not None
+        return runtime
+
+
+class NVCCCompiler(Compiler):
+    @staticmethod
+    def __version__() -> Tuple[int, int]:
+        _, version = get_nvcc_compiler()
+        major, minor = map(int, version.split("."))
+        return major, minor
+
+    @classmethod
+    def signature(cls) -> str:
+        return f"{get_nvcc_compiler()[0]}+{cls.__version__()}"
+
+    @classmethod
+    def flags(cls) -> List[str]:
+        cxx_flags = ["-fPIC", "-O3", "-fconcepts", "-Wno-deprecated-declarations", "-Wno-abi"]
+        return [
+            *super().flags(),
+            *[f"-I{d}" for d in cls.include_dirs()],
+            "-gencode=arch=compute_90a,code=sm_90a",
+            "-cubin",
+            "-O3",
+            "--expt-relaxed-constexpr",
+            "--expt-extended-lambda",
+            f"--compiler-options={','.join(cxx_flags)}",
+        ]
+
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        # Write the code
+        path = os.path.join(get_cache_dir(), name)
+        src_path = os.path.join(path, "kernel.cu")
+        put(src_path, code)
+        command = [get_nvcc_compiler()[0], src_path, "-o", target_path, *cls.flags()]
+        if int(os.getenv("DG_JIT_DEBUG", 0)) or int(os.getenv("DG_JIT_PRINT_COMPILER_COMMAND", 0)):
+            print(f"Compiling JIT runtime {name} with command {command}")
+
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.returncode != 0:
+            print(f"NVCC compilation failed: stdout: {result.stdout}, stderr: {result.stderr}")
+            assert False, f"Failed to compile {src_path}"
+
+
+class NVRTCCompiler(Compiler):
+    @staticmethod
+    def __version__() -> Tuple[int, int]:
+        res, major, minor = nvrtc.nvrtcVersion()
+        if res != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            # Failed to get the actual NVRTC version, use cuda-bindings version instead
+            major, minor = map(int, cuda.bindings.__version__.split(".")[:2])
+        return major, minor
+
+    @classmethod
+    def signature(cls) -> str:
+        return f"nvrtc+{cls.__version__()}"
+
+    @staticmethod
+    def include_dirs() -> List[str]:
+        if CUDA_HOME is None:
+            raise RuntimeError("CUDA_HOME is required for NVRTC compilation")
+        return [get_jit_include_dir(), os.path.join(CUDA_HOME, "include")]
+
+    @classmethod
+    def flags(cls) -> List[str]:
+        flags = [
+            *super().flags(),
+            *[f"-I{d}" for d in cls.include_dirs()],
+            "--gpu-architecture=sm_90a",
+            "-default-device",
+        ]
+        # NOTES: PCH is vital for compilation speed
+        if cls.__version__() >= (12, 8):
+            flags += ["--pch"]
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                flags += ["--pch-verbose=true"]
+        return flags
+
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        # Create program
+        code_bytes = bytes(code, "utf-8")
+        result, program = nvrtc.nvrtcCreateProgram(code_bytes, bytes(name, "utf-8"), 0, [], [])
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to create program: {result}"
+
+        # Compile
+        options = [bytes(flag, "utf-8") for flag in cls.flags()]
+        if int(os.getenv("DG_JIT_DEBUG", 0)) or int(os.getenv("DG_JIT_PRINT_COMPILER_COMMAND", 0)):
+            print(f"Compiling JIT runtime {name} with options: {options}")
+        compile_result = nvrtc.nvrtcCompileProgram(program, len(options), options)[0]
+
+        # Print compiler log
+        if int(os.getenv("DG_JIT_DEBUG", 0)) or compile_result != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            result, log_size = nvrtc.nvrtcGetProgramLogSize(program)
+            assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to get program log size: {result}"
+
+            log_bytes = bytes(log_size)
+            result = nvrtc.nvrtcGetProgramLog(program, log_bytes)[0]
+            assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to get program log: {result}"
+            print(f"Compiler log: {log_bytes.decode('utf-8')}")
+
+        # Exit if failed
+        assert compile_result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to compile program: {compile_result}"
+
+        # Create CUBIN
+        result, cubin_size = nvrtc.nvrtcGetCUBINSize(program)
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to get CUBIN size: {result}"
+        cubin_bytes = bytes(cubin_size)
+        result = nvrtc.nvrtcGetCUBIN(program, cubin_bytes)[0]
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, f"Failed to get CUBIN: {result}"
+
+        # Write into the file system
+        put(target_path, cubin_bytes)
+
+        # Destroy handler
+        assert (
+            nvrtc.nvrtcDestroyProgram(program)[0] == nvrtc.nvrtcResult.NVRTC_SUCCESS
+        ), f"Failed to destroy program: {result}"
+
+
+def build(name: str, runtime_cls: Type[Runtime], kwargs: Dict[str, Any] = None) -> Runtime:
+    compiler_cls = NVRTCCompiler if int(os.getenv("DG_JIT_USE_NVRTC", 0)) else NVCCCompiler
+    return compiler_cls.build(name, runtime_cls, kwargs)

ops/csrc/fp8/deep_gemm/jit/interleave_ffma.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 import argparse
 import mmap
@@ -22,11 +22,10 @@ import os
 import re
 import subprocess
 
-from ..utils import get_cuda_home
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
 
 
 def run_cuobjdump(file_path):
-    CUDA_HOME = get_cuda_home()
     command = [f"{CUDA_HOME}/bin/cuobjdump", "-sass", file_path]
     result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
     assert result.returncode == 0
@@ -57,7 +56,7 @@ def extract_ffma(sass):
                 collected.append((f"{arch_name}::{func_name}", current))
             current = []
 
-    if os.getenv("DG_PRINT_REG_REUSE", None):
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
         print(f"Found {len(collected)} FFMA segments")
     return collected
 
@@ -93,7 +92,7 @@ def parse_registers(line):
 
 
 def modify_segment(m, name, ffma_lines):
-    num_lines = len(ffma_lines)
+    num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2
     assert num_lines % 2 == 0
 
     le_bytes, new_le_bytes = [], []
@@ -120,7 +119,7 @@ def modify_segment(m, name, ffma_lines):
         dst_reg_set.add(dst_reg)
         new_le_bytes.append(low_hex.to_bytes(8, "little") + high_hex.to_bytes(8, "little"))
         last_reused, last_dst_reg = reused, dst_reg
-    if os.getenv("DG_PRINT_REG_REUSE", None):
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
         print(f" > segment `{name}` new reused list ({num_changed} changed): {reused_list}")
 
     # Find the offset
@@ -138,7 +137,7 @@ def modify_segment(m, name, ffma_lines):
 
 
 def process(path):
-    if os.getenv("DG_PRINT_REG_REUSE", None):
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
         print(f"Processing {path}")
     output = run_cuobjdump(path)
     segments = extract_ffma(output)

ops/csrc/fp8/deep_gemm/jit/runtime.py

@@ -12,26 +12,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 import ctypes
+import enum
 import os
-from typing import Optional
+import subprocess
+import time
+from typing import Any, Dict, Optional, Type
 
+import cuda.bindings.driver as cbd
 import paddle
-from paddle import Tensor
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
 
-from .template import map_ctype
+
+def get_num_math_warpgroups(block_m: int) -> int:
+    return 1 if block_m == 64 else 2
+
+
+def get_num_threads_per_sm(num_tma_threads: int, num_math_threads_per_group: int, block_m: int) -> int:
+    assert num_math_threads_per_group == 128, "Only support 128 threads per math group"
+    return get_num_math_warpgroups(block_m) * num_math_threads_per_group + num_tma_threads
+
+
+class GemmType(enum.Enum):
+    Normal = 0
+    GroupedContiguous = 1
+    GroupedMasked = 2
+
+    def __str__(self) -> str:
+        return {
+            0: "Normal",
+            1: "GroupedContiguous",
+            2: "GroupedMasked",
+        }[self.value]
 
 
 class Runtime:
     def __init__(self, path: str) -> None:
         self.path = path
         self.lib = None
-        self.args = None
-
+        self.kernel = None
         assert self.is_path_valid(self.path)
 
     @staticmethod
@@ -41,46 +64,314 @@ class Runtime:
             return False
 
         # Contains all necessary files
-        files = ["kernel.cu", "kernel.args", "kernel.so"]
+        files = ["kernel.cubin"]
         return all(os.path.exists(os.path.join(path, file)) for file in files)
 
-    def __call__(self, *args) -> int:
-        # Load SO file
-        if self.lib is None or self.args is None:
-            self.lib = ctypes.CDLL(os.path.join(self.path, "kernel.so"))
-            with open(os.path.join(self.path, "kernel.args"), "r") as f:
-                self.args = eval(f.read())
+    @staticmethod
+    def generate(kwargs: Dict[str, Any]) -> str:
+        raise NotImplementedError
 
-        # Check args and launch
-        assert len(args) == len(self.args), f"Expected {len(self.args)} arguments, got {len(args)}"
-        cargs = []
-        for arg, (name, dtype) in zip(args, self.args):
-            if isinstance(arg, Tensor):
-                assert arg.dtype == dtype, f"Expected tensor dtype `{dtype}` for `{name}`, got `{arg.dtype}`"
-            else:
-                assert isinstance(arg, dtype), f"Expected built-in type `{dtype}` for `{name}`, got `{type(arg)}`"
-            cargs.append(map_ctype(arg))
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult:
+        raise NotImplementedError
 
-        return_code = ctypes.c_int(0)
-        self.lib.launch(*cargs, ctypes.byref(return_code))
-        return return_code.value
+    def __call__(self, **kwargs) -> cbd.CUresult:
+        # Load CUBIN
+        if self.kernel is None:
+            start_time = time.time_ns()
+
+            # Load CUBIN
+            path = bytes(os.path.join(self.path, "kernel.cubin"), "utf-8")
+            result, self.lib = cbd.cuLibraryLoadFromFile(path, [], [], 0, [], [], 0)
+            assert result == cbd.CUresult.CUDA_SUCCESS, f"Failed to load library: {result}"
+
+            # Extract the kernel name
+            # TODO: use `cuda-bindings` API to do this (requires at least 12.8)
+            command = [f"{CUDA_HOME}/bin/cuobjdump", "-symbols", path]
+            result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+            assert result.returncode == 0
+            illegal_names = ["vprintf", "__instantiate_kernel", "__internal", "__assertfail"]
+            check_illegal = lambda line: any([name in line for name in illegal_names])
+            kernel_names = [
+                line.split()[-1]
+                for line in result.stdout.splitlines()
+                if line.startswith("STT_FUNC") and not check_illegal(line)
+            ]
+            assert len(kernel_names) == 1, f"Too many kernels in the library: {kernel_names}"
+
+            # Load kernel from the library
+            result, self.kernel = cbd.cuLibraryGetKernel(self.lib, bytes(kernel_names[0], encoding="utf-8"))
+            assert result == cbd.CUresult.CUDA_SUCCESS, f"Failed to load kernel: {result}"
+
+            end_time = time.time_ns()
+            elapsed_time = (end_time - start_time) / 1e6
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                print(f"Loading JIT runtime {self.path} took {elapsed_time:.2f} ms.")
+
+        # noinspection PyArgumentList
+        return self.launch(self.kernel, kwargs)
+
+    def __del__(self) -> None:
+        if self.lib is not None:
+            res = cbd.cuLibraryUnload(self.lib)[0]
+            if res != cbd.CUresult.CUDA_SUCCESS:
+                raise Exception(f"Failed to unload library {self.path}: {res}")
 
 
 class RuntimeCache:
     def __init__(self) -> None:
         self.cache = {}
 
-    def __getitem__(self, path: str) -> Optional[Runtime]:
+    def __setitem__(self, path: str, runtime: Runtime) -> None:
+        self.cache[path] = runtime
+
+    def get(
+        self,
+        path: str,
+        runtime_cls: Type[Runtime],
+        name: str = "",
+        kwargs: Dict[str, Any] = None,
+        force_enable_cache: bool = False,
+    ) -> Optional[Runtime]:
         # In Python runtime
         if path in self.cache:
             return self.cache[path]
 
         # Already compiled
-        if os.path.exists(path) and Runtime.is_path_valid(path):
-            runtime = Runtime(path)
+        use_cache = force_enable_cache or not int(os.getenv("DG_JIT_DISABLE_CACHE", 0))
+        if use_cache and os.path.exists(path) and Runtime.is_path_valid(path):
+            # Print heuristic for the first time
+            if name and (int(os.getenv("DG_JIT_DEBUG", 0)) or int(os.getenv("DG_PRINT_CONFIGS", 0))):
+                simplified_kwargs = dict()
+                for key, value in kwargs.items() if kwargs is not None else dict().items():
+                    value = f"paddle.Tensor<{value.dtype}>" if isinstance(value, paddle.Tensor) else value
+                    value = "cuda.bindings.driver.CUtensorMap" if isinstance(value, cbd.CUtensorMap) else value
+                    simplified_kwargs[key] = value
+                print(f"Put kernel {name} with {simplified_kwargs} into runtime cache")
+
+            runtime = runtime_cls(path)
             self.cache[path] = runtime
             return runtime
         return None
 
-    def __setitem__(self, path, runtime) -> None:
-        self.cache[path] = runtime
+
+def get_cache_key(kwargs, num_tma_threads, num_math_threads_per_group):
+    key_params = {
+        "NUM_TMA_MULTICAST": kwargs["NUM_TMA_MULTICAST"],
+        "NUM_SMS": kwargs["NUM_SMS"],
+        "BLOCK_M": kwargs["BLOCK_M"],
+        "SMEM_SIZE": kwargs["SMEM_SIZE"],
+        "STREAM": kwargs["STREAM"],
+        "num_tma_threads": num_tma_threads,
+        "num_math_threads_per_group": num_math_threads_per_group,
+    }
+    return hash(frozenset(key_params.items()))
+
+
+class KernelLaunchCache:
+    def __init__(self):
+        self.config_cache = {}
+        self.attr_cache = {}
+
+    @staticmethod
+    def create_attr(kwargs):
+        """Creates and caches property objects"""
+        attr_val = cbd.CUlaunchAttributeValue()
+        attr_val.clusterDim.x = kwargs["NUM_TMA_MULTICAST"]
+        attr_val.clusterDim.y = 1
+        attr_val.clusterDim.z = 1
+        attr = cbd.CUlaunchAttribute()
+        attr.id = cbd.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
+        attr.value = attr_val
+        return attr
+
+    @staticmethod
+    def create_config(kwargs, num_tma_threads, num_math_threads_per_group, attr):
+        """Creates and caches configuration objects"""
+        config = cbd.CUlaunchConfig()
+        config.numAttrs = 1
+        config.attrs = [attr]
+        config.gridDimX = kwargs["NUM_SMS"]
+        config.gridDimY = 1
+        config.gridDimZ = 1
+        config.blockDimX = get_num_threads_per_sm(num_tma_threads, num_math_threads_per_group, kwargs["BLOCK_M"])
+        config.blockDimY = 1
+        config.blockDimZ = 1
+        config.sharedMemBytes = kwargs["SMEM_SIZE"]
+        config.hStream = kwargs["STREAM"]
+        return config
+
+    def get_launch_config(self, kwargs, num_tma_threads, num_math_threads_per_group):
+        """Retrieves cached config or creates new instance"""
+        cache_key = get_cache_key(kwargs, num_tma_threads, num_math_threads_per_group)
+
+        if cache_key not in self.config_cache:
+            # 先检查属性缓存
+            attr_key = (kwargs["NUM_TMA_MULTICAST"],)
+            if attr_key not in self.attr_cache:
+                self.attr_cache[attr_key] = self.create_attr(kwargs)
+
+            # 创建新配置
+            config = self.create_config(kwargs, num_tma_threads, num_math_threads_per_group, self.attr_cache[attr_key])
+            self.config_cache[cache_key] = config
+
+        return self.config_cache[cache_key]
+
+
+launch_cache = KernelLaunchCache()
+
+
+class FP8GemmRuntime(Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    @staticmethod
+    def generate(kwargs: Dict[str, Any]) -> str:
+        code = f"""
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#include <string>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#include <deep_gemm/fp8_gemm.cuh>
+
+using namespace deep_gemm;
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&fp8_gemm_kernel<
+        {kwargs['N']},
+        {kwargs['K']},
+        {kwargs['BLOCK_M']},
+        {kwargs['BLOCK_N']},
+        {kwargs['BLOCK_K']},
+        {kwargs['BLOCK_N_PADDING']},
+        {kwargs['SWIZZLE_D_MODE']},
+        {kwargs['NUM_GROUPS']},
+        {kwargs['NUM_STAGES']},
+        {kwargs['NUM_TMA_THREADS']},
+        {kwargs['NUM_MATH_THREADS_PER_GROUP']},
+        {kwargs['NUM_TMA_MULTICAST']},
+        {'true' if kwargs['IS_TMA_MULTICAST_ON_A'] else 'false'},
+        GemmType::{kwargs['GEMM_TYPE']}
+      >);
+}};
+"""
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(f"Generated FP8 GEMM code:\n{code}")
+        return code
+
+    # noinspection PyMethodOverriding
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult:
+        num_tma_threads = 128
+        num_math_threads_per_group = 128
+
+        result = cbd.cuKernelSetAttribute(
+            cbd.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+            kwargs["SMEM_SIZE"],
+            kernel,
+            cbd.CUdevice(kwargs["DEVICE_INDEX"]),
+        )[0]
+        assert result == cbd.CUresult.CUDA_SUCCESS, f"Failed to set max dynamic shared memory size: {result}"
+        config = launch_cache.get_launch_config(kwargs, num_tma_threads, num_math_threads_per_group)
+
+        arg_values = (
+            kwargs["SCALES_B"].data_ptr(),
+            kwargs["GROUPED_LAYOUT"].data_ptr(),
+            kwargs["M"],
+            kwargs["TENSOR_MAP_A"],
+            kwargs["TENSOR_MAP_B"],
+            kwargs["TENSOR_MAP_SCALES_A"],
+            kwargs["TENSOR_MAP_D"],
+        )
+        arg_types = (
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+            None,
+            None,
+            None,
+            None,
+        )
+        ret = cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)
+        return ret
+
+
+class FP8WGradGemmRuntime(Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    @staticmethod
+    def generate(kwargs: Dict[str, Any]) -> str:
+        code = f"""
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#include <string>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#include <deep_gemm/fp8_wgrad_gemm.cuh>
+
+using namespace deep_gemm;
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&fp8_wgrad_gemm_kernel<
+        {kwargs['M']},
+        {kwargs['N']},
+        {kwargs['BLOCK_M']},
+        {kwargs['BLOCK_N']},
+        {kwargs['BLOCK_K']},
+        {kwargs['NUM_STAGES']},
+        {kwargs['NUM_LAST_STAGES']},
+        {kwargs['NUM_TMA_THREADS']},
+        {kwargs['NUM_MATH_THREADS_PER_GROUP']},
+        {kwargs['NUM_TMA_MULTICAST']},
+        {'true' if kwargs['IS_TMA_MULTICAST_ON_A'] else 'false'}
+      >);
+}};
+"""
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(f"Generated FP8 WGrad GEMM code:\n{code}")
+        return code
+
+    # noinspection PyMethodOverriding
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult:
+        num_tma_threads = 128
+        num_math_threads_per_group = 128
+
+        result = cbd.cuKernelSetAttribute(
+            cbd.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+            kwargs["SMEM_SIZE"],
+            kernel,
+            cbd.CUdevice(kwargs["DEVICE_INDEX"]),
+        )[0]
+        assert result == cbd.CUresult.CUDA_SUCCESS, f"Failed to set max dynamic shared memory size: {result}"
+        config = launch_cache.get_launch_config(kwargs, num_tma_threads, num_math_threads_per_group)
+
+        arg_values = (
+            kwargs["K"],
+            kwargs["TENSOR_MAP_A"],
+            kwargs["TENSOR_MAP_B"],
+            kwargs["TENSOR_MAP_SCALES_A"],
+            kwargs["TENSOR_MAP_SCALES_B"],
+            kwargs["TENSOR_MAP_D"],
+        )
+        arg_types = (
+            ctypes.c_uint32,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+        return cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)

ops/csrc/fp8/deep_gemm/jit/template.py

@@ -1,122 +0,0 @@
-# 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.
-
-# The file has been adapted from DeepSeek DeepEP project
-# Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
-
-import copy
-import ctypes
-import os
-from typing import Any, Dict, Iterable, Tuple
-
-import paddle
-from paddle import Tensor
-
-# Name map for Python `eval`
-typename_map: Dict[Any, str] = {
-    **{t: t.__name__ for t in (bool, int, float)},
-    paddle.int32: "paddle.int32",
-    paddle.float32: "paddle.float32",
-    paddle.bfloat16: "paddle.bfloat16",
-    paddle.float8_e4m3fn: "paddle.float8_e4m3fn",
-    paddle.device.cuda.Stream: "paddle.device.cuda.Stream",
-}
-# `ctype` map for Python casting
-ctype_map: Dict[Any, Any] = {
-    **{t: getattr(ctypes, f"c_{t.__name__}") for t in (bool, int, float)},
-    **{
-        t: ctypes.c_void_p
-        for t in (paddle.int32, paddle.float32, paddle.bfloat16, paddle.float8_e4m3fn, paddle.device.cuda.Stream)
-    },
-}
-
-
-# Type map for both Python API and source code usages
-genc_map = {
-    bool: ("bool", "bool"),
-    int: ("int", "int"),
-    float: ("float", "float"),
-    paddle.int32: ("void*", "int*"),
-    paddle.float32: ("void*", "float*"),
-    paddle.bfloat16: ("void*", "__nv_bfloat16*"),
-    paddle.float8_e4m3fn: ("void*", "__nv_fp8_e4m3*"),
-    paddle.device.cuda.Stream: ("void*", "cudaStream_t"),
-}
-
-
-def map_ctype(value: Any) -> Any:
-    ctype = ctype_map[value.dtype if isinstance(value, Tensor) else type(value)]
-    if isinstance(value, Tensor):
-        return ctype(value.data_ptr())
-    if isinstance(value, paddle.device.cuda.Stream):
-        return ctype(value.cuda_stream)
-    return ctype(value)
-
-
-def cpp_format(template: str, keys: Dict[str, Any]) -> str:
-    # We don't use `str.format` because it's not safe for C++ {} braces
-    new_template = copy.deepcopy(template)
-    for key, value in keys.items():
-        new_template = new_template.replace(f"{{{key}}}", f"{value}")
-    return new_template
-
-
-def generate(includes: Iterable[str], arg_defs: Iterable[Tuple], body: str) -> str:
-    # Common prefix
-    code = "// DeepGEMM auto-generated JIT CUDA source file\n\n"
-
-    # Includes
-    preload_sys_includes = ["<cuda.h>", "<cuda_fp8.h>", "<cuda_runtime.h>", "<iostream>"]
-    preload_package_includes = ['"cutlass/cutlass.h"']
-
-    assert isinstance(includes, list) or isinstance(includes, tuple)
-    sys_includes = sorted(
-        list(set(preload_sys_includes + [include for include in includes if include.startswith("<")]))
-    )
-    package_includes = sorted(
-        list(set(preload_package_includes + [include for include in includes if include.startswith('"')]))
-    )
-    code += "\n".join(f"#include {include}" for include in sys_includes) + "\n\n"
-    code += "\n".join(f"#include {include}" for include in package_includes) + "\n\n"
-
-    # Function signature
-    raw = "__raw_"
-    get_def = lambda n, t: f"{genc_map[t][0]} " + (raw if genc_map[t][0] != genc_map[t][1] else "") + n
-    code += f'extern "C" void launch('
-    code += ", ".join(
-        [get_def(*arg_def) for arg_def in arg_defs]
-        + [
-            "int& __return_code",
-        ]
-    )
-    code += ") {\n"
-
-    # Cast raw types
-    code += "    // Cast raw types (if needed)\n"
-    for arg_name, arg_type in arg_defs:
-        if genc_map[arg_type][0] != genc_map[arg_type][1]:
-            code += f"    auto {arg_name} = reinterpret_cast<{genc_map[arg_type][1]}>({raw}{arg_name});\n"
-
-    # Function body
-    code += "\n".join([(("    " if line else "") + line) for line in body.split("\n")])
-
-    # End the function
-    code += "}\n\n"
-
-    # Debug print
-    if os.getenv("DG_JIT_DEBUG", None):
-        print(f"Generated code:\n{code}")
-
-    return code

ops/csrc/fp8/deep_gemm/jit_kernels/__init__.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 from .gemm import gemm_fp8_fp8_bf16_nt
 from .m_grouped_gemm import (
@@ -28,3 +28,4 @@ from .utils import (
     get_num_sms,
     set_num_sms,
 )
+from .wgrad_gemm import k_grouped_wgrad_gemm_fp8_fp8_fp32_nt, wgrad_gemm_fp8_fp8_fp32_nt

ops/csrc/fp8/deep_gemm/jit_kernels/gemm.py

@@ -12,17 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
-import functools
+import math
+from functools import lru_cache
 from typing import Tuple
+import os
 
 import paddle
-from paddle import Tensor
 
-from .tuner import jit_tuner
+from ..jit import FP8GemmRuntime, build
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
 from .utils import (
     ceil_div,
     get_col_major_tma_aligned_tensor,
@@ -30,45 +38,56 @@ from .utils import (
     get_num_sms,
 )
 
-# C++ code templates
-includes = ('"deep_gemm/fp8_gemm.cuh"',)
-template = """
-using namespace deep_gemm;
-
-// Templated args from Python JIT call
-constexpr auto N = {N}, K = {K};
-constexpr auto BLOCK_M = {BLOCK_M};
-constexpr auto BLOCK_N = {BLOCK_N};
-constexpr auto kNumStages = {NUM_STAGES};
-constexpr auto kNumTMAMulticast = {NUM_TMA_MULTICAST};
-
-// Make a templated GEMM
-using GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, 1, kNumStages, kNumTMAMulticast, GemmType::Normal>;
-
-// Launch kernel
-auto tma_a_desc = GemmType::make_2d_tma_a_desc(lhs, m);
-auto tma_b_desc = GemmType::make_2d_tma_b_desc(rhs);
-auto tma_scales_a_desc = GemmType::make_2d_tma_scales_a_desc(lhs_scales, m);
-auto tma_d_desc = GemmType::make_2d_tma_d_desc(out, m);
-GemmType::run(out, rhs_scales, nullptr,
-              m,
-              tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
-              stream, num_sms, smem_size);
-"""
+global_empty_tensor = paddle.empty([0], dtype=paddle.int32)
+# Todo: Use default stream to accelerate CPU time. Optimize here if use multistream to launch gemm kernel.
+global_stream = paddle.device.current_stream().stream_base.cuda_stream
 
 
-def is_tma_multicast_legal(n: int, block_n: int, num_tma_multicast: int, num_sms: int) -> bool:
-    if num_tma_multicast == 1:
-        return True
-    return (n % (block_n * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
+def is_tma_multicast_legal(
+    shape_dim: int, block_dim: int, num_tma_multicast: int, num_sms: int, require_divisible: bool = False
+) -> bool:
+    divisible = ceil_div(shape_dim, block_dim) % num_tma_multicast == 0 or not require_divisible
+    return divisible and num_sms % num_tma_multicast == 0
 
 
-def get_smem_size(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128) -> int:
-    smem_d = block_m * block_n * 2
+def get_swizzle_mode(block_n: int) -> int:
+    elem_size = 2
+    for mode_bytes in (128, 64, 32):
+        if (block_n * elem_size) % mode_bytes == 0:
+            return mode_bytes
+    return 0
+
+
+def get_block_n_padding_for_smem_d(block_n: int) -> int:
+    # NOTES: padding is for solving bank conflicts, but wastes shared memory space
+    elem_size, requirement = 2, (4, 8)
+    bank_stride = (block_n * elem_size) // 4
+    padding = (requirement[0] - bank_stride) % requirement[1]
+    return (((padding + requirement[1]) if padding < 0 else padding) * 4) // elem_size
+
+
+def get_smem_config(
+    num_stages: int,
+    k: int,
+    block_m: int,
+    block_n: int,
+    block_k: int = 128,
+    is_fp32_out: bool = False,
+    is_wgrad: bool = False,
+) -> Tuple[int, int, int]:
+    assert block_k == 128
+
+    # Try swizzle first, as it does not waste shared memory
+    swizzle_mode = get_swizzle_mode(block_n)
+    block_n_padding = get_block_n_padding_for_smem_d(block_n) if swizzle_mode == 0 else 0
+
+    # NOTES: `scales_b` in a total manner or per-stage manner
+    smem_d = block_m * (block_n + block_n_padding) * (4 if is_fp32_out else 2)
     smem_a_per_stage = block_m * block_k
     smem_scales_a_per_stage = block_m * 4
     smem_b_per_stage = block_n * block_k
-    smem_scales_b = ceil_div(k, block_k) * 4
+    smem_scales_b_per_stage = ceil_div(block_n * 4, block_k) * block_k if is_wgrad else 0
+    smem_scales_b = ceil_div(k, block_k) * 4 if not is_wgrad else 0
     smem_barrier = num_stages * 8 * 2
 
     smem_size = 0
@@ -76,20 +95,50 @@ def get_smem_size(num_stages: int, k: int, block_m: int, block_n: int, block_k:
     smem_size += num_stages * smem_a_per_stage
     smem_size += num_stages * smem_scales_a_per_stage
     smem_size += num_stages * smem_b_per_stage
+    smem_size += num_stages * smem_scales_b_per_stage
     smem_size += ceil_div(smem_scales_b * (1 if block_k % block_n == 0 else 2), 8) * 8
     smem_size += smem_barrier
-    return smem_size
+
+    # Swizzle and padding are not compatible
+    assert int(swizzle_mode > 0) + int(block_n_padding > 0) <= 1
+
+    return smem_size, swizzle_mode, block_n_padding
 
 
+@lru_cache(maxsize=None)
 def get_best_configs(
-    m: int, n: int, k: int, num_groups: int, num_sms: int, is_grouped_contiguous: bool = False
-) -> Tuple[int, int, int, int, int]:
+    m: int,
+    n: int,
+    k: int,
+    num_groups: int,
+    num_sms: int,
+    is_grouped_contiguous: bool = False,
+    is_grouped_masked: bool = False,
+    is_fp32_out: bool = False,
+    is_wgrad: bool = False,
+) -> Tuple[int, int, int, int, Tuple[int, bool], Tuple[int, int, int]]:
     if not is_grouped_contiguous:
-        # TODO: for some cases, smaller M block is better, add them into tuning space
-        block_ms = (64 if m <= 64 else 128,)
+        block_ms = (
+            64,
+            128,
+        ) + ((256,) if not is_fp32_out else ())
     else:
         block_ms = (get_m_alignment_for_contiguous_layout(),)
-    block_ns = tuple(range(16, 129, 8))
+    block_ns = tuple(range(16, 129, 8)) + (
+        (
+            136,
+            152,
+        )
+        if is_wgrad
+        else (
+            144,
+            160,
+        )
+    )
+
+    # Avoid bank conflicts for FP32 output
+    if is_fp32_out:
+        block_ns = [x for x in block_ns if x % 16 == 8]
 
     fix_wave_saturate = lambda x: num_sms if x == 0 else x
     get_num_waves = lambda bm, bn: (ceil_div(ceil_div(m, bm) * ceil_div(n, bn) * num_groups, num_sms) if bm else None)
@@ -98,7 +147,8 @@ def get_best_configs(
     # Decide block sizes by waves
     best_block_m, best_block_n = None, None
     for block_m in block_ms:
-        for block_n in block_ns:
+        # NOTES: the block sizes cannot be too large, so at least one dim less than 128
+        for block_n in filter(lambda bn: block_m <= 128 or bn <= 128, block_ns):
             success = False
             num_waves, best_num_waves = get_num_waves(block_m, block_n), get_num_waves(best_block_m, best_block_n)
             if best_block_m is None or best_block_n is None:
@@ -109,79 +159,75 @@ def get_best_configs(
                 # Check last wave utilization
                 util = get_last_wave_util(block_m, block_n)
                 best_util = get_last_wave_util(best_block_m, best_block_n)
-                success = util > best_util or (
-                    util == best_util
-                    and (block_m > best_block_m or (block_m == best_block_m and block_n < best_block_n))
-                )
+                success = util > best_util
+                if util == best_util:
+                    # Case 1: same `block_m`, smaller `block_n` (wasted)
+                    success |= block_m == best_block_m and block_n < best_block_n
+                    # Case 2: same `block_n`, smaller `block_m` (wasted)
+                    success |= block_n == best_block_n and block_m < best_block_m
+                    # Case 3: different for both `block_m` and `block_n`, `block_n` larger is better
+                    success |= block_m != best_block_m and block_n > best_block_n
             best_block_m, best_block_n = (block_m, block_n) if success else (best_block_m, best_block_n)
     assert best_block_m is not None and best_block_n is not None
 
     # Always pick the longest one
     # NOTES: for double B scales, the best number of stages may be reduced
-    best_num_stages, best_smem_size, sm90_capacity = None, None, 232448
-    for num_stages in (6, 5, 4) if 128 % best_block_n != 0 else (8, 7, 6, 5, 4):
-        best_smem_size = get_smem_size(num_stages, k, best_block_m, best_block_n)
-        if best_smem_size <= sm90_capacity:
+    best_num_stages, best_smem_config, sm90_capacity = None, None, 232448
+    stage_candidates = tuple(filter(lambda s: s <= max(k // 128, 1), (8, 7, 6, 5, 4, 3, 2, 1)))
+    if 128 % best_block_n != 0 and 128 // math.gcd(128, best_block_n) <= 4:
+        # Unrolling both stages and `num_former_iters` will cause large code size
+        stage_candidates = tuple(filter(lambda s: s <= max(k // 128, 1), (4, 3, 2, 1)))
+    for num_stages in stage_candidates:
+        best_smem_config = get_smem_config(
+            num_stages, k, best_block_m, best_block_n, is_fp32_out=is_fp32_out, is_wgrad=is_wgrad
+        )
+        if best_smem_config[0] <= sm90_capacity:
             best_num_stages = num_stages
             break
+    assert best_smem_config is not None
     assert best_num_stages is not None
 
-    # Decide the number of TMA multicast
-    best_num_tma_multicast = 1
-    if m >= 1024 and is_tma_multicast_legal(n, best_block_n, 2, num_sms) and num_groups == 1:
-        best_num_tma_multicast = 2
+    # Decide the number of TMA multicasts and whether broadcast on A
+    best_tma_multicast_config = (1, True)
 
-    return best_block_m, best_block_n, best_num_stages, best_num_tma_multicast, best_smem_size
+    # Try to multicast on the larger block side first
+    # NOTES: currently, grouped masked GEMM only supports multicast on A and requires the number of blocks in the N-direction to be even
+    is_multicast_legal = {
+        "A": is_tma_multicast_legal(n, best_block_n, 2, num_sms, is_grouped_masked),
+        "B": is_tma_multicast_legal(m, best_block_m, 2, num_sms) and not is_grouped_masked,
+    }
+    for i in ("A", "B") if best_block_m > best_block_n else ("B", "A"):
+        if m >= 512 and is_multicast_legal[i]:
+            best_tma_multicast_config = (2, i == "A")
+            break
+
+    # Recompute the minimal number of SMs required
+    # NOTES: less L2 cache usage and less GPU frequency drop
+    num_waves = get_num_waves(best_block_m, best_block_n)
+    num_min_sms = ceil_div(ceil_div(m, best_block_m) * ceil_div(n, best_block_n) * num_groups, num_waves)
+    num_min_sms = ceil_div(num_min_sms, best_tma_multicast_config[0]) * best_tma_multicast_config[0]
+    assert num_min_sms <= num_sms
+
+    return num_min_sms, best_block_m, best_block_n, best_num_stages, best_tma_multicast_config, best_smem_config
 
 
-@functools.lru_cache()
-def auto_tuning_with_compilation(m, n, k):
-    global includes, template
-    num_sms = get_num_sms()
-    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(m, n, k, 1, num_sms)
-    runtime = jit_tuner.compile_and_tune(
-        m,
-        n,
-        k,
-        name="gemm_fp8_fp8_bf16_nt",
-        keys={
-            "BLOCK_M": block_m,
-            "BLOCK_N": block_n,
-            "K": k,
-            "N": n,
-            "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": num_tma_multicast,
-        },
-        space=(),
-        includes=includes,
-        arg_defs=(
-            ("lhs", paddle.float8_e4m3fn),
-            ("lhs_scales", paddle.float32),
-            ("rhs", paddle.float8_e4m3fn),
-            ("rhs_scales", paddle.float32),
-            ("out", paddle.bfloat16),
-            ("m", int),
-            ("stream", paddle.device.cuda.Stream),
-            ("num_sms", int),
-            ("smem_size", int),
-        ),
-        template=template,
-    )
-    return runtime, num_sms, smem_size
-
-
-def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor) -> None:
+def gemm_fp8_fp8_bf16_nt(
+    lhs: Tuple[paddle.Tensor, paddle.Tensor], rhs: Tuple[paddle.Tensor, paddle.Tensor], out: paddle.Tensor, num_sms: int = None
+) -> None:
     """
-    Do a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
-    LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
-    RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
-        this function will do a transposing with a set of slow Paddle operations.
+    Perform a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, and output tensors must be contiguous in dimension 1, i.e., strides[1] = 1.
+        The strides[0] of LHS and RHS must be a multiple of 16, and the strides[0] of output must be a multiple of 8.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
 
     Arguments:
         lhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[m, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[m, ⌈k / 128⌉]`.
-        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[n, k]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[n, k]`,
              the second element is an FP32 128x128 scaling tensor for RHS of shape `[⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[m, n]`, representing the result.
     """
@@ -190,27 +236,74 @@ def gemm_fp8_fp8_bf16_nt(lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor],
     m, k = lhs.shape
     n, k_ = rhs.shape
     m_, n_ = out.shape
-    assert n % 64 == 0 and k % 128 == 0
 
     # Type and shape checks
     assert m == m_ and n == n_ and k == k_
     assert n > 0 and k > 0
-    assert lhs_scales.shape == [m, (k + 127) // 128]
-    assert rhs_scales.shape == [(n + 127) // 128, (k + 127) // 128]
+    assert lhs_scales.shape == [m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [ceil_div(n, 128), ceil_div(k, 128)]
     assert lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
     assert rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
     assert out.dtype == paddle.bfloat16
-    assert lhs.is_contiguous() and rhs.is_contiguous() and out.is_contiguous()
+    assert lhs.strides[1] == 1 and out.strides[1] == 1 and rhs.strides[1] == 1
 
-    # LHS scales must be transposed for TMA load, but not for RHS scales
-    # NOTES: `get_tma_aligned_lhs_scales` may launch a kernel if not processed by previous kernels
+    # LHS scales must be transposed for TMA loads, but not for RHS scales
+    # NOTES: `get_col_major_tma_aligned_tensor` may launch a kernel if not processed by previous kernels
     lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
     assert rhs_scales.is_contiguous()
 
     # Do nothing if `m` is zero
     if m == 0:
         return
-    runtime, num_sms, smem_size = auto_tuning_with_compilation(m, n, k)
-    args = (lhs, lhs_scales, rhs, rhs_scales, out, m, paddle.device.cuda.current_stream(), num_sms, smem_size)
-    # Run the kernel.
-    runtime(*args)
+
+    # K must be aligned to 128
+    aligned_k = ceil_div(k, 128) * 128
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(m, n, k, 1, num_sms)
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(f"Auto-tuned gemm_fp8_fp8_bf16_nt as num_sms={num_sms}, block_m={block_m}, block_n={block_n}")
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(GemmType.Normal, lhs, m, k, lhs.strides[0], block_m, block_k, 1)
+    tensor_map_b = make_2d_tma_b_desc(GemmType.Normal, rhs, n, k, rhs.strides[0], block_n, block_k, 1)
+    tensor_map_d = make_2d_tma_d_desc(GemmType.Normal, out, m, n, out.strides[0], block_m, block_n, 1, smem_config[1])
+    tensor_map_scales_a = make_2d_tma_scales_desc(GemmType.Normal, lhs_scales, m, k, block_m, block_k, 1)
+
+    kwargs = {
+        # Templated arguments
+        "GEMM_TYPE": GemmType.Normal,
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": aligned_k,
+        "NUM_GROUPS": 1,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": global_empty_tensor,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": global_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)

ops/csrc/fp8/deep_gemm/jit_kernels/m_grouped_gemm.py

@@ -12,122 +12,72 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
-import functools
+from functools import reduce
 from typing import Tuple
+import os
 
 import paddle
-from paddle import Tensor
 
+from ..jit import FP8GemmRuntime, build
 from .gemm import get_best_configs
-from .tuner import jit_tuner
-from .utils import get_col_major_tma_aligned_tensor, get_num_sms
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
+from .utils import ceil_div, get_col_major_tma_aligned_tensor, get_num_sms
 
-# C++ code templates
-includes = ('"deep_gemm/fp8_gemm.cuh"',)
-template = """
-using namespace deep_gemm;
-
-// Templated args from Python JIT call
-constexpr auto N = {N}, K = {K};
-constexpr auto BLOCK_M = {BLOCK_M};
-constexpr auto BLOCK_N = {BLOCK_N};
-constexpr auto kNumStages = {NUM_STAGES};
-constexpr auto kNumTMAMulticast = {NUM_TMA_MULTICAST};
-
-// Make a templated grouped GEMM
-using GemmType = Gemm<N, K, BLOCK_M, BLOCK_N, 128, {NUM_GROUPS}, kNumStages, kNumTMAMulticast, GemmType::{GEMM_TYPE}>;
-
-// Launch kernel
-auto tma_a_desc = GemmType::make_2d_tma_a_desc(lhs, m);
-auto tma_b_desc = GemmType::make_2d_tma_b_desc(rhs);
-auto tma_scales_a_desc = GemmType::make_2d_tma_scales_a_desc(lhs_scales, m);
-auto tma_d_desc = GemmType::make_2d_tma_d_desc(out, m);
-GemmType::run(out, rhs_scales, grouped_layout,
-              m,
-              tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
-              stream, num_sms, smem_size);
-"""
-
-
-@functools.lru_cache()
-def auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms):
-    global includes, template
-    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(
-        m, n, k, 1, num_sms, is_grouped_contiguous=True
-    )
-    runtime = jit_tuner.compile_and_tune(
-        m,
-        n,
-        k,
-        name="m_grouped_gemm_fp8_fp8_bf16_nt",
-        keys={
-            "BLOCK_M": block_m,
-            "BLOCK_N": block_n,
-            "GEMM_TYPE": "GroupedContiguous",
-            "K": k,
-            "N": n,
-            "NUM_GROUPS": num_groups,
-            "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": num_tma_multicast,
-        },
-        space=(),
-        includes=includes,
-        arg_defs=(
-            ("lhs", paddle.float8_e4m3fn),
-            ("lhs_scales", paddle.float32),
-            ("rhs", paddle.float8_e4m3fn),
-            ("rhs_scales", paddle.float32),
-            ("out", paddle.bfloat16),
-            ("grouped_layout", paddle.int32),
-            ("m", int),
-            ("num_groups", int),
-            ("stream", paddle.device.cuda.Stream),
-            ("num_sms", int),
-            ("smem_size", int),
-        ),
-        template=template,
-    )
-    return runtime, num_sms, smem_size
+# Todo: Use default stream to accelerate CPU time. Optimize here if use multistream to launch gemm kernel.
+global_stream = paddle.device.current_stream().stream_base.cuda_stream
 
 
 def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
-    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, m_indices: Tensor
+    lhs: Tuple[paddle.Tensor, paddle.Tensor],
+    rhs: Tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    m_indices: paddle.Tensor,
+    num_sms: int = None,
 ) -> None:
     """
-    Do a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
-    LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
-    RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
-        this function will do a transposing with a set of slow Paddle operations.
-    On the M axis, inputs are grouped into several batches, of which batch sizes aligned to
-        `get_m_alignment_for_contiguous_layout()` (128).
+    Perform a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
+        On the M axis, inputs are grouped into several batches, of which batch sizes aligned to
+            `get_m_alignment_for_contiguous_layout()` (128).
 
     Arguments:
         lhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[m_sum, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[m_sum, ⌈k / 128⌉]`.
-        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[num_groups, n, k]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[num_groups, n, k]`,
              the second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[m_sum, n]`, representing the result.
-        m_indices: a tensor of shape `[m_sum]` with type `paddle.int32`.
-            `m_indices[i]` records the group which the j-th row of the LHS belong to,
+        m_indices: a tensor of shape `[m_sum]` with type `paddle.int`.
+            `m_indices[i]` records the group which the i-th row of the LHS belongs to,
             which means that the i-th row of the LHS matrix will be multiplied with `rhs[m_indices[i]]`.
             Values of `m_indices` in every-m-alignment-block must also be the same.
-            `-1` in this tensor indicates no RHS matrix selected, the kernel will skip the computation for that aligned block.
     """
     lhs, lhs_scales = lhs
     rhs, rhs_scales = rhs
     m, k = lhs.shape
     num_groups, n, k_ = rhs.shape
     m_, n_ = out.shape
-    m__ = m_indices.numel()
+    m_shape = m_indices.shape
+    m__ = reduce(lambda x, y: x * y, m_shape)
+
     # Type and shape checks
     assert m == m_ == m__ and k == k_ and n == n_
-    assert lhs_scales.shape == [m, (k + 127) // 128]
-    assert rhs_scales.shape == [num_groups, (n + 127) // 128, (k + 127) // 128]
+    assert lhs_scales.shape == [m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [num_groups, ceil_div(n, 128), ceil_div(k, 128)]
     assert lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
     assert rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
     assert out.dtype == paddle.bfloat16
@@ -142,44 +92,87 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
     # Do nothing if `m` is zero
     if m == 0:
         return
-    # Auto-tuning with compilation
-    global includes, template
-    num_sms = get_num_sms()
-    runtime, num_sms, smem_size = auto_tuning_with_compilation_grouped_gemm_contiguous(m, n, k, num_groups, num_sms)
 
-    args = (
-        lhs,
-        lhs_scales,
-        rhs,
-        rhs_scales,
-        out,
-        m_indices,
-        m,
-        num_groups,
-        paddle.device.cuda.current_stream(),
-        num_sms,
-        smem_size,
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(
+        m, n, k, 1, num_sms, is_grouped_contiguous=True
     )
-    runtime(*args)
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(f"Auto-tuned m_grouped_gemm_fp8_fp8_bf16_nt_contiguous as num_sms={num_sms}, block_m={block_m}, block_n={block_n}")
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(GemmType.GroupedContiguous, lhs, m, k, k, block_m, block_k, num_groups)
+    tensor_map_b = make_2d_tma_b_desc(GemmType.GroupedContiguous, rhs, n, k, k, block_n, block_k, num_groups)
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.GroupedContiguous, out, m, n, n, block_m, block_n, num_groups, smem_config[1]
+    )
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.GroupedContiguous, lhs_scales, m, k, block_m, block_k, num_groups
+    )
+
+    kwargs = {
+        # Templated arguments
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": k,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_GROUPS": num_groups,
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        "GEMM_TYPE": GemmType.GroupedContiguous,
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": m_indices,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": global_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("m_grouped_gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)
 
 
 def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
-    lhs: Tuple[Tensor, Tensor], rhs: Tuple[Tensor, Tensor], out: Tensor, masked_m: Tensor, expected_m: int
+    lhs: Tuple[paddle.Tensor, paddle.Tensor],
+    rhs: Tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    masked_m: paddle.Tensor,
+    expected_m: int,
+    num_sms: int = None,
 ) -> None:
     """
-    Do a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
-    LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
-    RHS and RHS scaling factors are required to be transposed.
-    The LHS scaling tensor requires TMA-aligned transposed format, if your input does not match the requirement,
-        this function will do a transposing with a set of slow Paddle operations.
-    Moreover, this alignment requirement is different with the contiguous-format kernel, as we require that each batch
-        should be separately transposed.
+    Perform a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
+        Moreover, this alignment requirement is different with the contiguous-format kernel, as we require that each batch
+            should be separately transposed.
 
     Arguments:
-        lhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[num_groups, m_max, k]`,
+        lhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[num_groups, m_max, k]`,
              the second element is an FP32 1x128 scaling tensor for LHS of shape `[num_groups, m_max, ⌈k / 128⌉]`.
-        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[num_groups, n, k]`.
-             the second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[num_groups, n, k]`.
+             The second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
         out: the BF16 output tensor of shape `[num_groups, m_max, n]`, representing the result.
         masked_m: a tensor of shape `[num_groups]`, `masked_m[i]` records actual rows of the `lhs[i]` matrix to compute
             in the i-th group.
@@ -191,14 +184,14 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
     num_groups, m, k = lhs.shape
     num_groups_, n, k_ = rhs.shape
     num_groups__, m_, n_ = out.shape
-    num_groups___ = masked_m.numel()
+    num_groups___ = masked_m.shape[0]
 
     # Type and shape checks
     assert num_groups == num_groups_ == num_groups__ == num_groups___
     assert m == m_ and n == n_ and k == k_
     assert expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0
-    assert lhs_scales.shape == [num_groups, m, (k + 127) // 128]
-    assert rhs_scales.shape == [num_groups, (n + 127) // 128, (k + 127) // 128]
+    assert lhs_scales.shape == [num_groups, m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [num_groups, ceil_div(n, 128), ceil_div(k, 128)]
     assert lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
     assert rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
     assert out.dtype == paddle.bfloat16
@@ -211,59 +204,62 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
     assert rhs_scales.is_contiguous()
 
     # Auto-tuning with compilation
-    global includes, template
-    num_sms = get_num_sms()
-    block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(
-        expected_m, n, k, num_groups, num_sms
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(
+        expected_m, n, k, num_groups, num_sms, is_grouped_masked=True
     )
-
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(f"Auto-tuned m_grouped_gemm_fp8_fp8_bf16_nt_masked as num_sms={num_sms}, block_m={block_m}, block_n={block_n}")
     # Extra checks for TMA store
     if num_groups > 1 and m > block_m:
         assert (
             m % block_m == 0
         ), f"For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})"
 
-    args = (
-        lhs,
-        lhs_scales,
-        rhs,
-        rhs_scales,
-        out,
-        masked_m,
-        m,
-        paddle.device.cuda.current_stream(),
-        num_sms,
-        smem_size,
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(GemmType.GroupedMasked, lhs, m, k, k, block_m, block_k, num_groups)
+    tensor_map_b = make_2d_tma_b_desc(GemmType.GroupedMasked, rhs, n, k, k, block_n, block_k, num_groups)
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.GroupedMasked, out, m, n, n, block_m, block_n, num_groups, smem_config[1]
     )
-    runtime = jit_tuner.compile_and_tune_group_gemm_masked(
-        name="m_grouped_gemm_fp8_fp8_bf16_nt",
-        keys={
-            "N": n,
-            "K": k,
-            "BLOCK_M": block_m,
-            "BLOCK_N": block_n,
-            "NUM_GROUPS": num_groups,
-            "NUM_STAGES": num_stages,
-            "NUM_TMA_MULTICAST": num_tma_multicast,
-            "GEMM_TYPE": "GroupedMasked",
-        },
-        space=(),
-        includes=includes,
-        arg_defs=(
-            ("lhs", paddle.float8_e4m3fn),
-            ("lhs_scales", paddle.float32),
-            ("rhs", paddle.float8_e4m3fn),
-            ("rhs_scales", paddle.float32),
-            ("out", paddle.bfloat16),
-            ("grouped_layout", paddle.int32),
-            ("m", int),
-            ("stream", paddle.device.cuda.Stream),
-            ("num_sms", int),
-            ("smem_size", int),
-        ),
-        template=template,
-        args=args,
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.GroupedMasked, lhs_scales, m, k, block_m, block_k, num_groups
     )
 
-    # Run the kernel
-    runtime(*args)
+    kwargs = {
+        # Templated arguments
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": k,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_GROUPS": num_groups,
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        "GEMM_TYPE": GemmType.GroupedMasked,
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": masked_m,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": paddle.device.cuda.current_stream().cuda_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("m_grouped_gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)

ops/csrc/fp8/deep_gemm/jit_kernels/runtime.py

@@ -0,0 +1,163 @@
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from typing import Any, Dict, Tuple
+
+import cuda.bindings.driver as cbd
+import paddle
+
+from ..jit.runtime import GemmType
+from .utils import get_tma_aligned_size
+
+# TODO Support dtype in Paddle
+tmap_type_map: Dict[Any, str] = {
+    paddle.int8: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    paddle.int16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    paddle.int32: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT32,
+    paddle.int64: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT64,
+    paddle.uint8: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.uint16:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    # paddle.uint32:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT32,
+    # paddle.uint64:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT64,
+    paddle.float32: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT32,
+    paddle.float16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
+    paddle.bfloat16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+    paddle.float8_e4m3fn: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.float8_e4m3fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    paddle.float8_e5m2: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.float8_e5m2fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+}
+
+swizzle_type_map = {
+    0: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,
+    32: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_32B,
+    64: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B,
+    128: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+}
+
+
+def make_2d_tma_copy_desc(
+    t: paddle.Tensor,
+    gmem_dims: Tuple[cbd.cuuint64_t, cbd.cuuint64_t],
+    gmem_outer_stride: cbd.cuuint64_t,
+    smem_dims: Tuple[cbd.cuuint32_t, cbd.cuuint32_t],
+    swizzle_type: cbd.CUtensorMapSwizzle,
+) -> cbd.CUtensorMap:
+    tensor_dtype = tmap_type_map[t.dtype]
+    res, tensor_map = cbd.cuTensorMapEncodeTiled(
+        tensor_dtype,
+        2,
+        t.data_ptr(),
+        gmem_dims,
+        (gmem_outer_stride,),
+        smem_dims,
+        (cbd.cuuint32_t(1), cbd.cuuint32_t(1)),
+        cbd.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE,
+        swizzle_type,
+        cbd.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
+        cbd.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    )
+    if res != cbd.CUresult.CUDA_SUCCESS:
+        raise Exception(f"Failed to encode tensor map: {res}")
+    return tensor_map
+
+
+def make_2d_tma_desc(
+    t: paddle.Tensor,
+    gmem_inner_dim: int,
+    gmem_outer_dim: int,
+    gmem_outer_stride: int,
+    smem_inner_dim: int,
+    smem_outer_dim: int,
+    swizzle_type: cbd.CUtensorMapSwizzle = cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+) -> cbd.CUtensorMap:
+    gmem_dim = (cbd.cuuint64_t(gmem_inner_dim), cbd.cuuint64_t(gmem_outer_dim))
+    smem_dim = (cbd.cuuint32_t(smem_inner_dim), cbd.cuuint32_t(smem_outer_dim))
+    return make_2d_tma_copy_desc(
+        t, gmem_dim, cbd.cuuint64_t(gmem_outer_stride * t.element_size()), smem_dim, swizzle_type
+    )
+
+
+def make_2d_tma_a_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_m: int,
+    shape_k: int,
+    m_stride: int,
+    block_m: int,
+    block_k: int,
+    num_groups: int,
+) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(
+        t, shape_k, shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1), m_stride, block_k, block_m
+    )
+
+
+def make_2d_tma_b_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_n: int,
+    shape_k: int,
+    n_stride: int,
+    block_n: int,
+    block_k: int,
+    num_groups: int,
+) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(
+        t, shape_k, shape_n * (num_groups if gemm_type != GemmType.Normal else 1), n_stride, block_k, block_n
+    )
+
+
+def make_2d_tma_d_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_m: int,
+    shape_n: int,
+    m_stride: int,
+    block_m: int,
+    block_n: int,
+    num_groups: int,
+    swizzle_mode: int,
+) -> cbd.CUtensorMap:
+    # Swizzling requires the inner box dim to be less or equal than `kSwizzleDMode`
+    # bytes, so `BLOCK_N * sizeof(T) / kSwizzleDMode` TMA stores are required
+    return make_2d_tma_desc(
+        t,
+        shape_n,
+        shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+        m_stride,
+        block_n if swizzle_mode == 0 else swizzle_mode // t.element_size(),
+        block_m,
+        swizzle_type_map[swizzle_mode],
+    )
+
+
+def make_2d_tma_scales_desc(
+    gemm_type: GemmType, t: paddle.Tensor, shape_mn: int, shape_k: int, block_mn: int, block_k: int, num_groups: int
+) -> cbd.CUtensorMap:
+    # Make TMA aligned to 16 bytes
+    shape_mn = get_tma_aligned_size(shape_mn, t.element_size())
+    return make_2d_tma_desc(
+        t,
+        shape_mn,
+        (shape_k + block_k - 1) // block_k * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+        shape_mn,
+        block_mn,
+        1,
+        cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,
+    )

ops/csrc/fp8/deep_gemm/jit_kernels/tuner.py

@@ -1,164 +0,0 @@
-# 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.
-
-# The file has been adapted from DeepSeek DeepEP project
-# Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
-
-import copy
-import os
-from typing import Any, Dict
-
-import paddle
-
-from ..jit import Runtime, build, cpp_format, generate
-
-
-class JITTuner:
-    def __init__(self) -> None:
-        self.tuned = {}
-
-    def compile_and_tune_group_gemm_masked(
-        self,
-        name: str,
-        keys: Dict[str, Any],
-        space: tuple,
-        includes: tuple,
-        arg_defs: tuple,
-        template: str,
-        args: tuple,
-    ) -> Runtime:
-        # NOTES: we always assume the space and template will not change
-        # We also assume the GPU device will not be changed
-        # NOTES: the function must have no accumulated side effects
-        keys = {k: keys[k] for k in sorted(keys.keys())}
-        signature = (name, f"{keys}")
-        if signature in self.tuned:
-            if os.getenv("DG_JIT_DEBUG", None):
-                print(f"Using cached JIT kernel {name} with keys {keys}")
-            return self.tuned[signature]
-
-        if os.getenv("DG_JIT_DEBUG", None):
-            print(f"Auto-tuning JIT kernel {name} with keys {keys}")
-
-        assert signature not in self.tuned
-        assert args is not None
-        space = (dict(),) if len(space) == 0 else space
-
-        kernels = []
-        for tuned_keys in space:
-            assert isinstance(tuned_keys, dict)
-            full_keys = copy.deepcopy(keys)
-            full_keys.update(tuned_keys)
-            code = generate(includes, arg_defs, cpp_format(template, full_keys))
-
-            # Illegal build must raise errors
-            kernels.append((build(name, arg_defs, code), tuned_keys))
-
-        best_runtime, best_time, best_keys = None, None, None
-        for runtime, tuned_keys in kernels:
-            if len(space) > 1:
-                # Check kernel validity
-                return_code = runtime(*args)
-                if return_code != 0:
-                    # Pass illegal kernels, e.g. insufficient shared memory capacity
-                    if os.getenv("DG_JIT_DEBUG", None):
-                        print(
-                            f"Illegal JIT kernel {name} with keys {keys} and tuned keys {tuned_keys}: error code {return_code}"
-                        )
-                    continue
-
-                # Measure performance with L2 flush and a large GEMM kernel before to reduce overhead between kernels
-                start_event = paddle.device.cuda.Event(enable_timing=True)
-                end_event = paddle.device.cuda.Event(enable_timing=True)
-                paddle.empty(int(256e6 // 4), dtype=paddle.int32).zero_()
-                paddle.randn((8192, 8192), dtype=paddle.float32, device="cuda") @ paddle.randn(
-                    (8192, 8192), dtype=paddle.float32
-                )
-                start_event.record()
-                for i in range(20):
-                    assert runtime(*args) == 0
-                end_event.record()
-                end_event.synchronize()
-                elapsed_time = start_event.elapsed_time(end_event)
-            else:
-                elapsed_time = 0
-
-            # Compare if better
-            if best_time is None or elapsed_time < best_time:
-                best_runtime, best_time, best_keys = runtime, elapsed_time, tuned_keys
-            if os.getenv("DG_JIT_DEBUG", None):
-                print(f"Tuned JIT kernel {name} with keys {keys} and tuned keys {tuned_keys} has time {elapsed_time}")
-        assert best_runtime is not None, f"Failed to tune JIT kernel {name} with keys {keys}"
-
-        # Cache the best runtime and return
-        if os.getenv("DG_JIT_DEBUG", None) or os.getenv("DG_PRINT_AUTOTUNE", None):
-            print(f"Best JIT kernel {name} with keys {keys} has tuned keys {best_keys} and time {best_time}")
-        self.tuned[signature] = best_runtime
-        return best_runtime
-
-    def compile_and_tune(
-        self,
-        m,
-        n,
-        k,
-        name: str,
-        keys: Dict[str, Any],
-        space: tuple,
-        includes: tuple,
-        arg_defs: tuple,
-        template: str,
-        # args: tuple,
-    ) -> Runtime:
-        # NOTES: we always assume the space and template will not change
-        # We also assume the GPU device will not be changed
-        # NOTES: the function must have no accumulated side effects
-        signature = (name, m, k, n)
-        if signature in self.tuned:
-            return self.tuned[signature]
-        # keys = {k: keys[k] for k in sorted(keys.keys())}
-        # signature = (name, f"{keys}")
-        # if signature in self.tuned:
-        #     return self.tuned[signature]
-        space = (dict(),) if len(space) == 0 else space
-
-        kernels = []
-        for tuned_keys in space:
-            assert isinstance(tuned_keys, dict)
-            full_keys = copy.deepcopy(keys)
-            full_keys.update(tuned_keys)
-            code = generate(includes, arg_defs, cpp_format(template, full_keys))
-
-            # Illegal build must raise errors
-            kernels.append((build(name, arg_defs, code), tuned_keys))
-
-        best_runtime, best_time, best_keys = None, None, None
-        for runtime, tuned_keys in kernels:
-            elapsed_time = 0
-
-            # Compare if better
-            if best_time is None or elapsed_time < best_time:
-                best_runtime, best_time, best_keys = runtime, elapsed_time, tuned_keys
-            if os.getenv("DG_JIT_DEBUG", None):
-                print(f"Tuned JIT kernel {name} with keys {keys} and tuned keys {tuned_keys} has time {elapsed_time}")
-        assert best_runtime is not None, f"Failed to tune JIT kernel {name} with keys {keys}"
-
-        # Cache the best runtime and return
-        if os.getenv("DG_JIT_DEBUG", None) or os.getenv("DG_PRINT_AUTOTUNE", None):
-            print(f"Best JIT kernel {name} with keys {keys} has tuned keys {best_keys} and time {best_time}")
-        self.tuned[signature] = best_runtime
-        return best_runtime
-
-
-jit_tuner = JITTuner()

ops/csrc/fp8/deep_gemm/jit_kernels/utils.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 import paddle
 from paddle import Tensor
@@ -109,6 +109,13 @@ def get_col_major_tma_aligned_tensor(x: Tensor) -> Tensor:
     assert x.dim() in (2, 3)
     remove_dim = False
     if x.dim() == 2:
+        m, n = x.shape
+
+        aligned_m = get_tma_aligned_size(m, x.element_size())
+
+        if x.strides[0] == 1 and x.strides[1] == aligned_m:
+            return x
+
         x, remove_dim = x.unsqueeze(0), True
 
     b, m, n = x.shape

ops/csrc/fp8/deep_gemm/jit_kernels/wgrad_gemm.py

@@ -0,0 +1,195 @@
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from typing import List, Tuple
+import os
+import paddle
+
+from ..jit import FP8WGradGemmRuntime, build
+from .gemm import get_best_configs
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
+from .utils import (
+    ceil_div,
+    get_col_major_tma_aligned_tensor,
+    get_num_sms,
+    get_tma_aligned_size,
+)
+
+
+def wgrad_gemm_fp8_fp8_fp32_nt(
+    lhs: Tuple[paddle.Tensor, paddle.Tensor], rhs: Tuple[paddle.Tensor, paddle.Tensor], out: paddle.Tensor, num_sms: int = None
+):
+    """
+    Perform a weight gradient GEMM with FP8 inputs and FP32 output, with 1x128 LHS scaling and 1x128 RHS scaling.
+        Results will be accumulated into the output tensor.
+
+    Requirements:
+        LHS, RHS, and output tensors must be contiguous in dimension 1, i.e., strides[1] = 1.
+        The strides[0] of LHS and RHS must be a multiple of 16, and the strides[0] of output must be a multiple of 4.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling and RHS scaling tensor require a TMA-aligned transposed format.
+            If your input does not match the requirement, this function will do a transposing with a set of slow PaddlePaddle operations.
+
+    Arguments:
+        lhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[m, k]`,
+             the second element is an FP32 1x128 scaling tensor for LHS of shape `[m, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[n, k]`,
+             the second element is an FP32 1x128 scaling tensor for RHS of shape `[n, ⌈k / 128⌉]`.
+        out: the FP32 output tensor of shape `[m, n]`, which will be accumulated.
+    """
+    lhs, lhs_scales = lhs
+    rhs, rhs_scales = rhs
+    m, k = lhs.shape
+    n, k_ = rhs.shape
+    m_, n_ = out.shape
+
+    # Type and shape checks
+    assert m == m_ and n == n_ and k == k_
+    assert n > 0 and m > 0
+    assert lhs_scales.shape == [m, ceil_div(k, 128)] or lhs_scales.shape == [ceil_div(k, 128), m]
+    assert rhs_scales.shape == [n, ceil_div(k, 128)] or rhs_scales.shape == [ceil_div(k, 128), n]
+    assert lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
+    assert rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
+    assert out.dtype == paddle.float32
+    assert lhs.strides[1] == 1 and out.strides[1] == 1 and rhs.strides[1] == 1
+
+    # LHS and RHS scales must be transposed for TMA load
+    # NOTES: `get_col_major_tma_aligned_tensor` may launch a kernel if not processed by previous kernels
+    def get_valid_scales(scales: paddle.Tensor, mn: int):
+        if scales.shape == [ceil_div(k, 128), mn]:
+            # For k-grouped GEMMs
+            scales = scales.transpose([1, 0])
+            assert get_tma_aligned_size(mn, 4) == scales.strides[1] == mn
+        else:
+            scales = get_col_major_tma_aligned_tensor(scales)
+        return scales
+
+    lhs_scales = get_valid_scales(lhs_scales, m)
+    rhs_scales = get_valid_scales(rhs_scales, n)
+
+    # Do nothing if `k` is zero
+    if k == 0:
+        return
+
+    # K must be aligned to 128
+    aligned_k = ceil_div(k, 128) * 128
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = get_best_configs(
+        m, n, aligned_k, 1, num_sms, is_fp32_out=True, is_wgrad=True
+    )
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(f"Auto-tuned wgrad_gemm_fp8_fp8_fp32_nt as num_sms={num_sms}, block_m={block_m}, block_n={block_n}")
+    num_last_stages = ceil_div(k, 128) % num_stages
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(GemmType.Normal, lhs, m, k, lhs.strides[0], block_m, block_k, 1)
+    tensor_map_b = make_2d_tma_b_desc(GemmType.Normal, rhs, n, k, rhs.strides[0], block_n, block_k, 1)
+    tensor_map_d = make_2d_tma_d_desc(GemmType.Normal, out, m, n, out.strides[0], block_m, block_n, 1, smem_config[1])
+    tensor_map_scales_a = make_2d_tma_scales_desc(GemmType.Normal, lhs_scales, m, k, block_m, block_k, 1)
+    tensor_map_scales_b = make_2d_tma_scales_desc(GemmType.Normal, rhs_scales, n, k, block_n, block_k, 1)
+
+    kwargs = {
+        # Templated arguments
+        "GEMM_TYPE": GemmType.Normal,
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": aligned_k,
+        "NUM_GROUPS": 1,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "NUM_STAGES": num_stages,
+        "NUM_LAST_STAGES": num_last_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        # Runtime arguments
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_SCALES_B": tensor_map_scales_b,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": paddle.device.current_stream().stream_base.cuda_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("wgrad_gemm_fp8_fp8_fp32_nt", FP8WGradGemmRuntime, kwargs)
+    runtime(**kwargs)
+
+
+def k_grouped_wgrad_gemm_fp8_fp8_fp32_nt(
+    lhs: Tuple[paddle.Tensor, paddle.Tensor],
+    rhs: Tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    batch_sizes: List[int],
+    num_sms: int = None,
+):
+    """
+    Perform a k-grouped weight gradient GEMM with FP8 inputs and FP32 output, with 1x128 LHS scaling and 1x128 RHS scaling.
+        Results will be accumulated into the output tensor.
+
+    Requirements:
+        This function handles multiple batches with varying k-dimensions, processing each batch sequentially.
+        Each batch's LHS, RHS, and output tensors must be contiguous.
+        The RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling and RHS scaling tensors require a TMA-aligned transposed format.
+
+    Arguments:
+        lhs: The first element is a flattened FP8 tensor (typed `paddle.bfloat16`) containing all batches of LHS data,
+                 and the flattened shape is `[sum(m * k for k in batch_sizes)]`, where m is the number of rows.
+             The second element is an FP32 scaling tensor for LHS with shape `[⌈k / 128⌉ for k in batch_sizes), m]`,
+                 representing the per-128-channel scaling factors.
+        rhs: The first element is a flattened FP8 tensor (typed `paddle.bfloat16`) containing all batches of RHS data,
+                 and the flattened shape is `[sum(n * k for k in batch_sizes)]`, where n is the number of rows.
+             The second element is an FP32 scaling tensor for RHS with shape `[⌈k / 128⌉ for k in batch_sizes), n]`,
+                 representing the per-128-channel scaling factors.
+        out: The FP32 output tensor of shape [num_batches, m, n], which will be accumulated.
+        batch_sizes: A list of integers specifying the k-dimension for each batch.
+    """
+    lhs, lhs_scales = paddle.view(lhs[0], [-1]), lhs[1]
+    rhs, rhs_scales = paddle.view(rhs[0], [-1]), rhs[1]
+    num_batches, m, n = out.shape
+
+    lhs_offset, rhs_offset, scales_offset = 0, 0, 0
+
+    for i in range(num_batches):
+        k = batch_sizes[i]
+        lhs_slice = paddle.view(lhs[lhs_offset : lhs_offset + m * k], (m, k))
+        rhs_slice = paddle.view(rhs[rhs_offset : rhs_offset + n * k], (n, k))
+        lhs_scales_slice = lhs_scales[scales_offset : scales_offset + ceil_div(k, 128)]
+        rhs_scales_slice = rhs_scales[scales_offset : scales_offset + ceil_div(k, 128)]
+        wgrad_gemm_fp8_fp8_fp32_nt((lhs_slice, lhs_scales_slice), (rhs_slice, rhs_scales_slice), out[i], num_sms)
+
+        lhs_offset += m * k
+        rhs_offset += n * k
+        scales_offset += ceil_div(k, 128)

ops/csrc/fp8/deep_gemm/utils.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
 
 import os
 import sys
@@ -22,51 +22,6 @@ import sys
 import paddle
 
 
-def bench(fn, num_warmups: int = 5, num_tests: int = 10, high_precision: bool = False):
-    # Flush L2 cache with 256 MB data
-    paddle.device.cuda.synchronize()
-    cache = paddle.empty(int(256e6 // 4), dtype=paddle.int32)
-    cache.zero_()
-
-    # Warmup
-    for _ in range(num_warmups):
-        fn()
-
-    # Add a large kernel to eliminate the CPU launch overhead
-    if high_precision:
-        x = paddle.randn((8192, 8192), dtype=paddle.float32)
-        y = paddle.randn((8192, 8192), dtype=paddle.float32)
-        x @ y
-
-    # Testing
-    start_event = paddle.device.cuda.Event(enable_timing=True)
-    end_event = paddle.device.cuda.Event(enable_timing=True)
-    start_event.record()
-    for i in range(num_tests):
-        fn()
-    end_event.record()
-    paddle.cuda.synchronize()
-
-    return start_event.elapsed_time(end_event) / num_tests
-
-
-def get_cuda_home():
-    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
-    if cuda_home:
-        return cuda_home
-
-    try:
-        which_cmd = "which nvcc"
-
-        nvcc_path = os.popen(which_cmd).read().strip()
-        if nvcc_path:
-            return os.path.dirname(os.path.dirname(nvcc_path))
-    except Exception:
-        pass
-
-    return None
-
-
 class empty_suppress:
     def __enter__(self):
         return self

ops/csrc/fp8/setup.py

@@ -92,7 +92,7 @@ if __name__ == "__main__":
         packages=["deep_gemm", "deep_gemm/jit", "deep_gemm/jit_kernels"],
         package_data={
             "deep_gemm": [
-                "include/deep_gemm/**/*",
+                "include/deep_gemm/*",
                 "include/cute/**/*",
                 "include/cutlass/**/*",
             ]

ops/csrc/fp8/tests/test_core.py

@@ -12,29 +12,36 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# The file has been adapted from DeepSeek DeepEP project
+# The file has been adapted from DeepSeek DeepGEMM project
 # Copyright (c) 2025 DeepSeek
-# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+import cuda.bindings.nvrtc as nvrtc
+
+print(f"NVRTC version: {nvrtc.nvrtcVersion()[1:]}")
 
 import random
-from typing import Tuple
+from typing import List, Tuple
 
 import deep_gemm
 import paddle
 from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
+from deep_gemm.jit_kernels.utils import get_m_alignment_for_contiguous_layout
 from paddle import Tensor
 
 
 def per_token_cast_to_fp8(x: Tensor) -> Tuple[Tensor, Tensor]:
-    assert x.dim() == 2 and x.shape[1] % 128 == 0
+    assert x.dim() == 2
     m, n = x.shape
+    pad_size = (128 - (n % 128)) % 128
+    x = paddle.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x
     x_view = paddle.view(x, (m, -1, 128))
     x_abs = paddle.abs(x_view).astype(paddle.float32)
     x_amax = paddle.amax(x_abs, axis=2)
     x_amax = paddle.view(x_amax, (m, -1))
     x_amax = paddle.clip(x_amax, min=1e-4)
     scaled_x = x_view * (448.0 / x_amax.unsqueeze(2))
-    scaled_x_converted = paddle.view(scaled_x.astype(paddle.float8_e4m3fn), (m, n))
+    scaled_x_converted = paddle.view(scaled_x.astype(paddle.float8_e4m3fn), (m, n + pad_size))[:, :n]
 
     x_amax_scaled = paddle.view((x_amax / 448.0), (m, -1))
 
@@ -71,26 +78,62 @@ def construct(m: int, k: int, n: int) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tens
     return x_fp8, y_fp8, out, ref_out
 
 
-def construct_grouped(
-    num_groups: int, m: int, k: int, n: int, is_masked: bool
-) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tensor, Tensor]:
-    # x_np = np.full((num_groups, m, k), 3)
-    # y_np = np.full((num_groups, n, k), 2)
-    # x=paddle.to_tensor(x_np).astype(paddle.bfloat16)
-    # y=paddle.to_tensor(y_np).astype(paddle.bfloat16)
-    x = paddle.randn((num_groups, m, k), dtype=paddle.bfloat16)
+def construct_contiguous_grouped(
+    num_groups: int, expected_m_per_group: int, k: int, n: int
+) -> Tuple[int, Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tensor, Tensor, Tensor]:
+
+    alignment = get_m_alignment_for_contiguous_layout()
+    group_ms = [int(expected_m_per_group * random.uniform(0.7, 1.3)) for _ in range(num_groups)]
+    m = sum([ceil_div(x, alignment) * alignment for x in group_ms])
+
+    x = paddle.randn((m, k), dtype=paddle.bfloat16)
     y = paddle.randn((num_groups, n, k), dtype=paddle.bfloat16)
-    out = paddle.empty((num_groups, m, n), dtype=paddle.bfloat16)
-    ref_out = paddle.einsum("gmk,gnk->gmn", x, y)
+    m_indices = paddle.empty([m], dtype=paddle.int32)
+    out = paddle.empty((m, n), dtype=paddle.bfloat16)
+    ref_out = paddle.randn((m, n), dtype=paddle.bfloat16)
+
+    start = 0
+    for i, group_m in enumerate(group_ms):
+        actual_end = start + group_m
+        aligned_end = start + ceil_div(group_m, alignment) * alignment
+        m_indices[start:actual_end] = i
+        m_indices[actual_end:aligned_end] = -1
+        ref_out[start:aligned_end] = x[start:aligned_end] @ y[i].t()
+        start = aligned_end
+
+    ref_out = paddle.where((m_indices == -1).unsqueeze(1), paddle.zeros_like(ref_out), ref_out)
 
     assert m % 4 == 0, f"TMA alignment error: {m}"
+    x_fp8 = per_token_cast_to_fp8(x)
+    y_fp8 = (
+        paddle.empty_like(y, dtype=paddle.float8_e4m3fn),
+        paddle.empty((num_groups, ceil_div(n, 128), k // 128), dtype=paddle.float32),
+    )
+    for i in range(num_groups):
+        # y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i])
+        y_fp8_0_i, y_fp8_1_i = per_block_cast_to_fp8(y[i])
+        paddle.assign(y_fp8_0_i, y_fp8[0][i])
+        paddle.assign(y_fp8_1_i, y_fp8[1][i])
+
+    return m, x_fp8, y_fp8, m_indices, out, ref_out
+
+
+def construct_masked_grouped(
+    num_groups: int, max_m: int, expected_m_per_group: int, k: int, n: int
+) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tensor, Tensor, Tensor]:
+    x = paddle.randn((num_groups, max_m, k), dtype=paddle.bfloat16)
+    y = paddle.randn((num_groups, n, k), dtype=paddle.bfloat16)
+    out = paddle.empty((num_groups, max_m, n), dtype=paddle.bfloat16)
+    ref_out = paddle.einsum("gmk,gnk->gmn", x, y)
+
+    assert max_m % 4 == 0, f"TMA alignment error: {max_m}"
     x_fp8 = (
         paddle.empty_like(x, dtype=paddle.float8_e4m3fn),
-        paddle.empty((num_groups, m, k // 128), dtype=paddle.float32),
+        paddle.empty((num_groups, max_m, k // 128), dtype=paddle.float32),
     )
     y_fp8 = (
         paddle.empty_like(y, dtype=paddle.float8_e4m3fn),
-        paddle.empty((num_groups, (n + 127) // 128, k // 128), dtype=paddle.float32),
+        paddle.empty((num_groups, ceil_div(n, 128), k // 128), dtype=paddle.float32),
     )
     for i in range(num_groups):
         # x_fp8[0][i], x_fp8[1][i] = per_token_cast_to_fp8(x[i])
@@ -102,25 +145,99 @@ def construct_grouped(
         paddle.assign(y_fp8_0_i, y_fp8[0][i])
         paddle.assign(y_fp8_1_i, y_fp8[1][i])
 
-    # For non-masked input, we must merge the group and M dims
-    if not is_masked:
-        x_fp8 = (paddle.view(x_fp8[0], (-1, k)), per_token_cast_to_fp8(paddle.view(x, (-1, k)))[1])
-        out, ref_out = paddle.view(out, (-1, n)), paddle.view(ref_out, (-1, n))
-
     # Transpose earlier so that the testing will not trigger transposing kernels
     x_fp8 = (x_fp8[0], get_col_major_tma_aligned_tensor(x_fp8[1]))
-    return x_fp8, y_fp8, out, ref_out
+
+    # Construct mask
+    masked_m = paddle.empty((num_groups,), dtype=paddle.int32)
+    for j in range(num_groups):
+        masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3))
+    assert masked_m.amax().item() <= max_m
+    return x_fp8, y_fp8, masked_m, out, ref_out
+
+
+def construct_wgrad(
+    m: int, k: int, n: int
+) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tensor, Tensor, Tensor]:
+    x = paddle.randn((m, k), dtype=paddle.bfloat16)
+    y = paddle.randn((n, k), dtype=paddle.bfloat16)
+    residual = paddle.randn((m, n), dtype=paddle.float32) * 10
+    out = residual.clone()
+    ref_out = residual + (x.astype(paddle.float32) @ y.astype(paddle.float32).t())
+
+    x_fp8 = per_token_cast_to_fp8(x)
+    y_fp8 = per_token_cast_to_fp8(y)
+
+    # NOTES: please do inplace add on the `out` later
+    return x_fp8, y_fp8, residual, out, ref_out
+
+
+def construct_k_grouped_wgrad(
+    m: int, n: int, k_sizes: List[int]
+) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tensor, Tensor, List[int]]:
+    num_groups, total_k = len(k_sizes), sum(k_sizes)
+
+    x_flat = paddle.empty((m * total_k,), dtype=paddle.bfloat16)
+    y_flat = paddle.empty((n * total_k,), dtype=paddle.bfloat16)
+    out = paddle.zeros((num_groups, m, n), dtype=paddle.float32)
+    ref_out = paddle.zeros((num_groups, m, n), dtype=paddle.float32)
+
+    # Fill tensors with data and compute reference output
+    x_offset, y_offset = 0, 0
+    for idx, k in enumerate(k_sizes):
+        x_chunk = paddle.randn((m, k), dtype=paddle.bfloat16)
+        y_chunk = paddle.randn((n, k), dtype=paddle.bfloat16)
+
+        paddle.assign(x_chunk.flatten(), output=x_flat[x_offset : x_offset + m * k])
+        paddle.assign(y_chunk.flatten(), output=y_flat[y_offset : y_offset + n * k])
+        ref_out[idx] = x_chunk.astype(paddle.float32) @ y_chunk.astype(paddle.float32).t()
+
+        x_offset += m * k
+        y_offset += n * k
+
+    x_fp8_flat = paddle.empty_like(x_flat, dtype=paddle.float8_e4m3fn)
+    y_fp8_flat = paddle.empty_like(y_flat, dtype=paddle.float8_e4m3fn)
+
+    total_scale_factors = sum(ceil_div(k, 128) for k in k_sizes)
+    x_scales = paddle.empty((total_scale_factors, m), dtype=paddle.float32)
+    y_scales = paddle.empty((total_scale_factors, n), dtype=paddle.float32)
+
+    # Cast to FP8 and prepare scale factors
+    x_offset, y_offset, scale_offset = 0, 0, 0
+    for k in k_sizes:
+        x_fp8_chunk, x_scale_chunk = per_token_cast_to_fp8(paddle.view(x_flat[x_offset : x_offset + m * k], (m, k)))
+        y_fp8_chunk, y_scale_chunk = per_token_cast_to_fp8(paddle.view(y_flat[y_offset : y_offset + n * k], (n, k)))
+
+        paddle.assign(x_fp8_chunk.flatten(), output=x_fp8_flat[x_offset : x_offset + m * k])
+        paddle.assign(y_fp8_chunk.flatten(), output=y_fp8_flat[y_offset : y_offset + n * k])
+
+        num_scales = ceil_div(k, 128)
+        paddle.assign(x_scale_chunk.T, output=x_scales[scale_offset : scale_offset + num_scales])
+        paddle.assign(y_scale_chunk.T, output=y_scales[scale_offset : scale_offset + num_scales])
+
+        x_offset += m * k
+        y_offset += n * k
+        scale_offset += num_scales
+
+    return (x_fp8_flat, x_scales), (y_fp8_flat, y_scales), out, ref_out, k_sizes
 
 
 def test_gemm() -> None:
     print("Testing GEMM:")
-    for m in (64,):
+    for m in (64, 128, 4096):
         for k, n in [
+            (576, 7168),
             (7168, 2112),
+            (1536, 24576),
+            (512, 32768),
+            (16384, 7168),
+            (7168, 4096),
+            (2048, 7168),
         ]:
             x_fp8, y_fp8, out, ref_out = construct(m, k, n)
             deep_gemm.gemm_fp8_fp8_bf16_nt(x_fp8, y_fp8, out)
             diff = calc_diff(out, ref_out)
+            print("diff:", diff)
             assert diff < 0.001, f"{m=}, {k=}, {n=}, {diff:.5f}"
 
     print()
@@ -129,42 +246,88 @@ def test_gemm() -> None:
 def test_m_grouped_gemm_contiguous() -> None:
     print("Testing grouped contiguous GEMM:")
 
-    for num_groups, m, k, n in ((4, 8192, 7168, 4096),):
-        # TODO: make a stronger test
-        x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=False)
-        m_indices = paddle.arange(0, num_groups, dtype=paddle.int32)
-        # m_indices = m_indices.unsqueeze(-1).expand(num_groups, m).contiguous().view(-1)
-        m_indices = paddle.flatten(paddle.expand(paddle.unsqueeze(m_indices, -1), shape=[num_groups, m]))
+    for num_groups, expected_m_per_group, k, n in (
+        (4, 8192, 7168, 4096),
+        (4, 8192, 2048, 7168),
+        (8, 4096, 7168, 4096),
+        (8, 4096, 2048, 7168),
+        (32, 256, 7168, 4096),
+        (32, 256, 2048, 7168),
+    ):
+        # NOTES: we should mask the unfilled part before calculating difference
+        m, x_fp8, y_fp8, m_indices, out, ref_out = construct_contiguous_grouped(num_groups, expected_m_per_group, k, n)
         deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(x_fp8, y_fp8, out, m_indices)
+        out = paddle.where((m_indices == -1).unsqueeze(1), paddle.zeros_like(out), out)
         diff = calc_diff(out, ref_out)
         print("diff:", diff)
-        assert diff < 0.001, f"m={m * num_groups}, {k=}, {n=}, {diff:.5f}"
+        assert diff < 0.001, f"{m=}, {k=}, {n=}, {diff:.5f}"
     print()
 
 
 def test_m_grouped_gemm_masked() -> None:
     print("Testing grouped masked GEMM:")
-
-    for num_groups, m in ((1, 1024),):
-        for k, n in ((7168, 4096),):
+    for num_groups, expected_m_per_group in ((1, 1024), (2, 512), (4, 256)):
+        for k, n in (
+            (7168, 4096),
+            (2048, 7168),
+        ):
             # Test correctness
-            masked_m_candidates = list(filter(lambda candidate: candidate <= m, (64, 128, 192, 256, 320, 384)))
             for i in range(10):
-                x_fp8, y_fp8, out, ref_out = construct_grouped(num_groups, m, k, n, is_masked=True)
-                masked_m = paddle.empty((num_groups,), dtype=paddle.int32)
-                for j in range(num_groups):
-                    masked_m[j] = random.choice(masked_m_candidates)
-                # expected_m = min(int(masked_m.float().mean()) + 1, m)
-                masked_m_float = paddle.cast(masked_m, "float32")
-                masked_m_mean = paddle.mean(masked_m_float)
-                masked_m_mean_int = paddle.cast(masked_m_mean, "int32")
-                expected_m = min(masked_m_mean_int + 1, m)
-                deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m)
+                x_fp8, y_fp8, masked_m, out, ref_out = construct_masked_grouped(
+                    num_groups, 4096, expected_m_per_group, k, n
+                )
+                deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m_per_group)
                 for j in range(num_groups):
                     diff = calc_diff(out[j, : masked_m[j].item()], ref_out[j, : masked_m[j].item()])
-                    print("diff:", diff)
-                    assert diff < 0.001, f"{m=}, {k=}, {n=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}"
+                    assert (
+                        diff < 0.001
+                    ), f"{expected_m_per_group=}, {k=}, {n=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}"
 
+            # noinspection PyShadowingNames
+            # def test_func():
+            #     deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(x_fp8, y_fp8, out, masked_m, expected_m_per_group)
+
+            # Test performance with fixed shapes
+            # noinspection PyUnboundLocalVariable
+            # valid_m = masked_m.sum().item()
+            # t = bench_kineto(test_func, "fp8_gemm", suppress_kineto_output=True)
+            # print(f" > Perf ({num_groups=}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}): {t * 1e6:4.0f} us | "
+            #       f"throughput: {2 * valid_m * n * k / t / 1e12:4.0f} TFLOPS, "
+            #       f"{(valid_m * k + num_groups * k * n + valid_m * n * 2) / 1e9 / t:4.0f} GB/s")
+    print()
+
+
+def test_wgrad_gemm():
+    print("Testing weight gradient GEMM:")
+
+    for k in (4096, 8192):
+        for m, n in ((7168, 2112), (1536, 24576), (512, 32768), (16384, 7168), (7168, 4096), (2048, 7168)):
+            # Test correctness
+            x_fp8, y_fp8, residual, out, ref_out = construct_wgrad(m, k, n)
+            deep_gemm.wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out)
+            diff = calc_diff(out, ref_out)
+            print("diff:", diff)
+            assert diff < 0.001, f"{m=}, {k=}, {n=}, {diff:.5f}"
+    print()
+
+
+def test_k_grouped_wgrad_gemm():
+    print("Testing grouped weight gradient GEMM:")
+
+    for num_groups, base_k in ((4, 4096), (4, 8192), (8, 4096)):
+        for m, n in ((7168, 4096), (2048, 7168)):
+            # Vary k sizes around base_k
+            k_sizes = [base_k + random.randint(-1, 1) * 128 for _ in range(num_groups - 1)]
+            k_sizes.append(base_k * num_groups - sum(k_sizes))
+
+            # Test correctness
+            x_fp8, y_fp8, out, ref_out, k_sizes = construct_k_grouped_wgrad(m, n, k_sizes)
+            deep_gemm.k_grouped_wgrad_gemm_fp8_fp8_fp32_nt(x_fp8, y_fp8, out, k_sizes)
+
+            for idx in range(num_groups):
+                diff = calc_diff(out[idx], ref_out[idx])
+                print("diff:", diff)
+                assert diff < 0.001, f"{num_groups=}, {m=}, {n=}, k={k_sizes[idx]}, batch={idx}, {diff:.5f}"
     print()
 
 
@@ -176,3 +339,6 @@ if __name__ == "__main__":
     test_gemm()
     test_m_grouped_gemm_contiguous()
     test_m_grouped_gemm_masked()
+
+    test_wgrad_gemm()
+    test_k_grouped_wgrad_gemm()

ops/csrc/fp8/tests/test_jit.py

@@ -0,0 +1,116 @@
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+import ctypes
+import os
+from typing import Any, Dict
+
+import cuda.bindings.driver as cbd
+import paddle
+from deep_gemm import jit
+
+# Essential debugging staffs
+os.environ["DG_JIT_DEBUG"] = os.getenv("DG_JIT_DEBUG", "1")
+os.environ["DG_JIT_DISABLE_CACHE"] = os.getenv("DG_JIT_DISABLE_CACHE", "1")
+
+
+class VectorAddRuntime(jit.Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    @staticmethod
+    def generate(kwargs: Dict[str, Any]) -> str:
+        return f"""
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#endif
+
+#include <cuda_fp8.h>
+#include <cuda_bf16.h>
+
+template <typename T>
+__global__ void vector_add(T* a, T* b, T* c, uint32_t n) {{
+    uint32_t i = blockDim.x * blockIdx.x + threadIdx.x;
+    if (i < n) {{
+        c[i] = a[i] + b[i];
+    }}
+}}
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&vector_add<{kwargs['T']}>);
+}}
+"""
+
+    # noinspection PyShadowingNames,PyMethodOverriding
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: Dict[str, Any]) -> cbd.CUresult:
+        assert kwargs["A"].shape == kwargs["B"].shape == kwargs["C"].shape
+        assert kwargs["A"].place == kwargs["B"].place == kwargs["C"].place
+        assert kwargs["A"].dim() == 1
+
+        config = cbd.CUlaunchConfig()
+        config.gridDimX = (kwargs["A"].numel() + 127) // 128
+        config.gridDimY = 1
+        config.gridDimZ = 1
+        config.blockDimX = 128
+        config.blockDimY = 1
+        config.blockDimZ = 1
+        config.hStream = kwargs["STREAM"]
+
+        arg_values = (
+            kwargs["A"].data_ptr(),
+            kwargs["B"].data_ptr(),
+            kwargs["C"].data_ptr(),
+            kwargs["A"].numel(),
+        )
+        arg_types = (
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+        )
+
+        return cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)[0]
+
+
+if __name__ == "__main__":
+    print("Generated code:")
+    kwargs = {"T": "float"}
+    code = VectorAddRuntime.generate(kwargs)
+    print(code)
+    print()
+
+    for compiler_name in ("NVCC", "NVRTC"):
+        # Get compiler
+        compiler_cls = getattr(jit, f"{compiler_name}Compiler")
+        print(f"Compiler: {compiler_name}, version: {compiler_cls.__version__()}")
+
+        # Build
+        print("Building ...")
+        func = compiler_cls.build("test_func", code, VectorAddRuntime, kwargs)
+
+        # Run and check
+        a = paddle.randn((1024,), dtype=paddle.float32)
+        b = paddle.randn((1024,), dtype=paddle.float32)
+        c = paddle.empty_like(a)
+        ret = func(A=a, B=b, C=c, STREAM=paddle.device.cuda.current_stream().cuda_stream)
+        assert ret == cbd.CUresult.CUDA_SUCCESS, ret
+        assert paddle.allclose(c, a + b).item()
+        print(f"JIT test for {compiler_name} passed\n")

paddlenlp/data/dist_dataloader.py

@@ -94,22 +94,22 @@ class DistDataLoader(paddle.io.DataLoader):
 
         if self._need_data:
             self._dataloader = paddle.io.DataLoader(
-                dataset,
-                feed_list,
-                places,
-                return_list,
-                batch_sampler,
-                batch_size,
-                shuffle,
-                drop_last,
-                collate_fn,
-                num_workers,
-                use_buffer_reader,
-                prefetch_factor,
-                use_shared_memory,
-                timeout,
-                worker_init_fn,
-                persistent_workers,
+                dataset=dataset,
+                feed_list=feed_list,
+                places=places,
+                return_list=return_list,
+                batch_sampler=batch_sampler,
+                batch_size=batch_size,
+                shuffle=shuffle,
+                drop_last=drop_last,
+                collate_fn=collate_fn,
+                num_workers=num_workers,
+                use_buffer_reader=use_buffer_reader,
+                prefetch_factor=prefetch_factor,
+                use_shared_memory=use_shared_memory,
+                timeout=timeout,
+                worker_init_fn=worker_init_fn,
+                persistent_workers=persistent_workers,
             )
 
             self._lazy_dataloader_iter = None

paddlenlp/trainer/trainer.py

@@ -170,6 +170,7 @@ from .utils.helper import (  # nested_truncate,
     nested_numpify,
     nested_truncate,
 )
+from .utils.load_hf_ckpt import load_huggingface_ckpt
 from .utils.sharding_io import ShardingIO
 
 DEFAULT_CALLBACKS = [DefaultFlowCallback]
@@ -1009,6 +1010,9 @@ class Trainer:
         if self.args.ignore_data_skip:
             self.timers and self.timers("read-data").start()
 
+        if self.args.resume_from_huggingface_ckpt is not None:
+            load_huggingface_ckpt(model, self.args.resume_from_huggingface_ckpt)
+
         for epoch in range(epochs_trained, num_train_epochs):
             if isinstance(train_dataloader, paddle.io.DataLoader) and isinstance(
                 train_dataloader.batch_sampler, DistributedBatchSampler
@@ -1651,6 +1655,7 @@ class Trainer:
                 batch_size=self.args.per_device_train_batch_size,
                 collate_fn=self.data_collator,
                 num_workers=self.args.dataloader_num_workers,
+                prefetch_factor=32,
                 **additional_configs,
             )
         else:
@@ -1663,6 +1668,7 @@ class Trainer:
                 batch_sampler=train_sampler,
                 collate_fn=self.data_collator,
                 num_workers=self.args.dataloader_num_workers,
+                prefetch_factor=32,
                 **additional_configs,
             )
 

paddlenlp/trainer/trainer_callback.py

@@ -20,12 +20,14 @@ Callbacks to use with the Trainer class and customize the training loop.
 """
 import dataclasses
 import json
+import os
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 import numpy as np
 from tqdm.auto import tqdm
 
+from paddlenlp.transformers.moe_utils import offload, reload
 from paddlenlp.utils.log import logger
 
 from .trainer_utils import IntervalStrategy, has_length
@@ -40,6 +42,8 @@ __all__ = [
     "ProgressCallback",
     "PrinterCallback",
     "EarlyStoppingCallback",
+    "StepFlexToken",
+    "FP8QuantWeightCallback",
 ]
 
 
@@ -534,14 +538,14 @@ class PrinterCallback(TrainerCallback):
 
     def on_log(self, args, state, control, logs=None, **kwargs):
         _ = logs.pop("total_flos", None)
-        if state.is_local_process_zero:
-            if type(logs) is dict:
-                logger.info(", ".join(f"{k}: {v}" for k, v in logs.items()))
-                metrics_dumper = kwargs.get("metrics_dumper", None)
-                if metrics_dumper is not None:
-                    metrics_dumper.append(logs)
-            else:
-                logger.info(logs)
+        # if state.is_local_process_zero:
+        if type(logs) is dict:
+            logger.info(", ".join(f"{k}: {v}" for k, v in logs.items()))
+            metrics_dumper = kwargs.get("metrics_dumper", None)
+            if metrics_dumper is not None:
+                metrics_dumper.append(logs)
+        else:
+            logger.info(logs)
 
 
 class EarlyStoppingCallback(TrainerCallback):
@@ -601,3 +605,69 @@ class EarlyStoppingCallback(TrainerCallback):
         self.check_metric_value(args, state, control, metric_value)
         if self.early_stopping_patience_counter >= self.early_stopping_patience:
             control.should_training_stop = True
+
+
+class StepFlexToken(TrainerCallback):
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        model = kwargs.pop("model")
+        if hasattr(model, "step_flex_token"):
+            model.step_flex_token(state.global_step)
+
+
+g_shard_bypass_dygraph_optimizer = int(os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0))
+
+
+def enable_in_dict_config(config, key):
+    """enable_in_dict_config"""
+    return key in config and config[key]
+
+
+skip_count = 0
+
+
+class FP8QuantWeightCallback(TrainerCallback):
+    """
+    FP8QuantWeightCallback
+    """
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        """
+        每个step开始前把专家参数quant成fp8q
+        """
+        model = kwargs["model"]
+        optimizer = kwargs["optimizer"]
+        global skip_count
+
+        if (not g_shard_bypass_dygraph_optimizer or skip_count == 0) and hasattr(model, "fp8_quant_weight"):
+            model.fp8_quant_weight(True, quant_transpose=True)
+            optimizer.clear_param_storage("moe_expert")
+            optimizer.clear_param_storage("rms_linear")
+            optimizer.clear_param_storage("memory_attn")
+            optimizer.clear_param_storage("attn_out_project")
+            optimizer.clear_param_storage("shared_expert")
+
+            self.moe_weights_name = []
+            for param in optimizer._inner_opt._parameter_list:
+                color = getattr(param, "color", -1)
+                if isinstance(color, dict) and color["color"] == "moe_expert":
+                    self.moe_weights_name.append(param.name)
+
+            for name in self.moe_weights_name:
+                offload(optimizer._master_weights[name])
+
+        skip_count += 1
+
+    def on_optimizer_begin(self, args, state, control, **kwargs):
+        model = kwargs["model"]
+        optimizer = kwargs["optimizer"]
+        global skip_count
+
+        if (not g_shard_bypass_dygraph_optimizer) and hasattr(model, "fp8_quant_weight"):
+            for name in self.moe_weights_name:
+                reload(optimizer._master_weights[name])

paddlenlp/trainer/training_args.py

@@ -885,6 +885,10 @@ class TrainingArguments:
         default=None,
         metadata={"help": "The path to a folder with a valid checkpoint for your model."},
     )
+    resume_from_huggingface_ckpt: Optional[str] = field(
+        default=None,
+        metadata={"help": "The path to a folder with a valid huggingface checkpoint for your model."},
+    )
     auto_parallel_resume_form_hybrid_parallel: Optional[bool] = field(
         default=False,
         metadata={"help": "Wether hybrid paralle checkpoints be loaded in auto parallel mode."},
@@ -2173,6 +2177,7 @@ class TrainingArguments:
         """
         Whether or not the current process should produce log.
         """
+        return True
         if self.enable_auto_parallel:
             return True
         elif self.log_on_each_node:

paddlenlp/trainer/utils/ckpt_converter.py

@@ -19,17 +19,17 @@ from functools import reduce
 from typing import List, Union
 
 import paddle
-from paddle.distributed.checkpoint.load_state_dict import (
+from paddle.distributed.fleet.utils.log_util import logger
+from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
     _load_state_dict,
     get_rank_to_read_files,
 )
-from paddle.distributed.checkpoint.metadata import (
+from paddle.distributed.flex_checkpoint.dcp.metadata import (
     LocalTensorIndex,
     LocalTensorMetadata,
     Metadata,
 )
-from paddle.distributed.checkpoint.utils import flatten_state_dict
-from paddle.distributed.fleet.utils.log_util import logger
+from paddle.distributed.flex_checkpoint.dcp.utils import flatten_state_dict
 
 MODEL_WEIGHT_SUFFIX = ".pdparams"
 OPTIMIZER_WEIGHT_SUFFIX = ".pdopt"
@@ -206,7 +206,7 @@ class CheckpointConverter:
                 global_offset = [0] * self.tp_degree
                 for item in shard_info:
                     tp_rank = item[0]["tp_rank"]
-                    state_name_with_tp_rank = state_name + "_tp" + "{:02d}".format(tp_rank)
+                    state_name_with_tp_rank = state_name + "_tp" + f"{tp_rank:02d}"
                     local_tensor_meta_data = LocalTensorMetadata((global_offset[tp_rank],), item[1], item[2])
                     local_tensor_index = LocalTensorIndex(state_name_with_tp_rank, (global_offset[tp_rank],))
                     global_offset[tp_rank] += item[1][0]
@@ -225,7 +225,7 @@ class CheckpointConverter:
                 renamed_state_dict = {}
                 (tp_rank, pp_rank, sharding_rank) = self.get_distribution_rank_from_file_name(file_name)
                 for state_name, state_value in state_dict.items():
-                    state_name_with_tp_rank = state_name + "_tp" + "{:02d}".format(tp_rank)
+                    state_name_with_tp_rank = state_name + "_tp" + f"{tp_rank:02d}"
                     renamed_state_dict[state_name_with_tp_rank] = state_value
 
                 source_state_dict_for_merge_sharding[file_name] = renamed_state_dict
@@ -235,7 +235,7 @@ class CheckpointConverter:
             sharding_metas_keys = []
             for i in range(self.tp_degree):
                 for j in range(self.pp_degree):
-                    sharding_metas_keys.append("tp{:02d}_pp{:02d}".format(i, j))
+                    sharding_metas_keys.append(f"tp{i:02d}_pp{j:02d}")
             for key in sharding_metas_keys:
                 param_meta = self.model_meta["sharding_metas"][key]["param_meta"]
                 for param_name, param_shape_and_dtype in param_meta.items():
@@ -253,7 +253,7 @@ class CheckpointConverter:
             all_param_meta = {}
             for i in range(self.tp_degree):
                 for j in range(self.pp_degree):
-                    key = "tp{:02d}_pp{:02d}".format(i, j)
+                    key = f"tp{i:02d}_pp{j:02d}"
                     param_meta = self.model_meta["sharding_metas"][key]["param_meta"]
                     for param_name, param_shape_and_dtype in param_meta.items():
                         all_param_meta[param_name] = param_shape_and_dtype
@@ -269,7 +269,7 @@ class CheckpointConverter:
             with paddle.base.dygraph.guard(place=paddle.CPUPlace()):
                 for key in cur_rank_need_load_model_state_keys:
                     for tp_rank in range(self.tp_degree):
-                        tp_rank_suffix = "_tp{:02d}".format(tp_rank)
+                        tp_rank_suffix = f"_tp{tp_rank:02d}"
                         optimizer_state_dict[key + ".moment1" + tp_rank_suffix] = paddle.zeros(
                             (param_flattened_shapes[key],), "float32"
                         )
@@ -353,7 +353,7 @@ class CheckpointConverter:
                 else:
                     concat_optimier_state_dict[opt_state_name_removed_tp_rank] = tp_tensors[0]
 
-            fake_file_name = "{:02d}".format(self.cur_rank) + ".distcp"
+            fake_file_name = f"{self.cur_rank:02d}" + ".distcp"
             local_tensor_meta_data = {}
             local_tensor_index = {}
             for k, v in concat_optimier_state_dict.items():
@@ -472,7 +472,7 @@ class CheckpointConverter:
                                 reshaped_v = v.reshape(shape)
                                 target_state_dict[k] = reshaped_v
 
-                fake_file_name = "{:02d}".format(self.cur_rank) + ".distcp"
+                fake_file_name = f"{self.cur_rank:02d}" + ".distcp"
                 local_tensor_meta_data = {}
                 local_tensor_index = {}
                 for k, v in target_state_dict.items():
@@ -911,7 +911,7 @@ class CheckpointConverter:
                 self.model_meta = json.load(file)
 
         (tp_rank, pp_rank, sharding_rank) = self.get_distribution_rank_from_file_name(file_name)
-        dist_strategy_key = "tp" + "{:02d}".format(tp_rank) + "_" + "pp" + "{:02d}".format(pp_rank)
+        dist_strategy_key = "tp" + f"{tp_rank:02d}" + "_" + "pp" + f"{pp_rank:02d}"
         # Map model weight names to their corresponding names of master_weights in the optimizer state.
         if file_name.endswith(OPTIMIZER_WEIGHT_SUFFIX):
             structure_name_mapping = self.model_meta["sharding_metas"][dist_strategy_key]["structure_name_mapping"]

paddlenlp/trainer/utils/load_hf_ckpt.py

@@ -0,0 +1,378 @@
+# 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 json
+import re
+import sys
+from collections import defaultdict
+from typing import List, Optional
+
+import paddle
+
+try:
+    from safetensors import safe_open
+except:
+    safe_open = None
+
+_LAYER_RE = re.compile(r"^_layers\.(\d+)\.(\d+)(?:\.(.*))?$")
+_EXPERT_W1_RE = re.compile(r"^mlp\.experts\.(\d+)\.w1(?:\.weight)?$")
+_EXPERT_W2_RE = re.compile(r"^mlp\.experts\.(\d+)\.w2(?:\.weight)?$")
+_SHARE_EXPERT_W1_RE = re.compile(r"^mlp\.shared_experts\.w1(?:\.weight)?$")
+_SHARE_EXPERT_W2_RE = re.compile(r"^mlp\.shared_experts\.w2(?:\.weight)?$")
+
+_EXPERT_W1_RE_v2 = re.compile(r"^mlp\.experts\.(\d+)\.gate_up_fused_proj(?:\.weight)?$")
+_SHARE_EXPERT_W1_RE_v2 = re.compile(r"^mlp\.shared_experts\.gate_up_fused_proj(?:\.weight)?$")
+_LAYER_RE_v2 = re.compile(r"_layers.deepseek_v2.layers\.(\d+)\.(.*)$")
+
+custom_name_map = {
+    "self_attn.input_layernorm.weight": "input_layernorm.weight",
+    "self_attn.fused_rms_norm_linear.rms_norm_weight": "input_layernorm.weight",
+    "self_attn.memory_recompute_att.kv_ln_weight": "self_attn.kv_a_layernorm.weight",
+    "self_attn.fused_rms_norm_linear.kv_down_weight": "self_attn.kv_a_proj_with_mqa.weight",
+    "self_attn.memory_recompute_att.kv_up_weight": "self_attn.kv_b_proj.weight",
+    "self_attn.memory_recompute_att.q_ln_weight": "self_attn.q_a_layernorm.weight",
+    "self_attn.fused_rms_norm_linear.q_down_weight": "self_attn.q_a_proj.weight",
+    "self_attn.memory_recompute_att.q_up_weight": "self_attn.q_b_proj.weight",
+}
+
+
+def paddle_name_to_hf_names_ds_v2(paddle_name: str) -> List[str]:
+    """
+    将Paddle模型参数名称转换为Hugging Face格式的名称列表
+
+    参数:
+        paddle_name: Paddle格式的参数名称
+
+    返回:
+        Hugging Face格式的参数名称列表(可能拆分多个参数)
+    """
+    if paddle_name == "_layers.deepseek_v2.embed_tokens.weight":
+        return ["model.embed_tokens.weight"]
+
+    if paddle_name == "_layers.deepseek_v2.norm.weight":
+        return ["model.norm.weight"]
+
+    if paddle_name == "_layers.lm_head.weight":
+        return ["lm_head.weight"]
+
+    m = _LAYER_RE_v2.match(paddle_name)
+    if not m:
+        print("not match here !!", paddle_name)
+        return []
+
+    rest = m.group(2) or ""
+    layer_id = m.group(1)
+    if rest in custom_name_map:
+        rest = custom_name_map[rest]
+    out_name = "model.layers." + layer_id + "." + rest
+
+    if rest == "mlp.gate_up_fused_proj.weight" or rest == "mlp.w1":
+        return [
+            "model.layers." + layer_id + ".mlp.gate_proj.weight",
+            "model.layers." + layer_id + ".mlp.up_proj.weight",
+        ]
+
+    if rest == "mlp.w2":
+        return ["model.layers." + layer_id + ".mlp.down_proj.weight"]
+
+    if rest == "mlp.shared_experts.gate_up_fused_proj.weight":
+        return [
+            "model.layers." + layer_id + ".mlp.shared_experts.gate_proj.weight",
+            "model.layers." + layer_id + ".mlp.shared_experts.up_proj.weight",
+        ]
+
+    if m := _EXPERT_W1_RE_v2.match(rest):
+        expert_id = m.group(1)
+        return [
+            "model.layers." + layer_id + ".mlp.experts." + expert_id + ".gate_proj.weight",
+            "model.layers." + layer_id + ".mlp.experts." + expert_id + ".up_proj.weight",
+        ]
+
+    if m := _EXPERT_W1_RE.match(rest):
+        expert_id = m.group(1)
+        return [
+            "model.layers." + layer_id + ".mlp.experts." + expert_id + ".gate_proj.weight",
+            "model.layers." + layer_id + ".mlp.experts." + expert_id + ".up_proj.weight",
+        ]
+
+    if m := _EXPERT_W2_RE.match(rest):
+        expert_id = m.group(1)
+        return ["model.layers." + layer_id + ".mlp.experts." + expert_id + ".down_proj.weight"]
+
+    if m := _SHARE_EXPERT_W1_RE.match(rest):
+        return [
+            "model.layers." + layer_id + ".mlp.shared_experts.gate_proj.weight",
+            "model.layers." + layer_id + ".mlp.shared_experts.up_proj.weight",
+        ]
+
+    if m := _SHARE_EXPERT_W2_RE.match(rest):
+        return ["model.layers." + layer_id + ".mlp.shared_experts.down_proj.weight"]
+
+    return [out_name]
+
+
+def paddle_name_to_hf_names(paddle_name: str) -> List[str]:
+    """
+    将Paddle模型参数名称转换为Hugging Face格式的名称列表
+
+    参数:
+        paddle_name: Paddle格式的参数名称
+
+    返回:
+        Hugging Face格式的参数名称列表（可能拆分多个参数）
+    """
+    if paddle_name == "_layers.local_shared_layers.DeepseekV2_shared_weight.embed_tokens.weight":
+        return ["model.embed_tokens.weight"]
+
+    if paddle_name == "_layers.deepseek_v2.embed_tokens.weight":
+        return ["model.embed_tokens.weight"]
+
+    m = _LAYER_RE.match(paddle_name)
+
+    if not m:
+        print("not match here !!", paddle_name)
+        return []
+    else:
+        rest = m.group(3) or ""
+
+    segment_id = int(m.group(1))
+    id_in_segment = int(m.group(2))
+
+    hf_prefix = _get_hf_prefix(segment_id, id_in_segment)
+
+    if rest in custom_name_map:
+        return [f"{hf_prefix}.{custom_name_map[rest]}"]
+
+    if expert_names := _handle_expert_weights(hf_prefix, rest):
+        return expert_names
+
+    if shared_mlp_names := _handle_shared_expert_weights(hf_prefix, rest):
+        return shared_mlp_names
+
+    if mlp_names := _handle_mlp_weights(hf_prefix, rest):
+        return mlp_names
+
+    if rest == "mlp.gate_up_fused_proj.weight" or rest == "mlp.w1":
+        return [hf_prefix + ".mlp.gate_proj.weight", hf_prefix + ".mlp.up_proj.weight"]
+
+    if rest == "mlp.w2":
+        return [hf_prefix + ".mlp.down_proj.weight"]
+
+    if rest == "mlp.shared_experts.gate_up_fused_proj.weight":
+        return [hf_prefix + ".mlp.shared_experts.gate_proj.weight", hf_prefix + ".mlp.shared_experts.up_proj.weight"]
+
+    if m := _EXPERT_W1_RE_v2.match(rest):
+        expert_id = m.group(1)
+        return [
+            hf_prefix + ".mlp.experts." + expert_id + ".gate_proj.weight",
+            hf_prefix + ".mlp.experts." + expert_id + ".up_proj.weight",
+        ]
+
+    if m := _EXPERT_W1_RE.match(rest):
+        expert_id = m.group(1)
+        return [
+            hf_prefix + ".mlp.experts." + expert_id + ".gate_proj.weight",
+            hf_prefix + ".mlp.experts." + expert_id + ".up_proj.weight",
+        ]
+
+    if m := _EXPERT_W2_RE.match(rest):
+        expert_id = m.group(1)
+        return [hf_prefix + ".mlp.experts." + expert_id + ".down_proj.weight"]
+
+    if m := _SHARE_EXPERT_W1_RE.match(rest):
+        return [hf_prefix + ".mlp.shared_experts.gate_proj.weight", hf_prefix + ".mlp.shared_experts.up_proj.weight"]
+
+    if m := _SHARE_EXPERT_W2_RE.match(rest):
+        return [hf_prefix + ".mlp.shared_experts.down_proj.weight"]
+
+    return [f"{hf_prefix}.{rest}"] if rest else [hf_prefix]
+
+
+def _get_hf_prefix(segment_id: int, id_in_segment: int) -> str:
+    """生成Hugging Face格式的层级前缀"""
+    # 特殊层级映射
+    # special_cases = {(0, 0): "model", (60, 2): "model.layers.61", (60, 3): "model"}
+    # special_cases = {(0, 0): "model", (28, 2): "model.layers.61", (28, 3): "model"}
+    # special_cases = {(0, 0): "model", (28, 2): "model.layers.61", (4, 1): "model"}
+    # special_cases = {(0, 0): "model",  (28, 2): "model", (28,3): "lm_head"}
+    special_cases = {(0, 0): "model", (60, 2): "model.layers.61", (60, 3): "model", (60, 4): "lm_head"}
+
+    if (segment_id, id_in_segment) in special_cases:
+        return special_cases[(segment_id, id_in_segment)]
+
+    # 通用层级计算
+    layer_idx = segment_id + id_in_segment - 1
+    return f"model.layers.{layer_idx}"
+
+
+def _handle_expert_weights(hf_prefix: str, rest: str) -> Optional[List[str]]:
+    if m := _EXPERT_W1_RE.match(rest):
+        expert_id = int(m.group(1))
+        return [
+            f"{hf_prefix}.mlp.experts.{expert_id}.gate_proj.weight",
+            f"{hf_prefix}.mlp.experts.{expert_id}.up_proj.weight",
+        ]
+
+    if m := _EXPERT_W2_RE.match(rest):
+        expert_id = int(m.group(1))
+        return [f"{hf_prefix}.mlp.experts.{expert_id}.down_proj.weight"]
+
+    return None
+
+
+def _handle_shared_expert_weights(hf_prefix: str, rest: str) -> Optional[List[str]]:
+    if _SHARE_EXPERT_W1_RE.match(rest):
+        return [
+            f"{hf_prefix}.mlp.shared_experts.gate_proj.weight",
+            f"{hf_prefix}.mlp.shared_experts.up_proj.weight",
+        ]
+
+    if _SHARE_EXPERT_W2_RE.match(rest):
+        return [f"{hf_prefix}.mlp.shared_experts.down_proj.weight"]
+
+    return None
+
+
+def _handle_mlp_weights(hf_prefix: str, rest: str) -> Optional[List[str]]:
+    if rest == "mlp.w1":
+        return [f"{hf_prefix}.mlp.gate_proj.weight", f"{hf_prefix}.mlp.up_proj.weight"]
+
+    if rest == "mlp.w2":
+        return [f"{hf_prefix}.mlp.down_proj.weight"]
+
+    return None
+
+
+def prepare_tensor(tensor, dst_shape, *, force_transpose=False):
+    if isinstance(tensor, list):
+        t = paddle.concat(
+            [
+                paddle.transpose(tensor[0], perm=[1, 0]).contiguous(),
+                paddle.transpose(tensor[1], perm=[1, 0]).contiguous(),
+            ],
+            axis=-1,
+        )
+        if t.shape != dst_shape:
+            print("base shape", tensor[0].shape, tensor[1].shape)
+            print("shape not match ", t.shape, dst_shape)
+            sys.exit()
+        return t
+
+    if force_transpose:
+        return tensor.T.contiguous()
+
+    if tensor.shape == dst_shape:
+        if len(tensor.shape) != 1:
+            print("attention same shape not transpose !!!!!!!!!!!!!!!!!!!!!!")
+        return tensor
+    if len(tensor.shape) == 2 and paddle.transpose(tensor, perm=[1, 0]).contiguous().shape == dst_shape:
+        return paddle.transpose(tensor, perm=[1, 0]).contiguous()
+
+    print("shape not match here")
+    sys.exit()
+
+
+def load_huggingface_ckpt(model, huggingface_ckpt_path):
+    ckpt_pre = huggingface_ckpt_path
+
+    # 1. 加载参数-文件映射表
+    weight_map_path = ckpt_pre + "/model.safetensors.index.json"
+    with open(weight_map_path, "r") as f:
+        weight_map = json.load(f)["weight_map"]
+
+    # 2. 创建反向索引：文件 -> 参数列表
+    file_to_params = defaultdict(list)
+    for param_name, filename in weight_map.items():
+        file_to_params[filename].append(param_name)
+
+    # 2. 收集模型需要的文件列表
+    required_files = set()
+    file_to_pd_param_name = defaultdict(list)
+    pd_param_name_to_file = defaultdict(list)
+    for pd_name, p in model.named_parameters():
+        hf_name = paddle_name_to_hf_names(pd_name)
+        if hf_name[0] in weight_map:
+            filename = weight_map[hf_name[0]]
+            required_files.add(filename)
+            file_to_pd_param_name[filename].append(pd_name)
+            pd_param_name_to_file[pd_name].append(filename)
+        else:
+            print(f"Warning: {pd_name} -> {hf_name[0]} not found in weight map")
+            import sys
+
+            sys.exit()
+
+        if len(hf_name) > 1:
+            if hf_name[1] in weight_map:
+                filename = weight_map[hf_name[1]]
+                required_files.add(filename)
+                file_to_pd_param_name[filename].append(pd_name)
+                if filename != pd_param_name_to_file[pd_name][0]:
+                    pd_param_name_to_file[pd_name].append(filename)
+            else:
+                print(f"Warning: {pd_name} -> {hf_name[1]} not found in weight map")
+
+    # 3. 按文件分组加载
+    check_list = []
+    print("Start load huggingface ckpt")
+    for i, filename in enumerate(required_files):
+        try:
+            with safe_open(ckpt_pre + filename, framework="paddle", device="cpu") as f:
+                # 加载该文件包含的所有参数
+                pd_params = file_to_pd_param_name[filename]
+                for pd_param in pd_params:
+                    if pd_param in check_list:
+                        continue
+
+                    hf_name = paddle_name_to_hf_names(pd_param)
+                    if len(hf_name) == 1:
+                        tensor = f.get_tensor(hf_name[0])
+
+                        force_transpose = False
+
+                        model.state_dict()[pd_param].set_value(
+                            paddle.cast(
+                                prepare_tensor(
+                                    tensor, model.state_dict()[pd_param].shape, force_transpose=force_transpose
+                                ),
+                                model.state_dict()[pd_param].dtype,
+                            )
+                        )
+                    else:
+                        files = pd_param_name_to_file[pd_param]
+                        if len(files) == 1:
+                            tensor0 = f.get_tensor(hf_name[0])
+                            tensor1 = f.get_tensor(hf_name[1])
+                        else:
+                            if weight_map[hf_name[0]] == filename:
+                                tensor0 = f.get_tensor(hf_name[0])
+                                with safe_open(
+                                    ckpt_pre + weight_map[hf_name[1]], framework="paddle", device="cpu"
+                                ) as f_other:
+                                    tensor1 = f_other.get_tensor(hf_name[1])
+                            else:
+                                with safe_open(
+                                    ckpt_pre + weight_map[hf_name[0]], framework="paddle", device="cpu"
+                                ) as f_other:
+                                    tensor0 = f_other.get_tensor(hf_name[0])
+                                tensor1 = f.get_tensor(hf_name[1])
+                        model.state_dict()[pd_param].set_value(
+                            prepare_tensor([tensor0, tensor1], model.state_dict()[pd_param].shape)
+                        )
+                    check_list.append(pd_param)
+
+        except Exception as e:
+            print(f"Error loading {filename}: {str(e)}")
+            raise

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -179,6 +179,20 @@ class DeepseekV2Config(PretrainedConfig):
         attention_dropout=0.0,
         speculate_model_type=False,
         using_flex_token=False,
+        use_dualpipev=False,
+        send_mtp_embed=False,
+        using_post_norm_recompute=False,
+        stepped_recompute_fwd_gate_up=False,
+        recompute_fwd_gate_up=0,
+        recompute_fa3=0,
+        is_split_group_gemm=False,
+        fakse_gate_restrict_balance=False,
+        adaptive_remained_O1_recompute_ratio=0,
+        offline_quant_expert_weight=True,
+        clear_origin_weight_when_offline_quant=True,
+        mlp_bwd_subbatch_rows=0,
+        mlp_fwd_subbatch_rows=0,
+        output_subbatch_rows=0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -227,6 +241,20 @@ class DeepseekV2Config(PretrainedConfig):
         self.speculate_model_type = speculate_model_type
         self.use_fp8 = False
         self.using_flex_token = using_flex_token
+        self.use_dualpipev = use_dualpipev
+        self.send_mtp_embed = send_mtp_embed
+        self.using_post_norm_recompute = using_post_norm_recompute
+        self.recompute_fwd_gate_up = recompute_fwd_gate_up
+        self.recompute_fa3 = recompute_fa3
+        self.stepped_recompute_fwd_gate_up = stepped_recompute_fwd_gate_up
+        self.is_split_group_gemm = is_split_group_gemm
+        self.fakse_gate_restrict_balance = fakse_gate_restrict_balance
+        self.adaptive_remained_O1_recompute_ratio = adaptive_remained_O1_recompute_ratio
+        self.offline_quant_expert_weight = offline_quant_expert_weight
+        self.clear_origin_weight_when_offline_quant = clear_origin_weight_when_offline_quant
+        self.mlp_bwd_subbatch_rows = mlp_bwd_subbatch_rows
+        self.mlp_fwd_subbatch_rows = mlp_fwd_subbatch_rows
+        self.output_subbatch_rows = output_subbatch_rows
 
         super().__init__(
             pad_token_id=pad_token_id,

paddlenlp/transformers/deepseek_v2/fp8_linear.py

@@ -18,6 +18,8 @@ original_linear = paddle.nn.functional.linear
 
 from typing import Literal, Optional
 
+import paddle.nn.functional as F
+
 # from ..linear_utils import RowParallelLinear as PD_RowParallelLinear
 from ..linear_utils import ColumnParallelLinear as PD_ColumnParallelLinear
 from ..linear_utils import (
@@ -32,6 +34,15 @@ try:
 except:
     pass
 
+try:
+    from paddle.incubate.nn.functional import swiglu
+except ImportError:
+
+    def swiglu(x, y=None):
+        if y is None:
+            x, y = paddle.chunk(x, chunks=2, axis=-1)
+        return F.silu(x) * y
+
 
 __all__ = [
     "Linear",

paddlenlp/transformers/deepseek_v2/kernel.py

@@ -59,6 +59,8 @@ def act_quant(x: paddle.Tensor, block_size: int = 128) -> Tuple[paddle.Tensor, p
             - The quantized tensor with dtype `paddle.float8_e4m3fn`.
             - A tensor of scaling factors with dtype `paddle.float32`.
     """
+    if not x.is_contiguous():
+        x = x.contiguous()
     assert x.is_contiguous(), "Input tensor must be contiguous"
     assert (
         x.shape[-1] % block_size == 0

paddlenlp/transformers/fused_a2a.py

@@ -72,78 +72,145 @@ def get_buffer(group: Group, hidden_bytes: int):
     return _buffer
 
 
+def fused_dispatch_forward_func(
+    x,
+    token_indices,
+    token_probs,
+    num_experts,
+    group,
+    previous_event=None,
+    async_finish=False,
+    allocate_on_comm_stream=False,
+):
+    """Forward pass of fused dispatch."""
+    # Calculate layout before actual dispatch
+    if isinstance(x, tuple):
+        buffer = get_buffer(group, get_hidden_bytes(x[0]))
+    else:
+        buffer = get_buffer(group, get_hidden_bytes(x))
+    (
+        num_tokens_per_rank,
+        num_tokens_per_rdma_rank,
+        num_tokens_per_expert,
+        is_token_in_rank,
+        previous_event_,
+    ) = buffer.get_dispatch_layout(
+        token_indices,
+        num_experts,
+        previous_event=previous_event,
+        async_finish=async_finish,
+        allocate_on_comm_stream=allocate_on_comm_stream,
+    )
+
+    assert token_probs.dtype == paddle.float32
+    # Do MoE dispatch
+    # NOTES: the CPU will wait for GPU's signal to arrive,
+    # so this is not compatible with CUDA graph
+    (recv_x, recv_token_indices, recv_token_probs, num_recv_tokens_per_expert_list, handle, event,) = buffer.dispatch(
+        x,
+        topk_idx=token_indices,
+        topk_weights=token_probs,
+        num_tokens_per_rank=num_tokens_per_rank,
+        num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
+        is_token_in_rank=is_token_in_rank,
+        num_tokens_per_expert=num_tokens_per_expert,
+        previous_event=previous_event,
+        async_finish=async_finish,
+        allocate_on_comm_stream=allocate_on_comm_stream,
+    )
+
+    states = dict()
+    states["dispatched_indices"] = recv_token_indices
+    states["tokens_per_expert"] = num_recv_tokens_per_expert_list
+    states["handle"] = handle
+
+    return recv_x, recv_token_probs, states, event
+
+
+def fused_dispatch_backward_func(
+    grad_output,
+    grad_token_probs,
+    group,
+    handle,
+    previous_event=None,
+    async_finish=False,
+    allocate_on_comm_stream=False,
+):
+    """Backward pass of fused dispatch."""
+    buffer = get_buffer(group, get_hidden_bytes(grad_output))
+
+    grad_x, grad_token_probs, event = buffer.combine(
+        grad_output.contiguous(),
+        handle,
+        topk_weights=grad_token_probs.cast(paddle.float32),
+        previous_event=previous_event,
+        async_finish=async_finish,
+        allocate_on_comm_stream=allocate_on_comm_stream,
+    )
+    return grad_x, None, grad_token_probs
+
+
+def fused_combine_forward_func(
+    x, group, states, previous_event=None, async_finish=False, allocate_on_comm_stream=False
+):
+    """Forward pass of fused combine."""
+    handle = states["handle"]
+    buffer = get_buffer(group, get_hidden_bytes(x))
+    combined_x, _, event = buffer.combine(
+        x,
+        handle=handle,
+        async_finish=async_finish,
+        previous_event=previous_event,
+        allocate_on_comm_stream=allocate_on_comm_stream,
+    )
+    return combined_x
+
+
+def fused_combine_backward_func(
+    grad_output, group, handle, previous_event=None, async_finish=False, allocate_on_comm_stream=False
+):
+    """Backward pass of fused combine."""
+    if isinstance(grad_output, tuple):
+        buffer = get_buffer(group, get_hidden_bytes(grad_output[0]))
+        grad_x, _, _, _, _, event = buffer.dispatch(
+            (grad_output[0].contiguous(), grad_output[1].contiguous()),
+            handle=handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+    else:
+        buffer = get_buffer(group, get_hidden_bytes(grad_output))
+        grad_x, _, _, _, _, event = buffer.dispatch(
+            grad_output.contiguous(),
+            handle=handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+    return grad_x
+
+
 class FusedDispatch(PyLayer):
     """Fused dispatch operation for MoE routing combining computation and communication."""
 
     @staticmethod
     def forward(ctx, x, token_indices, token_probs, num_experts, group, previous_event=None):
         """Forward pass of fused dispatch."""
-        # Calculate layout before actual dispatch
-        buffer = get_buffer(group, get_hidden_bytes(x))
-        (
-            num_tokens_per_rank,
-            num_tokens_per_rdma_rank,
-            num_tokens_per_expert,
-            is_token_in_rank,
-            previous_event,
-        ) = buffer.get_dispatch_layout(
-            token_indices,
-            num_experts,
-            previous_event=None,
-            async_finish=False,
-            allocate_on_comm_stream=False,
-        )
-
-        # Do MoE dispatch
-        # NOTES: the CPU will wait for GPU's signal to arrive,
-        # so this is not compatible with CUDA graph
-        (
-            recv_x,
-            recv_token_indices,
-            recv_token_probs,
-            num_recv_tokens_per_expert_list,
-            handle,
-            event,
-        ) = buffer.dispatch(
-            x,
-            topk_idx=token_indices,
-            topk_weights=token_probs.cast(paddle.float32),
-            num_tokens_per_rank=num_tokens_per_rank,
-            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
-            is_token_in_rank=is_token_in_rank,
-            num_tokens_per_expert=num_tokens_per_expert,
-            previous_event=None,
-            async_finish=False,
-            allocate_on_comm_stream=False,
+        recv_x, recv_token_probs, states, event = fused_dispatch_forward_func(
+            x, token_indices, token_probs, num_experts, group, previous_event
         )
 
         ctx.group = group
-        ctx.handle = handle
+        ctx.handle = states["handle"]
         ctx.event = event
-        tokens_per_expert = paddle.to_tensor(num_recv_tokens_per_expert_list)
-
-        states = dict()
-        states["dispatched_indices"] = recv_token_indices
-        states["tokens_per_expert"] = tokens_per_expert
-        states["handle"] = handle
 
         return recv_x, recv_token_probs, states
 
     @staticmethod
     def backward(ctx, grad_output, grad_token_probs):
         """Backward pass of fused dispatch."""
-        buffer = get_buffer(ctx.group, get_hidden_bytes(grad_output))
-        handle = ctx.handle
-
-        grad_x, grad_token_probs, event = buffer.combine(
-            grad_output.contiguous(),
-            handle,
-            topk_weights=grad_token_probs.cast(paddle.float32),
-            previous_event=None,
-            async_finish=False,
-            allocate_on_comm_stream=False,
-        )
-        return grad_x, None, grad_token_probs
+        return fused_dispatch_backward_func(grad_output, grad_token_probs, ctx.group, ctx.handle)
 
 
 class FusedCombine(PyLayer):
@@ -152,12 +219,9 @@ class FusedCombine(PyLayer):
     @staticmethod
     def forward(ctx, x, group, states, previous_event=None):
         """Forward pass of fused combine."""
-        handle = states["handle"]
-        buffer = get_buffer(group, get_hidden_bytes(x))
-        combined_x, _, event = buffer.combine(
-            x, handle=handle, async_finish=False, previous_event=None, allocate_on_comm_stream=False
-        )
-        ctx.handle = handle
+        combined_x = fused_combine_forward_func(x, group, states, previous_event)
+
+        ctx.handle = states["handle"]
         ctx.group = group
         ctx.previous_event = previous_event
 
@@ -166,15 +230,7 @@ class FusedCombine(PyLayer):
     @staticmethod
     def backward(ctx, grad_output):
         """Backward pass of fused combine."""
-        buffer = get_buffer(ctx.group, get_hidden_bytes(grad_output))
-        grad_x, _, _, _, _, event = buffer.dispatch(
-            grad_output.contiguous(),
-            handle=ctx.handle,
-            previous_event=ctx.previous_event,
-            async_finish=False,
-            allocate_on_comm_stream=False,
-        )
-        return grad_x
+        return fused_combine_backward_func(grad_output, ctx.group, ctx.handle, ctx.previous_event)
 
 
 if HAVE_DEEP_EP:
@@ -214,3 +270,96 @@ if HAVE_DEEP_EP:
 else:
     fused_dispatch = None
     fused_combine = None
+
+
+class DispatchNode:
+    def __init__(self, name="dispatch"):
+        self.name = name
+
+    def reset_statue(self):
+        self.handle = None
+
+    def forward(
+        self,
+        x,
+        token_indices,
+        token_probs,
+        num_experts,
+        group,
+        previous_event=None,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
+        """Forward pass of fused dispatch."""
+        recv_x, recv_token_probs, states, event = fused_dispatch_forward_func(
+            x,
+            token_indices,
+            token_probs,
+            num_experts,
+            group,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+
+        self.group = group
+        self.handle = states["handle"]
+        self.event = event
+
+        return recv_x, recv_token_probs, states
+
+    def backward(
+        self, grad_output, grad_token_probs, previous_event=None, async_finish=False, allocate_on_comm_stream=False
+    ):
+        """Backward pass of fused dispatch."""
+        out = fused_dispatch_backward_func(
+            grad_output,
+            grad_token_probs,
+            self.group,
+            self.handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        self.reset_statue()
+        return out
+
+
+class CombineNode:
+    def __init__(self, name="combine"):
+        self.name = name
+
+    def reset_statue(self):
+        self.handle = None
+
+    def forward(self, x, group, handle, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
+        """Forward pass of fused combine."""
+        states = dict()
+        states["handle"] = handle
+        combined_x = fused_combine_forward_func(
+            x,
+            group,
+            states,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+
+        self.handle = handle
+        self.group = group
+        self.previous_event = previous_event
+
+        return combined_x
+
+    def backward(self, grad_output, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
+        """Backward pass of fused combine."""
+        out = fused_combine_backward_func(
+            grad_output,
+            self.group,
+            self.handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        self.reset_statue()
+        return out

paddlenlp/transformers/moe_gate.py

@@ -270,7 +270,7 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
 
         group_scores = scores.reshape([0, n_group, -1]).max(axis=-1)  # [n, n_group]
         group_idx = paddle.topk(group_scores, k=topk_group, axis=-1, sorted=True)[1]  # [n, top_k_group]
-        group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.to_tensor(1.0), axis=-1)  # fmt:skip
+        group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.ones([], dtype="float32"), axis=-1)  # fmt:skip
         score_mask = (
             group_mask.unsqueeze(-1).expand([bsz_seq_len, n_group, n_experts // n_group]).reshape([bsz_seq_len, -1])
         )  # [n, e]
@@ -302,11 +302,11 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
 
         assert self.e_score_correction_bias is not None, "e_score_correction_bias is None"
         scores_for_choice = scores.reshape([bsz_seq_len, -1]) + self.e_score_correction_bias.unsqueeze(0)
-        group_scores = (
-            scores_for_choice.reshape([bsz_seq_len, self.n_group, -1]).topk(2, axis=-1)[0].sum(axis=-1)
-        )  # fmt:skip [n, n_group]
+        reshape_tmp_rst = scores_for_choice.reshape([bsz_seq_len, self.n_group, -1])
+        top_k = min(reshape_tmp_rst.shape[2], 2)
+        group_scores = reshape_tmp_rst.topk(top_k, axis=-1)[0].sum(axis=-1)  # fmt:skip [n, n_group]
         group_idx = paddle.topk(group_scores, k=topk_group, axis=-1, sorted=True)[1]  # [n, top_k_group]
-        group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.to_tensor(1.0, dtype="float32"), axis=-1)  # fmt:skip
+        group_mask = paddle.zeros_like(group_scores).put_along_axis(group_idx, paddle.ones([], dtype="float32"), axis=-1)  # fmt:skip
         score_mask = (
             group_mask.unsqueeze(-1).expand([bsz_seq_len, n_group, n_experts // n_group]).reshape([bsz_seq_len, -1])
         )  # [n, e]
@@ -370,9 +370,7 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
 
         _, top_idx = paddle.topk(mask1_rand, k=capacity, axis=0)  # Select top_capacity tokens
 
-        new_mask1 = mask1 * paddle.zeros_like(mask1).put_along_axis(
-            top_idx, paddle.to_tensor(1.0, dtype="float32"), axis=0
-        )
+        new_mask1 = mask1 * paddle.zeros_like(mask1).put_along_axis(top_idx, paddle.ones([], dtype="float32"), axis=0)
         mask1 = new_mask1
 
         # Compute locations in capacity buffer
@@ -496,7 +494,7 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
         top_gate = top_gate * self.routed_scaling_factor
 
         # get topk mask
-        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0, dtype="float32"), axis=1)
+        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.ones([], dtype="float32"), axis=1)
         if hasattr(self.config, "seq_aux") and self.config.seq_aux:
             l_aux = self._cal_seq_aux_loss(gates_ori, self.top_k, top_idx)
         else:
@@ -567,14 +565,25 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
             top_gate, top_idx = self._topk_noaux_tc(
                 gates, k=self.top_k, n_group=self.n_group, topk_group=self.topk_group
             )
+
             # norm gate to sum 1
-        if self.top_k > 1 and self.norm_topk_prob:
-            denominator = top_gate.sum(axis=-1, keepdim=True) + 1e-20
-            top_gate = top_gate / denominator
-        top_gate = top_gate * self.routed_scaling_factor
+        # if self.top_k > 1 and self.norm_topk_prob:
+        #     denominator = top_gate.sum(axis=-1, keepdim=True) + 1e-20
+        #     top_gate = top_gate / denominator
+        # top_gate = top_gate * self.routed_scaling_factor
 
         # get topk mask
-        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0), axis=1)
+        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.ones([], dtype="float32"), axis=1)
+
+        gates_masked = gates * mask
+        # if self.training:
+        gates_s = paddle.sum(gates_masked, axis=-1, keepdim=True)
+        denom_s = paddle.clip(gates_s, min=paddle.finfo(gates_masked.dtype).eps)
+
+        if self.norm_topk_prob:
+            gates_masked = gates_masked / denom_s
+        
+        gates_masked *= self.routed_scaling_factor
 
         if hasattr(self.config, "seq_aux") and self.config.seq_aux:
             l_aux = self._cal_seq_aux_loss(gates_ori, self.top_k, top_idx)
@@ -582,5 +591,5 @@ class PretrainedMoEGate(nn.Layer, MoEGateMixin):
             l_aux = self._cal_aux_loss(gates, mask)
 
         exp_counts = paddle.sum(mask.cast(paddle.int64), axis=0)
-        topk_masked_gates = paddle.zeros_like(gates).put_along_axis(top_idx, top_gate, axis=1)
-        return topk_masked_gates, mask, exp_counts, l_aux, l_zloss
+        # topk_masked_gates = paddle.zeros_like(gates).put_along_axis(top_idx, top_gate, axis=1)
+        return gates_masked, mask, exp_counts, l_aux, l_zloss

paddlenlp/transformers/moe_layer.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import os
 from typing import Any, List, Tuple
 
 import numpy as np
@@ -24,8 +25,50 @@ import paddle.distributed as dist
 from paddle import Tensor, nn
 from paddle.distributed.communication.group import Group
 
+from ..utils.log import logger
+from .fp8_utils import FP8GroupGemmMlpFunctionNode, extract_first_if_tuple
+from .fused_a2a import CombineNode, DispatchNode, get_buffer, get_hidden_bytes
 from .moe_gate import PretrainedMoEGate
-from .token_dispatcher import MoEFlexTokenDispatcher
+from .moe_utils import (
+    UnZipNode,
+    ZipNode,
+    merge_subbatch_cast,
+    offload,
+    reload,
+    tokens_zip_unique_add_with_subbatch,
+)
+from .token_dispatcher import MoEFlexTokenDispatcher, PreDispatchNode
+
+try:
+    import paddle.distributed.communication.deep_ep as deep_ep
+except ImportError:
+    deep_ep = None
+
+DSV3_USE_FP8_GEMM = os.getenv("DSV3_USE_FP8_GEMM", "False").lower() == "true"
+
+DSV3_USE_FP8_GROUP_GEMM = os.getenv("DSV3_USE_FP8_GROUP_GEMM", "False").lower() == "true"
+
+DSV3_USE_FP8_DISPATCH = os.getenv("DSV3_USE_FP8_DISPATCH", "False").lower() == "true"
+
+try:
+    import TokenDispatcherUtils as TDU
+except ImportError:
+    TDU = None
+
+
+def record_stream_for_multi_input(x):
+    if isinstance(x, (tuple, list)):
+        for i in range(len(x)):
+            x[i]._record_stream()
+    else:
+        x._record_stream()
+
+
+def stop_gradient_for_multi_input(x):
+    if isinstance(x, (tuple, list)):
+        x[0].stop_gradient = False
+    else:
+        x.stop_gradient = False
 
 
 def dispatching(x, dispatch_mask, scatter_index, num_experts, capacity):
@@ -162,6 +205,7 @@ class MoELayer(nn.Layer):
         capacity: int = 1.0,
         moe_group: str = "data",
         all_to_all_dropout=0.0,
+        using_post_norm_recompute=False,
     ):
         super().__init__()
 
@@ -176,8 +220,11 @@ class MoELayer(nn.Layer):
         except AttributeError:
             is_fleet_init = False
 
-        if is_fleet_init and dist.get_world_size() > 1 and moe_group == "data":
-            self.moe_group = dist.fleet.get_hybrid_communicate_group().get_data_parallel_group()
+        if is_fleet_init and dist.get_world_size() > 1:
+            if moe_group == "data":
+                self.moe_group = dist.fleet.get_hybrid_communicate_group().get_data_parallel_group()
+            elif moe_group == "expert":
+                self.moe_group = dist.fleet.get_hybrid_communicate_group().expert_parallel_group
             self.moe_rank = dist.get_rank(self.moe_group)
             self.moe_rank = 0 if self.moe_rank < 0 else self.moe_rank
             self.expert_parallel_degree = dist.get_world_size(self.moe_group)
@@ -206,8 +253,32 @@ class MoELayer(nn.Layer):
 
         self.gate = gate
         self.gate.group = self.moe_group
+        # for flex token moe layer
+        self.router = gate
+        self.ep_size = dist.get_world_size(self.moe_group)
+        self.moe_router_topk = gate.top_k
+        self.num_local_experts = moe_num_experts // self.ep_size
+        self.token_dispatcher = MoEFlexTokenDispatcher(
+            self.num_local_experts, self.moe_router_topk, self.moe_num_experts, self.moe_group
+        )
+        self.token_drop_steps = config.token_drop_steps
+        self.using_flex_token = False
+
+        self.using_post_norm_recompute = using_post_norm_recompute
         self._post_init()
 
+    def update_flex_token(self):
+        from paddlenlp.transformers.deepseek_v2 import get_global_step
+
+        if (not self.config.using_flex_token) or (get_global_step() < self.token_drop_steps):
+            self.using_flex_token = False
+            self.router.using_flex_token = False
+        else:
+            if not self.using_flex_token:
+                logger.info("Changing to flex token moe mode")
+            self.using_flex_token = True
+            self.router.using_flex_token = True
+
     def _parse_moe_expert_parallel(self, moe_num_experts, expert_parallel_degree):
         assert (
             moe_num_experts >= expert_parallel_degree
@@ -230,8 +301,35 @@ class MoELayer(nn.Layer):
                     # logger.info(f"expert param={p.name}, no-sync={p.no_sync}")
 
     def forward(
+        self,
+        hidden_states: paddle.Tensor,
+        probs=None,
+        routing_map=None,
+        capacity=None,
+        topk_weight=None,
+        topk_ids=None,
+        token_priority=None,
+        l_aux=None,
+        l_zloss=None,
+    ):
+        self.update_flex_token()
+
+        if self.using_flex_token:
+            return self.forward_flex_token(hidden_states, probs, routing_map, l_aux, l_zloss)
+        else:
+            return self.forward_drop_token(
+                hidden_states, capacity, topk_weight, topk_ids, token_priority, l_aux, l_zloss
+            )
+
+    def forward_drop_token(
         self,
         hidden_state: paddle.Tensor,
+        capacity=None,
+        topk_weight=None,
+        topk_ids=None,
+        token_priority=None,
+        l_aux=None,
+        l_zloss=None,
     ):
         """MoE Layer forward function
             1. Gate Forward.
@@ -253,7 +351,17 @@ class MoELayer(nn.Layer):
         # topk_ids    : sk
         # token_priority    : se
         # self.exp_counts  :
-        capacity, topk_weight, topk_ids, token_priority, l_aux, l_zloss = self.gate(hidden_state)
+        if self.using_post_norm_recompute:
+            assert (
+                capacity is not None
+                and topk_weight is not None
+                and topk_ids is not None
+                and token_priority is not None
+                and l_aux is not None
+                and l_zloss is not None
+            )
+        else:
+            capacity, topk_weight, topk_ids, token_priority, l_aux, l_zloss = self.gate(hidden_state)
 
         """MoE expert dispatch from: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py"""
         cnts = paddle.zeros([topk_ids.shape[0], len(self.experts)], dtype=topk_ids.dtype)
@@ -332,43 +440,781 @@ class MoELayer(nn.Layer):
 
         return final_out, l_aux, l_zloss
 
+    def forward_flex_token(self, hidden_states: paddle.Tensor, probs=None, routing_map=None, l_aux=None, l_zloss=None):
+        _, _, d_model = hidden_states.shape
+        # reshaped_input = hidden_states.reshape([-1, d_model])
+        if self.using_post_norm_recompute:
+            assert probs is not None and routing_map is not None and l_aux is not None and l_zloss is not None
+        else:
+            probs, routing_map, l_aux, l_zloss = self.router(hidden_states)
+        if DSV3_USE_FP8_GEMM:
+            if DSV3_USE_FP8_DISPATCH:
+                output = FusionMoe.apply(
+                    hidden_states,
+                    probs,
+                    routing_map,
+                    self,
+                    recompute_fwd_gate_up=self.config.recompute_fwd_gate_up,
+                    is_split_group_gemm=self.config.is_split_group_gemm,
+                )
+            else:
+                hidden_states, token_indices, token_probs = self.token_dispatcher.pre_dispatch(
+                    hidden_states, probs, routing_map
+                )
+                output = FusionMoe.apply(
+                    hidden_states,
+                    token_indices,
+                    token_probs,
+                    self,
+                    recompute_fwd_gate_up=self.config.recompute_fwd_gate_up,
+                    is_split_group_gemm=self.config.is_split_group_gemm,
+                )
+        else:
+            (
+                dispatched_input,
+                token_permuted_indices,
+                prob_permuted_indices,
+                dispatched_probs,
+            ) = self.token_dispatcher.token_permutation(hidden_states, probs, routing_map)
 
-class MoEFlexTokenLayer(nn.Layer):
-    def __init__(self, config, moe_num_experts, expert_class, expert_kwargs, gate, moe_group):
+            expert_output = self.expert_forward(dispatched_input)
+            output, _ = self.token_dispatcher.token_unpermutation(
+                expert_output, token_permuted_indices, prob_permuted_indices, dispatched_probs, None
+            )
+        return output, l_aux, l_zloss
 
-        super().__init__()
-        self.config = config
-        self.moe_group = moe_group
-        self.ep_size = dist.get_world_size(self.moe_group)
-        self.moe_router_topk = gate.top_k
-        self.moe_num_experts = moe_num_experts
-        self.num_local_experts = moe_num_experts // self.ep_size
-        self.token_dispatcher = MoEFlexTokenDispatcher(
-            self.num_local_experts, self.moe_router_topk, self.moe_num_experts, moe_group
-        )
+    def get_tokens_per_expert(self):
+        return self.token_dispatcher._comm_manager.tokens_per_expert_list
 
-        self.experts = nn.LayerList([expert_class(**expert_kwargs)] * self.num_local_experts)
-        self.router = gate
+    def set_tokens_per_expert(self, tokens_per_expert_list):
+        self.token_dispatcher._comm_manager.tokens_per_expert_list = tokens_per_expert_list
 
-    def expert_forward(self, dispatched_input, tokens_per_expert):
+    def expert_forward(self, dispatched_input):
         outputs = []
-        tokens_per_expert = tokens_per_expert.tolist()
-        # print(f"all tokens: {sum(tokens_per_expert)}, detail: {tokens_per_expert}")
-        chunks = paddle.split(dispatched_input, num_or_sections=tokens_per_expert, axis=0)
-        for chunk, expert in zip(chunks, self.experts):
+        chunks = paddle.split(dispatched_input, num_or_sections=self.get_tokens_per_expert(), axis=0)
+        for i, chunk in enumerate(chunks):
             chunk = chunk.contiguous()
             # assert chunk.shape[0] != 0, "Cannot dispatch empty input"
+            expert = self.experts[i + self.moe_rank * self.moe_num_experts_per_device]
             outputs += [expert(chunk)]
 
         return paddle.concat(outputs, axis=0)
 
-    def forward(self, hidden_states: paddle.Tensor):
+    def pre_dispatch_compute(self, hidden_states):
         _, _, d_model = hidden_states.shape
-        # reshaped_input = hidden_states.reshape([-1, d_model])
         probs, routing_map, l_aux, l_zloss = self.router(hidden_states)
-        (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation(
+        hidden_states, token_indices, token_probs = self.token_dispatcher.pre_dispatch(
             hidden_states, probs, routing_map
         )
-        expert_output = self.expert_forward(dispatched_input, tokens_per_expert)
-        output, _ = self.token_dispatcher.token_unpermutation(expert_output, None)
-        return output, l_aux, l_zloss
+        return l_aux, l_zloss, hidden_states, token_indices, token_probs
+
+    def post_dispatch_compute(self, hidden_states, dispatched_indices, dispatched_probs):
+        (global_input_tokens, token_permuted_indices, prob_permuted_indices) = self.token_dispatcher.post_dispatch(
+            hidden_states, dispatched_indices
+        )
+        return (global_input_tokens, token_permuted_indices, prob_permuted_indices)
+
+    def pre_combine_compute(self, hidden_states, token_permuted_indices, prob_permuted_indices, dispatched_probs):
+        hidden_states = self.token_dispatcher.pre_combine(
+            hidden_states, token_permuted_indices, prob_permuted_indices, dispatched_probs
+        )
+        return hidden_states
+
+    def post_combine_compute(self, hidden_states):
+        hidden_states = self.token_dispatcher.post_combine(hidden_states)
+        return hidden_states
+
+
+class Fp8DispatchQuantNode:
+    def __init__(self, token_dispatcher, name="fp8_dispatch_quant_node"):
+        self.token_dispatcher = token_dispatcher
+        self.pre_dispatch_node = PreDispatchNode(token_dispatcher)
+        self.name = name
+
+    @paddle.no_grad()
+    def forward(self, hidden_states, probs, routing_map):
+        # reshape
+        self.token_dispatcher.hidden_shape = hidden_states.shape
+        hs_2d = hidden_states.view([-1, self.token_dispatcher.hidden_shape[-1]])
+
+        if DSV3_USE_FP8_DISPATCH:
+            # quant
+            hs_fp8, hs_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+                hs_2d, output_scale_transpose=False, quant_method="1x128", input_transpose=False
+            )
+
+            # pre_dispatch
+            token_indices, token_probs = self.pre_dispatch_node.forward(routing_map, probs)
+
+            self.hidden_states_shape = hidden_states.shape
+            hs_fp8.stop_gradient = False
+            token_probs.stop_gradient = False
+            return (hs_fp8, hs_scale), token_indices, token_probs
+        else:
+            # pre_dispatch
+            token_indices, token_probs = self.pre_dispatch_node.forward(routing_map, probs)
+
+            self.hidden_states_shape = hidden_states.shape
+            hs_2d.stop_gradient = False
+            token_probs.stop_gradient = False
+            return hs_2d, token_indices, token_probs
+
+    @paddle.no_grad()
+    def backward(self, hs_grad, token_probs_grad):
+        # predispatch grad
+        probs_grad = self.pre_dispatch_node.backward(token_probs_grad)
+        token_probs_grad._record_stream()
+
+        # reshape_grad
+        hs_grad = hs_grad.view(self.hidden_states_shape)
+        hs_grad._record_stream()
+
+        return hs_grad, probs_grad, None
+
+
+class Fp8DispatchNode:
+    def __init__(self, token_dispatcher, name="fp8_dispatch_node"):
+        self.token_dispatcher = token_dispatcher
+        self.dispatch_act_node = DispatchNode(token_dispatcher)
+        self.name = name
+
+    @paddle.no_grad()
+    def forward(
+        self,
+        hs_2d,
+        token_indices,
+        token_probs,
+        previous_event=None,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
+        # dispatch
+        hs_2d_dispatched, dispatched_probs, states = self.dispatch_act_node.forward(
+            hs_2d,
+            token_indices,
+            token_probs,
+            self.token_dispatcher._comm_manager.num_experts,
+            self.token_dispatcher._comm_manager.group,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        self.token_dispatcher._comm_manager.handle = states["handle"]
+        self.token_dispatcher._comm_manager.tokens_per_expert = states["tokens_per_expert"]
+        dispatched_indices = states["dispatched_indices"]
+
+        stop_gradient_for_multi_input(hs_2d_dispatched)
+        dispatched_probs.stop_gradient = False
+        return hs_2d_dispatched, dispatched_indices, dispatched_probs
+
+    @paddle.no_grad()
+    def backward(
+        self,
+        hs_dispatched_grad,
+        dispatched_probs_grad,
+        previous_event=None,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
+        # dispatch grad
+        hs_grad, _, token_probs_grad = self.dispatch_act_node.backward(
+            hs_dispatched_grad,
+            dispatched_probs_grad,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        return hs_grad, token_probs_grad
+
+
+class Fp8CombineNode:
+    def __init__(self, token_dispatcher, name="fp8_combine_node"):
+        self.token_dispatcher = token_dispatcher
+        self.combine_node = CombineNode(token_dispatcher)
+        self.name = name
+
+    @paddle.no_grad()
+    def forward(self, hidden_states_out, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
+        # combine
+        output_combine = self.combine_node.forward(
+            hidden_states_out,
+            self.token_dispatcher._comm_manager.group,
+            self.token_dispatcher._comm_manager.handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        output_combine.stop_gradient = False
+        self.token_dispatcher._comm_manager.handle = None
+        return output_combine
+
+    @paddle.no_grad()
+    def backward(self, output_combine_grad, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
+        # combine grad -> fp8
+        hidden_states_out_grad = self.combine_node.backward(
+            output_combine_grad,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
+        )
+        return hidden_states_out_grad
+
+
+class Fp8CombineQuantNode:
+    def __init__(self, token_dispatcher, moe_group=None, name="fp8_combine_quant_node"):
+        self.token_dispatcher = token_dispatcher
+        self.name = name
+        self.moe_group = moe_group
+
+    @paddle.no_grad()
+    def forward(self, output_combine):
+        # post combine
+        output = output_combine.reshape(self.token_dispatcher.hidden_shape)
+        output_combine._record_stream()
+        self.output_combine_shape = output_combine.shape
+        output.stop_gradient = False
+        return output
+
+    @paddle.no_grad()
+    def backward(self, output_grad, event_to_wait=None):
+        # post combine grad
+        if DSV3_USE_FP8_DISPATCH:
+            if event_to_wait is not None:
+                assert self.moe_group is not None
+                event_to_wait.comm_stream_wait(self.moe_group.id)
+                buffer = get_buffer(self.token_dispatcher._comm_manager.group, get_hidden_bytes(output_grad))
+                custom_stream = paddle.device.Stream(stream_base=buffer.runtime.get_comm_stream())
+            else:
+                custom_stream = paddle.device.current_stream()
+            with paddle.device.stream_guard(custom_stream):
+                output_combine_grad = paddle.reshape(output_grad, [-1, output_grad.shape[-1]])
+                # output_combine_grad quant to fp8
+                output_combine_grad_fp8, output_combine_grad_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
+                    output_combine_grad, output_scale_transpose=False, quant_method="1x128", input_transpose=False
+                )
+                output_grad._record_stream()
+                quant_event = None
+                if event_to_wait is not None:
+                    quant_event = deep_ep.get_event_from_custom_stream(custom_stream.stream_base)
+            return (output_combine_grad_fp8, output_combine_grad_scale), quant_event
+        else:
+            output_combine_grad = paddle.reshape(output_grad, [-1, output_grad.shape[-1]])
+            return output_combine_grad, None
+
+
+class FusionMlpNode:
+    """
+    The FusedMoeLayer class includes operations for unzipping, expert computation, and zipping.
+    """
+
+    def __init__(
+        self,
+        custom_map,
+        max_topk,
+        recompute_fwd_gate_up=False,
+        is_split_group_gemm=True,
+        mlp_fwd_subbatch_rows=0,
+        mlp_bwd_subbatch_rows=0,
+        output_subbatch_rows=0,
+    ):
+        self.token_dispatcher = custom_map.token_dispatcher
+        self.experts = custom_map.experts
+        self.unzip_node = UnZipNode()
+        self.zip_node = ZipNode()
+        self.experts_group_gemm_node = FP8GroupGemmMlpFunctionNode(
+            custom_map,
+            recompute_fwd_gate_up=recompute_fwd_gate_up,
+            is_split_group_gemm=is_split_group_gemm,
+        )
+
+        self.seq_length = custom_map.config.seq_length
+        self.num_experts_per_tok = custom_map.config.num_experts_per_tok
+        self.adaptive_remained_O1_recompute_ratio = custom_map.config.adaptive_remained_O1_recompute_ratio
+
+        self.recompute_fwd_gate_up = recompute_fwd_gate_up
+        self.dispatched_indices = None
+        self.dispatched_probs = None
+        self.tokens_per_expert = None
+        self.padding_token_per_experts = None
+        self.router_topk = max_topk
+        self.mlp_fwd_subbatch_rows = mlp_fwd_subbatch_rows
+        self.mlp_bwd_subbatch_rows = mlp_bwd_subbatch_rows
+        self.output_subbatch_rows = output_subbatch_rows
+
+    def set_recompute_fwd_gate_up(self, recompute_fwd_gate_up):
+        self.experts_group_gemm_node.recompute_fwd_gate_up = recompute_fwd_gate_up
+
+    def reset_statue(self):
+        """
+        重置所有状态变量。
+
+        Args:
+            无。
+
+        Returns:
+            无。
+
+        """
+        self.dispatched_indices = None
+        self.dispatched_probs = None
+        self.tokens_per_expert = None
+        self.padding_token_per_experts = None
+        self.router_topk = None
+
+        del self.unzip_node
+        del self.zip_node
+        self.unzip_node = None
+        self.zip_node = None
+
+        self.experts_group_gemm_node.reset_statue()
+        self.experts_group_gemm_node = None
+
+    def prepare_env_subbatch(self, unzipped_tokens=None, unzipped_tokens_scale=None, is_fwd=True):
+        if is_fwd:
+            assert unzipped_tokens is not None and unzipped_tokens_scale is not None
+            self.experts_group_gemm_node.input_fp8 = unzipped_tokens
+            self.experts_group_gemm_node.input_scale = unzipped_tokens_scale
+            self.m_indices = self.experts_group_gemm_node.gen_m_indices(self.padding_token_per_experts)
+            self.experts_group_gemm_node.fwd_subbatch = True
+        else:
+            self.m_indices = (
+                self.experts_group_gemm_node.gen_m_indices(self.padding_token_per_experts)
+                if not hasattr(self, "m_indices")
+                else self.m_indices
+            )
+            self.experts_group_gemm_node.bwd_subbatch = True
+            reload(self.experts_group_gemm_node.input_fp8)
+            reload(self.experts_group_gemm_node.input_scale)
+
+    def gemm_forward_subbatch(
+        self,
+        unzipped_tokens,
+        unzipped_tokens_scale,
+        unzipped_probs,
+        map_unzipped_indices_to_zipped,
+        output,
+        total_zipped_tokens,
+        padding_token_per_experts,
+        start_idx=None,
+        end_idx=None,
+        output_subbatch_rows=None,
+    ):
+        if start_idx is None or end_idx is None:
+            start_idx = 0
+            end_idx = unzipped_tokens.shape[0]
+        start_idx = max(0, start_idx)
+        end_idx = min(unzipped_tokens.shape[0], end_idx)
+
+        expert_out = self.experts_group_gemm_node.forward(
+            (unzipped_tokens[start_idx:end_idx], unzipped_tokens_scale[start_idx:end_idx]),
+            unzipped_probs[start_idx:end_idx],
+            padding_token_per_experts,
+            m_indices=self.m_indices[start_idx:end_idx],
+        )
+
+        output = tokens_zip_unique_add_with_subbatch(
+            output,
+            expert_out,
+            map_unzipped_indices_to_zipped[start_idx:end_idx],
+            total_zipped_tokens,
+            subbatch_rows=output_subbatch_rows,
+        )
+        return output
+
+    def gemm_backward_subbatch(
+        self,
+        unzipped_grad,
+        map_unzipped_indices_to_zipped,
+        total_zipped_tokens,
+        output,
+        padding_token_per_experts,
+        start_idx=None,
+        end_idx=None,
+        output_subbatch_rows=None,
+        reset_status=False,
+    ):
+        def split_list_prefix(l, start, end):
+            prefix_sum = [0] * (len(l) + 1)
+            for i in range(len(l)):
+                prefix_sum[i + 1] = prefix_sum[i] + l[i]
+
+            result = []
+            for i in range(len(l)):
+                segment_start = prefix_sum[i]
+                segment_end = prefix_sum[i + 1]
+                overlap_start = max(start, segment_start)
+                overlap_end = min(end, segment_end)
+                selected = max(0, overlap_end - overlap_start)
+                result.append(selected)
+            return result
+
+        if start_idx is None or end_idx is None:
+            start_idx = 0
+            end_idx = extract_first_if_tuple(unzipped_grad).shape[0]
+
+        start_idx = max(0, start_idx)
+        end_idx = min(extract_first_if_tuple(unzipped_grad).shape[0], end_idx)
+
+        # m_indices = self.experts_group_gemm_node.gen_m_indices(self.tokens_per_expert)
+        unzipped_inp_grad = (
+            (unzipped_grad[0][start_idx:end_idx].contiguous(), unzipped_grad[1][start_idx:end_idx].contiguous())
+            if isinstance(unzipped_grad, tuple)
+            else unzipped_grad[start_idx:end_idx].contiguous()
+        )
+        unzipped_grad, unzipped_probs_grad = self.experts_group_gemm_node.backward(
+            unzipped_inp_grad,
+            self.unzipped_probs[start_idx:end_idx].contiguous(),
+            input_fp8_slice=self.experts_group_gemm_node.input_fp8[start_idx:end_idx].contiguous(),
+            input_scale_slice=self.experts_group_gemm_node.input_scale[start_idx:end_idx].contiguous(),
+            tokens_per_expert=split_list_prefix(padding_token_per_experts, start_idx, end_idx),
+            m_indices=self.m_indices[start_idx:end_idx].contiguous(),
+            reset_status=reset_status,
+        )
+
+        output = tokens_zip_unique_add_with_subbatch(
+            output,
+            unzipped_grad,
+            map_unzipped_indices_to_zipped[start_idx:end_idx],
+            zipped_rows=total_zipped_tokens,
+            subbatch_rows=output_subbatch_rows,
+        )
+
+        return output, unzipped_probs_grad
+
+    @paddle.no_grad()
+    def forward(self, hs_2d_dispatched, dispatched_indices, dispatched_probs):
+        """
+        对输入数据进行前向传播计算。
+
+        Args:
+            hs_fp8_dispatched (Tensor): 表示被分派到各个专家的输入数据。
+            dispatched_indices (Tensor):表示输入数据被分派到的专家索引。
+            dispatched_probs (Tensor): 表示输入数据被分派到各个专家的概率。
+
+        Returns:
+            Tensor: 经过前向传播计算后的输出数据。
+
+        """
+        self.tokens_per_expert = self.token_dispatcher._comm_manager.tokens_per_expert
+        self.dispatched_probs = dispatched_probs
+        num_experts = len(self.tokens_per_expert)
+        padding_token_per_experts = [(x + 127) // 128 * 128 for x in self.tokens_per_expert]
+        self.padding_token_per_experts = padding_token_per_experts
+        # 1 unzip
+        self.dispatched_indices = dispatched_indices.to(paddle.int32)
+
+        total_zipped_tokens = extract_first_if_tuple(hs_2d_dispatched).shape[0]
+        (unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_tokens_scale,) = self.unzip_node.forward(
+            hs_2d_dispatched,
+            self.dispatched_indices,
+            dispatched_probs,
+            topk=self.router_topk,
+            num_experts=num_experts,
+            tokens_per_expert=self.tokens_per_expert,
+        )
+        record_stream_for_multi_input(hs_2d_dispatched)
+        dispatched_indices._record_stream()
+        dispatched_probs._record_stream()
+
+        self.unzipped_probs = unzipped_probs.unsqueeze(-1)
+
+        if DSV3_USE_FP8_DISPATCH:
+            total_unzipped_tokens = extract_first_if_tuple(unzipped_tokens).shape[0]
+            # If adaptive O1 recompute is enabled, determine whether to enable recompute O1 based on the degree of imbalance
+            if self.recompute_fwd_gate_up == -1:
+                if (
+                    total_unzipped_tokens
+                    > self.seq_length * self.num_experts_per_tok * self.adaptive_remained_O1_recompute_ratio
+                ):
+                    # logger.debug(f"recompute_fwd_gate_up changed to True, Because the receives {unzipped_tokens.shape[0]} Tensors greater then {self.seq_length*self.num_experts_per_tok*self.adaptive_remained_O1_recompute_ratio}.")
+                    self.set_recompute_fwd_gate_up(True)
+                else:
+                    # logger.debug(f"recompute_fwd_gate_up changed to False, Because the receives {unzipped_tokens.shape[0]} Tensors less then {self.seq_length*self.num_experts_per_tok*self.adaptive_remained_O1_recompute_ratio}.")
+                    self.set_recompute_fwd_gate_up(False)
+
+            # if use_mlp_subbatch is enabled, then split the unzipped_tokens into subbatches
+            if self.mlp_fwd_subbatch_rows != 0 and total_unzipped_tokens > self.mlp_fwd_subbatch_rows * 2:
+                assert (
+                    self.experts_group_gemm_node.recompute_fwd_gate_up
+                ), "recompute_fwd_gate_up must be true when use_mlp_subbatch = True"
+                map_unzipped_indices_to_zipped = TDU.tokens_unzip_slice(
+                    extract_first_if_tuple(hs_2d_dispatched),
+                    zipped_expertwise_rowmap,
+                    num_experts,
+                    total_unzipped_tokens,
+                    0,
+                    total_unzipped_tokens + 1,
+                )
+                if isinstance(hs_2d_dispatched, tuple):
+                    hs_2d_dispatched[0]._clear_to_zero_allocation()
+                    hs_2d_dispatched[1]._clear_to_zero_allocation()
+                else:
+                    hs_2d_dispatched._clear_to_zero_allocation()
+
+                subbatch_rows = min((total_unzipped_tokens // num_experts) // 128 * 128, self.mlp_fwd_subbatch_rows)
+                nparts = (total_unzipped_tokens + subbatch_rows - 1) // subbatch_rows
+                output = paddle.empty([0, extract_first_if_tuple(hs_2d_dispatched).shape[-1]], dtype=paddle.float32)
+                self.prepare_env_subbatch(unzipped_tokens, unzipped_tokens_scale, True)
+                logger.info(
+                    f"Enable subbatch_forward!! total_zipped_tokens:{total_zipped_tokens}, total_unzipped_tokens:{total_unzipped_tokens}, nparts:{nparts}, subbatch_rows:{subbatch_rows}, output_sub_rows:{self.output_subbatch_rows}"
+                )
+                for i in range(nparts):
+                    start_idx = i * subbatch_rows
+                    end_idx = min(start_idx + subbatch_rows, total_unzipped_tokens)
+                    output = self.gemm_forward_subbatch(
+                        unzipped_tokens,
+                        unzipped_tokens_scale,
+                        unzipped_probs,
+                        map_unzipped_indices_to_zipped,
+                        output,
+                        total_zipped_tokens,
+                        padding_token_per_experts,
+                        start_idx=start_idx,
+                        end_idx=end_idx,
+                        output_subbatch_rows=self.output_subbatch_rows,
+                    )
+
+                output = merge_subbatch_cast(output, paddle.bfloat16)
+                output.stop_gradient = False
+                offload(self.experts_group_gemm_node.input_fp8)
+                offload(self.experts_group_gemm_node.input_scale)
+                return output
+
+            # 2 experts
+            expert_out = self.experts_group_gemm_node.forward(
+                (unzipped_tokens, unzipped_tokens_scale), unzipped_probs, padding_token_per_experts
+            )
+        else:
+            # If adaptive O1 recompute is enabled, determine whether to enable recompute O1 based on the degree of imbalance
+            if self.recompute_fwd_gate_up == -1:
+                if (
+                    unzipped_tokens.shape[0]
+                    > self.seq_length * self.num_experts_per_tok * self.adaptive_remained_O1_recompute_ratio
+                ):
+                    self.set_recompute_fwd_gate_up(True)
+                else:
+                    self.set_recompute_fwd_gate_up(False)
+
+            # 2 experts
+            expert_out = self.experts_group_gemm_node.forward(
+                unzipped_tokens, unzipped_probs, padding_token_per_experts
+            )
+
+        # 3 zip
+        if isinstance(hs_2d_dispatched, tuple):
+            hs_2d_dispatched[0]._clear_to_zero_allocation()
+            hs_2d_dispatched[1]._clear_to_zero_allocation()
+        else:
+            hs_2d_dispatched._clear_to_zero_allocation()
+        expert_out_tmp = expert_out.reshape([-1, expert_out.shape[-1]])
+
+        expert_out_zipped = self.zip_node.forward(
+            expert_out_tmp,
+            zipped_expertwise_rowmap,
+            self.dispatched_indices,
+            unzipped_probs,
+            total_zipped_tokens=total_zipped_tokens,
+            num_experts=num_experts,
+        )
+
+        expert_out_zipped.stop_gradient = False
+        return expert_out_zipped
+
+    @paddle.no_grad()
+    def backward(self, hidden_states_out_grad):
+        """
+        反向传播函数。
+
+        Args:
+            hidden_states_out_grad_fp8 (Tensor): 隐藏状态梯度。
+
+        Returns:
+            Tuple[Tensor, Tensor]: 包含两个元素，分别为hs_fp8_dispatched_grad和dispatched_probs_grad。
+                - hs_fp8_dispatched_grad (Tensor): 解压后的隐藏状态梯度。
+                - dispatched_probs_grad (Tensor): 分发概率梯度。
+
+        """
+        # zip_grad
+        unzipped_grad = self.zip_node.backward(
+            hidden_states_out_grad,
+            self.dispatched_indices,
+            self.dispatched_probs,
+            top_k=self.router_topk,
+            num_experts=len(self.tokens_per_expert),
+            tokens_per_expert=self.tokens_per_expert,
+        )
+        record_stream_for_multi_input(hidden_states_out_grad)
+
+        total_zipped_tokens = extract_first_if_tuple(hidden_states_out_grad).shape[0]
+        total_unzipped_tokens = extract_first_if_tuple(unzipped_grad).shape[0]
+        hidden_states_size = extract_first_if_tuple(hidden_states_out_grad).shape[-1]
+        num_experts = len(self.tokens_per_expert)
+        padding_token_per_experts = [(x + 127) // 128 * 128 for x in self.tokens_per_expert]
+
+        if self.mlp_bwd_subbatch_rows != 0 and total_unzipped_tokens > self.mlp_bwd_subbatch_rows * 2:
+            map_unzipped_indices_to_zipped = TDU.tokens_unzip_slice(
+                extract_first_if_tuple(hidden_states_out_grad),
+                self.unzip_node.zipped_expertwise_rowmap,
+                num_experts,
+                total_unzipped_tokens,
+                0,
+                total_unzipped_tokens + 1,
+            )
+            if isinstance(hidden_states_out_grad, tuple):
+                hidden_states_out_grad[0]._clear_to_zero_allocation()
+                hidden_states_out_grad[1]._clear_to_zero_allocation()
+            else:
+                hidden_states_out_grad._clear_to_zero_allocation()
+
+            subbatch_rows = min((total_unzipped_tokens // num_experts) // 128 * 128, self.mlp_bwd_subbatch_rows)
+            nparts = (total_unzipped_tokens + subbatch_rows - 1) // subbatch_rows
+            output = paddle.empty([0, hidden_states_size], dtype=paddle.float32)
+            probs_grad_list = []
+            self.prepare_env_subbatch(is_fwd=False)
+            logger.info(
+                f"Enable subbatch_backward!! total_zipped_tokens:{total_zipped_tokens}, total_unzipped_tokens:{total_unzipped_tokens}, nparts:{nparts}, subbatch_rows:{subbatch_rows}, output_sub_rows:{self.output_subbatch_rows}"
+            )
+            for i in range(nparts):
+                reset_status = True if i == nparts - 1 else False  # release saved status in the last part.
+                start_idx = i * subbatch_rows
+                end_idx = min(start_idx + subbatch_rows, total_unzipped_tokens)
+                output, probs_grad = self.gemm_backward_subbatch(
+                    unzipped_grad,
+                    map_unzipped_indices_to_zipped,
+                    total_zipped_tokens,
+                    output,
+                    padding_token_per_experts,
+                    start_idx=start_idx,
+                    end_idx=end_idx,
+                    output_subbatch_rows=self.output_subbatch_rows,
+                    reset_status=reset_status,
+                )
+                probs_grad_list.append(probs_grad)
+            if isinstance(unzipped_grad, tuple):
+                unzipped_grad[0]._clear_to_zero_allocation()
+                unzipped_grad[1]._clear_to_zero_allocation()
+            else:
+                unzipped_grad._clear_to_zero_allocation()
+            hs_dispatched_grad = merge_subbatch_cast(output, paddle.bfloat16)
+            dispatched_probs_grad = TDU.tokens_zip_prob_seq_subbatch(
+                probs_grad_list, self.unzip_node.zipped_expertwise_rowmap, self.dispatched_indices, subbatch_rows
+            )
+            self.reset_statue()
+            return hs_dispatched_grad, dispatched_probs_grad
+
+        if isinstance(hidden_states_out_grad, tuple):
+            hidden_states_out_grad[0]._clear_to_zero_allocation()
+            hidden_states_out_grad[1]._clear_to_zero_allocation()
+        else:
+            hidden_states_out_grad._clear_to_zero_allocation()
+
+        # expert_grad
+        expert_out, probs_grad = self.experts_group_gemm_node.backward(
+            unzipped_grad, self.unzipped_probs, padding_token_per_experts
+        )
+
+        hs_dispatched_grad, dispatched_probs_grad = self.unzip_node.backward(
+            expert_out,
+            total_zipped_tokens,
+            probs_grad,
+            self.dispatched_indices,
+            num_experts=num_experts,
+        )
+
+        self.reset_statue()
+        return hs_dispatched_grad, dispatched_probs_grad
+
+
+class FusionMoeNode:
+    def __init__(
+        self,
+        custom_map,
+        recompute_fwd_gate_up=False,
+        is_split_group_gemm=True,
+        mlp_fwd_subbatch_rows=0,
+        mlp_bwd_subbatch_rows=0,
+        output_subbatch_rows=0,
+        name="fusion_moe_node",
+    ):
+        self.token_dispatcher = custom_map.token_dispatcher
+        self.moe_router_topk = custom_map.moe_router_topk
+        self.dispatch_quant_node = Fp8DispatchQuantNode(self.token_dispatcher)
+        self.dispatch_node = Fp8DispatchNode(self.token_dispatcher)
+        self.mlp_node = FusionMlpNode(
+            custom_map,
+            self.moe_router_topk,
+            recompute_fwd_gate_up=recompute_fwd_gate_up,
+            is_split_group_gemm=is_split_group_gemm,
+            mlp_fwd_subbatch_rows=mlp_fwd_subbatch_rows,
+            mlp_bwd_subbatch_rows=mlp_bwd_subbatch_rows,
+            output_subbatch_rows=output_subbatch_rows,
+        )
+        self.combine_node = Fp8CombineNode(self.token_dispatcher)
+        self.combine_quant_node = Fp8CombineQuantNode(self.token_dispatcher, custom_map.moe_group)
+        self.name = name
+
+    @paddle.no_grad()
+    def forward(self, hidden_states, probs, routing_map):
+        if DSV3_USE_FP8_DISPATCH:
+            (hs_fp8, hs_scale), token_indices, token_probs = self.dispatch_quant_node.forward(
+                hidden_states, probs, routing_map
+            )
+            (
+                (hs_fp8_dispatched, hs_scale_dispatched),
+                dispatched_indices,
+                dispatched_probs,
+            ) = self.dispatch_node.forward((hs_fp8, hs_scale), token_indices, token_probs)
+            hidden_states_out = self.mlp_node.forward(
+                (hs_fp8_dispatched, hs_scale_dispatched), dispatched_indices, dispatched_probs
+            )
+            output_combine = self.combine_node.forward(hidden_states_out)
+            output = self.combine_quant_node.forward(output_combine)
+            output.stop_gradient = False
+            return output
+        else:
+            hs_2d_dispatched, dispatched_indices, dispatched_probs = self.dispatch_node.forward(
+                hidden_states, probs, routing_map
+            )
+            hidden_states_out = self.mlp_node.forward(hs_2d_dispatched, dispatched_indices, dispatched_probs)
+            output_combine = self.combine_node.forward(hidden_states_out)
+            output = self.combine_quant_node.forward(output_combine)
+            output.stop_gradient = False
+            return output
+
+    @paddle.no_grad()
+    def backward(self, output_grad):
+        output_combine_grad, _ = self.combine_quant_node.backward(output_grad)
+        hidden_states_out_grad = self.combine_node.backward(output_combine_grad)
+
+        hs_dispatched_grad, dispatched_probs_grad = self.mlp_node.backward(hidden_states_out_grad)
+
+        if DSV3_USE_FP8_DISPATCH:
+            hs_fp8_grad, token_probs_grad = self.dispatch_node.backward(hs_dispatched_grad, dispatched_probs_grad)
+            hs_grad, probs_grad, routing_map_grad = self.dispatch_quant_node.backward(hs_fp8_grad, token_probs_grad)
+            return hs_grad, probs_grad, routing_map_grad
+        else:
+            hs_bf16_grad, token_probs_grad = self.dispatch_node.backward(hs_dispatched_grad, dispatched_probs_grad)
+            return hs_bf16_grad, None, token_probs_grad
+
+
+class FusionMoe(paddle.autograd.PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden_states,
+        probs,
+        routing_map,
+        custom_map,
+        recompute_fwd_gate_up=False,
+        is_split_group_gemm=True,
+    ):
+        ctx.node = FusionMoeNode(
+            custom_map,
+            recompute_fwd_gate_up=recompute_fwd_gate_up,
+            is_split_group_gemm=is_split_group_gemm,
+        )
+        return ctx.node.forward(hidden_states, probs, routing_map)
+
+    @staticmethod
+    def backward(ctx, output_grad):
+        return ctx.node.backward(output_grad)

paddlenlp/transformers/moe_utils.py

@@ -14,15 +14,61 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-
+import numpy as np
 import paddle
 
+try:
+    import TokenDispatcherUtils as TDU
+except ImportError:
+    TDU = None
+
+from .fp8_utils import FP8LinearFunctionBase
+
+if not hasattr(paddle.Tensor, "_clear_to_zero_allocation"):
+
+    def _clear_to_zero_allocation(self):
+        """
+        _clear_to_zero_allocation
+        """
+        old_shape = self.shape
+        dst = paddle.empty([0], dtype=self.dtype)
+        dst_t = dst.value().get_tensor()
+        src_t = self.value().get_tensor()
+        src_t._share_data_with(dst_t)
+        src_t._set_dims(old_shape)
+
+    setattr(paddle.Tensor, "_clear_to_zero_allocation", _clear_to_zero_allocation)
+
+
+if not hasattr(paddle.Tensor, "_holder_size"):
+
+    def _holder_size(self):
+        """
+        _holder_size
+        """
+        if self._is_initialized():
+            return int(np.prod(self.shape)) * paddle.core.size_of_dtype(self.dtype)
+        else:
+            return 0
+
+    setattr(paddle.Tensor, "_holder_size", _holder_size)
+
+
+def topk_to_permuted_indices(x, num_tokens_per_expert_list, topk):
+    x = paddle.flatten(x)
+    prob_permuted_indices = paddle.concat(
+        [
+            paddle.tensor.search._restrict_nonzero(x == i, total_true_num)
+            for i, total_true_num in enumerate(num_tokens_per_expert_list)
+        ]
+    ).flatten()
+    token_permuted_indices = prob_permuted_indices // topk
+    return token_permuted_indices, prob_permuted_indices
+
 
 def permute(
     tokens,
-    routing_map,
-    num_out_tokens: Optional[int] = None,
+    token_permuted_indices,
     drop_and_pad: bool = False,
 ):
     """Permute the tokens and probs based on the mask.
@@ -32,35 +78,21 @@ def permute(
 
     Args:
         tokens (paddle.Tensor): The input token tensor, [num_tokens, hidden].
-        routing_map (paddle.Tensor): The sparse token to expert mapping, [num_tokens, num_experts].
-        num_out_tokens (int, optional): The number of output tokens. If None, it's set to
-                                        the number of input tokens.
         drop_and_pad (bool, optional): Whether or not the token dispatcher uses token-drop
                                        and pads the number of tokens to the expert capacity.
     """
     assert not drop_and_pad, "token-drop and pads is not supported"
-    num_tokens, hidden = tokens.shape
-    num_experts = routing_map.shape[1]
-
-    # mask [num_tokens, num_experts] -> [num_experts, num_tokens]
-    routing_map = routing_map.cast(paddle.bool).T.contiguous()
-
-    # Create a dense expert-to-token mapping from the sparse token-to-expert mapping
-    token_indices = paddle.arange(num_tokens).unsqueeze(0).expand([num_experts, -1])
-    sorted_indices = token_indices.masked_select(routing_map)
-
-    # use the mapping to permute the tokens
-    permuted_input = tokens.index_select(axis=0, index=sorted_indices)
-
-    return permuted_input, sorted_indices
+    # permuted_input = paddle.gather(tokens, token_permuted_indices)
+    permuted_input = tokens.index_select(axis=0, index=token_permuted_indices)
+    return permuted_input
 
 
 def unpermute(
     permuted_tokens: paddle.Tensor,
-    sorted_indices: paddle.Tensor,
+    token_permuted_indices: paddle.Tensor,
+    prob_permuted_indices: paddle.Tensor,
     restore_shape: paddle.shape,
     probs: paddle.Tensor = None,
-    routing_map: paddle.Tensor = None,
     drop_and_pad: bool = False,
 ):
     """
@@ -69,11 +101,9 @@ def unpermute(
 
     Args:
         permuted_tokens (paddle.Tensor): The permuted token tensor.
-        sorted_indices (paddle.Tensor): The indices used to sort the tokens.
+        token_permuted_indices (paddle.Tensor): The indices used to sort the tokens.
         restore_shape (paddle.shape): The shape of the unpermuted tensor.
         probs (paddle.Tensor, optional): The unpermuted probs tensor,
-        routing_map (paddle.Tensor, optional): Token to expert mapping, shape
-            [num_tokens, num_experts].
         drop_and_pad (bool, optional): Whether or not the token dispatcher uses token-drop
                                        and pads the number of tokens to the expert capacity.
 
@@ -82,20 +112,356 @@ def unpermute(
     """
     assert not drop_and_pad, "token-drop and pads is not supported"
     _, hidden = restore_shape
-
     if probs is not None:
-        assert routing_map is not None, "Mask must be provided to permute the probs."
-        permuted_probs = probs.T.contiguous().masked_select(routing_map.T.contiguous())
+        permuted_probs = paddle.gather(probs.flatten(), prob_permuted_indices)
         permuted_tokens = permuted_tokens * permuted_probs.unsqueeze(-1)
-
     # Create an output tensor filled with zeros
     output_tokens = paddle.zeros(restore_shape, dtype=permuted_tokens.dtype)
     # Scatter add the permuted_input back to the original positions
+
     output_tokens.put_along_axis_(
         axis=0,
-        indices=sorted_indices.unsqueeze(1).expand([-1, hidden]),
+        indices=token_permuted_indices.unsqueeze(1).expand([-1, hidden]),
         values=permuted_tokens,
         reduce="add",
         include_self=True,
     )
     return output_tokens
+
+
+class UnZipNode:
+    def __init__(self, name="unzip"):
+        self.name = name
+        self.unzipped_probs = None
+        self.zipped_expertwise_rowmap = None
+
+    def reset_statue(self):
+        self.unzipped_probs = None
+        self.zipped_expertwise_rowmap = None
+
+    @paddle.no_grad()
+    def forward(
+        self,
+        hs_2d_dispatched,
+        dispatched_indices,
+        dispatched_probs,
+        topk,
+        num_experts,
+        tokens_per_expert,
+    ):
+        if isinstance(hs_2d_dispatched, tuple):
+            with paddle.amp.auto_cast(False):
+                (
+                    unzipped_tokens,
+                    zipped_expertwise_rowmap,
+                    unzipped_probs,
+                    unzipped_scale,
+                ) = paddle.nn.functional.moe_permute(
+                    hs_2d_dispatched[0],
+                    hs_2d_dispatched[1],
+                    dispatched_indices,
+                    dispatched_probs,
+                    num_experts=num_experts,
+                    tokens_per_expert=tokens_per_expert,
+                    padding_alignment=128,
+                )
+        else:
+            with paddle.amp.auto_cast(False):
+                (
+                    unzipped_tokens,
+                    zipped_expertwise_rowmap,
+                    unzipped_probs,
+                    unzipped_scale,
+                ) = paddle.nn.functional.moe_permute(
+                    hs_2d_dispatched,
+                    None,
+                    dispatched_indices,
+                    dispatched_probs,
+                    num_experts=num_experts,
+                    tokens_per_expert=tokens_per_expert,
+                    padding_alignment=128,
+                )
+        self.unzipped_probs = unzipped_probs
+        self.zipped_expertwise_rowmap = zipped_expertwise_rowmap
+        return (unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_scale)
+
+    @paddle.no_grad()
+    def backward(self, dx, total_zipped_tokens, probs_grad, dispatched_indices, num_experts):
+        with paddle.amp.auto_cast(False):
+            weighted_zipped_tokens, probs_grad_zipped = paddle.nn.functional.moe_unpermute(
+                dx,
+                self.zipped_expertwise_rowmap,
+                dispatched_indices,
+                probs_grad,
+                total_zipped_tokens=total_zipped_tokens,
+                num_experts=num_experts,
+            )
+        self.reset_statue()
+        return weighted_zipped_tokens, probs_grad_zipped
+
+
+class ZipNode:
+    def __init__(self, name="zip"):
+        self.name = name
+
+    @paddle.no_grad()
+    def forward(
+        self, expert_out, zipped_expertwise_rowmap, routemap_topk, unzipped_probs, total_zipped_tokens, num_experts
+    ):
+        with paddle.amp.auto_cast(False):
+            expert_out_zipped, zipped_probs_topk = paddle.nn.functional.moe_unpermute(
+                expert_out, zipped_expertwise_rowmap, routemap_topk, unzipped_probs, total_zipped_tokens, num_experts
+            )
+        return expert_out_zipped
+
+    @paddle.no_grad()
+    def backward(
+        self,
+        grad_output,
+        dispatched_indices,
+        dispatched_probs,
+        top_k,
+        num_experts,
+        tokens_per_expert,
+    ):
+        if isinstance(grad_output, tuple):
+            with paddle.amp.auto_cast(False):
+                (
+                    unzipped_grad,
+                    zipped_expertwise_rowmap_grad,
+                    unzipped_probs_grad,
+                    unzipped_scale_grad,
+                ) = paddle.nn.functional.moe_permute(
+                    grad_output[0],
+                    grad_output[1],
+                    dispatched_indices,
+                    dispatched_probs,
+                    num_experts,
+                    tokens_per_expert,
+                    padding_alignment=128,
+                )
+                return (unzipped_grad, unzipped_scale_grad)
+        else:
+            with paddle.amp.auto_cast(False):
+                (
+                    unzipped_grad,
+                    zipped_expertwise_rowmap_grad,
+                    unzipped_probs_grad,
+                    unzipped_scale_grad,
+                ) = paddle.nn.functional.moe_permute(
+                    grad_output,
+                    None,
+                    dispatched_indices,
+                    dispatched_probs,
+                    num_experts,
+                    tokens_per_expert,
+                    padding_alignment=128,
+                )
+
+        return unzipped_grad
+
+
+class PermuteNode:
+    def __init__(self, token_dispatcher, name="permute"):
+        self.token_dispatcher = token_dispatcher
+        self.name = name
+
+    def reset_status(self):
+        self.token_permuted_indices = None
+        self.prob_permuted_indices = None
+
+    def forward(self, hidden_states, hidden_states_scale, dispatched_indices):
+        self.token_dispatcher._comm_manager.hidden_shape_before_permute = hidden_states.shape
+        self.hidden_shape_before_permute = hidden_states.shape
+        self.token_permuted_indices, self.prob_permuted_indices = topk_to_permuted_indices(
+            dispatched_indices,
+            self.token_dispatcher._comm_manager.tokens_per_expert,
+            self.token_dispatcher._comm_manager.router_topk,
+        )
+        hidden_states = permute(hidden_states, self.token_permuted_indices)
+        # permute scale
+        hidden_states_scale = permute(hidden_states_scale, self.token_permuted_indices)
+
+        return hidden_states, hidden_states_scale, self.token_permuted_indices, self.prob_permuted_indices
+
+    def backward(self, out_grad, dispatched_probs):
+        input_dtype = out_grad.dtype
+        hidden_states_grad = unpermute(
+            permuted_tokens=out_grad,
+            token_permuted_indices=self.token_permuted_indices,
+            prob_permuted_indices=self.prob_permuted_indices,
+            restore_shape=self.hidden_shape_before_permute,
+            probs=dispatched_probs,
+        )
+        self.reset_status()
+        return hidden_states_grad.to(input_dtype)
+
+
+class UnPermuteNode:
+    def __init__(self, token_dispatcher, name="unpermute"):
+        self.token_dispatcher = token_dispatcher
+        self.name = name
+
+    def reset_status(self):
+        self.token_permuted_indices = None
+        self.hidden_states = None
+        self.prob_permuted_indices = None
+        self.faltten_dispatched_probs = None
+        self.hidden = None
+        self.permuted_tokens = None
+        self.output_tokens = None
+
+    def forward(
+        self,
+        hidden_states,
+        token_permuted_indices,
+        prob_permuted_indices,
+        dispatched_probs,
+    ):
+        self.token_permuted_indices = token_permuted_indices
+        self.input_dtype = hidden_states.dtype
+        self.hidden_states = hidden_states
+        self.prob_permuted_indices = prob_permuted_indices
+        self.dispatched_probs_shape = dispatched_probs.shape
+        # permute
+        _, self.hidden = self.token_dispatcher._comm_manager.hidden_shape_before_permute
+
+        self.faltten_dispatched_probs = dispatched_probs.flatten()
+
+        self.permuted_probs = paddle.gather(self.faltten_dispatched_probs, self.prob_permuted_indices)
+        permuted_tokens = self.hidden_states * self.permuted_probs.unsqueeze(-1)
+        permuted_tokens = permuted_tokens.cast(self.hidden_states.dtype)
+
+        # Create an output tensor filled with zeros
+        output_tokens = paddle.zeros(
+            self.token_dispatcher._comm_manager.hidden_shape_before_permute, dtype=self.hidden_states.dtype
+        )
+        # Scatter add the permuted_input back to the original positions
+        output_tokens.put_along_axis_(
+            axis=0,
+            indices=self.token_permuted_indices.cast("int32").unsqueeze(1).expand([-1, self.hidden]),
+            values=permuted_tokens,
+            reduce="add",
+            include_self=True,
+        )
+        with paddle.base.device_guard("cpu"):
+            self.output_tokens = paddle.empty(shape=output_tokens.shape, dtype=output_tokens.dtype)
+
+        return output_tokens.to(self.input_dtype)
+
+    def backward(self, out_grad, out_grad_scale):
+        hidden_states_grad = paddle.gather(out_grad, self.token_permuted_indices)
+
+        output_tokens_grad = FP8LinearFunctionBase.dequantize_fp8_to_fp32(out_grad, out_grad_scale)
+        permuted_tokens = self.hidden_states * self.permuted_probs.unsqueeze(-1)
+        permuted_tokens = permuted_tokens.cast(self.hidden_states.dtype)
+
+        _, permuted_tokens_grad = paddle._C_ops.put_along_axis_grad(
+            self.output_tokens,
+            self.token_permuted_indices.cast("int32").unsqueeze(1).expand([-1, self.hidden]),
+            permuted_tokens,
+            self.output_tokens,
+            output_tokens_grad,
+            0,
+            "add",
+            True,
+        )
+
+        permuted_probs_grad = (permuted_tokens_grad * self.hidden_states).sum(axis=-1)
+
+        faltten_dispatched_probs_grad = paddle._C_ops.gather_grad(
+            self.faltten_dispatched_probs, self.prob_permuted_indices, permuted_probs_grad, 0
+        )
+
+        # dispatched_probs_grad = paddle._C_ops.flatten_grad(self.dispatched_probs, faltten_dispatched_probs_grad)
+        dispatched_probs_grad = faltten_dispatched_probs_grad.reshape(self.dispatched_probs_shape)
+
+        self.reset_status()
+        return hidden_states_grad, dispatched_probs_grad
+
+
+def tokens_zip_unique_add_with_subbatch(zipped, unzipped, index_unzipped, zipped_rows, subbatch_rows=None):
+    """
+    tokens_zip_unique_add_with_subbatch
+    """
+    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 merge_subbatch_cast(x, dtype):
+    if isinstance(x, (list, tuple)):
+        if len(x) == 1:
+            x = x[0]
+            return x.cast(dtype) if x.dtype != dtype else x
+        else:
+            return TDU.merge_subbatch_cast(x, dtype)
+    else:
+        return x.cast(dtype) if x.dtype != dtype else x
+
+
+def get_env_device():
+    """
+    Return the device name of running environment.
+    """
+    if paddle.is_compiled_with_cuda():
+        return "gpu"
+    elif "npu" in paddle.device.get_all_custom_device_type():
+        return "npu"
+    elif "mlu" in paddle.device.get_all_custom_device_type():
+        return "mlu"
+    elif "gcu" in paddle.device.get_all_custom_device_type():
+        return "gcu"
+    elif "intel_hpu" in paddle.device.get_all_custom_device_type():
+        return "intel_hpu"
+    elif paddle.is_compiled_with_rocm():
+        return "rocm"
+    elif paddle.is_compiled_with_xpu():
+        return "xpu"
+    return "cpu"
+
+
+def to_device(tensor, place=None):
+    if place is None:
+        place = get_env_device()
+
+    if isinstance(place, str):
+        place = paddle.device._convert_to_place(place)
+
+    if not tensor.place._equals(place):
+        new_t = tensor._copy_to(place, True)
+        dst_tensor = tensor.value().get_tensor()
+        src_tensor = new_t.value().get_tensor()
+        dst_tensor._share_data_with(src_tensor)
+
+    return tensor
+
+
+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"

paddlenlp/transformers/token_dispatcher.py

@@ -14,74 +14,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
 from typing import Optional, Tuple
 
 import paddle
 from paddle.distributed.communication.group import Group
 
 from .fused_a2a import fused_combine, fused_dispatch
-from .moe_utils import permute, unpermute
+from .moe_utils import permute, topk_to_permuted_indices, unpermute
 
 
-class _DispatchManager(ABC):
-    """
-    A manager class to handle dispatch and combine processes for MoE models.
-
-    DispatcherManager handles token dispatching according to the routing_map of format
-    [num_local_tokens, world_size, num_instances]. The routing_map is a 3D tensor where each
-    element indicates whether a token should be sent to a specific rank.
-
-    num_instances is the maximum number of tokens instances dispatched into a target rank, it
-    can be the number of local experts, or the size of sub_group.
-    """
-
-    @abstractmethod
-    def setup_metadata(self, routing_map: paddle.Tensor, probs: paddle.Tensor):
-        """Set up metadata of routing_map and probs."""
-        pass
-
-    @abstractmethod
-    def dispatch(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
-        """Dispatch the hidden_states according to the routing_map."""
-        pass
-
-    @abstractmethod
-    def combine(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
-        """Combine the hidden_states after expert processing."""
-        pass
-
-    @abstractmethod
-    def get_dispached_metadata(self) -> paddle.Tensor:
-        """Get the metadata of the dispatched hidden_states."""
-        pass
-
-    @abstractmethod
-    def get_permuted_hidden_states_by_experts(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
-        """Get the permuted hidden states by instances."""
-        pass
-
-    @abstractmethod
-    def get_restored_hidden_states_by_experts(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
-        """Get the restored hidden states by instances."""
-        pass
-
-
-class _DeepepManager(_DispatchManager):
+class _DeepepManager:
     """
     A manager class to handle fused all-to-all communication processes for MoE models using
     DeepEP backend. See https://github.com/deepseek-ai/deepep for more details.
 
     The workflow of the DeepEP dispatcher is:
-    (1) setup_metadata(): Process routing map and probabilities to prepare dispatch metadata
-    (2) dispatch():
+    (1) dispatch():
         - Use fused kernel to permute tokens and perform all-to-all communication in single step
-    (3) get_permuted_hidden_states_by_instances():
+    (2) get_permuted_hidden_states_by_instances():
         - Convert routing map and probabilities to multihot format
         - Permute tokens using fused kernel
-    (4) get_restored_hidden_states_by_instances():
+    (3) get_restored_hidden_states_by_instances():
         - Reverse permutation using fused kernel
-    (5) combine():
+    (4) combine():
         - Reverse process using fused kernel to unpermute and perform all-to-all in single step
 
     This implementation uses fused communication kernels (fused_dispatch/fused_combine) that
@@ -101,33 +56,22 @@ class _DeepepManager(_DispatchManager):
         self.num_experts = num_experts
         self.num_local_experts = num_local_experts
 
-        # Metadata
-        self.token_indices = None
-        self.token_probs = None
-        # Handle used for combine operation
         self.handle = None
 
         if fused_dispatch is None:
             raise ImportError("DeepEP is not supported in your paddlepaddle whl package.")
 
-    def setup_metadata(self, routing_map: paddle.Tensor, probs: paddle.Tensor):
-        num_tokens = routing_map.shape[0]
-
-        routing_map = routing_map.reshape([num_tokens, self.num_experts])
-        probs = probs.reshape([num_tokens, self.num_experts])
-        # Convert the format of routing map from multihot to indices.
-        self.token_probs, self.token_indices = paddle.topk(probs, self.router_topk, axis=-1)
-
-    def dispatch(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
+    def dispatch(
+        self, hidden_states: paddle.Tensor, token_indices: paddle.Tensor, token_probs: paddle.Tensor
+    ) -> paddle.Tensor:
         hidden_states, dispatched_probs, states = fused_dispatch(
-            hidden_states, self.token_indices, self.token_probs, self.num_experts, self.group
+            hidden_states, token_indices, token_probs, self.num_experts, self.group
         )
         self.handle = states["handle"]
-        self.tokens_per_expert = states["tokens_per_expert"]
-        self.dispatched_indices = states["dispatched_indices"]
-        self.dispatched_probs = dispatched_probs
+        self.tokens_per_expert_list = states["tokens_per_expert"]
+        dispatched_indices = states["dispatched_indices"]
 
-        return hidden_states
+        return hidden_states, dispatched_indices, dispatched_probs
 
     def _indices_to_multihot(self, indices, probs):
         """
@@ -154,57 +98,56 @@ class _DeepepManager(_DispatchManager):
         multihot_probs[row_indices, valid_indices] = probs[mask]
         return multihot_routing_map.cast(paddle.bool), multihot_probs
 
-    def get_dispached_metadata(self) -> paddle.Tensor:
-        return self.dispatched_indices, self.dispatched_probs
-
-    def get_number_of_tokens_per_expert(self) -> paddle.Tensor:
-        """
-        Get the number of tokens per expert.
-        """
-        return self.tokens_per_expert
-
     def combine(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
         hidden_states = fused_combine(hidden_states, self.group, self.handle)
-        # Release the handle after combine operation
-        self.handle = None
         return hidden_states
 
-    def get_permuted_hidden_states_by_experts(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
-        self.dispatched_routing_map, self.dispatched_probs = self._indices_to_multihot(
-            self.dispatched_indices, self.dispatched_probs
-        )
+    def get_permuted_hidden_states_by_experts(
+        self, hidden_states: paddle.Tensor, dispatched_indices: paddle.Tensor
+    ) -> paddle.Tensor:
         self.hidden_shape_before_permute = hidden_states.shape
-        hidden_states, self.reversed_mapping_for_combine = permute(
-            hidden_states,
-            self.dispatched_routing_map,
-            num_out_tokens=sum(self.tokens_per_expert),
+        token_permuted_indices, prob_permuted_indices = topk_to_permuted_indices(
+            dispatched_indices, self.tokens_per_expert_list, self.router_topk
         )
-        return hidden_states
+        hidden_states = permute(hidden_states, token_permuted_indices)
+        return hidden_states, token_permuted_indices, prob_permuted_indices
 
-    def get_restored_hidden_states_by_experts(self, hidden_states: paddle.Tensor) -> paddle.Tensor:
+    def get_restored_hidden_states_by_experts(
+        self,
+        hidden_states: paddle.Tensor,
+        token_permuted_indices: paddle.Tensor,
+        prob_permuted_indices: paddle.Tensor,
+        dispatched_probs: paddle.Tensor,
+    ) -> paddle.Tensor:
         input_dtype = hidden_states.dtype
-        assert self.dispatched_probs.dtype == paddle.float32, "DeepEP only supports float32 probs"
+        assert dispatched_probs.dtype == paddle.float32, "DeepEP only supports float32 probs"
         hidden_states = unpermute(
-            hidden_states,
-            self.reversed_mapping_for_combine,
+            permuted_tokens=hidden_states,
+            token_permuted_indices=token_permuted_indices,
+            prob_permuted_indices=prob_permuted_indices,
             restore_shape=self.hidden_shape_before_permute,
-            routing_map=self.dispatched_routing_map,
-            probs=self.dispatched_probs,
+            probs=dispatched_probs,
         )
         return hidden_states.to(input_dtype)
 
 
-class MoETokenDispatcher:
+class MoEFlexTokenDispatcher:
     """
-    MoE Token Dispatcher
+    Flexible token dispatcher for MoE models with Efficient-A2A communication kernels.
     """
 
-    def __init__(self, ep_group) -> None:
-        """
-        Initialize the MoE Token Dispatcher.
-        """
+    def __init__(self, num_local_experts: int, moe_router_topk: int, num_moe_experts: int, ep_group: Group):
         self._ep_group = ep_group
 
+        self.num_local_experts = num_local_experts
+        assert self.ep_size > 1, "Flex token dispatcher requires EP > 1"
+        self._comm_manager = _DeepepManager(
+            group=self.ep_group,
+            router_topk=moe_router_topk,
+            num_experts=num_moe_experts,
+            num_local_experts=self.num_local_experts,
+        )
+
     @property
     def ep_group(self):
         """Get expert model parallel group."""
@@ -215,70 +158,106 @@ class MoETokenDispatcher:
         """Get expert model parallel world_size."""
         return self.ep_group.world_size
 
-    @abstractmethod
-    def token_permutation(self, tokens: paddle.Tensor, probs: paddle.Tensor, routing_map: paddle.Tensor):
-        """Dispatch tokens to experts.
+    def pre_dispatch(self, hidden_states, probs, routing_map):
+        self.hidden_shape = hidden_states.shape
+        hidden_states = hidden_states.view([-1, self.hidden_shape[-1]])
+        num_tokens = routing_map.shape[0]
+        routing_map = routing_map.reshape([num_tokens, self._comm_manager.num_experts])
+        probs = probs.reshape([num_tokens, self._comm_manager.num_experts])
+        # Convert the format of routing map from multihot to indices.
+        token_probs, token_indices = paddle.topk(probs, self._comm_manager.router_topk, axis=-1)
+        return hidden_states, token_indices, token_probs
 
-        Args:
-            tokens (paddle.Tensor): Input tokens.
-            probs (paddle.Tensor): The routing probability tensor [num_tokens, num_experts].
-            routing_map (paddle.Tensor): Token to expert mapping tensor.
+    def post_dispatch(self, hidden_states, dispatched_indices):
+        (
+            global_input_tokens,
+            token_permuted_indices,
+            prob_permuted_indices,
+        ) = self._comm_manager.get_permuted_hidden_states_by_experts(hidden_states, dispatched_indices)
+        return (global_input_tokens, token_permuted_indices, prob_permuted_indices)
 
-        Returns:
-            paddle.Tensor: Tokens tensor.
-        """
-        raise NotImplementedError("Dispatch function not implemented.")
-
-    @abstractmethod
-    def token_unpermutation(self, expert_output: paddle.Tensor, bias: paddle.Tensor = None):
-        """Restores the expert output to its original ordering.
-
-        Args:
-            expert_output (paddle.Tensor): The output tensor from the expert models.
-            bias (paddle.Tensor): The bias tensor.
-
-        Returns:
-            (paddle.Tensor, paddle.Tensor): Unpermuted activation and optional bias.
-        """
-        raise NotImplementedError("Restore function not implemented.")
-
-
-class MoEFlexTokenDispatcher(MoETokenDispatcher):
-    """
-    Flexible token dispatcher for MoE models with Efficient-A2A communication kernels.
-    """
-
-    def __init__(self, num_local_experts: int, moe_router_topk: int, num_moe_experts: int, ep_group: Group):
-        super().__init__(ep_group)
-
-        self.num_local_experts = num_local_experts
-        assert self.ep_size > 1, "Flex token dispatcher requires EP > 1"
-        self._comm_manager = _DeepepManager(
-            group=self.ep_group,
-            router_topk=moe_router_topk,
-            num_experts=num_moe_experts,
-            num_local_experts=self.num_local_experts,
+    def pre_combine(self, hidden_states, token_permuted_indices, prob_permuted_indices, dispatched_probs):
+        hidden_states = self._comm_manager.get_restored_hidden_states_by_experts(
+            hidden_states, token_permuted_indices, prob_permuted_indices, dispatched_probs
         )
+        return hidden_states
+
+    def post_combine(self, hidden_states):
+        hidden_states = hidden_states.reshape(self.hidden_shape)
+        return hidden_states
 
     def token_permutation(
         self, hidden_states: paddle.Tensor, probs: paddle.Tensor, routing_map: paddle.Tensor
     ) -> Tuple[paddle.Tensor, paddle.Tensor]:
-        self.hidden_shape = hidden_states.shape
-        hidden_states = hidden_states.view([-1, self.hidden_shape[-1]])
+        hidden_states, token_indices, token_probs = self.pre_dispatch(hidden_states, probs, routing_map)
+        hidden_states, dispatched_indices, dispatched_probs = self._comm_manager.dispatch(
+            hidden_states, token_indices, token_probs
+        )
+        (global_input_tokens, token_permuted_indices, prob_permuted_indices) = self.post_dispatch(
+            hidden_states, dispatched_indices
+        )
 
-        self._comm_manager.setup_metadata(routing_map, probs)
-        hidden_states = self._comm_manager.dispatch(hidden_states)
-        global_input_tokens = self._comm_manager.get_permuted_hidden_states_by_experts(hidden_states)
-        tokens_per_expert = self._comm_manager.get_number_of_tokens_per_expert()
-
-        return global_input_tokens, tokens_per_expert
+        return (
+            global_input_tokens,
+            token_permuted_indices,
+            prob_permuted_indices,
+            dispatched_probs,
+        )
 
     def token_unpermutation(
-        self, hidden_states: paddle.Tensor, bias: Optional[paddle.Tensor] = None
+        self,
+        hidden_states: paddle.Tensor,
+        token_permuted_indices,
+        prob_permuted_indices,
+        dispatched_probs,
+        bias: Optional[paddle.Tensor] = None,
     ) -> Tuple[paddle.Tensor, Optional[paddle.Tensor]]:
         assert bias is None, "Bias is not supported in MoEFlexTokenDispatcher"
-        hidden_states = self._comm_manager.get_restored_hidden_states_by_experts(hidden_states)
+        hidden_states = self.pre_combine(
+            hidden_states, token_permuted_indices, prob_permuted_indices, dispatched_probs
+        )
         hidden_states = self._comm_manager.combine(hidden_states)
 
-        hidden_states = hidden_states.reshape(self.hidden_shape)
+        hidden_states = self.post_combine(hidden_states)
         return hidden_states, None
+
+
+class PreDispatchNode:
+    def __init__(self, token_dispatcher):
+        self.token_dispatcher = token_dispatcher
+        self.probs_origin_shape = None
+
+    def reset_status(self):
+        self.probs = None
+        self.reshaped_probs = None
+        self.token_indices = None
+
+    @paddle.no_grad()
+    def forward(self, routing_map, probs):
+        num_tokens = routing_map.shape[0]
+        self.probs_origin_shape = probs.shape
+        # routing_map = routing_map.reshape([num_tokens, token_dispatcher._comm_manager.num_experts])
+        self.probs = probs
+        reshaped_probs = probs.reshape([num_tokens, self.token_dispatcher._comm_manager.num_experts])
+        self.reshaped_probs = reshaped_probs
+        token_probs, token_indices = paddle.topk(
+            reshaped_probs, self.token_dispatcher._comm_manager.router_topk, axis=-1
+        )
+        self.token_indices = token_indices
+        token_probs.stop_gradient = False
+        return token_indices, token_probs
+
+    @paddle.no_grad()
+    def backward(self, token_probs_g):
+        probs_grad = paddle._C_ops.topk_grad(
+            self.reshaped_probs,
+            self.token_indices,
+            token_probs_g,
+            self.token_dispatcher._comm_manager.router_topk,
+            -1,
+            True,
+            True,
+        )
+        probs_reshape_g = paddle._C_ops.reshape_grad(self.probs, probs_grad)
+        self.reset_status()
+        return probs_reshape_g

paddlenlp/transformers/utils.py

@@ -1002,3 +1002,10 @@ def caculate_llm_per_token_flops(
     # 2 for mul + add in matmul
     # 1 for forward, 2 for backwards since we caluate gradients for input_x and input_y
     return 2 * (layer_num * (flops_per_transformer * 3 + flops_recompute_transformer) + 3 * flops_loggits) / seq_length
+
+
+def cast_if_needed(x, dtype):
+    """
+    cast_if_needed
+    """
+    return x.cast(dtype) if x.dtype != dtype else x

requirements-dev.txt

@@ -30,4 +30,5 @@ wget
 huggingface_hub>=0.19.2
 tiktoken
 tokenizers<=0.20.3; python_version<="3.8"
-tokenizers>=0.21,<0.22; python_version>"3.8"
+tokenizers>=0.21,<0.22; python_version>"3.8"
+cuda-python

slm/model_zoo/gpt-3/external_ops/fused_ln/layer_norm_cuda.cu

@@ -238,4 +238,15 @@ PD_BUILD_GRAD_OP(fused_rms_norm)
     ;
 
 
+PD_BUILD_OP(fused_rms_norm_grad_func)
+    .Inputs({"x", "scale", "invvar", "dy"})
+    .Outputs({"dx", "d_scale"})
+    .Attrs({"epsilon: float"})
+    .SetKernelFn(PD_KERNEL(RMSLnBwd))
+    .SetInferShapeFn(PD_INFER_SHAPE(RMSLnBwdInferShape))
+#ifdef CUSTOM_OP_WITH_SPMD
+    .SetInferSpmdFn(PD_INFER_SPMD_RULE(phi::distributed::RmsNormGradInferSpmd))
+#endif
+    ;
+
 // https://github.com/NVIDIA/apex/blob/85e9eddece9d4ac72b48c2407f8162f2173e1bf4/csrc/layer_norm_cuda_kernel.cu#L679

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.py

@@ -23,6 +23,7 @@ def get_gencode_flags():
     cc = prop.major * 10 + prop.minor
     return ["-gencode", "arch=compute_{0},code=sm_{0}".format(cc)]
 
+
 def run(func):
     p = multiprocessing.Process(target=func)
     p.start()

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_act_dequant.py

@@ -0,0 +1,111 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import FusedQuantOps as FQO 
+
+'''
+Swiglu Function:
+out = silu(x) * y when y is not None
+out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+'''
+
+def dequantize_fp8_to_bf16(fp8_tensor: paddle.Tensor, 
+                           scale: paddle.Tensor) -> paddle.Tensor:
+    expanded_scale = paddle.repeat_interleave(
+        scale, 
+        repeats=128, 
+        axis=-1
+    )
+    # 非规整情况，需要截断
+    expanded_scale = expanded_scale[:, :fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale)
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_swiglu_quant_result():
+    for width in [130, 4098, 7168]:
+        for height in [256, 1026, 4098]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            for padding in [False]:
+                pad_tag = "Padded: True" if padding is not None else "Padded: False"
+                print("-" * 20 + f"Testing with {pad_tag}" + "-" * 20)
+                x_fp8, scale = FQO.fused_act_quant(x, transpose_output=True, padding_last_dim_to_8x=padding, using_pow2_scaling=False)
+                fused_res, fused_scales = FQO.fused_act_dequant_transpose_act_quant(x_fp8,scale,padding_last_dim_to_8x=padding,using_pow2_scaling=False)
+                np_results=[]
+                golden_res = x
+                np_results.append(golden_res.astype("float").numpy())
+                if padding:
+                    dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                    np_results.append(dequanted_sliced_result[:, :height].numpy())
+                else:
+                    np_results.append(dequantize_fp8_to_bf16(fused_res, fused_scales).numpy())
+                nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+                print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+                try:
+                    np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                    print("+++++++ Passed ++++++++")
+                except AssertionError as err:
+                    print(err)
+                    compare_tensors(np_results[0], np_results[1])
+            
+def run():
+    verify_swiglu_quant_result()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_act_dequant.py

@@ -0,0 +1,91 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import FusedQuantOps as FQO 
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_act_dequant():
+    for width in [7168]:
+        for height in [4096, 16384, 32768]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            print("-" * 20 + f"Testing with {width} * {height}" + "-" * 20)
+            x_fp8, scale = FQO.fused_act_quant(x, transpose_output=False, padding_last_dim_to_8x=False, using_pow2_scaling=False)
+            dequant_result= FQO.fused_act_dequant(x_fp8,scale)
+            np_results=[]
+            golden_res = x
+            np_results.append(golden_res.astype("float").numpy())
+            np_results.append(dequant_result.astype("float").numpy())
+            nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+            print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+            try:
+                np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                print("+++++++ Passed ++++++++")
+            except AssertionError as err:
+                print(err)
+                compare_tensors(np_results[0], np_results[1])
+            print(np_results[0])
+            print("------------")
+            print( np_results[1])
+            
+def run():
+    verify_act_dequant()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_act_dequant_transpose_act_quant.py

@@ -0,0 +1,112 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+#import test_quant
+import FusedQuantOps as FQO 
+
+'''
+Swiglu Function:
+out = silu(x) * y when y is not None
+out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+'''
+
+def dequantize_fp8_to_bf16(fp8_tensor: paddle.Tensor, 
+                           scale: paddle.Tensor) -> paddle.Tensor:
+    expanded_scale = paddle.repeat_interleave(
+        scale, 
+        repeats=128, 
+        axis=-1
+    )
+    # 非规整情况，需要截断
+    expanded_scale = expanded_scale[:, :fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale)
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_swiglu_quant_result():
+    for width in [130, 4098, 7168]:
+        for height in [256, 1026, 4098]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            for padding in [False]:
+                pad_tag = "Padded: True" if padding is not None else "Padded: False"
+                print("-" * 20 + f"Testing with {pad_tag}" + "-" * 20)
+                x_fp8, scale = FQO.fused_act_quant(x, transpose_output=True, padding_last_dim_to_8x=padding, using_pow2_scaling=False)
+                fused_res, fused_scales = FQO.fused_act_dequant_transpose_act_quant(x_fp8,scale,padding_last_dim_to_8x=padding,using_pow2_scaling=False)
+                np_results=[]
+                golden_res = x
+                np_results.append(golden_res.astype("float").numpy())
+                if padding:
+                    dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                    np_results.append(dequanted_sliced_result[:, :height].numpy())
+                else:
+                    np_results.append(dequantize_fp8_to_bf16(fused_res, fused_scales).numpy())
+                nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+                print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+                try:
+                    np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                    print("+++++++ Passed ++++++++")
+                except AssertionError as err:
+                    print(err)
+                    compare_tensors(np_results[0], np_results[1])
+            
+def run():
+    verify_swiglu_quant_result()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_act_quant.py

@@ -0,0 +1,116 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+#import test_quant
+import FusedQuantOps as FQO 
+
+'''
+Swiglu Function:
+out = silu(x) * y when y is not None
+out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+'''
+
+def dequantize_fp8_to_bf16(fp8_tensor: paddle.Tensor, 
+                           scale: paddle.Tensor) -> paddle.Tensor:
+    expanded_scale = paddle.repeat_interleave(
+        scale, 
+        repeats=128, 
+        axis=-1
+    )
+    # 非规整情况，需要截断
+    expanded_scale = expanded_scale[:, :fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale)
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_swiglu_quant_result():
+    for width in [130, 4098, 7168]:
+        for height in [256, 1026, 4098]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            for transposing in [True,False]:
+                for padding in [True]:
+                    pad_tag = "Padded: True" if padding is not None else "Padded: False"
+                    print("-" * 20 + f"Testing swiglu with {pad_tag} and transposing: {transposing}" + "-" * 20)
+                    fused_res, fused_scales = FQO.fused_act_quant(x,transpose_output=transposing,padding_last_dim_to_8x=padding,using_pow2_scaling=False)
+                    np_results=[]
+                    golden_res = x if not transposing else x.T
+                    np_results.append(golden_res.astype("float").numpy())
+                    if padding:
+                        if transposing:
+                            dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                            np_results.append(dequanted_sliced_result[:, :height].numpy())
+                        else:
+                            dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                            np_results.append(dequanted_sliced_result[:, :width].numpy())
+                    else:
+                        np_results.append(dequantize_fp8_to_bf16(fused_res, fused_scales).numpy())
+                    nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+                    print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+                    try:
+                        np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                        print("+++++++ Passed ++++++++")
+                    except AssertionError as err:
+                        print(err)
+                        compare_tensors(np_results[0], np_results[1])
+            
+def run():
+    verify_swiglu_quant_result()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_spaq.py

@@ -0,0 +1,105 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+#import test_quant
+import FusedQuantOps as FQO 
+
+'''
+Swiglu Function:
+out = silu(x) * y when y is not None
+out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+'''
+
+def dequantize_fp8_to_bf16(fp8_tensor: paddle.Tensor, 
+                           scale: paddle.Tensor) -> paddle.Tensor:
+    expanded_scale = paddle.repeat_interleave(
+        scale, 
+        repeats=128, 
+        axis=-1
+    )
+    # 非规整情况，需要截断
+    expanded_scale = expanded_scale[:, :fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale)
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_swiglu_quant_result():
+    for width in [4096,7168]:
+        for height in [8192, 16384, 32768]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            prob = paddle.randn([height, 1]).astype("float32")
+            np_results=[]
+            golden_res = F.swiglu(x) * prob 
+            fused_res, fused_scales = FQO.fused_spaq(x,prob, using_pow2_scaling=False)
+            np_results.append(golden_res.astype("float").numpy())
+            np_results.append(dequantize_fp8_to_bf16(fused_res, fused_scales).numpy())
+            nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+            print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+            try:
+                np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                print("+++++++ Passed ++++++++")
+            except AssertionError as err:
+                print(err)
+                compare_tensors(np_results[0], np_results[1])
+
+def run():
+    verify_swiglu_quant_result()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_stack_transpose_quant.py

@@ -0,0 +1,56 @@
+import FusedQuantOps as FQO
+import numpy as np
+
+import paddle
+
+
+def restore_stack_quant(out, scale):
+    scale = paddle.repeat_interleave(scale, repeats=128, axis=0)
+    scale = paddle.repeat_interleave(scale, repeats=128, axis=1)
+    x = out.astype('float32') * scale
+    return x
+
+
+def test_fused_stack_transpose_quant(
+    num_experts, seq_len, hidden_size, transpose
+):
+    print(num_experts, seq_len, hidden_size, transpose)
+
+    x_vec = []
+    for _ in range(num_experts):
+        x = paddle.randn([seq_len, hidden_size], dtype='bfloat16')
+        x = paddle.clip(x, min=-50, max=50)
+        x_vec.append(x)
+
+    if transpose:
+        out, scale = FQO.fused_stack_transpose_quant(x_vec)
+    else:
+        out, scale = FQO.fused_stack_quant(x_vec)
+
+    x_fp32 = paddle.stack(x_vec).reshape([-1, hidden_size]).astype('float32')
+    x_restored = restore_stack_quant(out, scale)
+
+    if transpose:
+        x_restored = (
+            x_restored.reshape([num_experts, hidden_size, seq_len])
+            .transpose([0, 2, 1])
+            .reshape([-1, hidden_size])
+        )
+
+    np.testing.assert_allclose(
+        x_fp32, x_restored, rtol=0.01, atol=0.2
+    )  # 存在截断误差，atol=0.2，通常在1e-6
+
+
+def run():
+    for batch_size in [1, 4]:
+        for seq_len in [2048, 7168]:
+            for hidden_size in [128, 4096]:
+                for transpose in [False, True]:
+                    test_fused_stack_transpose_quant(
+                        batch_size, seq_len, hidden_size, transpose
+                    )
+
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_swiglu_act_quant.py

@@ -0,0 +1,124 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+#import test_quant
+import FusedQuantOps as FQO 
+
+'''
+Swiglu Function:
+out = silu(x) * y when y is not None
+out = silu(xs[0]) * xs[1] when y is None, where xs = paddle.chunk(x, 2, axis=-1)
+'''
+
+def dequantize_fp8_to_bf16(fp8_tensor: paddle.Tensor, 
+                           scale: paddle.Tensor) -> paddle.Tensor:
+    expanded_scale = paddle.repeat_interleave(
+        scale, 
+        repeats=128, 
+        axis=-1
+    )
+    # 非规整情况，需要截断
+    expanded_scale = expanded_scale[:, :fp8_tensor.shape[-1]]
+    return (fp8_tensor.astype('float32') * expanded_scale)
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def printany(te):
+    for i in range(te.shape[0]):
+        for j in range(te.shape[1]):
+            print(te[i][j], end=", ")
+        print()
+    print("-"*20)
+
+def verify_swiglu_quant_result():
+    for width in [130, 4098, 7168]:
+        for height in [256, 1026, 4098]:
+            print("#"*60 + f" Testing width:{width}, height:{height} " + "#"*60)
+            x= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            y= paddle.clip(paddle.randn([height, width]).astype("bfloat16"), min=-50, max=50)
+            for transposing in [True,False]:
+                for padding in [True]:
+                    for optional_y in [y,None]:
+                        y_tag = "is_combined: False" if optional_y is not None else "is_combined: True"
+                        pad_tag = "Padded: True" if padding is not None else "Padded: False"
+                        print("-" * 20 + f"Testing swiglu with {y_tag} , {pad_tag} and transposing: {transposing}" + "-" * 20)
+                        fused_res, fused_scales = FQO.fused_swiglu_act_quant(x,optional_y,transpose_output=transposing, to_e4m3=True, using_pow2_scaling=False, padding_last_dim_to_8x=padding)
+                        np_results=[]
+                        if optional_y is None:
+                            golden_res = F.swiglu(x) if not transposing else F.swiglu(x).T
+                            np_results.append(golden_res.astype("float").numpy())
+                        else:
+                            golden_res = F.swiglu(x,y) if not transposing else F.swiglu(x,y).T
+                            np_results.append(golden_res.astype("float").numpy())
+                        if padding:
+                            if transposing:
+                                dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                                np_results.append(dequanted_sliced_result[:, :height].numpy())
+                            else:
+                                rank = width//2 if optional_y is None else width
+                                dequanted_sliced_result = dequantize_fp8_to_bf16(fused_res, fused_scales)
+                                np_results.append(dequanted_sliced_result[:, :rank].numpy())
+                        else:
+                            np_results.append(dequantize_fp8_to_bf16(fused_res, fused_scales).numpy())
+                        nan_cnt_golden, nan_cnt_fused= np.sum(np.isnan(np_results[0])), np.sum(np.isnan(np_results[1]))
+                        print(f"Nan count of Golden result: {nan_cnt_golden}; Nan count of Fused result: {nan_cnt_fused}")
+                        try:
+                            np.testing.assert_allclose(np_results[0], np_results[1], rtol=0.01, atol=1) #存在截断误差，atol=1，通常在1e-6
+                            print("+++++++ Passed ++++++++")
+                        except AssertionError as err:
+                            print(err)
+                            compare_tensors(np_results[0], np_results[1])
+            
+def run():
+    verify_swiglu_quant_result()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_fused_swiglu_probs_bwd.py

@@ -0,0 +1,108 @@
+import paddle
+import FusedQuantOps as FQO
+from paddle.base import core
+import numpy as np
+
+REP=100
+
+from paddle.incubate.nn.functional import swiglu
+
+paddle.seed(42)
+# o1.shape = [ 8 * 4096, 2048 * 2], dtype = "bfloat16"
+# prob.shape = [8 * 4096, 1], dtype = "float32"
+# do2_p.shape = [8 * 4096, 2048], 
+
+seq_len = 4096
+topk = 8
+moe_intermediate_size = 2048
+
+'''
+o1 = paddle.rand( [seq_len * topk, moe_intermediate_size * 2], dtype="bfloat16")
+unzipped_probs = paddle.rand( [ seq_len * topk, 1], dtype="float32")
+do2_s = paddle.rand( [seq_len * topk, moe_intermediate_size], dtype="bfloat16")
+'''
+o1 = paddle.rand( [topk, seq_len, moe_intermediate_size * 2], dtype="bfloat16")
+unzipped_probs = paddle.rand( [ topk, seq_len, 1], dtype="float32")
+do2_s = paddle.rand( [topk, seq_len , moe_intermediate_size], dtype="bfloat16")
+
+
+def swiglu_grad(x, dz):
+    xs = paddle.split(x, 2, axis=-1)
+    lhs = xs[0] 
+    rhs = xs[1]
+    
+    sig = paddle.nn.functional.sigmoid(lhs)
+    tmp = sig * lhs
+    x0_grad = dz * rhs * sig * (1 + lhs - tmp)
+    x1_grad = dz * tmp
+    
+    x_grad = paddle.concat([x0_grad, x1_grad], axis=-1)
+    
+    return x_grad
+
+def fn_splits():
+    # do2: 前向从bfloat16-->float32，反向从float32-->bfloat16,do2 需要保持 bfloat16（因为 o2 是 bfloat16)
+    o2 = swiglu(o1)
+    o2_s = (o2 * unzipped_probs)
+    do2 = (do2_s.cast(paddle.float32) * unzipped_probs)
+    do2 = do2.cast(paddle.bfloat16)
+    
+    # 展开swiglu_grad(o1, do2)部分
+    xs = paddle.split(o1, 2, axis=-1)
+    lhs = xs[0]
+    rhs = xs[1]
+    sig = paddle.nn.functional.sigmoid(lhs)
+    tmp = sig * lhs
+    x0_grad = do2 * rhs * sig * (1 + lhs - tmp)
+    x1_grad = do2 * tmp
+    do1 = paddle.concat([x0_grad, x1_grad], axis=-1)
+    
+    probs_grad = (do2_s.cast(paddle.float32) * (o2.cast(paddle.float32))).sum(axis=-1)
+    return do1, probs_grad, o2_s
+def fn_fused():
+    return FQO.fused_swiglu_probs_bwd(o1, do2_s, unzipped_probs)
+
+#input: o1 unzipped_probs do2
+def fn_gold(): 
+    # do2: 前向从bfloat16-->float32，反向从float32-->bfloat16,do2 需要保持 bfloat16（因为 o2 是 bfloat16)
+    o2 = swiglu(o1)
+    o2_s = (o2 * unzipped_probs)
+    do2 = (do2_s.cast(paddle.float32) * unzipped_probs)
+    do2 = do2.cast(paddle.bfloat16)
+    do1, _ = paddle._C_ops.swiglu_grad(o1, None, do2)
+    probs_grad = (do2_s.cast(paddle.float32) * (o2.cast(paddle.float32))).sum(axis=-1)
+    return do1, probs_grad, o2_s
+
+core.nvprof_start()
+
+for i in range(REP):
+    core.nvprof_nvtx_push("original")
+    do1_gold, pg_gold, o2_s_gold= fn_gold()
+    core.nvprof_nvtx_pop()
+for i in range(REP):
+    core.nvprof_nvtx_push("fused")
+    do1, pg, o2_s = fn_fused()
+    core.nvprof_nvtx_pop()
+do1_splits, pg_splits, o2_s_splits= fn_splits()
+
+print(do1_gold.dtype)
+print(pg_gold.dtype)
+print(o2_s_gold.dtype)
+print("----------")
+print(do1.dtype)
+print(pg.dtype)
+print(o2_s.dtype)
+print("pg_gold", pg_gold.astype("float32").numpy())
+print("pg_splits", pg_splits.astype("float32").numpy())
+print("pg", pg.astype("float32").numpy())
+print("o2_s_gold", o2_s_gold.astype("float32").numpy())
+print("o2_s_splits", o2_s_splits.astype("float32").numpy())
+print("o2_s", o2_s.astype("float32").numpy())
+print("do1_gold", do1_gold.astype("float32").numpy())
+print("do1_splits", do1_splits.astype("float32").numpy())
+print("do1", do1.astype("float32").numpy())
+pg_gold = pg_gold.flatten()
+pg = pg.flatten()
+np.testing.assert_allclose(pg_gold.astype("float32").numpy(), pg.astype("float32").numpy())
+np.testing.assert_allclose(o2_s_gold.astype("float32").numpy(), o2_s.astype("float32").numpy())
+np.testing.assert_allclose(do1_gold.astype("float32").numpy(), do1.astype("float32").numpy())

tests/ops/test_fused_transpose_split_quant.py

@@ -0,0 +1,52 @@
+import FusedQuantOps as FQO
+import numpy as np
+
+import paddle
+
+
+def restore_transpose_split_quant(out, scale):
+    out = [t.astype('float32') for t in out]
+    out = paddle.concat(out, axis=1).transpose([1, 0])
+    scale = paddle.concat(scale, axis=0)
+    scale = paddle.repeat_interleave(scale, repeats=128, axis=0)
+    return out * scale
+
+
+def test_fused_transpose_split_quant(tokens_per_expert, seq_len, pow_2_scales):
+    print(tokens_per_expert, seq_len, pow_2_scales)
+
+    x = paddle.randn([sum(tokens_per_expert), seq_len], dtype='bfloat16')
+    x = paddle.clip(x, min=-50, max=50)
+
+    out, scale = [], []
+    for tokens in tokens_per_expert:
+        out.append(paddle.empty([seq_len, tokens], dtype='float8_e4m3fn'))
+        scale.append(paddle.empty([tokens//128, seq_len], dtype='float32'))
+
+    FQO.fused_transpose_split_quant(x, out, scale, pow_2_scales)
+
+    x_restore = restore_transpose_split_quant(out, scale)
+    x_cast = x.astype('float32')
+
+    np.testing.assert_allclose(x_cast, x_restore, rtol=0.01, atol=0.3)
+
+
+def run():
+    test_fused_transpose_split_quant([0, 0], 1024, False)
+    test_fused_transpose_split_quant([128, 2*128], 0, True)
+    test_fused_transpose_split_quant([128], 1, False)
+    test_fused_transpose_split_quant([0, 128, 0, 2*128], 127, True)
+    test_fused_transpose_split_quant([3*128, 4*128, 5*128], 233, False)
+    test_fused_transpose_split_quant(
+        [24*128, 128, 50*128, 16*128], 2162, True
+    )
+    test_fused_transpose_split_quant(
+        [7*128, 29*128, 3*128, 128*128, 13*128], 4000, False
+    )
+    test_fused_transpose_split_quant(
+        [18*128, 5*128, 24*128, 128, 6*128, 0, 27*128, 7*128], 7168, True
+    )
+
+
+if __name__ == '__main__':
+    run()

tests/ops/test_guided_unzip.py

@@ -0,0 +1,140 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def verify_tokens_unzip():
+    expert_num = 4
+    topk = 8
+    seqlen = 3
+    token_len = 8
+    tokens_zipped = [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3]
+    ]
+    routemap_topk = [ 
+        [-1,-1,0,1,-1,-1,-1,-1],
+        [1,-1,-1,-1,-1,-1,-1,-1],
+        [-1,0,-1,-1,-1,-1,-1,-1],
+    ]
+    probs_topk = [ 
+        [0,0,0.3,0.5,0,0,0,0],
+        [0.5,0,0,0,0,0,0,0],
+        [0,0.7,0,0,0,0,0,0],
+    ]
+    total_unzipped_tokens_num = 4
+    expected_unzipped_tokens= [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3],
+        [1,1,1,1,1,1,1,1],
+    ]
+    expected_unzipped_probs= [
+        0.3,
+        0.5,
+        0.7,
+        0.5
+    ]
+    expected_zipped_expertwise_rowmap= [
+        [0,3,-1,-1],
+        [-1,1,-1,-1],
+        [2,-1,-1,-1],
+    ]
+    expected_unzipped_expert_idx = [
+        0,
+        1,
+        0,
+        1 
+    ]
+    expected_guided_unzipped_probs= [ 
+        [0,0,0.3,0.5,0,0,0,0],
+        [0.5,0,0,0,0,0,0,0],
+        [0,0.7,0,0,0,0,0,0],
+        [0,0,0.3,0.5,0,0,0,0],
+    ]
+    for prec in ['bfloat16']:
+        tokens_zipped = paddle.to_tensor(tokens_zipped, dtype=prec)
+        routemap_topk = paddle.to_tensor(routemap_topk, dtype='int32')
+        probs_topk = paddle.to_tensor(probs_topk, dtype='bfloat16')
+        expected_unzipped_probs = paddle.to_tensor(expected_unzipped_probs, dtype='bfloat16')
+        expected_zipped_expertwise_rowmap = paddle.to_tensor(expected_zipped_expertwise_rowmap, dtype='int32')
+        expected_unzipped_tokens = paddle.to_tensor(expected_unzipped_tokens, dtype=prec)
+        expected_unzipped_expert_idx = paddle.to_tensor(expected_unzipped_expert_idx, dtype='int32')
+        expected_guided_unzipped_probs = paddle.to_tensor(expected_guided_unzipped_probs, dtype='bfloat16')
+        unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_expert_idx = TDU.tokens_unzip(tokens_zipped,routemap_topk, probs_topk,total_unzipped_tokens_num=total_unzipped_tokens_num, topk=topk, num_experts=expert_num)
+
+        # 用于probs的给定顺序增广
+        guided_unzipped_probs = TDU.probs_topk_guided_unzip(probs_topk,routemap_topk, zipped_expertwise_rowmap, total_unzipped_token_num=total_unzipped_tokens_num,num_experts=expert_num, topk=topk)
+
+        # 用于dout的给定顺序增广
+        guided_unzipped_dout = TDU.tokens_guided_unzip(probs_topk,zipped_expertwise_rowmap, total_unzipped_token_num=total_unzipped_tokens_num,num_experts=expert_num)
+
+        # ------------------------- 前向验证 ------------------------
+        print("-------- probs guided unzipped by customed op: ------------")
+        print(guided_unzipped_probs)
+        print("-------- probs guided unzipped expected ------------")
+        print(expected_unzipped_probs)
+        print("-------- tokens guided unzipped by customed op: ------------")
+        print(guided_unzipped_dout)
+        print("-------- tokens guided unzipped expected ------------")
+        print(expected_guided_unzipped_probs)
+
+    
+
+    
+def run():
+    verify_tokens_unzip()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_regroup.py

@@ -0,0 +1,121 @@
+
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def verify_regroup_token():
+    expert_num = 3
+    seqlen = 5 
+    h1 = 8
+    h2 = 4
+    tokens = [ 
+        [1,1,1,1,1,1,1,1],
+        [3,3,3,3,3,3,3,3],
+        [2,2,2,2,2,2,2,2],
+        [5,5,5,5,5,5,5,5],
+        [4,4,4,4,4,4,4,4],
+    ]
+    dout = [
+        [0.1,0.1,0.1,0.1],
+        [0.3,0.3,0.3,0.3],
+        [0.2,0.2,0.2,0.2],
+        [0.5,0.5,0.5,0.5],
+        [0.4,0.4,0.4,0.4],
+    ]
+    expert_idx = [0,1,0,2,1]
+    token_per_experts = [2,2,1]
+    expected_tokens= [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3],
+        [4,4,4,4,4,4,4,4],
+        [5,5,5,5,5,5,5,5],
+        [0,0,0,0,0,0,0,0]
+    ]
+    expected_dout = [
+        [0.1,0.1,0.1,0.1],
+        [0.2,0.2,0.2,0.2],
+        [0.3,0.3,0.3,0.3],
+        [0.4,0.4,0.4,0.4],
+        [0.5,0.5,0.5,0.5],
+        [0.,0.,0.,0.]
+    ]
+
+    tokens =  paddle.to_tensor(tokens, dtype="bfloat16")
+    dout =  paddle.to_tensor(dout, dtype="bfloat16")
+    expert_idx = paddle.to_tensor(expert_idx,dtype="int32")
+    expected_tokens = paddle.to_tensor(expected_tokens, dtype="bfloat16")
+    expected_dout = paddle.to_tensor(expected_dout, dtype="bfloat16")
+    result_tokens, result_dout= TDU.regroup_tokens(
+        tokens,
+        dout,
+        expert_idx,
+        expert_num=expert_num,
+        token_max_per_expert=2)
+    print("##### executed unpadded")
+
+    print("-------- result_padded Generated by fused op: ------------")
+    print(result_tokens)
+    print("-------- result_padded expected : ------------")
+    print(expected_tokens)
+    print("-------- dout_padded Generated by fused op: ------------")
+    print(result_dout)
+    print("-------- dout_padded expected : ------------")
+    print(expected_dout)
+    
+def run():
+    verify_regroup_token()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_stable_unzip.py

@@ -0,0 +1,81 @@
+import numpy as np
+import paddle
+import sys
+import time
+try:
+    import deep_gemm
+    import kitchen
+    import kitchen.quantization_subchannel_block_hybrid
+    from kitchen.quantization import QParams, ScalingType
+except:
+    pass
+
+from paddle.base import core
+
+from paddle.incubate.nn.functional import swiglu
+import TokenDispatcherUtils as TDU
+
+H1 = 7168
+H2 = 2048
+topk = 8
+
+
+topk_ind = np.load("topk_indice.npy")
+reci_x = paddle.randn( [ topk_ind.shape[0], H1], dtype="bfloat16")
+reci_x_fp8 = reci_x.cast("float8_e4m3fn")
+reci_x_scale = paddle.randn((reci_x.shape[0], int((H1 + 127) / 128)), dtype="float32")
+print("reci_x scale shape: ", reci_x_scale.shape)
+topk_ind_base = paddle.to_tensor(topk_ind, dtype="int32")
+probs = paddle.ones( topk_ind_base.shape, dtype="bfloat16") # uses ones as topk_ind_base???
+print( topk_ind.shape)
+print( "recv x", reci_x.shape)
+
+total_num = int((topk_ind != -1).astype("int64").sum())
+
+e0_num = int( (topk_ind == 0).astype("int64").sum() )
+e1_num = int((topk_ind == 1).astype("int64").sum())
+e2_num = int( (topk_ind == 2).astype("int64").sum())
+e3_num = int( (topk_ind == 3).astype("int64").sum())
+token_per_expert = [  e0_num, e1_num, e2_num, e3_num ]
+print("token per expert: ", token_per_expert)
+max_tokens = max(token_per_expert)
+
+
+
+def test_unzip_stable():
+    # ---------------------------- Forward --------------------
+    # ------------ unzip and preprocess --------------
+    # ----------- FP8 -----------------
+    print("############## FP8 ################")
+    unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_scales = TDU.tokens_unzip_stable(
+        reci_x_fp8,
+        reci_x_scale,
+        topk_ind_base, 
+        probs,
+        topk=topk, num_experts=4, max_tokens_per_expert=max_tokens)
+    
+    print("zipped_expertwise_rowmap_fp8: ", zipped_expertwise_rowmap)
+    np.savetxt("zipped_expertwise_rowmap_fp8.csv", zipped_expertwise_rowmap, delimiter=",", fmt='%d')
+    np.savetxt("topk_ind_fp8.csv", topk_ind, delimiter=",", fmt='%d')
+    np.savetxt("unzipped_scales_fp8.csv", unzipped_scales[:10, :], delimiter=",", fmt='%d')
+    # ----------- BF16 -----------------
+    print("############## BF16 ################")
+    unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, _= TDU.tokens_unzip_stable(
+        reci_x,
+        None,
+        topk_ind_base, 
+        probs,
+        topk=topk, num_experts=4, max_tokens_per_expert=max_tokens)
+    
+    print("zipped_expertwise_rowmap: ", zipped_expertwise_rowmap)
+    np.savetxt("zipped_expertwise_rowmap.csv", zipped_expertwise_rowmap, delimiter=",", fmt='%d')
+    np.savetxt("topk_ind.csv", topk_ind, delimiter=",", fmt='%d')
+
+# core.nvprof_enable_record_event()
+
+for i in range(1):
+    core.nvprof_nvtx_push("step")
+    if i == 1:
+      core.nvprof_start()
+    test_unzip_stable()
+    core.nvprof_nvtx_pop()

tests/ops/test_tokens_unzip.py

@@ -0,0 +1,186 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def check_expert_idx(unzipped_expert_idx, zipped_expertwise_rowmap, expert_num=4, zipped_tokens_num=3):
+    unzipped_expert_idx = unzipped_expert_idx.reshape((1, -1))
+    print(unzipped_expert_idx)
+    assert zipped_tokens_num == zipped_expertwise_rowmap.shape[0]
+    for token_idx in range(zipped_tokens_num):
+        for expert_idx in range(expert_num):
+            this_expert_row = zipped_expertwise_rowmap[token_idx][expert_idx]
+            if this_expert_row != -1:
+                expected_expert_idx = expert_idx
+                this_unzipped_expert_idx = unzipped_expert_idx[0][this_expert_row]
+                if this_unzipped_expert_idx != expected_expert_idx:
+                    print(f"第{token_idx}行第{expert_idx}号专家不匹配， rowmap里是{expected_expert_idx}, unzipped里是{this_unzipped_expert_idx}")
+                    return  
+                else:
+                    print(f"rowmap [{token_idx}][{expert_idx}] 指向{this_expert_row}, unzipped[{this_expert_row}] = {this_unzipped_expert_idx}")
+    print("检查结束，通过")
+
+def check_indices(dispatched_indices, zipped_expertwise_rowmap, unzipped_expert_idx):
+    topk = dispatched_indices.shape[1]
+    zipped_token_num = zipped_expertwise_rowmap.shape[0]
+    for token_idx in range(zipped_token_num):
+        for k in range(topk):
+            this_expert = dispatched_indices[token_idx][k]
+            if this_expert != -1: #有效专家
+                this_expert_row = zipped_expertwise_rowmap[token_idx][this_expert]
+                unzipped_expert = unzipped_expert_idx[this_expert_row]
+                if unzipped_expert != this_expert:
+                    print(f"unzipped[{this_expert_row}] 为{unzipped_expert}, 而原始数据[{token_idx}][{k}]为{this_expert}")
+                    return 
+                else:
+                    print(f"dispatched_indices[{token_idx}][{k}] 派发到专家{this_expert}, unzipped[{this_expert_row}]结果为{unzipped_expert}")
+    print("检查结束，通过")
+
+def verify_tokens_unzip():
+    expert_num = 4
+    topk = 8
+    seqlen = 3
+    token_len = 8
+    tokens_zipped = [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3]
+    ]
+    routemap_topk = [ 
+        [-1,-1,0,1,-1,-1,-1,-1],
+        [1,-1,-1,-1,-1,-1,-1,-1],
+        [-1,0,-1,-1,-1,-1,-1,-1],
+    ]
+    probs_topk = [ 
+        [0,0,0.3,0.5,0,0,0,0],
+        [0.5,0,0,0,0,0,0,0],
+        [0,0.7,0,0,0,0,0,0],
+    ]
+    total_unzipped_tokens_num = 4
+    expected_unzipped_tokens= [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3],
+        [1,1,1,1,1,1,1,1],
+    ]
+    expected_unzipped_probs= [
+        0.3,
+        0.5,
+        0.7,
+        0.5
+    ]
+    expected_zipped_expertwise_rowmap= [
+        [0,3,-1,-1],
+        [-1,1,-1,-1],
+        [2,-1,-1,-1],
+    ]
+    expected_unzipped_expert_idx = [
+        0,
+        1,
+        0,
+        1 
+    ]
+    # TODO: 加入fp8单测
+    for prec in ['bfloat16']:
+        tokens_zipped = paddle.to_tensor(tokens_zipped, dtype=prec)
+        routemap_topk = paddle.to_tensor(routemap_topk, dtype='int32')
+        probs_topk = paddle.to_tensor(probs_topk, dtype='bfloat16')
+        expected_unzipped_probs = paddle.to_tensor(expected_unzipped_probs, dtype='bfloat16')
+        expected_zipped_expertwise_rowmap = paddle.to_tensor(expected_zipped_expertwise_rowmap, dtype='int32')
+        expected_unzipped_tokens = paddle.to_tensor(expected_unzipped_tokens, dtype=prec)
+        expected_unzipped_expert_idx = paddle.to_tensor(expected_unzipped_expert_idx, dtype='int32')
+
+        unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_expert_idx = TDU.tokens_unzip(tokens_zipped,routemap_topk, probs_topk,total_unzipped_tokens_num=total_unzipped_tokens_num, topk=topk, num_experts=expert_num)
+        
+        unzipped_tokens_stable, zipped_expertwise_rowmap_stable, unzipped_probs_stable, unzipped_expert_idx_stable= TDU.tokens_unzip_stable(tokens_zipped,routemap_topk, probs_topk, topk=topk, num_experts=2,max_tokens_per_expert=2)
+
+        check_expert_idx(unzipped_expert_idx,zipped_expertwise_rowmap)
+        check_indices(routemap_topk, zipped_expertwise_rowmap, unzipped_expert_idx)
+        # ------------------------- 前向验证 ------------------------
+        print("-------- Tokens unzipped by customed op: ------------")
+        print(unzipped_tokens)
+        print("-------- Tokens expected : ------------")
+        print(expected_unzipped_tokens)
+        print("-------- Probs unzipped by customed op: ------------")
+        print(unzipped_probs)
+        print("-------- Probs expected: ------------")
+        print(expected_unzipped_probs)
+        print("-------- zipped expertwize rowmap by customed op: ------------")
+        print(zipped_expertwise_rowmap)
+        print("-------- rowmap expected: ------------")
+        print(expected_zipped_expertwise_rowmap)
+        print("-------- expert_idx unzipped by customed op: ------------")
+        print(unzipped_expert_idx)
+        print("-------- expert_idx expected: ------------")
+        print(expected_unzipped_expert_idx)
+
+        print("-------- Tokens unzipped by stable op: ------------")
+        print(unzipped_tokens_stable)
+
+        print("-------- Probs unzipped by stable op: ------------")
+        print(unzipped_probs_stable)
+
+        print("-------- zipped expertwize rowmap by stable op: ------------")
+        print(zipped_expertwise_rowmap_stable)
+
+        print("-------- expert_idx unzipped by stable op: ------------")
+        print(unzipped_expert_idx_stable)
+
+    
+
+    
+def run():
+    verify_tokens_unzip()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_tokens_weighted_zip.py

@@ -0,0 +1,138 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def verify_tokens_unzip():
+    expert_num = 4
+    topk = 8
+    seqlen = 3
+    token_len = 8
+    tokens_zipped = [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3]
+    ]
+    routemap_topk = [ 
+        [-1,-1,0,1,-1,-1,-1,-1],
+        [1,-1,-1,-1,-1,-1,-1,-1],
+        [-1,0,-1,-1,-1,-1,-1,-1],
+    ]
+    probs_topk = [ 
+        [0,0,0.5,0.5,0,0,0,0],
+        [1,0,0,0,0,0,0,0],
+        [0,1,0,0,0,0,0,0],
+    ]
+    total_unzipped_tokens_num = 4
+    expected_unzipped_tokens= [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3],
+        [1,1,1,1,1,1,1,1],
+    ]
+    expected_unzipped_probs= [
+        0.5,
+        1,
+        1,
+        0.5
+    ]
+    expected_zipped_expertwise_rowmap= [
+        [0,3,-1,-1],
+        [-1,1,-1,-1],
+        [2,-1,-1,-1],
+    ]
+    expected_unzipped_expert_idx = [
+        0,
+        1,
+        0,
+        1 
+    ]
+    tokens_zipped = paddle.to_tensor(tokens_zipped, dtype='bfloat16')
+    routemap_topk = paddle.to_tensor(routemap_topk, dtype='int32')
+    probs_topk = paddle.to_tensor(probs_topk, dtype='bfloat16')
+    expected_unzipped_probs = paddle.to_tensor(expected_unzipped_probs, dtype='bfloat16')
+    expected_zipped_expertwise_rowmap = paddle.to_tensor(expected_zipped_expertwise_rowmap, dtype='int32')
+    expected_unzipped_tokens = paddle.to_tensor(expected_unzipped_tokens, dtype='bfloat16')
+    expected_unzipped_expert_idx = paddle.to_tensor(expected_unzipped_expert_idx, dtype='int32')
+
+    unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_expert_idx = TDU.tokens_unzip(tokens_zipped,routemap_topk, probs_topk,total_unzipped_tokens_num=total_unzipped_tokens_num, topk=topk, num_experts=expert_num)
+
+    weighted_zipped_tokens = TDU.tokens_weighted_zip(unzipped_tokens, unzipped_probs, zipped_expertwise_rowmap, total_zipped_tokens=seqlen, num_experts=expert_num)
+    # ------------------------- 前向验证 ------------------------
+    print("-------- Tokens unzipped by customed op: ------------")
+    print(unzipped_tokens)
+    print("-------- Tokens expected : ------------")
+    print(expected_unzipped_tokens)
+    print("-------- Probs unzipped by customed op: ------------")
+    print(unzipped_probs)
+    print("-------- Probs expected: ------------")
+    print(expected_unzipped_probs)
+    print("-------- zipped expertwize rowmap by customed op: ------------")
+    print(zipped_expertwise_rowmap)
+    print("-------- rowmap expected: ------------")
+    print(expected_zipped_expertwise_rowmap)
+    print("-------- expert_idx unzipped by customed op: ------------")
+    print(unzipped_expert_idx)
+    print("-------- expert_idx expected: ------------")
+    print(expected_unzipped_expert_idx)
+    print("-------- weighted zipped by customed op: ------------")
+    print(weighted_zipped_tokens)
+    print("-------- weighted zipped expected: ------------")
+    print(tokens_zipped)
+
+    
+def run():
+    verify_tokens_unzip()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_tokens_zip.py

@@ -0,0 +1,149 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def verify_tokens_unzip():
+    expert_num = 4
+    topk = 8
+    seqlen = 3
+    token_len = 8
+    tokens_zipped = [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3]
+    ]
+    routemap_topk = [ 
+        [-1,-1,0,1,-1,-1,-1,-1],
+        [1,-1,-1,-1,-1,-1,-1,-1],
+        [-1,0,-1,-1,-1,-1,-1,-1],
+    ]
+    probs_topk = [ 
+        [0,0,0.5,0.5,0,0,0,0],
+        [1,0,0,0,0,0,0,0],
+        [0,1,0,0,0,0,0,0],
+    ]
+    total_unzipped_tokens_num = 4
+    expected_unzipped_tokens= [
+        [1,1,1,1,1,1,1,1],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3],
+        [1,1,1,1,1,1,1,1],
+    ]
+    expected_unzipped_probs= [
+        0.5,
+        1,
+        1,
+        0.5
+    ]
+    expected_zipped_expertwise_rowmap= [
+        [0,3,-1,-1],
+        [-1,1,-1,-1],
+        [2,-1,-1,-1],
+    ]
+    expected_unzipped_expert_idx = [
+        0,
+        1,
+        0,
+        1 
+    ]
+    tokens_simple_zipped = [
+        [2,2,2,2,2,2,2,2],
+        [2,2,2,2,2,2,2,2],
+        [3,3,3,3,3,3,3,3]
+    ]
+    tokens_zipped = paddle.to_tensor(tokens_zipped, dtype='bfloat16')
+    routemap_topk = paddle.to_tensor(routemap_topk, dtype='int32')
+    probs_topk = paddle.to_tensor(probs_topk, dtype='float32')
+    expected_unzipped_probs = paddle.to_tensor(expected_unzipped_probs, dtype='float32')
+    expected_zipped_expertwise_rowmap = paddle.to_tensor(expected_zipped_expertwise_rowmap, dtype='int32')
+    expected_unzipped_tokens = paddle.to_tensor(expected_unzipped_tokens, dtype='bfloat16')
+    expected_unzipped_expert_idx = paddle.to_tensor(expected_unzipped_expert_idx, dtype='int32')
+
+    unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, unzipped_expert_idx = TDU.tokens_unzip(tokens_zipped,routemap_topk, probs_topk,total_unzipped_tokens_num=total_unzipped_tokens_num, topk=topk, num_experts=expert_num)
+
+    # 本算子
+    zipped_tokens, zipped_probs_topk = TDU.tokens_zip(unzipped_tokens.astype("float32"), zipped_expertwise_rowmap, routemap_topk, unzipped_probs, total_zipped_tokens=seqlen, num_experts=expert_num)
+    # ------------------------- 前向验证 ------------------------
+    print("-------- Tokens unzipped by customed op: ------------")
+    print(unzipped_tokens)
+    print("-------- Tokens expected : ------------")
+    print(expected_unzipped_tokens)
+    print("-------- Probs unzipped by customed op: ------------")
+    print(unzipped_probs)
+    print("-------- Probs expected: ------------")
+    print(expected_unzipped_probs)
+    print("-------- zipped expertwize rowmap by customed op: ------------")
+    print(zipped_expertwise_rowmap)
+    print("-------- rowmap expected: ------------")
+    print(expected_zipped_expertwise_rowmap)
+    print("-------- expert_idx unzipped by customed op: ------------")
+    print(unzipped_expert_idx)
+    print("-------- expert_idx expected: ------------")
+    print(expected_unzipped_expert_idx)
+    print("-------- zipped by customed op: ------------")
+    print(zipped_tokens)
+    print("-------- zipped expected: ------------")
+    print(tokens_simple_zipped)
+
+    print("-------- zipped probs_topk by customed op: ------------")
+    print(zipped_probs_topk)
+    print("-------- probs_topk expected: ------------")
+    print(probs_topk)
+
+    
+def run():
+    verify_tokens_unzip()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_topk_to_multihot.py

@@ -0,0 +1,111 @@
+import numpy as np
+import paddle
+import paddle.incubate.nn.functional as F
+import TokenDispatcherUtils as TDU
+
+def compare_tensors(a, b):
+    # 形状一致性检查
+    assert a.shape == b.shape, "输入张量形状不一致"
+    
+    # 计算绝对差距
+    abs_diff = np.abs(a - b)
+    max_abs_val = np.max(abs_diff)
+    max_abs_flat_idx = np.argmax(abs_diff)
+    max_abs_idx = np.unravel_index(max_abs_flat_idx, a.shape)
+    
+    # 计算相对差距（防止除以零）
+    denominator = np.maximum(np.abs(a), np.abs(b))
+    rel_diff = np.divide(
+        abs_diff, 
+        denominator, 
+        out=np.zeros_like(abs_diff),
+        where=(denominator != 0)
+    )
+    max_rel_val = np.max(rel_diff)
+    max_rel_flat_idx = np.argmax(rel_diff)
+    max_rel_idx = np.unravel_index(max_rel_flat_idx, a.shape)
+    
+
+    # 打印结果
+    print("\n[最大绝对差距]" f"位置: {max_abs_idx}")
+    print(f"a[{max_abs_idx}] = {a[max_abs_idx]:.6g}" + f"\t b[{max_abs_idx}] = {b[max_abs_idx]:.6g}" + f"\t 绝对差值: {max_abs_val:.6g}\n")
+    
+    print("[最大相对差距]" f"位置: {max_rel_idx}")
+    print(f"a[{max_rel_idx}] = {a[max_rel_idx]:.6g}" + f"\t b[{max_rel_idx}] = {b[max_rel_idx]:.6g}" + f"\t 相对差值: {max_rel_val:.6g}\n")
+    print("周围元素比较-a:")
+    print(f"{a[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    print("周围元素比较-b:")
+    print(f"{b[max_rel_idx[0], (max_rel_idx[1] - 10):(max_rel_idx[1] + 10)]} ")
+    
+    # 返回结构化结果
+    return {
+        'max_absolute': {
+            'index': max_abs_idx,
+            'a_value': a[max_abs_idx],
+            'b_value': b[max_abs_idx],
+            'difference': max_abs_val
+        },
+        'max_relative': {
+            'index': max_rel_idx,
+            'a_value': a[max_rel_idx],
+            'b_value': b[max_rel_idx],
+            'difference': max_rel_val
+        }
+    }
+
+def verify_topk_to_multihot():
+    expert_num = 4
+    topk = 8
+    seqlen = 3
+    routemap_topk = [ 
+        [-1,-1,0,1,-1,-1,-1,-1],
+        [1,-1,-1,-1,-1,-1,-1,-1],
+        [-1,0,-1,-1,-1,-1,-1,-1],
+    ]
+    probs_topk = [ 
+        [0,0,0.3,0.5,0,0,0,0],
+        [0.5,0,0,0,0,0,0,0],
+        [0,0.7,0,0,0,0,0,0],
+    ]
+    expected_routemap_multihot= [
+        [1,1,0,0],
+        [0,1,0,0],
+        [1,0,0,0]
+    ]
+    expected_probs_multihot = [
+        [0.3, 0.5, 0,0],
+        [0,0.5,0,0],
+        [0.7,0,0,0]
+    ]
+    routemap_topk = paddle.to_tensor(routemap_topk, dtype="int32")
+    expected_routemap_multihot= paddle.to_tensor(expected_routemap_multihot, dtype="int32")
+    expected_probs_multihot= paddle.to_tensor(expected_probs_multihot, dtype="float32")
+    probs_topk = paddle.to_tensor(probs_topk,dtype="float32")
+    probs_topk.stop_gradient=False;
+    route_multihot, probs_multihot = TDU.fused_topk_to_multihot(routemap_topk, probs_topk,seqlen=seqlen, topk=topk, num_experts=expert_num)
+    probs_multihot.stop_gradient=False;
+    route_multihot_only, _= TDU.fused_topk_to_multihot(routemap_topk,None, seqlen=seqlen, topk=topk, num_experts=expert_num)
+
+    # ------------------------- 前向验证 ------------------------
+    print(f"probs_multihot stop gradient: {probs_multihot.stop_gradient}")
+    print("-------- Route multihot Generated by fused op: ------------")
+    print(route_multihot)
+    print(route_multihot_only)
+    print("-------- Route expected : ------------")
+    print(expected_routemap_multihot)
+    print("-------- Probs multihot Generated by fused op: ------------")
+    print(probs_multihot)
+    print("-------- Probs multihot expected: ------------")
+    print(expected_probs_multihot)
+    # ------------------------- 反向验证 ------------------------
+    recovered = TDU.fused_multihot_prob_backto_topk(routemap_topk, probs_multihot, seqlen=seqlen, topk=topk, num_experts=expert_num)
+    print("-------- gradProbs backward Generated by fused op: ------------")
+    print(recovered)
+    print("-------- Probs backward expected: ------------")
+    print(probs_topk)
+    
+def run():
+    verify_topk_to_multihot()
+
+if __name__ == "__main__":
+    run()

tests/ops/test_unzip_zip.py

@@ -0,0 +1,142 @@
+# 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
+import TokenDispatcherUtils as TDU
+
+
+def fabricate_dispatch_result(seqlen, token_length, topk, num_experts, data_type="bfloat32", broadcast_ratio=0.5):
+    tokens = paddle.randn([seqlen, token_length], dtype=data_type)
+
+    tokens_scale = paddle.empty([0])
+    if data_type == "float8_e4m3fn":
+        scale_cols = (token_length + 127) // 128
+        tokens_scale = paddle.randn([seqlen, scale_cols], dtype="float32")
+
+    # 计算每个token选择的专家数量，更集中在期望值附近
+    expected_experts = max(1, min(broadcast_ratio * num_experts, topk))
+
+    # 使用正态分布生成专家数量，确保集中在期望值附近
+    std_dev = max(1, expected_experts / 6)  # 标准差设为期望值的1/4，确保集中分布
+    experts_count = paddle.normal(expected_experts, std_dev, [seqlen])
+    # 四舍五入并裁剪到合理范围
+    # experts_count = paddle.clip(paddle.round(experts_count), 1, min(topk, num_experts))
+    experts_count = paddle.clip(paddle.round(experts_count), 1, min(topk, num_experts))
+    experts_count = paddle.cast(experts_count, "int32")
+
+    # 预分配结果数组
+    dispatched_indices = paddle.full([seqlen, topk], -1, dtype="int32")
+    dispatched_probs = paddle.zeros([seqlen, topk], dtype="float32")
+
+    # 批量生成专家索引和概率
+    for i in range(seqlen):
+        count = experts_count[i].item()
+
+        # 高效生成随机不重复专家索引
+        indices = paddle.randperm(num_experts)[:count]
+        dispatched_indices[i, :count] = indices
+
+        # 高效设置概率值
+        prob_value = 1.0 / count
+        dispatched_probs[i, :count] = paddle.full([count], prob_value, dtype=data_type)
+
+    # 高效计算每个专家的最大token数
+    # 创建one-hot编码
+    valid_indices = dispatched_indices.reshape([-1])
+    valid_mask = valid_indices >= 0
+    valid_experts = valid_indices[valid_mask]
+
+    # 使用histogram统计每个专家的token数
+    expert_counts = paddle.histogram(valid_experts, bins=num_experts, min=0, max=num_experts - 1)
+    expert_counts = paddle.cast(expert_counts, "int32")
+    expert_counts = list(expert_counts)
+    print("expert counts: ", expert_counts)
+
+    return (
+        tokens,
+        tokens_scale,
+        dispatched_indices,
+        dispatched_probs,
+        expert_counts,
+    )
+
+
+def tensor_max_abs_rel_err(a, b, eps=1e-8):
+    max_abs_err = paddle.max(paddle.abs(a - b))
+    denom = paddle.maximum(paddle.abs(a), paddle.abs(b))
+    denom = paddle.maximum(denom, paddle.to_tensor(eps, dtype=denom.dtype))
+    max_rel_err = paddle.max(paddle.abs(a - b) / denom)
+    return max_abs_err, max_rel_err
+
+
+def test_unzip_zip():
+    SEQLEN = 16384
+    TOKEN_LEN = 7168
+    for dt in ["bfloat16"]:
+        for expert_num in [4, 8, 16, 32]:
+            for topk in [4, 8, 12]:
+                print("###################################")
+                print("testing with {} experts and topk {}, datatype is {}".format(expert_num, topk, dt))
+                (
+                    tokens,
+                    tokens_scale,
+                    dispatched_indices,
+                    dispatched_probs,
+                    expert_tokens_count,
+                ) = fabricate_dispatch_result(
+                    SEQLEN,
+                    TOKEN_LEN,
+                    topk,
+                    expert_num,
+                    data_type=dt,
+                    broadcast_ratio=0.5,
+                )
+                if dt == "bfloat16":
+                    tokens_scale = None
+                (
+                    unzipped_tokens,
+                    zipped_expertwise_rowmap,
+                    unzipped_probs,
+                    unzipped_scales,
+                ) = TDU.tokens_unzip_stable(
+                    tokens,
+                    tokens_scale,
+                    dispatched_indices,
+                    dispatched_probs,
+                    topk=topk,
+                    num_experts=expert_num,
+                    tokens_per_expert=expert_tokens_count,
+                    padding_multiplex=128,
+                    fill_output=True,
+                )
+                tokens_recovered, probs_recovered = TDU.tokens_zip(
+                    (unzipped_tokens * unzipped_probs.unsqueeze(-1)).astype("bfloat16"),
+                    zipped_expertwise_rowmap,
+                    dispatched_indices,
+                    unzipped_probs,
+                    total_zipped_tokens=SEQLEN,
+                    num_experts=expert_num,
+                )
+                print("unzip-zip tokens 最大绝对误差：{}, 相对误差：{}".format(*tensor_max_abs_rel_err(tokens, tokens_recovered)))
+                print(
+                    "unzip-zip probs 最大绝对误差：{}, 相对误差：{}".format(
+                        *tensor_max_abs_rel_err(dispatched_probs, probs_recovered)
+                    )
+                )
+
+
+# core.nvprof_enable_record_event()
+
+if __name__ == "__main__":
+    test_unzip_zip()

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 $@
+